[feat] planner: pick dynamicemb HYBRID vs CACHING from topology budgets

[tiankongdeguiji]

Summary

Make tzrec's embedding sharding planner take HBM and host-memory limits from the topology and decide per dynamicemb table whether to use HYBRID (HBM + DRAM hash-partitioned) or CACHING (HBM cache + full DRAM backing) mode. Per the consensus on NVIDIA/recsys-examples#306 (closed), both placement modes are kept and the planner — not the user — picks the best one under the joint budget.

• Enumerator emits 20 ShardingOptions per dynamicemb table: { HYBRID, CACHING } × { cache_load_factor 0.1, ..., 1.0 }. The base ShardingOption is deep-copied once per variant; dynamicemb_options.caching is then flipped in-place on the already-fresh per-variant copy.
• Storage estimator is mode-aware: HYBRID DDR = (1 − cache_load_factor) · T, CACHING DDR = T (full backing store). HBM accounting and hash-table metadata stay HBM-side only, matching the existing tzrec convention.
• Perf model is fitted from an on-device dynamicemb sweep (experiments/sweep_20260513-161030/full_a10gpu1.json: 4M-row table, dim=128, adam, pow-law alpha=1.05, A10 GPU). Median fwd+bwd latency clustered into three regimes — HYBRID@1.0 = 0.80 ms (HBM-only fast path), CACHING@<1.0 = 2.63 ms (~3.3x slower), HYBRID@<1.0 = 5.44 ms (~6.8x slower). Inverting the linear bw model with torchrec defaults gives x_eff bases 0.28 (CACHING) and 0.11 (HYBRID), with x_eff = 1.0 at cache_load_factor = 1.0 reflecting the runtime kernel switch. A small +0.01·x tiebreaker preserves a strict DP ranking within each block.
• DynamicProgrammingProposer is now 2D: discretizes per-rank HBM (100 bins/dev) and per-machine DDR (50 bins/dev) and picks per-table options that minimize total perf under both budgets. HBM prune is per-device, DDR prune is per-machine (DDR is host-shared across co-located ranks). Per-axis bin count is capped at 1024 to avoid blow-up on multi-host worlds. Backtracking emits one proposal per HBM bin — the perf-best plan across all DDR bins at that level — so the proposal count is at most _DP_AXIS_BIN_CAP regardless of ddr_bins.
• Per-table planning log. _to_sharding_plan emits one INFO line per dynamicemb table on rank 0 so the picked mode is visible without grepping the runtime print. Numbers come from the storage estimator's shards[0].storage.{hbm,ddr} (true HBM footprint including hash-table metadata + counter + pipeline I/O; mode-aware DDR), with an HBM_ONLY override at cache_load_factor == 1.0 to reflect what the runtime actually does:

  [dynamicemb plan] …cat_0_emb: mode=CACHING cache_load_factor=0.20 local_hbm=1.620GiB local_dram=2.384GiB
  [dynamicemb plan] …cat_5_emb: mode=HBM_ONLY cache_load_factor=1.00 local_hbm=4.000GiB local_dram=0.000GiB
  

  Upstream _print_memory_consume is mode-blind (always dram = total − hbm) and strips metadata from HBM; this log line has the honest numbers.
• No new proto fields — budgets come straight from Topology.hbm_cap / Topology.ddr_cap minus the existing reservation.
• Version bump to 1.2.13.

Test plan

• tzrec/utils/dynamicemb_util_test.py — storage formula + empirical perf-model + direct wrapper tests + PlanLogLineTest covering HBM_ONLY-override / RANK gating / GiB unit / cache_load_factor>1.0 regression
• tzrec/utils/plan_util_test.py — variant emission + 2D DP unit tests including multi-table mixed-modes, multi-host (4 GPUs × 2 hosts) per-machine DDR prune, empty search space, and the headline DP tests
• CUDA-gated end-to-end enumerate + DP test (folded into plan_util_test.py)
• pre-commit run --files <changed> clean
• Verified the planning log against a 27-table 5-CACHING-table real model run; mode column matches the chosen caching flag, local_hbm / local_dram are the true footprints (the upstream runtime print under-reports both)
• Smoke test on a real dynamicemb config end-to-end via torchrun --nproc-per-node=2 -m tzrec.train_eval ...

🤖 Generated with Claude Code

[github-actions]

