FEA[DE]: replay_increment, writing a delta back at its source slots - #477
Merged
shijieliu merged 10 commits intoSep 10, 2026
Merged
Conversation
jiashuy
changed the title
Collaborator
Author
|
/build |
Contributor
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Collaborator
|
❌ Pipeline #65820599 -- failed
Result: 0/3 jobs passed |
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #65954897 -- failed
Result: 0/3 jobs passed |
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #65990206 -- failed
Result: 10/17 jobs passed |
…e slots incremental_dump produces a delta; there was no way to apply one. Replay is the other half of that pipeline: a training job dumps periodically, the delta is shipped, and a replica catches up without reloading a full checkpoint. Write-back is by slot. Every key goes to the slot and value row it held in the source, leaving the target layout-identical to it, which is what lets the two converge rather than merely agree on contents. A key is only ever probed inside its own home bucket, so that requires the two tables to share a layout; replay_increment compares the delta's meta (capacity, bucket_capacity, num_scores, world_size, and the table_options fields score_strategy / dim / dist_type) and raises before writing anything if they differ. score_strategy is compared by physical word order, since (TIMESTAMP, LFU) and (LFU, TIMESTAMP) are the same layout on device. The kernel enforces the same rule per key: a slot outside its key's home bucket, or held by another writer, leaves that key unplaced and the host raises rather than silently dropping it. A delta now carries the whole stored row. values is embeddings only, with the rest of the value row in optimizer_states and every score word in scores, both at the width the file checkpoint uses -- rowwise Adagrad reserves 16 bytes per row in the fused layout but fills one scalar, so dumping the runtime width would ship padding and disagree with DynamicEmbDump about what a row is. Timestamp score columns travel as an age, since %globaltimer is per-device. ReplayContent picks which of the three to write back: a serving replica wants the embedding alone, a training replica resuming from its source wants all three. Removals are split in two. A key leaves a table by being evicted to make room or by being erased explicitly, and only the second is a removal a replica has to perform -- an eviction is reproduced by whoever takes over the slot, which is in the same delta. They are retained in separate buffers, so a replay can tell them apart instead of erasing keys that were merely evicted. Retention is configured at different times for the same reason: retaining evictions swaps in a collecting insert kernel and costs a device sync per evicting insert, so it is a table option, while an erase already holds its keys and takes its own EvictedItemMode per call. EvictedItemMode is a flag so further kinds of retention can be added without revisiting call sites. The cache needs explicit invalidation. Writing a slot reproduces an eviction in the storage but not in the cache, which is a separate index the write does not reach -- and flush_cache pushes every cached key back down, so a survivor would be resurrected into storage by the next dump. Replay flushes first, then drops the delta's keys, its erased keys, and its evicted keys from the cache. MurmurHash3's finalizer had four copies. It is now one on each side of the language boundary: src/murmur_hash.cuh for device code, murmur3_fmix64 for host. The copy in test_hash_roundrobin_kuairand.py stays deliberately separate, since an oracle that calls the code under test stops being one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jiashuy
force-pushed
the
feat/dynamicemb-replay-increment
branch
from
September 4, 2026 06:30
0da9b02 to
f5af24e
Compare
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #66187052 -- failed
Result: 3/5 jobs passed |
Collaborator
Author
|
/build |
Collaborator
|
✅ Pipeline #66552826 -- success
Result: 15/17 jobs passed |
…by pg replay_increment kept only the keys a rank owns, recomputing ownership as hash(key) % world_size. It took both the rank and the world size from the process group the delta had been gathered over, which is not what the tables were sharded across. That contradicted the rest of the module. meta["world_size"] is documented as the global WORLD the source was sharded over, explicitly not the gather pg; _replay_compatibility compares it against self._shard_world_size; and owned_key_mask's own docstring says ownership follows the target's fan-out. So the compatibility check validated one number and the filter then used another. With pg a strict subgroup of WORLD -- an intra-node group over two nodes, say -- a replay filtered by 8 against a table sharded by 16, using a node-local rank. Every key was claimed by several ranks and none by the rank that owns it: duplicated writes and lost keys, silently. The existing distributed test passes intra_and_cross_node_pg()[0] but runs on a single node, where that group is WORLD, so nothing caught it. Ownership now comes from dist.get_rank() and _shard_world_size. Drop the pg argument from replay_increment at both the module and model level. Once the filter stops using it, nothing else does: replay is local -- filter by ownership, then write -- with no collective to scope. Leaving an inert parameter would invite the reading that passing a subgroup narrows the replay, which is precisely the bug. incremental_dump keeps its pg, where it really does scope the all_gather. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…w writes The module asked three layout questions of a storage -- how much capacity a table has, how wide a hash bucket is, how many score words a key carries -- and each answered by branching on hasattr(storage, "key_index_map") to tell a single-tier storage from a hybrid one. Four such branches, with the capacity one written out twice. HybridStorage already exposes its tiers as .tables; give DynamicEmbStorage the same property, returning its one state. The branches then have nothing left to decide: capacity sums over the tiers, bucket capacity and score-word count read tiers[0]. The distinction stops leaking into the caller instead of being gathered into one place. The obvious alternative -- putting these on the Storage interface -- is not available: Storage is implemented by users for external parameter servers, so adding methods to it would break them. While collecting the capacity accessors, a second bound turned out to be missing. meta["current_capacity"] is the key map's capacity, which is the modulus for choosing a home bucket and so decides whether a slot means the same thing in two tables. It is not what bounds a row write. The two are the same number everywhere except NO_EVICTION, whose key map is deliberately 1 / max_load_factor times its value buffer -- and rounding that up to whole buckets makes the map's capacity non-injective in the buffer's. With bucket_capacity=128, an init_capacity of 100 and of 128 both give a 256-slot key map, over 100 and 128 rows (measured, not derived). Those two tables passed the compatibility check, and a source row of 127 was then written past the end of a 100-row target: device memory, not a wrong answer. meta now also carries row_capacity, the value-buffer rows per tier, and replay_increment compares it. Per tier rather than summed, since slot_index routes each key to one tier and a matching total would say nothing about either. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ro its block The docstring opened with "a delta carries embeddings, not scores", which was true when it was written and stopped being true once DeltaDumpResult grew a scores column. A delta does carry the source's scores; this path runs when the caller did not ask for ReplayContent.SCORE, so they were never loaded. Stating the missing data as a property of the format rather than of the request reads as a limitation instead of the choice it is, so say which flag turns it on. Build the block with zeros rather than empty. The loop fills one column per configured strategy while the width comes from the score policy; the two agree for every strategy that exists, but they are derived from different places. A divergence would leave columns holding whatever the allocator returned, written into the table as scores -- silent and unreproducible. Zeroing costs nothing at this size and makes that outcome deterministic. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Three of the score-related tests asserted something weaker or narrower than their prose said. test_dump_scores_are_column_aligned claimed to check logical column order while configuring (TIMESTAMP, LFU) -- the compound policy's physical layout is always that, so the permutation was the identity and the test passed with score_dump_permutation deleted. Configure the reversed order instead, where the frequency lands in column 0 only if the permutation runs. Its frequency check was an ordering (hot outranks cold), which still holds if every count is off by the same amount; assert the counts, one and five. Its timestamp column was not checked at all, so add the one thing an age can be checked against without a clock: the hot keys were touched last, so theirs is the smaller. test_replay_scores_follow_the_content_flag inferred the replica's scores from which keys a threshold selected. Read the scores instead. That makes the ReplayContent.ALL case exact -- five and one, carried verbatim -- where before it only had to split the same way, and lets the ReplayContent.EMBEDDING case say what it means, that every restored key scores alike, rather than the roundabout "all of them or none". The parameter is keeps_source_scores now: both cases have scores, and which they have is the question. test_replay_without_embedding_needs_aligned_rows is deleted. Replaying scores without embeddings is not a case that arises, and the test only checked that the guard fires, not that the embeddings it guards are preserved. Also fixes two comments that said the hot keys were accessed six times. They are accessed five. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
/build |
Collaborator
|
✅ Pipeline #66690895 -- success
Result: 15/17 jobs passed |
…ntal_load --incremental_dump trained a model and printed how many keys each dump matched. It showed that incremental_dump returns something; it did not show what that something is for. It now writes each dump to its own file, and --incremental_load replays them into a second model and reports its loss. Two commands over one --save_dir, so the files are the whole handoff -- which is the shape delta replication has in practice, a training job and a serving replica that never share a process. The replica is built with training=False, so its value rows hold embeddings alone. A delta from a training model still carries optimizer state; replay writes what the target has room for and drops the rest. Two things beyond the deltas turned out to be load-bearing. The dense weights. Only embeddings travel in a delta, so without the trained MLP the replica's loss is a fresh model's -- 17.1 against 2.0 here, which reads like a broken replay and is not one. Embedding weights are not in state_dict, so saving the model's state_dict is exactly the part a delta does not carry. A dump after the training loop, not only the periodic ones. Training carries on past the last periodic dump and those updates have to ship too; without it the replica is out of date, and the forward disagrees for a reason that has nothing to do with replay. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jiashuy
force-pushed
the
feat/dynamicemb-replay-increment
branch
from
September 8, 2026 06:11
c039ed7 to
e6bb456
Compare
…ough Two ways a replay could raise after it had already changed the table. A key was installed at its slot by scatter_keys_at_slots, and only then checked against same_key to see whether the row it landed on had been its own. Omitting ReplayContent.EMBEDDING left nothing to write into that row, so a key that did not already own it would go on serving the previous occupant's vector -- exactly what the check existed to prevent, raised one step too late to prevent it. The check is gone, and so is the option it guarded. There is no way to write a key at a slot without deciding what its row holds, so the embedding always travels with the key and ReplayContent chooses only what comes along: OPTIMIZER_STATE, SCORE, or the empty EMBEDDING_ONLY. That removes the ordering problem rather than paying an extra lookup to fix it, and it is honest about a flag that could not be turned off safely anyway. Compatibility was checked inside the per-table loop, so a delta spanning a collection -- which is the normal shape -- could erase from and write to its first tables and then raise on a later one's metadata. The docstring and DynamicEmb_APIs both promise a mismatch raises "before anything is written"; that held for one table and not for several. Validate every table first, then write. The new test pins the second: two tables, the second incompatible, and the first must come out untouched. Confirmed by re-interleaving validation with writes and watching it fail -- a first attempt at that check was itself wrong, letting the mismatched table raise in the validating pass so that nothing was written and the test passed against the broken code. Reported in review. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
… run inc_dump kept the delta directory across runs while restarting its filename counter at zero. A shorter run over a directory an earlier one had left behind -- a dump that was interrupted, or produced and never loaded -- overwrote the low numbers and left the high ones in place. inc_load replays every file it matches, in order, so those stragglers landed last and wrote a previous run's embeddings over the current one's. Nothing downstream could catch it. Two runs of the same script share a config, so the stale slot indices address the same layout: replay_increment's compatibility check passes, the kernel's home-bucket check passes, and the writes apply cleanly. They are simply the wrong values. A dump run now owns its rank's sequence outright and removes the files matching it before writing. Per rank rather than the whole directory, since each rank already owns a disjoint set of filenames and clearing them separately needs no coordination between ranks. Verified by seeding delta_rank0_0003..0007 and running a dump that produces two files: the directory comes out holding 0000 and 0001 alone, and the load that follows reports the same loss as a run over a clean directory. Reported in review. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
replay_increment had grown into one long loop that resolved, validated and wrote each table in the same pass. Split it: _plan_replay resolves and checks every table and writes nothing, _apply_replay writes one already-checked table. "No table is written until every table has been checked" is now a property of the control flow rather than something to be read out of a loop, and flush_cache moves after planning so a rejected delta leaves the model exactly as it was. The job carries everything its write needs -- erased_keys, evicted_keys and the source's current_score -- so apply never reaches back into the delta and the second subscript into the delta's column-aligned lists is gone. Close one more way a delta could apply partway. Preflight validated only row counts, so a wrong score width, a wrong optimizer-state width, a short slot_index or a delta carrying no optimizer state at all (SGD dumped, rowwise Adagrad replayed -- the layout check does not compare optimizers) all raised down in the write path, by which time the collection's earlier tables were in. Every shape the write path depends on is checked up front instead. Docs: the design doc still said the delta carries embeddings only and never optimizer state, and listed five tests that no longer exist while missing five that do. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #66736375 -- failed
Result: 12/17 jobs passed |
shijieliu
reviewed
Sep 10, 2026
shijieliu
reviewed
Sep 10, 2026
shijieliu
approved these changes
Sep 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
incremental_dump produces a delta; there was no way to apply one. Replay is the other half of that pipeline: a training job dumps periodically, the delta is shipped, and a replica catches up without reloading a full checkpoint.
Write-back is by slot. Every key goes to the slot and value row it held in the source, leaving the target layout-identical to it, which is what lets the two converge rather than merely agree on contents. A key is only ever probed inside its own home bucket, so that requires the two tables to share a layout; replay_increment compares the delta's meta (capacity, bucket_capacity, num_scores, world_size, and the table_options fields score_strategy / dim / dist_type) and raises before writing anything if they differ. score_strategy is compared by physical word order, since (TIMESTAMP, LFU) and (LFU, TIMESTAMP) are the same layout on device. The kernel enforces the same rule per key: a slot outside its key's home bucket, or held by another writer, leaves that key unplaced and the host raises rather than silently dropping it.
A delta now carries the whole stored row. values is embeddings only, with the rest of the value row in optimizer_states and every score word in scores, both at the width the file checkpoint uses -- rowwise Adagrad reserves 16 bytes per row in the fused layout but fills one scalar, so dumping the runtime width would ship padding and disagree with DynamicEmbDump about what a row is. Timestamp score columns travel as an age, since %globaltimer is per-device. ReplayContent picks which of the three to write back: a serving replica wants the embedding alone, a training replica resuming from its source wants all three.
Removals are split in two. A key leaves a table by being evicted to make room or by being erased explicitly, and only the second is a removal a replica has to perform -- an eviction is reproduced by whoever takes over the slot, which is in the same delta. They are retained in separate buffers, so a replay can tell them apart instead of erasing keys that were merely evicted. Retention is configured at different times for the same reason: retaining evictions swaps in a collecting insert kernel and costs a device sync per evicting insert, so it is a table option, while an erase already holds its keys and takes its own EvictedItemMode per call. EvictedItemMode is a flag so further kinds of retention can be added without revisiting call sites.
The cache needs explicit invalidation. Writing a slot reproduces an eviction in the storage but not in the cache, which is a separate index the write does not reach -- and flush_cache pushes every cached key back down, so a survivor would be resurrected into storage by the next dump. Replay flushes first, then drops the delta's keys, its erased keys, and its evicted keys from the cache.
MurmurHash3's finalizer had four copies. It is now one on each side of the language boundary: src/murmur_hash.cuh for device code, murmur3_fmix64 for host. The copy in test_hash_roundrobin_kuairand.py stays deliberately separate, since an oracle that calls the code under test stops being one.
Description
Checklist