[chaofengw-nv]

Background

Fixes #1799. The pinned Moshi reference loads PersonaPlex with dtype=torch.float16. Each RMSNorm.alpha, stored as BF16 in the checkpoint, is therefore held in Half. _rms_norm then computes in FP32 with that Half value.

The family add_rms_norm instead embedded the checkpoint weight as an FP32 constant and never rounded it. Most BF16 weights are exact in Half, but small ones that land in the Half subnormal range are not. 5 of the 278,528 PersonaPlex RMSNorm weights differ this way, all in temporal layers 0, 1 and 3 norm1.

Exit Criteria

• When the norm runs on a lower-precision input, its weight is rounded to that input dtype before the FP32 computation, as in the reference.
• FP32 networks keep the FP32 weight.
• A predeclared GPU counterexample fails on the base and passes on this head through the norm's normal output, using exact comparison.
• Non-goals:
  • No change to the norm's FP32 operation grouping or reciprocal square root. Those are same-precision reorderings, treated as a numerical-alignment candidate rather than a defect, and are not part of this PR.
  • No claim that the full-model audio-parity gate passes.

Implementation

• graph_ops.add_rms_norm: when need_cast is set, the FP32 gamma constant is cast to the input dtype and back to FP32 before the multiply. The rounding is expressed in the graph, so it follows the network's actual dtype (Half or BF16). The FP32 norm computation is unchanged.
• New test families/personaplex/tests/test_rmsnorm_weight_dtype.py (markers gpu and trt; it skips without CUDA):
  • inputs x = [3, 1, 1, 1] and gamma = [-1.0207295417785645e-06, 1, 1, 1], the real BF16 checkpoint value of transformer.layers.0.norm1.alpha[65];
  • mean(x * x) is exactly 3 at every width, so neither the reduction order nor the operation grouping affects the result;
  • widths 4 and 4096 (the block repeated), with 3 fresh engine builds each, compared exactly with the values of the reference _rms_norm on CUDA: -1.7285346984863281e-06 and 0.5771484375 (the old graph gives -1.7881393432617188e-06);
  • an FP32 control checks that FP32 networks still use the unrounded weight.

Change categories

• Model or runtime behavior
• Public API
• ABI
• Bundle or artifact format
• Dependencies
• Documentation only
• CI or developer tooling

Validation

Commands and Results

• Original/fixed control on GB300. Isolated trees of main 9a8b545 and of this head, each with the identical new test (sha256 cda09fa8...). Command: python -m pytest -q -rA families/personaplex/tests/test_rmsnorm_weight_dtype.py.
  • Original: 6 failed, 1 passed (the FP32 control). Each Half failure has 1/4 or 1024/4096 mismatched elements, with a maximum absolute difference of 5.96e-08 (one Half subnormal step).
  • Fixed: 7 passed.
  • graph_ops.py in the fixed tree equals this head (7fd7ea81...).
• Family CPU tests on the fixed tree with CUDA hidden: python -m pytest -q -m "not gpu and not trt and not e2e" families/personaplex/tests: 13 passed, 4 skipped.
• Reference data (diagnostic). The reference FP16 run of audio-parity clip 7, temporal norm1/norm2 inputs captured at layers 0 and 31, steps 0, 27 and 28:
  • The reference CUDA _rms_norm reproduces every captured output exactly.
  • The family norm before this change differs in 8 elements over the 12 captures. Element 65 of layer 0 norm1 (step 0) is removed only by rounding the weight to Half.
  • The rest are associated with the FP32 operation grouping and reciprocal square root, which this PR does not change.
• ruff check on both files and ruff format --check on the new test: clean. git diff --check: clean.

Hardware, Environment, and Revisions

GB300 (SBSA), CUDA 13.0, TensorRT 11.2.0.113, PyTorch from the product container. Tested head 3f15493d00d61a21889e9c72bf275f57c20e5ddb. Reference data: checkpoint nvidia/personaplex-7b-v1 revision fdaf4090a61cb315c138a1faee287ffd6c716309, run with the pinned PersonaPlex reference environment in FP16.

