[feat] bump to 1.4.0: torch 2.13.0 / torchrec 1.8.0 / fbgemm 1.8.0

[tiankongdeguiji]

Cuts the 1.4.0 release: torch 2.13.0 / torchrec 1.8.0 / fbgemm_gpu 1.8.0, with
torch-tensorrt 2.13.0 (TensorRT 11.0) in the cu130 image. triton stays at 3.7.1 —
torch 2.13.0 pins the same version and the same upstream triton commit as 2.12.1,
so the repacked wheel carrying the H20 ptxas fix is still the right one.

The three bumps have to move together: torchrec 1.8.0 imports fbgemm_gpu.tbe.config
unguarded, so fbgemm-gpu 1.7.0 stops being a viable floor, and fbgemm 1.8.0 in turn
targets torch 2.13.

cu129 kept, on a rebuilt toolchain

cu129 was initially dropped here: fbgemm 1.8.0's tbe_training_backward .so overflows
the linker's 2 GiB PC32 range under CUDA 12.9, which packs the same SASS into a
noticeably larger .nv_fatbin than CUDA 13 does (cu130 is itself already at 96.3% of
that limit). Rebuilding the wheels with nvcc --compress-mode=size brings it back
under, at the cost of requiring an R550.54.14+ driver — so the variant stays.

It fills a real gap between the other two images: cu126 runs on R525 but stops at
sm_90, while cu130 reaches Blackwell but needs R580. cu129 covers cu130's
architectures on a 550 driver. The docs now lead with a driver table, since driver —
not compute capability — is the first thing to choose on:

image	min driver	SASS	TensorRT
-cu126	525.60.13	sm_50 … 90	no
-cu129	550.54.14	sm_75 … 120	no
-cu130	580.65.06	sm_75 … 120	yes

They also note that the op wheels ship no PTX, so an unlisted architecture cannot JIT
fall back, and that LD_LIBRARY_PATH=/usr/local/cuda-13/compat lets cu130 run on an
older driver via the bundled cuda-compat.

dynamicemb / fbgemm_gpu_hstu move out of extra.txt

Both now live in requirements/cu126.txt and requirements/cu130.txt, so the GPU
images and the new tzrec[cu126] / tzrec[cu130] extras carry them instead of
requiring a separate manual install. requirements/extra.txt keeps only the RTP
torch_fx_tool helper, which scripts/ci/ci_test.sh still installs on its own.
Both wheels are rebuilds of the same upstream source commit against the new
torch/CUDA (dynamicemb 6b94bbf, fbgemm_gpu_hstu fece651b), so no API change
follows them.

Compatibility audit

A source-level audit of every place TorchEasyRec copies, subclasses or monkeypatches
upstream internals — plan_util.py (enumerator/estimator/calculate_shard_storages),
dynamicemb_util.py (to_sharding_plan, HardwarePerfConfig.get_device_bw,
ShardPerfContext.build_shard_perf_contexts, the jagged-tensor validator),
optim/optimizer.py (apply_split_helper), dist_util.py (DMP, train pipelines,
mc-ebc input_dist) and export_util.py (RTP fbgemm op remaps, _all_keys_used_once)
— found no signature or import breakage. fbgemm's apply_split_helper body is
byte-identical between 1.7.0 and 1.8.0, and every torch.ops.fbgemm.* schema
TorchEasyRec calls is unchanged.

tzrec/acc/aot_utils.py drops _backport_pt178147_int_array_dedup:
pytorch/pytorch#178147 landed in 2.13.0, and the patch's guard only recognised its own
marker attribute, so leaving it in would have silently reinstated the older
codegen_int_array_var over the upstream fix.

Three behaviour regressions the bump surfaced

The first two were caught by CI on the first run; the third by a user report.

Dense graphs stopped being prunable. torch 2.13 taught
torch._library.utils.is_impure to resolve an OpOverloadPacket to its default
overload before checking schema mutability. FX records call_function targets as
packets, so the zero-user fbgemm.bounds_check_indices node a quantized lookup leaves
behind is now impure and eliminate_dead_code keeps it. In a dense graph that node is
unreachable from the output but still pins the whole sparse KJT subtree, so the dense
graph went on reading raw feature keys — which are absent from the sparse model's
output dict it is handed at export time, hence KeyError: '<feature>.values' across
every AOT/TRT export test. A dense graph is only ever the subgraph feeding its output,
so _eliminate_dense_dead_code prunes it by reachability, anchoring placeholders and
the output; the sparse-side callers keep default semantics because the lookup
legitimately belongs there.

