From f179f32c6db3cdda6bc32d5cf90ee83d3ad8ffcf Mon Sep 17 00:00:00 2001 From: JiaxinD Date: Wed, 21 Oct 2026 09:32:03 -0700 Subject: [PATCH] fix(llama): preserve BF16 checkpoint value range Stage BF16 weights and decoder constants in FP32 to avoid FP16 overflow and underflow before the runtime cast. Select normalization compute casts from the runtime tensor dtype so FP32 reductions and BF16 outputs remain explicit. Cover real BF16 safetensors, both builder prefixes, and normalization dtype boundaries. BF16 host weight arrays now use four bytes per element instead of two. Signed-off-by: JiaxinD --- families/llama/checkpoint_mapper.py | 2 +- .../llama/dual_profile_decoder_builder.py | 22 +- families/llama/graph_ops.py | 16 +- families/llama/standard_decoder_builder.py | 5 +- families/llama/tests/test_bf16_storage.py | 242 ++++++++++++++++++ 5 files changed, 256 insertions(+), 31 deletions(-) create mode 100644 families/llama/tests/test_bf16_storage.py diff --git a/families/llama/checkpoint_mapper.py b/families/llama/checkpoint_mapper.py index c931921d75..48d18cbf90 100644 --- a/families/llama/checkpoint_mapper.py +++ b/families/llama/checkpoint_mapper.py @@ -26,7 +26,7 @@ def _target_np_dtype(precision: str) -> np.dtype: """Map precision string to numpy dtype for weight storage.""" - if precision in ("fp16", "bf16"): + if precision == "fp16": return np.float16 return np.float32 diff --git a/families/llama/dual_profile_decoder_builder.py b/families/llama/dual_profile_decoder_builder.py index 1335e2ad1c..8082952843 100644 --- a/families/llama/dual_profile_decoder_builder.py +++ b/families/llama/dual_profile_decoder_builder.py @@ -70,16 +70,7 @@ def _const_in_work_dtype( work_np_dtype: np.dtype, work_trt_dtype: trt.DataType, ) -> trt.ITensor: - """Create a constant in work_np_dtype storage and cast it to work_trt_dtype. - - Needed for bf16 builds: the dual-profile builder stores bf16 weights - on disk as fp16 (work_np_dtype = np.float16), but the runtime tensor - must be bfloat16 to match the rest of the graph. ``add_constant`` - alone produces an fp16 constant — we need an explicit cast to - bfloat16 so layers like IRotaryEmbeddingLayer (which require all - inputs to share a dtype) accept it. fp16 / fp32 builds are no-ops - because work_np_dtype maps directly to work_trt_dtype. - """ + """Create a constant in storage dtype and cast it to runtime dtype.""" const = graph_ops.add_constant(network, shape, values, dtype=work_np_dtype) if const.dtype != work_trt_dtype: const = network.add_cast(const, work_trt_dtype).get_output(0) @@ -309,7 +300,7 @@ def build_dual_profile_decoder_engine( if precision == "fp16": work_np_dtype, work_trt_dtype = np.float16, trt.float16 elif precision == "bf16": - work_np_dtype, work_trt_dtype = np.float16, trt.bfloat16 + work_np_dtype, work_trt_dtype = np.float32, trt.bfloat16 else: work_np_dtype, work_trt_dtype = np.float32, trt.float32 @@ -427,14 +418,11 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False): interleaved=interleaved_rope, rope_scaling=config.raw.get("rope_scaling"), ) - # BF16 must round directly from the FP32 indexed table. Routing - # through FP16 storage would introduce FP16 -> BF16 double rounding. - rope_np_dtype = np.float32 if work_trt_dtype == trt.bfloat16 else work_np_dtype cos_half_table = _const_in_work_dtype( - network, cos_half_np.shape, cos_half_np, rope_np_dtype, work_trt_dtype + network, cos_half_np.shape, cos_half_np, work_np_dtype, work_trt_dtype ) sin_half_table = _const_in_work_dtype( - network, sin_half_np.shape, sin_half_np, rope_np_dtype, work_trt_dtype + network, sin_half_np.shape, sin_half_np, work_np_dtype, work_trt_dtype ) # Learned position embedding (GPT-2 / OPT / GPT-Neo / XGLM). @@ -490,8 +478,6 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False): ) hidden_state = pos_add.get_output(0) - # Make sure the main hidden stream is in the requested runtime dtype - # before entering the layer stack (BF16 mode stores fp16 constants). if hidden_state.dtype != work_trt_dtype: hidden_state = network.add_cast(hidden_state, work_trt_dtype).get_output(0) diff --git a/families/llama/graph_ops.py b/families/llama/graph_ops.py index f0adeaa3aa..243e38d8fc 100644 --- a/families/llama/graph_ops.py +++ b/families/llama/graph_ops.py @@ -125,13 +125,13 @@ def add_rms_norm( ) -> trt.ITensor: """RMSNorm: gamma * (x / sqrt(mean(x^2) + eps)). - FP32 precision boundary: when dtype != float32, casts to FP32 before + FP32 precision boundary: when input is not float32, casts to FP32 before norm computation for numerical stability, then casts back. TRT's native normalization API implements mean-centered LayerNorm, not RMSNorm, so this remains a manual shared implementation. """ - need_cast = (dtype != np.float32) + need_cast = (inp.dtype != trt.float32) output_dtype = inp.dtype if need_cast: inp = network.add_cast(inp, trt.float32).get_output(0) @@ -166,12 +166,12 @@ def add_rms_norm_per_head( ) -> trt.ITensor: """Per-head RMSNorm for [Sq, num_heads * head_dim] tensors. - FP32 precision boundary: when dtype != float32, casts to FP32 before + FP32 precision boundary: when input is not float32, casts to FP32 before norm computation for numerical stability, then casts back. ``sequence_length=None`` means runtime-dynamic Sq. ``gamma`` may be [num_heads * head_dim] or [head_dim] broadcast to heads. """ - need_cast = (dtype != np.float32) + need_cast = (inp.dtype != trt.float32) output_dtype = inp.dtype seq_dim = -1 if sequence_length is None else sequence_length reshape_in = network.add_shuffle(inp) @@ -223,10 +223,10 @@ def add_layer_norm( ) -> trt.ITensor: """LayerNorm: gamma * ((x - mean) / sqrt(var + eps)) + beta. - FP32 precision boundary: when dtype != float32, casts to FP32 before + FP32 precision boundary: when input is not float32, casts to FP32 before norm computation for numerical stability, then casts back. """ - need_cast = (dtype != np.float32) + need_cast = (inp.dtype != trt.float32) output_dtype = inp.dtype if need_cast: inp = network.add_cast(inp, trt.float32).get_output(0) @@ -273,8 +273,8 @@ def add_gelu_new( """GELU (tanh approximation): 0.5*x*(1+tanh(sqrt(2/pi)*(x+0.044715*x^3))). Constants are cast to ``inp.dtype`` so the elementwise ops are valid in - a STRONGLY_TYPED network when ``inp`` is bf16 (storage np_dtype is - fp16, runtime trt_dtype is bfloat16) or any other non-matching combo. + a STRONGLY_TYPED network when ``inp`` is bf16 or any other + non-matching combo. """ target_dtype = inp.dtype const_shape = (1,) * max(1, len(tuple(inp.shape))) diff --git a/families/llama/standard_decoder_builder.py b/families/llama/standard_decoder_builder.py index f94aca0e74..9264607879 100644 --- a/families/llama/standard_decoder_builder.py +++ b/families/llama/standard_decoder_builder.py @@ -181,7 +181,7 @@ def build_standard_decoder_engine( work_np_dtype = np.float16 work_trt_dtype = trt.float16 elif precision == "bf16": - work_np_dtype = np.float16 # stored as float16, TRT uses bfloat16 + work_np_dtype = np.float32 work_trt_dtype = trt.bfloat16 else: work_np_dtype = np.float32 @@ -341,9 +341,6 @@ def _cast_work_dtype(tensor: trt.ITensor) -> trt.ITensor: ) hidden_state = pos_add.get_output(0) - # In BF16 mode many embedding/position constants are still materialized from - # float16 storage. Normalize the decoder's main hidden stream back to the - # requested runtime dtype before entering the layer stack. if hidden_state.dtype != work_trt_dtype: hidden_state = network.add_cast(hidden_state, work_trt_dtype).get_output(0) diff --git a/families/llama/tests/test_bf16_storage.py b/families/llama/tests/test_bf16_storage.py new file mode 100644 index 0000000000..f8954ea80c --- /dev/null +++ b/families/llama/tests/test_bf16_storage.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU checkpoint and typed graph-prefix checks, without TensorRT execution.""" + +import importlib +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace + +import ml_dtypes +import numpy as np +import pytest +from safetensors.numpy import save_file + +from ..checkpoint_mapper import load_standard_weights +from ..config import ModelConfig + + +class _Tensor: + def __init__(self, values, dtype="fp32", shape=None): + storage = {"fp32": np.float32, "fp16": np.float16, "bf16": ml_dtypes.bfloat16} + with np.errstate(over="ignore", under="ignore", invalid="ignore"): + self.values = np.asarray(values).astype(storage.get(dtype, np.int32)).astype(np.float32) + self.dtype = dtype + self.shape = self.values.shape if shape is None else shape + + +def _layer(tensor): + return SimpleNamespace(get_output=lambda index: tensor) + + +class _PrefixComplete(Exception): + pass + + +class _Network: + def __init__(self, stop_at_epsilon=False): + self.constants = [] + self.inputs = {} + self.stop_at_epsilon = stop_at_epsilon + + def add_constant(self, shape, values): + values = values.reshape(shape) + self.constants.append(values.copy()) + if self.stop_at_epsilon and shape == (1, 1): + raise _PrefixComplete + dtype = "fp16" if values.dtype == np.float16 else "fp32" + return _layer(_Tensor(values, dtype)) + + def add_input(self, name, dtype, shape): + tensor = _Tensor(np.zeros(tuple(max(1, n) for n in shape)), dtype, shape) + self.inputs[name] = tensor + return tensor + + def add_cast(self, tensor, dtype): + return _layer(_Tensor(tensor.values, dtype, tensor.shape)) + + def add_elementwise(self, lhs, rhs, operation): + assert lhs.dtype == rhs.dtype, "mixed elementwise dtypes" + function = {"prod": np.multiply, "sum": np.add, "sub": np.subtract}[operation] + return _layer(_Tensor(function(lhs.values, rhs.values), lhs.dtype)) + + def add_reduce(self, tensor, operation, axes, keep_dims): + assert operation == "avg" + dimensions = tuple(i for i in range(tensor.values.ndim) if axes & (1 << i)) + return _layer( + _Tensor(tensor.values.mean(axis=dimensions, keepdims=keep_dims), tensor.dtype) + ) + + def add_unary(self, tensor, operation): + function = {"sqrt": np.sqrt, "recip": np.reciprocal}[operation] + return _layer(_Tensor(function(tensor.values), tensor.dtype)) + + def add_shuffle(self, tensor): + shuffle = SimpleNamespace(reshape_dims=tensor.shape) + shuffle.get_output = lambda index: _Tensor( + tensor.values.reshape(shuffle.reshape_dims), tensor.dtype + ) + return shuffle + + +@pytest.fixture +def family(monkeypatch): + network = _Network() + config = SimpleNamespace(add_optimization_profile=lambda profile: None) + builder = SimpleNamespace( + create_network=lambda flags: network, + create_builder_config=lambda: config, + create_optimization_profile=lambda: SimpleNamespace(set_shape=lambda *args: None), + ) + logger = type("Logger", (), {"VERBOSE": 1, "WARNING": 2, "__init__": lambda self, level: None}) + trt = SimpleNamespace( + float16="fp16", + float32="fp32", + bfloat16="bf16", + int32="int32", + Weights=lambda values: values, + Logger=logger, + Builder=lambda logger: builder, + NetworkDefinitionCreationFlag=SimpleNamespace(STRONGLY_TYPED=1), + ElementWiseOperation=SimpleNamespace(PROD="prod", SUM="sum", SUB="sub"), + ReduceOperation=SimpleNamespace(AVG="avg"), + UnaryOperation=SimpleNamespace(SQRT="sqrt", RECIP="recip"), + ) + name = "_llama_bf16_storage_test" + package = ModuleType(name) + package.__path__ = [str(Path(__file__).resolve().parents[1])] + monkeypatch.setitem(sys.modules, "tensorrt", trt) + monkeypatch.setitem(sys.modules, name, package) + try: + yield SimpleNamespace( + graph=importlib.import_module(name + ".graph_ops"), + standard=importlib.import_module(name + ".standard_decoder_builder"), + dual=importlib.import_module(name + ".dual_profile_decoder_builder"), + network=network, + ) + finally: + for key in tuple(sys.modules): + if key.startswith(name + "."): + del sys.modules[key] + + +def _checkpoint(tmp_path, tied, *, wide=True): + config = ModelConfig( + model_type="llama", + hidden_size=4, + vocab_size=4, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + intermediate_size=4, + rms_norm_eps=2**-30, + raw={"_decoder_engine_role": "decode"}, + ) + values = np.array( + [2**20, -(2**20), 2**-30, -(2**-30)] if wide else [1024, -1024, 0.125, -0.125], + dtype=ml_dtypes.bfloat16, + ) + matrix = np.tile(values, (4, 1)) + tensors = {"model.embed_tokens.weight": matrix} + if not tied: + tensors["lm_head.weight"] = matrix + for norm in ["input_layernorm", "post_attention_layernorm"]: + tensors[f"model.layers.0.{norm}.weight"] = np.ones(4, dtype=ml_dtypes.bfloat16) + for projection in [ + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + ]: + tensors[f"model.layers.0.{projection}.weight"] = matrix + save_file(tensors, tmp_path / "model.safetensors") + return config, matrix.astype(np.float32) + + +@pytest.mark.parametrize("precision", ["bf16", "fp16", "fp32"]) +@pytest.mark.parametrize("tied", [False, True]) +def test_checkpoint_preserves_requested_storage_range(tmp_path, precision, tied): + config, source = _checkpoint(tmp_path, tied, wide=precision != "fp16") + dtype = np.float16 if precision == "fp16" else np.float32 + with np.errstate(over="ignore", under="ignore"): + weights = load_standard_weights(tmp_path, config, precision=precision) + expected = source.astype(dtype) + for key in [ + "embedding", + "w_out", + "layer.0.w_q", + "layer.0.w_k", + "layer.0.w_v", + "layer.0.w_o", + "layer.0.w_gate", + "layer.0.w_up", + "layer.0.w_down", + ]: + assert weights[key].dtype == dtype + np.testing.assert_array_equal(weights[key], expected if key == "embedding" else expected.T) + + +@pytest.mark.parametrize("precision", ["bf16", "fp16", "fp32"]) +@pytest.mark.parametrize("builder", ["standard", "dual"]) +def test_decoder_prefix_preserves_constant_storage(family, tmp_path, precision, builder): + config, values = _checkpoint(tmp_path, True, wide=precision != "fp16") + weights = { + "embedding": values, + "position_embedding": values, + "final_norm": np.ones(4), + "w_out": values.T, + } + family.network.stop_at_epsilon = True + module = getattr(family, builder) + function = getattr( + module, f"build_{'standard' if builder == 'standard' else 'dual_profile'}_decoder_engine" + ) + kwargs = {} if builder == "standard" else {"native_kv_cache": False, "profile_mode": "decode"} + with np.errstate(over="ignore", under="ignore"): + with pytest.raises(_PrefixComplete): + function(config, weights, 4, precision=precision, position_type="learned", **kwargs) + dtype = np.float16 if precision == "fp16" else np.float32 + expected = values.astype(dtype) + for constant in family.network.constants[:2]: + assert constant.dtype == dtype + np.testing.assert_array_equal(constant, expected) + assert ( + family.network.inputs["cache_k_0"].dtype + == {"bf16": "bf16", "fp16": "fp16", "fp32": "fp32"}[precision] + ) + if precision == "bf16": + np.testing.assert_array_equal(family.network.constants[-1], [[2**-30]]) + + +@pytest.mark.parametrize( + "dtype,storage", [("bf16", np.float32), ("fp16", np.float16), ("fp32", np.float32)] +) +@pytest.mark.parametrize("norm", ["rms", "head", "layer"]) +def test_norm_uses_fp32_compute_with_independent_storage(family, dtype, storage, norm): + inputs = _Tensor([[1.0, 2.0, 3.0, 4.0]], dtype) + epsilon = _Tensor([[1e-5]], "fp32") + gamma = np.ones(4, dtype=np.float32) + kwargs = {"dtype": storage} + if norm == "rms": + result = family.graph.add_rms_norm(family.network, inputs, 4, gamma, epsilon, **kwargs) + expected = inputs.values / np.sqrt(np.mean(inputs.values**2) + 1e-5) + elif norm == "head": + result = family.graph.add_rms_norm_per_head( + family.network, inputs, 2, 2, gamma, epsilon, **kwargs + ) + reshaped = inputs.values.reshape(1, 2, 2) + expected = (reshaped / np.sqrt(np.mean(reshaped**2, axis=2, keepdims=True) + 1e-5)).reshape( + 1, 4 + ) + else: + result = family.graph.add_layer_norm( + family.network, inputs, 4, gamma, np.zeros(4), epsilon, **kwargs + ) + centered = inputs.values - inputs.values.mean(axis=1, keepdims=True) + expected = centered / np.sqrt(np.mean(centered**2) + 1e-5) + assert result.dtype == dtype + np.testing.assert_allclose(result.values, _Tensor(expected, dtype).values, rtol=1e-6, atol=1e-7)