[bugfix] preserve absent KJT weights in exported graphs

[eric-gecheng]

FBGEMM CUDA can return an undefined tensor for absent weights, causing native TorchScript inference to fail on a subsequent permutation. Insert FX guards after TorchRec MC and quantized permutations so unweighted inputs retain None in exported models.

[github-actions]

Review summary (static, multi-area)

Verdict: the fix is sound and the testing is unusually thorough. All findings are minor/style-level — see 3 inline comments on fx_util.py.

Cross-checked against the pinned deps (requirements/runtime.txt: torchrec 1.8.0 / torch 2.13 / fbgemm-gpu 1.8.0):

• Bug premise verified against FBGEMM sources: the CUDA permute_1D/2D_sparse_data kernels return a default-constructed (undefined) tensor in the optional-weights slot when weights are absent (CPU returns nullopt); the Python boundary converts undefined → None, masking it in eager but not in native TorchScript — matching the docstring exactly.
• Graph rewrite is correct: replace_all_uses_with + the restored.args re-pin maintains the users map properly (torch/fx args setter routes through _update_args_kwargs); meta["is_wrapped"] mirrors what FX's own tracer stamps (_symbolic_trace.py:1072) and is load-bearing for re-traces; chained permutations get guarded in topological order; the identity-permute paths (where source aliases the output) degrade to a benign no-op. The dedup check makes re-tracing idempotent, which matters because production re-traces guarded graphs (export_util.py:1833, :2237).
• Arg extraction matches reality: both torchrec call sites pass features positionally (mc_modules.py:418, quant/embedding_modules.py:536/543/945/952), and the kwargs fallback matches the real parameter name (exercised by the new test).
• All scripted export paths are covered: plain script (export_util.py:349-353), TRT, legacy AOT sparse (aot_utils.py:118-122), and unified AOTI (aot_utils.py:411) all route their final trace through fx_util.symbolic_trace.
• The native-TorchScript subprocess test genuinely reproduces the failure environment (loads raw fbgemm_gpu_py.so, asserts no Python-side fbgemm_gpu/tzrec import) and covers weighted/unweighted, all-empty lengths, int32/int64, and cached vs computed permute order. I verified the expected-permutation math by hand (MC [0,1,2,5,3,4] then quant [0,4,5,1,2,3] composes to [0,3,4,1,2,5] / keys ["a","d","e","b","c","f"]). Conventions check out: @mark_ci_scope("gpu","h20") at class level, parameterized_name_func, make_test_dir(), imports isort-clean, no E501, commit message explains root cause + fix per AGENTS.md.

Low-severity observations (no action strictly required):

1. torchrec private APIs: _mcc_lazy_init_inplace / _permute_kjt / KJT._weights are private, and these are the repo's first top-level private torchrec imports. Fine under the strict ==1.8.0 pin, but fx_util is imported by training-side modules too, so a future torchrec bump renaming either symbol breaks imports repo-wide, not just export. Worth remembering at the next dep bump (_mcc_lazy_init_inplace doesn't exist before torchrec 1.6).
2. QUANT_EMB=1 + ENABLE_AOT=2 is an untested combination: quantization runs before the unified-AOTI branch (export_util.py:277→331), so the guard node — which mutates its argument — would flow through torch.export (aot_utils.py:411-435). Structurally it's the same opaque torch.fx.wrapped call species as the _permute_kjt nodes already on that path, and any incompatibility would fail loudly at export time rather than misbehave at runtime; integration tests cover ENABLE_AOT=2 only without quant.
3. Redundant wraps in the test: torch.fx.wrap(_mcc_lazy_init_inplace) / torch.fx.wrap(_permute_kjt) at fx_util_test.py:42-43 are no-ops — torchrec already decorates both at definition (mc_modules.py:94, embedding_modules.py:130). Harmless (registration is idempotent), but droppable.
4. Subprocess timeout diagnostics: on Linux, subprocess.run(timeout=...) discards the child's captured output when it raises TimeoutExpired, so a hang would surface with no clue where the child was stuck (the non-zero-exit path nicely reports stdout+stderr). Borderline nit given 120s is generous.

Review logistics: five areas were dispatched (code quality, performance, security, documentation accuracy, test coverage). The first four completed; the test-coverage reviewer stalled and was stopped, so test adequacy was verified manually instead (permutation math, FX placeholder handling of the weights=None default, runner isolation, CI-lane scoping, torch.save/load(weights_only=True) compatibility).

🤖 Generated with Claude Code

[eric-gecheng]

Follow-up to the review summary: changes are pushed in 701fbc5, with replies on all three inline comments.

1. Private TorchRec APIs: moved the two private helper imports inside symbolic_trace, limiting their import-time dependency to tracing. The pinned TorchRec version and the necessary KJT._weights normalization remain unchanged.

2. Quantization and unified AOTI: added permanent torch.export regression cases with and without weights. Separately tested a CUDA INT8 MC-EBC model with shared tables and mixed SUM/MEAN pooling under QUANT_EMB=1 ENABLE_AOT=2: FX tracing, torch.export, and exported-model execution passed, and the guarded/unguarded ATen graphs were identical. The existing unified-AOT integration test does not explicitly disable EBC quantization; these additional checks target the guarded permutation path. Full AOTI binary compilation/package loading remains unverified because the matching local Torch 2.13 environment lacks the CUDA 13 development toolkit.

3. Test-side torch.fx.wrap registrations: retained them and documented why. Wrap registration is tied to the caller's globals; definition-site wrapping in TorchRec does not register these imported aliases in the test module. In an isolated probe, removing the two registrations changed the graph to two call_method("permute") nodes and produced zero guards, so they are not redundant.

4. Timeout diagnostics: TimeoutExpired already retains captured output, but the default test traceback did not display it. The test now reports decoded stdout/stderr and the child flushes stage/case progress. An injected timeout confirmed that both streams are included.

Validation: all 10 tests in tzrec.utils.fx_util_test passed; full pre-commit run -a passed; project-wide pyrefly check reported 0 errors.