[feat] delta dump supports zch - #627
Conversation
… upload ZCH tables remap raw feature ids to compact local rows before lookup, so the delta dump previously rejected them: the tracker would have written remapped rows as serving keys. The dump now resolves each tracked local row back to its raw id through the owning MCHManagedCollisionModule (inverting _mch_remapped_ids_mapping back to ZCH slots and reading _mch_sorted_raw_ids), drops the fallback row and empty slots, and emits raw ids as key_id with the untouched row as embedding, matching the raw-id contract the distributed export already uses for ZCH tables. torchrec fires post_lookup_tracker_fn only from the inner collection's compute_and_output_dist, while the sharded managed collision wrapper drives it through plain compute(), so ZCH lookups never reached the tracker; the dumper wraps the inner compute to record remapped ids. Co-Authored-By: Claude <noreply@anthropic.com>
…action torchrec fires post_lookup_tracker_fn and post_odist_tracker_fn only from the inner collection's compute_and_output_dist, but a sharded managed collision wrapper drives it through plain compute(), so the ZCH tracking wrapper added earlier fired only the lookup callback: trigger_compaction never ran, curr_compact_index stayed at 0, and record_lookup's raw per-batch ids accumulated O(steps x batch ids) until the dump. With dump_interval_steps=1000 or time-based dumping this is an OOM risk and silently defeats the auto_compact behavior. tracked_compute now fires post_odist_tracker_fn after each per-sharding record, mirroring torchrec's compute_and_output_dist pairing; the trigger_compaction guard makes repeated calls per batch idempotent. Co-Authored-By: Claude <noreply@anthropic.com>
_install_zch_tracking wraps the whole inner collection's compute, but _lookup_embeddings only routed to the ZCH path for FQNs present in _zch_modules, so a tracked table of a ZCH inner module missing from that dict would fall through to the plain path and emit its remapped local rows as serving key_ids, silently corrupting the dump. Resolve the mapping per tracked table of every ZCH inner module and raise if a tracked table has no MCH module, instead of relying on the two dicts staying consistent by construction. Co-Authored-By: Claude <noreply@anthropic.com>
… dump The ZCH delta path only published rows the tracker recorded during the window, which misses two lifecycle events: (1) an id admitted into the ZCH table after its only lookup was previously published never (pre- admission lookups hit the fallback row, which the dump drops), leaving its trained row unpublished until a later lookup; and (2) an evicted id keeps its stale vector in FeatureStore forever because MERGE has no delete while the model now serves that id the fallback row. Each dump now diffs the current (raw id -> row) set of every ZCH table against the previous snapshot: ids added since the snapshot are published with their trained rows, and ids dropped are republished with the fallback row the model serves for them. The snapshot resets after checkpoint restore so the restored baseline is not republished. Co-Authored-By: Claude <noreply@anthropic.com>
…ables The slot inverse used torch.empty, so rows no slot mapped to read garbage slot indices instead of failing, and a remapped row outside the local shard range (e.g. a resharded checkpoint) hit an opaque device-side index assert. _zch_slot_of_row now fills with -1 and checks both conditions, raising with the table fqn and offending rows. The non-MCH managed collision type check fired at construction for every table in every collision collection, killing training even when the dump did not track that table; like the export path it now warns and skips, while a tracked table still fails via the existing resolution check. The compute wrapper installed for ZCH tracking is restored in close() so it stops capturing the dumper once it is closed. Co-Authored-By: Claude <noreply@anthropic.com>
Drop a duplicated empty-ids guard and a dead tuple-type branch, hoist a per-loop feature_name lookup, and simplify the ZCH test helpers.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| raw_ids = zch_module._mch_sorted_raw_ids.to(device) | ||
| rows = zch_module._mch_remapped_ids_mapping.to(device) | ||
| valid_mask = raw_ids != torch.iinfo(torch.int64).max | ||
| return raw_ids[valid_mask], rows[valid_mask] - zch_module._output_global_offset |
There was a problem hiding this comment.
P1: Please validate these checkpoint-derived rows against the actual local weight shard before returning them. This lifecycle path bypasses _zch_slot_of_row(): a negative row silently indexes from the end at line 1089, while an oversized row can trigger a device-side assert on one rank. The export path already performs this bounds check.
There was a problem hiding this comment.
Valid — fixed in a43b8b2. _zch_current_rows now validates the remapped rows against the local shard row range and raises, matching the contract _zch_slot_of_row and the export path (_zch_table_to_dynamic) already enforce, so an out-of-range row fails loudly instead of indexing the weight from the end or hitting a device-side assert. Covered by test_zch_current_rows_fails_on_row_outside_local_shard.
| if touched is not None | ||
| else torch.empty(0, dtype=torch.long, device=device) | ||
| ) | ||
| publish_mask = ~torch.isin(raw_ids, prev) |
There was a problem hiding this comment.
P1: This set-only snapshot misses an ID that is evicted and re-admitted within one dump interval: it is present in both prev and raw_ids, and admission can occur after its last fallback-row lookup, so it may also be absent from touched_raw_ids. FeatureStore then keeps the pre-eviction vector. Please track a mapping/lifecycle generation (not only membership) and cover this cycle.
There was a problem hiding this comment.
Valid — fixed in a43b8b2. The snapshot now stores the (raw id, local row) binding rather than id membership alone, and the diff republishes any id whose bound row changed since the last dump. Re-admission always lands an id on a new row (its old slot was reassigned and the row reset via the eviction init_fn), so an evict→re-admit cycle is republished with its trained row. Covered by test_zch_delta_republishes_readmitted_id_with_new_row.
| feature_name=feature_name, | ||
| table_fqn=fqn, | ||
| key_ids=evicted, | ||
| embeddings=weight[mch._zch_size - 1].expand( |
There was a problem hiding this comment.
P1: This copies the fallback row only when a key is evicted, but that row continues training on later unmatched lookups while fallback-row tracker hits are dropped above. The explicit FeatureStore value therefore becomes stale after the next fallback update. This needs a delete/default-row update contract or a way to refresh previously evicted keys; otherwise serving diverges from the model.
There was a problem hiding this comment.
The divergence is real, but I don't think either proposed remedy is viable here. The FeatureStore writer is MERGE-only (FEATURE_STORE_WRITE_MODE = "MERGE"), so there is no delete/tombstone to emit, and refreshing previously evicted keys would mean republishing an unbounded, ever-growing set on every dump. Publishing the fallback-row snapshot at eviction is still strictly better than the pre-PR behavior, which kept the last-trained vector forever even though its row had been reset and reused. The re-admission fix (binding diff) also bounds the drift: once an id is re-admitted, its trained row is republished. I've documented the residual drift for permanently evicted ids in the _append_zch_delta_rows docstring.
| """ | ||
| num_rows = 0 | ||
| for fqn, mch in self._zch_modules.items(): | ||
| raw_ids, rows = self._zch_current_rows(mch, fqn, table_weights) |
There was a problem hiding this comment.
P2: Every dump now scans the full ZCH capacity, moves the previous snapshot CPU→GPU, runs table-wide membership checks, and copies the full resident set back to CPU. The documented/configured examples use zch_size: 1000000, so a small delta can still stall or OOM training. Can admission/eviction changes be tracked incrementally (or at least processed in bounded chunks, with assume_unique=True for these sets)?
There was a problem hiding this comment.
Partially addressed in a43b8b2. Both isin calls now pass assume_unique=True (the id sets are provably unique — one raw id per occupied slot), which skips the internal dedup/sort, and the publish diff uses searchsorted against the sorted snapshot. That said, the OOM framing overstates it: the per-dump working set is bounded by zch_size int64 indices (~8 MB at zch_size=1e6), so it's a few ms of GPU set ops, not a memory risk. Incremental admission/eviction tracking isn't available — MCH exposes no admission/eviction callback to hook. Happy to revisit if profiling shows this as a hot path.
| key_ids = local_ids + table_weight.shard_info.row_offset | ||
| return weight[local_ids].detach(), key_ids | ||
|
|
||
| def _lookup_zch_embeddings( |
There was a problem hiding this comment.
P2: Repository policy requires Google-style docstrings for private methods. Please document that ids are remapped local rows, that the returned keys are raw ZCH feature IDs, and the mapping-related exceptions this helper can raise.
There was a problem hiding this comment.
Valid — fixed in a43b8b2. Added a Google-style docstring to _lookup_zch_embeddings documenting that ids are remapped local rows, that the returned keys are raw ZCH feature ids (and why serving keys on raw id), the fallback-row exclusion, and the KeyError/ValueError the helper can raise.
|
Static review complete; no tests or builds were run. I left five inline comments: three P1 consistency/corruption risks around ZCH lifecycle state, fallback refresh, and row validation; plus P2 feedback on full-capacity dump cost and the required helper documentation. The raw-ID inversion and admission/eviction coverage are a solid foundation, but the P1 cases should be resolved before merge. |
The zch snapshot diff tracked only raw id membership, so an id evicted and re-admitted between two dumps stayed in both snapshots and was skipped even though re-admission moved it onto a reset row, leaving its pre-eviction vector in the FeatureStore. Snapshot the raw id to row binding and republish any id whose bound row changed. The lifecycle path also indexed the weight with checkpoint-derived rows without the bound check the lookup and export paths already apply, so add it there. Also document the private zch lookup helper and pass assume_unique to isin on the provably-unique id sets. Co-Authored-By: Claude <noreply@anthropic.com>
1 participant
No description provided.