[yuchenwang3]

The compiled gated-residual helpers can fuse away the low-precision rounding after division, sigmoid, and multiplication. Their outputs and gradients can then differ from the same gate chain evaluated as separate BF16/FP16 operations.

This keeps those operation boundaries for BF16/FP16 and retains the compiled helpers for FP32. The eager regions are also excluded from an outer torch.compile. Projections, weights, and the gate formulas are unchanged. The existing GR unit tolerates compilation drift; this makes the low-precision helpers agree with its eager reference exactly, at the cost of fusion in that path.

Validation:

• 15 CPU helper cases passed on macOS, PyTorch 2.13.0, checking exact low-precision outputs and input gradients across stream counts and row sizes, plus FP32 dispatch. The local runner loads the helper definitions without CUDA-only Transformer Engine imports.
• A CPU negative control on the original compiled _mix_and_reduce reproduces the difference: with 17 rows, four streams, and width 64, BF16 output/input-gradient maximum absolute differences from eager are 0.00390625/0.00244140625. FP16 differs too. These are helper results, not full-model error bounds.
• pre-commit run --all-files passed.
• 16 tests passed with no skips on NVIDIA L20C, Python 3.12.3, PyTorch 2.11.0+cu130 / CUDA 13.0. The low-precision helper checks passed on CUDA, and the existing real TELinear unit reported zero output/input-gradient differences for both combine modes at 64 and 256 rows.
• A CUDA negative control at four 2560-wide streams reproduced BF16 compiled/eager differences up to 0.0078125 in outputs and input gradients; patched helpers matched eager exactly. FP16 also reproduced differences.
• Interleaved helper measurements with 10 warmups, 40 samples per variant, and alternating order showed a cost at larger row counts: for BF16 at 4096 rows, forward+backward synchronized wall medians were 0.264 ms compiled versus 0.368 ms eager (1.39×). CUDA-event medians were 0.244 versus 0.347 ms. At 257 rows, wall medians were 0.242 versus 0.184 ms. These short helper timings are affected by dispatch and shared-host activity and do not measure full-module or model throughput.
• CP/TP and full actor/rollout log-probability parity have not been established by these tests.

Implementation and tests were generated with OpenAI Codex.

[yuchenwang3]

@hjh0119 Could you take a look at the gated-residual rounding fix? The GPU run passes all 16 tests, including BF16/FP16 helper checks and the TELinear unit. The description includes the timing cost of the eager path; does that tradeoff look reasonable?

[hjh0119]

Do we have additional empirical evidence supporting the necessity of casting to bf16?