fix(dynamicemb): traverse nn.Module children in check_emb_collection_modules (#355)

* fix(dynamicemb): traverse nn.Module children in check_emb_collection_modules

The function only checked direct attributes, missing EmbeddingCollection
wrapped inside DistributedModelParallel or other nn.Module containers.
Now recursively walks module.children() to find the embedding submodule.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* test(dynamicemb): verify get_dynamic_emb_module traverses DMP wrappers

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* refactor(examples): drop dynamicemb_utils, use get_dynamic_emb_module

check_emb_collection_modules now traverses nn.Module.children() and
guards against cycles, so the examples-side duplicate
find_dynamicemb_modules is redundant. Route both call sites (cache stats
hook + FILL_DYNAMICEMB_TABLES fill path) through dynamicemb's public
get_dynamic_emb_module and delete commons/utils/dynamicemb_utils.py.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* fix(dynamicemb/test): rewrite check_counter_table_checkpoint for fused MultiTableKVCounter

After the admission-counter fusion refactor (#343), `_admission_counter`
became a single MultiTableKVCounter (or None) instead of a list of
per-table KVCounters. The verification helper still iterated it as a
list, so the dump/load counter check was silently no-op on any branch
whose `get_dynamic_emb_module` could traverse into DMP wrappers and
raised `TypeError: 'MultiTableKVCounter' object is not iterable` on any
branch that could. PR #355's `children()` traversal surfaced the latter.

Rewrite the check against the current API:
- Iterate logical tables via `range(len(table._table_names))`.
- Export one table_id's (keys, frequency) with
  `cnt.table_._batched_export_keys_scores([freq_name], device, table_id)`.
- Look those keys up in the peer counter via
  `cnt_y.table_.lookup(keys, table_ids, score_arg)` and assert both
  `founds.all()` and that the returned score_out matches the exported
  frequency.
- Handle the None (no-admission) case symmetrically.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: JacoCheung

corelib/dynamicemb/dynamicemb/dump_load.py

@@ -17,7 +17,7 @@
 import os
 import warnings
 from collections import deque
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Set, Tuple
 
 import torch
 import torch.distributed as dist
@@ -50,31 +50,42 @@ def find_sharded_modules(
     return sharded_modules
 
 
-def check_emb_collection_modules(module: nn.Module, ret_list: List[nn.Module]):
+def check_emb_collection_modules(
+    module: nn.Module,
+    ret_list: List[nn.Module],
+    visited: Optional[Set[int]] = None,
+):
+    if visited is None:
+        visited = set()
+
+    mid = id(module)
+    if mid in visited:
+        return
+    visited.add(mid)
+
     if isinstance(module, BatchedDynamicEmbeddingTablesV2):
         ret_list.append(module)
-        return ret_list
+        return
 
     if isinstance(module, nn.Module):
-        if hasattr(module, "_emb_module"):
-            check_emb_collection_modules(module._emb_module, ret_list)
-
-        if hasattr(module, "_emb_modules"):
-            check_emb_collection_modules(module._emb_modules, ret_list)
-
-        if hasattr(module, "_lookups"):
-            tmp_module_list = module._lookups
-            for tmp_emb_module in tmp_module_list:
-                check_emb_collection_modules(tmp_emb_module, ret_list)
-
-    if isinstance(module, nn.ModuleList):
-        for i in range(len(module)):
-            tmp_emb_module = module[i]
-
-            if isinstance(tmp_emb_module, nn.Module):
-                check_emb_collection_modules(tmp_emb_module, ret_list)
-            else:
+        # Follow TorchRec/DynamicEmb internal attributes.
+        # _lookups is stored as a plain Python list (not nn.ModuleList) so
+        # nn.Module.children() cannot discover it — traverse it explicitly.
+        for attr in ("_lookups", "_emb_modules", "_emb_module"):
+            child = getattr(module, attr, None)
+            if child is None:
                 continue
+            if isinstance(child, (list, nn.ModuleList)):
+                for item in child:
+                    check_emb_collection_modules(item, ret_list, visited)
+            else:
+                check_emb_collection_modules(child, ret_list, visited)
+
+        # Recurse into nn.Module children so we can traverse through
+        # wrapper modules (DMP / DDP / Float16Module) that are not
+        # covered by the private attributes above.
+        for child in module.children():
+            check_emb_collection_modules(child, ret_list, visited)
 
 
 def get_dynamic_emb_module(model: nn.Module) -> List[nn.Module]:

corelib/dynamicemb/test/unit_tests/test_embedding_dump_load.py

@@ -28,6 +28,7 @@
     DynamicEmbTableOptions,
     FrequencyAdmissionStrategy,
 )
+from dynamicemb.batched_dynamicemb_tables import BatchedDynamicEmbeddingTablesV2
 from dynamicemb.dump_load import (
     DynamicEmbDump,
     DynamicEmbLoad,
@@ -122,6 +123,54 @@ def assert_batched_dynamicemb_storage_class(
                 )
 
 
+def assert_get_dynamic_emb_module_finds_submodules(model) -> None:
+    """Verify get_dynamic_emb_module discovers BatchedDynamicEmbeddingTablesV2.
+
+    Tests two paths:
+      1. Via find_sharded_modules + get_dynamic_emb_module (existing usage)
+      2. Via get_dynamic_emb_module directly on the DMP model (requires
+         children() traversal through wrapper modules, the fix for #353)
+
+    Both must return the same set of modules.
+    """
+    # Path 1: existing approach - find sharded modules first, then search each
+    via_sharded = []
+    for _, _, sharded_module in find_sharded_modules(model, ""):
+        via_sharded.extend(get_dynamic_emb_module(sharded_module))
+
+    # Path 2: search directly on the DMP wrapper (requires children() traversal)
+    via_dmp = get_dynamic_emb_module(model)
+
+    assert (
+        len(via_sharded) > 0
+    ), "find_sharded_modules + get_dynamic_emb_module found no modules"
+    assert (
+        len(via_dmp) > 0
+    ), "get_dynamic_emb_module on DMP model found no modules (children() traversal broken)"
+
+    # Every module found via either path must be BatchedDynamicEmbeddingTablesV2
+    for m in via_sharded:
+        assert isinstance(
+            m, BatchedDynamicEmbeddingTablesV2
+        ), f"Expected BatchedDynamicEmbeddingTablesV2, got {type(m)}"
+    for m in via_dmp:
+        assert isinstance(
+            m, BatchedDynamicEmbeddingTablesV2
+        ), f"Expected BatchedDynamicEmbeddingTablesV2, got {type(m)}"
+
+    # Both paths must discover the exact same set of modules (by identity)
+    ids_sharded = set(id(m) for m in via_sharded)
+    ids_dmp = set(id(m) for m in via_dmp)
+    assert ids_sharded == ids_dmp, (
+        f"Module sets differ: via_sharded has {len(ids_sharded)} modules, "
+        f"via_dmp has {len(ids_dmp)} modules"
+    )
+
+    # No duplicates in either result
+    assert len(via_sharded) == len(ids_sharded), "Duplicates in via_sharded path"
+    assert len(via_dmp) == len(ids_dmp), "Duplicates in via_dmp path"
+
+
 def update_scores(
     score_strategy: str,
     expect_scores: Dict[int, int],
@@ -354,32 +403,46 @@ def create_model(
 
 
 def check_counter_table_checkpoint(x, y):
-    device = torch.cuda.current_device()
+    device = torch.device(f"cuda:{torch.cuda.current_device()}")
     tables_x = get_dynamic_emb_module(x)
     tables_y = get_dynamic_emb_module(y)
+    assert len(tables_x) == len(tables_y)
 
     for table_x, table_y in zip(tables_x, tables_y):
-        for cnt_tx, cnt_ty in zip(
-            table_x._admission_counter, table_y._admission_counter
-        ):
-            assert cnt_tx.table_.size() == cnt_ty.table_.size()
-
-            for keys, named_scores, _ in cnt_tx._batched_export_keys_scores(
-                cnt_tx.table_.score_names_, torch.device(f"cuda:{device}")
+        cnt_x = table_x._admission_counter
+        cnt_y = table_y._admission_counter
+        if cnt_x is None:
+            assert cnt_y is None
+            continue
+        assert cnt_x.table_.size() == cnt_y.table_.size()
+
+        freq_name = cnt_x.score_name_
+        for table_id in range(len(table_x._table_names)):
+            for keys, named_scores, _ in cnt_x.table_._batched_export_keys_scores(
+                [freq_name], device, table_id
             ):
                 if keys.numel() == 0:
                     continue
-                freq_name = cnt_tx.table_.score_names_[0]
                 frequencies = named_scores[freq_name]
 
+                lookup_table_ids = torch.full(
+                    (keys.numel(),), table_id, dtype=torch.int64, device=device
+                )
                 score_arg_lookup = ScoreArg(
                     name=freq_name,
                     value=torch.zeros_like(frequencies),
                     policy=ScorePolicy.CONST,
                 )
-                _, founds, _ = cnt_ty.lookup(keys, score_arg_lookup)
-
-                assert torch.equal(frequencies, score_arg_lookup)
+                score_out, founds, _ = cnt_y.table_.lookup(
+                    keys, lookup_table_ids, score_arg_lookup
+                )
+                assert founds.all(), (
+                    f"counter keys missing from loaded table_id={table_id}: "
+                    f"{keys[~founds].tolist()}"
+                )
+                assert torch.equal(
+                    frequencies, score_out
+                ), f"counter frequency mismatch for table_id={table_id}"
 
 
 @click.command()
@@ -441,6 +504,8 @@ def test_model_load_dump(
         ),
     )
 
+    assert_get_dynamic_emb_module_finds_submodules(ref_model)
+
     expect_scores_collection: Dict[str, Dict[int, int]] = {}
     kjts, feature_names, all_kjts = generate_sparse_feature(
         num_embedding_collections=num_embedding_collections,

examples/commons/utils/dynamicemb_cache_stats.py

@@ -21,7 +21,7 @@
 
 import torch
 import torch.nn as nn
-from commons.utils.dynamicemb_utils import find_dynamicemb_modules
+from dynamicemb.dump_load import get_dynamic_emb_module
 
 
 class _CacheDebugHook:
@@ -64,7 +64,7 @@ def install_cache_debug_hooks(model: nn.Module) -> int:
         Number of hooks installed.
     """
     count = 0
-    modules = find_dynamicemb_modules(model)
+    modules = get_dynamic_emb_module(model)
     rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
     if rank == 0:
         print(f"[CACHE_DEBUG] Found {len(modules)} DynamicEmb module(s)", flush=True)

examples/commons/utils/dynamicemb_utils.py

@@ -159,9 +159,9 @@ def main():
     maybe_load_ckpts(trainer_args.ckpt_load_dir, model, dense_optimizer)
 
     if os.environ.get("FILL_DYNAMICEMB_TABLES", "0") == "1":
-        from commons.utils.dynamicemb_utils import find_dynamicemb_modules
+        from dynamicemb.dump_load import get_dynamic_emb_module
 
-        for dyn_module in find_dynamicemb_modules(model_train):
+        for dyn_module in get_dynamic_emb_module(model_train):
             if hasattr(dyn_module, "fill_tables"):
                 try:
                     dyn_module.fill_tables(load_factor=0.95)