From f5af24e59f7877142d00200688e763f820505bdc Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Wed, 2 Sep 2026 15:02:54 +0530 Subject: [PATCH 1/9] feat(dynamicemb): replay_increment, writing a delta back at its source 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) --- corelib/dynamicemb/DynamicEmb_APIs.md | 136 ++- .../docs/replay_increment_design.md | 400 +++++++ corelib/dynamicemb/dynamicemb/__init__.py | 13 +- .../dynamicemb/batched_dynamicemb_function.py | 2 +- .../dynamicemb/batched_dynamicemb_tables.py | 479 +++++++- .../dynamicemb/dynamicemb_config.py | 86 +- .../dynamicemb/dynamicemb/incremental_dump.py | 333 +++++- .../dynamicemb/dynamicemb/key_value_table.py | 1008 +++++++++++++++-- .../dynamicemb/dynamicemb/scored_hashtable.py | 104 +- corelib/dynamicemb/dynamicemb/types.py | 35 + corelib/dynamicemb/src/murmur_hash.cuh | 43 + .../src/sparse_block_bucketize_features.cu | 8 +- .../dynamicemb/src/table_operation/insert.cu | 59 + .../src/table_operation/kernels.cuh | 98 ++ .../dynamicemb/src/table_operation/table.cu | 9 + .../dynamicemb/src/table_operation/table.cuh | 10 + .../dynamicemb/src/table_operation/types.cuh | 12 +- .../test_distributed_dynamicemb.py | 109 +- .../incremental_dump/test_incremental_dump.sh | 1 + .../incremental_dump/test_replay_increment.py | 686 +++++++++++ .../test_insert_collect_evicted.py | 61 + .../test_hash_roundrobin_kuairand.py | 7 + .../unit_tests/test_hybrid_storage_export.py | 13 +- 23 files changed, 3479 insertions(+), 233 deletions(-) create mode 100644 corelib/dynamicemb/docs/replay_increment_design.md create mode 100644 corelib/dynamicemb/src/murmur_hash.cuh create mode 100644 corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py diff --git a/corelib/dynamicemb/DynamicEmb_APIs.md b/corelib/dynamicemb/DynamicEmb_APIs.md index f9fd1ac21..9e4ee324d 100644 --- a/corelib/dynamicemb/DynamicEmb_APIs.md +++ b/corelib/dynamicemb/DynamicEmb_APIs.md @@ -17,7 +17,9 @@ This document consists of two parts, one is the introduction to the API, which c - [DynamicEmbDump](#dynamicembdump) - [DynamicEmbLoad](#dynamicembload) - [incremental_dump](#incremental_dump) +- [replay_increment](#replay_increment) - [pop_evicted_keys](#pop_evicted_keys) +- [pop_erased_keys](#pop_erased_keys) - [get_score](#get_score) - [set_score](#set_score) - [Counter](#counter) @@ -787,7 +789,28 @@ The meaning of the threshold depends on the table's `score_strategy`: - ``table_names: List[str]`` -- the dumped table names. - ``keys: List[torch.Tensor]`` -- per-table matched keys on host. - - ``values: List[torch.Tensor]`` -- per-table matched values on host. + - ``values: List[torch.Tensor]`` -- per-table `[N, dim]` embeddings on + host. **Embeddings only**: the rest of the stored row is in + `optimizer_states`, and concatenating the two along dim 1 + reproduces the row as the table holds it. + - ``optimizer_states: List[Optional[torch.Tensor]]`` -- per-table + optimizer state on host, at the **same width the file checkpoint + uses** — narrower than the runtime row for rowwise Adagrad, whose + fused layout reserves 16 bytes per row but fills one scalar. + ``None`` for a table whose optimizer keeps no per-row state, e.g. + plain SGD. + - ``scores: List[torch.Tensor]`` -- per-table `[N, num_scores]` score + words on host, in the table's configured (logical) column order. + Timestamp columns hold an **age** (`current_score - score`), not a + raw timestamp, since `%globaltimer` is per device and resets across + boots; a consumer rebases them onto its own clock. Every other + column (LFU frequency, STEP, CUSTOMIZED, NO_EVICTION row) is + carried verbatim. + - ``erased_keys: List[Optional[torch.Tensor]]`` -- per-table keys an + explicit erase removed, for tables whose ``evicted_item_mode`` + erase asked to have recorded (each `erase` call passes its own + `EvictedItemMode`), so always present and possibly empty. These + are the removals `replay_increment` applies. - ``evicted_keys: List[Optional[torch.Tensor]]`` -- per-table retained evicted keys on host for tables with ``evicted_item_mode=RETAIN_KEY``, else ``None``. Returning them drains that table's retained-evicted @@ -800,6 +823,9 @@ The meaning of the threshold depends on the table's `score_strategy`: ``keys``; the storage slot each dumped key occupies (for ``replay_increment``). - ``"current_capacity": int`` -- the table's current capacity. + - ``"bucket_capacity": int`` -- slots per hash bucket. + - ``"num_scores": int`` -- score words per key; part of the + slot layout ``replay_increment`` compares against. - ``"world_size": int`` -- ranks the table was sharded across. - ``"table_options": DynamicEmbTableOptions`` -- the table config. """ @@ -807,6 +833,78 @@ The meaning of the threshold depends on the table's `score_strategy`: More usage please see [test](https://github.com/NVIDIA/recsys-examples/blob/main/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py) +## replay_increment + +**Background** +`incremental_dump` produces a delta; `replay_increment` is the other half of that pipeline — it writes the delta back into a model. The typical deployment is delta replication: a training job dumps periodically, the delta is shipped to a serving replica, and the replica replays it to catch up without reloading a full checkpoint. + +**Behavior** +For every table in the delta, `replay_increment` erases the table's `erased_keys` (so the target converges to the source) and then upserts `keys` / `values` at the slots in `meta["slot_index"]`. `evicted_keys` is never applied: the key that took an evicted key's slot is in the same delta and overwrites it. + +*Write-back is by slot.* Every key is written at the slot and value row it held in the source table, leaving the target layout-identical to it. A key can only be found inside its own home bucket, and that bucket is `hash(key) % capacity / bucket_capacity`, so this requires the target's layout to match the source's. `replay_increment` compares the delta's `meta` (`current_capacity`, `bucket_capacity`, `num_scores`, `world_size`, and the `table_options` fields `score_strategy` / `dim` / `dist_type`) against the target table, and a mismatch raises `ValueError` naming the first mismatching field, **before anything is written**. Configure the target to match the source, or rebuild it from a full checkpoint (`DynamicEmbLoad`) instead. + +The same rule is enforced per key inside the kernel — a slot that does not land in its key's home bucket raises rather than dropping the key. + +> **Precondition:** writing a key at its source slot **overwrites whatever occupies that slot** in the target. That is what makes a replica converge — if the source evicted key `B` to make room for `A`, a same-capacity replica must do the same. It also means a replayed table must be built *only* by loading/replaying from its source: a table that also takes independent writes can lose a key whose slot a delta key claims. + +*Scores and optimizer state are dumped but not yet replayed.* `DeltaDumpResult.scores` and `.optimizer_states` carry them, but `replay_increment` ignores both for now: a restored key is scored as if it had just been inserted into the target, and keeps its optimizer state only if it already occupies the target row. So the replica ranks its own future evictions by when it received each key rather than by how the source ranked it. Embeddings are unaffected — the two models can evict in different orders, but never disagree on the value of a key they both hold. NO_EVICTION is the exception and is exact: its score word is a value row, not a score. + +*Sharding.* Replay keeps only the keys this rank owns, recomputing ownership from the key with **this** model's world size. A delta gathered over a process group (`incremental_dump(..., pg)`) holds the whole group's keys and so fans out correctly when the same delta is handed to every rank. A per-rank delta (`pg=None`) holds only the producing rank's keys and should be replayed there — replaying it on another rank is not an error, every key simply belongs to someone else and is skipped, which `ReplayStats.skipped` reports. Like `incremental_dump`, only `roundrobin` and `hash_roundrobin` are supported; `continuous` raises `NotImplementedError`. + +*Optimizer state* is not part of a delta. A key that already occupies its target row keeps its optimizer state; a row taken over from another key (or a brand-new one) is reset to the table's initial optimizer state. + + ```python + #How to import + from dynamicemb import replay_increment + + #API arguments + def replay_increment( + model: torch.nn.Module, + deltas: Dict[str, "DeltaDumpResult"], + pg: Optional[dist.ProcessGroup] = None, + ) -> Dict[str, Dict[str, "ReplayStats"]]: + """Write incremental_dump results back into a model's dynamic embedding tables. + + Args: + model(nn.Module): The model containing dynamic embedding tables. + deltas(Dict[str, DeltaDumpResult]): `incremental_dump`'s return value, keyed by embedding-collection path. Collections or tables the model does not have are skipped with a warning. + pg(Optional[dist.ProcessGroup]): optional. The process group defining this model's shard fan-out. Defaults to the world the tables were created against. + + Returns + ------- + Dict[str, Dict[str, ReplayStats]]: + `{collection_path: {table_name: ReplayStats}}`, where `ReplayStats` has: + + - `upserted: int` -- keys written back at their source slot. + - `erased: int` -- keys actually removed before the upsert, counted as really removed (a delta can name keys this replica never held). + - `skipped: int` -- delta keys this rank does not own. + + Raises + ------ + ValueError: a target table's layout does not match the source's, or a delta is missing the per-key data replay needs. + NotImplementedError: a table is sharded with `dist_type="continuous"`. + """ + ``` + +Example — replicate a training model's deltas into a serving model: + + ```python + from dynamicemb import replay_increment + from dynamicemb.incremental_dump import get_score, incremental_dump + + threshold = get_score(train_model) # reference point for the next window + ... # train for a while + deltas = incremental_dump(train_model, threshold, pg) + threshold = {c: {r.table_names[i]: r.meta[i]["current_score"] + for i in range(len(r.table_names))} + for c, r in deltas.items()} # threshold for the NEXT dump + + stats = replay_increment(serve_model, deltas, pg) # raises if layouts differ + for collection, per_table in stats.items(): + for name, s in per_table.items(): + print(collection, name, s.upserted, "keys replayed at their source slot") + ``` + ## pop_evicted_keys **Background** @@ -815,6 +913,8 @@ When a table's last-tier storage is full, evicting a key drops it from the syste **Behavior** Returns, per table, the keys evicted since the previous call, deduplicated within a table. This is a read-and-clear (incremental) operation: each evicted key is reported exactly once across successive calls, and returning a table's keys drains its retained-evicted buffer on this rank. Tables without `evicted_item_mode=RETAIN_KEY` are omitted from the result. +Keys removed by an *explicit erase* are not here — they live in a separate buffer read by [pop_erased_keys](#pop_erased_keys). The two are kept apart because a consumer usually wants different things from them: an eviction says the table ran out of room, an erase says someone asked for the key to go. + ```python #How to import from dynamicemb import pop_evicted_keys @@ -843,9 +943,39 @@ Returns, per table, the keys evicted since the previous call, deduplicated withi Returns ------- Dict[str, Dict[str, torch.Tensor]]: - {collection_path: {table_name: keys}} where keys is a 1-D int64 host + {collection_path: {table_name: keys}} where keys is a 1-D host tensor of table-unique evicted keys. Empty dict if the model has no - retain-enabled dynamic embedding tables. + tables retaining evictions. + """ + ``` + +## pop_erased_keys + +**Background** +The counterpart of [pop_evicted_keys](#pop_evicted_keys) for the other way a key leaves a table: an explicit removal, rather than the table running out of room. + +Unlike evictions, this is **not** a table setting. Retaining evictions has to be decided up front — it swaps in a collecting insert kernel and costs an extra kernel output plus a device sync on every evicting insert. An erase already holds its keys, so recording them costs a copy and nothing has to be prepared; each `erase` call therefore passes its own `EvictedItemMode` and decides for itself. + +**Behavior** +Same arguments and the same read-and-clear semantics as `pop_evicted_keys`, draining the erased-key buffer instead. Tables whose mode retains no erases are omitted. + +Only these keys are removals a replica has to perform for itself, which is why `replay_increment` applies `erased_keys` and never `evicted_keys`. + + ```python + #How to import + from dynamicemb import pop_erased_keys + + #API arguments + def pop_erased_keys( + model: torch.nn.Module, + table_names: Optional[Dict[str, List[str]]] = None, + pg: Optional[dist.ProcessGroup] = None, + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Return (and clear) the keys an explicit erase removed, per table. + + Every table is included -- whether an erase was recorded is that + erase call's decision, not the table's, so there is nothing to + filter on. Arguments are identical to pop_evicted_keys. """ ``` diff --git a/corelib/dynamicemb/docs/replay_increment_design.md b/corelib/dynamicemb/docs/replay_increment_design.md new file mode 100644 index 000000000..177c3ed59 --- /dev/null +++ b/corelib/dynamicemb/docs/replay_increment_design.md @@ -0,0 +1,400 @@ +# `replay_increment` · Design & Change Document + +> Branch: `feat/dynamicemb-replay-increment` (on top of the merged +> `DeltaDumpResult` refactor, `docs/delta_dump_result_design.md`). +> Goal: write a `DeltaDumpResult` produced by `incremental_dump` back into a +> dynamic embedding model — the "replay" half of the delta pipeline. +> `docs/delta_dump_result_design.md` §8 left this to "the next round"; this is +> that round. + +--- + +## 1. Background + +`incremental_dump` already returns, per embedding collection, a +`DeltaDumpResult` whose per-table `meta` was designed for this feature: + +| `meta[i]` field | why replay needs it | +|---|---| +| `slot_index` | the exact key slot / value row each dumped key occupied | +| `current_capacity` | a source slot is only transferable when the target's capacity matches | +| `world_size` | the source's row-wise fan-out (key → rank modulo base) | +| `table_options` | `bucket_capacity` / `score_strategy` / `dim` / `dist_type` compatibility | +| `current_score` | the table score after the dump | + +Everything replay needs is already there; §3 records what deliberately is *not*. + +--- + +## 2. Confirmed Design Decisions + +| Decision | Conclusion | +|---|---| +| Write-back | **By slot.** Every key is written at the slot and value row it held in the source table, leaving the target layout-identical to it | +| Layout mismatch | **Raise**, before writing anything. Checked at table level (delta `meta` vs target) and enforced again per key in the kernel | +| Per-key scores | **Dumped, not yet replayed** — `DeltaDumpResult.scores`; a restored key is scored as if freshly inserted into the target (§3) | +| Optimizer state | **Dumped, not yet replayed** — `DeltaDumpResult.optimizer_states`, split out of the same value row as `values` (§3) | +| `evicted_keys` | **Never replayed.** Overwriting a slot already reproduces an eviction, so this list exists only for other consumers of the dump | +| `erased_keys` | A separate buffer and a separate list, applied whenever non-empty, **erase before upsert** — nothing takes over an erased key's slot, so no write reproduces it | +| What replay writes | `ReplayContent` flag: embedding / optimizer state / score, any combination, default all three | +| Optimizer state | Never carried in the delta. A slot whose previous occupant was the *same* key keeps its optimizer state; any other slot is re-initialised | +| `dist_type` | `roundrobin` / `hash_roundrobin` only — same restriction the dump side already enforces; `continuous` raises | + +--- + +## 3. What a delta carries + +`DeltaDumpResult` holds five column-aligned per-table lists: `keys`, `values`, +`optimizer_states`, `scores`, plus the two removal lists `evicted_keys` and +`erased_keys`. + +**`values` / `optimizer_states` are the two halves of one stored row.** +`load_from_flat_single_table` returns the table's own compact `[N, value_dim]` +row, so the optimizer state is the trailing `value_dim - dim` columns. + +That runtime width is **not** what gets dumped. Rowwise Adagrad reserves a fixed +16 bytes per row in the fused FBGEMM layout but only ever fills one accumulator +scalar, so the block is narrowed with the same +`truncate_optimizer_states_for_checkpoint` the file dump uses, and replay expands +it back with `pad_optimizer_states_from_checkpoint`. A delta and a checkpoint +then describe a row identically, and the padding is not shipped on every dump. +A table whose optimizer keeps no per-row state (plain SGD) reports `None`, not a +zero-width tensor — matching `export_keys_values_iter`. + +**`scores` carries every score word**, in the user's configured (logical) column +order. Timestamp columns are converted to an **age** (`cur_ts - score`) because +`%globaltimer` is per device and resets across boots, so a raw timestamp means +nothing anywhere but the rank that produced it; the baseline is +`meta["current_score"]`, and a consumer rebases against its own clock. Every +other column (LFU frequency, STEP, CUSTOMIZED, NO_EVICTION row) is verbatim. +This is the same convention the file `dump`/`load` already uses for single-column +LRU tables. + +**Replay does not consume either yet.** A restored key is scored as if it had +just been inserted into the target (`_fresh_score_block`), and keeps its +optimizer state only when it already occupies the target row. The consequence is +worth stating plainly, because it is the one place a replica is *not* a copy of +its source: **the two can evict in different orders.** What they never disagree +on is the embedding of a key they both hold — which is what a delta exists to +synchronise. Wiring the two fields into `replay_increment` closes that gap and is +the natural next step. + +**NO_EVICTION is the exception, and it is exact.** Its score word is not a score +at all but the value row (§12), so replay writes it verbatim from `slot_index`. +There is nothing to approximate: reproducing the source's rows is the point. + +`cur_ts` is sampled **once, before** the per-table dump loop and reused as both +the age baseline and `meta["current_score"]`. Sampling before (rather than after) +the dump makes the next window at-least-once instead of at-most-once: a key +touched *during* the dump is re-dumped next time rather than dropped. + +Storage-level `incremental_dump` returns +`(keys, values, slot_index, optimizer_states, scores)`, and +`_all_gather_dumped_keys_values` is generalised to gather an arbitrary list of +aligned tensors — the optimizer column is present or absent as a whole, and +every rank agrees, so the collective stays symmetric. + +--- + +## 4. Write-back by slot + +### 4.1 Why the target must have the same layout + +A key can only ever live in its own home bucket: + +``` +bucket_id = bkt_begin + (hash(key) % table_cap) / bucket_capacity +``` + +So writing key `K` into slot `S` is only *findable* if `S` lies inside `K`'s home +bucket — which requires the target's `table_cap` and `bucket_capacity` to equal +the source's. §4.2 checks that up front and §4.3 re-checks it per key. + +### 4.2 Table-level compatibility check (Python) + +A table is replayed only when **all** of these match; otherwise `replay_increment` +raises `ValueError` naming the first mismatch, before writing anything: + +- `meta["current_capacity"]` vs target `key_index_map.capacity(table_id)` +- `meta["table_options"].bucket_capacity` vs the target table's +- `meta["table_options"].score_strategy` vs the target's, compared by + **physical** word order — `(TIMESTAMP, LFU)` and `(LFU, TIMESTAMP)` are the + same on-device layout, and with no scores in flight the configured tuple order + no longer matters +- `meta["table_options"].dim` vs the target's embedding dim +- `meta["world_size"]` vs the target's `_shard_world_size` +- storage kind (single-tier vs hybrid) and `num_scores` + +### 4.3 Key-level enforcement (kernel) + +New primitive `table_scatter_keys_at_slots`, one thread per key: + +1. `slot < 0`, invalid key, or `slot >= table_cap` → `status[i] = -1`. +2. `home_bucket(key) != bucket_of(slot)` → `status[i] = -1` (the layout check + can still fail per key if the target rehashed). +3. Probe the home bucket for `key`; if it already lives at a *different* slot, + reclaim that slot (this is the only way a duplicate could appear). +4. Lock the target slot, write digest + all score words, unlock with the key. + `bucket_sizes` is incremented only when the slot did not previously hold a + valid key. +5. `status[i] = slot`, and `same_key[i]` records whether the slot's previous + occupant was this very key (used to decide optimizer-state handling). + +Any `status[i] == -1` raises: the table-level check already established that the +layouts match, so a slot outside its key's home bucket means they have diverged +anyway, and continuing would silently drop that key. + +### 4.4 Precondition: the target must be a replica + +Writing at the source's slot **overwrites whatever occupies it** in the target. +That is the convergent behaviour a replica needs: if the source evicted key `B` to +make room for `A` in slot `S`, a same-capacity replica must do the same, so +refusing to overwrite would leave the replica permanently diverged. + +The corollary is a precondition: a replayed table must be built *only* by +loading/replaying from its source. A table that also takes independent writes can +lose a key whose slot a delta key claims. + +### 4.5 `slot_index` unpacking + +Mirrors `_encode_slot_index` (`docs/delta_dump_result_design.md` §4): + +- single-tier, normal policy: `slot_index` **is** both key slot and value row +- single-tier, NO_EVICTION: high 32 bits = key slot, low 32 bits = value row; + the target's `no_eviction_next_index` counter is bumped past `max(row) + 1` +- `HybridStorage`: bit 63 selects the tier (0 = HBM, 1 = host), bits 0..62 are + the key slot; each tier is replayed into its own state + +--- + +## 5. Value writes and optimizer state + +The delta carries embeddings only (`values[:, :emb_dim]`), never optimizer +state. `store_to_flat_table_contiguous` copies `min(value_dim, input_dim)` +columns from the row base, so writing an `[N, emb_dim]` tensor touches the +embedding region and leaves the optimizer region untouched. Replay therefore +splits every batch in two: + +| group | write | +|---|---| +| slot previously held the **same key** (`same_key` / `founds`) | embedding columns only — the key's own optimizer state survives | +| new key, or the slot's previous occupant was a different key | embedding columns **plus** `initial_optim_state` for the optimizer region | + +Without the second group a new key would silently inherit the evicted key's +optimizer moments. + +--- + +## 6. Sharding and removals + +**Rank filtering.** Replay keeps only the keys the **target** rank owns, +recomputing ownership from the key with the target's own fan-out: + +``` +roundrobin : key % world_size == rank +hash_roundrobin : fmix64(key) % world_size == rank +continuous : NotImplementedError (range-based; not reconstructible per key) +``` + +The delta's scope depends on how it was dumped: `incremental_dump(..., pg)` +all-gathers, so every rank holds the *global* delta and the filter hands each +rank its share; `pg=None` leaves each rank holding only its own keys, so that +delta belongs on the rank that produced it (elsewhere the filter drops all of it, +reported as `skipped` rather than as an error). + +Because ownership is recomputed with the *target's* `world_size`, a globally +gathered delta reshards for free — an 8-rank dump can be replayed into a 4-rank +replica. `world_size == 1` skips filtering entirely. `meta["world_size"]` is +compared only to decide whether the table may be replayed at all, never to route +keys. + +**Removals.** A key leaves a table two ways, and the two go to **separate +retained buffers**: + +| | buffer | retained when | replayed? | +|---|---|---|---| +| evicted to make room | `evicted_keys` | `RETAIN_KEY` | never | +| removed by an explicit erase | `erased_keys` | the `erase` call's own `EvictedItemMode` | whenever non-empty | + +Splitting them is what lets replay be correct *and* cheap. An eviction needs no +action on the replica: the key that took the evicted one's slot is in the very +same delta and overwrites it. An erase does, because nothing takes over that +slot. With both in one buffer a replay could not tell them apart and would have +to erase keys that were merely evicted -- dropping live data whose new owner had +not yet been dumped. + +Splitting them also removed a piece of machinery. An earlier design carried a +`ReplayMode` enum in `meta` to say whether the removals were worth applying, +precisely because one buffer could not distinguish the two cases. With separate +buffers there is nothing left to declare: a non-empty `erased_keys` *is* the +signal, so the enum and the "an erase happened" flag it depended on are both +gone. + +The erase runs **before** the upsert, so a key erased and re-added inside the +same window survives -- it appears in both `keys` and `erased_keys`. + +Retention is configured at different times for the two, and for a reason. +Retaining evictions swaps in a collecting insert kernel and costs an extra kernel +output plus a device sync on *every* evicting insert, so it has to be a +table-level decision made up front (`DynamicEmbTableOptions.evicted_item_mode`). +An erase already holds its keys, so recording them costs a copy and nothing has +to be prepared -- which is why `erase` takes its own `EvictedItemMode` per call. + +That also settles a `HybridStorage` question cleanly: only the host tier is a +*last* tier, so only it retains evictions (the HBM tier spills into it rather +than really evicting). Erases are not tier-specific -- `erase_keys` runs against +both tiers, passing its mode down -- so a key erased out of HBM is recorded just +the same, and `pop_erased_keys` drains both tiers. + +--- + +## 7. Cache interaction + +For a module with a `DynamicEmbCache` in front of `DynamicEmbStorage`: + +1. `flush_cache(cache, storage)` — push dirty cache entries down, so the storage + copy is authoritative before it is overwritten. +2. Replay into the storage. +3. Erase the replayed keys from the cache, so the next lookup refetches. (A full + `cache.reset()` would also be correct but throws away unrelated entries.) + +--- + +## 8. API surface (three layers, symmetric with the dump side) + +```python +# model level -- dynamicemb/incremental_dump.py +def replay_increment( + model: torch.nn.Module, + deltas: Dict[str, DeltaDumpResult], + pg: Optional[dist.ProcessGroup] = None, +) -> Dict[str, Dict[str, ReplayStats]]: ... + +# module level -- BatchedDynamicEmbeddingTables +def replay_increment(self, delta, pg=None) -> Dict[str, ReplayStats]: ... + +# storage level -- DynamicEmbStorage / HybridStorage +def replay_increment(self, table_id, keys, values, slot_index, insert_score, + content=ReplayContent.ALL, timestamp=None + ) -> ReplayStats: ... +def erase_keys(self, table_id, keys) -> int: ... +``` + +`ReplayStats` (in `dynamicemb/types.py`, alongside the other shared dataclasses) +is a plain work report -- replay either succeeds or raises: + +```python +@dataclass +class ReplayStats: + upserted: int # keys written back at their source slot + erased: int # keys actually removed (REMOVE_THEN_OVERWRITE only) + skipped: int # keys filtered out as not owned by this rank +``` + +--- + +## 9. Changed Files + +| File | Change | +|---|---| +| `src/table_operation/kernels.cuh` | new `scatter_keys_at_slots_kernel` | +| `src/table_operation/insert.cu` | `table_scatter_keys_at_slots` implementation | +| `src/table_operation/table.cuh` / `table.cu` | declaration + pybind registration | +| `dynamicemb/scored_hashtable.py` | `LinearBucketTable.scatter_keys_at_slots` | +| `dynamicemb/key_value_table.py` | dump-side `slot_index` / `_split_value_row` / `_dump_score_block` (+ age transform); storage-level `replay_increment` / `erase_keys`; `_fresh_score_block`; generalised all-gather | +| `dynamicemb/batched_dynamicemb_tables.py` | module-level `replay_increment`; `meta` layout fields; single `cur_ts` | +| `dynamicemb/incremental_dump.py` | model-level `replay_increment`; `DeltaDumpResult` doc | +| `dynamicemb/types.py` | `ReplayStats` | +| `dynamicemb/dynamicemb_config.py` | `ReplayContent`, `EvictedItemMode` as a flag | +| `dynamicemb/__init__.py` | export `replay_increment`, `ReplayStats`, `pop_erased_keys`, `ReplayContent` | +| `dynamicemb/batched_dynamicemb_function.py` | split slot vs value-row indices in the HBM-direct prefetch (§12) | +| `DynamicEmb_APIs.md` | document both | +| `test/unit_tests/incremental_dump/test_replay_increment.py` | new | +| `test/unit_tests/test_no_eviction_row_indexing.py` | new (§12 regression) | + +--- + +## 10. Tests + +`test/unit_tests/incremental_dump/test_replay_increment.py` (single process, one +`BatchedDynamicEmbeddingTablesV2` as source and another as target): + +| Test | Covers | +|---|---| +| `test_replay_round_trip` | keys, embeddings and **slot_index** all round-trip; parametrised over TIMESTAMP / STEP / LFU / compound `(TIMESTAMP, LFU)` / NO_EVICTION | +| `test_dump_splits_the_value_row` | `values` + `optimizer_states` widths add up to the table's value row; a table with no per-row state reports `None`; parametrised over SGD / ROWWISE_ADAGRAD | +| `test_dump_scores_are_column_aligned` | `scores` is `[N, num_scores]` in logical order, and the LFU column is verbatim (keys touched six times outrank those touched once) | +| `test_replay_does_not_carry_scores` | the deliberate gap: a threshold that splits the source's table splits nothing on the replica, because every restored key carries the same fresh score | +| `test_replay_is_idempotent` | replaying the same delta twice is a no-op | +| `test_replay_rejects_layout_mismatch` | capacity mismatch raises **and leaves the target untouched** (no partial write) | +| `test_replay_rejects_score_strategy_mismatch` | a genuinely different score layout raises | +| `test_replay_accepts_swapped_score_order` | `(TIMESTAMP, LFU)` vs `(LFU, TIMESTAMP)` is the *same* physical layout, so it must replay rather than be rejected | +| `test_replay_rejects_missing_table_options` | a delta from an older version raises instead of writing blind | +| `test_replay_applies_erasures_and_ignores_evictions` | `erased_keys` is applied, erase before upsert, an erased-then-readmitted key survives; a non-empty `evicted_keys` changes nothing | +| `test_replay_content_restores_optimizer_state` | with `OPTIMIZER_STATE` the source's state lands even on a row the key is taking over, where the default would have initialised it | +| `test_replay_content_restores_scores` | with `SCORE` the replica reproduces the source's ranking -- the mirror of the test above it | +| `test_replay_without_embedding_needs_aligned_rows` | omitting `EMBEDDING` raises when a key does not already own its target row, and works once the replica is aligned | +| `test_erased_and_evicted_use_separate_buffers` | an erase and an eviction land in different buffers, neither leaks into the other, both are read-and-clear | +| `test_erase_retention_is_per_call_not_per_table` | the same table records one erase and not the next; a table that discards evictions still reports its erases | +| `test_replay_erased_counts_only_keys_actually_present` | `erased` reports removals that really happened, not removals asked for | +| `test_replay_tolerates_an_empty_removal_list` | nothing to remove is normal, not an error | +| `test_replay_optimizer_state` | a key already at its target slot keeps its optimizer state, a fresh row is initialised | +| `test_replay_with_caching` | a lookup after replay sees the replayed value, not a stale cached one | + +| Test | Covers | +|---|---| +| `test_replay_increment_distributed` | the sharding path: each rank keeps only its own keys (`upserted + skipped == total`), the ranks' shares sum to the whole delta, and the replica's dump matches the source's keys/values | + +Not covered yet: replaying into a **differently sized** world (it needs two world +sizes in one job). + +--- + +## 11. Fixed along the way: NO_EVICTION slot-vs-row in the forward path + +Adding the `NO_EVICTION` parametrisation to `test_replay_round_trip` surfaced a +pre-existing fault that had nothing to do with replay: a plain **second forward +pass** over already-resident keys crashed with `an illegal memory access` inside +`load_from_flat` (`dynamic_emb_op.cu:561`). It reproduced with every dump and +replay call removed, so it was not introduced here -- but it blocked the +`NO_EVICTION` coverage, so it is fixed in this change. + +**Cause.** NO_EVICTION is the one strategy where the value row is not the hash +slot: rows come from a per-table auto-increment counter (`key_value_table.py` +`_get_no_eviction_insert_scores`), and its key map is sized +`ceil(init_capacity / 0.5)` -- **twice** the value buffer +(`key_value_table.py:357-367` vs `:403-411`). `_prefetch_hbm_direct_path` +translated slot to row on the **insert** branch only: + +```python +new_indices = (score_arg.value.to(torch.int64) # the row + if state.no_eviction_next_index is not None else new_indices) +... +if h_num_missing == 0: + return indices, indices.clone(), None # the raw hash slot +``` + +so a batch whose keys all hit returned slots, the forward fed them to +`load_from_flat`, and roughly half of them (slots are uniform over a space twice +the buffer) addressed rows past the end of it. The helper for exactly this, +`_flat_row_indices_for_value_load`, was never called on that path. + +**Fix.** Carry the two index spaces separately instead of hoping they coincide. +`PrefetchState` gains `value_row_indices` / `update_value_row_indices`, `None` +meaning "same as the slot tensors" so nothing extra is allocated for the other +strategies. Consumers pick by what they address: + +| Consumer | Index space | +|---|---| +| `load_from_flat` (forward), `fused_update_for_flat_table` (backward) | **row** | +| `increment_counter` / `decrement_counter` (ref counter beside the hash table) | **slot** | + +That also corrects a second, quieter defect on the same line: the old code +overwrote `new_indices` with the row *before* `increment_counter`, so newly +inserted NO_EVICTION keys pinned the wrong counter entry while looked-up keys +pinned the right one. + +Regression coverage: `test/unit_tests/test_no_eviction_row_indexing.py` drives +resident keys through the forward and checks each row against its own key's DEBUG +pattern (`key % 100000`), so a wrong-row read fails on values, not merely on a +crash. diff --git a/corelib/dynamicemb/dynamicemb/__init__.py b/corelib/dynamicemb/dynamicemb/__init__.py index a7307205a..3e861faab 100644 --- a/corelib/dynamicemb/dynamicemb/__init__.py +++ b/corelib/dynamicemb/dynamicemb/__init__.py @@ -22,6 +22,7 @@ DynamicEmbScoreStrategy, DynamicEmbTableOptions, EvictedItemMode, + ReplayContent, ScoreStrategy, align_to_table_size, data_type_to_dtype, @@ -32,7 +33,12 @@ string_to_evict_strategy, ) from .embedding_admission import FrequencyAdmissionStrategy, KVCounter -from .incremental_dump import DeltaDumpResult, pop_evicted_keys +from .incremental_dump import ( + DeltaDumpResult, + pop_erased_keys, + pop_evicted_keys, + replay_increment, +) from .optimizer import EmbOptimType, OptimizerArgs from .types import ( BUCKET_ALIGNMENT, @@ -42,6 +48,7 @@ Counter, DynamicEmbInitializerArgs, DynamicEmbInitializerMode, + ReplayStats, ) from .utils import torch_to_dyn_emb @@ -64,6 +71,7 @@ "DynamicEmbEvictStrategy", "DynamicEmbScoreStrategy", "EvictedItemMode", + "ReplayContent", "ScoreStrategy", "BATCH_SIZE_PER_DUMP", "data_type_to_dyn_emb", @@ -74,7 +82,10 @@ "DynamicEmbDump", "DynamicEmbLoad", "DeltaDumpResult", + "ReplayStats", + "pop_erased_keys", "pop_evicted_keys", + "replay_increment", "EmbOptimType", "OptimizerArgs", ] diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py index 40aa12585..3baf42032 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_function.py @@ -691,7 +691,7 @@ def _prefetch_hbm_direct_path( # victims are dropped here (unlike the cache path, which spills # to storage). Retain them when configured -- this is the # forward-path analogue of the _insert_key_values collection. - if state.evicted_item_mode == EvictedItemMode.RETAIN_KEY: + if EvictedItemMode.RETAIN_KEY in state.evicted_item_mode: ( new_indices, num_evicted, diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py index 57f5d3ce8..470306ee5 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py @@ -20,8 +20,9 @@ from enum import Enum from functools import partial from itertools import accumulate -from typing import Deque, Dict, List, Optional, Tuple +from typing import Any, Callable, Deque, Dict, List, Optional, Tuple, Union +import numpy as np import torch # usort:skip import torch.distributed as dist from dynamicemb.batched_dynamicemb_function import ( @@ -36,7 +37,9 @@ DynamicEmbScoreStrategy, DynamicEmbTableOptions, EvictedItemMode, + ReplayContent, get_eviction_score_strategy, + get_physical_score_order, score_strategy_has_timestamp_column, warning_for_cstm_score, ) @@ -60,6 +63,8 @@ SGDDynamicEmbeddingOptimizer, get_optimizer_state_dim, ) +from dynamicemb.scored_hashtable import murmur3_fmix64 +from dynamicemb.types import ReplayStats from dynamicemb.utils import DTYPE_NUM_BYTES from dynamicemb_extensions import device_timestamp from fbgemm_gpu.split_table_batched_embeddings_ops_training import ( @@ -131,6 +136,45 @@ def find_files(root_path: str, table_name: str, suffix: str) -> Tuple[List[str], return files, len(files) +def owned_key_mask( + keys: torch.Tensor, + rank: int, + world_size: int, + dist_type: str, +) -> Optional[torch.Tensor]: + """Boolean mask selecting the keys *rank* owns under row-wise sharding. + + ``incremental_dump`` all-gathers within its process group, so every rank holds + the whole delta; replay keeps only its own shard. Ownership is recomputed from + the key with the **target's** fan-out, which is what lets a globally gathered + delta be replayed into a differently sized world. + + Returns ``None`` when no filtering is needed (single rank), so callers can + skip the mask entirely. + + Ownership is computed in ``uint64``, matching the device kernel for the + 64-bit index types everyone uses. The kernel actually takes the hash modulo + in ``make_unsigned_t``, so a 32-bit index type would truncate first + and disagree here -- for a world size that is not a power of two, where the + high bits reach the result. Not handled: 32-bit keys are not a configuration + this is built for. + """ + if world_size <= 1: + return None + if dist_type == "continuous": + raise NotImplementedError( + "replay_increment does not support dist_type 'continuous': its " + "key->rank mapping is range-based and cannot be reconstructed from a " + "key alone. Use 'roundrobin' or 'hash_roundrobin'." + ) + keys_np = keys.detach().cpu().numpy().astype(np.uint64, copy=False) + if dist_type == "hash_roundrobin": + owners = murmur3_fmix64(keys_np) % np.uint64(world_size) + else: # roundrobin + owners = keys_np % np.uint64(world_size) + return torch.from_numpy(owners == np.uint64(rank)) + + def get_loading_files( root_path: str, name: str, @@ -1449,38 +1493,74 @@ def pop_evicted_keys( ) -> Dict[str, Tensor]: """Return + clear this rank's retained evicted keys, per table. - Only tables configured with ``evicted_item_mode=RETAIN_KEY`` are included; - others are omitted. Returns ``{table_name: 1-D unique int64 keys on - device}``. The keys are this rank's local shard only (row-wise sharded, so - disjoint across ranks); cross-rank aggregation is the model-level - ``pop_evicted_keys``'s job via ``pg``. Clearing affects only this rank. + Only tables configured with ``evicted_item_mode=RETAIN_KEY`` are + included; others are omitted. Returns + ``{table_name: 1-D unique keys}`` on host. The keys are this rank's local + shard only (row-wise sharded, so disjoint across ranks); cross-rank + aggregation is the model-level ``pop_evicted_keys``'s job via ``pg``. + Clearing affects only this rank. + + Keys removed by an explicit erase are **not** here -- see + :meth:`pop_erased_keys`. """ + return self._pop_retained( + "pop_evicted_keys", + lambda mode: EvictedItemMode.RETAIN_KEY in mode, + table_names, + ) + + def pop_erased_keys( + self, table_names: Optional[List[str]] = None + ) -> Dict[str, Tensor]: + """Return + clear the keys an explicit erase removed, per table. + + The counterpart of :meth:`pop_evicted_keys`, kept apart from it because + the two mean different things to a consumer: an eviction is reproduced by + whoever takes over the slot, while an erase leaves the slot to nobody and + has to be replayed as a removal. + + Every table is included, unlike :meth:`pop_evicted_keys`: whether an + erase was recorded is that ``erase`` call's decision, not the table's, so + there is no configuration to filter on -- a table nobody asked to record + simply returns an empty tensor. + """ + return self._pop_retained("pop_erased_keys", lambda mode: True, table_names) + + def _pop_retained( + self, + method: str, + wanted: Callable[[EvictedItemMode], bool], + table_names: Optional[List[str]], + ) -> Dict[str, Tensor]: + """Drain one of the storage's retained-key buffers, per table.""" storage = self._storage - if not hasattr(storage, "pop_evicted_keys"): + if not hasattr(storage, method): return {} + pop = getattr(storage, method) result: Dict[str, Tensor] = {} for i, name in enumerate(self._table_names): - if ( - self._dynamicemb_options[i].evicted_item_mode - != EvictedItemMode.RETAIN_KEY - ): + if not wanted(self._dynamicemb_options[i].evicted_item_mode): continue if table_names is not None and name not in table_names: continue - result[name] = storage.pop_evicted_keys(i).cpu() # host tensor + result[name] = pop(i).cpu() # host tensor return result def incremental_dump( self, - named_thresholds: Dict[str, int] = None, + named_thresholds: Dict[str, int], pg: Optional[dist.ProcessGroup] = None, ) -> "DeltaDumpResult": """Dump keys/values (+ evicted keys + meta) whose score crosses the threshold. Returns a :class:`DeltaDumpResult` for this module (column-aligned lists by - table). ``meta[i]`` carries current_score / slot_index / current_capacity / - world_size / table_options; ``evicted_keys[i]`` is the retained evicted keys - for an ``evicted_item_mode=RETAIN_KEY`` table (drained here) else ``None``. + table). ``values[i]`` is embeddings only, with the rest of each stored + row in ``optimizer_states[i]`` and every score word in ``scores[i]``; + ``meta[i]`` carries current_score / slot_index / current_capacity / + bucket_capacity / num_scores / world_size / table_options; ``evicted_keys[i]`` is the keys this table + retained since the last dump -- evictions and explicit erases under + ``RETAIN_KEY`` -- and ``erased_keys[i]`` the keys an explicit erase + asked to have recorded. Both are drained here. The meaning of the threshold depends on the table's score strategy: @@ -1497,7 +1577,6 @@ def incremental_dump( """ from dynamicemb.incremental_dump import ( # lazy: avoid import cycle DeltaDumpResult, - _all_gather_evicted_keys, ) storage = self._storage @@ -1509,7 +1588,11 @@ def incremental_dump( if self._cache is not None and isinstance(storage, DynamicEmbStorage): flush_cache(self._cache, storage) res = DeltaDumpResult() - ts: Optional[int] = None + # One reference timestamp for the whole call: the age baseline for the + # dumped timestamp score columns, and the returned current_score. + # Sampled BEFORE the dump so a key touched *during* the dump falls into + # the next window (at-least-once) instead of being missed. + ts = device_timestamp() for table_name, threshold in named_thresholds.items(): if table_name not in self._table_names: warnings.warn( @@ -1520,27 +1603,30 @@ def incremental_dump( ) continue table_id = self._table_names.index(table_name) - keys_cat, values_cat, slot_index = storage.incremental_dump( - table_id, threshold, pg - ) + ( + keys_cat, + values_cat, + slot_index, + opt_states, + scores, + ) = storage.incremental_dump(table_id, threshold, pg, timestamp=ts) option = self._dynamicemb_options[table_id] if score_strategy_has_timestamp_column(option.score_strategy): - if ts is None: - ts = device_timestamp() current_score = ts else: current_score = self._scores[table_name] - # evicted keys: only retain-enabled tables; drain the buffer, aggregate - # within the SAME pg as keys/values, return on host. - if option.evicted_item_mode == EvictedItemMode.RETAIN_KEY and hasattr( - storage, "pop_evicted_keys" - ): - ev = storage.pop_evicted_keys(table_id) - if pg is not None: - ev = _all_gather_evicted_keys(ev, pg) - ev = ev.cpu() - else: - ev = None + # Drain both retained-key buffers, each gated by what its table + # actually retains, and aggregate within the SAME pg as keys/values. + ev = self._drain_retained( + storage, + "pop_evicted_keys", + EvictedItemMode.RETAIN_KEY in option.evicted_item_mode, + table_id, + pg, + ) + # Always drained: recording an erase is the erase call's decision, + # so there is no table setting to gate on here. + er = self._drain_retained(storage, "pop_erased_keys", True, table_id, pg) # current_capacity: DynamicEmbStorage has one key_index_map; a # HybridStorage sums its tiers. if hasattr(storage, "key_index_map"): @@ -1552,14 +1638,339 @@ def incremental_dump( res.table_names.append(table_name) res.keys.append(keys_cat) res.values.append(values_cat) + res.optimizer_states.append(opt_states) + res.scores.append(scores) res.evicted_keys.append(ev) + res.erased_keys.append(er) res.meta.append( { "current_score": current_score, "slot_index": slot_index, "current_capacity": current_capacity, + "bucket_capacity": self._bucket_capacity_of(storage), + "num_scores": self._num_scores_of(storage), "world_size": self._shard_world_size, "table_options": option, } ) return res + + def _replay_compatibility( + self, table_id: int, meta: Dict[str, Any] + ) -> Optional[str]: + """Why this table cannot take a slot-for-slot replay, or ``None``. + + A key is only ever probed inside its own home bucket, and that bucket is + ``hash(key) % table_capacity / bucket_capacity`` -- so a source slot is + only meaningful when the target's capacity and bucket layout match the + source's. Everything compared here comes from the delta's ``meta``. + + ``score_strategy`` is compared by its **physical** word order, not the + configured tuple: what a slot write depends on is how the score words are + laid out on device, and ``(TIMESTAMP, LFU)`` and ``(LFU, TIMESTAMP)`` are + the same layout -- the tuple order only ever decided checkpoint column + order. Two genuinely different strategies still differ physically and are + still rejected. + """ + storage = self._storage + option = self._dynamicemb_options[table_id] + src_options = meta.get("table_options") + if src_options is None: + return "delta carries no 'table_options' (dumped by an older version)" + + if hasattr(storage, "key_index_map"): + capacity = storage.key_index_map.capacity(table_id) + else: + capacity = sum(s.key_index_map.capacity(table_id) for s in storage.tables) + checks = [ + ("capacity", meta.get("current_capacity"), capacity), + ( + "bucket_capacity", + meta.get("bucket_capacity"), + self._bucket_capacity_of(storage), + ), + ("num_scores", meta.get("num_scores"), self._num_scores_of(storage)), + ("world_size", meta.get("world_size"), self._shard_world_size), + # Compared by physical layout -- see the note above. + ( + "score_strategy", + get_physical_score_order(src_options.score_strategy), + get_physical_score_order(option.score_strategy), + ), + ("dim", src_options.dim, option.dim), + ("dist_type", src_options.dist_type, option.dist_type), + ] + for name, src, dst in checks: + if src is None: + return f"delta carries no '{name}' (dumped by an older version)" + if src != dst: + return f"{name} mismatch (source {src} vs target {dst})" + return None + + def replay_increment( + self, + delta: "DeltaDumpResult", + pg: Optional[dist.ProcessGroup] = None, + content: ReplayContent = ReplayContent.ALL, + ) -> Dict[str, ReplayStats]: + """Write an ``incremental_dump`` delta back into this module's tables. + + Replay is exact in *layout*: every key is written at the slot and value + row it held in the source table, so the target ends up layout-identical + to it. A table whose layout does not match the source's is rejected with + a ``ValueError`` rather than written some other way -- see + :meth:`_replay_compatibility`. + + *content* selects which parts of each row are written -- embedding, + optimizer state, score -- and defaults to all three, i.e. the replica + ends up holding what the source held. Dropping ``SCORE`` leaves restored + keys scored as if freshly inserted here, so the replica orders its own + future evictions by when it received each key rather than by how the + source ranked it; dropping ``OPTIMIZER_STATE`` keeps the state of a key + already on its target row and initialises any other. See + :class:`ReplayContent`. + + ``delta.erased_keys`` is applied whenever it holds anything: those keys + were explicitly removed at the source and nothing else will remove them + here. ``delta.evicted_keys`` is never applied -- writing the delta's + slots already reproduces an eviction, because the key that took the + evicted one's slot is in the delta. A table retaining evictions for some + other consumer therefore costs a replay nothing. + + Only the keys this rank owns are replayed, with ownership recomputed + from the key using *this* model's world size. A delta gathered over a + process group therefore fans out correctly when handed to every rank; a + per-rank delta (``incremental_dump`` with ``pg=None``) should be replayed + on the rank that produced it, or the filter drops all of it -- visible as + ``ReplayStats.skipped``. + + Args: + delta: one collection's :class:`DeltaDumpResult`. Tables not present + in this module are skipped with a warning. + pg: process group used to size this model's shard fan-out. Defaults + to the world the tables were created against. + + Returns: + ``{table_name: ReplayStats}`` -- keys written / removed / skipped + per table. + + Raises: + ValueError: a table's layout does not match the source's, or the + delta is missing the per-key data replay needs. Raised before + anything is written. + TypeError: this module's storage is neither ``DynamicEmbStorage`` nor + ``HybridStorage``. + NotImplementedError: a table is sharded with + ``dist_type="continuous"`` across more than one rank. + RuntimeError: a key could not be written at its source slot even + though the metadata matched. Unlike the checks above this fires + mid-write, so the table may hold a partial replay. + """ + storage = self._storage + if not isinstance(storage, (DynamicEmbStorage, HybridStorage)): + raise TypeError( + f"replay_increment requires DynamicEmbStorage or HybridStorage, " + f"got {type(storage).__name__}" + ) + rank = dist.get_rank(group=pg) if dist.is_initialized() else 0 + world_size = ( + dist.get_world_size(group=pg) + if (pg is not None and dist.is_initialized()) + else self._shard_world_size + ) + if self._cache is not None and isinstance(storage, DynamicEmbStorage): + # Push dirty cache entries down first so the storage copy this replay + # is about to overwrite is the authoritative one. + flush_cache(self._cache, storage) + + ts = device_timestamp() + results: Dict[str, ReplayStats] = {} + for i, table_name in enumerate(delta.table_names): + if table_name not in self._table_names: + warnings.warn( + f"replay_increment: table_name '{table_name}' is not in this " + f"module (available: {self._table_names}); skipping.", + UserWarning, + stacklevel=2, + ) + continue + table_id = self._table_names.index(table_name) + meta = delta.meta[i] + option = self._dynamicemb_options[table_id] + stats = ReplayStats() + + keys = delta.keys[i] + values = delta.values[i] + opt_states = ( + delta.optimizer_states[i] + if ReplayContent.OPTIMIZER_STATE in content + else None + ) + scores = delta.scores[i] if ReplayContent.SCORE in content else None + slot_index = meta.get("slot_index") + widths = {"values": values.size(0)} + if opt_states is not None: + widths["optimizer_states"] = opt_states.size(0) + if scores is not None: + widths["scores"] = scores.size(0) + if any(w != keys.numel() for w in widths.values()): + raise ValueError( + f"replay_increment: delta columns for table '{table_name}' are " + f"not row-aligned (keys={keys.numel()}, " + + ", ".join(f"{k}={v}" for k, v in widths.items()) + + ")." + ) + + if slot_index is None: + raise ValueError( + f"replay_increment: delta for table '{table_name}' carries no " + "slot_index; it was produced by an incompatible version of " + "incremental_dump." + ) + mismatch = self._replay_compatibility(table_id, meta) + if mismatch is not None: + raise ValueError( + f"replay_increment: cannot replay table '{table_name}' -- " + f"{mismatch}. Replay writes every key back at the slot it " + "held in the source table, which is only meaningful when the " + "two tables share a layout; configure the target to match the " + "source, or rebuild it from a full checkpoint instead." + ) + + # Only ``erased_keys`` is ever replayed. ``evicted_keys`` is not a + # removal a replica has to perform: the key that took the evicted + # one's slot is in this very delta and overwrites it. That list + # exists for other consumers of the dump, and replay ignores it. + # An empty or absent list is normal -- nothing was erased, or rank + # filtering left this rank none of them. + erased = delta.erased_keys[i] + # Not replayed as a removal -- the delta's own writes reproduce an + # eviction in the storage -- but still needed to invalidate the + # cache, which those writes do not reach. See below. + evicted = delta.evicted_keys[i] + mask = owned_key_mask(keys, rank, world_size, option.dist_type) + if mask is not None: + stats.skipped = int(keys.numel() - mask.sum().item()) + keys, values = keys[mask], values[mask] + if opt_states is not None: + opt_states = opt_states[mask] + if scores is not None: + scores = scores[mask] + slot_index = slot_index[mask] + if erased is not None and erased.numel() > 0: + er_mask = owned_key_mask(erased, rank, world_size, option.dist_type) + erased = erased[er_mask] + if evicted is not None and evicted.numel() > 0: + ev_mask = owned_key_mask( + evicted, rank, world_size, option.dist_type + ) + evicted = evicted[ev_mask] + + # Erase first: a key erased and then re-inserted inside the same + # window appears in BOTH lists, and must survive the replay. + if erased is not None and erased.numel() > 0: + # Not recorded into this model's own erased buffer: these + # removals came from upstream, and re-reporting them would make + # a chained replica replay what it already received. A model + # that is itself a dump source for someone further down would + # want the opposite -- say so when that case turns up. + stats.erased = storage.erase_keys(table_id, erased) + + # Adopt the source's score bookkeeping BEFORE writing, so a + # restored key lands on the same scale the replica will later + # threshold its own incremental_dump against. Timestamp-based tables + # read their score off the device clock, so there is nothing to + # carry -- their restored keys are stamped with ``ts``. + current_score = meta.get("current_score") + if ( + current_score is not None + and not score_strategy_has_timestamp_column(option.score_strategy) + and table_name in self._scores + ): + self._scores[table_name] = current_score + stats.merge( + storage.replay_increment( + table_id, + keys, + values, + opt_states, + scores, + slot_index, + self._scores.get(table_name, 0), + content=content, + timestamp=ts, + ) + ) + if self._cache is not None: + # The cache is a second index that writing a slot does not reach, + # so everything the storage just stopped holding has to be dropped + # from it explicitly: + # upserted keys -- the next lookup must see the value just + # written into the storage, not the cached one; + # erased keys -- a cached copy would resurrect them; + # evicted keys -- these lost their slot to a delta key. Replay + # ignores them as removals, precisely because the write + # reproduces the eviction *in the storage* -- but a cached + # copy survives that, and ``flush_cache`` would write it back + # down on the next dump, undoing the eviction. + # + # A source table that does not retain evictions (``DISCARD``) + # cannot report that last group, so a caching replica that has to + # converge exactly wants ``evicted_item_mode=RETAIN_KEY``. + parts = [keys] + for extra in (erased, evicted): + if extra is not None and extra.numel() > 0: + parts.append(extra.to(keys.dtype)) + stale = torch.cat(parts) if len(parts) > 1 else keys + if stale.numel() > 0: + device = torch.device(f"cuda:{torch.cuda.current_device()}") + self._cache.key_index_map.erase( + stale.to(device=device), + torch.full( + (stale.numel(),), + table_id, + dtype=torch.int64, + device=device, + ), + ) + results[table_name] = stats + return results + + @staticmethod + def _drain_retained( + storage: Union[DynamicEmbStorage, HybridStorage], + method: str, + wanted: bool, + table_id: int, + pg: Optional[dist.ProcessGroup], + ) -> Optional[Tensor]: + """Drain one retained-key buffer for one table, gathered and on host. + + ``None`` -- rather than an empty tensor -- when the table does not retain + this kind of key at all, so a consumer can tell "nothing was removed" + from "this table does not record removals". + """ + from dynamicemb.incremental_dump import ( # lazy: avoid import cycle + _all_gather_evicted_keys, + ) + + if not wanted or not hasattr(storage, method): + return None + out = getattr(storage, method)(table_id) + if pg is not None: + out = _all_gather_evicted_keys(out, pg) + return out.cpu() + + @staticmethod + def _bucket_capacity_of(storage: Union[DynamicEmbStorage, HybridStorage]) -> int: + """The storage's hash-bucket capacity (HBM tier for a hybrid storage).""" + if hasattr(storage, "key_index_map"): + return storage.key_index_map.bucket_capacity_ + return storage.tables[0].key_index_map.bucket_capacity_ + + @staticmethod + def _num_scores_of(storage: Union[DynamicEmbStorage, HybridStorage]) -> int: + """Score words per key (HBM tier for a hybrid storage).""" + if hasattr(storage, "key_index_map"): + return storage.key_index_map.num_scores_ + return storage.tables[0].key_index_map.num_scores_ diff --git a/corelib/dynamicemb/dynamicemb/dynamicemb_config.py b/corelib/dynamicemb/dynamicemb/dynamicemb_config.py index a9a6ebe14..362063db5 100644 --- a/corelib/dynamicemb/dynamicemb/dynamicemb_config.py +++ b/corelib/dynamicemb/dynamicemb/dynamicemb_config.py @@ -110,20 +110,71 @@ class DynamicEmbEvictStrategy(enum.Enum): CUSTOMIZED = EvictStrategy.KCustomized -class EvictedItemMode(enum.Enum): - """How the *last-tier* storage handles an item it evicts. +class EvictedItemMode(enum.Flag): + """What a *last-tier* table keeps about an item on its way out. - - ``DISCARD`` (default): the evicted key is dropped (existing behavior, zero - overhead). - - ``RETAIN_KEY``: the evicted keys are retained so they can be read back with - ``pop_evicted_keys``. + - ``DISCARD``: nothing is kept. + - ``RETAIN_KEY``: the key is kept, to be read back later. - Additional modes (e.g. retaining values) can be added later without changing - this option's type. + What "on the way out" means is the caller's business, not this type's: + ``DynamicEmbTableOptions.evicted_item_mode`` sets it for keys the table + evicts to make room, and ``erase`` takes it per call for keys it removes. + (The name predates the second use.) + + A flag rather than a plain enum so further kinds of retention -- values, + scores -- can be added and combined without revisiting every call site. Test + for one with ``in``:: + + if EvictedItemMode.RETAIN_KEY in mode: ... + + ``DISCARD`` is the empty set rather than a peer of the others, so "discard + *and* retain" is not something the type can express: ``|`` absorbs it + (``DISCARD | RETAIN_KEY is RETAIN_KEY``), and any combination that cancels + out is ``DISCARD`` again. Test for it as emptiness -- ``not mode``, or + ``mode is EvictedItemMode.DISCARD`` -- and **not** with ``in``, which is + subset containment and so reports the empty set as present in everything:: + + EvictedItemMode.DISCARD in EvictedItemMode.RETAIN_KEY # True, always """ DISCARD = 0 - RETAIN_KEY = 1 + RETAIN_KEY = enum.auto() + + +class ReplayContent(enum.Flag): + """Which parts of a dumped row ``replay_increment`` writes back. + + A delta carries an embedding, the optimizer state that shares its value row, + and every score word (see :class:`DeltaDumpResult`). Which of them a replica + wants depends on what it is for: a serving replica needs the embedding and + nothing else, while a training replica that has to resume from the source's + exact state wants all three. + + Combine with ``|`` and test with ``in``:: + + replay_increment(model, deltas) # ALL, the default + replay_increment(model, deltas, content=ReplayContent.EMBEDDING) + replay_increment( + model, deltas, + content=ReplayContent.EMBEDDING | ReplayContent.SCORE, + ) + + The key itself is always written at its source slot -- that is what a replay + *is*, and the flags only choose what travels with it. Removals + (``DeltaDumpResult.erased_keys``) are likewise always applied; they are the + source telling the replica a key is gone, not a payload to opt out of. + + Omitting ``EMBEDDING`` is only meaningful for a replica already aligned with + its source, where every key still occupies the row it held there. A key that + lands on a row it did not already own has no embedding to keep, and serving + the previous occupant's vector under a new key would be silent corruption -- + so ``replay_increment`` raises instead. + """ + + EMBEDDING = enum.auto() + OPTIMIZER_STATE = enum.auto() + SCORE = enum.auto() + ALL = EMBEDDING | OPTIMIZER_STATE | SCORE class DynamicEmbScoreStrategy(enum.IntEnum): @@ -534,9 +585,20 @@ class DynamicEmbTableOptions: strategy.""" evicted_item_mode: EvictedItemMode = EvictedItemMode.DISCARD - """How the *last-tier* storage handles an item it evicts. ``DISCARD`` (default) - drops evicted keys with zero overhead. ``RETAIN_KEY`` retains the keys it - evicts so they can be read back with ``pop_evicted_keys``. Only the final tier + """How the *last-tier* storage handles an item it **evicts** to make room. + ``DISCARD`` (default) drops evicted keys with zero overhead. ``RETAIN_KEY`` + retains them, to be read back with ``pop_evicted_keys``. + + Evictions have to be decided here, up front, because retaining them swaps in + a collecting insert kernel and costs an extra kernel output plus a device + sync on every evicting insert -- a price the table pays for its whole life. + Removing a key with ``erase`` is the other way a key leaves, and it is not + covered by this setting: an erase already holds its keys, so recording them + costs a copy and nothing has to be prepared, which lets each ``erase`` call + take its own :class:`EvictedItemMode` instead. The two land in separate + buffers (``pop_evicted_keys`` / ``pop_erased_keys``). + + Only the final tier that truly discards a key records it -- intermediate cache / HBM tiers spill their evictions to the next tier and are NOT recorded. Records the (key, table_id) only, no value/score. Tables differing in this mode are not diff --git a/corelib/dynamicemb/dynamicemb/incremental_dump.py b/corelib/dynamicemb/dynamicemb/incremental_dump.py index 8139118a2..83e01bd8a 100644 --- a/corelib/dynamicemb/dynamicemb/incremental_dump.py +++ b/corelib/dynamicemb/dynamicemb/incremental_dump.py @@ -20,6 +20,8 @@ import torch import torch.distributed as dist from dynamicemb.dump_load import find_sharded_modules, get_dynamic_emb_module +from dynamicemb.dynamicemb_config import ReplayContent +from dynamicemb.types import ReplayStats from torch import nn @@ -38,14 +40,52 @@ class DeltaDumpResult: keys : List[torch.Tensor] Per-table matched keys on host (created/modified keys to upsert). values : List[torch.Tensor] - Per-table matched values on host, aligned with ``keys``. + Per-table ``[N, dim]`` embeddings on host, aligned with ``keys``. + **Embeddings only**: the rest of each stored row is in + ``optimizer_states``, and concatenating the two along dim 1 reproduces + the row as the table holds it. + optimizer_states : List[Optional[torch.Tensor]] + Per-table optimizer state on host, aligned with ``keys`` -- the trailing + part of the same value row as ``values``, at the **same width the file + checkpoint uses**, so a delta and a checkpoint describe a row + identically. That is narrower than the runtime row for rowwise Adagrad, + which reserves a fixed 16 bytes per row in the fused layout but fills + only one accumulator scalar; ``replay_increment`` pads it back out. + ``None`` for a table whose optimizer keeps no per-row state, e.g. plain + SGD. + scores : List[torch.Tensor] + Per-table ``[N, num_scores]`` score words on host, aligned with ``keys``, + in the table's configured (logical) column order. Timestamp columns hold + an **age** (``meta[i]["current_score"] - score``) rather than a raw + timestamp, because ``%globaltimer`` is per device and resets across + boots; a consumer rebases them onto its own clock. Every other column + (LFU frequency, STEP, CUSTOMIZED, NO_EVICTION row) is carried verbatim. evicted_keys : List[Optional[torch.Tensor]] - Per-table evicted keys on host (keys only, no value/score) retained since - the last ``incremental_dump``. ``evicted_keys[i]`` is ``None`` for a table - without ``evicted_item_mode=RETAIN_KEY``; otherwise it holds that table's - retained evicted keys, and returning them drains and releases the table's - retained-evicted-keys buffer (each evicted key is reported exactly once - across successive ``incremental_dump`` calls). + Per-table keys on host (keys only, no value/score) that the table + **evicted** to make room since the last ``incremental_dump``. ``None`` + unless the table was configured with ``evicted_item_mode=RETAIN_KEY`` -- + retaining evictions has to be decided up front, since it swaps in a + collecting insert kernel. + + ``replay_increment`` never applies this list: the key that took an + evicted key's slot is in the same delta and overwrites it, so the + eviction is reproduced by the write. It is here for other consumers of + the dump. + erased_keys : List[Optional[torch.Tensor]] + Per-table keys on host that an **explicit erase** removed since the last + ``incremental_dump`` *and asked to have recorded* -- that is each + ``erase`` call's decision (its own ``EvictedItemMode``), not a table + setting, so this is always populated, with an empty tensor when nothing + was recorded. ``replay_increment`` reads ``None`` as empty too, for + deltas assembled by hand. + + These are the removals a replica has to perform for itself -- nothing + takes over the slot, so no write reproduces them. Whether + ``replay_increment`` applies them is ``meta[i]["replay_mode"]``'s call. + + Both lists drain and release their buffer, so each key is reported exactly + once across successive ``incremental_dump`` calls, and both hold only keys + that were really in the table. meta : List[Dict[str, Any]] Per-table dump metadata, aligned with ``table_names``. A flat dict with: meta[i]["current_score"]: int -- table's score after this dump; @@ -56,6 +96,10 @@ class DeltaDumpResult: by ``replay_increment``. For NO_EVICTION tables it packs the key slot (high 32 bits) and value row (low 32 bits) into one int64. meta[i]["current_capacity"]: int -- table's current capacity (slots). + meta[i]["bucket_capacity"]: int -- slots per hash bucket; replay + compares it to decide whether the source slots are usable. + meta[i]["num_scores"]: int -- score words per key; part of the + slot layout replay compares against. meta[i]["world_size"]: int -- ranks the source table was sharded across at creation (global WORLD), used by replay to reconstruct key->rank. NOT the gather ``pg`` (a comm scope only). @@ -66,7 +110,10 @@ class DeltaDumpResult: table_names: List[str] = field(default_factory=list) keys: List[torch.Tensor] = field(default_factory=list) values: List[torch.Tensor] = field(default_factory=list) + optimizer_states: List[Optional[torch.Tensor]] = field(default_factory=list) + scores: List[torch.Tensor] = field(default_factory=list) evicted_keys: List[Optional[torch.Tensor]] = field(default_factory=list) + erased_keys: List[Optional[torch.Tensor]] = field(default_factory=list) meta: List[Dict[str, Any]] = field(default_factory=list) @@ -283,16 +330,10 @@ def incremental_dump( Returns ------- - Tuple: - Dict[str, Dict[str, Tuple[torch.Tensor, torch.Tensor]]]: - The first 'str' is the name of embedding collection. - The second 'str' is the name of embedding table. - The first tensor in the Tuple is matched keys on hosts. - The second tensor in the Tuple is matched values on hosts. - Dict[str, Dict[str, int]]: - The first 'str' is the name of embedding collection. - The second 'str' is the name of embedding table. - `int` is the current score after finishing the dumping process, which will be used as the score for the next forward pass, and can also be used as the input of the next incremental_dump. If input score_threshold is `int`, the Dict will contain all dynamic embedding tables' current score, otherwise only dumped tables' current score will be returned. + Dict[str, DeltaDumpResult]: + One :class:`DeltaDumpResult` per embedding collection, keyed by the + collection's module path. See that class for the per-table columns and + the ``meta`` keys. """ if isinstance(score_threshold, int): @@ -388,7 +429,10 @@ def incremental_dump( collection_result.table_names.extend(module_result.table_names) collection_result.keys.extend(module_result.keys) collection_result.values.extend(module_result.values) + collection_result.optimizer_states.extend(module_result.optimizer_states) + collection_result.scores.extend(module_result.scores) collection_result.evicted_keys.extend(module_result.evicted_keys) + collection_result.erased_keys.extend(module_result.erased_keys) collection_result.meta.extend(module_result.meta) ret[collection_path] = collection_result @@ -396,6 +440,161 @@ def incremental_dump( return ret +def replay_increment( + model: torch.nn.Module, + deltas: Dict[str, DeltaDumpResult], + pg: Optional[dist.ProcessGroup] = None, + content: ReplayContent = ReplayContent.ALL, +) -> Dict[str, Dict[str, ReplayStats]]: + """Write ``incremental_dump`` results back into a model's dynamic embedding tables. + + The inverse of :func:`incremental_dump`: it takes that call's + ``{collection_path: DeltaDumpResult}`` and restores every key's embedding and + score into *model*. Typical use is delta replication -- train on one job, + dump periodically, ship the delta, replay it into a serving replica. + + **Write-back is by slot.** Every key is written at the slot and value row it + occupied in the source table, leaving the target layout-identical to it. That + is only meaningful when the two tables share a layout, so the target is + checked against the delta's ``meta`` (capacity, bucket capacity, score + layout, dim, dist_type, world size) and a mismatch raises + :class:`ValueError` before anything is written. + + Writing at the source's slot **overwrites whatever occupies it** -- that is + what makes a replica converge (the source evicted that occupant to make + room). It also means a replayed table must be built *only* by + loading/replaying from its source: a table that also takes independent writes + can lose a key whose slot a delta key claims. + + **What gets written** is *content*'s call -- embedding, optimizer state, + score, or any combination (see :class:`ReplayContent`). The default is all + three, so the replica ends up holding what the source held. + + Dropping ``SCORE`` leaves a restored key scored as if it had just been + inserted here, so the replica orders its own future evictions by when it + received a key rather than by how the source ranked it -- the two can evict + in different orders, but never hold a wrong embedding for a key they share. + (NO_EVICTION is unaffected either way: its score word is a value row, not a + score, and is restored exactly.) + + **Sharding.** Replay always keeps only the keys this rank owns, recomputing + ownership from the key with *this* model's world size. What that filter does + depends on how the delta was produced: + + - ``incremental_dump(..., pg)`` all-gathers, so every rank holds the whole + group's keys; hand the same delta to every rank and each takes its share. + - ``incremental_dump(..., pg=None)`` leaves each rank with only its own keys; + replay it on the rank that produced it, where the filter is a no-op. + Replaying it on a *different* rank is not an error -- every key simply + belongs to someone else and is skipped, which ``ReplayStats.skipped`` + makes visible. + + Only ``roundrobin`` and ``hash_roundrobin`` tables can be replayed: + ``continuous`` has no per-key rank mapping to reconstruct, so it raises + whenever the fan-out is more than one rank (``incremental_dump`` already + refuses to dump such a table at all). + + **Optimizer state** is not part of a delta. A key that already occupies its + target row keeps its optimizer state; a row taken over from another key (or a + brand-new one) is reset to the table's initial optimizer state. This is what + happens when *content* omits ``OPTIMIZER_STATE``; including it writes the + state the source dumped, for every key. + + Args: + model (nn.Module): the model containing dynamic embedding tables. + deltas (Dict[str, DeltaDumpResult]): ``incremental_dump``'s return value, + keyed by embedding-collection path. Collections or tables that the + model does not have are skipped with a warning. + pg (Optional[dist.ProcessGroup]): process group defining this model's + shard fan-out. Defaults to the world the tables were created against. + content (ReplayContent): which parts of each dumped row to write back -- + embedding, optimizer state, score, or any combination. Defaults to + all three. + + Returns: + Dict[str, Dict[str, ReplayStats]]: + ``{collection_path: {table_name: ReplayStats}}`` -- keys written, + removed and skipped per table. Empty dict when the model has no + dynamic embedding tables. + + Raises: + ValueError: a target table's layout does not match the source's, or a + delta is missing the per-key data replay needs. Raised before + anything is written. + TypeError: a module's storage is neither ``DynamicEmbStorage`` nor + ``HybridStorage``. + NotImplementedError: a table is sharded with ``dist_type="continuous"`` + across more than one rank. + RuntimeError: a key could not be written at its source slot even though + the metadata matched. Unlike the checks above this fires mid-write, + so that table may hold a partial replay. + """ + collections_list: List[Tuple[str, str, nn.Module]] = find_sharded_modules(model, "") + if len(collections_list) == 0: + warnings.warn( + "Input model don't have any TorchREC ShardedEmbeddingCollection or " + "ShardedEmbeddingBagCollection module, can't replay increment!", + UserWarning, + ) + return {} + + collections_by_path = {path: module for path, _, module in collections_list} + for collection_path in deltas.keys(): + if collection_path not in collections_by_path: + warnings.warn( + f"sharded module '{collection_path}' present in the delta was not " + "found in the model; skipping it.", + UserWarning, + ) + + ret: Dict[str, Dict[str, ReplayStats]] = {} + for collection_path, delta in deltas.items(): + collection_module = collections_by_path.get(collection_path) + if collection_module is None: + continue + collection_stats: Dict[str, ReplayStats] = {} + for dynamic_emb_module in get_dynamic_emb_module(collection_module): + # Each module owns a subset of the collection's tables and skips the + # rest, so the delta can be handed to all of them unchanged. + module_delta = _select_tables(delta, dynamic_emb_module.table_names) + if not module_delta.table_names: + continue + collection_stats.update( + dynamic_emb_module.replay_increment( + module_delta, + pg=pg, + content=content, + ) + ) + ret[collection_path] = collection_stats + + if not ret: + warnings.warn( + "Input model don't have any Dynamic embedding tables, can't replay " + "increment!", + UserWarning, + ) + return ret + + +def _select_tables(delta: DeltaDumpResult, table_names: List[str]) -> DeltaDumpResult: + """The sub-delta covering only *table_names*, keeping all lists column-aligned.""" + wanted = set(table_names) + out = DeltaDumpResult() + for i, name in enumerate(delta.table_names): + if name not in wanted: + continue + out.table_names.append(name) + out.keys.append(delta.keys[i]) + out.values.append(delta.values[i]) + out.optimizer_states.append(delta.optimizer_states[i]) + out.scores.append(delta.scores[i]) + out.evicted_keys.append(delta.evicted_keys[i]) + out.erased_keys.append(delta.erased_keys[i]) + out.meta.append(delta.meta[i]) + return out + + def _all_gather_evicted_keys( keys: torch.Tensor, pg: Optional[dist.ProcessGroup] ) -> torch.Tensor: @@ -431,41 +630,23 @@ def _all_gather_evicted_keys( return torch.unique(torch.cat(parts)).cpu() -def pop_evicted_keys( +def _pop_retained_keys( model: torch.nn.Module, - table_names: Optional[Dict[str, List[str]]] = None, - pg: Optional[dist.ProcessGroup] = None, + method: str, + what: str, + table_names: Optional[Dict[str, List[str]]], + pg: Optional[dist.ProcessGroup], ) -> Dict[str, Dict[str, torch.Tensor]]: - """Return + clear the keys evicted (and retained) by last-tier storage, per table. - - Only tables created with ``evicted_item_mode=RETAIN_KEY`` are included; all other - tables are omitted from the result. This is the incremental "pop" of keys - evicted since the previous call -- the retained buffers are cleared on this - rank as they are read. - - Args: - model (nn.Module): the model containing dynamic embedding tables. - table_names (Optional[Dict[str, List[str]]]): optional filter, keyed by - embedding-collection path. ``{collection_path: [table_name, ...]}`` - pops only the listed collections/tables. ``None`` pops every - retain-enabled table in the model. - pg (Optional[dist.ProcessGroup]): optional process group. ``None`` returns - each rank's LOCAL evicted keys (row-wise sharded, hence disjoint across - ranks; zero communication). When given, keys are all_gathered within - ``pg`` so every rank in the group receives the group-wide union. - Clearing always affects only this rank's buffer, regardless of ``pg``. + """Shared body of :func:`pop_evicted_keys` / :func:`pop_erased_keys`. - Returns: - Dict[str, Dict[str, torch.Tensor]]: - ``{collection_path: {table_name: keys}}`` where ``keys`` is a 1-D - int64 tensor of table-unique evicted keys. Empty dict when the model - has no retain-enabled dynamic embedding tables (or none matched). + The two differ only in which buffer they drain; *method* names the module + method and *what* is the noun used in the "no such tables" warning. """ collections_list: List[Tuple[str, str, nn.Module]] = find_sharded_modules(model, "") if len(collections_list) == 0: warnings.warn( "Input model don't have any TorchREC ShardedEmbeddingCollection or " - "ShardedEmbeddingBagCollection module, can't pop evicted keys!", + f"ShardedEmbeddingBagCollection module, can't pop {what} keys!", UserWarning, ) return {} @@ -478,9 +659,9 @@ def pop_evicted_keys( collection_result: Dict[str, torch.Tensor] = {} for dynamic_emb_module in get_dynamic_emb_module(collection_module): - if not hasattr(dynamic_emb_module, "pop_evicted_keys"): + if not hasattr(dynamic_emb_module, method): continue - local = dynamic_emb_module.pop_evicted_keys(wanted) + local = getattr(dynamic_emb_module, method)(wanted) for tname, keys in local.items(): if pg is not None: keys = _all_gather_evicted_keys(keys, pg) @@ -490,3 +671,63 @@ def pop_evicted_keys( ret[collection_path] = collection_result return ret + + +def pop_evicted_keys( + model: torch.nn.Module, + table_names: Optional[Dict[str, List[str]]] = None, + pg: Optional[dist.ProcessGroup] = None, +) -> Dict[str, Dict[str, torch.Tensor]]: + """Return + clear the keys last-tier storage **evicted** to make room, per table. + + Only tables configured with ``evicted_item_mode=RETAIN_KEY`` are included; + all other tables are omitted from the result. This is the + incremental "pop" of keys evicted since the previous call -- the retained + buffers are cleared on this rank as they are read. + + Keys removed by an explicit erase are reported separately, by + :func:`pop_erased_keys`. The two are kept apart because a consumer usually + wants different things from them: an eviction says the table ran out of room, + an erase says someone asked for the key to go. + + Args: + model (nn.Module): the model containing dynamic embedding tables. + table_names (Optional[Dict[str, List[str]]]): optional filter, keyed by + embedding-collection path. ``{collection_path: [table_name, ...]}`` + pops only the listed collections/tables. ``None`` pops every + retain-enabled table in the model. + pg (Optional[dist.ProcessGroup]): optional process group. ``None`` returns + each rank's LOCAL evicted keys (row-wise sharded, hence disjoint across + ranks; zero communication). When given, keys are all_gathered within + ``pg`` so every rank in the group receives the group-wide union. + Clearing always affects only this rank's buffer, regardless of ``pg``. + + Returns: + Dict[str, Dict[str, torch.Tensor]]: + ``{collection_path: {table_name: keys}}`` where ``keys`` is a 1-D + tensor of table-unique evicted keys. Empty dict when the model has no + tables retaining evictions (or none matched). + """ + return _pop_retained_keys(model, "pop_evicted_keys", "evicted", table_names, pg) + + +def pop_erased_keys( + model: torch.nn.Module, + table_names: Optional[Dict[str, List[str]]] = None, + pg: Optional[dist.ProcessGroup] = None, +) -> Dict[str, Dict[str, torch.Tensor]]: + """Return + clear the keys an **explicit erase** removed, per table. + + The counterpart of :func:`pop_evicted_keys`, with the same arguments and the + same drain-on-read semantics. Every table is included, unlike + ``pop_evicted_keys``: whether an erase was recorded is that ``erase`` call's + decision, not the table's, so there is no configuration to filter on -- a + table nobody asked to record simply returns an empty tensor. + + Returns: + Dict[str, Dict[str, torch.Tensor]]: + ``{collection_path: {table_name: keys}}`` where ``keys`` is a 1-D + tensor of table-unique erased keys. Empty dict when the model has no + tables retaining erases (or none matched). + """ + return _pop_retained_keys(model, "pop_erased_keys", "erased", table_names, pg) diff --git a/corelib/dynamicemb/dynamicemb/key_value_table.py b/corelib/dynamicemb/dynamicemb/key_value_table.py index 5184666cd..c9271eb13 100644 --- a/corelib/dynamicemb/dynamicemb/key_value_table.py +++ b/corelib/dynamicemb/dynamicemb/key_value_table.py @@ -26,7 +26,10 @@ DynamicEmbScoreStrategy, DynamicEmbTableOptions, EvictedItemMode, + ReplayContent, + ScoreStrategy, align_to_table_size, + get_physical_score_order, score_dump_permutation, score_load_permutation, ) @@ -53,10 +56,11 @@ SCORE_TYPE, Cache, CopyMode, + ReplayStats, Storage, torch_dtype_to_np_dtype, ) -from dynamicemb_extensions import EvictStrategy, flagged_compact +from dynamicemb_extensions import EvictStrategy, device_timestamp, flagged_compact from dynamicemb_extensions import load_from_flat_table_contiguous as _load_contiguous from dynamicemb_extensions import load_from_flat_table_emb as _load_emb from dynamicemb_extensions import load_from_flat_table_value as _load_value @@ -73,53 +77,98 @@ # --------------------------------------------------------------------------- -def _all_gather_dumped_keys_values( - keys: Tensor, - values: Tensor, - slot_index: Tensor, +def _all_gather_dumped_columns( + columns: List[Tensor], pg: dist.ProcessGroup, -) -> Tuple[Tensor, Tensor, Tensor]: - """Gather (keys, values, slot_index) from all ranks into concatenated CPU tensors. +) -> List[Tensor]: + """Gather a set of row-aligned dump columns from all ranks into CPU tensors. - keys: (N,) int64 on device; values: (N, D) on device; slot_index: (N,) int64 - on device. Returns (out_keys_cpu, out_values_cpu, out_slot_index_cpu) with all - ranks' data in rank order, column-aligned across the three. + Every tensor in *columns* must live on the same device and share ``size(0)`` + (the per-rank dumped-key count); trailing dimensions may differ (e.g. keys + ``(N,)``, values ``(N, D)``, scores ``(N, W)``). Returns the same columns with + all ranks' rows concatenated in rank order, on host and still row-aligned. + + NCCL has no variable-length gather, so this is one size all_gather followed by + a padded all_gather per column. """ - device = keys.device + assert columns, "_all_gather_dumped_columns needs at least one column" + device = columns[0].device world_size = dist.get_world_size(group=pg) - n = keys.numel() + n = columns[0].size(0) d_count = torch.tensor([n], dtype=torch.long, device=device) gathered_counts = [torch.empty_like(d_count) for _ in range(world_size)] dist.all_gather(gathered_counts, d_count, group=pg) - max_n = max(c.item() for c in gathered_counts) - emb_dim = values.shape[1] - dtype_val = values.dtype - keys_pad = torch.zeros(max_n, dtype=torch.int64, device=device) - values_pad = torch.zeros(max_n, emb_dim, dtype=dtype_val, device=device) - slot_pad = torch.zeros(max_n, dtype=torch.int64, device=device) - if n > 0: - keys_pad[:n] = keys - values_pad[:n, :] = values - slot_pad[:n] = slot_index - gathered_keys = [torch.empty_like(keys_pad) for _ in range(world_size)] - gathered_values = [torch.empty_like(values_pad) for _ in range(world_size)] - gathered_slots = [torch.empty_like(slot_pad) for _ in range(world_size)] - dist.all_gather(gathered_keys, keys_pad, group=pg) - dist.all_gather(gathered_values, values_pad, group=pg) - dist.all_gather(gathered_slots, slot_pad, group=pg) - out_keys = torch.cat( - [gathered_keys[i][: gathered_counts[i].item()] for i in range(world_size)], - dim=0, - ).cpu() - out_values = torch.cat( - [gathered_values[i][: gathered_counts[i].item()] for i in range(world_size)], - dim=0, - ).cpu() - out_slots = torch.cat( - [gathered_slots[i][: gathered_counts[i].item()] for i in range(world_size)], - dim=0, - ).cpu() - return out_keys, out_values, out_slots + counts = [int(c.item()) for c in gathered_counts] + max_n = max(counts) + + out: List[Tensor] = [] + for col in columns: + padded = torch.zeros( + (max_n,) + tuple(col.shape[1:]), dtype=col.dtype, device=device + ) + if n > 0: + padded[:n] = col + gathered = [torch.empty_like(padded) for _ in range(world_size)] + dist.all_gather(gathered, padded, group=pg) + out.append( + torch.cat( + [gathered[i][: counts[i]] for i in range(world_size)], dim=0 + ).cpu() + ) + return out + + +def _timestamp_score_columns(score_strategy: ScoreStrategy) -> List[int]: + """Logical score-column indices holding a device timestamp. + + ``%globaltimer`` is per-device and resets across boots, so a raw timestamp is + meaningless anywhere but the rank that produced it. These columns are dumped + as an *age* (``cur_ts - score``) -- the same convention the file checkpoint + uses for single-column LRU tables. Every other column (LFU frequency, STEP, + CUSTOMIZED, NO_EVICTION row) is carried verbatim. + """ + if isinstance(score_strategy, tuple): + return [ + i + for i, s in enumerate(score_strategy) + if s == DynamicEmbScoreStrategy.TIMESTAMP + ] + return [0] if score_strategy == DynamicEmbScoreStrategy.TIMESTAMP else [] + + +def _scores_to_age( + score_strategy: ScoreStrategy, scores: Tensor, cur_ts: int +) -> Tensor: + """Convert timestamp columns to ages, for dumping. + + *scores* is a ``[N, W]`` block in logical (configured) column order. When + there is nothing to convert (no timestamp column, or an empty block) the + input is returned unchanged rather than copied -- callers must not mutate the + result in place. Otherwise a clone is returned. + """ + ts_cols = _timestamp_score_columns(score_strategy) + if not ts_cols or scores.numel() == 0: + return scores + out = scores.clone() + for c in ts_cols: + out[:, c] = torch.clamp(cur_ts - scores[:, c], min=0) + return out + + +def _age_to_scores( + score_strategy: ScoreStrategy, ages: Tensor, target_ts: int +) -> Tensor: + """Inverse of :func:`_scores_to_age`, for replay. + + Same aliasing note: returns *ages* itself when there is nothing to convert. + """ + ts_cols = _timestamp_score_columns(score_strategy) + if not ts_cols or ages.numel() == 0: + return ages + out = ages.clone() + for c in ts_cols: + out[:, c] = torch.clamp(target_ts - ages[:, c], min=0) + return out # --------------------------------------------------------------------------- @@ -261,14 +310,23 @@ class DynamicEmbTableState: # Name of the score column incremental_dump thresholds on. Equals # score_policy.name (for LruLfu this is the leading, timestamp, word). incremental_score_name: Optional[str] = None - # Retain-evicted-keys (last tier only): when RETAIN_KEY, each insert that - # evicts collects the victims' (key, table_id). Accumulated as GPU-tensor - # chunks and only concatenated + de-duplicated on pop_evicted_keys (no - # compaction while appending). One chunk per insert that evicted anything; - # drained on pop. + # Retained-key buffers (last tier only). A key leaves a table two ways, and + # the two are kept apart because consumers want different things from them: + # an eviction is reproduced by whoever overwrites its slot, while an erase + # leaves the slot to nobody and has to be replayed as a removal. + # + # evicted_*_chunks -- victims of an insert that had to make room + # erased_*_chunks -- keys an explicit erase removed + # + # Both accumulate as (key, table_id) GPU-tensor chunks, one chunk per + # operation that removed anything, and are only concatenated + de-duplicated + # on pop (no compaction while appending). Which buffers are fed is + # ``evicted_item_mode``'s call. evicted_item_mode: EvictedItemMode = EvictedItemMode.DISCARD evicted_key_chunks: List[torch.Tensor] = field(default_factory=list) evicted_tid_chunks: List[torch.Tensor] = field(default_factory=list) + erased_key_chunks: List[torch.Tensor] = field(default_factory=list) + erased_tid_chunks: List[torch.Tensor] = field(default_factory=list) def create_table_state( @@ -795,6 +853,143 @@ def _encode_slot_index( return (slot_indices << 32) | flat_rows +def _fresh_score_block( + state: DynamicEmbTableState, + table_id: int, + value_rows: torch.Tensor, + timestamp: int, + insert_score: int, +) -> torch.Tensor: + """``[N, num_scores]`` physical score block for keys written by replay. + + A delta carries embeddings, not scores, so a replayed key is scored as if it + had just been inserted here: recency words get *timestamp*, every other word + gets *insert_score*. The replica therefore ranks its restored keys by when it + received them rather than by how the source ranked them -- acceptable because + the score only orders future evictions, and a replica that evicts on its own + schedule still holds correct embeddings. + + *insert_score* is passed in rather than read off ``state.score``: that field + is populated by the forward pass, and replay is the one write path that can + run before a target has ever done a forward. The caller owns the table's + score bookkeeping and always has a value. + + The exception is NO_EVICTION, whose score word is not a score at all but the + value row. Reproducing the source's rows is the entire point of the replay, + so that word is written verbatim from *value_rows*. + """ + device = state.device + if state.no_eviction_next_index is not None: + return value_rows.to(device=device, dtype=SCORE_TYPE).view(-1, 1) + n = value_rows.numel() + num_scores = state.key_index_map.num_scores_ + block = torch.empty((n, num_scores), device=device, dtype=SCORE_TYPE) + physical = get_physical_score_order(state.options_list[table_id].score_strategy) + for word, strategy in enumerate(physical): + block[:, word] = ( + timestamp if strategy == DynamicEmbScoreStrategy.TIMESTAMP else insert_score + ) + return block + + +def _dump_score_block( + state: DynamicEmbTableState, + table_id: int, + slot_indices: torch.Tensor, + primary_scores: torch.Tensor, + cur_ts: int, +) -> torch.Tensor: + """``[N, num_scores]`` per-key score block for a dumped table. + + Columns are in the user's configured (logical) order -- the same order the + file checkpoint uses -- and timestamp columns are converted to an age + relative to *cur_ts* so they survive the trip to another device or host. + Single-score tables reuse the score column the dump already thresholded on; + multi-word layouts gather every word from the slots. + """ + device = state.device + num_scores = state.key_index_map.num_scores_ + score_strategy = state.options_list[table_id].score_strategy + if num_scores > 1: + block = state.key_index_map.gather_score_blocks( + table_id, slot_indices.to(device=device, dtype=torch.int64) + ) + perm = score_dump_permutation(score_strategy) + if perm != list(range(block.size(1))): + block = block[:, perm].contiguous() + else: + block = primary_scores.to(device=device, dtype=SCORE_TYPE).view(-1, 1) + return _scores_to_age(score_strategy, block, cur_ts) + + +def _replay_score_block( + state: DynamicEmbTableState, + table_id: int, + scores: torch.Tensor, + target_ts: int, +) -> torch.Tensor: + """Inverse of :func:`_dump_score_block`: logical ages -> physical scores. + + Takes the ``[N, num_scores]`` logical block carried in the delta, restores + timestamp columns against this device's *target_ts*, and permutes back into + the physical device layout. + + Both steps index by *column*, so a block that is not exactly as wide as this + table's score layout would silently transform or permute the wrong words -- + e.g. writing a frequency into the timestamp column. The width is validated + here rather than trusted. + """ + device = state.device + score_strategy = state.options_list[table_id].score_strategy + num_scores = state.key_index_map.num_scores_ + block = scores.to(device=device, dtype=SCORE_TYPE) + if block.dim() == 1: + block = block.view(-1, 1) + if block.dim() != 2 or block.size(1) != num_scores: + raise ValueError( + f"replay_increment: table {table_id} has {num_scores} score word(s) " + f"per key, but the delta carries a score block of shape " + f"{tuple(scores.shape)}." + ) + block = _age_to_scores(score_strategy, block, target_ts) + perm = score_load_permutation(score_strategy) + if perm != list(range(num_scores)): + block = block[:, perm] + return block.contiguous() + + +def _split_value_row( + state: DynamicEmbTableState, table_id: int, values: torch.Tensor +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Split a loaded ``[N, value_dim]`` row into (embedding, optimizer state). + + ``load_from_flat_single_table`` returns the table's own compact row, so the + optimizer state is the trailing ``value_dim - emb_dim`` columns. That runtime + width is not all payload: rowwise Adagrad reserves a fixed 16 bytes per row + in the fused FBGEMM layout but only ever fills one accumulator scalar. The + dumped block is therefore narrowed to the width the file checkpoint uses -- + the same :func:`truncate_optimizer_states_for_checkpoint` ``_dump_table`` + applies -- so a delta and a checkpoint describe a row identically, and the + padding is not paid for on every dump. Replay expands it back with + :func:`pad_optimizer_states_from_checkpoint`. + + Tables whose optimizer keeps no per-row state (e.g. plain SGD) get ``None`` + rather than a zero-width tensor, matching :func:`export_keys_values_iter`. + """ + emb_dim = state.table_emb_dims_cpu[table_id] + optim_state_dim = state.table_value_dims_cpu[table_id] - emb_dim + emb = values[:, :emb_dim].to(dtype=state.emb_dtype) + if optim_state_dim <= 0: + return emb, None + opt = values[:, -optim_state_dim:] + return ( + emb, + truncate_optimizer_states_for_checkpoint( + state.optimizer, emb_dim, opt + ).contiguous(), + ) + + def load_from_flat( state: DynamicEmbTableState, indices: torch.Tensor, @@ -1100,31 +1295,67 @@ def _append_evicted( state.evicted_tid_chunks.append(evicted_table_ids[:n].clone()) -def _pop_state_evicted_keys(state: DynamicEmbTableState, table_id: int) -> torch.Tensor: - """Pop (return + clear) the unique evicted keys retained for ``table_id``. +def _append_erased( + state: DynamicEmbTableState, + keys: torch.Tensor, + table_id: int, +) -> None: + """Append an explicit erase's keys to the state's erased-key buffer. - Concatenates all retained chunks, selects this table's rows, de-duplicates, - and rebuilds the chunks to keep only the OTHER tables' rows -- so this table's - retained keys are cleared while other tables' remain for their own pop. - Returns a 1-D tensor of unique keys on ``state.device``. + Unlike :func:`_append_evicted` there is no count to read back: the caller has + already masked the keys down to the ones the erase actually removed, so this + costs no extra device sync. """ - if not state.evicted_key_chunks: + if keys.numel() == 0: + return + state.erased_key_chunks.append(keys.clone()) + state.erased_tid_chunks.append( + torch.full((keys.numel(),), table_id, dtype=torch.int64, device=keys.device) + ) + + +def _pop_chunked_keys( + key_chunks: List[torch.Tensor], + tid_chunks: List[torch.Tensor], + table_id: int, + state: DynamicEmbTableState, +) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]: + """Drain one table's rows out of a (key, table_id) chunk buffer. + + Concatenates the chunks, selects this table's rows, de-duplicates, and + returns the rebuilt chunk lists holding only the OTHER tables' rows -- so + this table's keys are cleared while the rest stay for their own pop. + """ + if not key_chunks: # Match the non-empty path's dtype: chunks carry key_index_map.key_type # (== the table's index_type, which may be int32/uint32), so a hardcoded # int64 here would make callers that concat/compare across empty and # non-empty pops hit a dtype mismatch. - return torch.empty(0, dtype=state.key_index_map.key_type, device=state.device) - keys = torch.cat(state.evicted_key_chunks) - tids = torch.cat(state.evicted_tid_chunks) + empty = torch.empty(0, dtype=state.key_index_map.key_type, device=state.device) + return empty, [], [] + keys = torch.cat(key_chunks) + tids = torch.cat(tid_chunks) mask = tids == table_id out = torch.unique(keys[mask]) keep = ~mask if bool(keep.any()): - state.evicted_key_chunks = [keys[keep]] - state.evicted_tid_chunks = [tids[keep]] - else: - state.evicted_key_chunks = [] - state.evicted_tid_chunks = [] + return out, [keys[keep]], [tids[keep]] + return out, [], [] + + +def _pop_state_evicted_keys(state: DynamicEmbTableState, table_id: int) -> torch.Tensor: + """Pop (return + clear) the unique keys ``table_id`` had evicted.""" + out, state.evicted_key_chunks, state.evicted_tid_chunks = _pop_chunked_keys( + state.evicted_key_chunks, state.evicted_tid_chunks, table_id, state + ) + return out + + +def _pop_state_erased_keys(state: DynamicEmbTableState, table_id: int) -> torch.Tensor: + """Pop (return + clear) the unique keys an explicit erase removed.""" + out, state.erased_key_chunks, state.erased_tid_chunks = _pop_chunked_keys( + state.erased_key_chunks, state.erased_tid_chunks, table_id, state + ) return out @@ -1148,7 +1379,7 @@ def _insert_key_values( score_out_flat: Optional[torch.Tensor] = None if state.no_eviction_next_index is not None: score_out_flat = torch.empty(n, dtype=torch.int64, device=unique_keys.device) - if state.evicted_item_mode == EvictedItemMode.RETAIN_KEY: + if EvictedItemMode.RETAIN_KEY in state.evicted_item_mode: ( indices, num_evicted, @@ -1660,6 +1891,267 @@ def _load_key_values( store_to_flat_single_table(state, indices, table_id, values) +# --------------------------------------------------------------------------- +# Replay – write an incremental_dump delta back into a table +# --------------------------------------------------------------------------- + + +def _slice(t: Optional[torch.Tensor], start: int, end: int) -> Optional[torch.Tensor]: + """Slice an optional column, keeping ``None`` as ``None``.""" + return None if t is None else t[start:end] + + +def _replay_write_values( + state: DynamicEmbTableState, + table_id: int, + rows: torch.Tensor, + embeddings: torch.Tensor, + optimizer_states: Optional[torch.Tensor], + keep_optimizer: torch.Tensor, + content: ReplayContent, +) -> None: + """Write a replayed row's value columns, per :class:`ReplayContent`. + + A value row is an embedding followed by the optimizer state, and the two are + written together or not at all -- ``store_to_flat_single_table`` copies from + the row base, so there is no way to write the tail without the head. That + shapes the three cases: + + - **Both requested.** The delta carries the whole row; write it as it was + dumped. Nothing is inferred. + - **Embedding only** (the default for a serving replica). Rows that already + belonged to this same key keep their optimizer state, because writing just + the embedding columns leaves the tail untouched. Every other row is new to + this key, so its tail still holds the previous occupant's moments and has + to be reset to ``initial_optim_state``. + - **Optimizer state only.** Every row must already belong to this key -- the + caller has checked -- so the current embedding is read back and rewritten + unchanged ahead of the new tail. + """ + if rows.numel() == 0: + return + emb_dim_cfg = state.table_emb_dims_cpu[table_id] + optstate_dim = state.optimizer.get_state_dim(emb_dim_cfg) + embeddings = embeddings.to(dtype=state.emb_dtype) + want_emb = ReplayContent.EMBEDDING in content + want_opt = ReplayContent.OPTIMIZER_STATE in content and optstate_dim > 0 + + if not want_emb: + # Keep what is there: the caller guarantees every row already holds this + # key, so its embedding is the one the replica should go on serving. + embeddings = load_from_flat_single_table(state, rows, table_id)[:, :emb_dim_cfg] + + if optstate_dim == 0: + store_to_flat_single_table(state, rows, table_id, embeddings) + return + + if want_opt: + if optimizer_states is None: + raise ValueError( + f"replay_increment: table {table_id} keeps optimizer state per " + "row, but the delta carries none. Drop " + "ReplayContent.OPTIMIZER_STATE, or replay a delta dumped from a " + "table with the same optimizer." + ) + ckpt_dim = state.optimizer.get_ckpt_state_dim(emb_dim_cfg) + if optimizer_states.dim() != 2 or optimizer_states.size(1) != ckpt_dim: + raise ValueError( + f"replay_increment: table {table_id} dumps {ckpt_dim} " + f"optimizer-state column(s) per row, but the delta carries a " + f"block of shape {tuple(optimizer_states.shape)}." + ) + # Back to the runtime width the fused value row expects. + opt = pad_optimizer_states_from_checkpoint( + state.optimizer, + emb_dim_cfg, + optimizer_states.to(device=embeddings.device), + state.initial_optim_state, + state.emb_dtype, + embeddings.device, + ) + store_to_flat_single_table( + state, rows, table_id, torch.cat([embeddings, opt], dim=-1) + ) + return + + keep = keep_optimizer + if bool(keep.any()): + store_to_flat_single_table(state, rows[keep], table_id, embeddings[keep]) + fresh = torch.logical_not(keep) + if bool(fresh.any()): + fresh_emb = embeddings[fresh] + opt_states = torch.full( + (fresh_emb.size(0), optstate_dim), + state.initial_optim_state, + dtype=state.emb_dtype, + device=fresh_emb.device, + ) + store_to_flat_single_table( + state, + rows[fresh], + table_id, + torch.cat([fresh_emb, opt_states], dim=-1), + ) + + +def _replay_at_slots( + state: DynamicEmbTableState, + table_id: int, + keys: torch.Tensor, + embeddings: torch.Tensor, + optimizer_states: Optional[torch.Tensor], + scores: Optional[torch.Tensor], + slot_index: torch.Tensor, + timestamp: int, + insert_score: int, + content: ReplayContent, +) -> int: + """Write keys at the exact slots they held in the source table. + + Every key must land: the caller has already established that this table's + layout matches the source's, so a slot that does not fall inside its key's + home bucket means the two have diverged anyway and the replay would silently + lose that key. Raises rather than writing a partial result. + + Returns the number of keys written (always all of them, or it raised). + """ + n = keys.numel() + device = state.device + if n == 0: + return 0 + tids = torch.full((n,), table_id, dtype=torch.int64, device=device) + + if state.no_eviction_next_index is not None: + # NO_EVICTION packs key slot (high 32) and value row (low 32); every + # other strategy uses one index for both. + key_slots = (slot_index >> 32) & 0xFFFFFFFF + value_rows = slot_index & 0xFFFFFFFF + else: + key_slots = slot_index + value_rows = slot_index + + # The score words are written with the key, so they are chosen here: either + # restored from the delta or minted as if the key had just been inserted. + if ReplayContent.SCORE in content and scores is not None: + score_block = _replay_score_block(state, table_id, scores, timestamp) + else: + score_block = _fresh_score_block( + state, table_id, value_rows, timestamp, insert_score + ) + status, same_key = state.key_index_map.scatter_keys_at_slots( + keys, tids, key_slots, score_block + ) + placed = status >= 0 + num_unplaced = int(n - placed.sum().item()) + if num_unplaced: + example = int(keys[torch.logical_not(placed)][0].item()) + raise RuntimeError( + f"replay_increment: {num_unplaced} of {n} keys could not be written " + f"at their source slot on table {table_id} (e.g. key {example}). " + "Either the slot falls outside its key's home bucket -- where no " + "lookup would ever probe it, meaning the target's hash layout " + "differs from the source's despite matching metadata -- or the slot " + "was held by another writer, meaning something ran against this " + "table concurrently with the replay. Replaying past either would " + "silently drop those keys." + ) + + if ReplayContent.EMBEDDING not in content: + # Without an embedding to write, a row that did not already hold this key + # would keep the previous occupant's vector and serve it under the new + # key. Refuse rather than corrupt. + stale = torch.logical_not(same_key) + num_stale = int(stale.sum().item()) + if num_stale: + example = int(keys[stale][0].item()) + raise ValueError( + f"replay_increment: {num_stale} of {n} keys on table {table_id} " + f"do not already occupy their target row (e.g. key {example}), " + "so there is no embedding to keep. ReplayContent.EMBEDDING may " + "only be omitted for a replica already aligned with its source." + ) + _replay_write_values( + state, table_id, value_rows, embeddings, optimizer_states, same_key, content + ) + if state.no_eviction_next_index is not None: + # Keep the auto-increment counter ahead of every restored row so a later + # native insert cannot hand out a row that is already in use. + next_row = int(value_rows.max().item()) + 1 + cur = int(state.no_eviction_next_index[table_id].item()) + if next_row > cur: + state.no_eviction_next_index[table_id] = next_row + state.no_eviction_next_index_dev[table_id] = next_row + return n + + +def _replay_state_increment( + state: DynamicEmbTableState, + table_id: int, + keys: torch.Tensor, + values: torch.Tensor, + optimizer_states: Optional[torch.Tensor], + scores: Optional[torch.Tensor], + slot_index: torch.Tensor, + timestamp: int, + insert_score: int, + content: ReplayContent, +) -> ReplayStats: + """Replay one delta batch into a single table state (one storage tier).""" + stats = ReplayStats() + n = keys.numel() + if n == 0: + return stats + device = state.device + keys = keys.to(device=device, dtype=state.key_index_map.key_type) + embeddings = values.to(device=device, dtype=state.emb_dtype) + slots = slot_index.to(device=device, dtype=torch.int64) + stats.upserted = _replay_at_slots( + state, + table_id, + keys, + embeddings, + optimizer_states, + scores, + slots, + timestamp, + insert_score, + content, + ) + return stats + + +def _erase_state_keys( + state: DynamicEmbTableState, + table_id: int, + keys: torch.Tensor, + mode: EvictedItemMode = EvictedItemMode.DISCARD, +) -> int: + """Erase keys from one table state. Returns how many were actually present. + + The erase kernel reports the slot it cleared, or ``-1`` for a key that was + not there, so the count is the real number removed rather than the number + asked for -- a delta can legitimately name keys this replica never held. + + *mode* is this call's, not the table's: retaining an erase costs only a copy + of keys already in hand, so there is nothing to configure up front and each + caller can decide whether its removals are worth reporting. + """ + n = keys.numel() + if n == 0: + return 0 + device = state.device + keys = keys.to(device=device, dtype=state.key_index_map.key_type) + tids = torch.full((n,), table_id, dtype=torch.int64, device=device) + indices = state.key_index_map.erase(keys, tids, return_indices=True) + removed = indices >= 0 + if EvictedItemMode.RETAIN_KEY in mode: + # Only the keys that were really there: a caller may name keys this + # table never held, and reporting those as removed would make a replica + # replay removals that never happened. + _append_erased(state, keys[removed], table_id) + return int(removed.sum().item()) + + # --------------------------------------------------------------------------- # DynamicEmbCache – Cache interface (key-only find / insert_and_evict) # --------------------------------------------------------------------------- @@ -1952,10 +2444,18 @@ def pop_evicted_keys(self, table_id: int) -> torch.Tensor: """Return + clear this rank's unique evicted keys retained for ``table_id``. DynamicEmbStorage is always a last tier, so retained keys live on its - state. Returns an empty tensor when evicted_item_mode is DISCARD or nothing - was evicted for the table since the last pop.""" + state. Returns an empty tensor when the mode does not retain evictions or + nothing was evicted for the table since the last pop.""" return _pop_state_evicted_keys(self._state, table_id) + def pop_erased_keys(self, table_id: int) -> torch.Tensor: + """Return + clear the keys an explicit erase removed from ``table_id``. + + Separate from :meth:`pop_evicted_keys` because the two mean different + things downstream: an eviction is reproduced by whoever takes the slot, + an erase leaves the slot to nobody and has to be replayed as a removal.""" + return _pop_state_erased_keys(self._state, table_id) + # -- Storage interface -- def find( @@ -2136,12 +2636,21 @@ def incremental_dump( table_id: int, threshold: int, pg: Optional[dist.ProcessGroup], - ) -> Tuple[Tensor, Tensor, Tensor]: - """Dump keys, embeddings and slot_index for one table (score >= threshold). + timestamp: Optional[int] = None, + ) -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Tensor]: + """Dump one table's matched rows (score >= threshold). Multi-rank: all_gather so the result is concatenated from all ranks. - ``slot_index`` is the packed key-slot/value-row for precise replay (see - :func:`_encode_slot_index`), column-aligned with keys/values.""" + + Returns ``(keys, values, slot_index, optimizer_states, scores)``, all + column-aligned. ``values`` is embeddings only and ``optimizer_states`` + the rest of the same row (``None`` when the optimizer keeps no per-row + state), so concatenating the two reproduces the stored row. + ``slot_index`` is the packed key-slot/value-row used by replay (see + :func:`_encode_slot_index`). ``scores`` is the ``[N, num_scores]`` block + in logical column order with timestamp columns held as an age relative + to *timestamp*, which defaults to this device's clock.""" + timestamp = device_timestamp() if timestamp is None else timestamp state = self._state if state.options_list[table_id].dist_type == "continuous": raise NotImplementedError( @@ -2157,6 +2666,8 @@ def incremental_dump( all_keys: List[Tensor] = [] all_values: List[Tensor] = [] all_slots: List[Tensor] = [] + all_opts: List[Optional[Tensor]] = [] + all_scores: List[Tensor] = [] for s in states_to_dump: keys, named_scores, indices = s.key_index_map.incremental_dump( {s.incremental_score_name: threshold}, @@ -2164,51 +2675,163 @@ def incremental_dump( return_index=True, table_id=table_id, ) - emb_dim = s.table_emb_dims_cpu[table_id] scores_batch = named_scores[s.incremental_score_name] flat_rows = _flat_row_indices_from_slots_and_scores( s, indices, scores_batch ) values = load_from_flat_single_table(s, flat_rows, table_id) - value = values[:, :emb_dim].to(dtype=s.emb_dtype) + value, opt = _split_value_row(s, table_id, values) key = keys.to(s.device) if keys.device.type != "cuda" else keys slot = _encode_slot_index(s, indices, flat_rows) + score_block = _dump_score_block( + s, table_id, indices, scores_batch, timestamp + ) if not do_multi_rank_gather: value = value.cpu() key = key.cpu() if key.is_cuda else key slot = slot.cpu() + score_block = score_block.cpu() + opt = opt.cpu() if opt is not None else None all_keys.append(key) all_values.append(value) all_slots.append(slot) + all_opts.append(opt) + all_scores.append(score_block) device_for_gather = state.device emb_dim_t = state.table_emb_dims_cpu[table_id] + # Checkpoint width, matching what _split_value_row emits: an empty + # batch still has to agree on column count with the non-empty ranks it + # is all_gathered against. + opt_dim_t = ( + state.optimizer.get_ckpt_state_dim(emb_dim_t) + if state.table_value_dims_cpu[table_id] > emb_dim_t + else 0 + ) + num_scores = state.key_index_map.num_scores_ + has_opt = bool(all_opts) and all_opts[0] is not None if all_keys: keys_cat = torch.cat(all_keys) values_cat = torch.cat(all_values, dim=0) slots_cat = torch.cat(all_slots) + scores_cat = torch.cat(all_scores, dim=0) + opts_cat = torch.cat(all_opts, dim=0) if has_opt else None else: - if do_multi_rank_gather: - keys_cat = torch.empty(0, dtype=torch.int64, device=device_for_gather) - values_cat = torch.empty( - 0, emb_dim_t, dtype=state.emb_dtype, device=device_for_gather - ) - slots_cat = torch.empty(0, dtype=torch.int64, device=device_for_gather) - else: - keys_cat = torch.empty(0, dtype=torch.int64, device="cpu") - values_cat = torch.empty(0, emb_dim_t, dtype=state.emb_dtype) - slots_cat = torch.empty(0, dtype=torch.int64, device="cpu") + dev = device_for_gather if do_multi_rank_gather else "cpu" + keys_cat = torch.empty(0, dtype=torch.int64, device=dev) + values_cat = torch.empty(0, emb_dim_t, dtype=state.emb_dtype, device=dev) + slots_cat = torch.empty(0, dtype=torch.int64, device=dev) + scores_cat = torch.empty(0, num_scores, dtype=SCORE_TYPE, device=dev) + opts_cat = ( + torch.empty(0, opt_dim_t, dtype=state.emb_dtype, device=dev) + if opt_dim_t > 0 + else None + ) if do_multi_rank_gather: - keys_cat = keys_cat.to(device_for_gather) - values_cat = values_cat.to(device_for_gather) - slots_cat = slots_cat.to(device_for_gather) - keys_cat, values_cat, slots_cat = _all_gather_dumped_keys_values( - keys_cat, values_cat, slots_cat, pg + # Every rank runs the same table config, so ``opts_cat is None`` + # agrees across ranks and the column list below is the same length + # everywhere -- an asymmetric all_gather would hang. + columns = [keys_cat, values_cat, slots_cat, scores_cat] + if opts_cat is not None: + columns.append(opts_cat) + gathered = _all_gather_dumped_columns( + [c.to(device_for_gather) for c in columns], pg ) + keys_cat, values_cat, slots_cat, scores_cat = gathered[:4] + opts_cat = gathered[4] if opts_cat is not None else None elif keys_cat.device.type == "cuda": keys_cat = keys_cat.cpu() values_cat = values_cat.cpu() slots_cat = slots_cat.cpu() - return keys_cat, values_cat, slots_cat + scores_cat = scores_cat.cpu() + opts_cat = opts_cat.cpu() if opts_cat is not None else None + return keys_cat, values_cat, slots_cat, opts_cat, scores_cat + + # -- Replay -- + + def replay_increment( + self, + table_id: int, + keys: Tensor, + values: Tensor, + optimizer_states: Optional[Tensor], + scores: Optional[Tensor], + slot_index: Tensor, + insert_score: int, + content: ReplayContent = ReplayContent.ALL, + timestamp: Optional[int] = None, + ) -> ReplayStats: + """Write an ``incremental_dump`` delta back into one table. + + Every key is written at the exact slot and value row it held in the + source; a key that cannot be placed there raises. The table is never + grown: an expansion rehashes, which would move every key's home bucket + and invalidate the source slots. + + Args: + table_id: logical table within this storage. + keys / values / optimizer_states / scores: the delta's columns for + this table, row aligned. The last two may be ``None`` when + *content* does not ask for them. + slot_index: the source's packed key-slot/value-row. + insert_score: the score a restored key gets in every non-recency + score word when *content* omits ``SCORE``. + content: which parts of each row to write; see + :class:`ReplayContent`. + timestamp: reference clock. Rebases the delta's timestamp ages when + ``SCORE`` is requested, and stamps restored keys as + freshly-inserted when it is not. Defaults to this device's clock. + + + Writing is batched at ``threads_in_wave`` keys -- one full pass over the + device, since the scatter kernel runs one thread per key -- which both + keeps the machine busy and bounds the temporary buffers a replay needs + (embeddings + optimizer/score blocks + masks) independently of how large + the delta is. Not a knob: it is a property of the GPU, not of the call, + and ``flush_cache`` sizes its own batches the same way. + + Returns: + :class:`ReplayStats` with ``upserted`` set; ``erased`` / ``skipped`` + are the caller's to fill. + + Raises: + RuntimeError: a key could not be placed at its source slot. + """ + stats = ReplayStats() + n = keys.numel() + if n == 0: + return stats + timestamp = device_timestamp() if timestamp is None else timestamp + batch_size = self._state.threads_in_wave + for start in range(0, n, batch_size): + end = min(start + batch_size, n) + stats.merge( + _replay_state_increment( + self._state, + table_id, + keys[start:end], + values[start:end], + _slice(optimizer_states, start, end), + _slice(scores, start, end), + slot_index[start:end], + timestamp, + insert_score, + content, + ) + ) + return stats + + def erase_keys( + self, + table_id: int, + keys: Tensor, + mode: EvictedItemMode = EvictedItemMode.DISCARD, + ) -> int: + """Erase keys from the table (used to replay a delta's removals). + + Returns how many of them were actually present. *mode* decides whether + the keys actually removed are recorded for :meth:`pop_erased_keys`. + """ + return _erase_state_keys(self._state, table_id, keys, mode) # -- Export -- @@ -2289,7 +2912,10 @@ def __init__( optimizer: BaseDynamicEmbeddingOptimizer, ): # Only the host tier is a last tier here: the HBM tier spills its - # evictions into the host tier (insert_and_evict), so it never retains. + # evictions into the host tier (insert_and_evict), so it never evicts for + # real and has nothing to retain. Erases are not covered by this setting + # at all -- ``erase_keys`` runs against both tiers and passes its own + # mode down, so a key erased out of HBM is recorded just the same. self._hbm = create_table_state(hbm_options, optimizer) self._host = create_table_state( host_options, @@ -2370,6 +2996,22 @@ def pop_evicted_keys(self, table_id: int) -> torch.Tensor: nothing was evicted for the table since the last pop.""" return _pop_state_evicted_keys(self._host, table_id) + def pop_erased_keys(self, table_id: int) -> torch.Tensor: + """Return + clear the keys an explicit erase removed from ``table_id``. + + Drained from **both** tiers, unlike :meth:`pop_evicted_keys`: an erase + removes a key from whichever tier holds it, and the tiers hold disjoint + key sets, so concatenating them is a union rather than double counting. + See :meth:`DynamicEmbStorage.pop_erased_keys` for why erases and + evictions are kept in separate buffers.""" + drained = [_pop_state_erased_keys(s, table_id) for s in self.tables] + parts = [p for p in drained if p.numel() > 0] + if not parts: + return torch.empty( + 0, dtype=self._host.key_index_map.key_type, device=self._host.device + ) + return torch.unique(torch.cat(parts)) + # -- Two-tier find (with values) -- def find( @@ -2596,18 +3238,35 @@ def incremental_dump( table_id: int, threshold: int, pg: Optional[dist.ProcessGroup], - ) -> Tuple[Tensor, Tensor, Tensor]: - """Dump keys, embeddings and slot_index for one table (score >= threshold). - - Multi-rank: all_gather so the result is concatenated from all ranks. - ``slot_index`` is the packed key-slot/value-row for precise replay (see - :func:`_encode_slot_index`), column-aligned with keys/values.""" + timestamp: Optional[int] = None, + ) -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Tensor]: + """Dump one table's matched rows (score >= threshold), both tiers. + + Same columns as :meth:`DynamicEmbStorage.incremental_dump`; the two + tiers' results are concatenated, and ``slot_index`` bit 63 tags which + tier each key came from. + + Raises: + NotImplementedError: the table is sharded with + ``dist_type="continuous"`` (replay cannot reconstruct a key's + owning rank), or the two tiers disagree on ``num_scores`` (their + score blocks are concatenated, so they must share a width). + """ + timestamp = device_timestamp() if timestamp is None else timestamp states_to_dump = self.tables if states_to_dump[0].options_list[table_id].dist_type == "continuous": raise NotImplementedError( "incremental_dump with slot_index does not support dist_type " "'continuous' (replay cannot reconstruct the owning rank from a key)." ) + num_scores = states_to_dump[0].key_index_map.num_scores_ + if any(s.key_index_map.num_scores_ != num_scores for s in states_to_dump): + # The tiers' score blocks are concatenated into one column-aligned + # result, so they must agree on width (mirrors the dump/load guard). + raise NotImplementedError( + "incremental_dump is not supported for HybridStorage whose tiers " + "have different score-word counts." + ) do_multi_rank_gather = ( pg is not None and dist.is_initialized() @@ -2616,6 +3275,8 @@ def incremental_dump( all_keys = [] all_values = [] all_slots = [] + all_opts: List[Optional[Tensor]] = [] + all_scores = [] # tier index into self.tables: 0 = HBM tier, 1 = host tier. slot_index # packs this tier bit (bit 63) so replay can tell the two tiers apart. for tier, s in enumerate(states_to_dump): @@ -2625,51 +3286,176 @@ def incremental_dump( return_index=True, table_id=table_id, ) - emb_dim = s.table_emb_dims_cpu[table_id] scores_batch = named_scores[s.incremental_score_name] flat_rows = _flat_row_indices_from_slots_and_scores( s, indices, scores_batch ) values = load_from_flat_single_table(s, flat_rows, table_id) - value = values[:, :emb_dim].to(dtype=s.emb_dtype) + value, opt = _split_value_row(s, table_id, values) key = keys.to(s.device) if keys.device.type != "cuda" else keys slot = _encode_slot_index(s, indices, flat_rows, tier=tier) + score_block = _dump_score_block( + s, table_id, indices, scores_batch, timestamp + ) if not do_multi_rank_gather: value = value.cpu() key = key.cpu() if key.is_cuda else key slot = slot.cpu() + score_block = score_block.cpu() + opt = opt.cpu() if opt is not None else None all_keys.append(key) all_values.append(value) all_slots.append(slot) + all_opts.append(opt) + all_scores.append(score_block) device_for_gather = states_to_dump[0].device emb_dim_t = states_to_dump[0].table_emb_dims_cpu[table_id] + # Checkpoint width -- see the note in DynamicEmbStorage.incremental_dump. + st0 = states_to_dump[0] + opt_dim_t = ( + st0.optimizer.get_ckpt_state_dim(emb_dim_t) + if st0.table_value_dims_cpu[table_id] > emb_dim_t + else 0 + ) + has_opt = bool(all_opts) and all_opts[0] is not None if all_keys: keys_cat = torch.cat(all_keys) values_cat = torch.cat(all_values, dim=0) slots_cat = torch.cat(all_slots) + scores_cat = torch.cat(all_scores, dim=0) + opts_cat = torch.cat(all_opts, dim=0) if has_opt else None else: - if do_multi_rank_gather: - keys_cat = torch.empty(0, dtype=torch.int64, device=device_for_gather) - values_cat = torch.empty( - 0, emb_dim_t, dtype=self.embedding_dtype(), device=device_for_gather - ) - slots_cat = torch.empty(0, dtype=torch.int64, device=device_for_gather) - else: - keys_cat = torch.empty(0, dtype=torch.int64, device="cpu") - values_cat = torch.empty(0, emb_dim_t, dtype=self.embedding_dtype()) - slots_cat = torch.empty(0, dtype=torch.int64, device="cpu") + dev = device_for_gather if do_multi_rank_gather else "cpu" + dt = self.embedding_dtype() + keys_cat = torch.empty(0, dtype=torch.int64, device=dev) + values_cat = torch.empty(0, emb_dim_t, dtype=dt, device=dev) + slots_cat = torch.empty(0, dtype=torch.int64, device=dev) + scores_cat = torch.empty(0, num_scores, dtype=SCORE_TYPE, device=dev) + opts_cat = ( + torch.empty(0, opt_dim_t, dtype=dt, device=dev) + if opt_dim_t > 0 + else None + ) if do_multi_rank_gather: - keys_cat = keys_cat.to(device_for_gather) - values_cat = values_cat.to(device_for_gather) - slots_cat = slots_cat.to(device_for_gather) - keys_cat, values_cat, slots_cat = _all_gather_dumped_keys_values( - keys_cat, values_cat, slots_cat, pg + # Every rank runs the same table config, so ``opts_cat is None`` + # agrees across ranks and the column list below is the same length + # everywhere -- an asymmetric all_gather would hang. + columns = [keys_cat, values_cat, slots_cat, scores_cat] + if opts_cat is not None: + columns.append(opts_cat) + gathered = _all_gather_dumped_columns( + [c.to(device_for_gather) for c in columns], pg ) + keys_cat, values_cat, slots_cat, scores_cat = gathered[:4] + opts_cat = gathered[4] if opts_cat is not None else None elif keys_cat.device.type == "cuda": keys_cat = keys_cat.cpu() values_cat = values_cat.cpu() slots_cat = slots_cat.cpu() - return keys_cat, values_cat, slots_cat + scores_cat = scores_cat.cpu() + opts_cat = opts_cat.cpu() if opts_cat is not None else None + return keys_cat, values_cat, slots_cat, opts_cat, scores_cat + + # -- Replay -- + + def replay_increment( + self, + table_id: int, + keys: Tensor, + values: Tensor, + optimizer_states: Optional[Tensor], + scores: Optional[Tensor], + slot_index: Tensor, + insert_score: int, + content: ReplayContent = ReplayContent.ALL, + timestamp: Optional[int] = None, + ) -> ReplayStats: + """Write an ``incremental_dump`` delta back into one table. + + Each key is routed to the tier it came from -- ``slot_index`` bit 63 + selects HBM (0) or host (1) -- and written at that tier's slot. Nothing + is promoted or spilled between tiers: the point is to reproduce the + source's layout, and the source already decided which tier each key + belongs to. + + Args: + table_id: logical table within this storage. + keys / values / optimizer_states / scores: the delta's columns for + this table, row aligned. The last two may be ``None`` when + *content* does not ask for them. + slot_index: the source's tier-tagged key slot (bit 63 = tier). + insert_score: the score a restored key gets in every non-recency + score word when *content* omits ``SCORE``. + content: which parts of each row to write; see + :class:`ReplayContent`. + timestamp: reference clock. Rebases the delta's timestamp ages when + ``SCORE`` is requested, and stamps restored keys as + freshly-inserted when it is not. Defaults to this device's clock. + + + Batched at ``threads_in_wave`` keys per tier, as + :meth:`DynamicEmbStorage.replay_increment` explains. + + Returns: + :class:`ReplayStats` with ``upserted`` summed over both tiers. + + Raises: + RuntimeError: a key could not be placed at its source slot. + """ + stats = ReplayStats() + n = keys.numel() + if n == 0: + return stats + timestamp = device_timestamp() if timestamp is None else timestamp + + # Both tiers sit on the same device, so one size covers them. + batch_size = self.tables[0].threads_in_wave + tier_bit = slot_index.to(torch.int64) < 0 # bit 63 set -> host tier + key_slots = slot_index.to(torch.int64) & 0x7FFFFFFFFFFFFFFF + for tier, state in enumerate(self.tables): + sel = tier_bit if tier else torch.logical_not(tier_bit) + if not bool(sel.any()): + continue + sel_keys, sel_values = keys[sel], values[sel] + sel_opts = optimizer_states[sel] if optimizer_states is not None else None + sel_scores = scores[sel] if scores is not None else None + sel_slots = key_slots[sel] + sel_n = sel_keys.numel() + for start in range(0, sel_n, batch_size): + end = min(start + batch_size, sel_n) + stats.merge( + _replay_state_increment( + state, + table_id, + sel_keys[start:end], + sel_values[start:end], + _slice(sel_opts, start, end), + _slice(sel_scores, start, end), + sel_slots[start:end], + timestamp, + insert_score, + content, + ) + ) + return stats + + def erase_keys( + self, + table_id: int, + keys: Tensor, + mode: EvictedItemMode = EvictedItemMode.DISCARD, + ) -> int: + """Erase keys from both tiers (a key may live in either). + + Returns how many were actually present. Summing across tiers is right + here, not double counting: the tiers hold disjoint key sets, so a key + erased from one is reported as absent by the other. *mode* applies to + both, so a key is recorded wherever it was found. + """ + removed = 0 + for state in self.tables: + removed += _erase_state_keys(state, table_id, keys, mode) + return removed # -- Dump: write host first, then append HBM -- diff --git a/corelib/dynamicemb/dynamicemb/scored_hashtable.py b/corelib/dynamicemb/dynamicemb/scored_hashtable.py index 7bf66ee96..28c2c7f55 100644 --- a/corelib/dynamicemb/dynamicemb/scored_hashtable.py +++ b/corelib/dynamicemb/dynamicemb/scored_hashtable.py @@ -39,6 +39,7 @@ table_insert_collect_evicted, table_lookup, table_partition, + table_scatter_keys_at_slots, table_scatter_score_blocks, table_update_counter_with_layout, ) @@ -277,19 +278,40 @@ def uint64_to_int64(x): return x if x < (1 << 63) else x - (1 << 64) -def murmur3_hash_64bits(key: int) -> int: - """ """ - k = key & 0xFFFFFFFFFFFFFFFF +def murmur3_fmix64(keys): + """MurmurHash3's 64-bit finalizer -- the host twin of ``murmur3_fmix64`` in + ``src/murmur_hash.cuh``, and the only host copy of it. - k ^= k >> 33 - k = (k * 0xFF51AFD7ED558CCD) & 0xFFFFFFFFFFFFFFFF + Takes a Python int or anything ``np.asarray`` accepts, and returns the same + shape as ``uint64``. Only the avalanche step is here: callers narrow the + result themselves -- modulo the world size to pick an owning rank, masked to + a non-negative int64 to pick a hash bucket -- exactly as the two device + callers do. - k ^= k >> 33 - k = (k * 0xC4CEB9FE1A85EC53) & 0xFFFFFFFFFFFFFFFF + A key is a bit pattern here, not a magnitude, so a negative one is + reinterpreted rather than rejected -- ``numpy`` refuses to build a ``uint64`` + from a negative Python int, while ``astype`` on an array wraps the way C + would. Wrapping is likewise the algorithm and not an error for the + multiplies, which numpy is silent about for arrays but warns about for + scalars; the warning is turned off rather than left to depend on the input's + shape. + """ + if isinstance(keys, (int, np.integer)): + k = np.uint64(int(keys) & 0xFFFFFFFFFFFFFFFF) + else: + k = np.asarray(keys).astype(np.uint64, copy=False) + with np.errstate(over="ignore"): + k = k ^ (k >> np.uint64(33)) + k = k * np.uint64(0xFF51AFD7ED558CCD) + k = k ^ (k >> np.uint64(33)) + k = k * np.uint64(0xC4CEB9FE1A85EC53) + k = k ^ (k >> np.uint64(33)) + return k - k ^= k >> 33 - return k +def murmur3_hash_64bits(key: int) -> int: + """Scalar :func:`murmur3_fmix64`, for constants computed once at import.""" + return int(murmur3_fmix64(key)) class LinearBucketTable(ScoredHashTable): @@ -871,12 +893,25 @@ def erase( self, keys: torch.Tensor, table_ids: torch.Tensor, - ) -> None: + return_indices: bool = False, + ) -> Optional[torch.Tensor]: """ Erase Keys Args: table_ids: int32 tensor of same length as keys, identifying which logical table each key belongs to. + return_indices: also report which keys were actually present. Costs + one int64 tensor of ``keys.numel()``; callers that only need the + erase itself should leave it False. + + Returns: + ``None`` unless *return_indices*, else the slot each key was erased + from, or ``-1`` where the key was not in the table. """ + indices = ( + torch.empty(keys.numel(), dtype=self.index_type, device=keys.device) + if return_indices + else None + ) table_erase( self.table_storage_, self.table_bucket_offsets_, @@ -884,8 +919,10 @@ def erase( self.bucket_sizes, keys, table_ids, + indices=indices, num_scores=self.num_scores_, ) + return indices def load( self, @@ -1418,6 +1455,53 @@ def scatter_score_blocks( self.key_type_, ) + def scatter_keys_at_slots( + self, + keys: torch.Tensor, + table_ids: torch.Tensor, + slots: torch.Tensor, + scores: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Place ``keys`` (with their full score blocks) at exact ``slots``. + + The precise write-back behind ``replay_increment``: ``slots`` are + table-relative flat slot indices taken from the source table's + ``incremental_dump``, and are only usable when this table's capacity and + bucket_capacity match the source's -- a key can only ever be found inside + its own home bucket. + + Args: + keys: 1-D key tensor on device. + table_ids: int64 tensor aligned with ``keys``. + slots: int64 table-relative flat slot per key. + scores: ``[N, num_scores_]`` score words (physical column order). + + Returns: + (status, same_key): ``status[i]`` is the slot written, or ``-1`` when + the key was NOT placed because its home bucket does not contain that + slot. The caller treats that as a fatal layout divergence and raises + -- the key cannot be reached from anywhere else, so writing it + elsewhere would silently lose it. + ``same_key[i]`` is True when the slot already held this very key, so + its value row (and optimizer state) can be kept. + """ + if scores.dim() != 2 or scores.size(1) != self.num_scores_: + raise ValueError( + f"scatter_keys_at_slots expects [{keys.numel()}, " + f"{self.num_scores_}] scores, got {tuple(scores.shape)}" + ) + return table_scatter_keys_at_slots( + self.table_storage_, + self.table_bucket_offsets_, + self.bucket_capacity_, + self.bucket_sizes, + keys, + table_ids, + slots, + scores, + num_scores=self.num_scores_, + ) + def capacity(self, table_id: Optional[int] = None) -> int: """ Return the capacity (main + overflow) of the table, or a specific logical table. diff --git a/corelib/dynamicemb/dynamicemb/types.py b/corelib/dynamicemb/dynamicemb/types.py index 68184a579..4937ef1fb 100644 --- a/corelib/dynamicemb/dynamicemb/types.py +++ b/corelib/dynamicemb/dynamicemb/types.py @@ -146,6 +146,41 @@ class CopyMode(enum.Enum): VALUE = "value" +@dataclass +class ReplayStats: + """What a ``replay_increment`` did to one table. + + Replay is all-or-nothing for the keys this rank owns: each of them is written + back at the slot and value row it occupied in the source table, or the call + raises. ``skipped`` counts the rest of the delta -- the keys another rank + owns, which this rank deliberately ignores. + + Attributes + ---------- + upserted : int + Keys written back at their source slot. + erased : int + Keys removed before the upsert, counted as *actually removed* -- a delta + can name keys this replica never held. Counts only the delta's + ``erased_keys``; an eviction is reproduced by overwriting its slot and + costs no removal. + skipped : int + Delta keys this rank does not own (row-wise sharding) and so ignored. + Always 0 on a single rank, where the whole delta is this rank's. + """ + + upserted: int = 0 + erased: int = 0 + skipped: int = 0 + + def merge(self, other: "ReplayStats") -> "ReplayStats": + """Accumulate *other* into self (used to fold per-batch / per-tier runs).""" + self.upserted += other.upserted + self.erased += other.erased + self.skipped += other.skipped + return self + + # make it standalone to avoid recursive references. class Storage(abc.ABC, Generic[OptionsT, OptimizerT]): @abc.abstractmethod diff --git a/corelib/dynamicemb/src/murmur_hash.cuh b/corelib/dynamicemb/src/murmur_hash.cuh new file mode 100644 index 000000000..7bd6bfee0 --- /dev/null +++ b/corelib/dynamicemb/src/murmur_hash.cuh @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace dyn_emb { + +// MurmurHash3's 64-bit finalizer, and nothing else: the avalanche step that +// spreads a key's bits before something narrows them down. +// +// Two callers narrow it differently -- picking an owning rank takes it modulo +// the world size, picking a hash bucket masks it to a non-negative int64 -- and +// those tails belong to the callers, not here. Only the avalanche is shared, +// which is why this header pulls in nothing but . +// +// Anything that has to agree with this on the host (``murmur3_fmix64`` in +// dynamicemb/scored_hashtable.py) is a translation of exactly this function. +__host__ __device__ __forceinline__ uint64_t murmur3_fmix64(uint64_t key) { + key ^= key >> 33; + key *= UINT64_C(0xff51afd7ed558ccd); + key ^= key >> 33; + key *= UINT64_C(0xc4ceb9fe1a85ec53); + key ^= key >> 33; + return key; +} + +} // namespace dyn_emb diff --git a/corelib/dynamicemb/src/sparse_block_bucketize_features.cu b/corelib/dynamicemb/src/sparse_block_bucketize_features.cu index 99a05a8b7..df8940dc5 100644 --- a/corelib/dynamicemb/src/sparse_block_bucketize_features.cu +++ b/corelib/dynamicemb/src/sparse_block_bucketize_features.cu @@ -19,6 +19,7 @@ All rights reserved. # SPDX-License-Identifier: Apache-2.0 #include #include +#include "murmur_hash.cuh" #include "sparse_block_bucketize_features_utils.h" #include "utils.h" @@ -29,12 +30,7 @@ namespace dyn_emb { // Hash function for load-balanced key distribution // Uses MurmurHash3 finalizer for good avalanche properties __forceinline__ __device__ uint64_t hash_key(uint64_t key) { - key ^= key >> 33; - key *= 0xff51afd7ed558ccdULL; - key ^= key >> 33; - key *= 0xc4ceb9fe1a85ec53ULL; - key ^= key >> 33; - return key; + return murmur3_fmix64(key); } __forceinline__ __device__ int atomicAdd(int *address, int val) { diff --git a/corelib/dynamicemb/src/table_operation/insert.cu b/corelib/dynamicemb/src/table_operation/insert.cu index 92f104255..e47b168a8 100644 --- a/corelib/dynamicemb/src/table_operation/insert.cu +++ b/corelib/dynamicemb/src/table_operation/insert.cu @@ -458,4 +458,63 @@ void table_scatter_score_blocks(at::Tensor table_storage, DEMB_CUDA_KERNEL_LAUNCH_CHECK(); } +// Write (key, score words) at exact table-relative slots -- the write-back +// behind replay_increment. Returns (status, same_key): status is the slot +// written, or -1 when the slot is not inside the key's home bucket and the key +// therefore could not be placed (the caller raises: the key would be +// unreachable). same_key tells whether the slot already held this very key (so +// its value row / optimizer state can be kept). +std::tuple table_scatter_keys_at_slots( + at::Tensor table_storage, at::Tensor table_bucket_offsets, + int64_t bucket_capacity, at::Tensor bucket_sizes, at::Tensor keys, + at::Tensor table_ids, at::Tensor slots, at::Tensor scores, + int64_t num_scores) { + + int64_t num_total = keys.size(0); + auto status = torch::empty( + {num_total}, + torch::TensorOptions().dtype(torch::kInt64).device(keys.device())); + auto same_key = torch::empty( + {num_total}, + torch::TensorOptions().dtype(torch::kBool).device(keys.device())); + if (num_total == 0) + return {status, same_key}; + + auto key_type = get_data_type(keys); + auto bucket_sizes_ = get_pointer(bucket_sizes); + auto table_ids_ptr = table_ids.data_ptr(); + auto table_bucket_offsets_ptr = table_bucket_offsets.data_ptr(); + // Callers hand scores over as int64 (SCORE_TYPE) or uint64; reinterpret rather + // than convert -- score words are opaque bit patterns. + at::Tensor score_vals = scores.contiguous(); + if (score_vals.scalar_type() != torch::kUInt64) + score_vals = score_vals.view(torch::kUInt64); + const ScoreType *score_ptr = score_vals.data_ptr(); + + auto stream = at::cuda::getCurrentCUDAStream().stream(); + constexpr int BLOCK_SIZE = 256; + + DISPATCH_KEY_TYPE(key_type, KeyType, [&] { + auto keys_ = get_pointer(keys); + int64_t total_size = + sizeof(KeyType) + sizeof(DigestType) + num_scores * sizeof(ScoreType); + int64_t bucket_bytes = bucket_capacity * total_size; + int64_t num_buckets = + table_storage.numel() * table_storage.element_size() / bucket_bytes; + + using Bucket = LinearBucket; + using Table = LinearBucketTable; + auto table = Table(reinterpret_cast(table_storage.data_ptr()), + num_buckets, bucket_capacity, num_scores); + + scatter_keys_at_slots_kernel + <<<(num_total + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>( + table, table_bucket_offsets_ptr, bucket_sizes_, num_total, keys_, + table_ids_ptr, slots.data_ptr(), score_ptr, + status.data_ptr(), same_key.data_ptr()); + }); + DEMB_CUDA_KERNEL_LAUNCH_CHECK(); + return {status, same_key}; +} + } // namespace dyn_emb diff --git a/corelib/dynamicemb/src/table_operation/kernels.cuh b/corelib/dynamicemb/src/table_operation/kernels.cuh index cdb3f467a..a89c50a2f 100644 --- a/corelib/dynamicemb/src/table_operation/kernels.cuh +++ b/corelib/dynamicemb/src/table_operation/kernels.cuh @@ -1064,4 +1064,102 @@ __global__ void scatter_score_blocks_kernel( *b.scores(it, k) = vals[i * ns + k]; } +// Write (key, score words) at an EXACT table-relative slot -- the "precise +// write-back" primitive behind replay_increment. Unlike insert(), the slot is +// dictated by the caller (it comes from the source table's dump) instead of +// being probed for. +// +// A key is only ever findable inside its own home bucket, so a slot from a +// source table is usable only when this table's capacity/bucket_capacity match +// the source's. That is verified per key here; a key whose slot falls outside +// its home bucket is reported with status -1 and left untouched, and the caller +// raises -- placing it anywhere else would make it unreachable. +// +// status[i] = the slot written, or -1 when this key was not placed. +// same_key[i] = true when the slot's previous occupant was this very key, i.e. +// the value row already belongs to it (its optimizer state can be +// kept); false means the row must be re-initialised. +template +__global__ void scatter_keys_at_slots_kernel( + Table table, int64_t const *__restrict__ table_bucket_offsets, + int *__restrict__ bucket_sizes, int64_t batch, + typename Table::KeyType const *__restrict__ input_keys, + int64_t const *__restrict__ table_ids, int64_t const *__restrict__ slots, + ScoreType const *__restrict__ scores, int64_t *__restrict__ status, + bool *__restrict__ same_key) { + + using KeyType = typename Table::KeyType; + using Bucket = typename Table::BucketType; + using Iter = typename Bucket::Iterator; + + auto tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + + for (int64_t i = tid; i < batch; i += gridDim.x * blockDim.x) { + status[i] = -1; + same_key[i] = false; + + KeyType key = input_keys[i]; + int64_t slot = slots[i]; + if (!Bucket::is_valid(key) || slot < 0) + continue; + + int64_t bc = table.bucket_capacity(); + int64_t t_id = table_ids[i]; + int64_t bkt_begin = table_bucket_offsets[t_id]; + int64_t bkt_end = table_bucket_offsets[t_id + 1]; + int64_t table_cap = (bkt_end - bkt_begin) * bc; + if (table_cap <= 0 || slot >= table_cap) + continue; + + // The slot must land in the key's home bucket, otherwise no lookup would + // ever probe it. This is the per-key half of the layout-compatibility check. + int64_t hashcode = Table::hash(key); + int64_t home_local = (hashcode % table_cap) / bc; + if (home_local != slot / bc) + continue; + + int64_t bucket_id = bkt_begin + home_local; + Bucket bucket = table[bucket_id]; + int64_t ns = table.num_scores(); + Iter target = static_cast(slot % bc); + + // The key is not searched for first: a replayed table is written only by + // replay from one source, so a key is either absent or already at the slot + // the source gives, never at some other slot in this bucket. Probing to rule + // that out would cost a probe per key to find nothing. + // + // Take the target slot under the bucket's lock, as any writer must. One + // attempt, no spin: a replay has the table to itself -- it cannot run + // alongside a forward pass, whose prefetch state holds row indices this very + // write invalidates -- and a delta's slots are distinct, one slot having + // held one key in the source, so no other thread in this launch is aiming + // here either. Finding the slot locked, or changed underfoot, therefore does + // not mean "wait"; it means that assumption did not hold and this replay is + // not safe. Leave status[i] at -1 and let the host raise. + auto key_slot = + reinterpret_cast(bucket.keys(target)); + KeyType old_key = key_slot->load(cuda::std::memory_order_relaxed); + // Checked before the CAS, not folded into it: locking a slot that already + // reads LockedKey would swap it for itself and report success. + if (old_key == static_cast(Bucket::LockedKey)) + continue; + KeyType expected_key = old_key; + if (!bucket.try_lock(target, expected_key)) + continue; + + for (int64_t k = 0; k < ns; ++k) + *bucket.scores(target, k) = scores[i * ns + k]; + *bucket.digests(target) = Bucket::key_to_digest(key); + // Key last, with release: whoever reads this key afterwards sees the score + // words and digest that belong to it. + bucket.unlock(target, key); + + if (!Bucket::is_valid(old_key)) + atomicAdd(bucket_sizes + bucket_id, 1); + + status[i] = slot; + same_key[i] = (old_key == key); + } +} + } // namespace dyn_emb diff --git a/corelib/dynamicemb/src/table_operation/table.cu b/corelib/dynamicemb/src/table_operation/table.cu index 0399a1568..81ea3ca61 100644 --- a/corelib/dynamicemb/src/table_operation/table.cu +++ b/corelib/dynamicemb/src/table_operation/table.cu @@ -173,6 +173,15 @@ void bind_table_operation(py::module &m) { py::arg("num_scores"), py::arg("bkt_begin"), py::arg("slots"), py::arg("values"), py::arg("key_dtype")); + m.def("table_scatter_keys_at_slots", &dyn_emb::table_scatter_keys_at_slots, + "write (key, score words) at exact table-relative slots; returns " + "(status, same_key) where status is the slot written or -1 when the " + "key's home bucket does not contain that slot", + py::arg("table_storage"), py::arg("table_bucket_offsets"), + py::arg("bucket_capacity"), py::arg("bucket_sizes"), py::arg("keys"), + py::arg("table_ids"), py::arg("slots"), py::arg("scores"), + py::arg("num_scores") = 1); + m.def("bucketize_keys", &dyn_emb::bucketize_keys, "bucketize input keys into a dense tensor, and return the output keys, " "buckets offset, inverse indices. num_buckets must equal " diff --git a/corelib/dynamicemb/src/table_operation/table.cuh b/corelib/dynamicemb/src/table_operation/table.cuh index f5b9ea8f8..24e151e52 100644 --- a/corelib/dynamicemb/src/table_operation/table.cuh +++ b/corelib/dynamicemb/src/table_operation/table.cuh @@ -146,6 +146,16 @@ void table_scatter_score_blocks(at::Tensor table_storage, int64_t bkt_begin, at::Tensor slots, at::Tensor values, torch::Dtype key_dtype); +// Write-back for replay_increment: place (key, score words) at the given +// table-relative slots. Returns (status, same_key) -- status is the slot +// written, or -1 when the key's home bucket does not contain that slot (the +// caller raises), same_key marks slots that already held the same key. +std::tuple table_scatter_keys_at_slots( + at::Tensor table_storage, at::Tensor table_bucket_offsets, + int64_t bucket_capacity, at::Tensor bucket_sizes, at::Tensor keys, + at::Tensor table_ids, at::Tensor slots, at::Tensor scores, + int64_t num_scores = 1); + std::vector table_partition(at::Tensor storage, std::vector dtypes, int64_t bucket_capacity, diff --git a/corelib/dynamicemb/src/table_operation/types.cuh b/corelib/dynamicemb/src/table_operation/types.cuh index ff84cee79..7d7eeca6b 100644 --- a/corelib/dynamicemb/src/table_operation/types.cuh +++ b/corelib/dynamicemb/src/table_operation/types.cuh @@ -20,6 +20,8 @@ All rights reserved. # SPDX-License-Identifier: Apache-2.0 #include #include #include + +#include "../murmur_hash.cuh" #include #include #include @@ -122,13 +124,9 @@ struct LinearBucket { static constexpr uint64_t ReserveKeyMask = UINT64_C(0xFFFFFFFFFFFFFFFC); static __device__ __forceinline__ int64_t hash(uint64_t key) { - uint64_t k = key; - k ^= k >> 33; - k *= UINT64_C(0xff51afd7ed558ccd); - k ^= k >> 33; - k *= UINT64_C(0xc4ceb9fe1a85ec53); - k ^= k >> 33; - return static_cast(k & INT64_MAX); // avoid overflow + // Shared avalanche, local tail: the mask keeps the result a non-negative + // int64 so it can index a bucket without overflowing. + return static_cast(murmur3_fmix64(key) & INT64_MAX); } static __device__ __forceinline__ KeyType empty_key() { return EmptyKey; } diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py index 534281781..31c449a1a 100644 --- a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py +++ b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py @@ -29,7 +29,12 @@ DynamicEmbTableOptions, ) from dynamicemb.dynamicemb_config import ScoreStrategy -from dynamicemb.incremental_dump import get_score, incremental_dump, set_score +from dynamicemb.incremental_dump import ( + get_score, + incremental_dump, + replay_increment, + set_score, +) from dynamicemb.planner import ( DynamicEmbeddingEnumerator, DynamicEmbeddingShardingPlanner, @@ -106,6 +111,7 @@ def get_planner( batch_size: int, multi_hot_sizes: List[int], device, + training: bool = True, ): dict_const = {} for i in range(len(table_names)): @@ -121,6 +127,7 @@ def get_planner( dynamicemb_options=DynamicEmbTableOptions( global_hbm_for_values=1024**3, score_strategy=score_strategies[i], + training=training, initializer_args=DynamicEmbInitializerArgs( mode=DynamicEmbInitializerMode.DEBUG, ), @@ -427,3 +434,103 @@ def test_incremental_dump_api( dump_keys = dump_keys % 100000 dump_vals = dump_vals.to(dump_keys.dtype) assert torch.all(dump_keys.unsqueeze(1).expand(-1, dim) == dump_vals) + + +def _build_sharded_model( + table_names, eb_configs, score_strategies, batch, dim, device, training=True +): + """One row-wise sharded EmbeddingCollection over dynamic embedding tables. + + ``training=False`` builds a serving replica: the table option (not + ``nn.Module.training``) is what drops the optimizer, so the value row is the + embedding alone and there is no per-row state for a replay to restore. + """ + ebc = torchrec.EmbeddingCollection(device=torch.device("meta"), tables=eb_configs) + planner = get_planner( + table_names, + eb_configs, + [True] * len(table_names), + score_strategies, + batch, + [1] * len(table_names), + device, + training=training, + ) + sharder = DynamicEmbeddingCollectionSharder( + fused_params={"optimizer": EmbOptimType.SGD, "learning_rate": 0.1}, + use_index_dedup=False, + ) + plan = planner.collective_plan(ebc, [sharder], dist.GroupMember.WORLD) + return DistributedModelParallel( + module=ebc, device=device, sharders=[sharder], plan=plan + ) + + +@pytest.mark.parametrize("table_num, dim, local_batch", [(2, 8, 64)]) +def test_replay_increment_distributed( + request, table_num, dim, local_batch, backend_session +): + """A globally gathered delta replayed into a second sharded model: every rank + keeps only the keys it owns, and together they restore the whole delta.""" + local_rank = int(os.environ["LOCAL_RANK"]) + dist.get_world_size() + device = torch.device(f"cuda:{local_rank}") + prefix_path = "model" + + table_names = [f"t_{t}" for t in range(table_num)] + eb_configs = [ + torchrec.EmbeddingConfig( + name=table_names[t], + embedding_dim=dim, + num_embeddings=BATCH_SIZE_PER_DUMP * 8, + feature_names=[f"f_{t}"], + ) + for t in range(table_num) + ] + strategies = [DynamicEmbScoreStrategy.TIMESTAMP] * table_num + + train_model = _build_sharded_model( + table_names, eb_configs, strategies, local_batch, dim, device + ) + # A serving replica: no optimizer state in the value row, and eval forwards + # that neither insert nor admit, so everything it holds came from the delta. + serve_model = _build_sharded_model( + table_names, eb_configs, strategies, local_batch, dim, device, training=False + ) + serve_model.eval() + + unique_indices = [set({}) for _ in table_names] + sparse_feature = generate_sparse_feature( + [f"f_{t}" for t in range(table_num)], + [1] * table_num, + local_batch, + unique_indices, + [True] * table_num, + [BATCH_SIZE_PER_DUMP * 8] * table_num, + ) + trained = {k: v.values().clone() for k, v in train_model(sparse_feature).items()} + + pg = intra_and_cross_node_pg()[0] + deltas = incremental_dump(train_model, 0, pg) + stats = replay_increment(serve_model, deltas, pg) + + delta = deltas[prefix_path] + for j, table_name in enumerate(delta.table_names): + s = stats[prefix_path][table_name] + total = delta.keys[j].numel() + assert s.upserted + s.skipped == total + + # Row-wise sharding is a partition: the ranks' shares must sum to the + # whole delta, with no key landing on two ranks. + counted = torch.tensor([s.upserted], dtype=torch.int64, device=device) + dist.all_reduce(counted, group=pg) + assert int(counted.item()) == total, (int(counted.item()), total) + + # What the replica is for: serving the same batch and returning what the + # source would. Comparing forward outputs rather than a dump of the replica + # exercises the path a replica is actually used through -- input dist, + # lookup, pooling -- instead of re-reading the table the replay just wrote. + served = {k: v.values() for k, v in serve_model(sparse_feature).items()} + assert set(served.keys()) == set(trained.keys()) + for feature in trained: + torch.testing.assert_close(served[feature], trained[feature]) diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_incremental_dump.sh b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_incremental_dump.sh index 380b32527..282e3b7a7 100644 --- a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_incremental_dump.sh +++ b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_incremental_dump.sh @@ -2,5 +2,6 @@ set -e pytest test/unit_tests/incremental_dump/test_batched_dynamicemb_tables.py -s +pytest test/unit_tests/incremental_dump/test_replay_increment.py -s torchrun --nproc_per_node=1 -m pytest test/unit_tests/incremental_dump/test_distributed_dynamicemb.py -s torchrun --nproc_per_node=8 -m pytest test/unit_tests/incremental_dump/test_distributed_dynamicemb.py -s diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py new file mode 100644 index 000000000..ad74f8910 --- /dev/null +++ b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py @@ -0,0 +1,686 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Module-level tests for ``replay_increment``: writing an ``incremental_dump`` +delta back into another model's tables.""" + +from typing import List, Optional + +import pytest +import torch +from dynamicemb import ( + BATCH_SIZE_PER_DUMP, + DynamicEmbCheckMode, + DynamicEmbPoolingMode, + DynamicEmbScoreStrategy, + DynamicEmbTableOptions, + EmbOptimType, + EvictedItemMode, +) +from dynamicemb.batched_dynamicemb_tables import BatchedDynamicEmbeddingTablesV2 +from dynamicemb.dynamicemb_config import ReplayContent + +TABLE_NAME = "t_0" +DIM = 8 +DEFAULT_CAPACITY = BATCH_SIZE_PER_DUMP * 2 + + +@pytest.fixture +def current_device(): + assert torch.cuda.is_available() + return torch.cuda.current_device() + + +def make_model( + current_device: int, + score_strategy=DynamicEmbScoreStrategy.TIMESTAMP, + max_capacity: int = DEFAULT_CAPACITY, + init_capacity: Optional[int] = None, + bucket_capacity: int = 128, + caching: bool = False, + local_hbm_for_values: int = 1024**3, + evicted_item_mode: EvictedItemMode = EvictedItemMode.DISCARD, + optimizer: EmbOptimType = EmbOptimType.SGD, + **opt_params, +) -> BatchedDynamicEmbeddingTablesV2: + options = DynamicEmbTableOptions( + index_type=torch.int64, + embedding_dtype=torch.float32, + device_id=current_device, + dim=DIM, + max_capacity=max_capacity, + bucket_capacity=bucket_capacity, + safe_check_mode=DynamicEmbCheckMode.IGNORE, + # A cache only exists when the values do not fit in HBM + # (``total_memory > local_hbm``), so a caching test has to shrink this -- + # ``caching=True`` alone silently yields a plain HBM_ONLY table. + local_hbm_for_values=local_hbm_for_values, + score_strategy=score_strategy, + caching=caching, + evicted_item_mode=evicted_item_mode, + # Only when asked: pinning init to max disables rehash growth, which the + # other tests do not want. + **({} if init_capacity is None else {"init_capacity": init_capacity}), + ) + return BatchedDynamicEmbeddingTablesV2( + table_options=[options], + output_dtype=torch.float32, + table_names=[TABLE_NAME], + feature_table_map=[0], + pooling_mode=DynamicEmbPoolingMode.SUM, + use_index_dedup=False, + optimizer=optimizer, + **opt_params, + ) + + +def feature_from_keys(keys: List[int], device: torch.device): + """One SUM-pooled feature, one key per bag.""" + indices = torch.tensor(keys, dtype=torch.int64, device=device) + offsets = torch.arange(0, len(keys) + 1, dtype=torch.int64, device=device) + return indices, offsets + + +def touch(model, keys: List[int], device: torch.device, backward: bool = False): + indices, offsets = feature_from_keys(keys, device) + out = model(indices, offsets) + if backward: + out.sum().backward() + torch.cuda.synchronize() + return out + + +def dump_all(model): + """Everything in the table: threshold 0 matches every score.""" + return model.incremental_dump({TABLE_NAME: 0}) + + +def sorted_view(keys: torch.Tensor, *columns: torch.Tensor): + """Sort a delta's columns by key so two dumps can be compared row by row.""" + order = torch.argsort(keys) + return (keys[order],) + tuple(c[order] for c in columns) + + +def _mask_delta(delta, mask: torch.Tensor, table_id: int = 0) -> None: + """Keep only ``mask``'s rows, across every column-aligned list at once. + + A delta has five per-key columns plus ``meta["slot_index"]``; masking them + one by one in a test is how a stale column slips through unnoticed. + """ + delta.keys[table_id] = delta.keys[table_id][mask] + delta.values[table_id] = delta.values[table_id][mask] + delta.scores[table_id] = delta.scores[table_id][mask] + if delta.optimizer_states[table_id] is not None: + delta.optimizer_states[table_id] = delta.optimizer_states[table_id][mask] + delta.meta[table_id]["slot_index"] = delta.meta[table_id]["slot_index"][mask] + + +def export_optimizer_state(model, table_id: int = 0) -> dict: + """{key: optimizer state row} straight out of the storage.""" + device = torch.device(f"cuda:{torch.cuda.current_device()}") + out = {} + for keys, _, opt_states, _ in model._storage.export_keys_values( + device, 65536, table_id + ): + if opt_states is None: + continue + for k, o in zip(keys.cpu().tolist(), opt_states.cpu()): + out[k] = o.clone() + return out + + +@pytest.mark.parametrize( + "score_strategy", + [ + DynamicEmbScoreStrategy.TIMESTAMP, + DynamicEmbScoreStrategy.STEP, + DynamicEmbScoreStrategy.LFU, + (DynamicEmbScoreStrategy.TIMESTAMP, DynamicEmbScoreStrategy.LFU), + # NO_EVICTION exercises the packed slot_index (key slot in the high 32 + # bits, value row in the low 32) and the auto-increment row counter. + DynamicEmbScoreStrategy.NO_EVICTION, + ], +) +def test_replay_round_trip(current_device, score_strategy): + """A delta replayed into an identically configured model reproduces it + key-for-key, value-for-value, at the same slots.""" + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device, score_strategy=score_strategy) + dst = make_model(current_device, score_strategy=score_strategy) + + keys = list(range(1001, 1301)) + touch(src, keys, device) + delta = dump_all(src) + assert set(delta.keys[0].tolist()) == set(keys) + + stats = dst.replay_increment(delta)[TABLE_NAME] + assert stats.upserted == len(keys), stats + assert stats.skipped == 0 + + replayed = dump_all(dst) + src_keys, src_vals = sorted_view(delta.keys[0], delta.values[0]) + dst_keys, dst_vals = sorted_view(replayed.keys[0], replayed.values[0]) + assert torch.equal(src_keys, dst_keys) + torch.testing.assert_close(src_vals, dst_vals) + + # Precise replay means the same slot, so the source's slot_index round-trips. + _, src_slots = sorted_view(delta.keys[0], delta.meta[0]["slot_index"]) + _, dst_slots = sorted_view(replayed.keys[0], replayed.meta[0]["slot_index"]) + assert torch.equal(src_slots, dst_slots) + + # A lookup on the replica returns the replayed embeddings. Sort the lookup + # into the same key order as the delta rather than permuting the delta back + # into lookup order -- one comparison, no dependence on `keys` being sorted. + out = touch(dst, keys, device) + _, out_vals = sorted_view(torch.tensor(keys), out.cpu()) + torch.testing.assert_close(out_vals, src_vals) + + +@pytest.mark.parametrize( + "optimizer", [EmbOptimType.SGD, EmbOptimType.EXACT_ROWWISE_ADAGRAD] +) +def test_dump_splits_the_value_row(current_device, optimizer): + """``values`` and ``optimizer_states`` are the two halves of one stored row. + + They are dumped as separate tensors, so the thing worth pinning down is that + nothing is lost or duplicated between them: the widths must add up to the + table's value row, and a table with no per-row state must say so with + ``None`` rather than a zero-width tensor. + """ + device = torch.device(f"cuda:{current_device}") + kwargs = ( + {"learning_rate": 0.1, "eps": 1e-8} + if optimizer == EmbOptimType.EXACT_ROWWISE_ADAGRAD + else {"learning_rate": 0.1} + ) + model = make_model(current_device, optimizer=optimizer, **kwargs) + keys = list(range(1001, 1051)) + touch(model, keys, device, backward=True) + + delta = dump_all(model) + n = delta.keys[0].numel() + assert n == len(keys) + assert delta.values[0].shape == (n, DIM) + + value_dim = model._storage.value_dim(0) + opt = delta.optimizer_states[0] + if value_dim == DIM: + assert opt is None, "a table with no per-row state must report None" + return + + # The width is the file checkpoint's, not the runtime row's. Rowwise Adagrad + # differs between the two -- 16 bytes reserved per row, one scalar used -- so + # asserting against value_dim - DIM would pin the padding instead. + optimizer_obj = model._storage._state.optimizer + ckpt_dim = optimizer_obj.get_ckpt_state_dim(DIM) + assert opt is not None and opt.shape == (n, ckpt_dim) + if optimizer == EmbOptimType.EXACT_ROWWISE_ADAGRAD: + assert ckpt_dim < value_dim - DIM, ( + "this parametrisation is meant to cover the case where the dumped " + "width is narrower than the runtime row" + ) + + +def test_dump_scores_are_column_aligned(current_device): + """``scores`` carries every score word, one row per key, in logical order.""" + device = torch.device(f"cuda:{current_device}") + strategy = (DynamicEmbScoreStrategy.TIMESTAMP, DynamicEmbScoreStrategy.LFU) + model = make_model(current_device, score_strategy=strategy) + hot = list(range(1001, 1051)) + cold = list(range(2001, 2051)) + touch(model, cold, device) + for _ in range(5): + touch(model, hot, device) + + delta = dump_all(model) + keys, scores = sorted_view(delta.keys[0], delta.scores[0]) + assert scores.shape == (keys.numel(), len(strategy)) + + # Column 1 is the LFU frequency, carried verbatim: the keys touched six + # times must outrank those touched once. Column 0 is a timestamp held as an + # age, so it is not compared against a raw clock here. + hot_set = set(hot) + is_hot = torch.tensor([int(k) in hot_set for k in keys.tolist()]) + assert int(scores[is_hot, 1].min()) > int(scores[~is_hot, 1].max()) + + +@pytest.mark.parametrize( + "content, replica_ranks_like_source", + [(ReplayContent.EMBEDDING, False), (ReplayContent.ALL, True)], +) +def test_replay_scores_follow_the_content_flag( + current_device, content, replica_ranks_like_source +): + """Whether the replica inherits the source's ranking is ``SCORE``'s call. + + Same setup either way -- hot keys accessed six times, cold keys once, and a + threshold drawn between them that splits the source's table exactly in half. + Without ``SCORE`` every restored key carries the same fresh score, so the + same threshold splits nothing; with it the frequency column is carried + verbatim and the split is reproduced. The pair is what makes the limitation + legible: neither half means much alone. + """ + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device, score_strategy=DynamicEmbScoreStrategy.LFU) + dst = make_model(current_device, score_strategy=DynamicEmbScoreStrategy.LFU) + + hot = list(range(1001, 1051)) + cold = list(range(2001, 2051)) + touch(src, cold, device) + for _ in range(5): # drive the two groups' LFU counts apart + touch(src, hot, device) + + all_keys = set(hot) | set(cold) + threshold = 3 # between the cold keys' 1 access and the hot keys' 6 + assert set(src.incremental_dump({TABLE_NAME: threshold}).keys[0].tolist()) == set( + hot + ), "the source must rank hot keys above cold ones, or there is nothing to inherit" + + dst.replay_increment(dump_all(src), content=content) + assert set(dump_all(dst).keys[0].tolist()) == all_keys + + split = set(dst.incremental_dump({TABLE_NAME: threshold}).keys[0].tolist()) + if replica_ranks_like_source: + assert split == set(hot), "the replica must rank the way the source did" + else: + # Every key scored alike, so the threshold takes all of them or none -- + # which of the two depends on where the fresh score falls, and that is + # an implementation detail this deliberately does not pin down. + assert split in (set(), all_keys), f"got {len(split)} of {len(all_keys)}" + + +def test_replay_spans_multiple_batches(current_device): + """Replay is chunked, and every column has to be chunked the same way. + + The columns are not all the same rank -- ``keys`` / ``slot_index`` are 1-D, + ``values`` / ``optimizer_states`` / ``scores`` are ``[N, ...]`` -- so the + slicing has to cut dimension 0 in each case, or a batch would pair one key's + row with another key's slot. Nothing else here notices: batches are sized to + the device (``threads_in_wave``, hundreds of thousands of keys), so every + other test in this file fits in one and never crosses a boundary. Shrink it + so the chunking actually runs. + """ + device = torch.device(f"cuda:{current_device}") + src = make_model( + current_device, + optimizer=EmbOptimType.EXACT_ROWWISE_ADAGRAD, + learning_rate=0.1, + initial_accumulator_value=0.0, + ) + dst = make_model( + current_device, + optimizer=EmbOptimType.EXACT_ROWWISE_ADAGRAD, + learning_rate=0.1, + initial_accumulator_value=0.0, + ) + + keys = list(range(1001, 1101)) + touch(src, keys, device, backward=True) + delta = dump_all(src) + assert delta.values[0].dim() == 2 and delta.scores[0].dim() == 2 + assert delta.optimizer_states[0] is not None + + # 7 keys per batch: not a divisor of 100, so the last batch is short too. + dst._storage._state.threads_in_wave = 7 + dst.replay_increment(delta, content=ReplayContent.ALL) + + out = dump_all(dst) + src_k, src_v = sorted_view(delta.keys[0], delta.values[0]) + dst_k, dst_v = sorted_view(out.keys[0], out.values[0]) + assert torch.equal(src_k, dst_k) + torch.testing.assert_close(src_v, dst_v) + # Slicing the wrong dimension would still produce the right key set, so the + # per-key payload is what actually pins it down. + got, want = export_optimizer_state(dst), export_optimizer_state(src) + for k in keys: + torch.testing.assert_close(got[k], want[k]) + + +def test_replay_rejects_layout_mismatch(current_device): + """A target whose layout differs cannot take the source's slots, so the whole + table is rejected -- and nothing is written.""" + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device, max_capacity=DEFAULT_CAPACITY) + dst = make_model(current_device, max_capacity=DEFAULT_CAPACITY * 2) + + keys = list(range(1001, 1201)) + touch(src, keys, device) + delta = dump_all(src) + + with pytest.raises(ValueError, match="capacity mismatch"): + dst.replay_increment(delta) + + # Rejection must be total: no partial write is left behind. + assert dump_all(dst).keys[0].numel() == 0 + + +def test_replay_rejects_score_strategy_mismatch(current_device): + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device, score_strategy=DynamicEmbScoreStrategy.TIMESTAMP) + dst = make_model(current_device, score_strategy=DynamicEmbScoreStrategy.LFU) + touch(src, list(range(1001, 1051)), device) + delta = dump_all(src) + with pytest.raises(ValueError, match="score_strategy mismatch"): + dst.replay_increment(delta) + assert dump_all(dst).keys[0].numel() == 0, "rejection must be total" + + +def test_replay_accepts_swapped_score_order(current_device): + """Same two score words, opposite configured order, must still replay. + + The tuple order is only ever the checkpoint column order: ``(TIMESTAMP, LFU)`` + and ``(LFU, TIMESTAMP)`` are the same physical layout, so a source slot means + the same thing in both. With no scores in flight there is nothing left for + the order to break, and rejecting the pair would be over-strict. + """ + device = torch.device(f"cuda:{current_device}") + src = make_model( + current_device, + score_strategy=( + DynamicEmbScoreStrategy.TIMESTAMP, + DynamicEmbScoreStrategy.LFU, + ), + ) + dst = make_model( + current_device, + score_strategy=( + DynamicEmbScoreStrategy.LFU, + DynamicEmbScoreStrategy.TIMESTAMP, + ), + ) + keys = list(range(1001, 1051)) + touch(src, keys, device) + delta = dump_all(src) + + dst.replay_increment(delta) + src_keys, src_vals = sorted_view(delta.keys[0], delta.values[0]) + out = dump_all(dst) + dst_keys, dst_vals = sorted_view(out.keys[0], out.values[0]) + assert torch.equal(src_keys, dst_keys) + torch.testing.assert_close(src_vals, dst_vals) + + +def test_replay_rejects_missing_table_options(current_device): + """A delta without table_options must be rejected, not replayed with the + score-order / dim / dist_type checks quietly skipped.""" + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device) + dst = make_model(current_device) + touch(src, list(range(1001, 1051)), device) + delta = dump_all(src) + delta.meta[0].pop("table_options") + + with pytest.raises(ValueError, match="carries no 'table_options'"): + dst.replay_increment(delta) + assert dump_all(dst).keys[0].numel() == 0 + + +def test_replay_applies_erasures_and_ignores_evictions(current_device): + """``erased_keys`` is replayed; ``evicted_keys`` never is. + + An eviction needs no action on the replica -- the key that took the evicted + one's slot is in this very delta and overwrites it. An erase does, because + nothing takes over that slot. Feeding a non-empty ``evicted_keys`` here would + delete live keys if the two were ever confused. + """ + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device) + dst = make_model(current_device) + + keys = list(range(1001, 1101)) + touch(src, keys, device) + delta = dump_all(src) + dst.replay_increment(delta) + assert set(dump_all(dst).keys[0].tolist()) == set(keys) + + # A delta that erases half the keys and re-adds one of them. + dropped = keys[:50] + readmitted = dropped[0] + keep = keys[50:] + next_delta = dump_all(src) + mask = torch.tensor([int(k) == readmitted for k in next_delta.keys[0].tolist()]) + _mask_delta(next_delta, mask) + next_delta.erased_keys[0] = torch.tensor(dropped, dtype=torch.int64) + # The surviving keys, offered as evictions: replay must not touch them. + next_delta.evicted_keys[0] = torch.tensor(keep, dtype=torch.int64) + + stats = dst.replay_increment(next_delta)[TABLE_NAME] + assert stats.erased == len(dropped) + left = set(dump_all(dst).keys[0].tolist()) + assert readmitted in left, "an erased-then-readmitted key must survive" + assert left == set(keep) | { + readmitted + }, "evicted_keys must not have removed anything" + + # An absent list is not an error: a table may never have been erased from, + # and under row-wise sharding a rank may own none of the keys that were. + # ``None`` rather than an empty tensor, since a hand-built delta may do that. + next_delta.erased_keys[0] = None + next_delta.evicted_keys[0] = None + assert dst.replay_increment(next_delta)[TABLE_NAME].erased == 0 + assert set(dump_all(dst).keys[0].tolist()) == left + + +def test_replay_erased_counts_only_keys_actually_present(current_device): + """``erased`` reports removals that really happened, not removals asked for. + + A delta may name keys this replica never held -- e.g. it was built from a + wider source, or the key was already gone. Counting the request instead of + the effect would silently overstate what the replica did. + """ + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device) + dst = make_model(current_device) + + keys = list(range(1001, 1051)) + touch(src, keys, device) + delta = dump_all(src) + dst.replay_increment(delta) + + present = keys[:20] + absent = list(range(90001, 90031)) # never inserted anywhere + delta.erased_keys[0] = torch.tensor(present + absent, dtype=torch.int64) + # Drop the upsert side so the removals are not immediately undone. + _mask_delta(delta, torch.zeros(delta.keys[0].numel(), dtype=torch.bool)) + + stats = dst.replay_increment(delta)[TABLE_NAME] + assert stats.erased == len(present), ( + f"expected {len(present)} real removals out of " + f"{len(present) + len(absent)} requested, got {stats.erased}" + ) + assert set(dump_all(dst).keys[0].tolist()) == set(keys) - set(present) + + +def test_replay_optimizer_state(current_device): + """Without ``ReplayContent.OPTIMIZER_STATE``, the state is preserved for a key + already sitting at its target slot and initialised for one that is not.""" + device = torch.device(f"cuda:{current_device}") + src = make_model( + current_device, + optimizer=EmbOptimType.EXACT_ROWWISE_ADAGRAD, + learning_rate=0.1, + initial_accumulator_value=0.0, + ) + dst = make_model( + current_device, + optimizer=EmbOptimType.EXACT_ROWWISE_ADAGRAD, + learning_rate=0.1, + initial_accumulator_value=0.0, + ) + + keys = list(range(1001, 1101)) + touch(src, keys, device, backward=True) + delta = dump_all(src) + + # First replay: brand-new rows, so the optimizer state starts at its initial. + dst.replay_increment(delta, content=ReplayContent.EMBEDDING) + fresh_state = export_optimizer_state(dst) + assert fresh_state, "expected an optimizer state region for ROWWISE_ADAGRAD" + for key in keys: + assert float(fresh_state[key].abs().max()) == 0.0 + + # Train the replica so its optimizer state diverges from the initial value. + touch(dst, keys, device, backward=True) + trained_state = export_optimizer_state(dst) + assert any(float(trained_state[k].abs().max()) > 0.0 for k in keys) + + # Second replay of the same delta: every key is already in its target slot, + # so the embedding is overwritten but the optimizer state must survive. + dst.replay_increment(delta, content=ReplayContent.EMBEDDING) + after_state = export_optimizer_state(dst) + for key in keys: + torch.testing.assert_close(after_state[key], trained_state[key]) + + +def test_replay_content_restores_optimizer_state(current_device): + """With ``OPTIMIZER_STATE``, the source's state is written -- even onto a row + the key is taking over, where the default would have initialised it.""" + device = torch.device(f"cuda:{current_device}") + kwargs = dict( + optimizer=EmbOptimType.EXACT_ROWWISE_ADAGRAD, + learning_rate=0.1, + initial_accumulator_value=0.0, + ) + src = make_model(current_device, **kwargs) + dst = make_model(current_device, **kwargs) + + keys = list(range(1001, 1101)) + touch(src, keys, device, backward=True) + src_state = export_optimizer_state(src) + assert any(float(src_state[k].abs().max()) > 0.0 for k in keys), ( + "the source must have trained, or there is no state to distinguish " + "from the initial value" + ) + + # Brand-new rows on the target: without the flag these would be initialised. + dst.replay_increment(dump_all(src), content=ReplayContent.ALL) + got = export_optimizer_state(dst) + for key in keys: + torch.testing.assert_close(got[key], src_state[key]) + + +def test_replay_without_embedding_needs_aligned_rows(current_device): + """Omitting ``EMBEDDING`` is only legal for an already-aligned replica. + + A key landing on a row it does not already own has no embedding to keep, and + the row still holds the previous occupant's vector -- serving that under the + new key would be silent corruption, so replay refuses. + """ + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device) + dst = make_model(current_device) + keys = list(range(1001, 1051)) + touch(src, keys, device) + delta = dump_all(src) + + score_only = ReplayContent.SCORE + with pytest.raises(ValueError, match="do not already occupy their target row"): + dst.replay_increment(delta, content=score_only) + + # Align the replica first; then the same call is fine. + dst.replay_increment(delta) + dst.replay_increment(delta, content=score_only) + assert set(dump_all(dst).keys[0].tolist()) == set(keys) + + +def test_replay_does_not_leave_displaced_keys_in_the_cache(current_device): + """A key whose slot a delta key takes over must not survive in the cache. + + Writing a key at its source slot reproduces an eviction in the *storage* -- + the slot's new owner overwrites the old one. It does not reach the cache, + which is a separate index the slot write knows nothing about. So the + displaced key is in neither ``keys`` nor ``erased_keys``, and an + invalidation working from those two lists alone leaves it behind. + + Letting it survive is worse than a stale read: ``flush_cache`` pushes every + cached key back down, so the next dump would *resurrect* it into storage and + it would reappear in the dump -- which is what this asserts against. + """ + device = torch.device(f"cuda:{current_device}") + # One bucket, so any key may legally occupy any slot (see below), and room + # to spare so that nothing here depends on the table being full. + # 1024 rows * 8 dims * 4 bytes = 32 KiB of values; half in HBM leaves a real + # GPU cache in front of a host-backed store. + cfg = dict( + max_capacity=1024, + init_capacity=1024, + bucket_capacity=1024, + caching=True, + local_hbm_for_values=16384, + ) + src = make_model(current_device, **cfg) + dst = make_model(current_device, **cfg) + assert dst._cache is not None, "this test is meaningless without a cache" + + keys = list(range(1001, 1129)) + touch(src, keys, device) + delta = dump_all(src) + dst.replay_increment(delta) + assert set(dump_all(dst).keys[0].tolist()) == set(keys), "must start converged" + + # Pull one key into the replica's cache while it is still resident there. + victim = int(delta.keys[0][0].item()) + touch(dst, [victim], device) + + # Now hand the replica a delta in which a *different* key claims the + # victim's slot -- what the source produces when it evicts the victim and + # admits someone else in its place. + # + # Constructed rather than provoked: which key a table evicts, and whether it + # evicts at all, depends on the eviction policy, on ref-counter pinning and + # on how full the table happens to be -- none of which this test is about, + # and all of which vary by GPU. The single bucket makes the substitution + # legal, since a slot only has to lie in its key's home bucket. + taker = 9001 + delta.keys[0] = delta.keys[0].clone() + delta.keys[0][0] = taker + delta.evicted_keys[0] = torch.tensor([victim], dtype=delta.keys[0].dtype) + + dst.replay_increment(delta) + + # dump_all flushes the cache into storage first, so a surviving cached copy + # of the victim would reappear here. + left = set(dump_all(dst).keys[0].tolist()) + assert taker in left, "the taker must hold the slot it was given" + assert ( + victim not in left + ), f"key {victim} lost its slot to {taker} but survived in the cache" + + +def test_replay_with_caching(current_device): + """With a cache in front of the storage, a lookup after replay must see the + replayed values, not the stale cached ones. + + Only the cached layout is worth a case here: without a cache this is + test_replay_round_trip's closing lookup. The HBM budget has to be shrunk or + the values fit and no cache is built at all. + """ + device = torch.device(f"cuda:{current_device}") + src = make_model(current_device, caching=True, local_hbm_for_values=2048) + dst = make_model(current_device, caching=True, local_hbm_for_values=2048) + assert dst._cache is not None, "this test is meaningless without a cache" + + keys = list(range(1001, 1101)) + touch(src, keys, device) + delta = dump_all(src) + + # Warm the target's cache with different values for the same keys. + touch(dst, keys, device) + dst.replay_increment(delta) + + out = touch(dst, keys, device) + _, src_vals = sorted_view(delta.keys[0], delta.values[0]) + _, out_vals = sorted_view(torch.tensor(keys), out.cpu()) + torch.testing.assert_close(out_vals, src_vals) diff --git a/corelib/dynamicemb/test/unit_tests/retain_evicted_keys/test_insert_collect_evicted.py b/corelib/dynamicemb/test/unit_tests/retain_evicted_keys/test_insert_collect_evicted.py index 60b23c619..c2fb5ca55 100644 --- a/corelib/dynamicemb/test/unit_tests/retain_evicted_keys/test_insert_collect_evicted.py +++ b/corelib/dynamicemb/test/unit_tests/retain_evicted_keys/test_insert_collect_evicted.py @@ -336,6 +336,67 @@ def _storage_insert(storage, keys, dim, device): torch.cuda.synchronize() +def test_erased_and_evicted_use_separate_buffers(current_device): + """An erase and an eviction are recorded in different buffers. + + They mean different things to a consumer -- an eviction is the table running + out of room, an erase is somebody asking for the key to go, and only the + latter is a removal a replica has to perform for itself. Mixing them would + make a replay erase keys that were merely evicted, dropping live data. + """ + device = current_device + dim = 8 + # Same full-bucket overflow as test_storage_retain_end_to_end: 128 keys fill + # the single bucket, then 58 newer ones evict the 58 oldest. + storage = _retain_storage( + dim=dim, max_capacity=128, bucket_capacity=128, retain=True + ) + _storage_insert( + storage, torch.arange(1, 129, dtype=torch.int64, device=device), dim, device + ) + _storage_insert( + storage, torch.arange(1000, 1058, dtype=torch.int64, device=device), dim, device + ) + + evicted = storage.pop_evicted_keys(0) + assert evicted.numel() > 0, "a full-bucket overflow must have evicted something" + assert storage.pop_erased_keys(0).numel() == 0, "no erase happened yet" + + # Erase a key that is still resident, i.e. one that was NOT evicted. + resident = torch.tensor([1000], dtype=torch.int64, device=device) + assert not bool(torch.isin(resident, evicted).any()) + assert storage.erase_keys(0, resident, EvictedItemMode.RETAIN_KEY) == 1 + + erased = storage.pop_erased_keys(0) + assert erased.tolist() == resident.tolist() + # The erase must not have leaked into the eviction buffer, and both buffers + # are read-and-clear. + assert storage.pop_evicted_keys(0).numel() == 0 + assert storage.pop_erased_keys(0).numel() == 0 + + +def test_erase_retention_is_per_call_not_per_table(current_device): + """Whether an erase is recorded is the call's decision, not the table's. + + Retaining evictions has to be configured up front (it swaps the insert + kernel), but an erase already holds its keys, so nothing has to be prepared + -- which is why the same table can record one erase and not the next, and + why a table configured to DISCARD evictions can still report its erases. + """ + device = current_device + dim = 8 + storage = _retain_storage(dim=dim, retain=False) # DISCARD: no eviction retention + keys = torch.arange(1, 9, dtype=torch.int64, device=device) + _storage_insert(storage, keys, dim, device) + + assert storage.erase_keys(0, keys[:3], EvictedItemMode.DISCARD) == 3 + assert storage.pop_erased_keys(0).numel() == 0, "DISCARD must record nothing" + + assert storage.erase_keys(0, keys[3:6], EvictedItemMode.RETAIN_KEY) == 3 + assert storage.pop_erased_keys(0).tolist() == keys[3:6].tolist() + assert storage.pop_evicted_keys(0).numel() == 0 + + def test_storage_retain_end_to_end(current_device): """DynamicEmbStorage (a last tier) with evicted_item_mode=RETAIN_KEY: a full-bucket insert evicts; pop returns the unique evicted keys; those keys are gone from the diff --git a/corelib/dynamicemb/test/unit_tests/test_hash_roundrobin_kuairand.py b/corelib/dynamicemb/test/unit_tests/test_hash_roundrobin_kuairand.py index 3ebd5270a..fce9862a0 100644 --- a/corelib/dynamicemb/test/unit_tests/test_hash_roundrobin_kuairand.py +++ b/corelib/dynamicemb/test/unit_tests/test_hash_roundrobin_kuairand.py @@ -14,6 +14,13 @@ def hash_key_cpu(keys_np: np.ndarray) -> np.ndarray: + """Deliberately a second implementation, not an import. + + This oracle exists to check the CUDA bucketize kernel's key->rank mapping. + Calling the production ``murmur3_fmix64`` here would compare that host code + against itself and stop testing the device at all, so the duplication is the + point -- do not "unify" it away. + """ k = np.asarray(keys_np, dtype=np.uint64) k = k ^ (k >> np.uint64(33)) k = k * np.uint64(0xFF51AFD7ED558CCD) diff --git a/corelib/dynamicemb/test/unit_tests/test_hybrid_storage_export.py b/corelib/dynamicemb/test/unit_tests/test_hybrid_storage_export.py index 0656f2f6a..22a89f7d9 100644 --- a/corelib/dynamicemb/test/unit_tests/test_hybrid_storage_export.py +++ b/corelib/dynamicemb/test/unit_tests/test_hybrid_storage_export.py @@ -521,10 +521,21 @@ def test_hybrid_incremental_dump_slot_index_tier(): storage.insert(keys, tids, vals) torch.cuda.synchronize() - dump_keys, _, slot_index = storage.incremental_dump(0, 0, None) + ( + dump_keys, + dump_values, + slot_index, + dump_opts, + dump_scores, + ) = storage.incremental_dump(0, 0, None) assert dump_keys.numel() > 0 assert slot_index.numel() == dump_keys.numel() assert slot_index.dtype == torch.int64 + # values/optimizer_states split the same stored row, scores carry every word + assert dump_values.shape == (dump_keys.numel(), embedding_dim) + assert dump_scores.shape == (dump_keys.numel(), 1) + if dump_opts is not None: + assert dump_opts.size(0) == dump_keys.numel() # bit 63 = tier: the host tier sets it, so those entries are negative int64; # the HBM tier leaves it 0 (non-negative). Both tiers must be present. From d241b074381ce69d6b8fab9dc8ed56e225d87981 Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Mon, 7 Sep 2026 07:47:06 +0530 Subject: [PATCH 2/9] fix(dynamicemb): size replay's rank filter by the shard fan-out, not 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) --- corelib/dynamicemb/DynamicEmb_APIs.md | 8 ++++---- .../dynamicemb/docs/replay_increment_design.md | 2 +- .../dynamicemb/batched_dynamicemb_tables.py | 18 +++++++++--------- .../dynamicemb/dynamicemb/incremental_dump.py | 4 ---- .../test_distributed_dynamicemb.py | 2 +- 5 files changed, 15 insertions(+), 19 deletions(-) diff --git a/corelib/dynamicemb/DynamicEmb_APIs.md b/corelib/dynamicemb/DynamicEmb_APIs.md index 9e4ee324d..dddbd6fff 100644 --- a/corelib/dynamicemb/DynamicEmb_APIs.md +++ b/corelib/dynamicemb/DynamicEmb_APIs.md @@ -849,7 +849,7 @@ The same rule is enforced per key inside the kernel — a slot that does not lan *Scores and optimizer state are dumped but not yet replayed.* `DeltaDumpResult.scores` and `.optimizer_states` carry them, but `replay_increment` ignores both for now: a restored key is scored as if it had just been inserted into the target, and keeps its optimizer state only if it already occupies the target row. So the replica ranks its own future evictions by when it received each key rather than by how the source ranked it. Embeddings are unaffected — the two models can evict in different orders, but never disagree on the value of a key they both hold. NO_EVICTION is the exception and is exact: its score word is a value row, not a score. -*Sharding.* Replay keeps only the keys this rank owns, recomputing ownership from the key with **this** model's world size. A delta gathered over a process group (`incremental_dump(..., pg)`) holds the whole group's keys and so fans out correctly when the same delta is handed to every rank. A per-rank delta (`pg=None`) holds only the producing rank's keys and should be replayed there — replaying it on another rank is not an error, every key simply belongs to someone else and is skipped, which `ReplayStats.skipped` reports. Like `incremental_dump`, only `roundrobin` and `hash_roundrobin` are supported; `continuous` raises `NotImplementedError`. +*Sharding.* Replay keeps only the keys this rank owns, recomputing ownership from the key with **this** model's global world size — the fan-out the tables were sharded over. There is deliberately no process-group argument: replay is local (filter, then write), so a group could only narrow the modulus and mis-route every key. A delta gathered over a process group (`incremental_dump(..., pg)`) holds the whole group's keys and so fans out correctly when the same delta is handed to every rank. A per-rank delta (`pg=None`) holds only the producing rank's keys and should be replayed there — replaying it on another rank is not an error, every key simply belongs to someone else and is skipped, which `ReplayStats.skipped` reports. Like `incremental_dump`, only `roundrobin` and `hash_roundrobin` are supported; `continuous` raises `NotImplementedError`. *Optimizer state* is not part of a delta. A key that already occupies its target row keeps its optimizer state; a row taken over from another key (or a brand-new one) is reset to the table's initial optimizer state. @@ -861,14 +861,14 @@ The same rule is enforced per key inside the kernel — a slot that does not lan def replay_increment( model: torch.nn.Module, deltas: Dict[str, "DeltaDumpResult"], - pg: Optional[dist.ProcessGroup] = None, + content: ReplayContent = ReplayContent.ALL, ) -> Dict[str, Dict[str, "ReplayStats"]]: """Write incremental_dump results back into a model's dynamic embedding tables. Args: model(nn.Module): The model containing dynamic embedding tables. deltas(Dict[str, DeltaDumpResult]): `incremental_dump`'s return value, keyed by embedding-collection path. Collections or tables the model does not have are skipped with a warning. - pg(Optional[dist.ProcessGroup]): optional. The process group defining this model's shard fan-out. Defaults to the world the tables were created against. + content(ReplayContent): which parts of each dumped row to write back — embedding, optimizer state, score, or any combination. Defaults to all three. Returns ------- @@ -899,7 +899,7 @@ Example — replicate a training model's deltas into a serving model: for i in range(len(r.table_names))} for c, r in deltas.items()} # threshold for the NEXT dump - stats = replay_increment(serve_model, deltas, pg) # raises if layouts differ + stats = replay_increment(serve_model, deltas) # raises if layouts differ for collection, per_table in stats.items(): for name, s in per_table.items(): print(collection, name, s.upserted, "keys replayed at their source slot") diff --git a/corelib/dynamicemb/docs/replay_increment_design.md b/corelib/dynamicemb/docs/replay_increment_design.md index 177c3ed59..2e6607e74 100644 --- a/corelib/dynamicemb/docs/replay_increment_design.md +++ b/corelib/dynamicemb/docs/replay_increment_design.md @@ -271,7 +271,7 @@ def replay_increment( ) -> Dict[str, Dict[str, ReplayStats]]: ... # module level -- BatchedDynamicEmbeddingTables -def replay_increment(self, delta, pg=None) -> Dict[str, ReplayStats]: ... +def replay_increment(self, delta, content=ReplayContent.ALL) -> Dict[str, ReplayStats]: ... # storage level -- DynamicEmbStorage / HybridStorage def replay_increment(self, table_id, keys, values, slot_index, insert_score, diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py index 470306ee5..6d0a09a6d 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py @@ -1710,7 +1710,6 @@ def _replay_compatibility( def replay_increment( self, delta: "DeltaDumpResult", - pg: Optional[dist.ProcessGroup] = None, content: ReplayContent = ReplayContent.ALL, ) -> Dict[str, ReplayStats]: """Write an ``incremental_dump`` delta back into this module's tables. @@ -1747,8 +1746,6 @@ def replay_increment( Args: delta: one collection's :class:`DeltaDumpResult`. Tables not present in this module are skipped with a warning. - pg: process group used to size this model's shard fan-out. Defaults - to the world the tables were created against. Returns: ``{table_name: ReplayStats}`` -- keys written / removed / skipped @@ -1772,12 +1769,15 @@ def replay_increment( f"replay_increment requires DynamicEmbStorage or HybridStorage, " f"got {type(storage).__name__}" ) - rank = dist.get_rank(group=pg) if dist.is_initialized() else 0 - world_size = ( - dist.get_world_size(group=pg) - if (pg is not None and dist.is_initialized()) - else self._shard_world_size - ) + # Both global. Ownership follows how the table was sharded -- row-wise + # over the whole world, which is what ``_shard_world_size`` records and + # what ``meta["world_size"]`` is checked against. There is deliberately + # no process-group argument: replay is local (filter by ownership, then + # write), so a group could only narrow the modulus and mis-route every + # key, leaving each claimed by several ranks and the one that owns it + # claiming nothing. + rank = dist.get_rank() if dist.is_initialized() else 0 + world_size = self._shard_world_size if self._cache is not None and isinstance(storage, DynamicEmbStorage): # Push dirty cache entries down first so the storage copy this replay # is about to overwrite is the authoritative one. diff --git a/corelib/dynamicemb/dynamicemb/incremental_dump.py b/corelib/dynamicemb/dynamicemb/incremental_dump.py index 83e01bd8a..3e3dcbf25 100644 --- a/corelib/dynamicemb/dynamicemb/incremental_dump.py +++ b/corelib/dynamicemb/dynamicemb/incremental_dump.py @@ -443,7 +443,6 @@ def incremental_dump( def replay_increment( model: torch.nn.Module, deltas: Dict[str, DeltaDumpResult], - pg: Optional[dist.ProcessGroup] = None, content: ReplayContent = ReplayContent.ALL, ) -> Dict[str, Dict[str, ReplayStats]]: """Write ``incremental_dump`` results back into a model's dynamic embedding tables. @@ -505,8 +504,6 @@ def replay_increment( deltas (Dict[str, DeltaDumpResult]): ``incremental_dump``'s return value, keyed by embedding-collection path. Collections or tables that the model does not have are skipped with a warning. - pg (Optional[dist.ProcessGroup]): process group defining this model's - shard fan-out. Defaults to the world the tables were created against. content (ReplayContent): which parts of each dumped row to write back -- embedding, optimizer state, score, or any combination. Defaults to all three. @@ -562,7 +559,6 @@ def replay_increment( collection_stats.update( dynamic_emb_module.replay_increment( module_delta, - pg=pg, content=content, ) ) diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py index 31c449a1a..cad944aa8 100644 --- a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py +++ b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_distributed_dynamicemb.py @@ -512,7 +512,7 @@ def test_replay_increment_distributed( pg = intra_and_cross_node_pg()[0] deltas = incremental_dump(train_model, 0, pg) - stats = replay_increment(serve_model, deltas, pg) + stats = replay_increment(serve_model, deltas) delta = deltas[prefix_path] for j, table_name in enumerate(delta.table_names): From 4b5087eb5d7c7d7b2438e000d877a4b2686423c4 Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Mon, 7 Sep 2026 08:31:28 +0530 Subject: [PATCH 3/9] refactor(dynamicemb): read storage layout through tiers, and bound row 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) --- corelib/dynamicemb/DynamicEmb_APIs.md | 9 ++- .../docs/replay_increment_design.md | 15 +++- .../dynamicemb/batched_dynamicemb_tables.py | 72 ++++++++++++++----- .../dynamicemb/dynamicemb/incremental_dump.py | 9 ++- .../dynamicemb/dynamicemb/key_value_table.py | 16 +++++ 5 files changed, 97 insertions(+), 24 deletions(-) diff --git a/corelib/dynamicemb/DynamicEmb_APIs.md b/corelib/dynamicemb/DynamicEmb_APIs.md index dddbd6fff..bb66509f4 100644 --- a/corelib/dynamicemb/DynamicEmb_APIs.md +++ b/corelib/dynamicemb/DynamicEmb_APIs.md @@ -822,7 +822,12 @@ The meaning of the threshold depends on the table's `score_strategy`: - ``"slot_index": torch.Tensor`` -- int64 host tensor aligned with ``keys``; the storage slot each dumped key occupies (for ``replay_increment``). - - ``"current_capacity": int`` -- the table's current capacity. + - ``"current_capacity": int`` -- key-map slots: the modulus that + decides a key's home bucket. + - ``"row_capacity": Tuple[int, ...]`` -- value-buffer rows per + storage tier, which is what bounds a row write. Equal to + `current_capacity` except under NO_EVICTION, whose key map is + deliberately larger than its value buffer. - ``"bucket_capacity": int`` -- slots per hash bucket. - ``"num_scores": int`` -- score words per key; part of the slot layout ``replay_increment`` compares against. @@ -841,7 +846,7 @@ More usage please see [test](https://github.com/NVIDIA/recsys-examples/blob/main **Behavior** For every table in the delta, `replay_increment` erases the table's `erased_keys` (so the target converges to the source) and then upserts `keys` / `values` at the slots in `meta["slot_index"]`. `evicted_keys` is never applied: the key that took an evicted key's slot is in the same delta and overwrites it. -*Write-back is by slot.* Every key is written at the slot and value row it held in the source table, leaving the target layout-identical to it. A key can only be found inside its own home bucket, and that bucket is `hash(key) % capacity / bucket_capacity`, so this requires the target's layout to match the source's. `replay_increment` compares the delta's `meta` (`current_capacity`, `bucket_capacity`, `num_scores`, `world_size`, and the `table_options` fields `score_strategy` / `dim` / `dist_type`) against the target table, and a mismatch raises `ValueError` naming the first mismatching field, **before anything is written**. Configure the target to match the source, or rebuild it from a full checkpoint (`DynamicEmbLoad`) instead. +*Write-back is by slot.* Every key is written at the slot and value row it held in the source table, leaving the target layout-identical to it. A key can only be found inside its own home bucket, and that bucket is `hash(key) % capacity / bucket_capacity`, so this requires the target's layout to match the source's. `replay_increment` compares the delta's `meta` (`current_capacity`, `row_capacity`, `bucket_capacity`, `num_scores`, `world_size`, and the `table_options` fields `score_strategy` / `dim` / `dist_type`) against the target table, and a mismatch raises `ValueError` naming the first mismatching field, **before anything is written**. Configure the target to match the source, or rebuild it from a full checkpoint (`DynamicEmbLoad`) instead. The same rule is enforced per key inside the kernel — a slot that does not land in its key's home bucket raises rather than dropping the key. diff --git a/corelib/dynamicemb/docs/replay_increment_design.md b/corelib/dynamicemb/docs/replay_increment_design.md index 2e6607e74..bf558acdd 100644 --- a/corelib/dynamicemb/docs/replay_increment_design.md +++ b/corelib/dynamicemb/docs/replay_increment_design.md @@ -17,7 +17,8 @@ | `meta[i]` field | why replay needs it | |---|---| | `slot_index` | the exact key slot / value row each dumped key occupied | -| `current_capacity` | a source slot is only transferable when the target's capacity matches | +| `current_capacity` | a source slot is only transferable when the target's key-map capacity matches | +| `row_capacity` | how far a row write may reach, which under NO_EVICTION is *not* the key-map capacity | | `world_size` | the source's row-wise fan-out (key → rank modulo base) | | `table_options` | `bucket_capacity` / `score_strategy` / `dim` / `dist_type` compatibility | | `current_score` | the table score after the dump | @@ -115,7 +116,17 @@ the source's. §4.2 checks that up front and §4.3 re-checks it per key. A table is replayed only when **all** of these match; otherwise `replay_increment` raises `ValueError` naming the first mismatch, before writing anything: -- `meta["current_capacity"]` vs target `key_index_map.capacity(table_id)` +- `meta["current_capacity"]` vs target `key_index_map.capacity(table_id)` — the + modulus for choosing a home bucket, i.e. whether a *slot* means the same thing + in both tables +- `meta["row_capacity"]` vs the target's value-buffer rows per tier — a separate + bound on where a *row* write may land. The two are the same number everywhere + except NO_EVICTION, whose key map is `1 / max_load_factor` times its value + buffer; rounding that up to whole buckets makes the map's capacity + non-injective in the buffer's, so equal `current_capacity` does not imply + equal row counts. With `bucket_capacity=128`, `init_capacity` 100 and 128 both + give a 256-slot key map over 100 and 128 rows -- and a source row of 127 + written past the end of a 100-row target - `meta["table_options"].bucket_capacity` vs the target table's - `meta["table_options"].score_strategy` vs the target's, compared by **physical** word order — `(TIMESTAMP, LFU)` and `(LFU, TIMESTAMP)` are the diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py index 6d0a09a6d..2b805d7d9 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py @@ -1557,7 +1557,8 @@ def incremental_dump( table). ``values[i]`` is embeddings only, with the rest of each stored row in ``optimizer_states[i]`` and every score word in ``scores[i]``; ``meta[i]`` carries current_score / slot_index / current_capacity / - bucket_capacity / num_scores / world_size / table_options; ``evicted_keys[i]`` is the keys this table + row_capacity / bucket_capacity / num_scores / world_size / + table_options; ``evicted_keys[i]`` is the keys this table retained since the last dump -- evictions and explicit erases under ``RETAIN_KEY`` -- and ``erased_keys[i]`` the keys an explicit erase asked to have recorded. Both are drained here. @@ -1627,14 +1628,6 @@ def incremental_dump( # Always drained: recording an erase is the erase call's decision, # so there is no table setting to gate on here. er = self._drain_retained(storage, "pop_erased_keys", True, table_id, pg) - # current_capacity: DynamicEmbStorage has one key_index_map; a - # HybridStorage sums its tiers. - if hasattr(storage, "key_index_map"): - current_capacity = storage.key_index_map.capacity(table_id) - else: - current_capacity = sum( - s.key_index_map.capacity(table_id) for s in storage.tables - ) res.table_names.append(table_name) res.keys.append(keys_cat) res.values.append(values_cat) @@ -1646,7 +1639,8 @@ def incremental_dump( { "current_score": current_score, "slot_index": slot_index, - "current_capacity": current_capacity, + "current_capacity": self._capacity_of(storage, table_id), + "row_capacity": self._row_capacity_of(storage, table_id), "bucket_capacity": self._bucket_capacity_of(storage), "num_scores": self._num_scores_of(storage), "world_size": self._shard_world_size, @@ -1678,12 +1672,17 @@ def _replay_compatibility( if src_options is None: return "delta carries no 'table_options' (dumped by an older version)" - if hasattr(storage, "key_index_map"): - capacity = storage.key_index_map.capacity(table_id) - else: - capacity = sum(s.key_index_map.capacity(table_id) for s in storage.tables) checks = [ - ("capacity", meta.get("current_capacity"), capacity), + ( + "capacity", + meta.get("current_capacity"), + self._capacity_of(storage, table_id), + ), + ( + "row_capacity", + meta.get("row_capacity"), + self._row_capacity_of(storage, table_id), + ), ( "bucket_capacity", meta.get("bucket_capacity"), @@ -1961,16 +1960,51 @@ def _drain_retained( out = _all_gather_evicted_keys(out, pg) return out.cpu() + @staticmethod + def _row_capacity_of( + storage: Union[DynamicEmbStorage, HybridStorage], table_id: int + ) -> Tuple[int, ...]: + """Value-buffer rows per tier for one logical table. + + A second bound, distinct from :meth:`_capacity_of`. That one 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. This one is + how many rows the value buffer actually has, which is what a *row* write + is bounded by. + + The two coincide everywhere except NO_EVICTION, where the key map is + deliberately ``1 / max_load_factor`` times the value buffer -- and + rounding that up to a whole number of 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 while + leaving 100 and 128 rows: equal on :meth:`_capacity_of`, and a source + row of 127 written past the end of a 100-row target. Per tier rather + than summed, since ``slot_index`` routes each key to one tier and a + matching total would say nothing about either. + + Note the two meanings of ``tables`` in play: a storage's tiers, each of + which has its own ``tables`` of value buffers indexed by logical table. + """ + return tuple(s.tables[table_id].shape[0] for s in storage.tables) + + @staticmethod + def _capacity_of( + storage: Union[DynamicEmbStorage, HybridStorage], table_id: int + ) -> int: + """Slots one logical table has across the whole storage. + + Summed over the tiers, because capacity is the one property here that a + second tier adds to rather than duplicates -- the others below read a + single tier. + """ + return sum(s.key_index_map.capacity(table_id) for s in storage.tables) + @staticmethod def _bucket_capacity_of(storage: Union[DynamicEmbStorage, HybridStorage]) -> int: """The storage's hash-bucket capacity (HBM tier for a hybrid storage).""" - if hasattr(storage, "key_index_map"): - return storage.key_index_map.bucket_capacity_ return storage.tables[0].key_index_map.bucket_capacity_ @staticmethod def _num_scores_of(storage: Union[DynamicEmbStorage, HybridStorage]) -> int: """Score words per key (HBM tier for a hybrid storage).""" - if hasattr(storage, "key_index_map"): - return storage.key_index_map.num_scores_ return storage.tables[0].key_index_map.num_scores_ diff --git a/corelib/dynamicemb/dynamicemb/incremental_dump.py b/corelib/dynamicemb/dynamicemb/incremental_dump.py index 3e3dcbf25..d47bb83d4 100644 --- a/corelib/dynamicemb/dynamicemb/incremental_dump.py +++ b/corelib/dynamicemb/dynamicemb/incremental_dump.py @@ -95,7 +95,14 @@ class DeltaDumpResult: with ``keys[i]``; the storage slot each dumped key occupies, used by ``replay_increment``. For NO_EVICTION tables it packs the key slot (high 32 bits) and value row (low 32 bits) into one int64. - meta[i]["current_capacity"]: int -- table's current capacity (slots). + meta[i]["current_capacity"]: int -- key-map slots for this table. + The modulus for choosing a home bucket, so it decides whether a + slot means the same thing in two tables. + meta[i]["row_capacity"]: Tuple[int, ...] -- value-buffer rows per + storage tier. A second bound: ``current_capacity`` says where a + key may sit, this says how far a row write may reach. The two + coincide except under NO_EVICTION, whose key map is deliberately + larger than its value buffer. meta[i]["bucket_capacity"]: int -- slots per hash bucket; replay compares it to decide whether the source slots are usable. meta[i]["num_scores"]: int -- score words per key; part of the diff --git a/corelib/dynamicemb/dynamicemb/key_value_table.py b/corelib/dynamicemb/dynamicemb/key_value_table.py index c9271eb13..656452ad7 100644 --- a/corelib/dynamicemb/dynamicemb/key_value_table.py +++ b/corelib/dynamicemb/dynamicemb/key_value_table.py @@ -2307,6 +2307,16 @@ def __init__( def key_index_map(self): return self._state.key_index_map + @property + def tables(self) -> List[DynamicEmbTableState]: + """The tiers this storage is made of -- one, here. + + Mirrors :attr:`HybridStorage.tables` so that a caller asking a layout + question ("how much capacity does table i have across this storage?") + can iterate uniformly instead of branching on which storage it holds. + """ + return [self._state] + def fill_tables( self, load_factor: float = 0.95, @@ -2926,6 +2936,12 @@ def __init__( @property def tables(self) -> List[DynamicEmbTableState]: + """The tiers this storage is made of: HBM first, then host. + + Order matters to callers that want a single tier's property rather than + a sum -- bucket capacity and score-word count are read off ``tables[0]``, + the HBM tier. + """ return [self._hbm, self._host] # -- Score management -- From 15c21f8ec3be873946398c57b591d429e31b07ba Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Mon, 7 Sep 2026 12:56:19 +0530 Subject: [PATCH 4/9] docs(dynamicemb): correct why _fresh_score_block mints scores, and zero 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) --- .../dynamicemb/dynamicemb/key_value_table.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/corelib/dynamicemb/dynamicemb/key_value_table.py b/corelib/dynamicemb/dynamicemb/key_value_table.py index 656452ad7..9fee35435 100644 --- a/corelib/dynamicemb/dynamicemb/key_value_table.py +++ b/corelib/dynamicemb/dynamicemb/key_value_table.py @@ -862,12 +862,14 @@ def _fresh_score_block( ) -> torch.Tensor: """``[N, num_scores]`` physical score block for keys written by replay. - A delta carries embeddings, not scores, so a replayed key is scored as if it - had just been inserted here: recency words get *timestamp*, every other word - gets *insert_score*. The replica therefore ranks its restored keys by when it - received them rather than by how the source ranked them -- acceptable because - the score only orders future evictions, and a replica that evicts on its own - schedule still holds correct embeddings. + Used when the caller did not ask for ``ReplayContent.SCORE``. The delta does + carry the source's scores, but they were not loaded, so a replayed key is + scored as if it had just been inserted here instead: recency words get + *timestamp*, every other word gets *insert_score*. The replica then ranks its + restored keys by when it received them rather than by how the source ranked + them -- acceptable because the score only orders future evictions, and a + replica that evicts on its own schedule still holds correct embeddings. + Ask for ``SCORE`` to inherit the source's ranking instead. *insert_score* is passed in rather than read off ``state.score``: that field is populated by the forward pass, and replay is the one write path that can @@ -883,7 +885,12 @@ def _fresh_score_block( return value_rows.to(device=device, dtype=SCORE_TYPE).view(-1, 1) n = value_rows.numel() num_scores = state.key_index_map.num_scores_ - block = torch.empty((n, num_scores), device=device, dtype=SCORE_TYPE) + # Zeros, not ``empty``: the loop below fills one column per configured + # strategy, while the width comes from the score policy. The two agree for + # every strategy there is, but they are derived from different places, so a + # future divergence should leave a quiet zero rather than whatever the + # allocator handed back. + block = torch.zeros((n, num_scores), device=device, dtype=SCORE_TYPE) physical = get_physical_score_order(state.options_list[table_id].score_strategy) for word, strategy in enumerate(physical): block[:, word] = ( From 81831fb64acb2eef49db5a192ceec2598ddef3c0 Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Mon, 7 Sep 2026 19:44:57 +0530 Subject: [PATCH 5/9] test(dynamicemb): assert what the score tests claim to assert 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) --- .../incremental_dump/test_replay_increment.py | 103 +++++++++--------- 1 file changed, 51 insertions(+), 52 deletions(-) diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py index ad74f8910..0fb53bc1a 100644 --- a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py +++ b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py @@ -234,9 +234,16 @@ def test_dump_splits_the_value_row(current_device, optimizer): def test_dump_scores_are_column_aligned(current_device): - """``scores`` carries every score word, one row per key, in logical order.""" + """``scores`` is one row per key, one column per configured strategy. + + Configured as ``(LFU, TIMESTAMP)`` on purpose. The physical layout of the + compound policy is always ``(TIMESTAMP, LFU)``, so the *reversed* order is + what makes the columns tell logical apart from physical: with the natural + order the permutation is the identity and the test would pass even with + ``score_dump_permutation`` deleted. + """ device = torch.device(f"cuda:{current_device}") - strategy = (DynamicEmbScoreStrategy.TIMESTAMP, DynamicEmbScoreStrategy.LFU) + strategy = (DynamicEmbScoreStrategy.LFU, DynamicEmbScoreStrategy.TIMESTAMP) model = make_model(current_device, score_strategy=strategy) hot = list(range(1001, 1051)) cold = list(range(2001, 2051)) @@ -248,29 +255,36 @@ def test_dump_scores_are_column_aligned(current_device): keys, scores = sorted_view(delta.keys[0], delta.scores[0]) assert scores.shape == (keys.numel(), len(strategy)) - # Column 1 is the LFU frequency, carried verbatim: the keys touched six - # times must outrank those touched once. Column 0 is a timestamp held as an - # age, so it is not compared against a raw clock here. hot_set = set(hot) is_hot = torch.tensor([int(k) in hot_set for k in keys.tolist()]) - assert int(scores[is_hot, 1].min()) > int(scores[~is_hot, 1].max()) + + # Column 0 is the frequency, carried verbatim -- exact counts, not an + # ordering: one access each for the cold keys, five for the hot ones. A + # comparison would still hold if every count were off by the same amount. + assert set(scores[is_hot, 0].tolist()) == {5} + assert set(scores[~is_hot, 0].tolist()) == {1} + + # Column 1 is the timestamp, dumped as an age (cur_ts - score). The hot keys + # were touched last, so theirs is the smaller -- which pins the direction of + # the conversion, the part a raw clock cannot be compared against. + assert int(scores[is_hot, 1].max()) < int(scores[~is_hot, 1].min()) @pytest.mark.parametrize( - "content, replica_ranks_like_source", + "content, keeps_source_scores", [(ReplayContent.EMBEDDING, False), (ReplayContent.ALL, True)], ) def test_replay_scores_follow_the_content_flag( - current_device, content, replica_ranks_like_source + current_device, content, keeps_source_scores ): """Whether the replica inherits the source's ranking is ``SCORE``'s call. - Same setup either way -- hot keys accessed six times, cold keys once, and a - threshold drawn between them that splits the source's table exactly in half. - Without ``SCORE`` every restored key carries the same fresh score, so the - same threshold splits nothing; with it the frequency column is carried - verbatim and the split is reproduced. The pair is what makes the limitation - legible: neither half means much alone. + Same setup either way -- cold keys accessed once, hot keys five times -- and + the replica's own scores are read back to see which it got. With ``SCORE`` + the frequencies arrive verbatim; without it every restored key is scored as + if freshly inserted here, so they all score alike and the replica ranks by + when it received a key rather than by how the source ranked it. The pair is + what makes that a visible choice: neither half says much alone. """ device = torch.device(f"cuda:{current_device}") src = make_model(current_device, score_strategy=DynamicEmbScoreStrategy.LFU) @@ -282,23 +296,32 @@ def test_replay_scores_follow_the_content_flag( for _ in range(5): # drive the two groups' LFU counts apart touch(src, hot, device) - all_keys = set(hot) | set(cold) - threshold = 3 # between the cold keys' 1 access and the hot keys' 6 - assert set(src.incremental_dump({TABLE_NAME: threshold}).keys[0].tolist()) == set( - hot - ), "the source must rank hot keys above cold ones, or there is nothing to inherit" + delta = dump_all(src) + hot_set = set(hot) + src_keys, src_freq = sorted_view(delta.keys[0], delta.scores[0][:, 0]) + src_is_hot = torch.tensor([k in hot_set for k in src_keys.tolist()]) + assert set(src_freq[src_is_hot].tolist()) == {5} + assert set(src_freq[~src_is_hot].tolist()) == { + 1 + }, "the source must rank hot above cold, or there is nothing to inherit" - dst.replay_increment(dump_all(src), content=content) - assert set(dump_all(dst).keys[0].tolist()) == all_keys + dst.replay_increment(delta, content=content) - split = set(dst.incremental_dump({TABLE_NAME: threshold}).keys[0].tolist()) - if replica_ranks_like_source: - assert split == set(hot), "the replica must rank the way the source did" + # Read the replica's scores directly rather than inferring them from what a + # threshold selects: the frequency is what the flag does or does not carry. + replayed = dump_all(dst) + assert set(replayed.keys[0].tolist()) == set(hot) | set(cold) + keys, freq = sorted_view(replayed.keys[0], replayed.scores[0][:, 0]) + is_hot = torch.tensor([k in hot_set for k in keys.tolist()]) + + if keeps_source_scores: + assert set(freq[is_hot].tolist()) == {5} + assert set(freq[~is_hot].tolist()) == {1} else: - # Every key scored alike, so the threshold takes all of them or none -- - # which of the two depends on where the fresh score falls, and that is - # an implementation detail this deliberately does not pin down. - assert split in (set(), all_keys), f"got {len(split)} of {len(all_keys)}" + # One value for every key, whatever it is -- which value a fresh insert + # assigns is not this test's business, only that the source's ranking + # did not come along. + assert len(set(freq.tolist())) == 1, f"got {sorted(set(freq.tolist()))}" def test_replay_spans_multiple_batches(current_device): @@ -571,30 +594,6 @@ def test_replay_content_restores_optimizer_state(current_device): torch.testing.assert_close(got[key], src_state[key]) -def test_replay_without_embedding_needs_aligned_rows(current_device): - """Omitting ``EMBEDDING`` is only legal for an already-aligned replica. - - A key landing on a row it does not already own has no embedding to keep, and - the row still holds the previous occupant's vector -- serving that under the - new key would be silent corruption, so replay refuses. - """ - device = torch.device(f"cuda:{current_device}") - src = make_model(current_device) - dst = make_model(current_device) - keys = list(range(1001, 1051)) - touch(src, keys, device) - delta = dump_all(src) - - score_only = ReplayContent.SCORE - with pytest.raises(ValueError, match="do not already occupy their target row"): - dst.replay_increment(delta, content=score_only) - - # Align the replica first; then the same call is fine. - dst.replay_increment(delta) - dst.replay_increment(delta, content=score_only) - assert set(dump_all(dst).keys[0].tolist()) == set(keys) - - def test_replay_does_not_leave_displaced_keys_in_the_cache(current_device): """A key whose slot a delta key takes over must not survive in the cache. From e6bb4569c50ce18009a2fdee3d8305d0d8bf55c1 Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Tue, 8 Sep 2026 06:23:11 +0530 Subject: [PATCH 6/9] example(dynamicemb): show delta replication end to end with --incremental_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) --- corelib/dynamicemb/example/example.py | 116 +++++++++++++++++++++- corelib/dynamicemb/example/run_example.sh | 3 +- 2 files changed, 115 insertions(+), 4 deletions(-) diff --git a/corelib/dynamicemb/example/example.py b/corelib/dynamicemb/example/example.py index 82c9b22a2..0d8a668d5 100644 --- a/corelib/dynamicemb/example/example.py +++ b/corelib/dynamicemb/example/example.py @@ -1,5 +1,6 @@ import argparse import builtins +import glob import math import os import shutil @@ -26,7 +27,7 @@ get_sharded_table_capacity, get_table_value_bytes, ) -from dynamicemb.incremental_dump import get_score, incremental_dump +from dynamicemb.incremental_dump import get_score, incremental_dump, replay_increment from dynamicemb.optimizer import EmbOptimType from dynamicemb.planner import ( DynamicEmbeddingEnumerator, @@ -173,6 +174,13 @@ def parse_args(): parser.add_argument("--load", action="store_true") parser.add_argument("--dump", action="store_true") parser.add_argument("--incremental_dump", action="store_true") + parser.add_argument( + "--incremental_load", + action="store_true", + help="replay the delta files a previous --incremental_dump run wrote " + "into a second, eval-only model, and report its loss. Run " + "--incremental_dump first, with the same --save_dir.", + ) parser.add_argument("--caching", action="store_true") parser.add_argument("--prefetch_pipeline", action="store_true") parser.add_argument("--external_storage", action="store_true") @@ -944,8 +952,28 @@ def load(args, runtime: RuntimeContext): dist.barrier(device_ids=[runtime.local_rank]) +def _delta_dir(args): + return os.path.join(args.save_dir, "deltas") + + +def _dense_path(args, runtime: RuntimeContext) -> str: + return os.path.join(args.save_dir, f"dense_rank{runtime.rank}.pt") + + def inc_dump(args, runtime: RuntimeContext): - os.makedirs(args.save_dir, exist_ok=True) + """Train, and periodically dump what changed since the last dump. + + Each dump is written to its own file under ``--save_dir``, and the dense + weights alongside them at the end. A later ``--incremental_load`` run + replays those into a second model; the two are separate commands, so the + files are the whole handoff. Only the embeddings travel in a delta -- the + replica needs the trained MLP from somewhere too, or the loss it reports + would be a fresh model's. + + Each dump's ``meta[i]["current_score"]`` becomes the next dump's threshold, + so successive deltas partition the run rather than overlapping. + """ + os.makedirs(_delta_dir(args), exist_ok=True) train_dataset = MovieLensDataset(args.data_path, split="train") # Use global rank for proper data distribution across all processes train_sampler = DistributedSampler( @@ -961,6 +989,20 @@ def inc_dump(args, runtime: RuntimeContext): criterion = nn.MSELoss() undumped_score = get_score(model) + saved = 0 + + def save_delta(res): + nonlocal saved + # Per rank: incremental_dump without a process group leaves each rank + # holding only its own shard, and the replica is this rank's. + path = os.path.join( + _delta_dir(args), f"delta_rank{runtime.rank}_{saved:04d}.pt" + ) + # Not a weights-only payload: meta carries the source's + # DynamicEmbTableOptions, which replay compares against the target. + torch.save(res, path) + saved += 1 + return os.path.basename(path) for epoch in range(args.epochs): train_sampler.set_epoch(epoch) @@ -997,12 +1039,78 @@ def inc_dump(args, runtime: RuntimeContext): keys.size(0) for dr in res.values() for keys in dr.keys ) print( - f"Epoch {epoch+1}/{args.epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {loss.item():.4f}, dump number: {dump_number}" + f"Epoch {epoch+1}/{args.epochs}, Batch {batch_idx}/{len(train_loader)}, Loss: {loss.item():.4f}, dump number: {dump_number}, saved {save_delta(res)}" ) avg_loss = total_loss / len(train_loader) print(f"Epoch {epoch+1}/{args.epochs}, Average Loss: {avg_loss:.4f}") + # Training carried on past the last periodic dump, so ship those updates too + # or the replica is simply out of date. + last = save_delta(incremental_dump(model, undumped_score)) + # Embedding weights are not in state_dict (see load()), so this is the dense + # arch only -- exactly the part a delta does not carry. + torch.save({"model_state_dict": model.state_dict()}, _dense_path(args, runtime)) + print(f"rank {runtime.rank}: final dump saved {last}, plus dense weights") + + +def inc_load(args, runtime: RuntimeContext): + """Replay ``inc_dump``'s delta files into an eval-only model. + + The consuming half. The replica is built with ``training=False``, so its + value rows hold embeddings alone -- a delta from a training model still + carries optimizer state, and replay writes what the target has room for and + drops the rest. Replaying every delta in order leaves it holding each key at + the latest value the training run gave it. + + Note the loop: a replica catches up by applying every delta since it was + last current, not one big one. + """ + delta_paths = sorted( + glob.glob(os.path.join(_delta_dir(args), f"delta_rank{runtime.rank}_*.pt")) + ) + if not delta_paths: + raise FileNotFoundError( + f"no delta files under {_delta_dir(args)} -- run " + "`example.py --incremental_dump` first, with the same --save_dir" + ) + + test_dataset = MovieLensDataset(args.data_path, split="test") + test_sampler = DistributedSampler( + test_dataset, num_replicas=runtime.world_size, rank=runtime.rank, shuffle=False + ) + test_loader = build_dataloader(test_dataset, test_sampler, args) + + model = create_model(args, runtime, training=False) + model.to(runtime.device) + model.load_state_dict( + torch.load(_dense_path(args, runtime), weights_only=False)["model_state_dict"], + strict=False, + ) + + for path in delta_paths: + # weights_only=False for the reason given where these are written. + stats = replay_increment(model, torch.load(path, weights_only=False)) + for collection, per_table in stats.items(): + for name, st in per_table.items(): + print( + f"{os.path.basename(path)} {collection}/{name}: " + f"upserted={st.upserted} erased={st.erased} skipped={st.skipped}" + ) + + # Trained MLP over replayed embeddings: comparable with what the training + # run reported, which is what makes the number worth printing. + test_one_epoch(model, test_loader, nn.MSELoss(), 0, 1, runtime) + + dist.barrier(device_ids=[runtime.local_rank]) + # Only global rank 0 should clean up, not local rank 0 on each node + if runtime.rank == 0: + try: + shutil.rmtree(args.save_dir) + except Exception as e: + print(f"Warning: Failed to remove {args.save_dir}: {e}") + dist.barrier(device_ids=[runtime.local_rank]) + def main(): args = parse_args() @@ -1022,6 +1130,8 @@ def main(): load(args, runtime) if args.incremental_dump: inc_dump(args, runtime) + if args.incremental_load: + inc_load(args, runtime) finally: cleanup_runtime(runtime) diff --git a/corelib/dynamicemb/example/run_example.sh b/corelib/dynamicemb/example/run_example.sh index 30a054ea8..cbda335b6 100644 --- a/corelib/dynamicemb/example/run_example.sh +++ b/corelib/dynamicemb/example/run_example.sh @@ -3,4 +3,5 @@ torchrun --standalone --nproc_per_node=${NGPU} example.py --train "$@" torchrun --standalone --nproc_per_node=${NGPU} example.py --train --caching --prefetch_pipeline "$@" torchrun --standalone --nproc_per_node=${NGPU} example.py --load --dump "$@" -torchrun --standalone --nproc_per_node=${NGPU} example.py --incremental_dump "$@" \ No newline at end of file +torchrun --standalone --nproc_per_node=${NGPU} example.py --incremental_dump "$@" +torchrun --standalone --nproc_per_node=${NGPU} example.py --incremental_load "$@" From e465893dfbd45b453308fb92d14c369fb95b710a Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Tue, 8 Sep 2026 12:03:18 +0530 Subject: [PATCH 7/9] fix(dynamicemb): make replay reject before it writes, not partway through 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) --- .../dynamicemb/batched_dynamicemb_tables.py | 15 ++++- .../dynamicemb/dynamicemb_config.py | 44 ++++++------- .../dynamicemb/dynamicemb/key_value_table.py | 43 +++---------- .../incremental_dump/test_replay_increment.py | 64 ++++++++++++++++++- 4 files changed, 105 insertions(+), 61 deletions(-) diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py index 2b805d7d9..33eec7d47 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py @@ -1784,6 +1784,12 @@ def replay_increment( ts = device_timestamp() results: Dict[str, ReplayStats] = {} + # Validate every table before writing any of them. A delta usually spans + # a whole collection, so checking table i only when its turn comes would + # let a mismatch on the last one raise after the first few had already + # been erased from and written to -- a partly-applied delta, which is + # worse than a rejected one and is not what this promises. + planned = [] for i, table_name in enumerate(delta.table_names): if table_name not in self._table_names: warnings.warn( @@ -1795,8 +1801,6 @@ def replay_increment( continue table_id = self._table_names.index(table_name) meta = delta.meta[i] - option = self._dynamicemb_options[table_id] - stats = ReplayStats() keys = delta.keys[i] values = delta.values[i] @@ -1835,6 +1839,13 @@ def replay_increment( "two tables share a layout; configure the target to match the " "source, or rebuild it from a full checkpoint instead." ) + planned.append((i, table_name, table_id, keys, values, opt_states, scores)) + + for i, table_name, table_id, keys, values, opt_states, scores in planned: + meta = delta.meta[i] + option = self._dynamicemb_options[table_id] + slot_index = meta["slot_index"] + stats = ReplayStats() # Only ``erased_keys`` is ever replayed. ``evicted_keys`` is not a # removal a replica has to perform: the key that took the evicted diff --git a/corelib/dynamicemb/dynamicemb/dynamicemb_config.py b/corelib/dynamicemb/dynamicemb/dynamicemb_config.py index 362063db5..6bcddeed5 100644 --- a/corelib/dynamicemb/dynamicemb/dynamicemb_config.py +++ b/corelib/dynamicemb/dynamicemb/dynamicemb_config.py @@ -142,39 +142,37 @@ class EvictedItemMode(enum.Flag): class ReplayContent(enum.Flag): - """Which parts of a dumped row ``replay_increment`` writes back. + """What travels with a key that ``replay_increment`` writes back. - A delta carries an embedding, the optimizer state that shares its value row, - and every score word (see :class:`DeltaDumpResult`). Which of them a replica - wants depends on what it is for: a serving replica needs the embedding and - nothing else, while a training replica that has to resume from the source's - exact state wants all three. + The key and its embedding always do -- that is what a delta is for, and + there is no way to write a key at a slot without deciding what its row + holds. The flags choose what comes *along*: the optimizer state that shares + the value row, and the score words. + + Which of them a replica wants depends on what it is for: a serving replica + needs nothing beyond the embedding, while a training replica resuming from + its source's exact state wants both. Combine with ``|`` and test with ``in``:: - replay_increment(model, deltas) # ALL, the default - replay_increment(model, deltas, content=ReplayContent.EMBEDDING) - replay_increment( - model, deltas, - content=ReplayContent.EMBEDDING | ReplayContent.SCORE, - ) + replay_increment(model, deltas) # ALL, default + replay_increment(model, deltas, content=ReplayContent.SCORE) + replay_increment(model, deltas, + content=ReplayContent.EMBEDDING_ONLY) # nothing extra - The key itself is always written at its source slot -- that is what a replay - *is*, and the flags only choose what travels with it. Removals - (``DeltaDumpResult.erased_keys``) are likewise always applied; they are the - source telling the replica a key is gone, not a payload to opt out of. + ``EMBEDDING_ONLY`` is the empty set rather than a peer of the others, so it + cannot combine with them: ``|`` absorbs it. Test for it as emptiness -- + ``not content`` -- and **not** with ``in``, which is subset containment and + reports the empty set as present in everything. - Omitting ``EMBEDDING`` is only meaningful for a replica already aligned with - its source, where every key still occupies the row it held there. A key that - lands on a row it did not already own has no embedding to keep, and serving - the previous occupant's vector under a new key would be silent corruption -- - so ``replay_increment`` raises instead. + Removals (``DeltaDumpResult.erased_keys``) are always applied; they are the + source telling the replica a key is gone, not a payload to opt out of. """ - EMBEDDING = enum.auto() + EMBEDDING_ONLY = 0 OPTIMIZER_STATE = enum.auto() SCORE = enum.auto() - ALL = EMBEDDING | OPTIMIZER_STATE | SCORE + ALL = OPTIMIZER_STATE | SCORE class DynamicEmbScoreStrategy(enum.IntEnum): diff --git a/corelib/dynamicemb/dynamicemb/key_value_table.py b/corelib/dynamicemb/dynamicemb/key_value_table.py index 9fee35435..072517f65 100644 --- a/corelib/dynamicemb/dynamicemb/key_value_table.py +++ b/corelib/dynamicemb/dynamicemb/key_value_table.py @@ -1921,33 +1921,24 @@ def _replay_write_values( A value row is an embedding followed by the optimizer state, and the two are written together or not at all -- ``store_to_flat_single_table`` copies from - the row base, so there is no way to write the tail without the head. That - shapes the three cases: - - - **Both requested.** The delta carries the whole row; write it as it was - dumped. Nothing is inferred. - - **Embedding only** (the default for a serving replica). Rows that already - belonged to this same key keep their optimizer state, because writing just - the embedding columns leaves the tail untouched. Every other row is new to - this key, so its tail still holds the previous occupant's moments and has - to be reset to ``initial_optim_state``. - - **Optimizer state only.** Every row must already belong to this key -- the - caller has checked -- so the current embedding is read back and rewritten - unchanged ahead of the new tail. + the row base, so there is no way to write the tail without the head. The + embedding is therefore always written, and the only choice is what the tail + gets: + + - **With ``OPTIMIZER_STATE``**, the delta carries the whole row; write it as + it was dumped. Nothing is inferred. + - **Without it**, rows that already belonged to this same key keep their + state, because writing just the embedding columns leaves the tail + untouched. Every other row is new to this key, so its tail still holds the + previous occupant's moments and has to be reset to ``initial_optim_state``. """ if rows.numel() == 0: return emb_dim_cfg = state.table_emb_dims_cpu[table_id] optstate_dim = state.optimizer.get_state_dim(emb_dim_cfg) embeddings = embeddings.to(dtype=state.emb_dtype) - want_emb = ReplayContent.EMBEDDING in content want_opt = ReplayContent.OPTIMIZER_STATE in content and optstate_dim > 0 - if not want_emb: - # Keep what is there: the caller guarantees every row already holds this - # key, so its embedding is the one the replica should go on serving. - embeddings = load_from_flat_single_table(state, rows, table_id)[:, :emb_dim_cfg] - if optstate_dim == 0: store_to_flat_single_table(state, rows, table_id, embeddings) return @@ -2063,20 +2054,6 @@ def _replay_at_slots( "silently drop those keys." ) - if ReplayContent.EMBEDDING not in content: - # Without an embedding to write, a row that did not already hold this key - # would keep the previous occupant's vector and serve it under the new - # key. Refuse rather than corrupt. - stale = torch.logical_not(same_key) - num_stale = int(stale.sum().item()) - if num_stale: - example = int(keys[stale][0].item()) - raise ValueError( - f"replay_increment: {num_stale} of {n} keys on table {table_id} " - f"do not already occupy their target row (e.g. key {example}), " - "so there is no embedding to keep. ReplayContent.EMBEDDING may " - "only be omitted for a replica already aligned with its source." - ) _replay_write_values( state, table_id, value_rows, embeddings, optimizer_states, same_key, content ) diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py index 0fb53bc1a..9edbf868b 100644 --- a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py +++ b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py @@ -272,7 +272,7 @@ def test_dump_scores_are_column_aligned(current_device): @pytest.mark.parametrize( "content, keeps_source_scores", - [(ReplayContent.EMBEDDING, False), (ReplayContent.ALL, True)], + [(ReplayContent.EMBEDDING_ONLY, False), (ReplayContent.ALL, True)], ) def test_replay_scores_follow_the_content_flag( current_device, content, keeps_source_scores @@ -435,6 +435,64 @@ def test_replay_accepts_swapped_score_order(current_device): torch.testing.assert_close(src_vals, dst_vals) +def _two_table_model(current_device, capacities): + """A module whose one storage holds two logical tables, sized separately.""" + options = [ + DynamicEmbTableOptions( + index_type=torch.int64, + embedding_dtype=torch.float32, + device_id=current_device, + dim=DIM, + max_capacity=cap, + bucket_capacity=128, + safe_check_mode=DynamicEmbCheckMode.IGNORE, + local_hbm_for_values=1024**3, + score_strategy=DynamicEmbScoreStrategy.TIMESTAMP, + ) + for cap in capacities + ] + return BatchedDynamicEmbeddingTablesV2( + table_options=options, + output_dtype=torch.float32, + table_names=["t_0", "t_1"], + feature_table_map=[0, 1], + pooling_mode=DynamicEmbPoolingMode.SUM, + use_index_dedup=False, + optimizer=EmbOptimType.SGD, + learning_rate=0.1, + ) + + +def test_replay_rejects_a_multi_table_delta_without_writing_any(current_device): + """A mismatch on one table must leave the others untouched. + + A delta normally spans a whole collection, so validating each table only when + its turn comes would let the last one's mismatch raise after the earlier ones + had already been written -- a partly applied delta, harder to recover from + than a rejected one. Here the first table matches and the second does not. + """ + device = torch.device(f"cuda:{current_device}") + src = _two_table_model(current_device, [DEFAULT_CAPACITY, DEFAULT_CAPACITY]) + dst = _two_table_model(current_device, [DEFAULT_CAPACITY, DEFAULT_CAPACITY * 2]) + + keys = list(range(1001, 1101)) + indices = torch.tensor(keys * 2, dtype=torch.int64, device=device) + offsets = torch.arange(0, len(keys) * 2 + 1, dtype=torch.int64, device=device) + src(indices, offsets) + torch.cuda.synchronize() + + delta = src.incremental_dump({"t_0": 0, "t_1": 0}) + assert all(k.numel() == len(keys) for k in delta.keys), "both tables must dump" + + with pytest.raises(ValueError, match="capacity mismatch"): + dst.replay_increment(delta) + + after = dst.incremental_dump({"t_0": 0, "t_1": 0}) + assert all( + k.numel() == 0 for k in after.keys + ), "t_0 must not have been written before t_1 was rejected" + + def test_replay_rejects_missing_table_options(current_device): """A delta without table_options must be rejected, not replayed with the score-order / dim / dist_type checks quietly skipped.""" @@ -548,7 +606,7 @@ def test_replay_optimizer_state(current_device): delta = dump_all(src) # First replay: brand-new rows, so the optimizer state starts at its initial. - dst.replay_increment(delta, content=ReplayContent.EMBEDDING) + dst.replay_increment(delta, content=ReplayContent.EMBEDDING_ONLY) fresh_state = export_optimizer_state(dst) assert fresh_state, "expected an optimizer state region for ROWWISE_ADAGRAD" for key in keys: @@ -561,7 +619,7 @@ def test_replay_optimizer_state(current_device): # Second replay of the same delta: every key is already in its target slot, # so the embedding is overwritten but the optimizer state must survive. - dst.replay_increment(delta, content=ReplayContent.EMBEDDING) + dst.replay_increment(delta, content=ReplayContent.EMBEDDING_ONLY) after_state = export_optimizer_state(dst) for key in keys: torch.testing.assert_close(after_state[key], trained_state[key]) From 6c9647a1683f9be5f16973aac91201c0ff750981 Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Tue, 8 Sep 2026 12:12:35 +0530 Subject: [PATCH 8/9] fix(example): clear a rank's old deltas before a new incremental_dump 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) --- corelib/dynamicemb/example/example.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/corelib/dynamicemb/example/example.py b/corelib/dynamicemb/example/example.py index 0d8a668d5..d182de2e6 100644 --- a/corelib/dynamicemb/example/example.py +++ b/corelib/dynamicemb/example/example.py @@ -974,6 +974,19 @@ def inc_dump(args, runtime: RuntimeContext): so successive deltas partition the run rather than overlapping. """ os.makedirs(_delta_dir(args), exist_ok=True) + # A dump run owns its rank's sequence outright. Numbering restarts at zero + # here, so a shorter run over a directory left by a longer one -- an earlier + # dump that was interrupted, or never loaded -- would overwrite the low + # numbers and leave the high ones behind. inc_load replays everything it + # matches, in order, so those stragglers would land last and write a previous + # run's embeddings over this one's. Nothing downstream can tell: same config + # means the same layout, so the stale slots validate and apply cleanly. + # Per rank rather than the whole directory, since each rank owns its own + # files and this then needs no coordination between them. + for stale in glob.glob( + os.path.join(_delta_dir(args), f"delta_rank{runtime.rank}_*.pt") + ): + os.remove(stale) train_dataset = MovieLensDataset(args.data_path, split="train") # Use global rank for proper data distribution across all processes train_sampler = DistributedSampler( From 8b43ecdb7d233098fc76b121300086f4fda2c9c1 Mon Sep 17 00:00:00 2001 From: Jiashu Yao Date: Tue, 8 Sep 2026 13:18:23 +0530 Subject: [PATCH 9/9] refactor(dynamicemb): make a replay job describe its own work 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 --- corelib/dynamicemb/DynamicEmb_APIs.md | 6 +- .../docs/replay_increment_design.md | 62 +-- .../dynamicemb/batched_dynamicemb_tables.py | 436 ++++++++++++------ .../dynamicemb/dynamicemb/incremental_dump.py | 13 +- .../incremental_dump/test_replay_increment.py | 60 ++- 5 files changed, 387 insertions(+), 190 deletions(-) diff --git a/corelib/dynamicemb/DynamicEmb_APIs.md b/corelib/dynamicemb/DynamicEmb_APIs.md index bb66509f4..2bec7e3d2 100644 --- a/corelib/dynamicemb/DynamicEmb_APIs.md +++ b/corelib/dynamicemb/DynamicEmb_APIs.md @@ -846,13 +846,13 @@ More usage please see [test](https://github.com/NVIDIA/recsys-examples/blob/main **Behavior** For every table in the delta, `replay_increment` erases the table's `erased_keys` (so the target converges to the source) and then upserts `keys` / `values` at the slots in `meta["slot_index"]`. `evicted_keys` is never applied: the key that took an evicted key's slot is in the same delta and overwrites it. -*Write-back is by slot.* Every key is written at the slot and value row it held in the source table, leaving the target layout-identical to it. A key can only be found inside its own home bucket, and that bucket is `hash(key) % capacity / bucket_capacity`, so this requires the target's layout to match the source's. `replay_increment` compares the delta's `meta` (`current_capacity`, `row_capacity`, `bucket_capacity`, `num_scores`, `world_size`, and the `table_options` fields `score_strategy` / `dim` / `dist_type`) against the target table, and a mismatch raises `ValueError` naming the first mismatching field, **before anything is written**. Configure the target to match the source, or rebuild it from a full checkpoint (`DynamicEmbLoad`) instead. +*Write-back is by slot.* Every key is written at the slot and value row it held in the source table, leaving the target layout-identical to it. A key can only be found inside its own home bucket, and that bucket is `hash(key) % capacity / bucket_capacity`, so this requires the target's layout to match the source's. `replay_increment` compares the delta's `meta` (`current_capacity`, `row_capacity`, `bucket_capacity`, `num_scores`, `world_size`, and the `table_options` fields `score_strategy` / `dim` / `dist_type`) against the target table, and checks that every column has the shape the write path needs. Both happen for **all** of the delta's tables before any of them is written, so a mismatch anywhere raises `ValueError` naming the first offending field with the collection untouched — a delta spans a whole collection, and a partly applied one is worse to recover from than a rejected one. Configure the target to match the source, or rebuild it from a full checkpoint (`DynamicEmbLoad`) instead. The same rule is enforced per key inside the kernel — a slot that does not land in its key's home bucket raises rather than dropping the key. > **Precondition:** writing a key at its source slot **overwrites whatever occupies that slot** in the target. That is what makes a replica converge — if the source evicted key `B` to make room for `A`, a same-capacity replica must do the same. It also means a replayed table must be built *only* by loading/replaying from its source: a table that also takes independent writes can lose a key whose slot a delta key claims. -*Scores and optimizer state are dumped but not yet replayed.* `DeltaDumpResult.scores` and `.optimizer_states` carry them, but `replay_increment` ignores both for now: a restored key is scored as if it had just been inserted into the target, and keeps its optimizer state only if it already occupies the target row. So the replica ranks its own future evictions by when it received each key rather than by how the source ranked it. Embeddings are unaffected — the two models can evict in different orders, but never disagree on the value of a key they both hold. NO_EVICTION is the exception and is exact: its score word is a value row, not a score. +*Scores and optimizer state travel under `ReplayContent`,* which asks for both by default. Drop `SCORE` and a restored key is scored as if it had just been inserted into the target, so the replica ranks its own future evictions by when it received each key rather than by how the source ranked it — the two can evict in different orders, but never disagree on the value of a key they both hold. Drop `OPTIMIZER_STATE` and a key keeps its state only if it already occupies the target row. NO_EVICTION is unaffected by `SCORE` either way and is exact: its score word is a value row, not a score. *Sharding.* Replay keeps only the keys this rank owns, recomputing ownership from the key with **this** model's global world size — the fan-out the tables were sharded over. There is deliberately no process-group argument: replay is local (filter, then write), so a group could only narrow the modulus and mis-route every key. A delta gathered over a process group (`incremental_dump(..., pg)`) holds the whole group's keys and so fans out correctly when the same delta is handed to every rank. A per-rank delta (`pg=None`) holds only the producing rank's keys and should be replayed there — replaying it on another rank is not an error, every key simply belongs to someone else and is skipped, which `ReplayStats.skipped` reports. Like `incremental_dump`, only `roundrobin` and `hash_roundrobin` are supported; `continuous` raises `NotImplementedError`. @@ -873,7 +873,7 @@ The same rule is enforced per key inside the kernel — a slot that does not lan Args: model(nn.Module): The model containing dynamic embedding tables. deltas(Dict[str, DeltaDumpResult]): `incremental_dump`'s return value, keyed by embedding-collection path. Collections or tables the model does not have are skipped with a warning. - content(ReplayContent): which parts of each dumped row to write back — embedding, optimizer state, score, or any combination. Defaults to all three. + content(ReplayContent): what travels with each key besides its embedding, which always does — optimizer state, score words, or both. Defaults to both; `ReplayContent.EMBEDDING_ONLY` for neither. Returns ------- diff --git a/corelib/dynamicemb/docs/replay_increment_design.md b/corelib/dynamicemb/docs/replay_increment_design.md index bf558acdd..be3ebd358 100644 --- a/corelib/dynamicemb/docs/replay_increment_design.md +++ b/corelib/dynamicemb/docs/replay_increment_design.md @@ -33,12 +33,12 @@ Everything replay needs is already there; §3 records what deliberately is *not* |---|---| | Write-back | **By slot.** Every key is written at the slot and value row it held in the source table, leaving the target layout-identical to it | | Layout mismatch | **Raise**, before writing anything. Checked at table level (delta `meta` vs target) and enforced again per key in the kernel | -| Per-key scores | **Dumped, not yet replayed** — `DeltaDumpResult.scores`; a restored key is scored as if freshly inserted into the target (§3) | -| Optimizer state | **Dumped, not yet replayed** — `DeltaDumpResult.optimizer_states`, split out of the same value row as `values` (§3) | +| Per-key scores | Dumped in `DeltaDumpResult.scores`, replayed under `ReplayContent.SCORE`; without it a restored key is scored as if freshly inserted into the target (§3) | +| Optimizer state | Dumped in `DeltaDumpResult.optimizer_states`, split out of the same value row as `values`, replayed under `ReplayContent.OPTIMIZER_STATE` (§3) | | `evicted_keys` | **Never replayed.** Overwriting a slot already reproduces an eviction, so this list exists only for other consumers of the dump | | `erased_keys` | A separate buffer and a separate list, applied whenever non-empty, **erase before upsert** — nothing takes over an erased key's slot, so no write reproduces it | -| What replay writes | `ReplayContent` flag: embedding / optimizer state / score, any combination, default all three | -| Optimizer state | Never carried in the delta. A slot whose previous occupant was the *same* key keeps its optimizer state; any other slot is re-initialised | +| What replay writes | The key and its embedding always; `ReplayContent` chooses what comes along (optimizer state / score), default both | +| Rejection | Every table validated -- metadata *and* column shapes -- before any is written, so a delta applies whole or not at all | | `dist_type` | `roundrobin` / `hash_roundrobin` only — same restriction the dump side already enforces; `continuous` raises | --- @@ -71,8 +71,10 @@ other column (LFU frequency, STEP, CUSTOMIZED, NO_EVICTION row) is verbatim. This is the same convention the file `dump`/`load` already uses for single-column LRU tables. -**Replay does not consume either yet.** A restored key is scored as if it had -just been inserted into the target (`_fresh_score_block`), and keeps its +**Whether replay consumes either is `ReplayContent`'s call**, and it consumes +both by default. Without `SCORE` a restored key is scored as if it had +just been inserted into the target (`_fresh_score_block`); without +`OPTIMIZER_STATE` it keeps its optimizer state only when it already occupies the target row. The consequence is worth stating plainly, because it is the one place a replica is *not* a copy of its source: **the two can evict in different orders.** What they never disagree @@ -180,19 +182,20 @@ Mirrors `_encode_slot_index` (`docs/delta_dump_result_design.md` §4): ## 5. Value writes and optimizer state -The delta carries embeddings only (`values[:, :emb_dim]`), never optimizer -state. `store_to_flat_table_contiguous` copies `min(value_dim, input_dim)` -columns from the row base, so writing an `[N, emb_dim]` tensor touches the -embedding region and leaves the optimizer region untouched. Replay therefore -splits every batch in two: +A value row is an embedding followed by the optimizer state, and the write copies +from the row base (`store_to_flat_single_table` copies `min(value_dim, +input_dim)` columns), so there is no way to write the tail without the head. The +embedding is therefore *always* written and the only question is what the tail +gets — which is what `ReplayContent.OPTIMIZER_STATE` decides: -| group | write | +| | write | |---|---| -| slot previously held the **same key** (`same_key` / `founds`) | embedding columns only — the key's own optimizer state survives | -| new key, or the slot's previous occupant was a different key | embedding columns **plus** `initial_optim_state` for the optimizer region | +| **with** `OPTIMIZER_STATE` | the delta carries the whole row: embedding plus the dumped state, widened back from checkpoint width to runtime width (`pad_optimizer_states_from_checkpoint`). Nothing is inferred, and a delta that carries no state for a table that keeps some raises | +| **without** it, slot previously held the **same key** (`same_key` / `founds`) | embedding columns only — writing just the head leaves the key's own state in place | +| **without** it, any other slot | embedding columns **plus** `initial_optim_state` — the tail still holds the previous occupant's moments, which a new key must not inherit | -Without the second group a new key would silently inherit the evicted key's -optimizer moments. +A table whose optimizer keeps no per-row state (`get_state_dim == 0`) writes the +embedding and ignores the flag. --- @@ -278,17 +281,17 @@ For a module with a `DynamicEmbCache` in front of `DynamicEmbStorage`: def replay_increment( model: torch.nn.Module, deltas: Dict[str, DeltaDumpResult], - pg: Optional[dist.ProcessGroup] = None, + content: ReplayContent = ReplayContent.ALL, ) -> Dict[str, Dict[str, ReplayStats]]: ... # module level -- BatchedDynamicEmbeddingTables def replay_increment(self, delta, content=ReplayContent.ALL) -> Dict[str, ReplayStats]: ... # storage level -- DynamicEmbStorage / HybridStorage -def replay_increment(self, table_id, keys, values, slot_index, insert_score, - content=ReplayContent.ALL, timestamp=None - ) -> ReplayStats: ... -def erase_keys(self, table_id, keys) -> int: ... +def replay_increment(self, table_id, keys, values, optimizer_states, scores, + slot_index, insert_score, content=ReplayContent.ALL, + timestamp=None) -> ReplayStats: ... +def erase_keys(self, table_id, keys, mode=EvictedItemMode.DISCARD) -> int: ... ``` `ReplayStats` (in `dynamicemb/types.py`, alongside the other shared dataclasses) @@ -335,22 +338,25 @@ class ReplayStats: | `test_replay_round_trip` | keys, embeddings and **slot_index** all round-trip; parametrised over TIMESTAMP / STEP / LFU / compound `(TIMESTAMP, LFU)` / NO_EVICTION | | `test_dump_splits_the_value_row` | `values` + `optimizer_states` widths add up to the table's value row; a table with no per-row state reports `None`; parametrised over SGD / ROWWISE_ADAGRAD | | `test_dump_scores_are_column_aligned` | `scores` is `[N, num_scores]` in logical order, and the LFU column is verbatim (keys touched six times outrank those touched once) | -| `test_replay_does_not_carry_scores` | the deliberate gap: a threshold that splits the source's table splits nothing on the replica, because every restored key carries the same fresh score | -| `test_replay_is_idempotent` | replaying the same delta twice is a no-op | +| `test_replay_scores_follow_the_content_flag` | parametrised pair: with `SCORE` the source's LFU frequencies arrive verbatim, without it every restored key is scored as if freshly inserted here | +| `test_replay_spans_multiple_batches` | with `threads_in_wave` shrunk so the chunking actually runs, every column is cut along dimension 0 — 1-D `keys` / `slot_index` and 2-D `values` / `optimizer_states` / `scores` stay paired | | `test_replay_rejects_layout_mismatch` | capacity mismatch raises **and leaves the target untouched** (no partial write) | | `test_replay_rejects_score_strategy_mismatch` | a genuinely different score layout raises | | `test_replay_accepts_swapped_score_order` | `(TIMESTAMP, LFU)` vs `(LFU, TIMESTAMP)` is the *same* physical layout, so it must replay rather than be rejected | | `test_replay_rejects_missing_table_options` | a delta from an older version raises instead of writing blind | | `test_replay_applies_erasures_and_ignores_evictions` | `erased_keys` is applied, erase before upsert, an erased-then-readmitted key survives; a non-empty `evicted_keys` changes nothing | | `test_replay_content_restores_optimizer_state` | with `OPTIMIZER_STATE` the source's state lands even on a row the key is taking over, where the default would have initialised it | -| `test_replay_content_restores_scores` | with `SCORE` the replica reproduces the source's ranking -- the mirror of the test above it | -| `test_replay_without_embedding_needs_aligned_rows` | omitting `EMBEDDING` raises when a key does not already own its target row, and works once the replica is aligned | -| `test_erased_and_evicted_use_separate_buffers` | an erase and an eviction land in different buffers, neither leaks into the other, both are read-and-clear | -| `test_erase_retention_is_per_call_not_per_table` | the same table records one erase and not the next; a table that discards evictions still reports its erases | +| `test_replay_rejects_a_multi_table_delta_without_writing_any` | a layout mismatch on the *second* table leaves the first unwritten | +| `test_replay_rejects_a_malformed_column_before_writing_any_table` | the same, for a bad column rather than bad metadata: a wrong score width, a wrong optimizer-state width, a short `slot_index` | +| `test_erased_and_evicted_use_separate_buffers`¹ | an erase and an eviction land in different buffers, neither leaks into the other, both are read-and-clear | +| `test_erase_retention_is_per_call_not_per_table`¹ | the same table records one erase and not the next; a table that discards evictions still reports its erases | | `test_replay_erased_counts_only_keys_actually_present` | `erased` reports removals that really happened, not removals asked for | -| `test_replay_tolerates_an_empty_removal_list` | nothing to remove is normal, not an error | | `test_replay_optimizer_state` | a key already at its target slot keeps its optimizer state, a fresh row is initialised | | `test_replay_with_caching` | a lookup after replay sees the replayed value, not a stale cached one | +| `test_replay_does_not_leave_displaced_keys_in_the_cache` | a key whose slot a delta key takes over is dropped from the cache too, or `flush_cache` would resurrect it into storage on the next dump | + +¹ in `test/unit_tests/retain_evicted_keys/test_insert_collect_evicted.py`, with +the rest of the retention coverage. | Test | Covers | |---|---| diff --git a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py index 33eec7d47..5a5a8d789 100644 --- a/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py +++ b/corelib/dynamicemb/dynamicemb/batched_dynamicemb_tables.py @@ -17,6 +17,7 @@ import warnings from collections import deque from copy import deepcopy +from dataclasses import dataclass from enum import Enum from functools import partial from itertools import accumulate @@ -175,6 +176,20 @@ def owned_key_mask( return torch.from_numpy(owners == np.uint64(rank)) +def owned_keys( + keys: Optional[Tensor], rank: int, world_size: int, dist_type: str +) -> Optional[Tensor]: + """:func:`owned_key_mask` applied, tolerating a list that is absent or empty. + + Removal lists are optional and often empty, so the caller would otherwise + repeat that guard at every use. + """ + if keys is None or keys.numel() == 0: + return keys + mask = owned_key_mask(keys, rank, world_size, dist_type) + return keys if mask is None else keys[mask] + + def get_loading_files( root_path: str, name: str, @@ -494,6 +509,33 @@ def F(x): print(output) +@dataclass +class _ReplayJob: + """One table's share of a delta, resolved and checked, ready to write. + + Produced by the planning pass and consumed by the applying one, so that "no + table is written until every table has been checked" is a property of the + control flow rather than something to be read out of a long loop. + """ + + table_id: int # where this table sits among the module's own + name: str + keys: Tensor + values: Tensor + optimizer_states: Optional[Tensor] + scores: Optional[Tensor] + slot_index: Tensor + # Only ``erased_keys`` is ever replayed as a removal. ``evicted_keys`` is + # not one a replica has to perform -- the key that took the evicted one's + # slot is in this very delta and overwrites it -- but the cache is a second + # index that write does not reach, so it is carried for invalidation. + erased_keys: Optional[Tensor] + evicted_keys: Optional[Tensor] + # The source's score bookkeeping, adopted before writing so a restored key + # lands on the scale the replica will later threshold against. + current_score: Optional[int] + + class BatchedDynamicEmbeddingTablesV2(nn.Module): """ Dynamic Embedding uses a GPU-optimized scored hash table backend. @@ -1719,14 +1761,13 @@ def replay_increment( a ``ValueError`` rather than written some other way -- see :meth:`_replay_compatibility`. - *content* selects which parts of each row are written -- embedding, - optimizer state, score -- and defaults to all three, i.e. the replica - ends up holding what the source held. Dropping ``SCORE`` leaves restored - keys scored as if freshly inserted here, so the replica orders its own - future evictions by when it received each key rather than by how the - source ranked it; dropping ``OPTIMIZER_STATE`` keeps the state of a key - already on its target row and initialises any other. See - :class:`ReplayContent`. + The key and its embedding always travel; *content* selects what comes + along, and defaults to both, i.e. the replica ends up holding what the + source held. Dropping ``SCORE`` leaves restored keys scored as if freshly + inserted here, so the replica orders its own future evictions by when it + received each key rather than by how the source ranked it; dropping + ``OPTIMIZER_STATE`` keeps the state of a key already on its target row + and initialises any other. See :class:`ReplayContent`. ``delta.erased_keys`` is applied whenever it holds anything: those keys were explicitly removed at the source and nothing else will remove them @@ -1768,68 +1809,71 @@ def replay_increment( f"replay_increment requires DynamicEmbStorage or HybridStorage, " f"got {type(storage).__name__}" ) - # Both global. Ownership follows how the table was sharded -- row-wise - # over the whole world, which is what ``_shard_world_size`` records and - # what ``meta["world_size"]`` is checked against. There is deliberately - # no process-group argument: replay is local (filter by ownership, then - # write), so a group could only narrow the modulus and mis-route every - # key, leaving each claimed by several ranks and the one that owns it - # claiming nothing. - rank = dist.get_rank() if dist.is_initialized() else 0 - world_size = self._shard_world_size + # Two passes on purpose. A delta usually spans a whole collection, so + # anything checked only when its table's turn came would raise after the + # earlier tables had been erased from and written to -- a partly applied + # delta, worse to recover from than a rejected one, and not what this + # promises. Planning touches nothing; applying validates nothing. + plan = self._plan_replay(delta, content, storage) + if self._cache is not None and isinstance(storage, DynamicEmbStorage): - # Push dirty cache entries down first so the storage copy this replay - # is about to overwrite is the authoritative one. + # After planning, so a rejected delta leaves the model exactly as it + # was. Pushing dirty cache entries down first makes the storage copy + # this replay is about to overwrite the authoritative one. flush_cache(self._cache, storage) ts = device_timestamp() - results: Dict[str, ReplayStats] = {} - # Validate every table before writing any of them. A delta usually spans - # a whole collection, so checking table i only when its turn comes would - # let a mismatch on the last one raise after the first few had already - # been erased from and written to -- a partly-applied delta, which is - # worse than a rejected one and is not what this promises. - planned = [] + return {job.name: self._apply_replay(job, content, storage, ts) for job in plan} + + def _plan_replay( + self, + delta: "DeltaDumpResult", + content: ReplayContent, + storage: Union[DynamicEmbStorage, HybridStorage], + ) -> List["_ReplayJob"]: + """Resolve and check every table in *delta*, writing nothing. + + Raises on the first table that cannot be replayed, so a caller that sees + an exception knows the collection is untouched. + """ + plan: List[_ReplayJob] = [] for i, table_name in enumerate(delta.table_names): if table_name not in self._table_names: warnings.warn( f"replay_increment: table_name '{table_name}' is not in this " f"module (available: {self._table_names}); skipping.", UserWarning, - stacklevel=2, + stacklevel=3, ) continue table_id = self._table_names.index(table_name) meta = delta.meta[i] - keys = delta.keys[i] - values = delta.values[i] - opt_states = ( - delta.optimizer_states[i] - if ReplayContent.OPTIMIZER_STATE in content - else None - ) - scores = delta.scores[i] if ReplayContent.SCORE in content else None slot_index = meta.get("slot_index") - widths = {"values": values.size(0)} - if opt_states is not None: - widths["optimizer_states"] = opt_states.size(0) - if scores is not None: - widths["scores"] = scores.size(0) - if any(w != keys.numel() for w in widths.values()): - raise ValueError( - f"replay_increment: delta columns for table '{table_name}' are " - f"not row-aligned (keys={keys.numel()}, " - + ", ".join(f"{k}={v}" for k, v in widths.items()) - + ")." - ) - if slot_index is None: raise ValueError( f"replay_increment: delta for table '{table_name}' carries no " "slot_index; it was produced by an incompatible version of " "incremental_dump." ) + job = _ReplayJob( + table_id=table_id, + name=table_name, + keys=delta.keys[i], + values=delta.values[i], + optimizer_states=( + delta.optimizer_states[i] + if ReplayContent.OPTIMIZER_STATE in content + else None + ), + scores=delta.scores[i] if ReplayContent.SCORE in content else None, + slot_index=slot_index, + erased_keys=delta.erased_keys[i], + evicted_keys=delta.evicted_keys[i], + current_score=meta.get("current_score"), + ) + self._check_delta_shapes(job, content, storage) + mismatch = self._replay_compatibility(table_id, meta) if mismatch is not None: raise ValueError( @@ -1839,112 +1883,202 @@ def replay_increment( "two tables share a layout; configure the target to match the " "source, or rebuild it from a full checkpoint instead." ) - planned.append((i, table_name, table_id, keys, values, opt_states, scores)) + plan.append(job) + return plan - for i, table_name, table_id, keys, values, opt_states, scores in planned: - meta = delta.meta[i] - option = self._dynamicemb_options[table_id] - slot_index = meta["slot_index"] - stats = ReplayStats() - - # Only ``erased_keys`` is ever replayed. ``evicted_keys`` is not a - # removal a replica has to perform: the key that took the evicted - # one's slot is in this very delta and overwrites it. That list - # exists for other consumers of the dump, and replay ignores it. - # An empty or absent list is normal -- nothing was erased, or rank - # filtering left this rank none of them. - erased = delta.erased_keys[i] - # Not replayed as a removal -- the delta's own writes reproduce an - # eviction in the storage -- but still needed to invalidate the - # cache, which those writes do not reach. See below. - evicted = delta.evicted_keys[i] - mask = owned_key_mask(keys, rank, world_size, option.dist_type) - if mask is not None: - stats.skipped = int(keys.numel() - mask.sum().item()) - keys, values = keys[mask], values[mask] - if opt_states is not None: - opt_states = opt_states[mask] - if scores is not None: - scores = scores[mask] - slot_index = slot_index[mask] - if erased is not None and erased.numel() > 0: - er_mask = owned_key_mask(erased, rank, world_size, option.dist_type) - erased = erased[er_mask] - if evicted is not None and evicted.numel() > 0: - ev_mask = owned_key_mask( - evicted, rank, world_size, option.dist_type - ) - evicted = evicted[ev_mask] - - # Erase first: a key erased and then re-inserted inside the same - # window appears in BOTH lists, and must survive the replay. - if erased is not None and erased.numel() > 0: - # Not recorded into this model's own erased buffer: these - # removals came from upstream, and re-reporting them would make - # a chained replica replay what it already received. A model - # that is itself a dump source for someone further down would - # want the opposite -- say so when that case turns up. - stats.erased = storage.erase_keys(table_id, erased) - - # Adopt the source's score bookkeeping BEFORE writing, so a - # restored key lands on the same scale the replica will later - # threshold its own incremental_dump against. Timestamp-based tables - # read their score off the device clock, so there is nothing to - # carry -- their restored keys are stamped with ``ts``. - current_score = meta.get("current_score") - if ( - current_score is not None - and not score_strategy_has_timestamp_column(option.score_strategy) - and table_name in self._scores - ): - self._scores[table_name] = current_score - stats.merge( - storage.replay_increment( - table_id, - keys, - values, - opt_states, - scores, - slot_index, - self._scores.get(table_name, 0), - content=content, - timestamp=ts, - ) + def _check_delta_shapes( + self, + job: "_ReplayJob", + content: ReplayContent, + storage: Union[DynamicEmbStorage, HybridStorage], + ) -> None: + """Every shape the write path depends on, checked before it runs. + + The storage layer checks the widths again -- it is callable on its own, + without this module's planning pass -- but by then a table's predecessors + have been written, so a malformed column on the last table of a + collection would leave the rest of it advanced. + """ + option = self._dynamicemb_options[job.table_id] + n = job.keys.numel() + + rows = {"values": job.values.size(0), "slot_index": job.slot_index.numel()} + if job.optimizer_states is not None: + rows["optimizer_states"] = job.optimizer_states.size(0) + if job.scores is not None: + rows["scores"] = job.scores.size(0) + if any(r != n for r in rows.values()): + raise ValueError( + f"replay_increment: delta columns for table '{job.name}' are not " + f"row-aligned (keys={n}, " + + ", ".join(f"{k}={v}" for k, v in rows.items()) + + ")." + ) + + if job.values.dim() != 2 or job.values.size(1) != option.dim: + raise ValueError( + f"replay_increment: table '{job.name}' has dim {option.dim}, but " + f"the delta's embeddings have shape {tuple(job.values.shape)}." ) - if self._cache is not None: - # The cache is a second index that writing a slot does not reach, - # so everything the storage just stopped holding has to be dropped - # from it explicitly: - # upserted keys -- the next lookup must see the value just - # written into the storage, not the cached one; - # erased keys -- a cached copy would resurrect them; - # evicted keys -- these lost their slot to a delta key. Replay - # ignores them as removals, precisely because the write - # reproduces the eviction *in the storage* -- but a cached - # copy survives that, and ``flush_cache`` would write it back - # down on the next dump, undoing the eviction. - # - # A source table that does not retain evictions (``DISCARD``) - # cannot report that last group, so a caching replica that has to - # converge exactly wants ``evicted_item_mode=RETAIN_KEY``. - parts = [keys] - for extra in (erased, evicted): - if extra is not None and extra.numel() > 0: - parts.append(extra.to(keys.dtype)) - stale = torch.cat(parts) if len(parts) > 1 else keys - if stale.numel() > 0: - device = torch.device(f"cuda:{torch.cuda.current_device()}") - self._cache.key_index_map.erase( - stale.to(device=device), - torch.full( - (stale.numel(),), - table_id, - dtype=torch.int64, - device=device, - ), + + if job.scores is not None: + num_scores = self._num_scores_of(storage) + # A single-word block may arrive 1-D; the storage layer reshapes it, + # so accept the same two spellings here. + width = job.scores.size(1) if job.scores.dim() == 2 else 1 + if job.scores.dim() > 2 or width != num_scores: + raise ValueError( + f"replay_increment: table '{job.name}' has {num_scores} score " + f"word(s) per key, but the delta carries a score block of " + f"shape {tuple(job.scores.shape)}." + ) + + if ReplayContent.OPTIMIZER_STATE in content: + optimizer = storage.tables[0].optimizer + # Only meaningful where the table keeps per-row state at all; one + # without it ignores the column rather than rejecting it. + if optimizer.get_state_dim(option.dim) > 0: + ckpt_dim = optimizer.get_ckpt_state_dim(option.dim) + if job.optimizer_states is None: + # A delta dumped from a table whose optimizer keeps no state + # -- SGD into rowwise Adagrad, say. Nothing in the layout + # check catches it, and the write path raises, so without + # this the collection's earlier tables would already be in. + raise ValueError( + f"replay_increment: table '{job.name}' keeps optimizer " + "state per row, but the delta carries none. Drop " + "ReplayContent.OPTIMIZER_STATE, or replay a delta " + "dumped from a table with the same optimizer." ) - results[table_name] = stats - return results + if ( + job.optimizer_states.dim() != 2 + or job.optimizer_states.size(1) != ckpt_dim + ): + raise ValueError( + f"replay_increment: table '{job.name}' dumps {ckpt_dim} " + f"optimizer-state column(s) per row, but the delta " + f"carries a block of shape " + f"{tuple(job.optimizer_states.shape)}." + ) + + def _apply_replay( + self, + job: "_ReplayJob", + content: ReplayContent, + storage: Union[DynamicEmbStorage, HybridStorage], + ts: int, + ) -> ReplayStats: + """Write one planned table. Everything here has already been checked.""" + option = self._dynamicemb_options[job.table_id] + stats = ReplayStats() + keys, values = job.keys, job.values + opt_states, scores, slot_index = ( + job.optimizer_states, + job.scores, + job.slot_index, + ) + + # An empty or absent list is normal -- nothing was erased, or rank + # filtering left this rank none of them. + erased, evicted = job.erased_keys, job.evicted_keys + + # Both global. Ownership follows how the table was sharded -- row-wise + # over the whole world, which is what ``_shard_world_size`` records and + # what ``meta["world_size"]`` is checked against. There is deliberately + # no process-group argument: replay is local (filter by ownership, then + # write), so a group could only narrow the modulus and mis-route every + # key, leaving each claimed by several ranks and the one that owns it + # claiming nothing. + rank = dist.get_rank() if dist.is_initialized() else 0 + mask = owned_key_mask(keys, rank, self._shard_world_size, option.dist_type) + if mask is not None: + stats.skipped = int(keys.numel() - mask.sum().item()) + keys, values = keys[mask], values[mask] + if opt_states is not None: + opt_states = opt_states[mask] + if scores is not None: + scores = scores[mask] + slot_index = slot_index[mask] + ws, dt = self._shard_world_size, option.dist_type + erased = owned_keys(erased, rank, ws, dt) + evicted = owned_keys(evicted, rank, ws, dt) + + # Erase first: a key erased and then re-inserted inside the same window + # appears in BOTH lists, and must survive the replay. + if erased is not None and erased.numel() > 0: + # Not recorded into this model's own erased buffer: these removals + # came from upstream, and re-reporting them would make a chained + # replica replay what it already received. A model that is itself a + # dump source for someone further down would want the opposite -- + # say so when that case turns up. + stats.erased = storage.erase_keys(job.table_id, erased) + + # Before writing, so a restored key lands on the same scale the replica + # will later threshold its own incremental_dump against. Timestamp-based + # tables read their score off the device clock, so there is nothing to + # carry -- their restored keys are stamped with ``ts``. + if ( + job.current_score is not None + and not score_strategy_has_timestamp_column(option.score_strategy) + and job.name in self._scores + ): + self._scores[job.name] = job.current_score + + stats.merge( + storage.replay_increment( + job.table_id, + keys, + values, + opt_states, + scores, + slot_index, + self._scores.get(job.name, 0), + content=content, + timestamp=ts, + ) + ) + self._invalidate_cache(job.table_id, keys, erased, evicted) + return stats + + def _invalidate_cache( + self, + table_id: int, + keys: Tensor, + erased: Optional[Tensor], + evicted: Optional[Tensor], + ) -> None: + """Drop from the cache everything the storage just stopped holding. + + The cache is a second index that writing a slot does not reach, so three + groups have to go explicitly: + + - upserted keys, so the next lookup sees the value just written into the + storage rather than the cached one; + - erased keys, whose cached copy would resurrect them; + - evicted keys, which lost their slot to a delta key. Replay ignores + those as removals precisely because the write reproduces the eviction + *in the storage* -- but a cached copy survives that, and + ``flush_cache`` would write it back down on the next dump, undoing the + eviction. + + A source table that does not retain evictions (``DISCARD``) cannot report + that last group, so a caching replica that has to converge exactly wants + ``evicted_item_mode=RETAIN_KEY``. + """ + if self._cache is None: + return + parts = [keys] + for extra in (erased, evicted): + if extra is not None and extra.numel() > 0: + parts.append(extra.to(keys.dtype)) + stale = torch.cat(parts) if len(parts) > 1 else keys + if stale.numel() == 0: + return + device = torch.device(f"cuda:{torch.cuda.current_device()}") + self._cache.key_index_map.erase( + stale.to(device=device), + torch.full((stale.numel(),), table_id, dtype=torch.int64, device=device), + ) @staticmethod def _drain_retained( diff --git a/corelib/dynamicemb/dynamicemb/incremental_dump.py b/corelib/dynamicemb/dynamicemb/incremental_dump.py index d47bb83d4..6de177e15 100644 --- a/corelib/dynamicemb/dynamicemb/incremental_dump.py +++ b/corelib/dynamicemb/dynamicemb/incremental_dump.py @@ -472,9 +472,10 @@ def replay_increment( loading/replaying from its source: a table that also takes independent writes can lose a key whose slot a delta key claims. - **What gets written** is *content*'s call -- embedding, optimizer state, - score, or any combination (see :class:`ReplayContent`). The default is all - three, so the replica ends up holding what the source held. + **The key and its embedding always travel**; *content* chooses what comes + along -- the optimizer state that shares the value row, the score words, or + both (see :class:`ReplayContent`). The default is both, so the replica ends + up holding what the source held. Dropping ``SCORE`` leaves a restored key scored as if it had just been inserted here, so the replica orders its own future evictions by when it @@ -511,9 +512,9 @@ def replay_increment( deltas (Dict[str, DeltaDumpResult]): ``incremental_dump``'s return value, keyed by embedding-collection path. Collections or tables that the model does not have are skipped with a warning. - content (ReplayContent): which parts of each dumped row to write back -- - embedding, optimizer state, score, or any combination. Defaults to - all three. + content (ReplayContent): what travels with each key besides its + embedding -- optimizer state, score words, or both. Defaults to + both; ``ReplayContent.EMBEDDING_ONLY`` for neither. Returns: Dict[str, Dict[str, ReplayStats]]: diff --git a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py index 9edbf868b..6795e5168 100644 --- a/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py +++ b/corelib/dynamicemb/test/unit_tests/incremental_dump/test_replay_increment.py @@ -436,7 +436,12 @@ def test_replay_accepts_swapped_score_order(current_device): def _two_table_model(current_device, capacities): - """A module whose one storage holds two logical tables, sized separately.""" + """A module whose one storage holds two logical tables, sized separately. + + Rowwise Adagrad rather than SGD so the rows carry optimizer state: a table + without any lets a malformed optimizer column through untouched, and these + tests are about columns being rejected. + """ options = [ DynamicEmbTableOptions( index_type=torch.int64, @@ -458,8 +463,9 @@ def _two_table_model(current_device, capacities): feature_table_map=[0, 1], pooling_mode=DynamicEmbPoolingMode.SUM, use_index_dedup=False, - optimizer=EmbOptimType.SGD, + optimizer=EmbOptimType.EXACT_ROWWISE_ADAGRAD, learning_rate=0.1, + initial_accumulator_value=0.0, ) @@ -493,6 +499,56 @@ def test_replay_rejects_a_multi_table_delta_without_writing_any(current_device): ), "t_0 must not have been written before t_1 was rejected" +@pytest.mark.parametrize( + "broken", ["scores", "optimizer_states", "no_optimizer_states", "slot_index"] +) +def test_replay_rejects_a_malformed_column_before_writing_any_table( + current_device, broken +): + """A bad column on a later table must not let an earlier one through. + + Row counts are not the only shape the write path depends on: the score block + and the optimizer-state block each have a width, and slot_index a length. + Those are checked again down in the storage layer, which is callable on its + own -- but by then this table's predecessors have been written, so the + collection ends up half advanced. The check has to happen up front. + """ + device = torch.device(f"cuda:{current_device}") + src = _two_table_model(current_device, [DEFAULT_CAPACITY, DEFAULT_CAPACITY]) + dst = _two_table_model(current_device, [DEFAULT_CAPACITY, DEFAULT_CAPACITY]) + + keys = list(range(1001, 1101)) + indices = torch.tensor(keys * 2, dtype=torch.int64, device=device) + offsets = torch.arange(0, len(keys) * 2 + 1, dtype=torch.int64, device=device) + src(indices, offsets) + torch.cuda.synchronize() + + delta = src.incremental_dump({"t_0": 0, "t_1": 0}) + # Damage the second table only: the first is well formed and would be + # written first if the check came too late. + if broken == "scores": + delta.scores[1] = delta.scores[1].repeat(1, 3) + elif broken == "optimizer_states": + delta.optimizer_states[1] = torch.zeros( + delta.keys[1].numel(), 7, dtype=delta.values[1].dtype + ) + elif broken == "no_optimizer_states": + # What a delta from a stateless optimizer looks like -- SGD dumped, + # rowwise Adagrad replayed. The layout check does not compare + # optimizers, so only this catches it before the write path does. + delta.optimizer_states[1] = None + else: + delta.meta[1]["slot_index"] = delta.meta[1]["slot_index"][:-1] + + with pytest.raises(ValueError): + dst.replay_increment(delta) + + after = dst.incremental_dump({"t_0": 0, "t_1": 0}) + assert all( + k.numel() == 0 for k in after.keys + ), "t_0 must not have been written before t_1's column was rejected" + + def test_replay_rejects_missing_table_options(current_device): """A delta without table_options must be rejected, not replayed with the score-order / dim / dist_type checks quietly skipped."""