Restore should be in finally. If _orig_build_shard_perf_contexts raises, sharding_option.cache_params is permanently left as the boosted x_eff clone. Downstream consumers — the storage estimator and _to_sharding_plan at line 435 (cache_params=sharding_option.cache_params) — then see the wrong load factor on every other option that touches the same ShardingOption. Suggested:

Suggested change

	result = _orig_build_shard_perf_contexts(
	cls,
	config,
	shard_sizes,
	sharding_option,
	topology,
	constraints,
	sharder,
	*args,
	**kwargs,
	)
	sharding_option.cache_params = original_cache_params
	return result
	try:
	result = _orig_build_shard_perf_contexts(
	cls,
	config,
	shard_sizes,
	sharding_option,
	topology,
	constraints,
	sharder,
	*args,
	**kwargs,
	)
	finally:
	sharding_option.cache_params = original_cache_params
	return result

[github-actions]

Redundant double deep-copy of dynamicemb_options. The caller (line 1035) attaches dynamicemb_options to base_option before calling this function, so copy.deepcopy(base_option) on line 848 already deep-copies dynamicemb_options into opt. Line 851 then deep-copies the original again and replaces it. With 20 variants per table this doubles a non-trivial copy. Either drop line 851 and just set opt.dynamicemb_options.caching = caching_mode, or stop attaching dynamicemb_options on the caller side before this call.

[github-actions]

Docstring says "per-rank DDR" but the implementation enforces per-machine DDR (lines 349–361 sum DDR across local_world_size co-located ranks and the per-option prune at line 391 compares against max_machine_ddr). The inline comment at 350–352 has the right framing — the class docstring should match: "per-rank HBM and per-machine DDR budgets."

[github-actions]

Sharp discontinuity at x=1.0. HYBRID@x=0.99 → 0.1199; HYBRID@x=1.00 → 1.0 — an ~8× jump for a 1% ratio change. The DP will reliably prefer x=1.0 over x=0.99 by a huge perf margin, then prefer CACHING@x=0.5 (0.285) over HYBRID@x=0.9 (0.119) even on workloads where HYBRID@0.9 is plainly faster in reality. If the empirical sweep really shows a step at x=1.0 because the runtime drops the host tier, please call that out as "x=1.0 = HBM-only kernel, not the same algorithm as x=0.99" — and consider whether the enumerator should emit a discrete mode={HBM_ONLY, HYBRID, CACHING} axis rather than packing the discontinuity into the same x knob.

[github-actions]

2D DP cost scales with world size — consider clamping total bins. hbm_bins = 100 · num_devices and ddr_bins = 50 · num_devices. On a 16-GPU world the inner triple-loop is 1600 × 800 × option_count pure-Python iterations per table; with 100 dynamicemb tables (20 options each) that's a few-billion-op DP, and dp + backtrack allocate ~256M tuples. Even on smaller worlds the bin counts multiply faster than necessary — a per-device discretization isn't really needed since the budget is a single scalar per axis. Consider clamping totals (e.g. hbm_bins = min(hbm_bins_per_device · num_devices, 2_000)) or vectorizing the transition with NumPy. The defaults are fine for 2–4 GPUs but worth a guardrail before the next user hits this.

[github-actions]

Missing the headline multi-table 2D-DP scenario. All three DynamicProgrammingProposerTest cases use a single table, so the cross-table DP transition at plan_util.py:434–463 is exercised only at table_i == 0. The PR's value prop is "pick CACHING for one table and HYBRID for another under joint HBM+DDR budgets" — a 2-table fixture where global DDR can only hold one CACHING but plenty of HBM is available would directly cover that. Same gap for the mem_bins_per_device legacy alias (untested) and the all-options-pruned / empty-search-space degenerate paths.

[github-actions]

Direct unit test for dynamicemb_calculate_shard_storages and _dynamicemb_aware_build_shard_perf_contexts is missing. Both have non-trivial branches — admission_counter HBM accounting (dynamicemb_util.py:663–672), the compute_device == "cpu" and not is_inference DDR branch (line 697), and the cache_params swap+restore (line 502–514) — that are reachable only via the E2E test, which is gated on has_dynamicemb and cuda. On a CPU dev box none of this is covered; a fake-ed _orig_build_shard_perf_contexts would let you verify the boost is applied and (more importantly) that the original cache_params is restored even when the wrapped call raises.

[github-actions]

Review summary

Solid PR overall — the 2D DP rewrite is clean, the mode-aware storage formula is well-tested, and the empirical perf-model docstring is refreshingly honest about its provenance. A few things worth addressing before merge (inline above):