The cu130 torch-tensorrt wheel has no C++ runtime. torch_tensorrt-2.13.0+cu130
on the pytorch-wheels index is a 620 KB Python-only build (2.12.1's was 3.87 MB with
libtorchtrt.so and the _C extension), so ENABLED_FEATURES.torch_tensorrt_runtime
is False and setup_engine() falls back to a pure-Python engine object. 2.13 also
dropped the engine is None guard from forward(), so handing that object to
tensorrt::execute_engine under torch.jit.trace failed as unsupported input type: PyObject rather than anything diagnosable. The plain PyPI wheel does carry the
runtime, so the Dockerfile takes it from there instead of the pytorch-wheels index.

An unusable hstu wheel took down every entry point. hstu/library.py only loads its
.so when torch.cuda.is_available() is true, while hstu/hstu_ops_gpu.py registers
fakes for fbgemm::hstu_varlen_fwd_* unconditionally — so on a GPU image whose driver
cannot serve the bundled CUDA runtime, the fake registration raises RuntimeError. The
guard in cutlass_hstu_attention.py caught only ImportError, and since main.py
imports aot_utils (and thus this module) at module scope, python -m tzrec.train_eval
died at import. This is exposed by the packaging change above: before 1.4.0 the wheel
lived in extra.txt and was simply absent from the images, so the import failed with
ImportError and the narrow guard sufficed. It now catches any exception and degrades
to the PyTorch/Triton kernel. Verified that a container with no usable GPU imports and
runs on CPU.

One latent bug on master, fixed here

tzrec/ops/benchmarks/hstu_attention_bench.py builds its CLI with click, but nothing
ever required it — pyre-check happened to pull it in transitively, so the import
resolved by accident. Replacing pyre-check with pyrefly drops that edge, and this is the
first image built from the new requirements, so pyrefly check reports the module as
missing. click is now a declared dependency. This is not caused by the bump: master
will hit it too as soon as its image is rebuilt.

Other

Pre-commit hooks move to ruff v0.16.4 and codespell v2.4.3 (mdformat deliberately
left at 0.7.22). Docs are updated for the new stack, the dropped cu129 image, and the
tzrec[cu130] / tzrec[cu126] install commands. The DLC tutorial's generated dlc submit command now uses the PAI official torcheasyrec image the tutorial already
tells you to select, instead of tzrec-devel, and passes the 535 driver its resource
section calls for. The PPU image is untouched and still tracks the vendor 1.2 stack.

Test Plan

• pre-commit run -a clean; pyrefly check reports 0 errors on the merged branch.
• Full CI green on the pre-merge revision: Code Style, PyTyping, Build Test, Unit Test
  (GPU), Unit Test CPU and Unit Test H20 all pass.
• The 10 GPU-lane tests that failed before the fixes were re-run in the
  tzrec-test:1.4-cu130 container: 8 pass with the dense-graph fix, and the 2 TRT ones
  pass once torch-tensorrt carries its runtime.
• python tzrec/tests/run.py --scope h20 on real 8×H20 hardware
  (torch 2.13.0+cu130 / torchrec 1.8.0 / fbgemm 1.8.0+cu130 / triton 3.7.1):
  94 tests, all pass — this is what covers the CUTLASS/WGMMA HSTU kernels.
• Image sanity for 1.4-cu130-u4 (and the same checks on the cu129 image): torch 2.13.0+cu130, torchrec 1.8.0, fbgemm
  1.8.0+cu130, triton 3.7.1, torch-tensorrt 2.13.0 with TensorRT 11.0.0.114 and
  runtime=True, plus master's graphlearn 1.3.9 / pyfg 1.0.6 / feature-store-py 2.2.10,
  and click / dynamicemb / hstu / faiss all importing. pyrefly check also runs clean
  inside that image.

Workflow YAMLs point at tzrec-test:1.4-u4 for this validation round and are reverted
to tzrec-devel:1.4 once the images are promoted. The registry rejects overwriting an
existing tag, which is why the validation images carry a -u<n> suffix; promotion
publishes them to tzrec-devel:1.4-<device> without it.

🤖 Generated with Claude Code

