From 987ffec558cd68ee9bf055f02f0f483c4f9a1557 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 21 Aug 2026 10:52:46 +0800 Subject: [PATCH 01/16] [feat] tombstone evicted dynamic embedding keys in delta dump Evicted dynamicemb keys were silently skipped at dump time, so their stale FeatureStore rows kept serving forever and the processor could never reclaim the memory. The dump now drains pop_evicted_keys after the tracker pass (capturing flush-induced evictions in the same window), subtracts keys already republished as real rows, and uploads the rest as all-zero tombstones through the existing chunk/quant path; the polling processor detects the zeros and deletes the key from NvEmbeddings. The new dump_evicted_tombstones switch (default true) gates the pass and auto-arms dynamicemb RETAIN_KEY recording for exactly those runs, degrading to a one-time warning on old dynamicemb builds. Also fix the trailing-comma JSON error in .pyre_configuration that broke pyre startup. --- .pyre_configuration | 2 +- tzrec/main.py | 13 +- tzrec/protos/train.proto | 7 + tzrec/utils/delta_embedding_dump.py | 167 ++++++++++- tzrec/utils/delta_embedding_dump_test.py | 361 ++++++++++++++++++++++- tzrec/utils/dynamicemb_util.py | 48 ++- 6 files changed, 586 insertions(+), 12 deletions(-) diff --git a/.pyre_configuration b/.pyre_configuration index f2ca8ba3a..bf7649a6d 100644 --- a/.pyre_configuration +++ b/.pyre_configuration @@ -7,7 +7,7 @@ "tzrec/utils/load_class.py", "tzrec/utils/filesystem_util.py", "tzrec/tools/convert_easyrec_config_to_tzrec_config.py", - "tzrec/ops/triton/*.py", + "tzrec/ops/triton/*.py" ], "site_package_search_strategy": "all", "source_directories": [ 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..d1aa306bd 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -108,6 +108,13 @@ 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. 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..da3b50eb3 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -162,6 +162,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 +677,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 +709,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: @@ -764,9 +782,23 @@ def clear(self) -> None: 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. + only the delta against the restored state. Retained evicted keys are + popped and discarded for the same reason: they predate the restored + baseline and must not be tombstoned by the next dump (the retain + buffer is process-local GPU memory and never survives a restart). """ 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]: @@ -911,12 +943,28 @@ def dump(self, global_step: int) -> Optional[str]: table_weights = self._collect_table_weights() dynamic_modules = self._collect_dynamic_modules() table_chunks: List[pa.Table] = [] + flushed_module_ids: Set[int] = set() + published_dynamic_key_ids: Dict[str, torch.Tensor] = {} num_rows = self._append_model_delta_rows( table_chunks, global_step=global_step, table_weights=table_weights, dynamic_modules=dynamic_modules, + flushed_module_ids=flushed_module_ids, + published_dynamic_key_ids=( + published_dynamic_key_ids if self._dump_evicted_tombstones else None + ), ) + if self._dump_evicted_tombstones: + # Must follow the tracker pass: its flush() can evict backing-store + # rows, and those evictions have to be captured by the same pop. + num_rows += self._append_dynamic_evicted_rows( + table_chunks, + global_step=global_step, + dynamic_modules=dynamic_modules, + published_key_ids=published_dynamic_key_ids, + flushed_module_ids=flushed_module_ids, + ) output_path: Optional[str] = None if write_local and (num_rows > 0 or self._world_size > 1): # Multi-rank shard sets stay complete even for an empty rank so @@ -965,12 +1013,32 @@ 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, + published_dynamic_key_ids: Optional[Dict[str, torch.Tensor]] = None, ) -> int: + """Append real rows for the tracker's touched ids. + + 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. + published_dynamic_key_ids: When given, records the dynamic keys + published as real rows here (founds-filtered, cpu int64), + keyed by table FQN, so the tombstone pass can subtract them + from the evicted keys. + + 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() + if flushed_module_ids is None: + flushed_module_ids = set() for fqn, unique_rows in self._tracker.get_unique(_CONSUMER).items(): ids = unique_rows.ids if ids.numel() == 0: @@ -983,6 +1051,13 @@ def _append_model_delta_rows( dynamic_modules=dynamic_modules, flushed_module_ids=flushed_module_ids, ) + if published_dynamic_key_ids is not None and fqn in dynamic_modules: + published_ids = key_ids.detach().cpu().to(torch.int64) + if fqn in published_dynamic_key_ids: + published_ids = torch.cat( + [published_dynamic_key_ids[fqn], published_ids] + ) + published_dynamic_key_ids[fqn] = published_ids feature_name = _feature_name( self._tracker.fqn_to_feature_names.get(fqn, []) ) @@ -1130,6 +1205,94 @@ def _lookup_dynamic_embeddings( ) return values[founds, :emb_dim].detach(), ids[founds] + 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_dynamic_evicted_rows( + self, + table_chunks: List[pa.Table], + global_step: int, + dynamic_modules: Dict[str, nn.Module], + published_key_ids: Dict[str, torch.Tensor], + flushed_module_ids: Set[int], + ) -> int: + """Append zero-row tombstones for dynamicemb keys evicted since the last dump. + + The authoritative eviction source is ``pop_evicted_keys`` (read-and-clear + semantics), not the tracker lookup's ``founds`` mask, which cannot tell + an evicted key from one never admitted. Keys already published as real + rows by this dump's tracker pass are subtracted, so an evicted-then- + reinserted key keeps its fresh row instead of racing a tombstone in the + MERGE upload. Each rank drains only its local shard (row-wise sharding + keeps 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. + dynamic_modules: Dynamic embedding modules keyed by table FQN; + tables without tracker rows are still visited, since their + eviction buffer must be drained too. + published_key_ids: Dynamic keys published as real rows this dump, + keyed by table FQN. + flushed_module_ids: Modules already flushed this dump, shared with + the tracker pass so each module flushes at most once. + + Returns: + Number of tombstone rows appended. + """ + num_rows = 0 + for fqn, dynamic_module in dynamic_modules.items(): + table_name = fqn.rsplit(".", maxsplit=1)[-1] + 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" + ) + continue + # flush() can evict backing-store rows; flush before the pop so + # those evictions join this drain (no-op for non-caching modules + # the tracker pass already flushed). + if id(dynamic_module) not in flushed_module_ids: + dynamic_module.flush() + flushed_module_ids.add(id(dynamic_module)) + 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)" + ) + continue + evicted = evicted.detach().cpu().to(torch.int64) + if evicted.numel() == 0: + continue + published = published_key_ids.get(fqn) + if published is not None and published.numel() > 0: + evicted = evicted[~torch.isin(evicted, published)] + if evicted.numel() == 0: + continue + table_id = dynamic_module.table_names.index(table_name) + # pyre-ignore [29] + emb_dim = dynamic_module._dynamicemb_options[table_id].dim + num_rows += self._append_table_chunk( + table_chunks, + global_step=global_step, + feature_name=_feature_name( + self._tracker.fqn_to_feature_names.get(fqn, []) + ), + table_fqn=fqn, + key_ids=evicted, + embeddings=torch.zeros((evicted.numel(), emb_dim), dtype=torch.float32), + source="dynamicemb_evicted", + ) + return num_rows + def _collect_table_shard_infos(self) -> Dict[str, _TableShardInfo]: table_shard_infos: Dict[str, _TableShardInfo] = {} for module_fqn, module in self._tracker.tracked_modules.items(): diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 42d8de25f..bb2e03c3e 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -78,6 +78,7 @@ ) from torchrec.types import DataType +from tzrec.protos import feature_pb2 from tzrec.protos.train_pb2 import ( DeltaEmbeddingDumpConfig, DeltaEmbeddingQuantType, @@ -86,6 +87,7 @@ 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, @@ -97,11 +99,16 @@ _validate_table_shard_info, validate_delta_embedding_dump_config, ) -from tzrec.utils.dynamicemb_util import has_dynamicemb +from tzrec.utils.dynamicemb_util import ( + 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 +125,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: @@ -1640,6 +1655,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 +1687,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={}), @@ -2174,6 +2191,310 @@ def test_lookup_dynamic_embeddings_flushes_module_once_per_dump(self): # once per dump rather than once per table. dynamic_module.flush.assert_called_once_with() + _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): + return SimpleNamespace( + table_names=[self._DYN_TABLE_FQN.rsplit(".", maxsplit=1)[-1]], + _dynamicemb_options=[SimpleNamespace(dim=2)], + flush=mock.MagicMock(), + pop_evicted_keys=mock.MagicMock( + return_value={ + self._DYN_TABLE_FQN.rsplit(".", maxsplit=1)[-1]: torch.tensor( + evicted_keys, dtype=torch.int64 + ) + } + ), + ) + + def test_append_dynamic_evicted_rows_publishes_zero_tombstones(self): + dumper = self._eviction_dumper() + dynamic_module = self._eviction_module([7, 9]) + table_chunks = [] + + num_rows = dumper._append_dynamic_evicted_rows( + table_chunks, + global_step=5, + dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, + published_key_ids={}, + flushed_module_ids=set(), + ) + + 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"]) + + def test_append_dynamic_evicted_rows_subtracts_republished_keys(self): + # 7 and 9 were evicted but reinserted and published as real rows this + # dump; a tombstone would race the fresh row in the MERGE upload. + dumper = self._eviction_dumper() + dynamic_module = self._eviction_module([7, 8, 9]) + table_chunks = [] + + num_rows = dumper._append_dynamic_evicted_rows( + table_chunks, + global_step=5, + dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, + published_key_ids={ + self._DYN_TABLE_FQN: torch.tensor([7, 9], dtype=torch.int64) + }, + flushed_module_ids=set(), + ) + + self.assertEqual(num_rows, 1) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), [8]) + + def test_append_dynamic_evicted_rows_tombstones_tracked_missing_ids(self): + # 102 was tracked but find() missed it (founds=False), so it is absent + # from the published ids; its eviction still gets a tombstone. + dumper = self._eviction_dumper() + dynamic_module = self._eviction_module([102]) + table_chunks = [] + + num_rows = dumper._append_dynamic_evicted_rows( + table_chunks, + global_step=5, + dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, + published_key_ids={ + self._DYN_TABLE_FQN: torch.tensor([101, 103], dtype=torch.int64) + }, + flushed_module_ids=set(), + ) + + self.assertEqual(num_rows, 1) + table = pa.concat_tables(table_chunks) + self.assertEqual(table["key_id"].to_pylist(), [102]) + + def test_append_dynamic_evicted_rows_skips_discard_and_empty_tables(self): + dumper = self._eviction_dumper() + discard_module = SimpleNamespace( + table_names=["discard_table"], + _dynamicemb_options=[SimpleNamespace(dim=2)], + flush=mock.MagicMock(), + pop_evicted_keys=mock.MagicMock(return_value={}), + ) + empty_module = SimpleNamespace( + table_names=["empty_table"], + _dynamicemb_options=[SimpleNamespace(dim=2)], + flush=mock.MagicMock(), + pop_evicted_keys=mock.MagicMock( + return_value={"empty_table": torch.tensor([], dtype=torch.int64)} + ), + ) + old_module = SimpleNamespace( # dynamicemb without pop_evicted_keys + table_names=["old_table"], + _dynamicemb_options=[SimpleNamespace(dim=2)], + flush=mock.MagicMock(), + ) + dynamic_modules = { + "model.ec.embeddings.discard_table": discard_module, + "model.ec.embeddings.empty_table": empty_module, + "model.ec.embeddings.old_table": old_module, + } + table_chunks = [] + + with mock.patch("tzrec.utils.delta_embedding_dump.logger") as log: + num_rows = dumper._append_dynamic_evicted_rows( + table_chunks, + global_step=5, + dynamic_modules=dynamic_modules, + published_key_ids={}, + flushed_module_ids=set(), + ) + + self.assertEqual(num_rows, 0) + self.assertEqual(table_chunks, []) + # 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: + dumper._append_dynamic_evicted_rows( + [], + global_step=6, + dynamic_modules=dynamic_modules, + published_key_ids={}, + flushed_module_ids=set(), + ) + log.warning.assert_not_called() + + def test_append_dynamic_evicted_rows_flushes_module_once(self): + dumper = self._eviction_dumper() + dynamic_module = SimpleNamespace( + table_names=["dyn_a", "dyn_b"], + _dynamicemb_options=[SimpleNamespace(dim=2), SimpleNamespace(dim=2)], + flush=mock.MagicMock(), + pop_evicted_keys=mock.MagicMock( + side_effect=lambda names: { + name: torch.tensor([5], dtype=torch.int64) for name in names + } + ), + ) + flushed_module_ids = set() + table_chunks = [] + for table_name in ("dyn_a", "dyn_b"): + dumper._append_dynamic_evicted_rows( + table_chunks, + global_step=5, + dynamic_modules={f"model.ec.embeddings.{table_name}": dynamic_module}, + published_key_ids={}, + flushed_module_ids=flushed_module_ids, + ) + # Both tables share the module; flush() flushes all tables, so the + # tombstone pass flushes it at most once per dump. + dynamic_module.flush.assert_called_once_with() + + # A module already flushed by the tracker pass is not flushed again. + dynamic_module.flush.reset_mock() + dumper._append_dynamic_evicted_rows( + [], + global_step=6, + dynamic_modules={"model.ec.embeddings.dyn_a": dynamic_module}, + published_key_ids={}, + flushed_module_ids={id(dynamic_module)}, + ) + dynamic_module.flush.assert_not_called() + + def test_append_dynamic_evicted_rows_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 + ) + dynamic_module = self._eviction_module([11]) + table_chunks = [] + + num_rows = dumper._append_dynamic_evicted_rows( + table_chunks, + global_step=5, + dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, + published_key_ids={}, + flushed_module_ids=set(), + ) + + 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_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() + + +@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, + ) + class DeltaEmbeddingDumpShardedIntegrationTest(MultiProcessTestBase): def __init__(self, methodName="runTest") -> None: @@ -2332,6 +2653,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 +2668,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 +2702,13 @@ 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 tiny initial table capacity evicts keys during training; they + # must reach the shards as tombstones or the pop_evicted_keys drain + # was never exercised. + 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..b99c59069 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,50 @@ _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 _dynamicemb_effective_cache_ratio( cache_load_factor: Optional[float], @@ -294,6 +338,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, From c7251304c182173a5e61fbc4a99b031fc27a41f5 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 21 Aug 2026 22:29:15 +0800 Subject: [PATCH 02/16] [bugfix] make dump_evicted_tombstones opt-in RETAIN_KEY keeps every evicted key in per-rank GPU memory until the dumper drains it at the next dump, so the buffer scales with eviction volume over the whole dump interval and has no upstream size cap. Defaulting the switch to true silently imposed that HBM growth on every existing delta-dump job, which can OOM on high-churn tables or long dump intervals. Flip the default to false so retention is only armed when a user opts into tombstones, and enable the flag explicitly in the multi-GPU dynamicemb integration test that covers the drain. Co-Authored-By: Claude Fable 5 --- tzrec/protos/train.proto | 9 ++++++--- tzrec/utils/delta_embedding_dump_test.py | 2 ++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index d1aa306bd..a5c230b66 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -112,9 +112,12 @@ message DeltaEmbeddingDumpConfig { // 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. Requires a dynamicemb - // build with EvictedItemMode; older builds degrade to a warning. - optional bool dump_evicted_tombstones = 7 [default = true]; + // (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 = false]; } message TrainConfig { diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index bb2e03c3e..098aea876 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -2638,6 +2638,8 @@ def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): dump_cfg.dump_interval_steps = 1 dump_cfg.output_dir = dump_dir dump_cfg.file_prefix = "delta_embedding" + # Tombstones are opt-in; enable them so the eviction drain is covered. + dump_cfg.dump_evicted_tombstones = True new_config_path = os.path.join(self.test_dir, "new_pipeline.config") config_util.save_message(pipeline_config, new_config_path) From 4e1f8484a5dcf940c1b439460c644491731dc2a5 Mon Sep 17 00:00:00 2001 From: gecheng Date: Fri, 21 Aug 2026 23:05:33 +0800 Subject: [PATCH 03/16] [bugfix] make evicted-key tombstone integration test deterministic The tombstone assertion could never pass on any dynamicemb build. init_capacity_per_rank only seeds rehash growth, so it never evicts, and the fixture's max_capacity values (1M/100k) exceeded the mock unique-id counts, which are drawn from [0, num_embeddings) with num_embeddings == max_capacity -- shrinking the capacity alone shrinks the id space in lockstep, so the tables could never overflow. Prepare the mock data against the original large id space first, then shrink max_capacity to 1024 so each rank faces far more candidate keys than its 512-row capacity and evicts within the first steps. Also gate the assertion on _dynamicemb_pop_evicted_keys_supported (until now dead code) because production degrades to a warning on old dynamicemb builds without pop_evicted_keys, and fix the comment that blamed the initial capacity for the evictions. Co-Authored-By: Claude Fable 5 --- tzrec/utils/delta_embedding_dump_test.py | 35 +++++++++++++++++------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 098aea876..135dfeffd 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -92,6 +92,7 @@ _DELTA_DUMP_SCHEMA, DeltaEmbeddingDumper, ModelDeltaTracker, + _dynamicemb_pop_evicted_keys_supported, _local_table_weight, _table_shard_info_from_config, _TableShardInfo, @@ -2616,11 +2617,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: @@ -2632,6 +2644,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 @@ -2704,13 +2717,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 tiny initial table capacity evicts keys during training; they + # The shrunken max_capacity forces evictions during training; they # must reach the shards as tombstones or the pop_evicted_keys drain - # was never exercised. - self.assertTrue( - dumped_tombstone_rows, - "no evicted-key tombstones dumped; pop_evicted_keys path not exercised", - ) + # 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): From e1f66b5bf6462097b20e079a18660aaa20cd8c6e Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 23 Aug 2026 14:17:22 +0800 Subject: [PATCH 04/16] [chore] revert the .pyre_configuration trailing-comma fix Repairing the trailing comma made pyre start for the first time since 5a198ad (#154) and surface ~1200 type errors that accumulated on master while the invalid config silently aborted every check. Reviving the lane is a separate cleanup; keep this PR scoped to tombstones. Co-Authored-By: Claude Fable 5 --- .pyre_configuration | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pyre_configuration b/.pyre_configuration index bf7649a6d..f2ca8ba3a 100644 --- a/.pyre_configuration +++ b/.pyre_configuration @@ -7,7 +7,7 @@ "tzrec/utils/load_class.py", "tzrec/utils/filesystem_util.py", "tzrec/tools/convert_easyrec_config_to_tzrec_config.py", - "tzrec/ops/triton/*.py" + "tzrec/ops/triton/*.py", ], "site_package_search_strategy": "all", "source_directories": [ From 60a184425ac6241dfd4091082edfce562abeb382 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 23 Aug 2026 17:33:50 +0800 Subject: [PATCH 05/16] [feat] publish evicted-key tombstones in bounded chunks by default Deletion semantics are required for delta-dump correctness, so the switch defaults back to on. The eviction drain used to materialize the whole evicted-count-by-dim zero matrix at once and retain it in the Arrow chunks and upload queue; it now emits tombstones in 65536-row chunks whose embeddings alias one chunk-sized zero buffer (Arrow wraps the row-slice views zero-copy), so host memory stays capped regardless of eviction volume. --- tzrec/protos/train.proto | 8 ++++-- tzrec/utils/delta_embedding_dump.py | 36 ++++++++++++++++-------- tzrec/utils/delta_embedding_dump_test.py | 33 ++++++++++++++++++++-- 3 files changed, 61 insertions(+), 16 deletions(-) diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index a5c230b66..7219e6fd0 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -115,9 +115,11 @@ message DeltaEmbeddingDumpConfig { // (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 = false]; + // Tombstone rows are emitted in bounded chunks that share one zero + // buffer, so host memory for them stays capped regardless of eviction + // volume. 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 da3b50eb3..478d67f02 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -69,6 +69,9 @@ ) _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 _ShardedEmbeddingModule = Union[ ShardedEmbeddingCollection, ShardedEmbeddingBagCollection ] @@ -1233,6 +1236,10 @@ def _append_dynamic_evicted_rows( reinserted key keeps its fresh row instead of racing a tombstone in the MERGE upload. Each rank drains only its local shard (row-wise sharding keeps shards disjoint), so no cross-rank collective is needed. + Tombstone rows are appended in chunks of at most + ``_TOMBSTONE_CHUNK_ROWS`` rows whose embeddings all alias one + chunk-sized zero buffer, so host memory for the zeros stays bounded + regardless of eviction volume. Args: table_chunks: List to append the per-table parquet chunks to. @@ -1280,17 +1287,24 @@ def _append_dynamic_evicted_rows( table_id = dynamic_module.table_names.index(table_name) # pyre-ignore [29] emb_dim = dynamic_module._dynamicemb_options[table_id].dim - num_rows += self._append_table_chunk( - table_chunks, - global_step=global_step, - feature_name=_feature_name( - self._tracker.fqn_to_feature_names.get(fqn, []) - ), - table_fqn=fqn, - key_ids=evicted, - embeddings=torch.zeros((evicted.numel(), emb_dim), dtype=torch.float32), - source="dynamicemb_evicted", - ) + # 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. + chunk_rows = min(_TOMBSTONE_CHUNK_ROWS, evicted.numel()) + zero_rows = torch.zeros((chunk_rows, emb_dim), dtype=torch.float32) + for start in range(0, evicted.numel(), chunk_rows): + key_chunk = evicted[start : start + chunk_rows] + num_rows += self._append_table_chunk( + table_chunks, + global_step=global_step, + feature_name=_feature_name( + self._tracker.fqn_to_feature_names.get(fqn, []) + ), + table_fqn=fqn, + key_ids=key_chunk, + embeddings=zero_rows[: key_chunk.numel()], + source="dynamicemb_evicted", + ) return num_rows def _collect_table_shard_infos(self) -> Dict[str, _TableShardInfo]: diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 135dfeffd..5f904fcc3 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -2417,6 +2417,37 @@ def test_append_dynamic_evicted_rows_quant_bytes_pin_tombstone_format(self): # Bitwise +0.0, not -0.0: NvEmbeddings dequantizes to fp16 0x0000. self.assertFalse(np.signbit(decoded).any()) + def test_append_dynamic_evicted_rows_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() + evicted_keys = list(range(10)) + dynamic_module = self._eviction_module(evicted_keys) + table_chunks = [] + + with mock.patch("tzrec.utils.delta_embedding_dump._TOMBSTONE_CHUNK_ROWS", 3): + num_rows = dumper._append_dynamic_evicted_rows( + table_chunks, + global_step=5, + dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, + published_key_ids={}, + flushed_module_ids=set(), + ) + + 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(), evicted_keys) + 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_clear_discards_retained_evicted_keys(self): dumper = object.__new__(DeltaEmbeddingDumper) dumper._dump_evicted_tombstones = True @@ -2651,8 +2682,6 @@ def test_dynamicemb_multi_gpu_delta_dump_writes_uniform_shards(self): dump_cfg.dump_interval_steps = 1 dump_cfg.output_dir = dump_dir dump_cfg.file_prefix = "delta_embedding" - # Tombstones are opt-in; enable them so the eviction drain is covered. - dump_cfg.dump_evicted_tombstones = True new_config_path = os.path.join(self.test_dir, "new_pipeline.config") config_util.save_message(pipeline_config, new_config_path) From f9bbbdabba55b6e897d37fccb37454da8b058c58 Mon Sep 17 00:00:00 2001 From: gecheng Date: Sun, 23 Aug 2026 17:34:01 +0800 Subject: [PATCH 06/16] [bugfix] disarm evicted-key retention after training set_auto_retain_evicted_keys() flips process-global state at model-build time and was never restored, so any dynamicemb table built later in the same process (e.g. export/predict) inherited RETAIN_KEY with no dumper to drain the eviction buffer and grew GPU memory without bound. Wrap the training body in try/finally so the flag resets when training ends, whether it succeeds or fails. --- tzrec/main.py | 67 ++++++++++++++++++++++++++++----------------------- 1 file changed, 37 insertions(+), 30 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 9cd95f323..963da092d 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -907,36 +907,43 @@ def train_and_evaluate( with open(os.path.join(pipeline_config.model_dir, "version"), "w") as f: f.write(tzrec_version + "\n") - if delta_embedding_dumper is not None: - delta_embedding_dumper.start() - # when slice batch by sample cost, data on all workers may not be balanced - check_all_workers_data_status = data_config.HasField("batch_cost_size") - _train_and_evaluate( - model, - optimizer, - train_dataloader, - eval_dataloader, - [sparse_lr, dense_lr, *part_lrs], - pipeline_config.model_dir, - train_config=train_config, - eval_config=pipeline_config.eval_config, - ckpt_manager=ckpt_manager, - skip_steps=skip_steps, - ckpt_path=ckpt_path, - check_all_workers_data_status=check_all_workers_data_status, - ignore_restore_optimizer=ignore_restore_optimizer, - dataloader_state=dataloader_state, - delta_embedding_dumper=delta_embedding_dumper, - pipeline_config_path=os.path.join(pipeline_config.model_dir, "pipeline.config"), - dense_ema=dense_ema, - export_config=pipeline_config.export_config, - ) - # Drain background uploads only after training succeeds. A training failure - # terminates the whole job (torchrun tears down every rank) and pending - # in-memory deltas are intentionally abandoned: the restarted run re-dumps - # from the latest checkpoint, so there is nothing to roll back or undo. - if delta_embedding_dumper is not None: - delta_embedding_dumper.close() + try: + if delta_embedding_dumper is not None: + delta_embedding_dumper.start() + # when slice batch by sample cost, data on all workers may not be balanced + check_all_workers_data_status = data_config.HasField("batch_cost_size") + _train_and_evaluate( + model, + optimizer, + train_dataloader, + eval_dataloader, + [sparse_lr, dense_lr, *part_lrs], + pipeline_config.model_dir, + train_config=train_config, + eval_config=pipeline_config.eval_config, + ckpt_manager=ckpt_manager, + skip_steps=skip_steps, + ckpt_path=ckpt_path, + check_all_workers_data_status=check_all_workers_data_status, + ignore_restore_optimizer=ignore_restore_optimizer, + dataloader_state=dataloader_state, + delta_embedding_dumper=delta_embedding_dumper, + pipeline_config_path=os.path.join( + pipeline_config.model_dir, "pipeline.config" + ), + dense_ema=dense_ema, + export_config=pipeline_config.export_config, + ) + # Drain background uploads only after training succeeds. A training failure + # terminates the whole job (torchrun tears down every rank) and pending + # in-memory deltas are intentionally abandoned: the restarted run re-dumps + # from the latest checkpoint, so there is nothing to roll back or undo. + if delta_embedding_dumper is not None: + delta_embedding_dumper.close() + finally: + # Evicted-key retention is process-global; disarm so anything running + # after training in this process never sees it left armed. + dynamicemb_util.set_auto_retain_evicted_keys(False) if is_local_rank_zero: logger.info("Train and Evaluate Finished.") From be850429177d26fc942a0e8fed0ceb7a09fa0657 Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 10:15:54 +0800 Subject: [PATCH 07/16] [chore] bump version to 1.3.17 --- tzrec/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tzrec/version.py b/tzrec/version.py index c770e6356..2b03e2c71 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.16" +__version__ = "1.3.17" From 7f0f8e97bc6856ff8047fe82d432ae5ad0154452 Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 10:21:51 +0800 Subject: [PATCH 08/16] [chore] bump version to 1.3.18 --- tzrec/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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" From 11041ac838d01ced5c72bdca0421865691166b6b Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 14:38:34 +0800 Subject: [PATCH 09/16] [refactor] drop ineffective post-training evicted-key retention reset train_and_evaluate runs in a one-shot subprocess that exits after training, so the process-global flag reset could never affect anything the process had already built; tests reset the flag in their own cleanup. Also trims over-detailed comments around the delta dump tombstone code. --- tzrec/main.py | 67 +++++++++++++---------------- tzrec/protos/train.proto | 6 +-- tzrec/utils/delta_embedding_dump.py | 29 +++++-------- 3 files changed, 43 insertions(+), 59 deletions(-) diff --git a/tzrec/main.py b/tzrec/main.py index 963da092d..9cd95f323 100644 --- a/tzrec/main.py +++ b/tzrec/main.py @@ -907,43 +907,36 @@ def train_and_evaluate( with open(os.path.join(pipeline_config.model_dir, "version"), "w") as f: f.write(tzrec_version + "\n") - try: - if delta_embedding_dumper is not None: - delta_embedding_dumper.start() - # when slice batch by sample cost, data on all workers may not be balanced - check_all_workers_data_status = data_config.HasField("batch_cost_size") - _train_and_evaluate( - model, - optimizer, - train_dataloader, - eval_dataloader, - [sparse_lr, dense_lr, *part_lrs], - pipeline_config.model_dir, - train_config=train_config, - eval_config=pipeline_config.eval_config, - ckpt_manager=ckpt_manager, - skip_steps=skip_steps, - ckpt_path=ckpt_path, - check_all_workers_data_status=check_all_workers_data_status, - ignore_restore_optimizer=ignore_restore_optimizer, - dataloader_state=dataloader_state, - delta_embedding_dumper=delta_embedding_dumper, - pipeline_config_path=os.path.join( - pipeline_config.model_dir, "pipeline.config" - ), - dense_ema=dense_ema, - export_config=pipeline_config.export_config, - ) - # Drain background uploads only after training succeeds. A training failure - # terminates the whole job (torchrun tears down every rank) and pending - # in-memory deltas are intentionally abandoned: the restarted run re-dumps - # from the latest checkpoint, so there is nothing to roll back or undo. - if delta_embedding_dumper is not None: - delta_embedding_dumper.close() - finally: - # Evicted-key retention is process-global; disarm so anything running - # after training in this process never sees it left armed. - dynamicemb_util.set_auto_retain_evicted_keys(False) + if delta_embedding_dumper is not None: + delta_embedding_dumper.start() + # when slice batch by sample cost, data on all workers may not be balanced + check_all_workers_data_status = data_config.HasField("batch_cost_size") + _train_and_evaluate( + model, + optimizer, + train_dataloader, + eval_dataloader, + [sparse_lr, dense_lr, *part_lrs], + pipeline_config.model_dir, + train_config=train_config, + eval_config=pipeline_config.eval_config, + ckpt_manager=ckpt_manager, + skip_steps=skip_steps, + ckpt_path=ckpt_path, + check_all_workers_data_status=check_all_workers_data_status, + ignore_restore_optimizer=ignore_restore_optimizer, + dataloader_state=dataloader_state, + delta_embedding_dumper=delta_embedding_dumper, + pipeline_config_path=os.path.join(pipeline_config.model_dir, "pipeline.config"), + dense_ema=dense_ema, + export_config=pipeline_config.export_config, + ) + # Drain background uploads only after training succeeds. A training failure + # terminates the whole job (torchrun tears down every rank) and pending + # in-memory deltas are intentionally abandoned: the restarted run re-dumps + # from the latest checkpoint, so there is nothing to roll back or undo. + if delta_embedding_dumper is not None: + delta_embedding_dumper.close() if is_local_rank_zero: logger.info("Train and Evaluate Finished.") diff --git a/tzrec/protos/train.proto b/tzrec/protos/train.proto index 7219e6fd0..5cbdfe6b0 100644 --- a/tzrec/protos/train.proto +++ b/tzrec/protos/train.proto @@ -115,10 +115,8 @@ message DeltaEmbeddingDumpConfig { // (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. - // Tombstone rows are emitted in bounded chunks that share one zero - // buffer, so host memory for them stays capped regardless of eviction - // volume. Requires a dynamicemb build with EvictedItemMode; older builds - // degrade to a warning. + // Requires a dynamicemb build with EvictedItemMode; older builds degrade + // to a warning. optional bool dump_evicted_tombstones = 7 [default = true]; } diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 478d67f02..56d00b7d5 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -782,13 +782,10 @@ 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. Retained evicted keys are - popped and discarded for the same reason: they predate the restored - baseline and must not be tombstoned by the next dump (the retain - buffer is process-local GPU memory and never survives a restart). + 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: @@ -1229,17 +1226,13 @@ def _append_dynamic_evicted_rows( ) -> int: """Append zero-row tombstones for dynamicemb keys evicted since the last dump. - The authoritative eviction source is ``pop_evicted_keys`` (read-and-clear - semantics), not the tracker lookup's ``founds`` mask, which cannot tell - an evicted key from one never admitted. Keys already published as real - rows by this dump's tracker pass are subtracted, so an evicted-then- - reinserted key keeps its fresh row instead of racing a tombstone in the - MERGE upload. Each rank drains only its local shard (row-wise sharding - keeps shards disjoint), so no cross-rank collective is needed. - Tombstone rows are appended in chunks of at most - ``_TOMBSTONE_CHUNK_ROWS`` rows whose embeddings all alias one - chunk-sized zero buffer, so host memory for the zeros stays bounded - regardless of eviction volume. + Evictions come from ``pop_evicted_keys`` (read-and-clear semantics), + not the tracker lookup's ``founds`` mask, which cannot tell an evicted + key from one never admitted. Keys already published as real rows this + dump are subtracted, so an evicted-then-reinserted key keeps its fresh + row instead of racing a tombstone in the MERGE upload. 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. From f198e522ad7e81fc6d5cd1fa496ac761d432c609 Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 14:48:17 +0800 Subject: [PATCH 10/16] [perf] materialize dynamic key ids on host once per dump DynamicEmb key ids were copied D2H twice: once for the tombstone pass's published-key map and again in _append_table_chunk. Materialize the host copy once and reuse it for both. The torch.cat merge branch was dead code since get_unique() yields each FQN once per dump and the map starts empty. --- tzrec/utils/delta_embedding_dump.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 56d00b7d5..91362c840 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -1052,12 +1052,10 @@ def _append_model_delta_rows( flushed_module_ids=flushed_module_ids, ) if published_dynamic_key_ids is not None and fqn in dynamic_modules: - published_ids = key_ids.detach().cpu().to(torch.int64) - if fqn in published_dynamic_key_ids: - published_ids = torch.cat( - [published_dynamic_key_ids[fqn], published_ids] - ) - published_dynamic_key_ids[fqn] = published_ids + # Materialize once on host; _append_table_chunk reuses it and + # its D2H copy becomes a no-op. + key_ids = key_ids.detach().cpu().to(torch.int64) + published_dynamic_key_ids[fqn] = key_ids feature_name = _feature_name( self._tracker.fqn_to_feature_names.get(fqn, []) ) From f535be1c86324809bc43cfbddd7e846de51af2ff Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 16:01:32 +0800 Subject: [PATCH 11/16] [perf] resolve dynamic delta rows and tombstones in one lookup The dump previously ran two dynamicemb passes: a tracker lookup that published found rows, then a pop_evicted_keys drain whose tombstones subtracted the keys the first pass had published. Merging the tracker ids with the popped evicted ids into a single post-flush find resolves each key exactly once: found keys publish their current embedding (evict-then-reinsert), missing keys from the evicted set publish a zero tombstone (insert-then-evict), and only never-admitted keys stay skipped with a warning. The second GPU lookup, the published-key subtraction, and the spurious missing-id warning for insert-then-evict keys are gone. --- tzrec/utils/delta_embedding_dump.py | 314 +++++++++++++---------- tzrec/utils/delta_embedding_dump_test.py | 298 ++++++++++----------- 2 files changed, 328 insertions(+), 284 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index 91362c840..b21a8bb9b 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -943,28 +943,13 @@ def dump(self, global_step: int) -> Optional[str]: table_weights = self._collect_table_weights() dynamic_modules = self._collect_dynamic_modules() table_chunks: List[pa.Table] = [] - flushed_module_ids: Set[int] = set() - published_dynamic_key_ids: Dict[str, torch.Tensor] = {} num_rows = self._append_model_delta_rows( table_chunks, global_step=global_step, table_weights=table_weights, dynamic_modules=dynamic_modules, - flushed_module_ids=flushed_module_ids, - published_dynamic_key_ids=( - published_dynamic_key_ids if self._dump_evicted_tombstones else None - ), + dump_evicted_tombstones=self._dump_evicted_tombstones, ) - if self._dump_evicted_tombstones: - # Must follow the tracker pass: its flush() can evict backing-store - # rows, and those evictions have to be captured by the same pop. - num_rows += self._append_dynamic_evicted_rows( - table_chunks, - global_step=global_step, - dynamic_modules=dynamic_modules, - published_key_ids=published_dynamic_key_ids, - flushed_module_ids=flushed_module_ids, - ) output_path: Optional[str] = None if write_local and (num_rows > 0 or self._world_size > 1): # Multi-rank shard sets stay complete even for an empty rank so @@ -1014,9 +999,19 @@ def _append_model_delta_rows( table_weights: Dict[str, _TableWeight], dynamic_modules: Dict[str, nn.Module], flushed_module_ids: Optional[Set[int]] = None, - published_dynamic_key_ids: Optional[Dict[str, torch.Tensor]] = None, + dump_evicted_tombstones: bool = False, ) -> int: - """Append real rows for the tracker's touched ids. + """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. @@ -1025,44 +1020,49 @@ def _append_model_delta_rows( 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. - published_dynamic_key_ids: When given, records the dynamic keys - published as real rows here (founds-filtered, cpu int64), - keyed by table FQN, so the tombstone pass can subtract them - from the evicted keys. + 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. if flushed_module_ids is None: flushed_module_ids = set() - for fqn, unique_rows in self._tracker.get_unique(_CONSUMER).items(): - ids = unique_rows.ids - if ids.numel() == 0: + 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, - ) - if published_dynamic_key_ids is not None and fqn in dynamic_modules: - # Materialize once on host; _append_table_chunk reuses it and - # its D2H copy becomes a no-op. - key_ids = key_ids.detach().cpu().to(torch.int64) - published_dynamic_key_ids[fqn] = key_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, @@ -1075,14 +1075,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] @@ -1166,42 +1159,131 @@ 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 into a + single ``find`` so 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] - table_id = dynamic_module.table_names.index(table_name) + # 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)) device = torch.device(f"cuda:{torch.cuda.current_device()}") - ids = ids.to(device=device, dtype=torch.int64) + ids = tracker_ids.to(device=device, dtype=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.to(device=device, dtype=torch.int64) + ids = torch.cat([ids, evicted_ids]) + ids = ids.unique(sorted=True) + if ids.numel() == 0: + return 0 + table_id = dynamic_module.table_names.index(table_name) 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()): + missing = ~founds + if evicted_ids is not None: + evicted_mask = torch.isin(ids, evicted_ids) + tombstone_ids = ids[missing & evicted_mask] + never_admitted = missing & ~evicted_mask + else: + tombstone_ids = ids[torch.zeros_like(founds)] + never_admitted = missing + if bool(never_admitted.any().item()): logger.warning( "Skip %s missing dynamic embedding ids for table %s.", - int((~founds).sum().item()), + int(never_admitted.sum().item()), fqn, ) - return values[founds, :emb_dim].detach(), ids[founds] + feature_name = _feature_name(self._tracker.fqn_to_feature_names.get(fqn, [])) + num_rows = 0 + if bool(founds.any().item()): + num_rows += self._append_table_chunk( + table_chunks, + global_step=global_step, + feature_name=feature_name, + table_fqn=fqn, + key_ids=ids[founds], + 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, + ) + 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.""" @@ -1214,88 +1296,48 @@ def _warn_no_retain_table_once(self, table_fqn: str, reason: str) -> None: "keys will never be reclaimed." ) - def _append_dynamic_evicted_rows( + def _append_tombstone_chunks( self, table_chunks: List[pa.Table], global_step: int, - dynamic_modules: Dict[str, nn.Module], - published_key_ids: Dict[str, torch.Tensor], - flushed_module_ids: Set[int], + fqn: str, + tombstone_ids: torch.Tensor, + emb_dim: int, ) -> int: - """Append zero-row tombstones for dynamicemb keys evicted since the last dump. + """Append zero-row tombstones for evicted keys in bounded chunks. - Evictions come from ``pop_evicted_keys`` (read-and-clear semantics), - not the tracker lookup's ``founds`` mask, which cannot tell an evicted - key from one never admitted. Keys already published as real rows this - dump are subtracted, so an evicted-then-reinserted key keeps its fresh - row instead of racing a tombstone in the MERGE upload. Row-wise - sharding keeps rank-local shards disjoint, so no cross-rank collective - is needed. + 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. Args: table_chunks: List to append the per-table parquet chunks to. global_step: Current training step. - dynamic_modules: Dynamic embedding modules keyed by table FQN; - tables without tracker rows are still visited, since their - eviction buffer must be drained too. - published_key_ids: Dynamic keys published as real rows this dump, - keyed by table FQN. - flushed_module_ids: Modules already flushed this dump, shared with - the tracker pass so each module flushes at most once. + 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.zeros((chunk_rows, emb_dim), dtype=torch.float32) num_rows = 0 - for fqn, dynamic_module in dynamic_modules.items(): - table_name = fqn.rsplit(".", maxsplit=1)[-1] - 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" - ) - continue - # flush() can evict backing-store rows; flush before the pop so - # those evictions join this drain (no-op for non-caching modules - # the tracker pass already flushed). - if id(dynamic_module) not in flushed_module_ids: - dynamic_module.flush() - flushed_module_ids.add(id(dynamic_module)) - 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)" - ) - continue - evicted = evicted.detach().cpu().to(torch.int64) - if evicted.numel() == 0: - continue - published = published_key_ids.get(fqn) - if published is not None and published.numel() > 0: - evicted = evicted[~torch.isin(evicted, published)] - if evicted.numel() == 0: - continue - table_id = dynamic_module.table_names.index(table_name) - # pyre-ignore [29] - emb_dim = dynamic_module._dynamicemb_options[table_id].dim - # 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. - chunk_rows = min(_TOMBSTONE_CHUNK_ROWS, evicted.numel()) - zero_rows = torch.zeros((chunk_rows, emb_dim), dtype=torch.float32) - for start in range(0, evicted.numel(), chunk_rows): - key_chunk = evicted[start : start + chunk_rows] - num_rows += self._append_table_chunk( - table_chunks, - global_step=global_step, - feature_name=_feature_name( - self._tracker.fqn_to_feature_names.get(fqn, []) - ), - table_fqn=fqn, - key_ids=key_chunk, - embeddings=zero_rows[: key_chunk.numel()], - source="dynamicemb_evicted", - ) + 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", + ) return num_rows def _collect_table_shard_infos(self) -> Dict[str, _TableShardInfo]: diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 5f904fcc3..e51fc9d1b 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -189,24 +189,23 @@ def forward(self, features: KeyedJaggedTensor) -> torch.Tensor: class _FakeDynamicTables: - def __init__(self) -> None: + def __init__(self, founds=None, values=None) -> None: self.ids = None self.table_ids = None self.copy_mode = None + 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, - ) + 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 @@ -1859,7 +1858,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])) @@ -1890,7 +1888,6 @@ def test_lookup_fails_on_ids_outside_table_rows(self): ), ) }, - dynamic_modules={}, ) def test_zch_lookup_binds_held_ids_to_their_rows(self): @@ -1965,7 +1962,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): @@ -2106,7 +2102,6 @@ def test_lookup_handles_empty_ids(self): ), ) }, - dynamic_modules={}, ) self.assertEqual(embeddings.shape, (0, 2)) self.assertEqual(key_ids.shape, (0,)) @@ -2131,17 +2126,16 @@ 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): + def test_append_dynamic_rows_publishes_found_rows_only(self): from dynamicemb.types import CopyMode torch.cuda.set_device(0) - dumper = object.__new__(DeltaEmbeddingDumper) + dumper = self._eviction_dumper() fake_tables = _FakeDynamicTables() dynamic_module = SimpleNamespace( table_names=["dyn_table"], @@ -2149,28 +2143,37 @@ 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]), + 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"], @@ -2181,16 +2184,31 @@ 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, + 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" @@ -2212,31 +2230,39 @@ def _eviction_dumper(self, quant_type=None): dumper._warned_no_retain_tables = set() return dumper - def _eviction_module(self, evicted_keys): + def _eviction_module(self, evicted_keys, founds, values): + table_name = self._DYN_TABLE_FQN.rsplit(".", maxsplit=1)[-1] return SimpleNamespace( - table_names=[self._DYN_TABLE_FQN.rsplit(".", maxsplit=1)[-1]], + table_names=[table_name], _dynamicemb_options=[SimpleNamespace(dim=2)], + tables=_FakeDynamicTables(founds=founds, values=values), flush=mock.MagicMock(), pop_evicted_keys=mock.MagicMock( - return_value={ - self._DYN_TABLE_FQN.rsplit(".", maxsplit=1)[-1]: torch.tensor( - evicted_keys, dtype=torch.int64 - ) - } + return_value={table_name: torch.tensor(evicted_keys, dtype=torch.int64)} ), ) - def test_append_dynamic_evicted_rows_publishes_zero_tombstones(self): + @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_append_dynamic_rows_tombstones_evicted_keys(self): + torch.cuda.set_device(0) dumper = self._eviction_dumper() - dynamic_module = self._eviction_module([7, 9]) + dynamic_module = self._eviction_module( + [7, 9], + founds=[False, False], + values=[[0.0, 0.0], [0.0, 0.0]], + ) table_chunks = [] - num_rows = dumper._append_dynamic_evicted_rows( + num_rows = dumper._append_dynamic_rows( table_chunks, global_step=5, - dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, - published_key_ids={}, + 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, ) self.assertEqual(num_rows, 2) @@ -2251,87 +2277,105 @@ def test_append_dynamic_evicted_rows_publishes_zero_tombstones(self): ) self.assertEqual(table["feature_name"].to_pylist(), ["user_id", "user_id"]) - def test_append_dynamic_evicted_rows_subtracts_republished_keys(self): - # 7 and 9 were evicted but reinserted and published as real rows this - # dump; a tombstone would race the fresh row in the MERGE upload. + @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_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]) + 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 = [] - num_rows = dumper._append_dynamic_evicted_rows( + num_rows = dumper._append_dynamic_rows( table_chunks, global_step=5, - dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, - published_key_ids={ - self._DYN_TABLE_FQN: torch.tensor([7, 9], dtype=torch.int64) - }, + 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, 1) + self.assertEqual(num_rows, 3) table = pa.concat_tables(table_chunks) - self.assertEqual(table["key_id"].to_pylist(), [8]) + 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"], + ) - def test_append_dynamic_evicted_rows_tombstones_tracked_missing_ids(self): - # 102 was tracked but find() missed it (founds=False), so it is absent - # from the published ids; its eviction still gets a tombstone. + @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_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]) + 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 = [] - num_rows = dumper._append_dynamic_evicted_rows( + num_rows = dumper._append_dynamic_rows( table_chunks, global_step=5, - dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, - published_key_ids={ - self._DYN_TABLE_FQN: torch.tensor([101, 103], dtype=torch.int64) - }, + 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, 1) + self.assertEqual(num_rows, 3) table = pa.concat_tables(table_chunks) - self.assertEqual(table["key_id"].to_pylist(), [102]) + self.assertEqual(table["key_id"].to_pylist(), [101, 103, 102]) + self.assertEqual( + table["source"].to_pylist(), + ["model_delta_tracker", "model_delta_tracker", "dynamicemb_evicted"], + ) - def test_append_dynamic_evicted_rows_skips_discard_and_empty_tables(self): + def test_pop_evicted_key_ids_warns_on_discard_and_old_builds(self): dumper = self._eviction_dumper() discard_module = SimpleNamespace( - table_names=["discard_table"], - _dynamicemb_options=[SimpleNamespace(dim=2)], - flush=mock.MagicMock(), - pop_evicted_keys=mock.MagicMock(return_value={}), + pop_evicted_keys=mock.MagicMock(return_value={}) ) empty_module = SimpleNamespace( - table_names=["empty_table"], - _dynamicemb_options=[SimpleNamespace(dim=2)], - flush=mock.MagicMock(), pop_evicted_keys=mock.MagicMock( return_value={"empty_table": torch.tensor([], dtype=torch.int64)} - ), - ) - old_module = SimpleNamespace( # dynamicemb without pop_evicted_keys - table_names=["old_table"], - _dynamicemb_options=[SimpleNamespace(dim=2)], - flush=mock.MagicMock(), + ) ) - dynamic_modules = { + 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, } - table_chunks = [] with mock.patch("tzrec.utils.delta_embedding_dump.logger") as log: - num_rows = dumper._append_dynamic_evicted_rows( - table_chunks, - global_step=5, - dynamic_modules=dynamic_modules, - published_key_ids={}, - flushed_module_ids=set(), - ) + popped = { + fqn: dumper._pop_evicted_key_ids( + fqn, module, fqn.rsplit(".", maxsplit=1)[-1] + ) + for fqn, module in modules.items() + } - self.assertEqual(num_rows, 0) - self.assertEqual(table_chunks, []) + 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) @@ -2341,53 +2385,13 @@ def test_append_dynamic_evicted_rows_skips_discard_and_empty_tables(self): 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: - dumper._append_dynamic_evicted_rows( - [], - global_step=6, - dynamic_modules=dynamic_modules, - published_key_ids={}, - flushed_module_ids=set(), - ) + 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_dynamic_evicted_rows_flushes_module_once(self): - dumper = self._eviction_dumper() - dynamic_module = SimpleNamespace( - table_names=["dyn_a", "dyn_b"], - _dynamicemb_options=[SimpleNamespace(dim=2), SimpleNamespace(dim=2)], - flush=mock.MagicMock(), - pop_evicted_keys=mock.MagicMock( - side_effect=lambda names: { - name: torch.tensor([5], dtype=torch.int64) for name in names - } - ), - ) - flushed_module_ids = set() - table_chunks = [] - for table_name in ("dyn_a", "dyn_b"): - dumper._append_dynamic_evicted_rows( - table_chunks, - global_step=5, - dynamic_modules={f"model.ec.embeddings.{table_name}": dynamic_module}, - published_key_ids={}, - flushed_module_ids=flushed_module_ids, - ) - # Both tables share the module; flush() flushes all tables, so the - # tombstone pass flushes it at most once per dump. - dynamic_module.flush.assert_called_once_with() - - # A module already flushed by the tracker pass is not flushed again. - dynamic_module.flush.reset_mock() - dumper._append_dynamic_evicted_rows( - [], - global_step=6, - dynamic_modules={"model.ec.embeddings.dyn_a": dynamic_module}, - published_key_ids={}, - flushed_module_ids={id(dynamic_module)}, - ) - dynamic_module.flush.assert_not_called() - - def test_append_dynamic_evicted_rows_quant_bytes_pin_tombstone_format(self): + 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 @@ -2395,15 +2399,14 @@ def test_append_dynamic_evicted_rows_quant_bytes_pin_tombstone_format(self): dumper = self._eviction_dumper( DeltaEmbeddingQuantType.DELTA_EMBEDDING_QUANT_INT8 ) - dynamic_module = self._eviction_module([11]) table_chunks = [] - num_rows = dumper._append_dynamic_evicted_rows( + num_rows = dumper._append_tombstone_chunks( table_chunks, global_step=5, - dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, - published_key_ids={}, - flushed_module_ids=set(), + fqn=self._DYN_TABLE_FQN, + tombstone_ids=torch.tensor([11], dtype=torch.int64), + emb_dim=2, ) self.assertEqual(num_rows, 1) @@ -2417,28 +2420,27 @@ def test_append_dynamic_evicted_rows_quant_bytes_pin_tombstone_format(self): # Bitwise +0.0, not -0.0: NvEmbeddings dequantizes to fp16 0x0000. self.assertFalse(np.signbit(decoded).any()) - def test_append_dynamic_evicted_rows_chunks_large_evictions(self): + 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() - evicted_keys = list(range(10)) - dynamic_module = self._eviction_module(evicted_keys) + 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_dynamic_evicted_rows( + num_rows = dumper._append_tombstone_chunks( table_chunks, global_step=5, - dynamic_modules={self._DYN_TABLE_FQN: dynamic_module}, - published_key_ids={}, - flushed_module_ids=set(), + 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(), evicted_keys) + 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. From 032928084f1a9be1add3157af013d4590b232954 Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 16:34:54 +0800 Subject: [PATCH 12/16] [perf] batch the merged dynamic delta lookup The merged tracker+evicted find ran in one shot, materializing a [N, dim] GPU values buffer sized by the dump interval's eviction volume. Merge and dedup the id sets on host and look up in batches of _FIND_BATCH_ROWS, appending each batch's rows and tombstones as they resolve, so the buffer stays capped at one batch. Co-Authored-By: Claude Fable 5 --- tzrec/utils/delta_embedding_dump.py | 86 +++++++++++++----------- tzrec/utils/delta_embedding_dump_test.py | 72 ++++++++++++++++++-- 2 files changed, 115 insertions(+), 43 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index b21a8bb9b..ff26141dd 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -72,6 +72,9 @@ # 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 ] @@ -1199,10 +1202,13 @@ def _append_dynamic_rows( ) -> int: """Append one merged lookup's real rows and tombstones for a table. - The tracker ids and this dump's evicted keys are merged into a - single ``find`` so 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. + 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. @@ -1231,58 +1237,62 @@ def _append_dynamic_rows( if id(dynamic_module) not in flushed_module_ids: dynamic_module.flush() flushed_module_ids.add(id(dynamic_module)) - device = torch.device(f"cuda:{torch.cuda.current_device()}") - ids = tracker_ids.to(device=device, dtype=torch.int64) + table_id = dynamic_module.table_names.index(table_name) + # pyre-ignore [29] + emb_dim = dynamic_module._dynamicemb_options[table_id].dim + 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.to(device=device, dtype=torch.int64) + evicted_ids = evicted.cpu().to(torch.int64) ids = torch.cat([ids, evicted_ids]) ids = ids.unique(sorted=True) if ids.numel() == 0: return 0 - table_id = dynamic_module.table_names.index(table_name) - 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) - missing = ~founds - if evicted_ids is not None: - evicted_mask = torch.isin(ids, evicted_ids) - tombstone_ids = ids[missing & evicted_mask] - never_admitted = missing & ~evicted_mask - else: - tombstone_ids = ids[torch.zeros_like(founds)] - never_admitted = missing - if bool(never_admitted.any().item()): - logger.warning( - "Skip %s missing dynamic embedding ids for table %s.", - int(never_admitted.sum().item()), - fqn, - ) + evicted_mask = torch.isin(ids, evicted_ids) if evicted_ids is not None else None + device = torch.device(f"cuda:{torch.cuda.current_device()}") feature_name = _feature_name(self._tracker.fqn_to_feature_names.get(fqn, [])) num_rows = 0 - if bool(founds.any().item()): + 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) + missing = ~founds + if evicted_mask is not None: + batch_evicted = evicted_mask[start : start + _FIND_BATCH_ROWS] + tombstone_ids = gpu_ids[missing & batch_evicted] + num_never_admitted += int((missing & ~batch_evicted).sum().item()) + else: + tombstone_ids = gpu_ids.new_empty(0) + num_never_admitted += int(missing.sum().item()) num_rows += self._append_table_chunk( table_chunks, global_step=global_step, feature_name=feature_name, table_fqn=fqn, - key_ids=ids[founds], + key_ids=gpu_ids[founds], 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, - ) + 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.", + num_never_admitted, + fqn, + ) return num_rows def _warn_no_retain_table_once(self, table_fqn: str, reason: str) -> None: diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index e51fc9d1b..d7d883764 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -189,10 +189,12 @@ def forward(self, features: KeyedJaggedTensor) -> torch.Tensor: class _FakeDynamicTables: - def __init__(self, founds=None, values=None) -> 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 @@ -201,9 +203,26 @@ def __init__(self, founds=None, values=None) -> None: ) def find(self, ids, table_ids, copy_mode): - self.ids = ids.detach().clone() - self.table_ids = table_ids.detach().clone() + 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 @@ -2230,12 +2249,12 @@ def _eviction_dumper(self, quant_type=None): dumper._warned_no_retain_tables = set() return dumper - def _eviction_module(self, evicted_keys, founds, values): + 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), + 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)} @@ -2347,6 +2366,49 @@ def test_append_dynamic_rows_tombstones_tracked_then_evicted_key(self): ["model_delta_tracker", "model_delta_tracker", "dynamicemb_evicted"], ) + @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_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("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( From 3c525f601f90051d5fc25e9b2a70987643f17ad9 Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 19:57:11 +0800 Subject: [PATCH 13/16] [bugfix] fix CPU/CUDA device mismatch in evicted-key tombstone mask evicted_mask was built with torch.isin on the host-merged ids, but find() returns founds on the lookup device, so missing & batch_evicted raised a device mismatch whenever tombstoning ran on evicted keys. Build the mask on the lookup device instead. The existing _append_dynamic_rows tests covered this path but were gated on has_dynamicemb, which is absent locally and in CI, so they never ran; drop that gate for the fully-faked tests and stub dynamicemb in sys.modules so they run on any CUDA machine. Co-Authored-By: Claude Fable 5 --- tzrec/utils/delta_embedding_dump.py | 6 +- tzrec/utils/delta_embedding_dump_test.py | 162 ++++++++++++----------- 2 files changed, 92 insertions(+), 76 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index ff26141dd..a10ab4524 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -1250,8 +1250,12 @@ def _append_dynamic_rows( ids = ids.unique(sorted=True) if ids.numel() == 0: return 0 - evicted_mask = torch.isin(ids, evicted_ids) if evicted_ids is not None else None device = torch.device(f"cuda:{torch.cuda.current_device()}") + # find() returns its founds on the lookup device, so keep the mask + # there too; a CPU mask & CUDA founds would raise a device mismatch. + evicted_mask = ( + torch.isin(ids, evicted_ids).to(device) 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 diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index d7d883764..24fa14ea4 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 @@ -228,6 +229,17 @@ def find(self, ids, table_ids, copy_mode): 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}") @@ -2147,12 +2159,9 @@ def test_row_wise_lookup_requires_shard_metadata(self): }, ) - @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_append_dynamic_rows_publishes_found_rows_only(self): - from dynamicemb.types import CopyMode - torch.cuda.set_device(0) dumper = self._eviction_dumper() fake_tables = _FakeDynamicTables() @@ -2164,15 +2173,18 @@ def test_append_dynamic_rows_publishes_found_rows_only(self): ) table_chunks = [] - 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, - ) + 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() @@ -2187,7 +2199,6 @@ def test_append_dynamic_rows_publishes_found_rows_only(self): ["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_append_dynamic_rows_flushes_module_once_per_dump(self): @@ -2202,31 +2213,32 @@ def test_append_dynamic_rows_flushes_module_once_per_dump(self): ) flushed_module_ids = set() - for table_name in ("dyn_a", "dyn_b"): + 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=5, - fqn=f"model.ec.embeddings.{table_name}", + global_step=6, + fqn="model.ec.embeddings.dyn_a", dynamic_module=dynamic_module, tracker_ids=torch.tensor([101, 102, 103]), - flushed_module_ids=flushed_module_ids, + flushed_module_ids={id(dynamic_module)}, 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" @@ -2261,7 +2273,6 @@ def _eviction_module(self, evicted_keys, founds=None, values=None, rows=None): ), ) - @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_append_dynamic_rows_tombstones_evicted_keys(self): @@ -2274,15 +2285,16 @@ def test_append_dynamic_rows_tombstones_evicted_keys(self): ) table_chunks = [] - 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, - ) + 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, + ) self.assertEqual(num_rows, 2) # flush() precedes the pop so flush-induced evictions join the drain. @@ -2296,7 +2308,6 @@ def test_append_dynamic_rows_tombstones_evicted_keys(self): ) self.assertEqual(table["feature_name"].to_pylist(), ["user_id", "user_id"]) - @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_append_dynamic_rows_publishes_fresh_row_for_reinserted_key(self): @@ -2312,15 +2323,16 @@ def test_append_dynamic_rows_publishes_fresh_row_for_reinserted_key(self): ) table_chunks = [] - 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, - ) + 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) @@ -2333,7 +2345,6 @@ def test_append_dynamic_rows_publishes_fresh_row_for_reinserted_key(self): ["model_delta_tracker", "model_delta_tracker", "dynamicemb_evicted"], ) - @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_append_dynamic_rows_tombstones_tracked_then_evicted_key(self): @@ -2348,15 +2359,16 @@ def test_append_dynamic_rows_tombstones_tracked_then_evicted_key(self): ) table_chunks = [] - 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, - ) + 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) @@ -2366,7 +2378,6 @@ def test_append_dynamic_rows_tombstones_tracked_then_evicted_key(self): ["model_delta_tracker", "model_delta_tracker", "dynamicemb_evicted"], ) - @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_append_dynamic_rows_batches_merged_lookup(self): @@ -2381,17 +2392,18 @@ def test_append_dynamic_rows_batches_merged_lookup(self): ) table_chunks = [] - 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, - ) + 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( From 3b31598289a7a293dde5591817d6371a77db1c7c Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 20:12:35 +0800 Subject: [PATCH 14/16] [feat] reject non-zero dynemb eval initializers under tombstone dump Evicted-key tombstones publish constant-zero rows for evicted keys, so a dynamicemb feature whose eval_initializer_args resolves to non-zero would make the same missing key yield different values on the dump consumer and in eval lookups. Validate at plan time in build_dynamicemb_constraints, where the per-feature eval initializer is consumed, reusing the existing auto-retain switch that is armed exactly when delta dump with dump_evicted_tombstones is enabled; unset or CONSTANT 0.0 keeps the historical default and stays allowed. Co-Authored-By: Claude Fable 5 --- tzrec/utils/delta_embedding_dump_test.py | 56 ++++++++++++++++++++++++ tzrec/utils/dynamicemb_util.py | 33 ++++++++++++++ 2 files changed, 89 insertions(+) diff --git a/tzrec/utils/delta_embedding_dump_test.py b/tzrec/utils/delta_embedding_dump_test.py index 24fa14ea4..4b2523c72 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -102,6 +102,7 @@ validate_delta_embedding_dump_config, ) from tzrec.utils.dynamicemb_util import ( + _validate_eval_initializer_for_tombstones, build_dynamicemb_constraints, has_dynamicemb, set_auto_retain_evicted_keys, @@ -2568,6 +2569,52 @@ def test_clear_without_tombstones_keeps_eviction_buffer(self): 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.", @@ -2603,6 +2650,15 @@ def test_auto_retain_off_keeps_discard_mode(self): 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): def __init__(self, methodName="runTest") -> None: diff --git a/tzrec/utils/dynamicemb_util.py b/tzrec/utils/dynamicemb_util.py index b99c59069..281ba22d4 100644 --- a/tzrec/utils/dynamicemb_util.py +++ b/tzrec/utils/dynamicemb_util.py @@ -97,6 +97,38 @@ def _arm_evicted_key_retention(demb_opt_kwargs: Dict[str, Any]) -> None: 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], caching: bool, @@ -290,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 From 41b6a783074d013d286a3263211245c11625cc0b Mon Sep 17 00:00:00 2001 From: gecheng Date: Mon, 24 Aug 2026 22:07:27 +0800 Subject: [PATCH 15/16] [perf] stop round-tripping dynamic delta key ids across devices The merged lookup key ids crossed the device boundary three times per dump: the tracker ids went D2H for the host-side merge, the merged ids went H2D for the find, and the found ids went D2H again when building the parquet chunks. On top of that the full evicted mask was copied H2D once per dump, sized by the whole dump interval's key count rather than the _FIND_BATCH_ROWS batch cap. Copy only the per-batch bool found mask D2H and select both the found and tombstone key ids from the CPU batch ids, so key ids cross devices exactly once and the evicted mask never leaves the CPU. The never-admitted counts now sum on CPU as well, dropping two GPU syncs per batch; only the embedding values still move D2H. Co-Authored-By: Claude Fable 5 --- tzrec/utils/delta_embedding_dump.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index a10ab4524..e35dbf6e4 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -1251,11 +1251,10 @@ def _append_dynamic_rows( if ids.numel() == 0: return 0 device = torch.device(f"cuda:{torch.cuda.current_device()}") - # find() returns its founds on the lookup device, so keep the mask - # there too; a CPU mask & CUDA founds would raise a device mismatch. - evicted_mask = ( - torch.isin(ids, evicted_ids).to(device) if evicted_ids is not None else None - ) + # 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 @@ -1267,20 +1266,23 @@ def _append_dynamic_rows( gpu_ids, table_ids, CopyMode.EMBEDDING ) founds = founds.to(dtype=torch.bool) - missing = ~founds + # 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 = gpu_ids[missing & batch_evicted] - num_never_admitted += int((missing & ~batch_evicted).sum().item()) + tombstone_ids = batch_ids[missing_cpu & batch_evicted] + num_never_admitted += int((missing_cpu & ~batch_evicted).sum()) else: - tombstone_ids = gpu_ids.new_empty(0) - num_never_admitted += int(missing.sum().item()) + 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=gpu_ids[founds], + key_ids=batch_ids[founds_cpu], embeddings=values[founds, :emb_dim].detach(), source="model_delta_tracker", ) From 4c055288e9fe8cf3afb20e1ca708c9a0f353010b Mon Sep 17 00:00:00 2001 From: gecheng Date: Tue, 25 Aug 2026 00:00:16 +0800 Subject: [PATCH 16/16] [perf] share one quantized zero buffer across INT8 tombstone chunks _append_tombstone_chunks keeps all Arrow chunks aliasing one chunk-sized zero buffer, but only FP32 preserved that: under INT8, _append_table_chunk re-quantized every chunk and _quantize_quint8_rowwise_f16 allocates a fresh output on each call, so each chunk held its own buffer and retained tombstone memory grew as O(evictions x (embedding_dim + 4)) instead of staying one chunk-sized buffer. Quantize the chunk-sized zero buffer once in _append_tombstone_chunks and pass the row slices as pre-quantized bytes so every chunk aliases the same storage. Adds an INT8 multi-chunk buffer-sharing test alongside the existing FP32 one. Co-Authored-By: Claude Fable 5 --- tzrec/utils/delta_embedding_dump.py | 76 +++++++++++++++++------- tzrec/utils/delta_embedding_dump_test.py | 30 ++++++++++ 2 files changed, 84 insertions(+), 22 deletions(-) diff --git a/tzrec/utils/delta_embedding_dump.py b/tzrec/utils/delta_embedding_dump.py index e35dbf6e4..3b5cfef1d 100644 --- a/tzrec/utils/delta_embedding_dump.py +++ b/tzrec/utils/delta_embedding_dump.py @@ -1325,7 +1325,9 @@ def _append_tombstone_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. + 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. @@ -1341,7 +1343,30 @@ def _append_tombstone_chunks( 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.zeros((chunk_rows, emb_dim), dtype=torch.float32) + 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] @@ -1353,6 +1378,7 @@ def _append_tombstone_chunks( key_ids=key_chunk, embeddings=zero_rows[: key_chunk.numel()], source="dynamicemb_evicted", + pre_quantized=pre_quantized, ) return num_rows @@ -1443,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, " @@ -1460,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 4b2523c72..771e8aa96 100644 --- a/tzrec/utils/delta_embedding_dump_test.py +++ b/tzrec/utils/delta_embedding_dump_test.py @@ -2525,6 +2525,36 @@ def test_append_tombstone_chunks_chunks_large_evictions(self): ] 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