[feat] add model delta tracker

[eric-gecheng]

No description provided.

[github-actions]

Code Review: delta embedding dump

Reviewed across correctness, performance, tests, docs, and multi-process safety. This is a well-structured and unusually thorough PR — the step-boundary logic (maybe_dump/final_dump/dump + _tracker.step() ordering, the i_step >= 0 guard, the boundary-skip to avoid overwriting real shards with empty ones) is correct and backed by targeted tests, and the per-rank empty-shard design to keep shard sets uniform is a nice touch. The dump-time work (model.modules() walks, .cpu()/.item(), parquet build) is correctly confined behind the dump interval, so per-step training overhead is just tracker.step() with auto_compact=True — the right call.

A few items worth considering (3 inline comments above + the below):

Tests — coverage of the validation surface, shard-info collection, row-wise key-id conversion, and the dynamicemb path is strong. Two pure-Python branches are only exercised via GPU-gated integration tests and are easy to unit-test directly:

• _local_table_weight — its 4 dispatch branches and both raise paths (single-shard assertion, unsupported-type TypeError) have no direct coverage.
• _lookup_embeddings — the out-of-range valid_mask filtering (silent row dropping) and the empty-ids early return are untested; only the all-in-range happy path is covered.

Docs — docs/source/usage/train.md enumerates the other train_config options but has no entry for the new delta_embedding_dump_config. Worth a short subsection documenting the four fields and the enforced constraints (CUDA-only, MC/ZCH unsupported → use dynamicemb, column-wise sharding unsupported, one parquet shard per rank), ideally cross-referenced with the dynamicemb doc.

Minor / nice-to-have

• The shard-info precedence machinery (_merge_shard_info / _merge_table_shard_info / _shard_info_quality) is correct but the hardest part of the file to follow — a one-line comment stating the intended precedence (shard metadata > sharding-plan placement > config attrs > tensor shape) would help future maintainers.
• _lookup_dynamic_embeddings couples to several private dynamicemb internals (_dynamicemb_options[...].dim, the 8-tuple tables.find() return unpacked positionally); a brief comment on the discarded slots would harden it against upstream changes.

[github-actions]

Code Review Summary — delta embedding dump

Reviewed by five focused passes (quality, performance, tests, docs, multi-process safety). This is a well-engineered feature: the atomic temp-file + os.replace write path, per-rank empty shards to keep step dirs uniform, the reentrant pause_tracking guard around eval, and clear() after restore are all handled carefully, and the explanatory comments genuinely document the why behind the subtle distributed decisions. Tests are substantial.

The findings worth acting on are posted inline. The most important is a possible NCCL hang in final_dump's collective all_reduce when a rank consumes zero training batches (main.py:541) — please look at that one first. Two other inline comments flag a related lagging-rank data-loss/ragged-shard edge case and an untested silent-drop path.

A few additional items, in priority order — none blocking:

• Missing user docs (medium). Every other train_config field is documented in docs/source/usage/train.md (incl. recent additions like save_checkpoints_timestamp_*, keep_checkpoint_max), but delta_embedding_dump_config has no entry anywhere under docs/. Worth a short section covering the four fields, the CUDA-only / no-MC-ZCH / no-column-wise constraints, and the single- vs multi-GPU output layout & parquet schema.
• Integration coverage gap (medium). The main.py lifecycle wiring (pause_tracking during eval, clear after restore, maybe_dump/final_dump) is only exercised by the GPU + dynamicemb e2e test. The common case — a FUSED/EBC row-wise table through test_train_eval with dump enabled — isn't covered; the 2-GPU test calls dump() directly and bypasses main.py. A FUSED-table e2e (GPU-gated, no dynamicemb) would cover the real lifecycle.
• Recompute per dump (low, easy win). _collect_table_shard_infos() (and the model.modules() walk + placement regex it drives) runs on every dump() via _collect_table_weights, but shard topology is fixed after construction. Caching it in __init__ (where _validate_supported_table_sharding already computes it once) removes a full module-tree walk per dump.
• Fragile dynamicemb coupling (low). _lookup_dynamic_embeddings (L577–580) positionally unpacks an 8-tuple from tables.find(...) and reads private _dynamicemb_options[table_id].dim; a dynamicemb version bump could break this silently. A length check + a comment pinning the expected API would harden it.

Nice work overall.