https://claude.ai/code/session_01BT6xK7cczW2iAD8uMNxvuU

[github-actions]

This dlc submit command installs tzrec[cu126], but the --worker_image below (line 151) is the bare tzrec-devel:${TZREC_DOCKER_VERSION} tag, which scripts/promote_docker.sh defines as an alias of the cu130 image. Before this PR the command was device-agnostic (pip install tzrec==), so the bare tag was harmless; now the extras are device-specific — tzrec[cu126] pulls the cu12 builds of faiss, dynamicemb and fbgemm_gpu_hstu, which pip would install over the cu130 image's preinstalled cu130 versions, breaking them at import time (no CUDA 12 runtime in that image).

Suggest pinning the worker image to tzrec-devel:${TZREC_DOCKER_VERSION}-cu126 here to match the cu126 PAI image the rest of this tutorial uses.

[github-actions]

The fix itself looks correct (anchors only placeholders/output; all 4 dense call sites converted, the 3 sparse-side ones correctly keep default DCE), but there is no CPU-runnable regression guard for it, so the next torch impurity change would only be caught by the GPU-lane e2e tests.

The existing CPU tests are structurally blind to this regression: they trace sparse embeddings as leaf modules and don't quantize, so the dead subtree contains only pure nodes that plain DCE removes on any torch version; and their serving_data is built from dict(batch.to_dict()), which still carries raw <feature>.values/.lengths keys — an unpruned dense graph runs fine in-test and KeyErrors only in real serving. (_run_dense_graph_sanity_check also skips running the pruned graph.)

A cheap CPU-lane tripwire: after build_dense_graph_module/split_model, assert no node in the dense graph does a getitem on a raw feature key (or has an fbgemm-target call_function), and/or run the scripted model on serving-shaped input built strictly from the dense_graph_config keys. The export_distributed_embedding path has the same gap and additionally no e2e test in any lane.

[github-actions]

Code review summary

Reviewed statically across five dimensions (code quality, performance, test coverage, documentation accuracy, security). Overall the bump is in very good shape.

Verified clean:

• The dense-graph DCE fix is correct: is_impure_node anchors only placeholder/output, so DCE erases exactly the nodes unreachable from the dense output — restoring the pre-2.13 pruning behavior. All 4 dense call sites were converted (export_rtp_model, split_model, export_distributed_embedding, build_dense_graph_module, which also covers the CPU and online dense paths); the 3 sparse-side eliminate_dead_code() calls correctly keep default semantics, since the live sparse lookup chain legitimately contains zero-user mutating ops.
• The _backport_pt178147_int_array_dedup removal is complete — no dangling references — and its removal condition (torch >= 2.13.0) is met by the new pin.
• cu129 cleanup is thorough: zero remaining references to cu129, torch 2.12.x, torchrec/fbgemm 1.7.0, or the 1.3 image tags anywhere in the repo (code, docs, scripts, workflows).
• Wheel pins are consistent across requirements/, the Dockerfile, and docs (dynamicemb 20260824.6b94bbf, fbgemm_gpu_hstu 20260823.fece651b, triton 3.7.1, torch-tensorrt 2.13.0 cu130-only). The torch_tensorrt==2.13.0 PyPI install is exactly pinned, and torch 2.13.0+cu130 satisfies its torch==2.13.0 dependency under PEP 440 local-version semantics, so no second torch/CUDA stack is pulled. No new performance or supply-chain risks introduced.

Two inline comments:

1. docs/source/quick_start/dlc_tutorial.md — the dlc submit example installs tzrec[cu126] but uses the bare tzrec-devel tag, which is the cu130 alias — a device mismatch this PR introduces (the old device-agnostic install was harmless here).
2. tzrec/utils/export_util.py — the pruning regression has no CPU-runnable guard; the existing CPU tests are structurally blind to it (leaf-module tracing, no quantization, serving data still carrying raw feature keys).

Two minor notes (not blocking):

• All 9 workflows pin the transient tzrec-test:1.4-u2 images; the PR body says a follow-up reverts them to tzrec-devel:1.4 after promotion — worth making sure that follow-up lands, since nothing in the repo flips them back automatically.
• docs/source/quick_start/local_tutorial.md: now that install extras are device-specific, a one-line note that the -cu126/-cpu images need tzrec[cu126] instead of the shown tzrec[cu130] would prevent copy-paste mismatches.