[feat]distributed export supports quant

[eric-gecheng]

No description provided.

[github-actions]

Code Review Summary

Reviewed the distributed-embedding export + INT8 quant feature across code quality, performance, tests, docs, and multi-process safety. Overall this is thoughtful, well-structured work: the quant path has clear actionable errors (odd embedding_dim, inconsistent shard metadata), the sparse-tensor extraction is well-tested with a numeric INT8 round-trip check, and the single-rank forcing / barrier removal is internally consistent (non-zero ranks exit before joining a process group, so no missing-sync bug). Most findings are inline; a few cross-cutting notes:

Functional (worth addressing)

• export_distributed_embedding silently drops the data_input_path and additional_export_config overrides that the other two export impls honor. Details inline.

Test coverage — the largest gap

• The core of the feature (export_distributed_embedding: sharded build → FX trace → sparse/dense graph split → artifact emission) has no integration test, and the new checkpoint_util INPUT_TILE=3 remap (create_local_plan) and _make_dynamicemb_input_tile_user_view are untested despite being pure filesystem/metadata logic that's easy to cover. The existing utils.test_export(env_str=...) harness (used in rank_integration_test.py) makes a USE_DISTRIBUTED_EMBEDDING=1 (+QUANT=INT8, +dynamicemb) variant straightforward — this would run on the gpu lane and close the biggest risk. This compounds with the H20 CI pull_request trigger being disabled (flagged inline).

Maintainability / robustness (minor)

• export_distributed_embedding is ~345 lines and carries multi-rank scaffolding (is_rank_zero guards, shard-merge) that is currently unreachable because export is forced to single-rank — consider decomposing and/or noting it's forward-looking.
• init_process_group() here lacks the dist.is_initialized() guard that export_model_normal has; combined with multi-tower models (MatchModel/TDM) that call export_model more than once per process, USE_DISTRIBUTED_EMBEDDING=1 would hit a double-init crash. Matches the pre-existing RTP pattern, so likely an unsupported combo — a guard or explicit check would make it robust.

Note on docstrings: several new private helpers in export_util.py/acc/utils.py lack docstrings. AGENTS.md asks for them on private functions too, so worth a consistency pass — but I verified with ruff 0.15.4 that this is not CI-blocking (ruff's D rules don't flag leading-underscore functions; ruff check passes on all changed files).

Nice touches: _dedup_key_files_by_realpath to avoid double-counting aliased shards, the temp-dir symlink view instead of mutating a possibly read-only checkpoint dir, and thorough cross-shard JSON consistency validation.

[github-actions]

Code review — distributed export INT8 quant

Solid change overall. The cross-shard consistency checks in _merge_sharded_embedding_json, the copy.deepcopy(info) fix (previously the first shard's dict was mutated by reference), and the is_distributed_sparse_quant() fail-loud on unsupported values are all good. The embedding.py guard removal is behavior-preserving (every non-mc EC key is already classified as seq or query), and the fx_mark_seq_ec_jt docstring update matches. Two inline comments cover the main items (a DIST_QUANT/model_acc.json mismatch on the non-distributed path, and missing finite-value guarding + transient memory in the quant core). A few smaller notes:

• Quant round-trip tests are lossless-by-construction. Every quant test row uses values that are exactly the row min/max ([[-2,2],[-1,1]], [[-3,3],[-4,4]], …), which quantize to exactly {0, 255} and dequantize losslessly — so atol=5e-3 never actually exercises rounding error. Adding a case with an interior value (dim ≥ 3, asymmetric/wide range) would verify real quantization error stays bounded. Relatedly, the constant-row branch (value_range == 0 → scale defaults to 1.0, export_util.py:1893/1897) is a numerically important path that no test hits.

• _merge_sharded_embedding_json inconsistency checks: only the row_bytes mismatch is tested; the other ~7 raise branches (incl. the quant-vs-non-quant storage_shape presence mismatch) are uncovered. A small @parameterized.expand table over (field_to_corrupt, expected_regex) would close these cleanly.

• Minor: the new private helpers _remove_torch_prefix, _value_nbytes, _prepare_sparse_export_values, _normalized_distributed_sparse_quant lack docstrings. AGENTS.md asks for them on private functions too — not ruff-enforced (D103 doesn't fire on _-prefixed names), so not CI-blocking, just a convention nit.

[tiankongdeguiji]

bump up version to 1.3.2