From 8a14ef074ba34eab7d29eaf166af3c30237961fd Mon Sep 17 00:00:00 2001 From: Yizhuo Zhang Date: Fri, 16 Oct 2026 19:25:43 -0700 Subject: [PATCH 1/4] test(phi4): export the checkpoint tokenizer for AIPerf Signed-off-by: Yizhuo Zhang --- .../tests/aiperf/prepare_tokenizer.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) create mode 100644 families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py diff --git a/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py b/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py new file mode 100644 index 0000000000..949463d6c1 --- /dev/null +++ b/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Export Phi4's standard tokenizer without its model-only remote configuration.""" + +import argparse +import importlib.metadata +import json +from pathlib import Path + + +def export_tokenizer(model_dir: Path, output_dir: Path) -> None: + from transformers import AutoTokenizer, GPT2TokenizerFast, PreTrainedTokenizerFast + + config = json.loads((model_dir / "config.json").read_text()) + if config.get("model_type") not in {"phi4mm", "phi4_multimodal", "phi4multimodal"}: + raise ValueError("model directory must contain the Phi4 multimodal checkpoint") + tokenizer = GPT2TokenizerFast.from_pretrained(model_dir, local_files_only=True) + probes = ("hello", "What is the main object in this image? Answer in one word.", + "<|user|>\n<|image_1|>\nQuestion?<|end|>\n<|assistant|>\n", "找到图片中的狗。") + identities = [ + {"text": text, "add_special_tokens": special, + "token_ids": tokenizer.encode(text, add_special_tokens=special)} + for text in probes for special in (False, True) + ] + vocabulary = tokenizer.get_vocab() + # Keep the checkpoint's GPT-4o regex/backend unchanged. A model-specific + # GPT2 wrapper in a newer client can rebuild it with a different regex. + exported = PreTrainedTokenizerFast( + tokenizer_object=tokenizer.backend_tokenizer, + model_max_length=tokenizer.model_max_length, + padding_side=tokenizer.padding_side, + truncation_side=tokenizer.truncation_side, + clean_up_tokenization_spaces=tokenizer.clean_up_tokenization_spaces, + **tokenizer.special_tokens_map, + ) + output_dir.mkdir(parents=True, exist_ok=False) + exported.save_pretrained(output_dir) + restored = AutoTokenizer.from_pretrained(output_dir, use_fast=True, local_files_only=True) + if restored.get_vocab() != vocabulary: + raise RuntimeError("exported tokenizer changed the vocabulary") + for probe in identities: + if restored.encode(probe["text"], add_special_tokens=probe["add_special_tokens"]) != probe["token_ids"]: + raise RuntimeError("exported tokenizer changed token IDs or framing") + audit = {"source": str(model_dir.resolve()), + "transformers": importlib.metadata.version("transformers"), + "tokenizer_class": type(tokenizer).__name__, "probes": identities} + (output_dir / "trtmc-tokenizer-export.json").write_text(json.dumps(audit, indent=2) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", type=Path, required=True, help="local Phi4 multimodal checkpoint") + parser.add_argument("--output-dir", type=Path, required=True, help="new directory for client tokenizer files") + args = parser.parse_args() + export_tokenizer(args.model_dir, args.output_dir) + + +if __name__ == "__main__": + main() From f7b1c229ef33bb84027fb90b69c38956330be969 Mon Sep 17 00:00:00 2001 From: Yizhuo Zhang Date: Fri, 16 Oct 2026 19:53:56 -0700 Subject: [PATCH 2/4] fix(phi4): preserve the exported chat template Signed-off-by: Yizhuo Zhang --- families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py b/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py index 949463d6c1..9cb1798562 100644 --- a/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py +++ b/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py @@ -33,6 +33,7 @@ def export_tokenizer(model_dir: Path, output_dir: Path) -> None: padding_side=tokenizer.padding_side, truncation_side=tokenizer.truncation_side, clean_up_tokenization_spaces=tokenizer.clean_up_tokenization_spaces, + chat_template=tokenizer.chat_template, **tokenizer.special_tokens_map, ) output_dir.mkdir(parents=True, exist_ok=False) @@ -40,6 +41,8 @@ def export_tokenizer(model_dir: Path, output_dir: Path) -> None: restored = AutoTokenizer.from_pretrained(output_dir, use_fast=True, local_files_only=True) if restored.get_vocab() != vocabulary: raise RuntimeError("exported tokenizer changed the vocabulary") + if restored.chat_template != tokenizer.chat_template: + raise RuntimeError("exported tokenizer changed the chat template") for probe in identities: if restored.encode(probe["text"], add_special_tokens=probe["add_special_tokens"]) != probe["token_ids"]: raise RuntimeError("exported tokenizer changed token IDs or framing") From 37d982b64e82e7c2f56b1116649763a83eac31ca Mon Sep 17 00:00:00 2001 From: Yizhuo Zhang Date: Fri, 16 Oct 2026 21:18:19 -0700 Subject: [PATCH 3/4] fix(phi4): release each staged engine before building next Reuse the file-backed bundle writer to stage prefill, decode and vision plans as soon as they are serialized. Avoid retaining prior engine blobs during subsequent builds without changing their graphs, weights or public named sections. A lifecycle regression fails the prior retention behavior. Signed-off-by: Yizhuo Zhang --- families/phi4_multimodal/model.py | 13 +++- .../tests/test_plan_lifetime.py | 78 +++++++++++++++++++ 2 files changed, 87 insertions(+), 4 deletions(-) create mode 100644 families/phi4_multimodal/tests/test_plan_lifetime.py diff --git a/families/phi4_multimodal/model.py b/families/phi4_multimodal/model.py index a1782e9d03..04877516fa 100644 --- a/families/phi4_multimodal/model.py +++ b/families/phi4_multimodal/model.py @@ -343,6 +343,9 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None: config.raw["_model_dir"] = str(model_dir) model = _Phi4MultimodalModel() weights = model.load_weights(str(model_dir), config) + writer.set_header(family="phi4_multimodal", task=request.task, backend=request.backend) + # Stage each serialized plan before building the next network. The + # bundle writer keeps sections on disk and publishes only at finish. config.raw["_decoder_engine_role"] = "prefill" prefill = model.build_engine( config, @@ -352,6 +355,8 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None: quant_ctx=None, verbose=request.verbose, ) + writer.add_bytes("prefill.plan", prefill) + del prefill config.raw["_decoder_engine_role"] = "decode" decode = model.build_engine( config, @@ -361,12 +366,16 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None: quant_ctx=None, verbose=request.verbose, ) + writer.add_bytes("engine.plan", decode) + del decode config.raw.pop("_decoder_engine_role", None) vision = model.build_vision_engine( str(model_dir), config, weights, precision=precision, verbose=request.verbose ) if vision is None: raise RuntimeError("Phi-4 Multimodal vision build returned no engine") + writer.add_bytes("vision.plan", vision) + del vision vl = model.get_vl_config(config) or {} runtime = { "tensor_parallel_size": 1, @@ -386,10 +395,6 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None: }, } runtime.update(vl) - writer.set_header(family="phi4_multimodal", task=request.task, backend=request.backend) - writer.add_bytes("engine.plan", decode) - writer.add_bytes("prefill.plan", prefill) - writer.add_bytes("vision.plan", vision) runtime.update(_tokenizer_runtime_contract(model_dir)) writer.add_json("runtime.json", runtime) for filename in ( diff --git a/families/phi4_multimodal/tests/test_plan_lifetime.py b/families/phi4_multimodal/tests/test_plan_lifetime.py new file mode 100644 index 0000000000..b5452a260b --- /dev/null +++ b/families/phi4_multimodal/tests/test_plan_lifetime.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Keep serialized engines out of RAM while building the next network.""" + +import gc +import json +from types import SimpleNamespace + +import pytest + + +def test_each_plan_is_staged_and_released_before_the_next_build(tmp_path, monkeypatch): + pytest.importorskip("tensorrt") + from families.phi4_multimodal import model + from families.phi4_multimodal.edge_llm import dispatch + from tensorrt_model_connect.bundle_writer import BundleWriter + + live = set() + built = [] + + class Plan(bytes): + def __new__(cls, role): + value = super().__new__(cls, role.encode()) + value.role = role + live.add(role) + return value + + def __del__(self): + live.remove(self.role) + + def next_plan(role): + gc.collect() + assert not live, f"previous serialized engine retained during {role}: {live}" + built.append(role) + return Plan(role) + + class Model: + @staticmethod + def load_weights(*_args): + return {} + + @staticmethod + def build_engine(config, *_args, **_kwargs): + return next_plan(config.raw["_decoder_engine_role"]) + + @staticmethod + def build_vision_engine(*_args, **_kwargs): + return next_plan("vision") + + @staticmethod + def get_vl_config(_config): + return {"image_token_id": 3, "vision_output_dim": 4} + + monkeypatch.setattr(model, "_Phi4MultimodalModel", Model) + monkeypatch.setattr(model, "_tokenizer_runtime_contract", lambda _directory: {}) + monkeypatch.setattr(dispatch, "build", lambda request, writer, native: native(request, writer)) + (tmp_path / "config.json").write_text(json.dumps({ + "model_type": "phi4mm", "hidden_size": 4, "num_hidden_layers": 1, + "vocab_size": 8, "num_attention_heads": 1, "num_key_value_heads": 1, + })) + request = SimpleNamespace( + mtp_seq_len=None, dynamic_kv_cache=False, image_height=None, image_width=None, + video_num_frames=None, max_batch_size=1, context_parallel_size=1, + task="vision_language_generation", tensor_parallel_size=1, quantization=None, + fp32_layers=(), model_dir=tmp_path, precision="fp16", max_sequence_length=768, + verbose=False, backend="tensorrt", + ) + writer = BundleWriter(tmp_path / "model.bundle") + try: + model.build(request, writer) + writer.finish() + finally: + writer.abort() + gc.collect() + assert built == ["prefill", "decode", "vision"] + assert not live + assert (tmp_path / "model.bundle").is_file() From eed140a43ed1ec49488e08b5a6e30acf572a0b6f Mon Sep 17 00:00:00 2001 From: Yizhuo Zhang Date: Sat, 17 Oct 2026 13:37:19 -0700 Subject: [PATCH 4/4] test(phi4): reduce reference allocator fragmentation Default the family pytest environment to expandable PyTorch segments while respecting caller allocator settings. A matched 21.8 GiB allocator-capacity control fails the original FP32 eager reference with the default allocator and completes with identical original text and all 30 token IDs with expandable segments. No reference operators, precision, requests or numerical gates change. Signed-off-by: Yizhuo Zhang --- families/phi4_multimodal/tests/conftest.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 families/phi4_multimodal/tests/conftest.py diff --git a/families/phi4_multimodal/tests/conftest.py b/families/phi4_multimodal/tests/conftest.py new file mode 100644 index 0000000000..fe700c1183 --- /dev/null +++ b/families/phi4_multimodal/tests/conftest.py @@ -0,0 +1,13 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Resource configuration for the original FP32 eager family reference.""" + +import os + + +def pytest_configure(config): + # The eager reference's varying allocations fragment a nearly full GPU. + # Keep its precision/operators unchanged and preserve caller configuration. + if "PYTORCH_ALLOC_CONF" not in os.environ: + os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")