[bugfix] avoid tensor-sized temporaries when trunc_normal_ inits large embedding tables

[tiankongdeguiji]

Problem

nn.init.trunc_normal_ in torch>=2.12 was rewritten with rejection sampling (pytorch/pytorch#174997). The new implementation materializes several temporaries as large as the tensor itself: the bounds check mask = (result < lo) | (result > hi) alone allocates up to three tensor-sized bool masks, and any resample iteration adds two fp32 buffers of full tensor size.

torchrec applies a table's init_fn to the full local shard at DMP sharding time (ShardedEmbeddingBagCollection._initialize_torch_state -> reset_parameters). For a very large table this transient allocation is huge — e.g. a 150M-row x dim-64 ZCH table with a fused_uvm_caching kernel needs an 8.94 GiB first mask (and ~27 GiB peak at that line) — and it OOMs GPUs that ran the exact same config fine on torch<=2.11, where trunc_normal_ was fully in-place (uniform_ -> erfinv_ -> mul_ -> add_ -> clamp_, ~0 extra bytes).

Fix

• Add tzrec/utils/init_util.py:
  • create_init_fn(expr) centralizes the eval of config init_fn expressions (previously inline eval(f"partial({expr})") at three sites) with an explicit {partial, nn, torch} namespace, and swaps nn.init.trunc_normal_ for an env-switchable trunc_normal_.
  • trunc_normal_ defers to native nn.init.trunc_normal_ by default. Setting USE_INPLACE_TRUNC_NORMAL=1 initializes with the in-place inverse-CDF implementation of torch<=2.11 (~0 transient HBM), restoring for memory-constrained huge-table jobs exactly the initialization their configs trained with on torch<=2.11.
• The zch.threshold_filtering_func eval in feature.py previously resolved partial/nn from module globals kept alive by # NOQA imports; it now gets an explicit namespace (partial, nn, torch, and the three *_threshold_filter helpers) instead of leaking module globals into config eval.
• Bump version to 1.3.4.

An earlier revision dispatched on a numel > 2**27 threshold; review correctly noted that a fixed element cutoff neither bounds native temporaries just below it (~1.4 GiB possible) nor avoids silently reintroducing the torch<=2.11 endpoint-clamp tail defect above it (a uniform_ sample landing exactly on an interval endpoint clamps to a/b after erfinv_; measured ~2^-24 per element, ~450 exact -2.0 outliers on a 150M x 64 table). The explicit opt-in keeps native statistics for everyone by default and makes the memory/statistics trade a deliberate per-job choice.

Test Plan

• tzrec/utils/init_util_test.py (7 tests): expression parsing and trunc_normal_ swap (kwargs preserved, non-trunc-normal init fns untouched), default path calls native nn.init.trunc_normal_ with exact argument forwarding, USE_INPLACE_TRUNC_NORMAL=1 path calls the in-place implementation with exact argument forwarding, statistics through the public API with nonzero mean (mean=0.5, std=0.0125), truncation at ±1σ, and generator reproducibility.
• tzrec/features/id_feature_test.py: added a partial(dynamic_threshold_filter, ...) case for zch.threshold_filtering_func to guard the explicit eval namespace.
• Ran tzrec.utils.init_util_test (7), tzrec.features.id_feature_test (58), tzrec.modules.embedding_test (52), tzrec.models.wide_and_deep_test (3, covers wide_init_fn) — all pass; pre-commit clean.
• End-to-end: built an IdFeature with init_fn: "nn.init.trunc_normal_,mean=0.0,std=0.0125" and verified the config-created init_fn calls native torch by default and the in-place path with USE_INPLACE_TRUNC_NORMAL=1, with correct statistics on both.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Lu52Dws7J8GLgr7Hzviisn

[github-actions]

Minor test gap: the in-place path's add_(mean) step is only exercised with mean=0.0, so a transcription error in the mean handling wouldn't be caught. Since the whole point of _inplace_trunc_normal_ is to faithfully reproduce the torch<=2.11 result, consider asserting the mean of a nonzero-mean fill (e.g. mean=0.5) here.

[github-actions]

Low priority: norm_cdf((a - mean) / std) divides by std with no guard, so a degenerate config (std=0, or a >= b) silently fills a large table with NaNs on this path. It's a malformed-config edge case (trusted config), but a std > 0 / a < b assertion would fail loudly instead of poisoning the embedding, and keep behavior consistent with the small-tensor path.

[github-actions]

Review summary

Clean, well-scoped bugfix. The root cause (torch>=2.12's rejection-sampling trunc_normal_ materializing tensor-sized bool masks + fp32 resample buffers at DMP sharding time) is correctly diagnosed, and the in-place inverse-CDF fallback (uniform_ -> erfinv_ -> mul_ -> add_ -> clamp_) faithfully reproduces the torch<=2.11 path with ~0 transient HBM. Verified:

• The identity swap keys on init_fn.func is nn.init.trunc_normal_, which correctly matches both nn.init. and torch.nn.init. spellings (same object).
• All three prior eval(f"partial({expr})") sites are converted to create_init_fn; the only other init path (dynamicemb) is a separate, out-of-scope TODO.
• Tightening the threshold_filtering_func eval to an explicit namespace is a genuine improvement (strictly narrower named-symbol surface). Note it is not a sandbox — CPython still auto-injects __builtins__ when the globals dict omits it — but that's unchanged from the prior module-globals eval and fine for trusted config.
• The numel > 2**27 threshold is evaluated per-shard (torchrec applies init_fn to the local shard), so it engages for the large-table case as intended, with no regression on the common small-tensor path.

Left two low-priority inline nits (a nonzero-mean test assertion, and a degenerate-config std>0 guard). Neither is a blocker — ship as-is.

One thing worth a quick self-check before merge: the docstring/PR-body reference to pytorch/pytorch#174997 and the exact 2.11/2.12 version boundary can't be verified statically; a wrong upstream reference tends to rot silently for whoever later tries to confirm when this workaround can be dropped.

[github-actions]

low/up round to 0/1 for the motivating std=0.0125, so uniform_ can produce a representable endpoint; erfinv_ makes it infinite and the final clamp turns it into an exact ±2 (160σ) outlier. On 9.6B FP32 elements, hundreds of such hits are likely, and FP16/BF16 are worse—this is the numerical defect pytorch/pytorch#174997 fixed. Please retain native sampling in bounded chunks (or explicitly exclude dtype endpoints) and add extreme-tail/low-precision coverage.

[github-actions]

This cutoff can still permit the same OOM class: at 2**27 FP32 elements the bounds expression alone can peak at ~384 MiB of masks, and a resample/narrow-truncation step adds about 1 GiB of full-size buffers. A shard at or below the cutoff can therefore fail with tight HBM headroom. Please use bounded chunks, or derive dispatch from an explicit transient-byte budget and dtype/parameters instead of a device-agnostic element count.

[github-actions]

Static review summary

The init-function centralization and narrower eval namespace are clean improvements. Focused review covered code quality, performance, test coverage, documentation accuracy, and security. I left two inline concerns:

• The inverse-CDF fallback reintroduces endpoint-clamped tail outliers, so it is not statistically equivalent to the native sampler.
• The fixed element cutoff can still allow substantial native temporaries and OOM below the boundary.

No additional noteworthy findings. Tests and builds were not run, per request.

[tiankongdeguiji]

Addressed the review findings in daf9e27:

1. Endpoint-clamp outliers of the inverse-CDF path (init_util.py:57) — confirmed empirically: on torch 2.11 (identical algorithm), a 2^30-element fp32 CUDA fill with std=0.0125 produces 51 elements exactly at -2.0 (~2^-24/element, single-sided since uniform_ includes from), i.e. ~450 outliers on a 150M×64 table. Resolved together with the threshold concern below: the in-place path is no longer the default for any size. trunc_normal_ now defers to native nn.init.trunc_normal_ unless USE_INPLACE_TRUNC_NORMAL=1 is set, which opts a memory-constrained job into exactly the torch<=2.11 initialization behavior (defect included, as documented in the docstring) that such configs were trained with for years.

2. 2**27 cutoff can still OOM below the boundary / ignores dtype (init_util.py:92) — agreed; the fixed element cutoff was a heuristic, not a bound. The threshold is removed entirely in favor of the explicit env switch, so the memory/statistics trade is a deliberate per-job choice rather than a silent size-dependent behavior change.

3. add_(mean) untested with nonzero mean (init_util_test.py:50) — fixed; the statistics test now fills with mean=0.5, std=0.0125 through the public API and asserts the mean.

4. std=0 / a>=b degenerate configs "silently fill with NaNs" (init_util.py:55) — not reproducible; the failure is already loud on both paths: std=0.0 raises ZeroDivisionError: float division by zero (Python float division in norm_cdf), and a > b raises RuntimeError: uniform_ expects to return a [from, to) range. a == b fills the constant a, identical to torch<=2.11 native behavior. No guard added.

🤖 Generated with Claude Code