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/aiperf/prepare_tokenizer.py b/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py new file mode 100644 index 0000000000..9cb1798562 --- /dev/null +++ b/families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py @@ -0,0 +1,64 @@ +#!/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, + chat_template=tokenizer.chat_template, + **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") + 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") + 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() 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") 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()