Should fix

• dynamicemb_util.py:515-527 — missing try/finally around the cache_params swap; an exception in the wrapped estimator permanently corrupts the option's cache load factor.
• dynamicemb_util.py:56-89 — the x_eff step from 0.12 → 1.0 at x=1.00 means HYBRID@0.99 and CACHING@0.5 both look much worse than HYBRID@1.0 by construction. If x=1.0 is genuinely a different code path (HBM-only kernel), modeling it as a separate mode axis is cleaner than a discontinuity in the same knob.
• plan_util.py:846-853 — redundant double-deepcopy of dynamicemb_options (20×/table).

Worth a follow-up

• plan_util.py:434-463 — 2D DP state blows up on multi-node worlds (bins scale with world_size); consider an absolute cap on total bins, or vectorize with NumPy.
• Test gap on the multi-table DP case where one table wants CACHING and another HYBRID under a shared DDR budget — that's the headline scenario and currently isn't covered by a unit test. Same for dynamicemb_calculate_shard_storages and _dynamicemb_aware_build_shard_perf_contexts (only reachable via the CUDA-gated E2E test) and the mem_bins_per_device legacy alias.

Docs

• DynamicProgrammingProposer docstring says "per-rank DDR" but the implementation budgets DDR per-machine — worth aligning the wording with the inline comment on lines 350-352.
• dynamicemb_calculate_shard_storages Args block (not in diff, around line 620): the caching_ratio description still reads "ratio of HBM to DDR memory for UVM caching" — under the new CACHING semantics that's stale; host always holds the full backing store and caching_ratio only sizes the HBM hot-row cache.
• PR description mentions TZREC_CACHING_HIT_RATE_MULTIPLIER and x_eff = min(1.0, x · multiplier); the code uses fitted constants instead (no env var). Worth updating the description so it doesn't mislead future readers.

No security/multi-process concerns: the monkey-patches run under the import lock and each rank is its own process.

[github-actions]

Storage estimator ↔ runtime ↔ log are not coherent at cache_load_factor >= 1.0.

The log overrides dram_bytes = 0 here because the runtime drops the host tier when the table fits HBM. But the storage estimator at _calculate_dynamicemb_table_storage_specific_size (line ~340) still returns total_value_memory DDR for CACHING mode regardless of cache_ratio — i.e. CACHING@1.0 reports a full backing-store DDR cost that the 2D DP then uses for pruning. Result: the DP can mark CACHING@1.0 plans infeasible on DDR-tight topologies even though they would actually fit at runtime, and the log will then report local_dram=0.000GiB for any plan that does make it through.

Suggest aligning the storage formula with the runtime: when cache_ratio >= 1.0, force CACHING's DDR to 0 too. Then the storage estimator, the DP, and this log all agree.

[github-actions]

When the caller passes both mem_bins_per_device=X and hbm_bins_per_device=Y, the legacy alias silently wins — test_legacy_mem_bins_per_device_alias even codifies this. That's the surprising direction; a user supplying the new name explicitly while a config layer injects the legacy one will get the legacy value with no warning. The class docstring also reads as a one-way alias ("Accepts legacy mem_bins_per_device as an alias").

Suggest either raising TypeError when both are non-None, or letting hbm_bins_per_device win and emitting a DeprecationWarning for mem_bins_per_device. Whatever the choice, document it in the Args: block.

[github-actions]

The class docstring (line 255) says backtracking "yields Pareto-optimal proposals," but this loop actually enumerates every reachable (h, d) cell at the last table — up to hbm_bins × ddr_bins ≈ 800k+ proposals at the cap. Downstream consumers iterate propose()/feedback() once per proposal, so a planner run on a large world can churn through hundreds of thousands of plans rather than a small Pareto frontier.

Two minor things compound this:

• The sort key hd[0] + hd[1] weighs the two axes equally even though hbm_bin_size and ddr_bin_size differ — so "decreasing total memory" isn't really sorted by memory.
• last_back[h][d][0] >= 0 is the only filter; nothing prunes dominated cells.

