From 30faad4b1a58d1b6acb57bc6ebf5ca28a3c06401 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Sun, 18 Oct 2026 10:44:08 +0000 Subject: [PATCH] fix(personaplex): preserve residual input rounding Preserve the Half linear-output boundary before attention and MLP residual additions. The previous fused path can add the FP32 GEMM accumulator to the residual before rounding, unlike Native Half linear followed by Half addition. Sum the already rounded branch outputs in FP32 and restore the model dtype in the standard, dual-profile and TP decoders. Half weights, hidden states and KV caches stay Half; the FP32 graph is unchanged. The actual decoder-layer regression fails numerically on all three original Half builds and passes on all six corrected Half/Float controls, with identical assertions and zero errors or skips. Whole-model acceptance remains separate. Fixes #1796 Signed-off-by: chaofengw --- families/personaplex/decoder_tp_builder.py | 16 +-- families/personaplex/default_decoder.py | 20 ++- .../default_dual_profile_decoder.py | 18 +-- families/personaplex/graph_ops.py | 15 ++ .../tests/test_residual_precision.py | 136 ++++++++++++++++++ 5 files changed, 169 insertions(+), 36 deletions(-) create mode 100644 families/personaplex/tests/test_residual_precision.py diff --git a/families/personaplex/decoder_tp_builder.py b/families/personaplex/decoder_tp_builder.py index 7adf423e75..690ac0995e 100644 --- a/families/personaplex/decoder_tp_builder.py +++ b/families/personaplex/decoder_tp_builder.py @@ -418,9 +418,7 @@ def _add_tp_decoder_layer( else: norm2 = attn["normed"] else: - residual1 = network.add_elementwise( - hidden, attn_out, trt.ElementWiseOperation.SUM - ).get_output(0) + residual1 = graph_ops.add_residual_sum(network, hidden, attn_out) norm2 = _apply_norm( network, residual1, @@ -447,17 +445,11 @@ def _add_tp_decoder_layer( mlp_out = add_all_reduce_sum(network, mlp_out, tp_size) if parallel_residual: - sum_attn = network.add_elementwise( - hidden, attn_out, trt.ElementWiseOperation.SUM - ).get_output(0) - residual2 = network.add_elementwise( - sum_attn, mlp_out, trt.ElementWiseOperation.SUM - ).get_output(0) + sum_attn = graph_ops.add_residual_sum(network, hidden, attn_out) + residual2 = graph_ops.add_residual_sum(network, sum_attn, mlp_out) post_attn_tensor = sum_attn else: - residual2 = network.add_elementwise( - residual1, mlp_out, trt.ElementWiseOperation.SUM - ).get_output(0) + residual2 = graph_ops.add_residual_sum(network, residual1, mlp_out) post_attn_tensor = residual1 return { diff --git a/families/personaplex/default_decoder.py b/families/personaplex/default_decoder.py index 272a57a595..db34332ea4 100644 --- a/families/personaplex/default_decoder.py +++ b/families/personaplex/default_decoder.py @@ -667,10 +667,10 @@ def _add_decoder_layer( else: norm2 = attn["normed"] else: - residual1 = network.add_elementwise(hidden, attn_out, trt.ElementWiseOperation.SUM) + residual1 = graph_ops.add_residual_sum(network, hidden, attn_out) norm2 = _apply_norm( network, - residual1.get_output(0), + residual1, hidden_size, weights[f"{prefix}.post_attn_norm"], weights.get(f"{prefix}.post_attn_norm_beta"), @@ -709,19 +709,15 @@ def _add_decoder_layer( # Final residual connection if parallel_residual: - sum_attn = network.add_elementwise(hidden, attn_out, trt.ElementWiseOperation.SUM) - residual2 = network.add_elementwise( - sum_attn.get_output(0), mlp_out, trt.ElementWiseOperation.SUM - ) - post_attn_tensor = sum_attn.get_output(0) + sum_attn = graph_ops.add_residual_sum(network, hidden, attn_out) + residual2 = graph_ops.add_residual_sum(network, sum_attn, mlp_out) + post_attn_tensor = sum_attn else: - residual2 = network.add_elementwise( - residual1.get_output(0), mlp_out, trt.ElementWiseOperation.SUM - ) - post_attn_tensor = residual1.get_output(0) + residual2 = graph_ops.add_residual_sum(network, residual1, mlp_out) + post_attn_tensor = residual1 return { - "hidden": residual2.get_output(0), + "hidden": residual2, "post_attn": post_attn_tensor, "present_k": present_k, "present_v": present_v, diff --git a/families/personaplex/default_dual_profile_decoder.py b/families/personaplex/default_dual_profile_decoder.py index 19fca8d359..53f706fe32 100644 --- a/families/personaplex/default_dual_profile_decoder.py +++ b/families/personaplex/default_dual_profile_decoder.py @@ -527,12 +527,10 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False): else: norm2 = normed else: - residual1 = network.add_elementwise( - hidden_state, attn_out, trt.ElementWiseOperation.SUM - ) + residual1 = graph_ops.add_residual_sum(network, hidden_state, attn_out) norm2 = _norm_multi( network, - residual1.get_output(0), + residual1, hidden, weights[f"{prefix}.post_attn_norm"], weights.get(f"{prefix}.post_attn_norm_beta"), @@ -567,15 +565,11 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False): # Final residual. if parallel_residual: - sum_attn = network.add_elementwise(hidden_state, attn_out, trt.ElementWiseOperation.SUM) - residual2 = network.add_elementwise( - sum_attn.get_output(0), mlp_out, trt.ElementWiseOperation.SUM - ) + sum_attn = graph_ops.add_residual_sum(network, hidden_state, attn_out) + residual2 = graph_ops.add_residual_sum(network, sum_attn, mlp_out) else: - residual2 = network.add_elementwise( - residual1.get_output(0), mlp_out, trt.ElementWiseOperation.SUM - ) - hidden_state = residual2.get_output(0) + residual2 = graph_ops.add_residual_sum(network, residual1, mlp_out) + hidden_state = residual2 # ---- Final norm + LM head ------------------------------------------- final_norm = weights.get("final_norm") diff --git a/families/personaplex/graph_ops.py b/families/personaplex/graph_ops.py index d2af56e069..f5306432b4 100644 --- a/families/personaplex/graph_ops.py +++ b/families/personaplex/graph_ops.py @@ -23,6 +23,21 @@ def _cast_back_to_trt_dtype( return network.add_cast(tensor, target_dtype).get_output(0) +def add_residual_sum( + network: trt.INetworkDefinition, residual: trt.ITensor, update: trt.ITensor +) -> trt.ITensor: + """Add already rounded branch outputs and restore the residual dtype. + + Summing in FP32 keeps the producer's Half rounding boundary distinct from + the residual add. Otherwise GEMM fusion can add the unrounded accumulator. + """ + output_dtype = residual.dtype + residual = _cast_back_to_trt_dtype(network, residual, trt.float32) + update = _cast_back_to_trt_dtype(network, update, trt.float32) + total = network.add_elementwise(residual, update, trt.ElementWiseOperation.SUM).get_output(0) + return _cast_back_to_trt_dtype(network, total, output_dtype) + + def _add_matrix_multiply_with_fp32_accumulation( network: trt.INetworkDefinition, lhs: trt.ITensor, diff --git a/families/personaplex/tests/test_residual_precision.py b/families/personaplex/tests/test_residual_precision.py new file mode 100644 index 0000000000..c510fb0e23 --- /dev/null +++ b/families/personaplex/tests/test_residual_precision.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Keep the linear-output rounding boundary inside the actual decoder layer.""" + +import numpy as np +import pytest + +trt = pytest.importorskip("tensorrt") +torch = pytest.importorskip("torch") + +from families.personaplex import graph_ops # noqa: E402 +from families.personaplex.default_decoder import _add_decoder_layer # noqa: E402 + + +@pytest.mark.gpu +@pytest.mark.trt +@pytest.mark.parametrize("rebuild", range(3)) +@pytest.mark.parametrize("precision", ["fp16", "fp32"]) +def test_decoder_residual_rounds_the_projection_first(precision, rebuild): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + del rebuild + width, mlp_size, heads, cache_length = 1024, 2816, 16, 18 + np_dtype = np.float16 if precision == "fp16" else np.float32 + trt_dtype = trt.float16 if precision == "fp16" else trt.float32 + torch_dtype = torch.float16 if precision == "fp16" else torch.float32 + rng = np.random.default_rng(42) + + def dense(): + return ( + rng.integers(-128, 129, size=(width, width), dtype=np.int16).astype(np.float32) / 1024 + ) + + w_v, w_o = dense(), dense() + w_v[:, :2] = 0 + w_v[0, 0], w_v[1, 1] = 1.0, 1.0 + w_o[:, 0] = 0 + w_o[0, 0], w_o[1, 0] = 1.0, 2**-12 + alpha = np.ones(width, np.float32) + alpha[0] = -1.0 + weights = { + "L.input_norm": alpha, + "L.post_attn_norm": np.ones(width, np.float32), + "L.w_q": np.zeros((width, width), np.float32), + "L.w_k": np.zeros((width, width), np.float32), + "L.w_v": w_v, + "L.w_o": w_o, + "L.w_gate": np.zeros((width, mlp_size), np.float32), + "L.w_up": np.zeros((width, mlp_size), np.float32), + "L.w_down": np.zeros((mlp_size, width), np.float32), + } + logger = trt.Logger(trt.Logger.ERROR) + builder = trt.Builder(logger) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + hidden = network.add_input("hidden", trt_dtype, (1, width)) + cache_k = network.add_input("cache_k", trt_dtype, (cache_length, width)) + cache_v = network.add_input("cache_v", trt_dtype, (cache_length, width)) + mask = network.add_input("mask", trt.float32, (1, cache_length + 1)) + if mask.dtype != trt_dtype: + mask = network.add_cast(mask, trt_dtype).get_output(0) + position = graph_ops.add_constant(network, (1,), np.zeros(1, np.int32), dtype=np.int32) + epsilon = graph_ops.add_constant(network, (1, 1), np.array([[1e-8]], np.float32)) + layer = _add_decoder_layer( + network=network, + hidden=hidden, + cache_k=cache_k, + cache_v=cache_v, + attention_mask=mask, + position_id=position, + attention_scale=1 / 8, + eps_tensor=epsilon, + eps=1e-8, + weights=weights, + prefix="L", + hidden_size=width, + attention_size=width, + kv_attention_size=width, + mlp_size=mlp_size, + num_heads=heads, + num_kv_heads=heads, + head_dim=64, + max_cache_length=cache_length, + position_type="learned", + dtype=np_dtype, + ) + for name in ("hidden", "present_k", "present_v"): + tensor = layer[name] + tensor.name = "output_" + name + network.mark_output(tensor) + config = builder.create_builder_config() + config.clear_flag(trt.BuilderFlag.TF32) + plan = builder.build_serialized_network(network, config) + assert plan is not None + runtime = trt.Runtime(logger) + engine = runtime.deserialize_cuda_engine(plan) + assert engine is not None + execution = engine.create_execution_context() + assert execution is not None + feeds = { + "hidden": torch.ones((1, width), dtype=torch_dtype, device="cuda"), + "cache_k": torch.zeros((cache_length, width), dtype=torch_dtype, device="cuda"), + "cache_v": torch.zeros((cache_length, width), dtype=torch_dtype, device="cuda"), + "mask": torch.full((1, cache_length + 1), -10000.0, device="cuda"), + } + feeds["hidden"][0, 0] = -1.0 + feeds["mask"][0, -1] = 0.0 + outputs = {} + for index in range(engine.num_io_tensors): + name = engine.get_tensor_name(index) + if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT: + tensor = feeds[name] + else: + tensor = torch.empty( + tuple(engine.get_tensor_shape(name)), dtype=torch_dtype, device="cuda" + ) + outputs[name] = tensor + assert execution.set_tensor_address(name, tensor.data_ptr()) + assert execution.execute_async_v3(torch.cuda.current_stream().cuda_stream) + torch.cuda.synchronize() + # RMSNorm gives an all-one vector; one valid key makes context exactly V. + # Native rounds linear(V) to the model dtype before adding the residual. + normalized = torch.ones((1, width), dtype=torch_dtype, device="cuda") + native_v = torch.nn.functional.linear( + normalized, torch.from_numpy(w_v.T.copy()).to("cuda", torch_dtype) + ) + native_projection = torch.nn.functional.linear( + native_v, torch.from_numpy(w_o.T.copy()).to("cuda", torch_dtype) + ) + expected = (native_projection + feeds["hidden"])[0, 0].cpu() + actual = outputs["output_hidden"][0, 0].cpu() + actual_dtype = outputs["output_hidden"].dtype + del execution, engine, runtime + assert float(expected) == (0.0 if precision == "fp16" else 2**-12) + assert actual_dtype == torch_dtype + torch.testing.assert_close(actual, expected, rtol=0, atol=0)