Not Run / Remaining Gaps

• BF16 networks, tensor-parallel builds and the full model are not covered by this control. The BF16 and TP builders reach the same add_rms_norm, but no separate BF16 or TP evidence was collected.
• Full model (diagnostic, not a pass). The personaplex-7b audio-parity profile, original 20 clips, true-FP16 temporal/depth core, against the unchanged Native FP16 outputs and gate. The tree is main plus the current heads of fix(personaplex): align FP16 inference with Native #1751, fix(personaplex): keep RMSNorm epsilon in FP32 #1780, fix(personaplex): preserve Half embedding additions #1781, fix(personaplex): avoid timing discarded audio decode #1783, fix(personaplex): round the FP16 depth input like the reference #1789, fix(personaplex): preserve reference RoPE precision #1791, fix(personaplex): round the SwiGLU gate product once #1794, fix(personaplex): build the FP16 profile with a Half core #1795, fix(personaplex): preserve residual input rounding #1797, fix(personaplex): match Native HQ audio resampling #1798 and this PR, on one GB300 with one bundle build.
  • Result: 18/20, the same pass count as the same tree without this PR.
  • Failures: clip 7 at 5.58 dB with RMS ratio 0.851 (unchanged), and clip 11 at 6.42 dB with RMS ratio 0.846 (6.90 dB without this PR).
  • Effect of this PR: 5 of the 20 clips' audio changes (1, 10, 11, 12, 18), and every clip keeps its pass/fail status.
  • The same-batch reference FP16 run reproduces the frozen reference PCM bit for bit on 20/20 clips, with equal output frame counts.
  • Output samples and rate are equal on 20/20 clips. Same-batch task P50 (median server-side model call, n=20 each, one TensorRT build): reference 2125.46 ms, TensorRT 1474.17 ms.
  • Clip 7 is not attributed to this change. The unmodified reference itself gives 5.58 dB on clip 7 against its own frozen output when only its attention backend changes from cuDNN to the efficient or math kernel.
• ThorX, Internal CI premerge and the family e2e test on GPU were not run.

Contributor Self-Review

• I have completed a self-review of this change.

I checked every caller of add_rms_norm in the family: graph_blocks.apply_norm from the default and TP decoders, and utils.norm_multi from the dual-profile decoder. They all pass the decoder's work dtype, so FP32 networks keep need_cast false and are unchanged. I also confirmed that the cast pair rounds subnormals as the reference does: the captured row and the frozen test both match exactly.

Notes For Future Readers

• The reference computes x * (alpha * rsqrt(var)) in FP32, while this graph computes (x * (1 / sqrt(var))) * alpha, also in FP32. This same-precision reordering accounts for most of the remaining 1-ULP norm differences on GB300. It is a numerical-alignment candidate, not a confirmed defect: POW(var, -0.5) with the reference grouping still differs from CUDA rsqrt in 11 of 66,560 synthetic cases (2 of them new), and no full-model improvement has been shown. It is deliberately left out of this PR.
• Independent of fix(personaplex): round the FP16 depth input like the reference #1789, fix(personaplex): preserve reference RoPE precision #1791, fix(personaplex): round the SwiGLU gate product once #1794 and fix(personaplex): build the FP16 profile with a Half core #1795; touches only add_rms_norm and a new test.

Risk level

• Low
• Medium
• High

Risk rationale: one weight-rounding step in FP16/BF16 RMSNorm. With this checkpoint, only 5 weights change, and only in the Half subnormal range. FP32 builds are unchanged. It rebuilds PersonaPlex bundles.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[chaofengw-nv]

Superseded by #1751. This change and its regression coverage are preserved in commit 93d1acb of the consolidated PR. Original issue links and acceptance criteria remain tracked there.