Either tighten to the Pareto frontier (drop (h, d) if some (h', d') with h' ≤ h, d' ≤ d, perf' ≤ perf exists), cap len(self._proposal_list), or update the docstring to say "all reachable plans, larger memory first."

[github-actions]

The wrapper hard-codes the first six positional parameters (config, shard_sizes, sharding_option, topology, constraints, sharder) and forwards everything else via *args, **kwargs. If torchrec renames sharding_option, reorders these six, or splits one into multiple, the wrapper either silently mis-dispatches (the boost goes into the wrong object) or raises an opaque TypeError deep inside the planner with no link back to this file.

Worth capturing inspect.signature(_orig_build_shard_perf_contexts) at patch time, asserting the first six parameter names match, and otherwise logging a warning + installing a pass-through. A pinned/asserted torchrec version range somewhere visible would also help.

[github-actions]

(PR #508 review R5) looks like an in-flight review artifact — PR numbers and ad-hoc reviewer ticket IDs aren't useful once merged (git blame is the canonical pointer). Suggest dropping the parenthetical or replacing with an actual follow-up issue number.

[github-actions]

The dynamicemb_options parameter is never read — the function relies entirely on base_option.dynamicemb_options via copy.deepcopy(base_option) and opt.dynamicemb_options.caching = .... The caller also already assigns sharding_option.dynamicemb_options = dynamicemb_options immediately before invocation, so the argument is doubly redundant.

Either drop the parameter, or use it as the canonical source so callers don't have to mutate base_option first.

[github-actions]

This per-machine DDR pruning is the load-bearing reason for the 2D rewrite, but it's not exercised in tests — _make_topology in plan_util_test.py accepts local_world_size but every test falls through to the default (local_world_size or num_devices), so per_host == num_devices in every case. A bug where max_machine_ddr collapses to ddr_total or the window stride is off-by-one wouldn't be caught.

Worth a multi-host test: e.g. num_devices=4, local_world_size=2, ddr_per_device=500 (so each machine has 1000 DDR, total 2000), then assert an option whose shard.ddr=1500 is pruned even though 1500 < ddr_total.

[github-actions]

This ~25-line log block has zero coverage in dynamicemb_util_test.py / plan_util_test.py. Three independent decision points worth a test each:

• cache_load_factor >= 1.0 HBM_ONLY override (float-precision regressions would silently mis-label).
• int(os.environ.get("RANK", 0)) == 0 gating (a regression to os.environ.get("RANK") == "0" would fail on the typical single-rank case where RANK is unset).
• 1 << 30 GiB divisor (a regression to 1e9 would not be caught).

A mock.patch.object(logger, "info") test driving a synthetic ShardingOption through _to_sharding_plan would cover all three cheaply.

[github-actions]

Code review summary (5 reviewer agents)

Substantial, well-tested planner change. The 2D DP refactor is clean, the empirical-fit docstring on _dynamicemb_effective_cache_ratio is unusually informative, and the new try/finally-restored monkey-patch on ShardPerfContext is properly tested for the exception path. Inline comments below capture the items worth a look before merge — they cluster into:

Correctness / coherence

• Storage estimator vs. runtime vs. planning log disagree at cache_load_factor >= 1.0 (CACHING DDR reported as full backing in the DP but as 0 in the log).
• DP "emits Pareto-optimal proposals" docstring overstates what the code does (it emits every reachable (h, d) cell, up to ~800k at the cap).

Robustness

• ShardPerfContext.build_shard_perf_contexts monkey-patch hardcodes a 6-positional signature with no version assert.
• mem_bins_per_device legacy kwarg silently wins over the new hbm_bins_per_device when both are passed — and the test codifies it. Worth a DeprecationWarning or TypeError.

Test gaps (also flagged inline where relevant)

• Per-machine DDR pruning via local_world_size is the motivation for the 2D rewrite but no test ever uses local_world_size != num_devices.
• The new per-table [dynamicemb plan] ... log block (incl. HBM_ONLY override + rank-0 gating + GiB formatting) is not covered.
• _DP_AXIS_BIN_CAP is never exercised (all DP tests use num_devices ≤ 2 with bins=20).

Minor

• (PR #508 review R5) looks like a stray reviewer-ticket reference in a shipping comment.
• _emit_dynamicemb_variants(base_option, dynamicemb_options) has a dead second parameter.

Performance note (no inline)

• The pure-Python 4-level inner DP loop is acceptable for a one-shot startup call at the stated 30-table/16-GPU workload (~25-70 s; nested-list dp + backtrack peaks at a few GB), but lives within ~2× of unacceptable on heavier worlds. The (PR #508 review R5) TODO is the right follow-up — switching dp / backtrack to NumPy arrays drops both time and memory by ~10×.

🤖 Generated with Claude Code