From ecab800ac14c8a01b50849775855a7d812f48a78 Mon Sep 17 00:00:00 2001 From: chaofengw Date: Mon, 21 Sep 2026 11:06:00 +0000 Subject: [PATCH] fix: stabilize Community GPU family validation Signed-off-by: chaofengw --- families/convbert/model.py | 21 ++++ .../convbert/tests/test_tokenizer_contract.py | 40 +++++++ families/deberta/model.py | 21 ++++ .../deberta/tests/test_tokenizer_contract.py | 40 +++++++ families/deepseek_ocr/model.py | 26 ++++- families/deepseek_ocr/tests/test_build.py | 77 ++++++++++++- families/deepseek_v2/model.py | 103 ++++++++++-------- .../deepseek_v2/tests/test_router_contract.py | 102 +++++++++++++++++ families/dinov3/requirements.txt | 4 +- families/glm/model.py | 16 +-- families/glm/tests/test_weight_memory.py | 69 ++++++++++++ families/gpt_oss/model.py | 98 ++++++++++------- families/gpt_oss/tests/test_weight_memory.py | 95 ++++++++++++++++ tools/ci/e2e.py | 34 +++++- tools/tests/test_new_ci.py | 23 +++- 15 files changed, 656 insertions(+), 113 deletions(-) create mode 100644 families/convbert/tests/test_tokenizer_contract.py create mode 100644 families/deberta/tests/test_tokenizer_contract.py create mode 100644 families/glm/tests/test_weight_memory.py create mode 100644 families/gpt_oss/tests/test_weight_memory.py diff --git a/families/convbert/model.py b/families/convbert/model.py index ed199d342e..be93ce5ea2 100644 --- a/families/convbert/model.py +++ b/families/convbert/model.py @@ -262,6 +262,8 @@ def _positive_int(value: object, name: str) -> int: def _tokenizer_runtime_contract(model_dir: Path) -> dict[str, object]: """Resolve this family's exact native-tokenizer framing.""" + import tempfile + from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained( @@ -269,6 +271,25 @@ def _tokenizer_runtime_contract(model_dir: Path) -> dict[str, object]: trust_remote_code=True, use_fast=True, ) + tokenizer_path = model_dir / "tokenizer.json" + if not tokenizer_path.is_file(): + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=model_dir, + prefix=".trtmc-convbert-tokenizer-", + suffix=".json", + delete=False, + ) as output: + temporary_path = Path(output.name) + tokenizer.backend_tokenizer.save(str(temporary_path)) + if not temporary_path.is_file() or temporary_path.stat().st_size == 0: + raise RuntimeError("tokenizer conversion did not create tokenizer.json") + temporary_path.replace(tokenizer_path) + temporary_path = None + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) default_ids = list(tokenizer.encode("hello")) plain_ids = list(tokenizer.encode("hello", add_special_tokens=False)) if default_ids == plain_ids: diff --git a/families/convbert/tests/test_tokenizer_contract.py b/families/convbert/tests/test_tokenizer_contract.py new file mode 100644 index 0000000000..6cedf8b767 --- /dev/null +++ b/families/convbert/tests/test_tokenizer_contract.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""ConvBERT tokenizer bundle regression coverage.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("tensorrt", reason="TensorRT is required for family builder tests") + +from families.convbert import model + + +def test_runtime_contract_materializes_missing_fast_tokenizer( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class Backend: + def save(self, path: str) -> None: + Path(path).write_text('{"model":{"type":"WordPiece"}}', encoding="utf-8") + + class Tokenizer: + backend_tokenizer = Backend() + + @staticmethod + def encode(_text: str, add_special_tokens: bool = True) -> list[int]: + return [101, 42, 102] if add_special_tokens else [42] + + auto = SimpleNamespace(from_pretrained=lambda *_args, **_kwargs: Tokenizer()) + monkeypatch.setitem(sys.modules, "transformers", SimpleNamespace(AutoTokenizer=auto)) + + contract = model._tokenizer_runtime_contract(tmp_path) + + assert (tmp_path / "tokenizer.json").is_file() + assert contract["tokenizer_prefix_ids"] == [101] + assert contract["tokenizer_suffix_ids"] == [102] diff --git a/families/deberta/model.py b/families/deberta/model.py index 0e04c33f60..5ee6863672 100644 --- a/families/deberta/model.py +++ b/families/deberta/model.py @@ -629,6 +629,8 @@ def _positive_int(value: object, name: str) -> int: def _tokenizer_runtime_contract(model_dir: Path) -> dict[str, object]: """Resolve this family's exact native-tokenizer framing.""" + import tempfile + from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained( @@ -636,6 +638,25 @@ def _tokenizer_runtime_contract(model_dir: Path) -> dict[str, object]: trust_remote_code=True, use_fast=True, ) + tokenizer_path = model_dir / "tokenizer.json" + if not tokenizer_path.is_file(): + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + dir=model_dir, + prefix=".trtmc-deberta-tokenizer-", + suffix=".json", + delete=False, + ) as output: + temporary_path = Path(output.name) + tokenizer.backend_tokenizer.save(str(temporary_path)) + if not temporary_path.is_file() or temporary_path.stat().st_size == 0: + raise RuntimeError("tokenizer conversion did not create tokenizer.json") + temporary_path.replace(tokenizer_path) + temporary_path = None + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) default_ids = list(tokenizer.encode("hello")) plain_ids = list(tokenizer.encode("hello", add_special_tokens=False)) if default_ids == plain_ids: diff --git a/families/deberta/tests/test_tokenizer_contract.py b/families/deberta/tests/test_tokenizer_contract.py new file mode 100644 index 0000000000..6b84dc5fea --- /dev/null +++ b/families/deberta/tests/test_tokenizer_contract.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""DeBERTa tokenizer bundle regression coverage.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("tensorrt", reason="TensorRT is required for family builder tests") + +from families.deberta import model + + +def test_runtime_contract_materializes_missing_fast_tokenizer( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + class Backend: + def save(self, path: str) -> None: + Path(path).write_text('{"model":{"type":"BPE"}}', encoding="utf-8") + + class Tokenizer: + backend_tokenizer = Backend() + + @staticmethod + def encode(_text: str, add_special_tokens: bool = True) -> list[int]: + return [1, 42, 2] if add_special_tokens else [42] + + auto = SimpleNamespace(from_pretrained=lambda *_args, **_kwargs: Tokenizer()) + monkeypatch.setitem(sys.modules, "transformers", SimpleNamespace(AutoTokenizer=auto)) + + contract = model._tokenizer_runtime_contract(tmp_path) + + assert (tmp_path / "tokenizer.json").is_file() + assert contract["tokenizer_prefix_ids"] == [1] + assert contract["tokenizer_suffix_ids"] == [2] diff --git a/families/deepseek_ocr/model.py b/families/deepseek_ocr/model.py index 7a5af80f20..a448f4f38b 100644 --- a/families/deepseek_ocr/model.py +++ b/families/deepseek_ocr/model.py @@ -37,6 +37,7 @@ from typing import TYPE_CHECKING +import gc import sys from pathlib import Path @@ -1343,6 +1344,8 @@ def _build_sam_attention( def _load_vision_weights( model_dir: str, + *, + precision: str, ) -> dict[str, np.ndarray]: """Load all vision pipeline weights from safetensors.""" readers = _open_safetensors(Path(model_dir)) @@ -1474,7 +1477,13 @@ def _load_vision_weights( # --- View separator --- vw["view_sep"] = _load_tensor(readers, "model.view_seperator").astype(np.float32) - return vw + if precision == "fp16": + return {name: np.ascontiguousarray(value, dtype=np.float16) for name, value in vw.items()} + if precision == "fp32": + return vw + raise ValueError( + f"Unsupported DeepSeek-OCR vision precision {precision!r}; expected fp32 or fp16" + ) def _build_deepseek_ocr_vision_engine( @@ -1493,7 +1502,6 @@ def _build_deepseek_ocr_vision_engine( """ print("[trtmc build] Building DeepSeek-OCR-2 vision engine (native TRT) ...", file=sys.stderr) - vw = _load_vision_weights(model_dir) if precision == "fp16": work_np_dtype, work_trt_dtype = np.float16, trt.float16 elif precision == "fp32": @@ -1502,6 +1510,7 @@ def _build_deepseek_ocr_vision_engine( raise ValueError( f"Unsupported DeepSeek-OCR vision precision {precision!r}; expected fp32 or fp16" ) + vw = _load_vision_weights(model_dir, precision=precision) # SAM config sam_hidden = 768 @@ -2042,6 +2051,9 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: verbose=request.verbose, parallel_config=parallel, ) + writer.add_bytes("prefill.plan", prefill) + del prefill + gc.collect() config.raw["_decoder_engine_role"] = "decode" decode = model.build_engine( config, @@ -2052,9 +2064,15 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: verbose=request.verbose, parallel_config=parallel, ) - config.raw.pop("_decoder_engine_role", None) writer.add_bytes("engine.plan", decode) - writer.add_bytes("prefill.plan", prefill) + del decode + gc.collect() + config.raw.pop("_decoder_engine_role", None) + # Vision weights are an independent checkpoint subset. Release the decoder + # arrays before loading it so three full model representations do not + # overlap during a single-device bundle build. + weights.clear() + gc.collect() vision = model.build_vision_engine( str(model_dir), config, weights, precision=precision, verbose=request.verbose ) diff --git a/families/deepseek_ocr/tests/test_build.py b/families/deepseek_ocr/tests/test_build.py index b63fd4ca83..68339b3517 100644 --- a/families/deepseek_ocr/tests/test_build.py +++ b/families/deepseek_ocr/tests/test_build.py @@ -75,7 +75,7 @@ def build_engine(self, loaded_config, weights, max_length, **options): def build_vision_engine(self, model_dir, loaded_config, weights, **options): assert model_dir == str(tmp_path) assert loaded_config is config - assert weights == {"weights": True} + assert weights == {} assert options == {"precision": "fp32", "verbose": False} writer.events.append("build:vision") vision_calls.append((model_dir, loaded_config, weights, options)) @@ -137,3 +137,78 @@ def get_vl_config(self, loaded_config): ] assert writer.json["runtime.json"]["tensor_parallel_size"] == 2 assert "_decoder_engine_role" not in config.raw + + +def test_single_device_build_streams_each_decoder_plan_before_the_next_build( + monkeypatch, tmp_path +) -> None: + config = SimpleNamespace( + model_type="deepseek_vl_v2", + max_position_embeddings=4096, + num_hidden_layers=2, + vocab_size=32, + bos_token_id=1, + eos_token_id=2, + hidden_size=8, + raw={}, + ) + writer = _Writer() + + class FakeModel: + def load_weights(self, *_args, **_kwargs): + return {"decoder": object()} + + def build_engine(self, _config, weights, _max_length, **_options): + role = config.raw["_decoder_engine_role"] + assert weights + writer.events.append(f"build:{role}") + if role == "decode": + assert "write:prefill.plan" in writer.events + return role.encode() + + def build_vision_engine(self, _model_dir, _config, weights, **_options): + assert weights == {} + writer.events.append("build:vision") + return b"vision" + + def get_vl_config(self, _config): + return { + "image_token_id": 7, + "vision_output_dim": 8, + "prefill_max_length": 64, + } + + monkeypatch.setattr( + model_module, "ModelConfig", SimpleNamespace(from_dir=lambda _model_dir: config) + ) + monkeypatch.setattr(model_module, "_DeepseekOcrModel", FakeModel) + request = SimpleNamespace( + model_dir=tmp_path, + backend="trt", + dynamic_kv_cache=False, + task="vision_language_generation", + precision="fp16", + max_sequence_length=4096, + tensor_parallel_size=1, + context_parallel_size=1, + quantization=None, + fp32_layers=(1, 2), + image_height=None, + image_width=None, + video_num_frames=None, + max_batch_size=1, + verbose=False, + ) + + model_module.build(request, writer) + + assert writer.events == [ + "write:header", + "build:prefill", + "write:prefill.plan", + "build:decode", + "write:engine.plan", + "build:vision", + "write:vision.plan", + ] + assert "_decoder_engine_role" not in config.raw diff --git a/families/deepseek_v2/model.py b/families/deepseek_v2/model.py index 55f8d3505b..51aee0cfb1 100644 --- a/families/deepseek_v2/model.py +++ b/families/deepseek_v2/model.py @@ -948,7 +948,7 @@ def _stack_expert_weights( ) -def _add_native_routed_experts( +def _add_routed_experts( network: trt.INetworkDefinition, inp: trt.ITensor, weights: WeightDict, @@ -960,56 +960,63 @@ def _add_native_routed_experts( scaled_weights: trt.ITensor, dtype: np.dtype, ) -> trt.ITensor: - """Routed-expert output computed by the native TensorRT MoE layer. - - ``set_gated_weights`` takes one stacked tensor per projection covering every - expert, in the orientation this family already stores them: gate and up as - ``[experts, hidden, intermediate]`` and down as - ``[experts, intermediate, hidden]``. No transposition is required. - - The layer applies the routing scores itself, so its output is the weighted - sum over the selected experts and only those experts are evaluated. - """ + """Compute only the selected experts with portable TensorRT graph layers.""" w_gate = _stack_expert_weights(weights, prefix, n_routed_experts, "w_gate", dtype) w_up = _stack_expert_weights(weights, prefix, n_routed_experts, "w_up", dtype) w_down = _stack_expert_weights(weights, prefix, n_routed_experts, "w_down", dtype) - # IMoELayer requires rank-3 hidden states [batch, tokens, hidden]; the - # decoder carries rank-2 [tokens, hidden]. - def _with_batch_dim(tensor: trt.ITensor, last_dim: int) -> trt.ITensor: - shuffle = network.add_shuffle(tensor) - shuffle.reshape_dims = (1, -1, last_dim) - return shuffle.get_output(0) - rank = len(tuple(inp.shape)) - if rank == 2: - moe_hidden = _with_batch_dim(inp, hidden_size) - moe_indices = _with_batch_dim(top_indices, num_experts_per_tok) - moe_scores = _with_batch_dim(scaled_weights, num_experts_per_tok) - elif rank == 3: - moe_hidden, moe_indices, moe_scores = inp, top_indices, scaled_weights - else: - raise ValueError(f"DeepSeek-V2 MoE expects rank-2 or rank-3 hidden states, got rank {rank}") - - moe = network.add_moe(moe_hidden, moe_indices, moe_scores) - if moe is None: - raise RuntimeError( - "TensorRT rejected addMoE for this build; the per-expert path " - "should have been selected instead" - ) - moe.set_gated_weights( - graph_ops.add_constant(network, w_gate.shape, w_gate, dtype=dtype), - graph_ops.add_constant(network, w_up.shape, w_up, dtype=dtype), - graph_ops.add_constant(network, w_down.shape, w_down, dtype=dtype), - trt.MoEActType.SILU, - ) - routed_out = moe.get_output(0) - - if rank == 2: - restore = network.add_shuffle(routed_out) - restore.reshape_dims = (-1, hidden_size) - routed_out = restore.get_output(0) - return routed_out + if rank != 2: + raise ValueError(f"DeepSeek-V2 MoE expects rank-2 hidden states, got rank {rank}") + + hidden = network.add_shuffle(inp) + hidden.reshape_dims = (-1, 1, 1, hidden_size) + + def selected(values: np.ndarray) -> trt.ITensor: + packed = graph_ops.add_constant(network, values.shape, values, dtype=dtype) + if packed.dtype != inp.dtype: + packed = network.add_cast(packed, inp.dtype).get_output(0) + return network.add_gather(packed, top_indices, 0).get_output(0) + + gate = network.add_matrix_multiply( + hidden.get_output(0), + trt.MatrixOperation.NONE, + selected(w_gate), + trt.MatrixOperation.NONE, + ).get_output(0) + up = network.add_matrix_multiply( + hidden.get_output(0), + trt.MatrixOperation.NONE, + selected(w_up), + trt.MatrixOperation.NONE, + ).get_output(0) + sigmoid = network.add_activation(gate, trt.ActivationType.SIGMOID).get_output(0) + swish = network.add_elementwise(gate, sigmoid, trt.ElementWiseOperation.PROD).get_output(0) + gated = network.add_elementwise(swish, up, trt.ElementWiseOperation.PROD).get_output(0) + down = network.add_matrix_multiply( + gated, + trt.MatrixOperation.NONE, + selected(w_down), + trt.MatrixOperation.NONE, + ).get_output(0) + experts = network.add_shuffle(down) + experts.reshape_dims = (-1, num_experts_per_tok, hidden_size) + routing = network.add_shuffle(scaled_weights) + routing.reshape_dims = (-1, num_experts_per_tok, 1) + routing_weights = routing.get_output(0) + if routing_weights.dtype != experts.get_output(0).dtype: + routing_weights = network.add_cast( + routing_weights, experts.get_output(0).dtype + ).get_output(0) + weighted = network.add_elementwise( + experts.get_output(0), routing_weights, trt.ElementWiseOperation.PROD + ).get_output(0) + return network.add_reduce( + weighted, + trt.ReduceOperation.SUM, + 1 << 1, + keep_dims=False, + ).get_output(0) def _add_moe_with_shared_experts( @@ -1035,7 +1042,7 @@ def _add_moe_with_shared_experts( 1. Router logits -> softmax/sigmoid -> top-k selection 2. Scale weights: renormalize (norm_topk_prob=True) or multiply by routed_scaling_factor (norm_topk_prob=False) - 3. Native TensorRT MoE routed-expert output. + 3. TensorRT graph layers compute the selected routed experts. 4. Compute shared expert output (always active) 5. Final = routed_output + shared_output """ @@ -1055,7 +1062,7 @@ def _add_moe_with_shared_experts( routed_scaling_factor=routed_scaling_factor, ) - result = _add_native_routed_experts( + result = _add_routed_experts( network, inp, weights, diff --git a/families/deepseek_v2/tests/test_router_contract.py b/families/deepseek_v2/tests/test_router_contract.py index f77be3eae6..a07f9f2387 100644 --- a/families/deepseek_v2/tests/test_router_contract.py +++ b/families/deepseek_v2/tests/test_router_contract.py @@ -121,6 +121,108 @@ def test_non_finite_router_score_bias_is_rejected() -> None: ) +def test_selected_expert_graph_matches_numpy(monkeypatch: pytest.MonkeyPatch) -> None: + """The portable selected-expert graph preserves DeepSeek SwiGLU routing.""" + + class Tensor: + def __init__(self, data: np.ndarray): + self.data = np.asarray(data) + self.shape = self.data.shape + self.dtype = self.data.dtype + + class Layer: + def __init__(self, data: np.ndarray): + self.output = Tensor(data) + + def get_output(self, _index: int) -> Tensor: + return self.output + + class Shuffle(Layer): + @property + def reshape_dims(self): + return self.output.shape + + @reshape_dims.setter + def reshape_dims(self, shape) -> None: + self.output = Tensor(self.output.data.reshape(shape)) + + class Network: + @staticmethod + def add_shuffle(tensor: Tensor) -> Shuffle: + return Shuffle(tensor.data) + + @staticmethod + def add_cast(tensor: Tensor, dtype) -> Layer: + return Layer(tensor.data.astype(dtype)) + + @staticmethod + def add_gather(data: Tensor, indices: Tensor, _axis: int) -> Layer: + return Layer(np.take(data.data, indices.data, axis=0)) + + @staticmethod + def add_matrix_multiply(left: Tensor, _left_op, right: Tensor, _right_op) -> Layer: + return Layer(np.matmul(left.data, right.data)) + + @staticmethod + def add_activation(tensor: Tensor, _operation) -> Layer: + return Layer(1.0 / (1.0 + np.exp(-tensor.data))) + + @staticmethod + def add_elementwise(left: Tensor, right: Tensor, _operation) -> Layer: + return Layer(left.data * right.data) + + @staticmethod + def add_reduce(tensor: Tensor, _operation, _axes: int, keep_dims: bool) -> Layer: + return Layer(np.sum(tensor.data, axis=1, keepdims=keep_dims)) + + monkeypatch.setattr( + model.graph_ops, + "add_constant", + lambda _network, _shape, values, dtype: Tensor(np.asarray(values, dtype=dtype)), + ) + hidden = np.array([[0.25, -0.5], [1.0, 0.75]], dtype=np.float16) + indices = np.array([[0, 2], [1, 0]], dtype=np.int32) + routing = np.array([[0.7, 0.3], [0.4, 0.6]], dtype=np.float32) + weights = {} + for expert in range(3): + scale = float(expert + 1) + weights[f"layer.0.expert.{expert}.w_gate"] = ( + np.array([[0.2, -0.1, 0.3], [0.4, 0.5, -0.2]], dtype=np.float16) * scale + ) + weights[f"layer.0.expert.{expert}.w_up"] = ( + np.array([[0.3, 0.2, -0.4], [-0.1, 0.6, 0.5]], dtype=np.float16) * scale + ) + weights[f"layer.0.expert.{expert}.w_down"] = ( + np.array([[0.5, -0.2], [0.1, 0.4], [-0.3, 0.2]], dtype=np.float16) / scale + ) + + actual = model._add_routed_experts( + Network(), + Tensor(hidden), + weights, + "layer.0", + hidden_size=2, + n_routed_experts=3, + num_experts_per_tok=2, + top_indices=Tensor(indices), + scaled_weights=Tensor(routing), + dtype=np.float16, + ).data + + expected = np.zeros_like(hidden) + for token in range(hidden.shape[0]): + for route in range(indices.shape[1]): + expert = int(indices[token, route]) + gate = hidden[token] @ weights[f"layer.0.expert.{expert}.w_gate"] + up = hidden[token] @ weights[f"layer.0.expert.{expert}.w_up"] + activated = gate / (1.0 + np.exp(-gate)) + down = (activated * up) @ weights[f"layer.0.expert.{expert}.w_down"] + expected[token] += routing[token, route] * down + + assert actual.dtype == np.float16 + np.testing.assert_allclose(actual, expected, rtol=2e-3, atol=2e-3) + + @pytest.mark.parametrize( ("scoring_func", "n_routed_experts", "n_group"), ( diff --git a/families/dinov3/requirements.txt b/families/dinov3/requirements.txt index 718a93c119..18b482b5e6 100644 --- a/families/dinov3/requirements.txt +++ b/families/dinov3/requirements.txt @@ -2,4 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 Pillow -timm>=1.0 +# The public parity oracle was qualified against this exact release. Keep its +# numerical reference fixed instead of resolving a different timm at runtime. +timm==1.0.28 diff --git a/families/glm/model.py b/families/glm/model.py index 1e29e4f95f..49c7a7b150 100644 --- a/families/glm/model.py +++ b/families/glm/model.py @@ -8,8 +8,6 @@ import json from typing import TYPE_CHECKING -import os -from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path import numpy as np @@ -122,20 +120,10 @@ def _load_layer(layer_idx: int) -> tuple[int, WeightDict, int, int]: return layer_idx, layer, q_raw.shape[0], intermediate - layer_results: list[tuple[int, WeightDict, int, int] | None] = [None] * num_layers - max_workers = min(8, max(1, os.cpu_count() or 1)) - with ThreadPoolExecutor(max_workers=max_workers) as executor: - futures = [executor.submit(_load_layer, i) for i in range(num_layers)] - for future in as_completed(futures): - layer_idx, layer, attention_size, mlp_size = future.result() - layer_results[layer_idx] = (layer_idx, layer, attention_size, mlp_size) - attention_size = 0 mlp_size = 0 - for result in layer_results: - if result is None: - continue - _layer_idx, layer, layer_attention_size, layer_mlp_size = result + for layer_idx in range(num_layers): + _layer_idx, layer, layer_attention_size, layer_mlp_size = _load_layer(layer_idx) weights.update(layer) if attention_size == 0: attention_size = layer_attention_size diff --git a/families/glm/tests/test_weight_memory.py b/families/glm/tests/test_weight_memory.py new file mode 100644 index 0000000000..d594a13e53 --- /dev/null +++ b/families/glm/tests/test_weight_memory.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GLM checkpoint loading memory regression coverage.""" + +from __future__ import annotations + +import threading +import time + +import numpy as np +import pytest + +pytest.importorskip("tensorrt", reason="TensorRT is required for family builder tests") + +from families.glm import model +from families.glm.config import ModelConfig + + +def test_layers_are_converted_one_at_a_time(monkeypatch: pytest.MonkeyPatch) -> None: + """Large temporary projection arrays cannot overlap across decoder layers.""" + hidden, vocab, intermediate = 4, 7, 6 + lock = threading.Lock() + active = 0 + maximum = 0 + + def load(_readers, name: str) -> np.ndarray: + nonlocal active, maximum + with lock: + active += 1 + maximum = max(maximum, active) + try: + time.sleep(0.005) + if name == "model.embed_tokens.weight": + shape = (vocab, hidden) + elif name.endswith("gate_up_proj.weight"): + shape = (2 * intermediate, hidden) + elif name.endswith("down_proj.weight"): + shape = (hidden, intermediate) + elif name.endswith(".weight") and ".self_attn." in name: + shape = (hidden, hidden) + else: + shape = (hidden,) + return np.ones(shape, dtype=np.float32) + finally: + with lock: + active -= 1 + + monkeypatch.setattr(model, "_open_safetensors", lambda _path: object()) + monkeypatch.setattr(model, "_load_tensor", load) + monkeypatch.setattr( + model, + "_has_tensor", + lambda _readers, name: name != "lm_head.weight", + ) + config = ModelConfig( + model_type="glm", + vocab_size=vocab, + hidden_size=hidden, + intermediate_size=intermediate, + num_hidden_layers=2, + num_attention_heads=1, + num_key_value_heads=1, + ) + + weights = model._GlmModel().load_weights("/unused", config, precision="fp16") + + assert weights["layer.1.w_down"].shape == (intermediate, hidden) + assert maximum == 1 diff --git a/families/gpt_oss/model.py b/families/gpt_oss/model.py index a12a23c4da..113f52773d 100644 --- a/families/gpt_oss/model.py +++ b/families/gpt_oss/model.py @@ -62,6 +62,8 @@ def load_weights( self, model_dir: str, config: ModelConfig, + *, + precision: str = "fp32", ) -> WeightDict: """Load GPT-OSS weights via AutoModelForCausalLM (handles MXFP4 dequant). @@ -82,10 +84,18 @@ def load_weights( torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ) - state = {k: v.float().cpu().numpy() for k, v in model.state_dict().items()} + state = model.state_dict() del model gc.collect() + target_np_dtype = np.float16 if precision in {"fp16", "bf16"} else np.float32 + target_torch_dtype = torch.float16 if target_np_dtype == np.float16 else torch.float32 + + def take(name: str) -> np.ndarray: + """Convert one tensor and release its dequantized source storage.""" + tensor = state.pop(name) + return tensor.detach().to(device="cpu", dtype=target_torch_dtype).numpy() + hidden = config.hidden_size vocab = config.vocab_size num_layers = config.num_hidden_layers @@ -94,11 +104,11 @@ def load_weights( weights = WeightDict() # Embedding - embedding = state["model.embed_tokens.weight"] + embedding = take("model.embed_tokens.weight") assert embedding.shape == (vocab, hidden), ( f"Embedding shape {embedding.shape} != ({vocab}, {hidden})" ) - weights["embedding"] = embedding.astype(np.float32) + weights["embedding"] = np.ascontiguousarray(embedding, dtype=target_np_dtype) attention_size = 0 moe_intermediate = 0 @@ -108,26 +118,22 @@ def load_weights( hf = f"model.layers.{layer_idx}" # RMSNorm (no biases) - weights[f"{prefix}.input_norm"] = state[f"{hf}.input_layernorm.weight"].astype( - np.float32 - ) - weights[f"{prefix}.post_attn_norm"] = state[ - f"{hf}.post_attention_layernorm.weight" - ].astype(np.float32) + weights[f"{prefix}.input_norm"] = take(f"{hf}.input_layernorm.weight") + weights[f"{prefix}.post_attn_norm"] = take(f"{hf}.post_attention_layernorm.weight") # --- Attention projections (with biases) --- - q_raw = state[f"{hf}.self_attn.q_proj.weight"] - k_raw = state[f"{hf}.self_attn.k_proj.weight"] - v_raw = state[f"{hf}.self_attn.v_proj.weight"] - o_raw = state[f"{hf}.self_attn.o_proj.weight"] + q_raw = take(f"{hf}.self_attn.q_proj.weight") + k_raw = take(f"{hf}.self_attn.k_proj.weight") + v_raw = take(f"{hf}.self_attn.v_proj.weight") + o_raw = take(f"{hf}.self_attn.o_proj.weight") if attention_size == 0: attention_size = q_raw.shape[0] - q_t = _transpose_2d(q_raw, "q_proj") - k_t = _transpose_2d(k_raw, "k_proj") - v_t = _transpose_2d(v_raw, "v_proj") - o_t = _transpose_2d(o_raw, "o_proj") + q_t = _transpose_2d(q_raw, "q_proj", precision=precision) + k_t = _transpose_2d(k_raw, "k_proj", precision=precision) + v_t = _transpose_2d(v_raw, "v_proj", precision=precision) + o_t = _transpose_2d(o_raw, "o_proj", precision=precision) # Keep compact GQA/MQA K/V @@ -137,27 +143,27 @@ def load_weights( weights[f"{prefix}.w_o"] = o_t # Attention biases - weights[f"{prefix}.q_bias"] = state[f"{hf}.self_attn.q_proj.bias"].astype(np.float32) - weights[f"{prefix}.o_bias"] = state[f"{hf}.self_attn.o_proj.bias"].astype(np.float32) + weights[f"{prefix}.q_bias"] = take(f"{hf}.self_attn.q_proj.bias") + weights[f"{prefix}.o_bias"] = take(f"{hf}.self_attn.o_proj.bias") - weights[f"{prefix}.k_bias"] = state[f"{hf}.self_attn.k_proj.bias"].astype(np.float32) - weights[f"{prefix}.v_bias"] = state[f"{hf}.self_attn.v_proj.bias"].astype(np.float32) + weights[f"{prefix}.k_bias"] = take(f"{hf}.self_attn.k_proj.bias") + weights[f"{prefix}.v_bias"] = take(f"{hf}.self_attn.v_proj.bias") # Attention sinks (per-head learned parameter for softmax normalization) sinks_key = f"{hf}.self_attn.sinks" if sinks_key in state: - weights[f"{prefix}.sinks"] = state[sinks_key].astype(np.float32) + weights[f"{prefix}.sinks"] = take(sinks_key) # --- Router --- - router_w = state[f"{hf}.mlp.router.weight"] # [num_experts, hidden] - weights[f"{prefix}.router"] = _transpose_2d(router_w, "router") - weights[f"{prefix}.router_bias"] = state[f"{hf}.mlp.router.bias"].astype(np.float32) + router_w = take(f"{hf}.mlp.router.weight") # [num_experts, hidden] + weights[f"{prefix}.router"] = _transpose_2d(router_w, "router", precision=precision) + weights[f"{prefix}.router_bias"] = take(f"{hf}.mlp.router.bias") # --- Packed expert weights --- - gate_up = state[f"{hf}.mlp.experts.gate_up_proj"] - gate_up_bias = state[f"{hf}.mlp.experts.gate_up_proj_bias"] - down = state[f"{hf}.mlp.experts.down_proj"] - down_bias = state[f"{hf}.mlp.experts.down_proj_bias"] + gate_up = take(f"{hf}.mlp.experts.gate_up_proj") + gate_up_bias = take(f"{hf}.mlp.experts.gate_up_proj_bias") + down = take(f"{hf}.mlp.experts.down_proj") + down_bias = take(f"{hf}.mlp.experts.down_proj_bias") # gate_up_proj is [E, hidden, 2*inter] with INTERLEAVED # gate/up columns: gate=even indices, up=odd indices. @@ -168,13 +174,13 @@ def load_weights( gu = gate_up[e_idx] # [hidden, 2*inter] # Interleaved: gate = columns 0,2,4,... up = columns 1,3,5,... weights[f"{prefix}.expert.{e_idx}.w_gate"] = np.ascontiguousarray( - gu[:, ::2], dtype=np.float32 + gu[:, ::2], dtype=target_np_dtype ) weights[f"{prefix}.expert.{e_idx}.w_up"] = np.ascontiguousarray( - gu[:, 1::2], dtype=np.float32 + gu[:, 1::2], dtype=target_np_dtype ) weights[f"{prefix}.expert.{e_idx}.w_down"] = np.ascontiguousarray( - down[e_idx], dtype=np.float32 + down[e_idx], dtype=target_np_dtype ) if moe_intermediate == 0: @@ -183,27 +189,37 @@ def load_weights( # Per-expert biases (also interleaved for gate_up) for e_idx in range(num_experts): gu_b = gate_up_bias[e_idx] # [2*inter] - weights[f"{prefix}.expert.{e_idx}.gate_bias"] = gu_b[::2].astype(np.float32) - weights[f"{prefix}.expert.{e_idx}.up_bias"] = gu_b[1::2].astype(np.float32) - weights[f"{prefix}.expert.{e_idx}.down_bias"] = down_bias[e_idx].astype(np.float32) + weights[f"{prefix}.expert.{e_idx}.gate_bias"] = np.ascontiguousarray( + gu_b[::2], dtype=target_np_dtype + ) + weights[f"{prefix}.expert.{e_idx}.up_bias"] = np.ascontiguousarray( + gu_b[1::2], dtype=target_np_dtype + ) + weights[f"{prefix}.expert.{e_idx}.down_bias"] = np.ascontiguousarray( + down_bias[e_idx], dtype=target_np_dtype + ) + del gate_up, gate_up_bias, down, down_bias # Final norm final_key = "model.norm.weight" if final_key in state: - weights["final_norm"] = state[final_key].astype(np.float32) + weights["final_norm"] = take(final_key) else: - weights["final_norm"] = np.ones(hidden, dtype=np.float32) + weights["final_norm"] = np.ones(hidden, dtype=target_np_dtype) # LM head lm_key = "lm_head.weight" if lm_key in state: - weights["w_out"] = _transpose_2d(state[lm_key], "lm_head") + weights["w_out"] = _transpose_2d(take(lm_key), "lm_head", precision=precision) else: - weights["w_out"] = _transpose_2d(embedding.copy(), "embedding_tied") + weights["w_out"] = _transpose_2d(embedding, "embedding_tied", precision=precision) lm_bias_key = "lm_head.bias" if lm_bias_key in state: - weights["lm_head_bias"] = state[lm_bias_key].astype(np.float32) + weights["lm_head_bias"] = take(lm_bias_key) + + state.clear() + gc.collect() # Metadata weights["_attention_size"] = attention_size # type: ignore[assignment] @@ -1067,7 +1083,7 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: config.raw["_resolved_build_precision"] = precision config.raw["_parallel_build_enabled"] = parallel.enabled config.raw["_quantized_build_requested"] = False - weights = model.load_weights(str(model_dir), config) + weights = model.load_weights(str(model_dir), config, precision=precision) writer.set_header(family="gpt_oss", task=request.task, backend=request.backend) if parallel.enabled: for rank in range(parallel.tp_size): diff --git a/families/gpt_oss/tests/test_weight_memory.py b/families/gpt_oss/tests/test_weight_memory.py new file mode 100644 index 0000000000..4fb9f3cca4 --- /dev/null +++ b/families/gpt_oss/tests/test_weight_memory.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPT-OSS checkpoint conversion memory regression coverage.""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +pytest.importorskip("tensorrt", reason="TensorRT is required for family builder tests") + +from families.gpt_oss import model +from families.gpt_oss.config import ModelConfig + + +class _Tensor: + def __init__(self, values: np.ndarray): + self.values = np.asarray(values) + + def detach(self): + return self + + def to(self, *, device: str, dtype): + assert device == "cpu" + return _Tensor(self.values.astype(dtype)) + + def numpy(self) -> np.ndarray: + return self.values + + +def test_fp16_conversion_releases_source_state_without_fp32_mirror( + monkeypatch: pytest.MonkeyPatch, +) -> None: + hidden, vocab, experts, intermediate = 2, 3, 2, 2 + prefix = "model.layers.0" + + def tensor(*shape: int) -> _Tensor: + return _Tensor(np.arange(np.prod(shape), dtype=np.float32).reshape(shape)) + + state = { + "model.embed_tokens.weight": tensor(vocab, hidden), + f"{prefix}.input_layernorm.weight": tensor(hidden), + f"{prefix}.post_attention_layernorm.weight": tensor(hidden), + f"{prefix}.self_attn.q_proj.weight": tensor(hidden, hidden), + f"{prefix}.self_attn.k_proj.weight": tensor(hidden, hidden), + f"{prefix}.self_attn.v_proj.weight": tensor(hidden, hidden), + f"{prefix}.self_attn.o_proj.weight": tensor(hidden, hidden), + f"{prefix}.self_attn.q_proj.bias": tensor(hidden), + f"{prefix}.self_attn.k_proj.bias": tensor(hidden), + f"{prefix}.self_attn.v_proj.bias": tensor(hidden), + f"{prefix}.self_attn.o_proj.bias": tensor(hidden), + f"{prefix}.mlp.router.weight": tensor(experts, hidden), + f"{prefix}.mlp.router.bias": tensor(experts), + f"{prefix}.mlp.experts.gate_up_proj": tensor(experts, hidden, 2 * intermediate), + f"{prefix}.mlp.experts.gate_up_proj_bias": tensor(experts, 2 * intermediate), + f"{prefix}.mlp.experts.down_proj": tensor(experts, intermediate, hidden), + f"{prefix}.mlp.experts.down_proj_bias": tensor(experts, hidden), + "model.norm.weight": tensor(hidden), + } + + class FakeModel: + def state_dict(self): + return state + + loader = SimpleNamespace(from_pretrained=lambda *_args, **_kwargs: FakeModel()) + monkeypatch.setitem( + sys.modules, + "torch", + SimpleNamespace(bfloat16="bfloat16", float16=np.float16, float32=np.float32), + ) + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace(AutoModelForCausalLM=loader), + ) + config = ModelConfig( + model_type="gpt_oss", + vocab_size=vocab, + hidden_size=hidden, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + raw={"num_local_experts": experts, "num_experts_per_tok": 1}, + ) + + weights = model._GptOssModel().load_weights("/unused", config, precision="fp16") + + assert state == {} + arrays = [value for value in weights.values() if isinstance(value, np.ndarray)] + assert arrays + assert {value.dtype for value in arrays} == {np.dtype(np.float16)} diff --git a/tools/ci/e2e.py b/tools/ci/e2e.py index 7da373648c..db5cb02d81 100644 --- a/tools/ci/e2e.py +++ b/tools/ci/e2e.py @@ -7,6 +7,7 @@ import json import re +import signal import shlex import tempfile import xml.etree.ElementTree as ET @@ -159,6 +160,34 @@ def _require_e2e_junit( raise CiError(f"{family} E2E result validation failed: " + "; ".join(problems)) +def _require_e2e_process_result( + report: Path, + family: str, + requested_testcases: tuple[str, ...], + returncode: int, +) -> None: + """Preserve process termination details when pytest cannot write JUnit.""" + if returncode and not report.is_file(): + if returncode < 0: + try: + cause = signal.Signals(-returncode).name + except ValueError: + cause = f"signal {-returncode}" + raise CiError( + f"{family} E2E pytest was terminated by {cause} before writing JUnit report" + ) + cause = f"exit code {returncode}" + if 128 < returncode <= 192: + try: + cause += f" ({signal.Signals(returncode - 128).name})" + except ValueError: + pass + raise CiError(f"{family} E2E pytest failed with {cause} before writing JUnit report") + _require_e2e_junit(report, family, requested_testcases) + if returncode: + raise CiError(f"{family} E2E pytest failed with exit code {returncode}") + + class E2ERunner: def __init__(self, context: CiContext): self.context = context @@ -320,13 +349,12 @@ def _run(self, families: tuple[str, ...], testcases: tuple[str, ...] = ()) -> No check=False, limit=self.context.env.get("TRTMC_E2E_TIMEOUT", "12h"), ) - _require_e2e_junit( + _require_e2e_process_result( e2e_junit, family, tuple(requested_testcases), + completed.returncode, ) - if completed.returncode: - raise CiError(f"{family} E2E pytest failed with exit code {completed.returncode}") def _family_testcases(self, family: str) -> tuple[str, ...]: manifests = self.context.repository / "families" / family / "tests/manifests" diff --git a/tools/tests/test_new_ci.py b/tools/tests/test_new_ci.py index a8fa291693..0c224b9c75 100644 --- a/tools/tests/test_new_ci.py +++ b/tools/tests/test_new_ci.py @@ -20,7 +20,12 @@ from tools.ci.context import CiContext from tools.ci.container import CiContainer from tools.ci.docker_image import DockerImageManager -from tools.ci.e2e import E2ERunner, _require_e2e_junit, _require_passing_junit +from tools.ci.e2e import ( + E2ERunner, + _require_e2e_junit, + _require_e2e_process_result, + _require_passing_junit, +) from tools.ci.package import ( SourceArchiveValidator, WheelArchiveValidator, @@ -398,6 +403,22 @@ def test_e2e_junit_fails_closed_for_invalid_result_cardinality_or_skip( _require_e2e_junit(report, "family", requested) +@pytest.mark.parametrize( + ("returncode", "message"), + ((-9, "terminated by SIGKILL"), (137, r"exit code 137 \(SIGKILL\)")), +) +def test_e2e_process_failure_preserves_signal_when_junit_is_missing( + tmp_path: Path, returncode: int, message: str +) -> None: + with pytest.raises(CiError, match=message): + _require_e2e_process_result( + tmp_path / "missing.xml", + "family", + ("family-case",), + returncode, + ) + + def test_e2e_rejects_multiple_family_environments(tmp_path: Path) -> None: context = RecordingContext(tmp_path, {})