diff --git a/tzrec/main.py b/tzrec/main.py index 6da197bfa..9cd95f323 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -79,7 +79,12 @@ from tzrec.protos.model_pb2 import Kernel as KernelProto from tzrec.protos.model_pb2 import ModelConfig from tzrec.protos.train_pb2 import TrainConfig -from tzrec.utils import checkpoint_util, config_util, predict_util +from tzrec.utils import ( + checkpoint_util, + config_util, + dynamicemb_util, + predict_util, +) from tzrec.utils.delta_embedding_dump import DeltaEmbeddingDumper from tzrec.utils.dist_util import ( DistributedModelParallel, @@ -689,6 +694,12 @@ def train_and_evaluate( is_local_rank_zero = int(os.environ.get("LOCAL_RANK", 0)) == 0 acc_utils.allow_tf32(train_config) enable_delta_embedding_dump = train_config.HasField("delta_embedding_dump_config") + # Arm evicted-key retention before feature/model building so the dump's + # pop_evicted_keys drain has a buffer to consume on every dynamicemb table. + dynamicemb_util.set_auto_retain_evicted_keys( + enable_delta_embedding_dump + and train_config.delta_embedding_dump_config.dump_evicted_tombstones + ) data_config = pipeline_config.data_config # Build feature diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 02b960fdf..5cbdfe6b0 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -108,6 +108,16 @@ message DeltaEmbeddingDumpConfig { // even embedding_dim. optional DeltaEmbeddingQuantType quant_type = 6 [default = DELTA_EMBEDDING_QUANT_NONE]; + // Publish all-zero tombstone rows for dynamicemb keys evicted since the + // last dump; the polling processor detects them and deletes the key from + // NvEmbeddings, reclaiming its memory (FeatureStore has no delete). While + // enabled, delta dump mode also arms dynamicemb's evicted-key recording + // (evicted_item_mode=RETAIN_KEY) automatically; the retained keys live in + // per-rank GPU memory and grow with eviction volume until each dump + // drains them, so keep the dump interval short on high-churn tables. + // Requires a dynamicemb build with EvictedItemMode; older builds degrade + // to a warning. + optional bool dump_evicted_tombstones = 7 [default = true]; } message TrainConfig { diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 15135c834..3b5cfef1d 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -69,6 +69,12 @@ ) _CONSUMER = "delta_embedding_dump" +# Rows per tombstone chunk: chunks alias one zero buffer of this size so an +# eviction drain never materializes a single huge zero matrix. +_TOMBSTONE_CHUNK_ROWS = 65536 +# Rows per merged-lookup batch: caps the find's [batch, dim] GPU values +# buffer, which would otherwise grow with the dump interval's eviction volume. +_FIND_BATCH_ROWS = 65536 _ShardedEmbeddingModule = Union[ ShardedEmbeddingCollection, ShardedEmbeddingBagCollection ] @@ -162,6 +168,22 @@ def _feature_name(feature_names: Iterable[str]) -> str: return ",".join(names) +def _dynamicemb_pop_evicted_keys_supported() -> bool: + """Return whether dynamicemb tables can report their evicted keys. + + Only used to skip tombstone test coverage on dynamicemb builds without + eviction retention; ``dump`` itself degrades to a warning there instead + of failing fast, because the tombstone switch defaults to on. + """ + try: + from dynamicemb.batched_dynamicemb_tables import ( + BatchedDynamicEmbeddingTablesV2, + ) + except ImportError: + return False + return hasattr(BatchedDynamicEmbeddingTablesV2, "pop_evicted_keys") + + def _metadata_shard_info(metadata: Optional[ShardMetadata]) -> _TableShardInfo: if metadata is None: return _TableShardInfo() @@ -661,6 +683,7 @@ def __init__( self._model = model self._config = config self._quant_type = config.quant_type + self._dump_evicted_tombstones = bool(config.dump_evicted_tombstones) self._schema = ( _DELTA_DUMP_QUANT_SCHEMA if self._quant_type == DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8 @@ -692,6 +715,7 @@ def __init__( auto_compact=True, ) self._zch_modules = self._tracker.zch_modules + self._warned_no_retain_tables: Set[str] = set() self._table_shard_infos = self._collect_table_shard_infos() self._validate_supported_table_sharding(self._table_shard_infos) if self._quant_type == DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8: @@ -761,12 +785,23 @@ def __init__( def clear(self) -> None: """Clear tracked sparse ids, usually after restore-time dummy steps. - The restored model state is the baseline the serving side already - holds, so admit/evict events recorded during the dummy forwards are - dropped together with the tracked ids; the next dump then reports - only the delta against the restored state. + The restored state is the baseline the serving side already holds, so + events recorded during the dummy forwards must not show up in the next + dump; retained evicted keys are popped and discarded for the same + reason. """ self._tracker.clear(_CONSUMER) + if not self._dump_evicted_tombstones: + return + popped_module_ids: Set[int] = set() + for dynamic_module in self._collect_dynamic_modules().values(): + # One module hosts several tables; a single pop drains them all. + if id(dynamic_module) in popped_module_ids: + continue + popped_module_ids.add(id(dynamic_module)) + pop_fn = getattr(dynamic_module, "pop_evicted_keys", None) + if pop_fn is not None: + pop_fn() @contextmanager def pause_tracking(self) -> Iterator[None]: @@ -916,6 +951,7 @@ def dump(self, global_step: int) -> Optional[str]: global_step=global_step, table_weights=table_weights, dynamic_modules=dynamic_modules, + dump_evicted_tombstones=self._dump_evicted_tombstones, ) output_path: Optional[str] = None if write_local and (num_rows > 0 or self._world_size > 1): @@ -965,31 +1001,71 @@ def _append_model_delta_rows( global_step: int, table_weights: Dict[str, _TableWeight], dynamic_modules: Dict[str, nn.Module], + flushed_module_ids: Optional[Set[int]] = None, + dump_evicted_tombstones: bool = False, ) -> int: + """Append real rows for touched ids and tombstones for evicted keys. + + Dynamic tables merge the tracker's ids with this dump's evicted + keys into one post-flush lookup: found keys publish their current + embeddings, missing keys from the evicted set publish zero + tombstones, and missing keys that were never admitted are skipped + with a warning. One lookup resolves both orderings within a dump + interval -- an evicted-then-reinserted key publishes exactly its + fresh row, an insert-then-evict key exactly one tombstone. Row-wise + sharding keeps rank-local shards disjoint, so no cross-rank + collective is needed. + + Args: + table_chunks: List to append the per-table parquet chunks to. + global_step: Current training step. + table_weights: Local weight shards keyed by table FQN. + dynamic_modules: Dynamic embedding modules keyed by table FQN. + flushed_module_ids: Modules already flushed this dump; created + locally when None so direct callers keep working. + dump_evicted_tombstones: When True, dynamic tables drain their + evicted-key buffers and publish tombstones; tables without + tracker rows are still visited so the buffers are drained. + + Returns: + Number of rows appended. + """ num_rows = 0 - # A dynamic module hosting multiple tables is shared across their FQN - # keys; flush() flushes the whole module, so track which - # modules were already flushed this dump and skip the redundant repeats. - flushed_module_ids: Set[int] = set() - for fqn, unique_rows in self._tracker.get_unique(_CONSUMER).items(): - ids = unique_rows.ids - if ids.numel() == 0: + if flushed_module_ids is None: + flushed_module_ids = set() + unique_rows_by_fqn = self._tracker.get_unique(_CONSUMER) + fqns = list(unique_rows_by_fqn) + if dump_evicted_tombstones: + fqns += [fqn for fqn in dynamic_modules if fqn not in unique_rows_by_fqn] + for fqn in fqns: + unique_rows = unique_rows_by_fqn.get(fqn) + dynamic_module = dynamic_modules.get(fqn) + if dynamic_module is not None: + tracker_ids = ( + unique_rows.ids + if unique_rows is not None + else torch.empty(0, dtype=torch.int64) + ) + num_rows += self._append_dynamic_rows( + table_chunks, + global_step=global_step, + fqn=fqn, + dynamic_module=dynamic_module, + tracker_ids=tracker_ids, + flushed_module_ids=flushed_module_ids, + dump_evicted_tombstones=dump_evicted_tombstones, + ) continue - ids = ids.unique(sorted=True) - embeddings, key_ids = self._lookup_embeddings( - fqn, - ids, - table_weights=table_weights, - dynamic_modules=dynamic_modules, - flushed_module_ids=flushed_module_ids, - ) - feature_name = _feature_name( - self._tracker.fqn_to_feature_names.get(fqn, []) - ) + if unique_rows is None or unique_rows.ids.numel() == 0: + continue + ids = unique_rows.ids.unique(sorted=True) + embeddings, key_ids = self._lookup_embeddings(fqn, ids, table_weights) num_rows += self._append_table_chunk( table_chunks, global_step=global_step, - feature_name=feature_name, + feature_name=_feature_name( + self._tracker.fqn_to_feature_names.get(fqn, []) + ), table_fqn=fqn, key_ids=key_ids, embeddings=embeddings, @@ -1002,14 +1078,7 @@ def _lookup_embeddings( fqn: str, ids: torch.Tensor, table_weights: Dict[str, _TableWeight], - dynamic_modules: Dict[str, nn.Module], - flushed_module_ids: Optional[Set[int]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - dynamic_module = dynamic_modules.get(fqn) - if dynamic_module is not None: - return self._lookup_dynamic_embeddings( - dynamic_module, fqn, ids, flushed_module_ids - ) if fqn not in table_weights: raise KeyError(f"Embedding table {fqn} not found in sharded model.") table_weight = table_weights[fqn] @@ -1093,42 +1162,225 @@ def _lookup_zch_embeddings( embeddings[held_mask] = weight[rows].detach() return embeddings, raw_ids - def _lookup_dynamic_embeddings( + def _pop_evicted_key_ids( + self, fqn: str, dynamic_module: nn.Module, table_name: str + ) -> Optional[torch.Tensor]: + """Pop this dump's evicted keys for one dynamic table. + + Args: + fqn: Dynamic table FQN, for warnings. + dynamic_module: Dynamic embedding module hosting the table. + table_name: Table name within the module. + + Returns: + The evicted key ids (read-and-clear semantics), or None when + the build or the table cannot retain evicted keys. + """ + pop_fn = getattr(dynamic_module, "pop_evicted_keys", None) + if pop_fn is None: + self._warn_no_retain_table_once( + fqn, "this dynamicemb build has no pop_evicted_keys" + ) + return None + evicted = pop_fn([table_name]).get(table_name) + if evicted is None: + self._warn_no_retain_table_once( + fqn, "the table does not retain evicted keys (DISCARD mode)" + ) + return None + return evicted.detach() + + def _append_dynamic_rows( self, - dynamic_module: nn.Module, + table_chunks: List[pa.Table], + global_step: int, fqn: str, - ids: torch.Tensor, - flushed_module_ids: Optional[Set[int]] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + dynamic_module: nn.Module, + tracker_ids: torch.Tensor, + flushed_module_ids: Set[int], + dump_evicted_tombstones: bool, + ) -> int: + """Append one merged lookup's real rows and tombstones for a table. + + The tracker ids and this dump's evicted keys are merged and + deduplicated on host, then resolved by a post-flush ``find`` batched + in ``_FIND_BATCH_ROWS`` slices so the GPU values buffer stays capped + regardless of eviction volume. Each key publishes exactly one row: + its current embedding when present, a zero tombstone when absent + but evicted, and nothing (with a warning) when it was never + admitted. + + Args: + table_chunks: List to append the per-table parquet chunks to. + global_step: Current training step. + fqn: Dynamic table FQN. + dynamic_module: Dynamic embedding module hosting the table. + tracker_ids: Touched ids recorded by the tracker for this table. + flushed_module_ids: Modules already flushed this dump; a module + hosting several tables flushes once, not once per table. + dump_evicted_tombstones: Whether to drain this table's evicted + keys and publish tombstones for the missing ones. + + Returns: + Number of rows appended. + """ try: from dynamicemb.types import CopyMode except ImportError as exc: raise RuntimeError( "dynamicemb is required to dump dynamic embedding values." ) from exc - # flush() flushes the whole module; only the first table of a - # multi-table module needs it within a dump. - if flushed_module_ids is None or id(dynamic_module) not in flushed_module_ids: - dynamic_module.flush() - if flushed_module_ids is not None: - flushed_module_ids.add(id(dynamic_module)) table_name = fqn.rsplit(".", maxsplit=1)[-1] + # flush() both syncs the caching store for the lookup and can evict + # backing-store rows; flush before the pop and the find so both see + # the post-flush state. + if id(dynamic_module) not in flushed_module_ids: + dynamic_module.flush() + flushed_module_ids.add(id(dynamic_module)) table_id = dynamic_module.table_names.index(table_name) - device = torch.device(f"cuda:{torch.cuda.current_device()}") - ids = ids.to(device=device, dtype=torch.int64) - table_ids = torch.full_like(ids, table_id, dtype=torch.int64) - _, _, _, _, _, founds, _, values = dynamic_module.tables.find( - ids, table_ids, CopyMode.EMBEDDING - ) + # pyre-ignore [29] emb_dim = dynamic_module._dynamicemb_options[table_id].dim - founds = founds.to(dtype=torch.bool) - if not bool(founds.all().item()): + ids = tracker_ids.cpu().to(torch.int64) + evicted_ids: Optional[torch.Tensor] = None + if dump_evicted_tombstones: + evicted = self._pop_evicted_key_ids(fqn, dynamic_module, table_name) + if evicted is not None and evicted.numel() > 0: + evicted_ids = evicted.cpu().to(torch.int64) + ids = torch.cat([ids, evicted_ids]) + ids = ids.unique(sorted=True) + if ids.numel() == 0: + return 0 + device = torch.device(f"cuda:{torch.cuda.current_device()}") + # The evicted mask stays on CPU: per-batch tombstone and never-admitted + # selections reuse the CPU batch ids, so only the found mask and the + # embedding values cross the device boundary. + evicted_mask = torch.isin(ids, evicted_ids) if evicted_ids is not None else None + feature_name = _feature_name(self._tracker.fqn_to_feature_names.get(fqn, [])) + num_rows = 0 + num_never_admitted = 0 + for start in range(0, ids.numel(), _FIND_BATCH_ROWS): + batch_ids = ids[start : start + _FIND_BATCH_ROWS] + gpu_ids = batch_ids.to(device=device) + table_ids = torch.full_like(gpu_ids, table_id, dtype=torch.int64) + _, _, _, _, _, founds, _, values = dynamic_module.tables.find( + gpu_ids, table_ids, CopyMode.EMBEDDING + ) + founds = founds.to(dtype=torch.bool) + # Only the compact found mask moves D2H; found and tombstone key + # ids reuse the CPU batch ids instead of round-tripping GPU ids. + founds_cpu = founds.cpu() + missing_cpu = ~founds_cpu + if evicted_mask is not None: + batch_evicted = evicted_mask[start : start + _FIND_BATCH_ROWS] + tombstone_ids = batch_ids[missing_cpu & batch_evicted] + num_never_admitted += int((missing_cpu & ~batch_evicted).sum()) + else: + tombstone_ids = batch_ids.new_empty(0) + num_never_admitted += int(missing_cpu.sum()) + num_rows += self._append_table_chunk( + table_chunks, + global_step=global_step, + feature_name=feature_name, + table_fqn=fqn, + key_ids=batch_ids[founds_cpu], + embeddings=values[founds, :emb_dim].detach(), + source="model_delta_tracker", + ) + num_rows += self._append_tombstone_chunks( + table_chunks, + global_step=global_step, + fqn=fqn, + tombstone_ids=tombstone_ids, + emb_dim=emb_dim, + ) + if num_never_admitted > 0: logger.warning( "Skip %s missing dynamic embedding ids for table %s.", - int((~founds).sum().item()), + num_never_admitted, fqn, ) - return values[founds, :emb_dim].detach(), ids[founds] + return num_rows + + def _warn_no_retain_table_once(self, table_fqn: str, reason: str) -> None: + """Warn once per table why its evicted keys cannot be tombstoned.""" + if table_fqn in self._warned_no_retain_tables: + return + self._warned_no_retain_tables.add(table_fqn) + logger.warning( + "Delta embedding dump cannot tombstone evicted keys for table " + f"{table_fqn}: {reason}; stale FeatureStore rows for its evicted " + "keys will never be reclaimed." + ) + + def _append_tombstone_chunks( + self, + table_chunks: List[pa.Table], + global_step: int, + fqn: str, + tombstone_ids: torch.Tensor, + emb_dim: int, + ) -> int: + """Append zero-row tombstones for evicted keys in bounded chunks. + + One zero buffer backs every chunk: _append_table_chunk wraps + row-slice views zero-copy, so Arrow chunks share this storage and + retained tombstone zeros stay one chunk-sized buffer regardless of + eviction volume. Under INT8 quantization the zero buffer is + quantized once up front and shared the same way, because + quantizing per chunk would allocate a separate buffer per chunk. + + Args: + table_chunks: List to append the per-table parquet chunks to. + global_step: Current training step. + fqn: Dynamic table FQN. + tombstone_ids: Evicted key ids to tombstone. + emb_dim: The table's embedding dimension. + + Returns: + Number of tombstone rows appended. + """ + if tombstone_ids.numel() == 0: + return 0 + feature_name = _feature_name(self._tracker.fqn_to_feature_names.get(fqn, [])) + chunk_rows = min(_TOMBSTONE_CHUNK_ROWS, tombstone_ids.numel()) + zero_rows: torch.Tensor = torch.zeros( + (chunk_rows, emb_dim), dtype=torch.float32 + ) + pre_quantized = False + if self._quant_type == DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8: + try: + quantized_zero = distributed_quantize_embeddings( + zero_rows, + emb_dim, + feature_name, + DISTRIBUTED_SPARSE_SUPPORTED_QUANT_FORMATS[0], + ) + except ValueError as e: + # quant_util errors address the distributed-export DIST_QUANT + # switch; delta dump quantization is toggled by quant_type. + raise ValueError( + "Delta embedding dump INT8 quantization failed for " + f"feature '{feature_name}' (table '{fqn}'): {e}. " + "Disable delta dump quantization by setting " + "delta_embedding_dump_config.quant_type to " + "DELTA_EMBEDDING_QUANT_NONE." + ) from e + zero_rows = torch.from_numpy(quantized_zero) + pre_quantized = True + num_rows = 0 + for start in range(0, tombstone_ids.numel(), chunk_rows): + key_chunk = tombstone_ids[start : start + chunk_rows] + num_rows += self._append_table_chunk( + table_chunks, + global_step=global_step, + feature_name=feature_name, + table_fqn=fqn, + key_ids=key_chunk, + embeddings=zero_rows[: key_chunk.numel()], + source="dynamicemb_evicted", + pre_quantized=pre_quantized, + ) + return num_rows def _collect_table_shard_infos(self) -> Dict[str, _TableShardInfo]: table_shard_infos: Dict[str, _TableShardInfo] = {} @@ -1217,9 +1469,13 @@ def _append_table_chunk( key_ids: torch.Tensor, embeddings: torch.Tensor, source: str, + pre_quantized: bool = False, ) -> int: key_ids_cpu = key_ids.detach().cpu().to(torch.int64).contiguous() - embeddings_cpu = embeddings.detach().cpu().to(torch.float32).contiguous() + if pre_quantized: + embeddings_cpu = embeddings.detach().cpu().contiguous() + else: + embeddings_cpu = embeddings.detach().cpu().to(torch.float32).contiguous() if embeddings_cpu.dim() != 2: raise ValueError( "delta embedding dump expects a 2-D embedding tensor, " @@ -1234,25 +1490,27 @@ def _append_table_chunk( f"key_ids={num_rows}, embeddings={embeddings_cpu.size(0)}." ) if self._quant_type == DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8: - emb_dim = embeddings_cpu.size(1) - try: - quantized = distributed_quantize_embeddings( - embeddings_cpu, - emb_dim, - feature_name, - DISTRIBUTED_SPARSE_SUPPORTED_QUANT_FORMATS[0], - ) - except ValueError as e: - # quant_util errors address the distributed-export DIST_QUANT - # switch; delta dump quantization is toggled by quant_type. - raise ValueError( - "Delta embedding dump INT8 quantization failed for " - f"feature '{feature_name}' (table '{table_fqn}'): {e}. " - "Disable delta dump quantization by setting " - "delta_embedding_dump_config.quant_type to " - "DELTA_EMBEDDING_QUANT_NONE." - ) from e - embeddings_cpu = torch.from_numpy(quantized) + if not pre_quantized: + emb_dim = embeddings_cpu.size(1) + try: + quantized = distributed_quantize_embeddings( + embeddings_cpu, + emb_dim, + feature_name, + DISTRIBUTED_SPARSE_SUPPORTED_QUANT_FORMATS[0], + ) + except ValueError as e: + # quant_util errors address the distributed-export + # DIST_QUANT switch; delta dump quantization is + # toggled by quant_type. + raise ValueError( + "Delta embedding dump INT8 quantization failed for " + f"feature '{feature_name}' (table '{table_fqn}'): {e}. " + "Disable delta dump quantization by setting " + "delta_embedding_dump_config.quant_type to " + "DELTA_EMBEDDING_QUANT_NONE." + ) from e + embeddings_cpu = torch.from_numpy(quantized) value_type = pa.uint8() else: value_type = pa.float32() diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 42d8de25f..771e8aa96 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -12,10 +12,11 @@ import glob import os import shutil +import sys import tempfile import unittest from contextlib import contextmanager -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from typing import Dict, List from unittest import mock @@ -78,6 +79,7 @@ ) from torchrec.types import DataType +from tzrec.protos import feature_pb2 from tzrec.protos.train_pb2 import ( DeltaEmbeddingDumpConfig, DeltaEmbeddingQuantType, @@ -86,10 +88,12 @@ from tzrec.tests import utils as test_utils from tzrec.utils import config_util from tzrec.utils.delta_embedding_dump import ( + _CONSUMER, _DELTA_DUMP_QUANT_SCHEMA, _DELTA_DUMP_SCHEMA, DeltaEmbeddingDumper, ModelDeltaTracker, + _dynamicemb_pop_evicted_keys_supported, _local_table_weight, _table_shard_info_from_config, _TableShardInfo, @@ -97,11 +101,17 @@ _validate_table_shard_info, validate_delta_embedding_dump_config, ) -from tzrec.utils.dynamicemb_util import has_dynamicemb +from tzrec.utils.dynamicemb_util import ( + _validate_eval_initializer_for_tombstones, + build_dynamicemb_constraints, + has_dynamicemb, + set_auto_retain_evicted_keys, +) from tzrec.utils.feature_store_delta_uploader import ( FEATURE_STORE_EMBEDDING_TYPE_FLOAT, FEATURE_STORE_EMBEDDING_TYPE_UINT8, ) +from tzrec.utils.quant_util import dequantize_quint8_rowwise_f16 from tzrec.utils.test_util import gpu_unavailable, make_test_dir, mark_ci_scope from tzrec.utils.zch_util import register_post_zch_event_tracker_fn @@ -118,6 +128,14 @@ _SHARED_EBC_EMBEDDING_DIM = 4 _SHARED_EC_EMBEDDING_DIM = 8 +try: + from dynamicemb.dynamicemb_config import EvictedItemMode + + _HAS_EVICTED_ITEM_MODE = True +except ImportError: + EvictedItemMode = None + _HAS_EVICTED_ITEM_MODE = False + class _DeltaDumpEBCModel(nn.Module): def __init__(self) -> None: @@ -173,27 +191,56 @@ def forward(self, features: KeyedJaggedTensor) -> torch.Tensor: class _FakeDynamicTables: - def __init__(self) -> None: + def __init__(self, founds=None, values=None, rows=None) -> None: self.ids = None self.table_ids = None self.copy_mode = None + self.call_count = 0 + self._rows = rows + self._founds = founds if founds is not None else [True, False, True] + self._values = ( + values + if values is not None + else [[1.0, 2.0, 20.0], [3.0, 4.0, 40.0], [5.0, 6.0, 60.0]] + ) def find(self, ids, table_ids, copy_mode): - self.ids = ids.detach().clone() - self.table_ids = table_ids.detach().clone() - self.copy_mode = copy_mode - founds = torch.tensor([True, False, True], device=ids.device) - values = torch.tensor( - [ - [1.0, 2.0, 20.0], - [3.0, 4.0, 40.0], - [5.0, 6.0, 60.0], - ], - device=ids.device, + self.call_count += 1 + ids = ids.detach().clone() + self.ids = ids if self.ids is None else torch.cat([self.ids, ids]) + table_ids = table_ids.detach().clone() + self.table_ids = ( + table_ids + if self.table_ids is None + else torch.cat([self.table_ids, table_ids]) ) + self.copy_mode = copy_mode + if self._rows is not None: + dim = len(next(iter(self._rows.values()))) + founds = torch.tensor( + [int(key.item()) in self._rows for key in ids], device=ids.device + ) + values = torch.tensor( + [self._rows.get(int(key.item()), [0.0] * dim) for key in ids], + device=ids.device, + ) + return None, None, None, None, None, founds, None, values + founds = torch.tensor(self._founds, device=ids.device) + values = torch.tensor(self._values, device=ids.device) return None, None, None, None, None, founds, None, values +def _dynamicemb_stub_modules() -> Dict[str, ModuleType]: + # The _append_dynamic_rows tests fake the whole dynamic module; the only + # real dynamicemb dependency is the CopyMode import, so stubbing it lets + # them run (and cover device placement) without dynamicemb installed. + types_mod = ModuleType("dynamicemb.types") + types_mod.CopyMode = SimpleNamespace(EMBEDDING="EMBEDDING") + mod = ModuleType("dynamicemb") + mod.types = types_mod + return {"dynamicemb": mod, "dynamicemb.types": types_mod} + + def _build_sharded_delta_dump_model(rank: int, world_size: int, ctx): torch.manual_seed(2026) device = torch.device(f"cuda:{rank}") @@ -1640,6 +1687,7 @@ def test_multi_gpu_dump_writes_empty_shard_when_rank_has_no_delta(self): dumper._retain_local_dump = False dumper._quant_type = DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_NONE dumper._schema = _DELTA_DUMP_SCHEMA + dumper._dump_evicted_tombstones = False with ( mock.patch.object(dumper, "_collect_table_weights", return_value={}), mock.patch.object(dumper, "_collect_dynamic_modules", return_value={}), @@ -1671,6 +1719,7 @@ def test_single_gpu_dump_skips_file_when_rank_has_no_delta(self): dumper._retain_local_dump = False dumper._quant_type = DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_NONE dumper._schema = _DELTA_DUMP_SCHEMA + dumper._dump_evicted_tombstones = False with ( mock.patch.object(dumper, "_collect_table_weights", return_value={}), mock.patch.object(dumper, "_collect_dynamic_modules", return_value={}), @@ -1841,7 +1890,6 @@ def test_row_wise_lookup_outputs_global_key_ids(self): ), ) }, - dynamic_modules={}, ) torch.testing.assert_close(embeddings, weight[[0, 2]]) torch.testing.assert_close(key_ids, torch.tensor([32, 34])) @@ -1872,7 +1920,6 @@ def test_lookup_fails_on_ids_outside_table_rows(self): ), ) }, - dynamic_modules={}, ) def test_zch_lookup_binds_held_ids_to_their_rows(self): @@ -1947,7 +1994,6 @@ def _lookup_zch(self, dumper, table_fqn, table_weights, ids): table_fqn, torch.tensor(ids, dtype=torch.int64), table_weights=table_weights, - dynamic_modules={}, ) def test_zch_delta_publishes_touched_held_id_with_row(self): @@ -2088,7 +2134,6 @@ def test_lookup_handles_empty_ids(self): ), ) }, - dynamic_modules={}, ) self.assertEqual(embeddings.shape, (0, 2)) self.assertEqual(key_ids.shape, (0,)) @@ -2113,17 +2158,13 @@ def test_row_wise_lookup_requires_shard_metadata(self): ), ) }, - dynamic_modules={}, ) - @unittest.skipUnless(has_dynamicemb, "dynamicemb is not installed; skipping.") @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") @mark_ci_scope("gpu") - def test_lookup_dynamic_embeddings_filters_missing_ids(self): - from dynamicemb.types import CopyMode - + def test_append_dynamic_rows_publishes_found_rows_only(self): torch.cuda.set_device(0) - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._eviction_dumper() fake_tables = _FakeDynamicTables() dynamic_module = SimpleNamespace( table_names=["dyn_table"], @@ -2131,28 +2172,39 @@ def test_lookup_dynamic_embeddings_filters_missing_ids(self): flush=mock.MagicMock(), _dynamicemb_options=[SimpleNamespace(dim=2)], ) + table_chunks = [] - embeddings, key_ids = dumper._lookup_dynamic_embeddings( - dynamic_module, - "model.ec.embeddings.dyn_table", - torch.tensor([101, 102, 103]), - ) + with mock.patch.dict(sys.modules, _dynamicemb_stub_modules()): + from dynamicemb.types import CopyMode + num_rows = dumper._append_dynamic_rows( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + dynamic_module=dynamic_module, + tracker_ids=torch.tensor([101, 102, 103]), + flushed_module_ids=set(), + dump_evicted_tombstones=False, + ) + + self.assertEqual(num_rows, 2) dynamic_module.flush.assert_called_once_with() self.assertIs(fake_tables.copy_mode, CopyMode.EMBEDDING) torch.testing.assert_close(fake_tables.ids.cpu(), torch.tensor([101, 102, 103])) torch.testing.assert_close(fake_tables.table_ids.cpu(), torch.tensor([0, 0, 0])) - torch.testing.assert_close(key_ids.cpu(), torch.tensor([101, 103])) - torch.testing.assert_close( - embeddings.cpu(), torch.tensor([[1.0, 2.0], [5.0, 6.0]]) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), [101, 103]) + self.assertEqual(table["embedding"].to_pylist(), [[1.0, 2.0], [5.0, 6.0]]) + self.assertEqual( + table["source"].to_pylist(), + ["model_delta_tracker", "model_delta_tracker"], ) - @unittest.skipUnless(has_dynamicemb, "dynamicemb is not installed; skipping.") @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") @mark_ci_scope("gpu") - def test_lookup_dynamic_embeddings_flushes_module_once_per_dump(self): + def test_append_dynamic_rows_flushes_module_once_per_dump(self): torch.cuda.set_device(0) - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._eviction_dumper() # One module hosting two tables, reachable under both table_name keys. dynamic_module = SimpleNamespace( table_names=["dyn_a", "dyn_b"], @@ -2162,17 +2214,480 @@ def test_lookup_dynamic_embeddings_flushes_module_once_per_dump(self): ) flushed_module_ids = set() - for table_name in ("dyn_a", "dyn_b"): - dumper._lookup_dynamic_embeddings( - dynamic_module, - f"model.ec.embeddings.{table_name}", - torch.tensor([101, 102, 103]), - flushed_module_ids, + with mock.patch.dict(sys.modules, _dynamicemb_stub_modules()): + for table_name in ("dyn_a", "dyn_b"): + dumper._append_dynamic_rows( + [], + global_step=5, + fqn=f"model.ec.embeddings.{table_name}", + dynamic_module=dynamic_module, + tracker_ids=torch.tensor([101, 102, 103]), + flushed_module_ids=flushed_module_ids, + dump_evicted_tombstones=False, + ) + + # Both tables share the module; flush() flushes all tables, so it + # runs once per dump rather than once per table. + dynamic_module.flush.assert_called_once_with() + # A module already flushed this dump is not flushed again. + dynamic_module.flush.reset_mock() + dumper._append_dynamic_rows( + [], + global_step=6, + fqn="model.ec.embeddings.dyn_a", + dynamic_module=dynamic_module, + tracker_ids=torch.tensor([101, 102, 103]), + flushed_module_ids={id(dynamic_module)}, + dump_evicted_tombstones=False, + ) + dynamic_module.flush.assert_not_called() + + _DYN_TABLE_FQN = "model.ec.embeddings.dyn_table" + + def _eviction_dumper(self, quant_type=None): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._rank = 0 + dumper._world_size = 1 + dumper._quant_type = ( + quant_type or DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_NONE + ) + dumper._schema = ( + _DELTA_DUMP_QUANT_SCHEMA + if dumper._quant_type == DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8 + else _DELTA_DUMP_SCHEMA + ) + dumper._tracker = SimpleNamespace( + fqn_to_feature_names={self._DYN_TABLE_FQN: ["user_id"]} + ) + dumper._warned_no_retain_tables = set() + return dumper + + def _eviction_module(self, evicted_keys, founds=None, values=None, rows=None): + table_name = self._DYN_TABLE_FQN.rsplit(".", maxsplit=1)[-1] + return SimpleNamespace( + table_names=[table_name], + _dynamicemb_options=[SimpleNamespace(dim=2)], + tables=_FakeDynamicTables(founds=founds, values=values, rows=rows), + flush=mock.MagicMock(), + pop_evicted_keys=mock.MagicMock( + return_value={table_name: torch.tensor(evicted_keys, dtype=torch.int64)} + ), + ) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") + @mark_ci_scope("gpu") + def test_append_dynamic_rows_tombstones_evicted_keys(self): + torch.cuda.set_device(0) + dumper = self._eviction_dumper() + dynamic_module = self._eviction_module( + [7, 9], + founds=[False, False], + values=[[0.0, 0.0], [0.0, 0.0]], + ) + table_chunks = [] + + with mock.patch.dict(sys.modules, _dynamicemb_stub_modules()): + num_rows = dumper._append_dynamic_rows( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + dynamic_module=dynamic_module, + tracker_ids=torch.empty(0, dtype=torch.int64), + flushed_module_ids=set(), + dump_evicted_tombstones=True, ) - # Both tables share the module; flush() flushes all tables, so it runs - # once per dump rather than once per table. + self.assertEqual(num_rows, 2) + # flush() precedes the pop so flush-induced evictions join the drain. dynamic_module.flush.assert_called_once_with() + dynamic_module.pop_evicted_keys.assert_called_once_with(["dyn_table"]) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), [7, 9]) + self.assertEqual(table["embedding"].to_pylist(), [[0.0, 0.0], [0.0, 0.0]]) + self.assertEqual( + table["source"].to_pylist(), ["dynamicemb_evicted", "dynamicemb_evicted"] + ) + self.assertEqual(table["feature_name"].to_pylist(), ["user_id", "user_id"]) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") + @mark_ci_scope("gpu") + def test_append_dynamic_rows_publishes_fresh_row_for_reinserted_key(self): + # 7 and 9 were evicted but reinserted, so the merged lookup finds them + # and each publishes exactly one fresh row; 8 stayed evicted and + # publishes one tombstone. + torch.cuda.set_device(0) + dumper = self._eviction_dumper() + dynamic_module = self._eviction_module( + [7, 8, 9], + founds=[True, False, True], + values=[[1.0, 2.0, 0.0], [3.0, 4.0, 0.0], [5.0, 6.0, 0.0]], + ) + table_chunks = [] + + with mock.patch.dict(sys.modules, _dynamicemb_stub_modules()): + num_rows = dumper._append_dynamic_rows( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + dynamic_module=dynamic_module, + tracker_ids=torch.tensor([7, 9]), + flushed_module_ids=set(), + dump_evicted_tombstones=True, + ) + + self.assertEqual(num_rows, 3) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), [7, 9, 8]) + self.assertEqual( + table["embedding"].to_pylist(), [[1.0, 2.0], [5.0, 6.0], [0.0, 0.0]] + ) + self.assertEqual( + table["source"].to_pylist(), + ["model_delta_tracker", "model_delta_tracker", "dynamicemb_evicted"], + ) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") + @mark_ci_scope("gpu") + def test_append_dynamic_rows_tombstones_tracked_then_evicted_key(self): + # 102 was tracked but find() misses it because it was evicted after + # the touch; the merged lookup tombstones it directly. + torch.cuda.set_device(0) + dumper = self._eviction_dumper() + dynamic_module = self._eviction_module( + [102], + founds=[True, False, True], + values=[[1.0, 2.0, 0.0], [3.0, 4.0, 0.0], [5.0, 6.0, 0.0]], + ) + table_chunks = [] + + with mock.patch.dict(sys.modules, _dynamicemb_stub_modules()): + num_rows = dumper._append_dynamic_rows( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + dynamic_module=dynamic_module, + tracker_ids=torch.tensor([101, 102, 103]), + flushed_module_ids=set(), + dump_evicted_tombstones=True, + ) + + self.assertEqual(num_rows, 3) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), [101, 103, 102]) + self.assertEqual( + table["source"].to_pylist(), + ["model_delta_tracker", "model_delta_tracker", "dynamicemb_evicted"], + ) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for dynamicemb.") + @mark_ci_scope("gpu") + def test_append_dynamic_rows_batches_merged_lookup(self): + # Merged set [101, 102, 103, 200] looked up 2 ids per batch: 102 was + # tracked then evicted (tombstone), 103 was tracked but never admitted + # (one warning across both batches), 200 was evicted then reinserted and + # publishes its fresh row from the second batch. + torch.cuda.set_device(0) + dumper = self._eviction_dumper() + dynamic_module = self._eviction_module( + [102, 200], rows={101: [1.0, 2.0, 0.0], 200: [5.0, 6.0, 0.0]} + ) + table_chunks = [] + + with mock.patch.dict(sys.modules, _dynamicemb_stub_modules()): + with mock.patch("tzrec.utils.delta_embedding_dump._FIND_BATCH_ROWS", 2): + with mock.patch("tzrec.utils.delta_embedding_dump.logger") as log: + num_rows = dumper._append_dynamic_rows( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + dynamic_module=dynamic_module, + tracker_ids=torch.tensor([101, 102, 103]), + flushed_module_ids=set(), + dump_evicted_tombstones=True, + ) + + self.assertEqual(dynamic_module.tables.call_count, 2) + torch.testing.assert_close( + dynamic_module.tables.ids.cpu(), torch.tensor([101, 102, 103, 200]) + ) + self.assertEqual(num_rows, 3) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), [101, 102, 200]) + self.assertEqual( + table["embedding"].to_pylist(), [[1.0, 2.0], [0.0, 0.0], [5.0, 6.0]] + ) + self.assertEqual( + table["source"].to_pylist(), + ["model_delta_tracker", "dynamicemb_evicted", "model_delta_tracker"], + ) + self.assertEqual(log.warning.call_count, 1) + + def test_pop_evicted_key_ids_warns_on_discard_and_old_builds(self): + dumper = self._eviction_dumper() + discard_module = SimpleNamespace( + pop_evicted_keys=mock.MagicMock(return_value={}) + ) + empty_module = SimpleNamespace( + pop_evicted_keys=mock.MagicMock( + return_value={"empty_table": torch.tensor([], dtype=torch.int64)} + ) + ) + old_module = SimpleNamespace() # dynamicemb without pop_evicted_keys + modules = { + "model.ec.embeddings.discard_table": discard_module, + "model.ec.embeddings.empty_table": empty_module, + "model.ec.embeddings.old_table": old_module, + } + + with mock.patch("tzrec.utils.delta_embedding_dump.logger") as log: + popped = { + fqn: dumper._pop_evicted_key_ids( + fqn, module, fqn.rsplit(".", maxsplit=1)[-1] + ) + for fqn, module in modules.items() + } + + self.assertIsNone(popped["model.ec.embeddings.discard_table"]) + self.assertIsNone(popped["model.ec.embeddings.old_table"]) + self.assertIsNotNone(popped["model.ec.embeddings.empty_table"]) + self.assertEqual(popped["model.ec.embeddings.empty_table"].numel(), 0) + # DISCARD tables and pop-less builds warn once each; an empty eviction + # buffer is silent. + self.assertEqual(log.warning.call_count, 2) + warned = " ".join(str(call) for call in log.warning.call_args_list) + self.assertIn("model.ec.embeddings.discard_table", warned) + self.assertIn("model.ec.embeddings.old_table", warned) + self.assertNotIn("model.ec.embeddings.empty_table", warned) + # The one-time warnings do not repeat on the next dump. + with mock.patch("tzrec.utils.delta_embedding_dump.logger") as log: + for fqn, module in modules.items(): + dumper._pop_evicted_key_ids( + fqn, module, fqn.rsplit(".", maxsplit=1)[-1] + ) + log.warning.assert_not_called() + + def test_append_tombstone_chunks_quant_bytes_pin_tombstone_format(self): + # A zero row through INT8 quantization is exactly [codes=0,0][fp16 + # scale 1.0 little-endian 0x00,0x3C=60][fp16 offset 0.0] -- the byte + # contract the processor's tombstone detector is pinned to (the raw + # bytes are NOT all zero). + dumper = self._eviction_dumper( + DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8 + ) + table_chunks = [] + + num_rows = dumper._append_tombstone_chunks( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + tombstone_ids=torch.tensor([11], dtype=torch.int64), + emb_dim=2, + ) + + self.assertEqual(num_rows, 1) + table = pa.concat_tables(table_chunks) + rows = table["embedding"].to_pylist() + self.assertEqual(rows, [[0, 0, 0, 60, 0, 0]]) + decoded = dequantize_quint8_rowwise_f16( + np.asarray(rows, dtype=np.uint8), emb_dim=2 + ) + np.testing.assert_array_equal(decoded, np.zeros((1, 2), dtype=np.float32)) + # Bitwise +0.0, not -0.0: NvEmbeddings dequantizes to fp16 0x0000. + self.assertFalse(np.signbit(decoded).any()) + + def test_append_tombstone_chunks_chunks_large_evictions(self): + # Evictions beyond the chunk size split into ordered chunks whose + # embeddings all alias one chunk-sized zero buffer. + dumper = self._eviction_dumper() + tombstone_ids = torch.tensor(list(range(10)), dtype=torch.int64) + table_chunks = [] + + with mock.patch("tzrec.utils.delta_embedding_dump._TOMBSTONE_CHUNK_ROWS", 3): + num_rows = dumper._append_tombstone_chunks( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + tombstone_ids=tombstone_ids, + emb_dim=2, + ) + + self.assertEqual(num_rows, 10) + # 3 + 3 + 3 + 1 rows: chunk boundaries at the constant, order kept. + self.assertEqual([chunk.num_rows for chunk in table_chunks], [3, 3, 3, 1]) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), list(range(10))) + self.assertEqual(table["embedding"].to_pylist(), [[0.0, 0.0]] * 10) + # Every chunk's embedding values wrap the same storage, so retained + # tombstone zeros stay one chunk regardless of eviction count. + addresses = [ + chunk.column("embedding").chunk(0).values.buffers()[1].address + for chunk in table_chunks + ] + self.assertEqual(len(set(addresses)), 1) + + def test_append_tombstone_chunks_int8_chunks_share_one_buffer(self): + # INT8 tombstones quantize the zero buffer once up front; every chunk + # aliases that one buffer, so retained tombstone zeros stay one chunk + # regardless of eviction count instead of one allocation per chunk. + dumper = self._eviction_dumper( + DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8 + ) + tombstone_ids = torch.tensor(list(range(10)), dtype=torch.int64) + table_chunks = [] + + with mock.patch("tzrec.utils.delta_embedding_dump._TOMBSTONE_CHUNK_ROWS", 3): + num_rows = dumper._append_tombstone_chunks( + table_chunks, + global_step=5, + fqn=self._DYN_TABLE_FQN, + tombstone_ids=tombstone_ids, + emb_dim=2, + ) + + self.assertEqual(num_rows, 10) + self.assertEqual([chunk.num_rows for chunk in table_chunks], [3, 3, 3, 1]) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), list(range(10))) + self.assertEqual(table["embedding"].to_pylist(), [[0, 0, 0, 60, 0, 0]] * 10) + addresses = [ + chunk.column("embedding").chunk(0).values.buffers()[1].address + for chunk in table_chunks + ] + self.assertEqual(len(set(addresses)), 1) + + def test_clear_discards_retained_evicted_keys(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._dump_evicted_tombstones = True + tracker = SimpleNamespace(clear=mock.MagicMock()) + dumper._tracker = tracker + dynamic_module = SimpleNamespace( + table_names=["dyn_a", "dyn_b"], + pop_evicted_keys=mock.MagicMock(return_value={}), + ) + old_module = SimpleNamespace(table_names=["old_table"]) + with mock.patch.object( + dumper, + "_collect_dynamic_modules", + return_value={ + "model.ec.embeddings.dyn_a": dynamic_module, + "model.ec.embeddings.dyn_b": dynamic_module, + "model.ec.embeddings.old_table": old_module, + }, + ): + dumper.clear() + + tracker.clear.assert_called_once_with(_CONSUMER) + # One module hosts several tables; pop() drains them all, once. + dynamic_module.pop_evicted_keys.assert_called_once_with() + + def test_clear_without_tombstones_keeps_eviction_buffer(self): + dumper = object.__new__(DeltaEmbeddingDumper) + dumper._dump_evicted_tombstones = False + dumper._tracker = SimpleNamespace(clear=mock.MagicMock()) + dynamic_module = SimpleNamespace( + table_names=["dyn_table"], + pop_evicted_keys=mock.MagicMock(return_value={}), + ) + with mock.patch.object( + dumper, + "_collect_dynamic_modules", + return_value={"model.ec.embeddings.dyn_table": dynamic_module}, + ): + dumper.clear() + + dumper._tracker.clear.assert_called_once_with(_CONSUMER) + dynamic_module.pop_evicted_keys.assert_not_called() + + +class DynamicembTombstoneEvalInitializerTest(unittest.TestCase): + def setUp(self): + set_auto_retain_evicted_keys(False) + self.addCleanup(set_auto_retain_evicted_keys, False) + + def _cfg(self, **init_kwargs): + cfg = feature_pb2.DynamicEmbedding(max_capacity=1024) + if init_kwargs: + cfg.eval_initializer_args.CopyFrom( + feature_pb2.DynamicEmbInitializerArgs(**init_kwargs) + ) + return cfg + + def test_unset_eval_initializer_is_allowed(self): + set_auto_retain_evicted_keys(True) + _validate_eval_initializer_for_tombstones(self._cfg(), "dyn_table") + + def test_constant_zero_eval_initializer_is_allowed(self): + set_auto_retain_evicted_keys(True) + _validate_eval_initializer_for_tombstones( + self._cfg(mode="CONSTANT"), "dyn_table" + ) + _validate_eval_initializer_for_tombstones( + self._cfg(mode="CONSTANT", value=0.0), "dyn_table" + ) + + def test_nonzero_constant_eval_initializer_rejected(self): + set_auto_retain_evicted_keys(True) + with self.assertRaisesRegex(ValueError, "dyn_table"): + _validate_eval_initializer_for_tombstones( + self._cfg(mode="CONSTANT", value=0.5), "dyn_table" + ) + + def test_nonconstant_eval_initializer_rejected(self): + set_auto_retain_evicted_keys(True) + with self.assertRaisesRegex(ValueError, "mode=NORMAL"): + _validate_eval_initializer_for_tombstones( + self._cfg(mode="NORMAL"), "dyn_table" + ) + + def test_validation_inactive_when_tombstones_off(self): + _validate_eval_initializer_for_tombstones( + self._cfg(mode="CONSTANT", value=0.5), "dyn_table" + ) + + +@unittest.skipUnless( + has_dynamicemb and _HAS_EVICTED_ITEM_MODE, + "dynamicemb without EvictedItemMode is not installed; skipping.", +) +@mark_ci_scope("gpu") +class DynamicembUtilAutoRetainTest(unittest.TestCase): + def setUp(self): + set_auto_retain_evicted_keys(False) + self.addCleanup(set_auto_retain_evicted_keys, False) + self.dynamicemb_cfg = feature_pb2.DynamicEmbedding(max_capacity=1024) + self.emb_config = EmbeddingBagConfig( + name="dyn_table", + num_embeddings=1024, + embedding_dim=8, + feature_names=["user_id"], + ) + + def test_auto_retain_arms_evicted_item_mode(self): + set_auto_retain_evicted_keys(True) + + constraints = build_dynamicemb_constraints(self.dynamicemb_cfg, self.emb_config) + + self.assertEqual( + constraints.dynamicemb_options.evicted_item_mode, + EvictedItemMode.RETAIN_KEY, + ) + + def test_auto_retain_off_keeps_discard_mode(self): + constraints = build_dynamicemb_constraints(self.dynamicemb_cfg, self.emb_config) + + self.assertEqual( + constraints.dynamicemb_options.evicted_item_mode, + EvictedItemMode.DISCARD, + ) + + def test_auto_retain_rejects_nonzero_eval_initializer(self): + set_auto_retain_evicted_keys(True) + self.dynamicemb_cfg.eval_initializer_args.CopyFrom( + feature_pb2.DynamicEmbInitializerArgs(mode="CONSTANT", value=0.5) + ) + + with self.assertRaisesRegex(ValueError, "dump_evicted_tombstones"): + build_dynamicemb_constraints(self.dynamicemb_cfg, self.emb_config) class DeltaEmbeddingDumpShardedIntegrationTest(MultiProcessTestBase): @@ -2295,11 +2810,22 @@ def tearDown(self): @mark_ci_scope("gpu") def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): world_size = int(os.getenv("TEST_NPROC_PER_NODE", "2")) - pipeline_config = config_util.load_pipeline_config( - "tzrec/tests/configs/multi_tower_din_fg_dynamicemb_mock.config" - ) + config_path = "tzrec/tests/configs/multi_tower_din_fg_dynamicemb_mock.config" + # Prepare mock data against the config's original (large) id space + # before shrinking max_capacity: a dynamicemb feature's mock ids are + # drawn from [0, num_embeddings) and num_embeddings == max_capacity, + # so data generated after the shrink could never overflow the tables. + prepared_config = test_utils.load_config_for_test( + config_path, self.test_dir, user_id="user_id", item_id="item_id" + ) + pipeline_config = config_util.load_pipeline_config(config_path) + pipeline_config.train_input_path = prepared_config.train_input_path + pipeline_config.eval_input_path = prepared_config.eval_input_path + pipeline_config.data_config.num_workers = 2 # Admit every id immediately so the find() lookup returns embeddings - # for the touched ids (default frequency admission would hide them). + # for the touched ids (default frequency admission would hide them), + # and shrink the dynamic tables far below the prepared id space so + # they overflow and evict keys deterministically. for feature_config in pipeline_config.feature_configs: feature_type = feature_config.WhichOneof("feature") if feature_type is None: @@ -2311,6 +2837,7 @@ def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): admission = feature.dynamicemb.WhichOneof("admission_strategy") if admission is not None: feature.dynamicemb.ClearField(admission) + feature.dynamicemb.max_capacity = 1024 dump_dir = os.path.abspath(os.path.join(self.test_dir, "delta_dump")) dump_cfg = pipeline_config.train_config.delta_embedding_dump_config @@ -2332,6 +2859,7 @@ def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): self.assertTrue(step_dirs, f"no delta dump produced under {dump_dir}") dumped_real_rows = False + dumped_tombstone_rows = False for step_dir in step_dirs: shards = sorted(glob.glob(os.path.join(step_dir, "*.parquet"))) # Every rank writes a shard even with no delta, so each step dir @@ -2346,15 +2874,33 @@ def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): self.assertEqual(table.schema, _DELTA_DUMP_SCHEMA) if table.num_rows == 0: continue - dumped_real_rows = True self.assertEqual(set(table["world_size"].to_pylist()), {world_size}) - self.assertEqual( - set(table["source"].to_pylist()), {"model_delta_tracker"} + table_fqns = table["table_fqn"].to_pylist() + key_ids = table["key_id"].to_pylist() + sources = table["source"].to_pylist() + embeddings = table["embedding"].to_pylist() + self.assertLessEqual( + set(sources), {"model_delta_tracker", "dynamicemb_evicted"} ) # dynamic lookup must return a real embedding vector per id. - self.assertTrue( - all(len(emb) > 0 for emb in table["embedding"].to_pylist()) - ) + self.assertTrue(all(len(emb) > 0 for emb in embeddings)) + real_keys = { + (fqn, key) + for fqn, key, source in zip(table_fqns, key_ids, sources) + if source == "model_delta_tracker" + } + if real_keys: + dumped_real_rows = True + for fqn, key, source, emb in zip( + table_fqns, key_ids, sources, embeddings + ): + if source != "dynamicemb_evicted": + continue + dumped_tombstone_rows = True + # A tombstone is an all-zero row for a key this dump did + # not also publish as a real row (reinserted keys win). + self.assertEqual(emb, [0.0] * len(emb)) + self.assertNotIn((fqn, key), real_keys) # If no rank ever dumped a real row, the flush()/find() lookup path was # not actually exercised and the test would be vacuous. @@ -2362,6 +2908,15 @@ def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): dumped_real_rows, "no dynamic delta rows dumped; flush()/find() path not exercised", ) + # The shrunken max_capacity forces evictions during training; they + # must reach the shards as tombstones or the pop_evicted_keys drain + # was never exercised. Old dynamicemb builds lack pop_evicted_keys + # and degrade to a warning, so only assert where it exists. + if _dynamicemb_pop_evicted_keys_supported(): + self.assertTrue( + dumped_tombstone_rows, + "no evicted-key tombstones dumped; pop_evicted_keys path not exercised", + ) class DeltaEmbeddingDumpZchIntegrationTest(unittest.TestCase): diff --git a/tzrec/utils/dynamicemb_util.py b/tzrec/utils/dynamicemb_util.py index 3e374088f..281ba22d4 100644 --- a/tzrec/utils/dynamicemb_util.py +++ b/tzrec/utils/dynamicemb_util.py @@ -12,7 +12,7 @@ import dataclasses import math import os -from typing import Any, List, Optional, Tuple, Type, cast +from typing import Any, Dict, List, Optional, Tuple, Type, cast import torch from torchrec.distributed.embedding_types import EmbeddingComputeKernel @@ -52,6 +52,82 @@ _DYNAMICEMB_HYBRID_X_EFF_BASE = 0.11 _DYNAMICEMB_X_EFF_TIEBREAK = 0.01 +# Armed by train() when delta dump publishes evicted-key tombstones: that +# dump's pop_evicted_keys drain is the only consumer of the evicted-key +# buffer, so retention is enabled exactly when a consumer exists (an +# unconsumed buffer would grow without bound in GPU memory). +_auto_retain_evicted_keys = False +_warned_missing_evicted_item_mode = False + + +def set_auto_retain_evicted_keys(enabled: bool) -> None: + """Toggle RETAIN_KEY eviction recording for dynamicemb tables built later. + + Takes effect at plan/shard time (``build_dynamicemb_constraints``), which + cannot see the train config, hence this module-level switch. + + Args: + enabled: Whether new dynamicemb tables should retain evicted keys. + """ + global _auto_retain_evicted_keys + _auto_retain_evicted_keys = bool(enabled) + + +def _arm_evicted_key_retention(demb_opt_kwargs: Dict[str, Any]) -> None: + """Add evicted_item_mode=RETAIN_KEY to table options when auto-retain is on. + + Older dynamicemb builds have no EvictedItemMode; warn once and keep the + default DISCARD mode rather than raise, so the default-on tombstone dump + cannot break existing jobs -- tombstones are then silently missing. + + Args: + demb_opt_kwargs: Keyword arguments for DynamicEmbTableOptions. + """ + global _warned_missing_evicted_item_mode + try: + from dynamicemb.dynamicemb_config import EvictedItemMode + except ImportError: + if not _warned_missing_evicted_item_mode: + _warned_missing_evicted_item_mode = True + logger.warning( + "dynamicemb lacks EvictedItemMode; evicted-key tombstones " + "will be missing from delta embedding dumps." + ) + return + demb_opt_kwargs["evicted_item_mode"] = EvictedItemMode.RETAIN_KEY + + +def _validate_eval_initializer_for_tombstones( + dynamicemb_cfg: feature_pb2.DynamicEmbedding, table_name: str +) -> None: + """Reject non-zero eval initializers when evicted-key tombstones are armed. + + Tombstone dumps publish constant-zero rows for evicted keys, so a + non-zero eval initializer would make the same missing key resolve to + different values on the dump consumer and in eval lookups. + + Args: + dynamicemb_cfg: The feature's dynamic embedding config. + table_name: Embedding table name, for the error message. + + Raises: + ValueError: if eval_initializer_args is set and not CONSTANT 0.0. + """ + if not _auto_retain_evicted_keys: + return + if not dynamicemb_cfg.HasField("eval_initializer_args"): + return + init_cfg = dynamicemb_cfg.eval_initializer_args + mode = init_cfg.mode if init_cfg.HasField("mode") else "CONSTANT" + if mode != "CONSTANT" or init_cfg.value != 0.0: + raise ValueError( + f"dynamic embedding table {table_name} sets eval_initializer_args " + f"(mode={mode}, value={init_cfg.value}), but delta embedding dump " + "with dump_evicted_tombstones publishes constant-zero tombstones " + "for evicted keys; leave eval_initializer_args unset or set it " + "to CONSTANT 0.0." + ) + def _dynamicemb_effective_cache_ratio( cache_load_factor: Optional[float], @@ -246,6 +322,7 @@ def build_dynamicemb_constraints( "dynamicemb is not installed; required by features with " "`dynamicemb { }` set." ) + _validate_eval_initializer_for_tombstones(dynamicemb_cfg, emb_config.name) embedding_dim = emb_config.embedding_dim num_embeddings = emb_config.num_embeddings @@ -294,6 +371,8 @@ def build_dynamicemb_constraints( demb_opt_kwargs = {} if dynamicemb_cfg.HasField("bucket_capacity"): demb_opt_kwargs["bucket_capacity"] = dynamicemb_cfg.bucket_capacity + if _auto_retain_evicted_keys: + _arm_evicted_key_retention(demb_opt_kwargs) dynamicemb_options = dynamicemb.DynamicEmbTableOptions( max_capacity=dynamicemb_cfg.max_capacity, diff --git a/tzrec/version.py b/tzrec/version.py index 2b03e2c71..7c19cbf95 100644 --- a/tzrec/version.py +++ b/tzrec/version.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.3.17" +__version__ = "1.3.18"