diff --git a/docs/source/features/checkpoint-loading.md b/docs/source/features/checkpoint-loading.md index ec4da35c5492..75d0ce644b85 100644 --- a/docs/source/features/checkpoint-loading.md +++ b/docs/source/features/checkpoint-loading.md @@ -6,9 +6,10 @@ The PyTorch backend provides a flexible and extensible infrastructure for loadin 1. [Overview](#overview) 2. [Core Components](#core-components) 3. [Built-in Checkpoint Formats](#built-in-checkpoint-formats) -4. [Checkpoint I/O Policies](#checkpoint-io-policies) -5. [Using Checkpoint Loaders](#using-checkpoint-loaders) -6. [Creating Custom Checkpoint Loaders](#creating-custom-checkpoint-loaders) +4. [Using Checkpoint Loaders](#using-checkpoint-loaders) +5. [Creating Custom Checkpoint Loaders](#creating-custom-checkpoint-loaders) +6. [Checkpoint I/O Policies](#checkpoint-io-policies) +7. [Startup Telemetry](#startup-telemetry) ## Overview @@ -76,6 +77,15 @@ Transforms weights between different naming conventions and applies model-specif ## Built-in Checkpoint Formats +Select a format with `checkpoint_format` in the LLM API or the YAML config: + +| `checkpoint_format` | Checkpoint it loads | +|---|---| +| `HF` (default) | Hugging Face checkpoint | +| `mistral` | Mistral native checkpoint | +| `mistral_large_3` | Mistral Large 3 native checkpoint | +| `MX` | ModelExpress peer-to-peer (P2P) weight transfer with Hugging Face checkpoint as fallback | + ### HuggingFace Format Currently, HF checkpoint loader is the primary built-in format, supporting: @@ -83,197 +93,27 @@ Currently, HF checkpoint loader is the primary built-in format, supporting: - **Weights loading** (`.safetensors/.bin/.pth`) - Loading HF compatible weights from disk - **Configuration parser** - Parsing HF stored configuration information to TRTLLM `ModelConfig` object - **Weights Mapping** - Converting HF weights into TRTLLM compatible representation +- **Weights Streaming** - Support experimental weight streaming options (see [Checkpoint I/O Policies](#checkpoint-io-policies)) -### ModelExpress (MX) Loading Path - -The PyTorch backend can use ModelExpress (MX) for peer-to-peer weight transfer -from a running TensorRT-LLM source instance before falling back to Hugging Face -checkpoint loading. Selecting MX does not require an MX-specific on-disk -checkpoint or conversion of the Hugging Face checkpoint. For installation, MX -service deployment, and configuration details, see -[ModelExpress (MX) Checkpoint Loading](./model-express.md). +### Mistral Native Format -## Checkpoint I/O Policies +`checkpoint_format: mistral` loads checkpoints in Mistral's own layout: a +`params.json` config and `consolidated*.safetensors` weights. -Checkpoint format describes how weights are represented and mapped. The -experimental `checkpoint_io_policy` independently controls how file-backed -checkpoint bytes reach host memory: - -- `auto` (default) selects rank-striped read-ahead for compatible built-in - PyTorch/HF loads and selects native I/O for other configurations. -- `native` always preserves the existing checkpoint loader and its synchronous - SafeTensors prefetch when the load group is eligible. -- `rank_striped_read_ahead` divides SafeTensors files into fixed extents. - Node-local ranks issue disjoint background `pread` requests into the Linux - page cache while the existing mapping, transformation, and H2D path runs. - The page cache can drop those pages before materialization reads them (for - example, when a Lustre client cancels aged DLM locks); when the node then - takes major page faults during materialization, the readers re-read their - extents. After the first refresh, the minimum delay from a completed refresh - to the next refresh is 10 seconds, then 20, 40, and at most 60 seconds. - Faults are still sampled every 5 seconds during the cooldown; a qualifying - burst queues one refresh for when the cooldown ends. The delay limits - refresh frequency, not the total number of refreshes. A warning is logged - once when the delay reaches 60 seconds. Before each refresh, the reader - rechecks available host memory against the full checkpoint size and startup - headroom. If admission fails or the memory probe fails, the refresh remains - pending and the check is retried on later polls. A warning is logged once - per deferral episode. This is a conservative check, not a memory reservation: - resident pages charged to a cgroup can cause deferral even when only a small - part of the checkpoint is missing. Normal weight materialization continues - while refresh is deferred. After every rank finishes - materialization, speculative read-ahead work that is still queued or active - can no longer improve first-token latency and is stopped. - A worker already blocked in a synchronous `pread` must finish its - current 8 MiB read before shutdown, so degraded storage can delay the join. - -```yaml -checkpoint_format: HF -load_format: auto -checkpoint_io_policy: rank_striped_read_ahead -``` +### Mistral Large 3 Format -The optimized path currently requires identical policy configuration across -all ranks, the automatically constructed built-in HF loader -(`checkpoint_loader` must be unset), `load_format: auto`, SafeTensors, and an -active MPI model-load communicator for distributed jobs. Static incompatibility -such as MX, a custom or explicitly provided loader, a non-automatic -load format, or known partial-model loading selects native I/O before any -rank-striped communicator or reader setup. An explicit incompatible -`rank_striped_read_ahead` request emits a warning instead of failing startup; -`auto` records the native selection at info level. - -Sessionless compatibility is a separate structural fallback: - -- Direct `load_weights()` calls use native I/O because they cannot keep - `open_weight_session()` alive through model materialization. Current callers - include separate draft/MTP checkpoint loading and the GMS restore path; the - primary built-in HF path remains eligible for rank-striped read-ahead. This - structural fallback is reported at info level even for an explicit - `rank_striped_read_ahead` request. - -Checkpoint-dependent eligibility remains a coordinated preflight. Lazy Kimi -loading, raw-weight caching, layer overrides, insufficient host-memory -headroom, and `.bin`/`.pth` select native materialization before readers start -or the model is mutated. These pre-activation runtime fallbacks log at info -level for `auto` and warn for an explicit `rank_striped_read_ahead` request. -With a valid node communicator, fallback preserves native prefetch collectives -on that load group; if communicator setup failed or its size did not match the -model-load mapping, native loading safely skips prefetch. -Memory admission reserves cgroup-aware startup headroom. Reader setup failures -also clean up and fall back before mapping; mapping, transformation, or H2D -failure after activation never retries a partially mutated model. A later -advisory read failure keeps successfully materialized weights. - -Read-ahead is intentionally not TP/PP/CP/EP-selective: each active load group -may warm one logical checkpoint copy per node, although the work stops when -materialization completes. The 64-reader budget is per active load group per -node, not a host-wide arbitration mechanism across colocated independent -TRT-LLM instances. Policy logs distinguish requested, selected, activated, and -effective policy, so `auto` selection and native fallback remain observable; -activation logs also report the local reader assignment. - -### CI startup experiment - -Pre-merge perf-sanity runs that upload telemetry explicitly assign eligible -built-in PyTorch/HF launches to the default `auto` policy, so merge decisions do -not depend on an unrelated CI identifier. Post-merge telemetry runs use the -numeric root Jenkins build number modulo 2 to assign `native` to even builds and -`auto` to odd builds before generated server configs are written. This gives a -50/50 assignment across consecutive root builds without adding CI stages or -model launches; successful observations may not be evenly split. Every rank -and node in one distributed startup receives the same concrete policy. The -experiment version is `checkpoint-io-v4-postmerge-50-auto-50-native`, distinct -from the previous 75/25 assignment. -Non-telemetry runs, invalid or missing post-merge build identities, -incompatible configurations, and purpose-built configs with an explicit -`checkpoint_io_policy` are not assigned and retain their normal configuration, -whose checkpoint I/O default remains `auto`. - -Set `TRTLLM_PERF_SANITY_CHECKPOINT_IO_POLICY=auto` or `native` to reproduce an -experiment arm without bucket assignment. Uploaded rows record the experiment -version, bucket, assigned arm, assignment source, PR number or root build -number, actual checkpoint loader/source metadata, -requested/selected/activated/effective policy, and a bounded fallback category -and reason. The primary comparison is intent-to-treat across post-merge rows -whose assignment source is `postmerge_build_number`: an `auto` assignment that -executes native I/O remains an auto fallback, not a native control. Pre-merge -rows use `premerge_default` and are excluded from that comparison. -Multiple client rows can share one server startup; use -`s_startup_observation_id` to group them and filter -`b_startup_observation_primary_row:true` to count each startup once. - -To compare rank-striped loading against native loading, separately report -confirmed effective `rank_striped_read_ahead` observations and native controls -with matching runtime-image identity, GPU, model/checkpoint, and parallelism -configuration. Keep auto fallbacks and unknown effective policies separate; -do not infer activation from the assigned arm. Retain root build IDs when -grouping repeated observations. This effective-policy comparison is -observational, not a controlled head-to-head benchmark. - -Keep results from the old 75/25 split separate from the new 50/50 split using -the experiment version, not bucket numbers alone. To check the split, use the -assigned policy; to compare loading performance, use the loader that actually -ran. - -Even builds use native loading, while odd builds use auto. Before comparing -performance, check that both groups use comparable runtime images, hardware, -CI configurations, and cache conditions—otherwise, those differences could be -mistaken for a loader improvement. +`checkpoint_format: mistral_large_3` reads the same layout as `mistral`. It +uses the Mistral Large 3 weight mapper for every architecture in the +checkpoint, including the multimodal `PixtralForConditionalGeneration` +wrapper. -The primary model's finalized policy is also published as -`startup_metrics.model_loader.checkpoint_io_policy` in `/server_info`. -This detached snapshot is captured after weight-session finalization, before a -separate draft checkpoint can reuse the loader. It preserves activation even -when a subsequent advisory read failure makes the effective policy native. -HF-derived formats such as Mistral and MX also report observed native disk -loads; this does not enable rank-striped read-ahead for those formats. -GMS writers using the built-in HF disk path capture the native policy after -target materialization, before draft loading; GMS readers remain unknown. -Unsupported, preloaded, or unobserved paths report an explicit unknown (`null`), -not an inferred native policy. The collector prefers this structured primary -status; only older runtimes without the field use `legacy_logs` fallback. -Draft statuses do not override the structured primary status. - -Use the following fields to qualify a comparison (with `ctx_` or `gen_` -prefixes for disaggregated server roles): - -- `s_checkpoint_io_policy_source`: `server_info`, `legacy_logs`, or `unknown`. -- `b_checkpoint_io_policy_complete`: every expected server provided one valid - structured primary policy. Legacy log attribution is not considered complete. -- `b_startup_metrics_timing_complete`: every expected server provided finite, - nonnegative preparation, population, finalization, and total model-loading - times. The existing collection-complete flag only checks successful fetches - and server counts; it does not imply these fields were present. - -CI records include image provenance so startup comparisons can account for -runtime-environment differences instead of attributing them to the loader. -This describes the benchmark environment, not which images customers deploy. - -The existing `s_runtime_image`, `s_runtime_image_source`, -`s_runtime_image_digest`, and `s_runtime_image_identity_strength` fields -describe the original OCI/base-image reference selected for the test -container, not the Jenkins dispatcher image. A valid digest-pinned reference -reports `digest_pinned` for that OCI reference only; a tagged reference reports -`reference_only` with an unknown digest. Missing provenance remains `unknown`. - -Enroot may launch a derived fat `.sqsh` with dependencies and TRT-LLM already -installed. These fields do not identify or verify that derived artifact or -the environment after package installation. A tag hash or `.sqsh` cache -filename is not a content digest. Group comparisons by the recorded base -reference as well as model/configuration, GPU and parallelism, but do not -treat matching base references as proof of identical runtime environments. -Final launch-artifact provenance remains a separate follow-up. - -These additions do not change loader policy or add startup collectives. Timing -remains rank-0-based per server, and failed launches are not comprehensively -persisted. Native counterfactual eligibility, checkpoint fingerprints, -slowest-rank timing, and process-launch-to-first-token measurement remain -separate follow-ups. +### ModelExpress (MX) Loading Path -This policy remains separate from ModelStreamer, MX, GMS, or snapshot -integrations. Those systems may change the source or bypass raw loading without -requiring a new combined checkpoint format. +ModelExpress (MX) is a P2P weight transfer service that lets one GPU transfer +its loaded weights to another, skipping slow weight loading on the second GPU from storage. +This checkpoint option checks for MX service first for P2P weight transfer, and +falls back to loading from HuggingFace checkpoint if MX is not available. +For MX service details, see [ModelExpress (MX) Checkpoint Loading](./model-express.md). ## Using Checkpoint Loaders @@ -367,6 +207,8 @@ class CustomCheckpointLoader(BaseCheckpointLoader): def get_default_config_loader(self) -> BaseConfigLoader: return CustomConfigLoader() + + ... ``` ### Step 2: Create the Checkpoint Weight Loader @@ -517,3 +359,132 @@ class CustomWeightMapper(BaseWeightMapper) ``` By setting the model name, the registered mapper will be associated with the specific model. + +## Checkpoint I/O Policies + +The experimental `checkpoint_io_policy` option controls how the built-in HF +loader reads SafeTensors files from storage into GPU memory. It provides +options to speed up weight loading via efficient streaming from storage. + +- `native`: It reads the checkpoint files into the OS page cache if there is + enough CPU memory. It then opens each tensor via mmap and loads them to GPU. + There could be many read contentions. +- `rank_striped_read_ahead`: It divides the checkpoint files into pieces and + reads each piece into the OS page cache while copying tensors to GPU at the + same time. This assumes most tensors being copied are from the OS page cache. + It runs only when the whole checkpoint fits in the available host memory, + minus some headroom. Otherwise it falls back to `native`. +- `auto` (default): use `rank_striped_read_ahead` when the load supports it, + and `native` I/O otherwise. + +Apart from built-in HF checkpoint and built-in checkpoint loader, to enable +non-native I/O policy, the following must all be true: +- `load_format` is `auto`. +- All of the model is set to be loaded (no partial model loading). +- The checkpoint uses SafeTensors files. + +### Re-reads in `rank_striped_read_ahead` + +The OS can drop read-ahead pages from the page cache before their tensors are +copied to GPU, for example when a Lustre client cancels its aged DLM locks. +When the node then takes major page faults while copying tensors to GPU, the +reader threads read their pieces into the page cache again. Each such re-read +is a refresh. + +- Wait between refreshes: after the first refresh, the next refresh starts at + least 10 seconds after the previous one completes. The wait then grows to 20 + and 40 seconds, up to 60 seconds. A warning is logged once when the wait + reaches 60 seconds. The wait limits how often refreshes run, not how many + refreshes run. +- Fault checks during the wait: the reader still checks page faults every 5 + seconds. A burst of faults queues one refresh, which starts when the wait + ends. +- Memory check: before each refresh, the reader checks again that the + available host memory can hold the full checkpoint plus the startup + headroom. If the check fails, or the reader cannot read the memory status, + the refresh stays queued and the reader checks again at later polls. A + warning is logged once each time a refresh starts to wait on memory. This + check is conservative and does not reserve memory: pages that are already in + memory and charged to a cgroup can defer a refresh even when only a small + part of the checkpoint is missing from the page cache. +- Tensor copies to GPU continue while a refresh waits. + +### Logs + +Each rank logs one line per checkpoint read (draft checkpoint loading has +its own line): + +```text +Checkpoint I/O policy: requested=auto, selected=rank_striped_read_ahead, activated=True, effective=rank_striped_read_ahead, fallback_reason=none. +``` + +- `requested`: the configured policy. +- `selected`: the policy chosen by TRTLLM before loading started. Various checks like +config incompatibility could lead to `native` fallback. +- `activated`: whether the background reader threads in `rank_striped_read_ahead` +started. +- `effective`: the final policy that read the checkpoint. If there is a runtime issue +like communication error, it could fall back to `native`. +- `fallback_reason`: why loading switched to native I/O. + +The `/server_info` response also reports these five fields in +`startup_metrics.model_loader`: `checkpoint_io_policy` for the target +checkpoint read, and `draft_checkpoint_io_policy` for a separate draft +checkpoint read. The value is `null` when the built-in HF weight loader did not +read the checkpoint from storage, for example with preloaded weights, a custom +weight loader, or a GMS reader. + +## Startup Telemetry + +CI perf-sanity runs that upload telemetry record how each server loaded its +checkpoint, so startup comparisons can account for the loader that ran. + +The primary model's finalized policy is also published as +`startup_metrics.model_loader.checkpoint_io_policy` in `/server_info`. +This detached snapshot is captured after weight-session finalization, before a +separate draft checkpoint can reuse the loader. It preserves activation even +when a subsequent advisory read failure makes the effective policy native. +HF-derived formats such as Mistral and MX also report observed native disk +loads; this does not enable rank-striped read-ahead for those formats. +GMS writers using the built-in HF disk path capture the native policy after +target materialization, before draft loading; GMS readers remain unknown. +Unsupported, preloaded, or unobserved paths report an explicit unknown (`null`), +not an inferred native policy. The collector prefers this structured primary +status; only older runtimes without the field use `legacy_logs` fallback. +Draft statuses do not override the structured primary status. + +Use the following fields to qualify a comparison (with `ctx_` or `gen_` +prefixes for disaggregated server roles): + +- `s_checkpoint_io_policy_source`: `server_info`, `legacy_logs`, or `unknown`. +- `b_checkpoint_io_policy_complete`: every expected server provided one valid + structured primary policy. Legacy log attribution is not considered complete. +- `b_startup_metrics_timing_complete`: every expected server provided finite, + nonnegative preparation, population, finalization, and total model-loading + times. The existing collection-complete flag only checks successful fetches + and server counts; it does not imply these fields were present. + +CI records include image provenance so startup comparisons can account for +runtime-environment differences instead of attributing them to the loader. +This describes the benchmark environment, not which images customers deploy. + +The existing `s_runtime_image`, `s_runtime_image_source`, +`s_runtime_image_digest`, and `s_runtime_image_identity_strength` fields +describe the original OCI/base-image reference selected for the test +container, not the Jenkins dispatcher image. A valid digest-pinned reference +reports `digest_pinned` for that OCI reference only; a tagged reference reports +`reference_only` with an unknown digest. Missing provenance remains `unknown`. + +Enroot may launch a derived fat `.sqsh` with dependencies and TRT-LLM already +installed. These fields do not identify or verify that derived artifact or +the environment after package installation. A tag hash or `.sqsh` cache +filename is not a content digest. Group comparisons by the recorded base +reference as well as model/configuration, GPU and parallelism, but do not +treat matching base references as proof of identical runtime environments. +Final launch-artifact provenance remains a separate follow-up. + +These additions do not change loader policy or add startup collectives. Timing +remains rank-0-based per server, and failed launches are not comprehensively +persisted. Native counterfactual eligibility, checkpoint fingerprints, +slowest-rank timing, and process-launch-to-first-token measurement remain +separate follow-ups. diff --git a/docs/source/features/sleep-wakeup.md b/docs/source/features/sleep-wakeup.md index 89d18330a4a2..01d1e40c341c 100644 --- a/docs/source/features/sleep-wakeup.md +++ b/docs/source/features/sleep-wakeup.md @@ -355,7 +355,7 @@ Principal implementation locations are: - Workers: `tensorrt_llm/executor/base_worker.py` and `tensorrt_llm/executor/ray/gpu_worker.py` - Allocation tagging: `tensorrt_llm/_torch/pyexecutor/py_executor_creator.py` - and `tensorrt_llm/_torch/pyexecutor/model_loader.py` + and `tensorrt_llm/_torch/models/checkpoints/memory_init.py` - Python VMM wrapper: `tensorrt_llm/_torch/virtual_memory.py` - CUDA VMM: `cpp/tensorrt_llm/runtime/virtualMemory.cpp` and `cpp/include/tensorrt_llm/runtime/virtualMemory.h` diff --git a/docs/source/llm-api/index.md b/docs/source/llm-api/index.md index 8023aea59977..c7e868161146 100644 --- a/docs/source/llm-api/index.md +++ b/docs/source/llm-api/index.md @@ -182,8 +182,8 @@ The `py_executor.model_engine_creation_seconds` timing includes the total `model | `draft_checkpoint_preparation_seconds` | Checkpoint preparation time for draft weights loaded as part of the model loader. | | `draft_weight_population_seconds` | Weight population time for draft weights loaded as part of the model loader. | | `draft_checkpoint_finalization_seconds` | Checkpoint finalization time for draft weights loaded as part of the model loader. | -| `post_load_processing_seconds` | Time spent in format-specific hooks and model finalization, including post-load weight transformation, quantization and memory cleanup. | -| `total_model_loading_seconds` | Overall model construction and loading interval measured after checkpoint configuration validation. It includes the named phases below. | +| `post_load_processing_seconds` | Time spent in format-specific hooks and model finalization, including post-load weight transformation, quantization, memory cleanup and other work. | +| `total_model_loading_seconds` | Overall model loading interval, from loading and validating the checkpoint configuration to the end of post-load processing. It includes the named phases above. | `trtllm-serve` exposes the same rank-0 payload in the `startup_metrics` field of the `GET /server_info` response: diff --git a/docs/source/torch/features/checkpoint_loading.md b/docs/source/torch/features/checkpoint_loading.md deleted file mode 100644 index 2a54905f3a48..000000000000 --- a/docs/source/torch/features/checkpoint_loading.md +++ /dev/null @@ -1,332 +0,0 @@ -# Checkpoint Loading - -The PyTorch backend provides a flexible and extensible infrastructure for loading model checkpoints from different sources and formats, such as HuggingFace (HF) or custom formats, by implementing required components like the checkpoint's weight loader, mapper, and configuration parser. - -## Table of Contents -1. [Overview](#overview) -2. [Core Components](#core-components) -3. [Built-in Checkpoint Formats](#built-in-checkpoint-formats) -4. [Using Checkpoint Loaders](#using-checkpoint-loaders) -5. [Creating Custom Checkpoint Loaders](#creating-custom-checkpoint-loaders) - -## Overview - -The checkpoint loading design is built around a plugin-like architecture that is separated into four distinct components: - -- **Checkpoint Loaders**: Orchestrates the loading process for specific formats. -- **Config Loaders**: Handles model configuration parsing and validation. -- **Weight Loaders**: Manages the actual loading of model weights from storage into memory. -- **Weight Mappers**: Maps and transforms loaded weights to the TRTLLM model's definition. - -This modular design allows for easy extension to support new checkpoint formats while maintaining backward compatibility and performance optimizations. By separating checkpoint loading into four subcomponents, users can leverage existing implementations and introduce custom, checkpoint-specific components. - -To support a new checkpoint format, you must implement all four components. -If the format shares components with an existing framework (such as HF), you only need to implement the components that differ. - -## Core Components - -### BaseCheckpointLoader - -The `BaseCheckpointLoader` is the central interface for all checkpoint loading operations. It provides a unified API regardless of the underlying checkpoint format. This interface is responsible for holding and exposing all objects required for the loading and parsing process. - -**Key Methods:** -- `load_config(checkpoint_dir, **kwargs)`: Loads and returns a `ModelConfig` object -- `load_weights(checkpoint_dir, **kwargs)`: Loads and returns a dictionary of weights -- `get_initialized_weight_mapper(model, config)`: Returns a weight mapper initialized at runtime for the model -- `cleanup()`: Releases resources and cleans up internal state - -### BaseConfigLoader - -Loads model configurations from checkpoint directories and parses them into a TRTLLM `ModelConfig`: - -```python -from tensorrt_llm._torch.models.checkpoints.base_config_loader import BaseConfigLoader - -class CustomConfigLoader(BaseConfigLoader): - def load(self, checkpoint_dir: str, **kwargs) -> ModelConfig: - # Load and parse configuration from your custom format - pretrained_config = self._get_pretrained_config(checkpoint_dir, **kwargs) - - return ModelConfig(pretrained_config=pretrained_config, - ...) - - def _get_pretrained_config(self, checkpoint_dir, **kwargs): - ... - -``` - -### BaseWeightLoader - -Handles the loading of model weights from storage: - -```python -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import BaseWeightLoader - -class CustomWeightLoader(BaseWeightLoader): - def load_weights(self, checkpoint_dir: str) -> dict[str, Any]: - # Load weights from your custom format - # Return a dictionary mapping parameter names to tensors - return weights_dict -``` - -### BaseWeightMapper - -Transforms weights between different naming conventions and applies model-specific transformations to the TRTLLM model object. - -## Built-in Checkpoint Formats - -### HuggingFace Format - -Currently, the HF checkpoint loader is the primary built-in format and supports: - -- **Weights loading** (`.safetensors, .bin, .pth`): Load HF-compatible weights from disk -- **Configuration parser** - Parse configuration information stored by HF into a TRTLLM `ModelConfig` object -- **Weights Mapping** - Convert HF weights into a TRTLLM-compatible representation - -## Using Checkpoint Loaders - -### Basic Usage - -There are two main approaches for using checkpoint loading objects - -The first approach is through the llm-api, as shown in the following example: - -```python -from tensorrt_llm import LLM - -hf_model_dir = "llama-models-v2/llama-v2-13b-hf" - -llm = LLM(model=hf_model_dir) -``` - -In this example, the `HfCheckpointLoader` is selected by default. - -To explicitly set the checkpoint loader, specify the required checkpoint-specific loader: - -```python -from tensorrt_llm import LLM -from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import HfCheckpointLoader - -hf_model_dir = "llama-models-v2/llama-v2-13b-hf" - -llm = LLM(model=hf_model_dir, - checkpoint_loader=HfCheckpointLoader()) -``` - -Similarly, to use a basic checkpoint loader with a specific subcomponent, provide the desired subcomponent as needed: - -```python -from tensorrt_llm import LLM -from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import HfCheckpointLoader - -hf_model_dir = "llama-models-v2/llama-v2-13b-hf" - -llm = LLM(model=hf_model_dir, - checkpoint_loader=HfCheckpointLoader(weight_loader=MyCustomWeightLoader())) -``` - -In the second approach, you can directly use the individual checkpoint loading components: - -```python -from tensorrt_llm._torch.models.checkpoints.hf.gemma3_weight_mapper import \ - Gemma3HfWeightMapper -from tensorrt_llm._torch.models.modeling_gemma3 import Gemma3ForCausalLM - -gemma3 = Gemma3ForCausalLM(model_config) -weight_mapper = Gemma3HfWeightMapper() -weight_mapper.init_model_and_config(gemma3, model_config) -gemma3.load_weights(hf_gemma3.state_dict(), weight_mapper) -``` -## Creating Custom Checkpoint Loaders - -To support a new checkpoint format, implement all four components. This section provides minimal templates for each. - -### When to Create Custom Components - -- **Complete New Format**: Implement all four components to support a new checkpoint format -- **Custom Weight Storage**: Implement only a custom weight loader if you have a unique weight storage format (such as a custom binary format or database storage) -- **Custom Configuration**: Implement only a custom config loader if your configuration format cannot be parsed by existing loaders -- **Custom Weight Mapping**: Implement only a custom weight mapper if your model has unique weight naming or transformation requirements that are checkpoint-specific - -### Step 1: Create the Checkpoint Loader - -```python -from typing import Optional -from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader -from tensorrt_llm._torch.models.checkpoints.base_config_loader import BaseConfigLoader -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import BaseWeightLoader -from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper -from tensorrt_llm._torch.models.modeling_utils import register_checkpoint_loader - -@register_checkpoint_loader("CUSTOM_FORMAT") -class CustomCheckpointLoader(BaseCheckpointLoader): - def __init__(self, - *, - weight_loader: Optional[BaseWeightLoader] = None, - weight_mapper: Optional[BaseWeightMapper] = None, - config_loader: Optional[BaseConfigLoader] = None): - self._weight_loader = weight_loader or self.get_default_weight_loader() - self._config_loader = config_loader or self.get_default_config_loader() - self._weight_mapper = weight_mapper - self._checkpoint_format = "CUSTOM_FORMAT" # Set the checkpoint format name - - def get_default_weight_loader(self) -> BaseWeightLoader: - return CustomWeightLoader() - - def get_default_config_loader(self) -> BaseConfigLoader: - return CustomConfigLoader() -``` - -### Step 2: Create the Checkpoint Weight Loader - -```python -from typing import Any -from tensorrt_llm._torch.models.checkpoints.base_weight_loader import BaseWeightLoader -from tensorrt_llm._torch.models.modeling_utils import register_checkpoint_weight_loader - -@register_checkpoint_weight_loader("CUSTOM_FORMAT") -class CustomWeightLoader(BaseWeightLoader): - def load_weights(self, checkpoint_dir: str, **kwargs) -> dict[str, Any]: - """ - Load weights from your custom format. - - Args: - checkpoint_dir: Directory containing checkpoint files - **kwargs: Additional loading parameters - - Returns: - Dictionary mapping parameter names to tensors - """ - weights = {} # Implement your custom weight loading logic here - - # Examples: - # - Load from custom binary files - # - Load from databases - # - Load from compressed archives - # - Apply custom preprocessing - - return weights -``` - -### Step 3: Create the Checkpoint Config Loader - -```python -from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.checkpoints.base_config_loader import BaseConfigLoader -from tensorrt_llm._torch.models.modeling_utils import register_config_loader - -@register_config_loader("CUSTOM_FORMAT") -class CustomConfigLoader(BaseConfigLoader): - def load(self, checkpoint_dir: str, **kwargs) -> ModelConfig: - """ - Load and parse configuration from your custom format. - - Args: - checkpoint_dir: Directory containing configuration files - **kwargs: Additional loading parameters - - Returns: - ModelConfig object containing parsed configuration - """ - # Load your custom configuration format here - # Examples: - # - Parse YAML/TOML files - # - Convert from proprietary formats - - pretrained_config = self._load_pretrained_config(checkpoint_dir, **kwargs) - - return ModelConfig( - pretrained_config=pretrained_config, - # Add other ModelConfig parameters as needed - ) - - def _load_pretrained_config(self, checkpoint_dir: str, **kwargs): - """Load the raw configuration from your custom format.""" - # Implement as needed - pass -``` - -### Step 4: Create the Checkpoint Weight Mapper - -```python -from torch import nn -from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper -from tensorrt_llm._torch.models.modeling_utils import register_mapper - -@register_mapper("CUSTOM_FORMAT") -class CustomWeightMapper(BaseWeightMapper): - def __init__(self): - super().__init__() - # Define any weight transformation callbacks - self._callbacks = [ - # Add your custom weight transformation functions - # self._custom_transform_function, - ] - - def map_weights(self) -> None: - """ - Define mappings between source and target weight names. - """ - self.mapping.update({ - # Map source names to target names - # 'target_module_name': ['source_param1', 'source_param2'], - # For example: 'qkv_proj': ['q_proj', 'k_proj', 'v_proj'] - }) - - def apply_callbacks(self, module: nn.Module, module_name: str, - module_names_breakdown: list[str], - weights: dict) -> list[dict]: - """ - Apply weight transformations for modules that require special handling. - - Args: - module: The target module - module_name: The specific module name being processed - module_names_breakdown: Module path components - weights: Source weights dictionary - - Returns: - List of transformed weight dictionaries - """ - module_weights = [] - - for new_name in self._mapping[module_name]: - # Filter weights for this specific parameter - fw = self.filter_weights( - '.'.join(module_names_breakdown + [new_name]), weights) - - # Apply transformation callbacks - for callback in self._callbacks: - fw = callback(module, new_name, fw) - - module_weights.append(fw) - - return module_weights - - def should_skip_module(self, module_name: str) -> bool: - """ - Define which modules should be skipped during loading. - """ - # Add logic to skip specific modules based on your requirements - # Examples: - # - Skip LoRA-specific modules - # - Skip temporary/auxiliary modules - - return super().should_skip_module(module_name) -``` - -Note: When creating a custom mapper, you can define either a checkpoint-format-specific mapper. For example: - -```python -@register_mapper("CUSTOM_FORMAT") -class CustomWeightMapper(BaseWeightMapper) -``` - -Alternatively, you can define a checkpoint-model-specific mapper. For example: - -```python -@register_mapper("CUSTOM_FORMAT", "Gemma3ForCausalLM") -class CustomWeightMapper(BaseWeightMapper) -``` - -By setting the model name, the registered mapper will be associated with the specific model. diff --git a/tensorrt_llm/_torch/memory/gpu_memory_backend.py b/tensorrt_llm/_torch/memory/gpu_memory_backend.py index 8a0a74dced64..dcdf7ca11803 100644 --- a/tensorrt_llm/_torch/memory/gpu_memory_backend.py +++ b/tensorrt_llm/_torch/memory/gpu_memory_backend.py @@ -24,11 +24,12 @@ Design notes: - We deliberately avoid the upstream `setup_gms()` monkey-patch entry - point. TRT-LLM owns the integration policy in `ModelLoader`; this + point. TRT-LLM owns the integration policy in + `models/checkpoints/gms/model_loading.py`; this backend is a thin adapter over the GMS library's primitive API. - The protocol has a single point of contact per concern, so when the upstream GMS Python API drifts, only the methods on `GMSBackend` - need to change — the call sites in `model_loader.py` are stable. + need to change — the call sites in `gms/model_loading.py` are stable. Operating modes: - **RW (Read-Write)**: First worker loads weights via the normal @@ -120,8 +121,9 @@ class GMSBackend: This adapter calls the GMS library's stable per-call primitives (`GMSClientMemoryManager` + the helpers in `gpu_memory_service.client.torch.allocator` and `.module`) and - composes them into the methods TRT-LLM's `ModelLoader` invokes for - the `LoadFormat.GMS` branch. We intentionally do **not** use the + composes them into the methods that TRT-LLM's GMS loading code + (`models/checkpoints/gms/model_loading.py`) invokes for + `LoadFormat.GMS`. We intentionally do **not** use the upstream `setup_gms()` monkey-patch — TRT-LLM owns the integration policy and we keep the dependency surface narrow and explicit. """ diff --git a/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py index 0456c7d8530f..e16248fcbe53 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py @@ -2,8 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod -from collections.abc import Iterator -from contextlib import contextmanager from typing import Any from torch import nn @@ -73,12 +71,6 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, mapping=mapping, **kwargs) - @contextmanager - def open_weight_session(self, checkpoint_dir: str, mapping: Mapping, - **kwargs) -> Iterator[dict[str, Any]]: - """Keep loader-specific work alive while weights are materialized.""" - yield self.load_weights(checkpoint_dir, mapping=mapping, **kwargs) - def is_weights_preloaded(self) -> bool: """Whether the last load wrote weights directly into the model.""" return False @@ -141,3 +133,21 @@ def get_initialized_weight_mapper(self, model: nn.Module, weight_mapper.init_model_and_config(model, config) self.weight_mapper = weight_mapper return weight_mapper + + def get_initialized_draft_weight_mapper( + self, model: nn.Module, target_weight_mapper: BaseWeightMapper | None + ) -> BaseWeightMapper | None: + """Return the weight mapper for the separate draft checkpoint of ``model``. + + Each draft config initializes its own architecture mapper so + replacement MTP weights use the draft's quantization and MoE backend. + """ + if model.draft_config is None: + # Drafters without a separate config reuse the target mapper. + return target_weight_mapper + draft_model_arch = model.draft_config.pretrained_config.architectures[0] + draft_weight_mapper = AutoCheckpointMapper.get(self.checkpoint_format, + draft_model_arch) + draft_weight_mapper.init_model_and_config(model.draft_model, + model.draft_config) + return draft_weight_mapper diff --git a/tensorrt_llm/_torch/models/checkpoints/gms/__init__.py b/tensorrt_llm/_torch/models/checkpoints/gms/__init__.py new file mode 100644 index 000000000000..467079831e16 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/gms/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tensorrt_llm/_torch/models/checkpoints/gms/model_loading.py b/tensorrt_llm/_torch/models/checkpoints/gms/model_loading.py new file mode 100644 index 000000000000..9b18ce2b1655 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/gms/model_loading.py @@ -0,0 +1,485 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GMS helpers that `ModelLoader` calls for ``LoadFormat.GMS``. + +GMS (GPU Memory Service) keeps weight tensors in a node-shared GPU memory pool, +so several TRT-LLM instances can share the same weight bytes without copies. +The `GMSBackend` client itself lives in `tensorrt_llm._torch.memory`. +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Optional + +import torch + +from tensorrt_llm._errors import ConfigInputError +from tensorrt_llm._torch.models.checkpoints import memory_init, model_hooks +from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import ( + checkpoint_io_policy_metadata, +) +from tensorrt_llm._torch.models.checkpoints.mx import model_loading as mx_loading +from tensorrt_llm._torch.models.modeling_utils import ModelLoaderMetricNames, timing_metric +from tensorrt_llm._torch.weight_sharing import ( + IdentityCheckPolicy, + PostTransformQualificationDecision, + SourceIdentity, + check_weight_sharing_compatibility, +) +from tensorrt_llm.checkpoint_format import CheckpointFormat +from tensorrt_llm.llmapi.llm_args import LoadFormat +from tensorrt_llm.logger import logger + +if TYPE_CHECKING: + from tensorrt_llm._torch.memory import GMSBackend + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader + from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper + from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM + from tensorrt_llm.llmapi.llm_args import GmsConfig, TorchLlmArgs + from tensorrt_llm.mapping import Mapping + + +@dataclass +class GmsLoadResult: + """What a GMS load hands back to the model loader.""" + + # Whether the checkpoint loader filled the model directly (e.g. MX P2P). + # Readers report False. + weights_preloaded: bool + # Weight mapper the writer used to fill the model; None for readers and + # for writers whose checkpoint loader preloaded the weights. + weight_mapper: Optional["BaseWeightMapper"] + # Connected GMS client; the model loader keeps it until cleanup. + gms_backend: "GMSBackend" + + +def load_gms_weights( + model: "DecoderModelForCausalLM", + *, + gms_config: "GmsConfig", + mapping: "Mapping", + llm_args: "TorchLlmArgs", + checkpoint_dir: str, + checkpoint_loader: "BaseCheckpointLoader", + model_config: "ModelConfig", + is_meta_init: bool, + loads_draft_weights: bool, + speculative_model: Optional[str], + qualification: PostTransformQualificationDecision, + source_identity: Optional[SourceIdentity], + metrics: dict[str, float], + startup_metadata: dict[str, Any], +) -> GmsLoadResult: + """Fill ``model`` with weights for ``LoadFormat.GMS``, including all post-load hooks. + + GMS gives this rank one of two roles: + + - RW (writer): brings the model up inside the GMS pool and commits it, so + readers get the post-post_load layout. + - RO (reader): zero-copy maps the committed weights from the pool into the + model; no disk I/O and no per-instance copies. + + Both roles run the post-load hooks themselves, so the generic post-load + step in `ModelLoader.load()` must skip them. + + Args: + model: model built on meta or CUDA tensors + gms_config: GMS daemon socket, mode and tag + mapping: parallel mapping of this rank + llm_args: LLM args of this load + checkpoint_dir: checkpoint directory of the target model + checkpoint_loader: checkpoint loader of this load + model_config: final model config + is_meta_init: whether ``model`` was built on meta tensors + loads_draft_weights: whether a separate draft checkpoint is loaded + speculative_model: draft checkpoint path from the speculative config + qualification: MX post-transform profile decision for this load + source_identity: identity this rank built for its own model + metrics: model loader timing metrics, updated in place + startup_metadata: model loader startup metadata; a writer records the + checkpoint I/O policy of its target read in place + + Returns: + The load result for the model loader to keep. + + Raises: + RuntimeError: GMS connection fails, the writer gets no weights, or + the backend role is unset. + """ + gms_backend = _connect_gms_backend(gms_config, mapping) + try: + # `is_rw` is `Optional[bool]` on the GPU memory backend + # protocol: True = RW lock granted, False = RO lock granted, + # None = unset. After a successful `connect()` it must be + # True or False; `None` is treated as an adapter bug and + # surfaces as a hard error rather than silently falling + # through to the RO path. + if gms_backend.is_rw is True: + weights_preloaded, weight_mapper = _load_gms_writer_weights( + gms_backend, + model, + mapping=mapping, + llm_args=llm_args, + checkpoint_dir=checkpoint_dir, + checkpoint_loader=checkpoint_loader, + model_config=model_config, + is_meta_init=is_meta_init, + loads_draft_weights=loads_draft_weights, + speculative_model=speculative_model, + qualification=qualification, + source_identity=source_identity, + metrics=metrics, + startup_metadata=startup_metadata, + ) + return GmsLoadResult( + weights_preloaded=weights_preloaded, + weight_mapper=weight_mapper, + gms_backend=gms_backend, + ) + if gms_backend.is_rw is False: + _load_gms_reader_weights( + gms_backend, + model, + checkpoint_dir=checkpoint_dir, + checkpoint_loader=checkpoint_loader, + qualification=qualification, + source_identity=source_identity, + metrics=metrics, + ) + return GmsLoadResult( + weights_preloaded=False, weight_mapper=None, gms_backend=gms_backend + ) + raise RuntimeError( + f"GMS backend connected but lock state is unset " + f"(is_rw={gms_backend.is_rw!r}); expected True (RW) " + "or False (RO). This indicates a bug in the GMS " + "adapter or a protocol violation." + ) + # Use BaseException to catch all possible interruption to ensure GMS is always cleaned up + except BaseException: + gms_backend.cleanup() + raise + + +def _load_gms_writer_weights( + gms_backend: "GMSBackend", + model: "DecoderModelForCausalLM", + *, + mapping: "Mapping", + llm_args: "TorchLlmArgs", + checkpoint_dir: str, + checkpoint_loader: "BaseCheckpointLoader", + model_config: "ModelConfig", + is_meta_init: bool, + loads_draft_weights: bool, + speculative_model: Optional[str], + qualification: PostTransformQualificationDecision, + source_identity: Optional[SourceIdentity], + metrics: dict[str, float], + startup_metadata: dict[str, Any], +) -> tuple[bool, Optional["BaseWeightMapper"]]: + """Load and post-process weights inside the GMS pool, then commit them. + + The pool opens BEFORE meta-tensor materialization and to(cuda), so every + CUDA allocation needed to bring the model up (meta materialization, + to(cuda), checkpoint load, draft load, post_load hooks) lands in the GMS + pool. After the pool is closed, ``finalize_write`` commits the + post-post_load layout for RO peers to map zero-copy. This mirrors the + upstream Dynamo TRT-LLM shim pattern. + + Returns: + ``weights_preloaded`` as reported by the checkpoint loader, and the + weight mapper used to fill the model (None when preloaded). + """ + + weight_mapper: Optional["BaseWeightMapper"] = None + with gms_backend.mem_pool_scope(torch.device("cuda")): + # Materialize meta tensors inside the pool. `ModelLoader.load()` skips + # `memory_init.allocate_model_on_cuda` for GMS so it never ran. + if is_meta_init: + memory_init.init_meta_tensors_on_cuda(model) + + # Catch-all for any tensors not on CUDA yet. + model.to("cuda") + + # load checkpoint weights + weight_source = ( + model.llm_checkpoint_dir if hasattr(model, "llm_checkpoint_dir") else checkpoint_dir + ) + # Pass model= as the LoadFormat.AUTO target read does: the HF loader + # uses it to keep a context-only worker's unused tensors lazy, and MX + # writes P2P weights straight into the model's parameters. + load_weights_kwargs = { + "mapping": mapping, + "model": model, + "source_identity": source_identity, + } + if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: + load_weights_kwargs.update( + mx_loading.get_mx_load_weights_kwargs( + model_config=model_config, + llm_args=llm_args, + qualification=qualification, + ) + ) + with timing_metric(ModelLoaderMetricNames.CHECKPOINT_PREPARATION_SECONDS.value, metrics): + weights = checkpoint_loader.load_weights(weight_source, **load_weights_kwargs) + + weights_preloaded = checkpoint_loader.is_weights_preloaded() + _check_gms_writer_weights_loaded(checkpoint_loader, weights, weights_preloaded) + if weights: + weight_mapper = checkpoint_loader.get_initialized_weight_mapper(model, model_config) + with timing_metric(ModelLoaderMetricNames.WEIGHT_POPULATION_SECONDS.value, metrics): + model_hooks.call_load_weights(model.load_weights, weights, weight_mapper) + if torch.cuda.is_available(): # CPU guard + torch.cuda.synchronize() + # Take the policy before the draft read with the same loader resets it. + startup_metadata["checkpoint_io_policy"] = checkpoint_io_policy_metadata(checkpoint_loader) + + if loads_draft_weights: + _load_gms_writer_draft_weights( + model, + checkpoint_loader, + speculative_model=speculative_model, + mapping=mapping, + target_weight_mapper=weight_mapper, + metrics=metrics, + ) + + # run post-load hooks + with timing_metric(ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, metrics): + # Run post_load hooks INSIDE the pool so any tensors they create + # or rebind (fused QKV, quantization scales, derived aliases) land + # in GMS and become part of the committed layout that RO peers + # receive. + checkpoint_loader.post_load_apply(model, weights_preloaded=weights_preloaded) + mx_loading.run_post_load_walks( + checkpoint_loader, + model, + weights_preloaded=weights_preloaded, + qualification=qualification, + ) + _move_stray_params_into_gms_pool(gms_backend, model) + mx_loading.post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + qualification=qualification, + source_identity=source_identity, + ) + + # Pool closed. Commit the post-post_load layout. + with timing_metric(ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, metrics): + gms_backend.finalize_write(model) + logger.info( + "LoadFormat.GMS (RW): loaded and committed weights via %s", + checkpoint_loader.checkpoint_format, + ) + return weights_preloaded, weight_mapper + + +def _load_gms_writer_draft_weights( + model: "DecoderModelForCausalLM", + checkpoint_loader: "BaseCheckpointLoader", + *, + speculative_model: Optional[str], + mapping: "Mapping", + target_weight_mapper: Optional["BaseWeightMapper"], + metrics: dict[str, float], +) -> None: + """Load draft/MTP weights from ``speculative_model`` into the one-engine model.""" + + draft_load_kwargs = mx_loading.get_mx_draft_load_weights_kwargs(checkpoint_loader) + with timing_metric(ModelLoaderMetricNames.DRAFT_CHECKPOINT_PREPARATION_SECONDS.value, metrics): + draft_weights = checkpoint_loader.load_weights( + speculative_model, mapping=mapping, **draft_load_kwargs + ) + + draft_weight_mapper = checkpoint_loader.get_initialized_draft_weight_mapper( + model, target_weight_mapper + ) + with timing_metric(ModelLoaderMetricNames.DRAFT_WEIGHT_POPULATION_SECONDS.value, metrics): + model_hooks.call_load_weights(model.load_draft_weights, draft_weights, draft_weight_mapper) + if torch.cuda.is_available(): # CPU guard + torch.cuda.synchronize() + + +def _load_gms_reader_weights( + gms_backend: "GMSBackend", + model: "DecoderModelForCausalLM", + *, + checkpoint_dir: str, + checkpoint_loader: "BaseCheckpointLoader", + qualification: PostTransformQualificationDecision, + source_identity: Optional[SourceIdentity], + metrics: dict[str, float], +) -> None: + """Zero-copy map the committed GMS weights into ``model``. + + Weights come from a GMS donor that has already committed the + post-post_load layout, so the reader passes ``weights_preloaded=True`` to + the checkpoint-loader hooks. This prevents duplicate transforms while + leaving any reader-side publication policy to the format-specific hook. + + Hook order: + + 1. ``post_load_apply``: format-specific apply work (e.g. MX preshard + markers). + 2. Per-module ``setup_aliases``: creates structural aliases BEFORE + ``materialize_module`` walks the final module tree (including + ``draft_model`` for spec dec). + 3. SourceIdentity gate: STRICT pre-materialize compatibility check (GMS + has no disk fallback). + 4. ``materialize_module``: zero-copy bind GMS pool storage onto the model + parameters. + 5. Mark materialized weights as transformed, then recompute Python-side + derived state without re-running one-shot transforms. + 6. ``post_load_publish``: apply the format-specific reader publication + policy. + """ + + with timing_metric(ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, metrics): + checkpoint_loader.post_load_apply(model, weights_preloaded=True) + model_hooks.setup_aliases(model) + check_gms_source_identity(source_identity, gms_backend) + gms_backend.materialize_module(model) + model_hooks.mark_weights_transformed(model) + model_hooks.walk_cache_state(model) + mx_loading.post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=True, + qualification=qualification, + source_identity=source_identity, + ) + logger.info("LoadFormat.GMS (RO): materialized weights") + + +def check_gms_locality_domain_compat(load_format: LoadFormat, config: "ModelConfig") -> None: + """Raise when LoadFormat is GMS but the model uses locality domain localized weights. + + Raises: + ConfigInputError: locality domain execution is enabled in ``config``. + """ + if load_format == LoadFormat.GMS: + locality_domain_policy = getattr(config, "locality_domain_policy", None) + if getattr(locality_domain_policy, "enabled", False): + raise ConfigInputError( + "LoadFormat.GMS is incompatible with locality domain localized weights. " + "GMS shares registered parameters, while locality domain execution " + "releases those parameters after creating partition-local " + "weight shards." + ) + + +def _connect_gms_backend(gms_config: "GmsConfig", mapping: "Mapping") -> "GMSBackend": + """Create a GMS client and connect it to the GMS daemon. + + After a successful connect, ``is_rw`` tells the role this rank got: True + for the writer (RW) and False for a reader (RO). + + Raises: + RuntimeError: the client cannot connect to the daemon. + """ + # Imported lazily so the optional `gpu_memory_service` dependency only + # loads when GMS is actually selected. + from tensorrt_llm._torch.memory import GMSBackend + + gms_backend = GMSBackend( + socket_path=gms_config.socket_path, + mapping=mapping, + mode=gms_config.mode, + tag=gms_config.tag, + ) + if not gms_backend.connect(): + raise RuntimeError(f"Failed to connect to GMS at {gms_config.socket_path}") + return gms_backend + + +def _check_gms_writer_weights_loaded( + checkpoint_loader: "BaseCheckpointLoader", + weights: Optional[dict[str, Any]], + weights_preloaded: bool, +) -> None: + """Refuse to commit a model that the checkpoint loader did not fill. + + ``weights`` from ``checkpoint_loader.load_weights`` may be: + + - a non-empty dict: the standard mapping pipeline fills the model. + - empty or None with ``weights_preloaded`` True: the loader filled the + model directly (e.g. MX P2P), so no mapping is needed. + - empty or None with ``weights_preloaded`` False: nothing filled the + model, so it would be committed uninitialized. + + Per-tensor partial loads (missing keys in a non-empty dict) are caught + later by the strict checks in ``model.load_weights``. + + Raises: + RuntimeError: the loader returned no weights and did not preload. + """ + if weights or weights_preloaded: + return + raise RuntimeError( + f"GMS RW: checkpoint loader " + f"'{checkpoint_loader.checkpoint_format}' " + f"returned empty weights ({weights!r}) and " + "did not preload the model. Refusing to " + "commit an unpopulated model to the GMS " + "pool." + ) + + +def _move_stray_params_into_gms_pool( + gms_backend: "GMSBackend", model: "DecoderModelForCausalLM" +) -> None: + """Move parameters allocated outside the GMS pool into it, then free the originals. + + This is a last-mile sweep after the post-load hooks. It catches tensors + from C++ ops that bypassed the active torch allocator (e.g. native + cudaMalloc). + """ + gms_backend.move_untracked_params(model) + # Safe with active GMS mappings: GMSBackend.connect() installed upstream's + # patch_empty_cache(), which makes torch.cuda.empty_cache() skip GMS-backed + # VMM regions. This frees the non-GMS originals dropped by + # move_untracked_params (tensor.data was rebound to GMS-backed + # replacements) before commit, so the cached size does not show as live + # in memory accounting. + torch.cuda.empty_cache() + + +def check_gms_source_identity( + local_source_identity: Optional[SourceIdentity], gms_backend: "GMSBackend" +) -> None: + """Pre-materialize SourceIdentity gate for the GMS read-only path. + + GMS has no disk-fallback path, so a missing or mismatched identity + raises under `IdentityCheckPolicy.STRICT` rather than falling back. + + Args: + local_source_identity: identity this reader built for its own model + gms_backend: connected GMS backend (RO role) + + Raises: + SourceIdentityMismatchError: the writer's identity is missing or + incompatible with this reader's identity. + """ + check_weight_sharing_compatibility( + local_source_identity, + gms_backend.get_source_identity(), + IdentityCheckPolicy.STRICT, + ) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index 5f85aa7958aa..e7c94e3a8842 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -1,8 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from collections.abc import Iterator -from contextlib import contextmanager from typing import Any, Optional from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ @@ -61,7 +59,16 @@ def cleanup(self) -> None: def get_default_weight_loader(self) -> HfWeightLoader: return HfWeightLoader() - def _uses_context_only_lazy_loading(self, model: Any) -> bool: + def uses_context_only_lazy_loading(self, model: Any) -> bool: + """Whether reading for this model keeps the tensors it does not own lazy. + + Only the built-in HF checkpoint and weight loaders do this, and only + for a context-only worker. Such a read must not read the whole + checkpoint ahead. + + Args: + model: target model passed to the read, or None + """ return (type(self) is HfCheckpointLoader and type(self.weight_loader) is HfWeightLoader and getattr(getattr(model, "model", None), @@ -70,7 +77,7 @@ def _uses_context_only_lazy_loading(self, model: Any) -> bool: def load_weights(self, checkpoint_dir: str, mapping: Mapping, **kwargs) -> dict[str, Any]: """Keep omitted decoder tensors lazy on a context-only worker.""" - if self._uses_context_only_lazy_loading(kwargs.get("model")): + if self.uses_context_only_lazy_loading(kwargs.get("model")): # The model remapper drops unowned keys before touching their # slices. Eager loading or read-ahead would read those bytes anyway. loader = self.weight_loader @@ -88,28 +95,6 @@ def load_weights(self, checkpoint_dir: str, mapping: Mapping, return weights return super().load_weights(checkpoint_dir, mapping=mapping, **kwargs) - @contextmanager - def open_weight_session(self, checkpoint_dir: str, mapping: Mapping, - **kwargs) -> Iterator[dict[str, Any]]: - """Delegate the optimized session only for the built-in HF path.""" - if (type(self) is HfCheckpointLoader - and type(self.weight_loader) is HfWeightLoader - and self.weight_loader.checkpoint_io_policy - != CheckpointIOPolicy.NATIVE and - not self._uses_context_only_lazy_loading(kwargs.get("model"))): - with self.weight_loader.open_weight_session(checkpoint_dir, - mapping=mapping, - **kwargs) as weights: - yield weights - return - - # MX, Mistral, and custom subclasses retain their polymorphic - # load_weights implementations. - with super().open_weight_session(checkpoint_dir, - mapping=mapping, - **kwargs) as weights: - yield weights - def get_default_config_loader(self) -> HfConfigLoader: return HfConfigLoader() @@ -143,3 +128,31 @@ def config_loader(self) -> Optional[BaseConfigLoader]: @property def checkpoint_format(self) -> CheckpointFormat: return self._checkpoint_format + + +def checkpoint_io_policy_metadata( + checkpoint_loader: BaseCheckpointLoader +) -> dict[str, str | bool | None] | None: + """Snapshot observed policy for loads using the built-in HF weight loader. + + Args: + checkpoint_loader: Loader whose completed weight load is observed. + + Returns: + A detached policy snapshot, or None for unsupported or unobserved loads. + """ + if (not isinstance(checkpoint_loader, HfCheckpointLoader) + or type(checkpoint_loader.weight_loader) is not HfWeightLoader + or checkpoint_loader.is_weights_preloaded()): + return None + status = checkpoint_loader.weight_loader.last_checkpoint_io_status + if status.effective not in (CheckpointIOPolicy.NATIVE, + CheckpointIOPolicy.RANK_STRIPED_READ_AHEAD): + return None + return { + "requested": status.requested, + "selected": status.selected, + "activated": status.activated, + "effective": status.effective, + "fallback_reason": status.fallback_reason, + } diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 2c403c5b0a1a..921ca035de18 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -103,6 +103,14 @@ class _CheckpointIOStatus: activated: bool = False effective: str = "none" fallback_reason: str | None = None + # Time `open_weight_session` took to hand out the weights: policy checks + # across ranks, file mapping and starting the background readers, or the + # native fallback load. + preparation_seconds: float = 0.0 + # Time `open_weight_session` took to close: waiting for every rank and + # stopping the background readers. Always 0 without an active + # read-ahead session. + finalization_seconds: float = 0.0 @register_checkpoint_weight_loader(CheckpointFormat.MX) @@ -177,6 +185,28 @@ def checkpoint_io_policy(self) -> CheckpointIOPolicy: def last_checkpoint_io_status(self) -> _CheckpointIOStatus: return replace(self._last_checkpoint_io_status) + def record_session_timings(self, + *, + preparation_seconds: float, + finalization_seconds: float = 0.0) -> None: + """Record how long the last weight session took to start and to finish. + + This loader's ``open_weight_session`` calls this, and + ``last_checkpoint_io_status`` reports the times. The status object is + replaced, not changed in + place: a subclass that does not call ``__init__`` shares the + class-level default status with all other such subclasses. + + Args: + preparation_seconds: time the session took to hand out the weights + finalization_seconds: time the session took to close; 0 when the + session has nothing to finish + """ + self._last_checkpoint_io_status = replace( + self._last_checkpoint_io_status, + preparation_seconds=preparation_seconds, + finalization_seconds=finalization_seconds) + def _reset_checkpoint_io_status(self) -> None: self._last_checkpoint_io_status = _CheckpointIOStatus( requested=self._requested_checkpoint_io_policy, @@ -499,70 +529,76 @@ def open_weight_session(self, mapping: Mapping, use_consolidated: bool = False, **kwargs) -> Iterator[dict[str, Any]]: - """Keep opt-in read-ahead alive through model materialization.""" + """Keep opt-in read-ahead alive through model materialization. + + Copy the yielded weights into the model inside the ``with`` block. + Every rank must enter and leave the block, because starting and + finishing read-ahead are collective. Afterwards, + `last_checkpoint_io_status` reports the policy that took effect and + how long the session took to start and to finish. + """ + preparation_started = time.perf_counter() self._reset_checkpoint_io_status() + status = self._last_checkpoint_io_status + + # prepare the weights: native load, native fallback or read-ahead + session: RankStripedReadAheadSession | None = None if self._checkpoint_io_policy == CheckpointIOPolicy.NATIVE: weights = self._load_weights_native(checkpoint_dir, mapping, use_consolidated, **kwargs) - self._last_checkpoint_io_status.effective = CheckpointIOPolicy.NATIVE + status.effective = CheckpointIOPolicy.NATIVE self._log_checkpoint_io_status() - yield weights - return - - active_communicator = self._active_communicator() - mapping_error = None - communicator_size = (None if active_communicator is None else - active_communicator.Get_size()) - if (communicator_size is not None - and mapping.world_size != communicator_size): - mapping_error = RuntimeError( - "Rank-striped read-ahead requires mapping.world_size to " - "match the active MPI communicator size " - f"({communicator_size}); received {mapping.world_size}.") - coordinated_mapping_error = coordinate_error( - active_communicator, - "rank-striped mapping validation", - mapping_error, - ) - if coordinated_mapping_error is not None: - self._last_checkpoint_io_status.selected = CheckpointIOPolicy.NATIVE - weights = self._fallback_to_native( - checkpoint_dir, - mapping, - use_consolidated, - str(coordinated_mapping_error), - active_communicator=active_communicator, - allow_native_prefetch=False, - **kwargs, - ) - yield weights - return - if mapping.world_size > 1 and active_communicator is None: - self._last_checkpoint_io_status.selected = CheckpointIOPolicy.NATIVE - weights = self._fallback_to_native( - checkpoint_dir, - mapping, - use_consolidated, - "distributed rank-striped read-ahead requires an active MPI communicator", - # Ray workers intentionally disable MPI. Preserve the native - # path's per-process read-ahead in that supported mode. - allow_native_prefetch=mpi_disabled(), - **kwargs, + else: + active_communicator = self._active_communicator() + mapping_error = None + communicator_size = (None if active_communicator is None else + active_communicator.Get_size()) + if (communicator_size is not None + and mapping.world_size != communicator_size): + mapping_error = RuntimeError( + "Rank-striped read-ahead requires mapping.world_size to " + "match the active MPI communicator size " + f"({communicator_size}); received {mapping.world_size}.") + coordinated_mapping_error = coordinate_error( + active_communicator, + "rank-striped mapping validation", + mapping_error, ) - yield weights - return - - weights, session = self._start_rank_striped_read_ahead( - checkpoint_dir, - mapping, - use_consolidated, - active_communicator, - **kwargs, - ) - if session is None: - yield weights - return + if coordinated_mapping_error is not None: + status.selected = CheckpointIOPolicy.NATIVE + weights = self._fallback_to_native( + checkpoint_dir, + mapping, + use_consolidated, + str(coordinated_mapping_error), + active_communicator=active_communicator, + allow_native_prefetch=False, + **kwargs, + ) + elif mapping.world_size > 1 and active_communicator is None: + status.selected = CheckpointIOPolicy.NATIVE + weights = self._fallback_to_native( + checkpoint_dir, + mapping, + use_consolidated, + "distributed rank-striped read-ahead requires an active MPI communicator", + # Ray workers intentionally disable MPI. Preserve the native + # path's per-process read-ahead in that supported mode. + allow_native_prefetch=mpi_disabled(), + **kwargs, + ) + else: + weights, session = self._start_rank_striped_read_ahead( + checkpoint_dir, + mapping, + use_consolidated, + active_communicator, + **kwargs, + ) + preparation_seconds = time.perf_counter() - preparation_started + self.record_session_timings(preparation_seconds=preparation_seconds) + # hand the weights to the caller's with-block, then finish read-ahead body_error = None try: yield weights @@ -570,39 +606,58 @@ def open_weight_session(self, body_error = error raise finally: - try: - read_error = session.finish(body_error) - except Exception as error: - status = self._last_checkpoint_io_status - status.effective = "none" - if body_error is None: - status.fallback_reason = str(error) - self._log_checkpoint_io_status() - raise - status.fallback_reason = ( - "model materialization failed: " - f"{type(body_error).__name__}: {body_error}") - self._log_checkpoint_io_status() - # Preserve the specific local error (for example, CUDA OOM) instead - # of finish()'s rank-coordinated wrapper. - logger.error( - "Rank-striped model materialization failed; preserving " - "the original exception. Coordinated session result: " - f"{type(error).__name__}: {error}") - else: - status = self._last_checkpoint_io_status - if read_error is None: - status.effective = CheckpointIOPolicy.RANK_STRIPED_READ_AHEAD - else: - # Read-ahead is advisory. The native mmap/materialization - # path succeeded, so never retry a partially mutated model. - status.effective = CheckpointIOPolicy.NATIVE - status.fallback_reason = str(read_error) - logger.warning( - "Rank-striped read-ahead degraded after activation; " - "keeping the successfully materialized model: " - f"{read_error}") + if session is not None: + finalization_started = time.perf_counter() + try: + self._finish_read_ahead_session(session, body_error) + finally: + self.record_session_timings( + preparation_seconds=preparation_seconds, + finalization_seconds=time.perf_counter() - + finalization_started) + + def _finish_read_ahead_session(self, session: RankStripedReadAheadSession, + body_error: BaseException | None) -> None: + """Finish read-ahead after the caller's with-block and record the outcome. + + Called by `open_weight_session` when its with-block ends. + + Raises: + Exception: finishing failed and the with-block itself did not + fail. When the with-block failed, its error is kept instead. + """ + status = self._last_checkpoint_io_status + try: + read_error = session.finish(body_error) + except Exception as error: + status.effective = "none" + if body_error is None: + status.fallback_reason = str(error) self._log_checkpoint_io_status() + raise + status.fallback_reason = ( + "model materialization failed: " + f"{type(body_error).__name__}: {body_error}") + self._log_checkpoint_io_status() + # Preserve the specific local error (for example, CUDA OOM) instead + # of finish()'s rank-coordinated wrapper. + logger.error( + "Rank-striped model materialization failed; preserving " + "the original exception. Coordinated session result: " + f"{type(error).__name__}: {error}") + return + + if read_error is None: + status.effective = CheckpointIOPolicy.RANK_STRIPED_READ_AHEAD + else: + # Read-ahead is advisory. The native mmap/materialization + # path succeeded, so never retry a partially mutated model. + status.effective = CheckpointIOPolicy.NATIVE + status.fallback_reason = str(read_error) + logger.warning("Rank-striped read-ahead degraded after activation; " + "keeping the successfully materialized model: " + f"{read_error}") + self._log_checkpoint_io_status() @staticmethod def _active_communicator(): diff --git a/tensorrt_llm/_torch/models/checkpoints/memory_init.py b/tensorrt_llm/_torch/models/checkpoints/memory_init.py new file mode 100644 index 000000000000..0f937a8b629d --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/memory_init.py @@ -0,0 +1,247 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Model creation on meta tensors and the later allocation of real CUDA memory. + +`ModelLoader.load()` first builds the model with meta tensors, then calls +this file to give every tensor CUDA memory before checkpoint weights arrive. +""" + +import copy +import functools +import traceback +import warnings +from typing import TYPE_CHECKING, Optional + +import torch + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_auto import AutoModelForCausalLM +from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM, MetaInitMode +from tensorrt_llm._torch.virtual_memory import RestoreMode +from tensorrt_llm._torch.virtual_memory import scope as virtual_memory_scope +from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType, LoadFormat +from tensorrt_llm.logger import logger + +if TYPE_CHECKING: + from tensorrt_llm._torch.virtual_memory import _MultiPoolProxy + +# Weight memory tag -> the engine memory tag that holds weights already on +# CUDA before the weight allocation stage. +_WEIGHT_TAG_TO_ENGINE_TAG = { + ExecutorMemoryType.MODEL_WEIGHTS_MAIN: ExecutorMemoryType.MODEL_ENGINE_MAIN, + ExecutorMemoryType.MODEL_WEIGHTS_DRAFT: ExecutorMemoryType.MODEL_ENGINE_DRAFT, +} + + +def create_model_with_meta_tensors( + config: ModelConfig, +) -> tuple[DecoderModelForCausalLM, ModelConfig, bool]: + """Try to build the LLM model with meta tensors. + + This changes every uninitialized tensors in the model to meta tensors. + Meta tensor: tensor with shape and dtype info, but no GPU allocation. + This avoids many costly CPU tensor allocation inside torch.nn.Module.__init__(). + + Note input ``config`` may be in place for some models. + + Returns: + - The model + - The config it was built from (the model may have changed it): the + copy after meta init, the original ``config`` after the fallback + - bool: Whether the model was built with meta tensors. + """ + # Config will be modified in-place for some models, like Qwen2. + # To ensure the fallback code in the except block uses the original, un-corrupt config, + # we make a deepcopy here. + config_copy = copy.deepcopy(config) + try: + with MetaInitMode(): + model = AutoModelForCausalLM.from_config(config_copy) + return model, config_copy, True + except Exception: + logger.info(f"Fallback to regular model init: {traceback.format_exc(limit=10)}") + model = AutoModelForCausalLM.from_config(config) + return model, config, False + + +def allocate_model_on_cuda( + model: DecoderModelForCausalLM, + *, + is_meta_init: bool, + load_format: LoadFormat, + model_weights_memory_tag: Optional[ExecutorMemoryType], + model_weights_restore_mode: Optional[RestoreMode], +) -> Optional["_MultiPoolProxy"]: + """Move every model tensor to CUDA. + + This is called before checkpoint weights are loaded. + + If `model_weights_memory_tag` is not None, it allocates model weights and + buffers in their respective virtual memory scope for memory management at + runtime. + + Args: + model: model built by `create_model_with_meta_tensors` + is_meta_init: whether ``model`` was built with meta tensors + load_format: load format of this load + model_weights_memory_tag: virtual memory tag for the weights, or None + model_weights_restore_mode: restore mode of that tag + + Returns: + The memory pools of the weight scope, which must stay alive, or None + when there is no weight memory tag. + """ + weight_pool = None + if model_weights_memory_tag is not None: + weight_pool = _allocate_weights_in_tagged_scope( + model, + is_meta_init=is_meta_init, + load_format=load_format, + model_weights_memory_tag=model_weights_memory_tag, + model_weights_restore_mode=model_weights_restore_mode, + ) + elif is_meta_init: + init_meta_tensors_on_cuda(model) + + # Ensure everything is at least on CUDA + # No-op if worked as expected + model.to("cuda") + + logger.info( + f"Use {get_rank_model_storage(model) / (1024**3):.2f} GB on CUDA for model weights." + ) + return weight_pool + + +def init_meta_tensors_on_cuda(model: torch.nn.Module) -> None: + """Allocate empty CUDA storage for every meta tensor of ``model``. + + Tensors that are not meta stay as they are. + """ + # Shared tensors map to one CUDA tensor. + meta_to_cuda_tensor: dict[torch.Tensor, torch.Tensor] = {} + model._apply( + functools.partial(_init_meta_tensor_on_cuda, meta_to_cuda_tensor=meta_to_cuda_tensor) + ) + + +def _init_meta_tensor_on_cuda( + tensor: torch.Tensor, meta_to_cuda_tensor: dict[torch.Tensor, torch.Tensor] +) -> torch.Tensor: + """Return an empty CUDA tensor for a meta tensor; keep other tensors as is. + + Called through ``model._apply`` by `init_meta_tensors_on_cuda`. + """ + if tensor.device != torch.device("meta"): + return tensor + if tensor not in meta_to_cuda_tensor: + meta_to_cuda_tensor[tensor] = torch.empty_like(tensor, device="cuda") + return meta_to_cuda_tensor[tensor] + + +def _allocate_weights_in_tagged_scope( + model: DecoderModelForCausalLM, + *, + is_meta_init: bool, + load_format: LoadFormat, + model_weights_memory_tag: ExecutorMemoryType, + model_weights_restore_mode: Optional[RestoreMode], +) -> "_MultiPoolProxy": + """Allocate model weights in the tagged virtual memory scope. + + Note model buffers are not allocated in the tagged virtual memory scope. + They are allocated normally in the current virtual memory scope. + """ + # Old tensor -> its CUDA tensor; each CUDA tensor also maps to itself. + old_to_cuda_tensor: dict[torch.Tensor, torch.Tensor] = {} + + # Allocate buffers to the outer virtual_memory_scope, + # but parameters (weights) to the dedicated inner virtual_memory_scope. + _apply_to_buffers_only( + model, functools.partial(_allocate_buffer_on_cuda, old_to_cuda_tensor=old_to_cuda_tensor) + ) + + # Weights that are not meta already hold values. Keep them when the load + # does not overwrite every weight (not AUTO or DUMMY), or when meta init + # left them real on purpose (e.g. created with ``ones`` or ``full``). + need_initialized_weights = load_format not in (LoadFormat.AUTO, LoadFormat.DUMMY) + copy_real_weight_values = need_initialized_weights or is_meta_init + with virtual_memory_scope(model_weights_memory_tag, model_weights_restore_mode) as pool: + model._apply( + functools.partial( + _allocate_weight_on_cuda, + old_to_cuda_tensor=old_to_cuda_tensor, + copy_real_weight_values=copy_real_weight_values, + model_weights_memory_tag=model_weights_memory_tag, + ) + ) + return pool + + +def _apply_to_buffers_only(model: torch.nn.Module, fn): + """Apply *fn* to every buffer in *model*, skipping parameters.""" + for module in model.modules(): + for key, buf in module._buffers.items(): + if buf is not None: + module._buffers[key] = fn(buf) + + +def _allocate_buffer_on_cuda( + tensor: torch.Tensor, old_to_cuda_tensor: dict[torch.Tensor, torch.Tensor] +) -> torch.Tensor: + """Return the CUDA tensor for a buffer: empty for meta, a copy otherwise.""" + if tensor not in old_to_cuda_tensor: + if tensor.device == torch.device("meta"): + cuda_tensor = torch.empty_like(tensor, device="cuda") + else: + cuda_tensor = tensor.cuda() + old_to_cuda_tensor[tensor] = cuda_tensor + old_to_cuda_tensor[cuda_tensor] = cuda_tensor + return old_to_cuda_tensor[tensor] + + +def _allocate_weight_on_cuda( + tensor: torch.Tensor, + old_to_cuda_tensor: dict[torch.Tensor, torch.Tensor], + *, + copy_real_weight_values: bool, + model_weights_memory_tag: ExecutorMemoryType, +) -> torch.Tensor: + """Return a new CUDA tensor for a weight, keeping its values when asked.""" + if tensor not in old_to_cuda_tensor: + cuda_tensor = torch.empty_like(tensor, device="cuda") + if tensor.device != torch.device("meta") and copy_real_weight_values: + if tensor.is_cuda: + warnings.warn( + f"A weight tensor of shape {tensor.shape} is already allocated on CUDA device before " + f"the weight allocation stage. This will cause extra CUDA memory usage in the " + f"'{_WEIGHT_TAG_TO_ENGINE_TAG[model_weights_memory_tag]}' scope." + ) + cuda_tensor.copy_(tensor) + old_to_cuda_tensor[tensor] = cuda_tensor + old_to_cuda_tensor[cuda_tensor] = cuda_tensor + return old_to_cuda_tensor[tensor] + + +def get_rank_model_storage(model) -> int: + """Compute total CUDA memory used by the model in the current rank""" + total_bytes = 0 + for _, param in model.named_parameters(): + if param.device.type == "cuda" and param.device.index == torch.cuda.current_device(): + total_bytes += param.element_size() * param.nelement() + for _, buf in model.named_buffers(): + if buf.device.type == "cuda" and buf.device.index == torch.cuda.current_device(): + total_bytes += buf.element_size() * buf.nelement() + return total_bytes diff --git a/tensorrt_llm/_torch/models/checkpoints/model_hooks.py b/tensorrt_llm/_torch/models/checkpoints/model_hooks.py new file mode 100644 index 000000000000..250954acd1c9 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/model_hooks.py @@ -0,0 +1,109 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Run the weight-loading hooks that models and their modules define. + +`call_load_weights` copies checkpoint tensors into a model through its +``load_weights()``. The walks then call one post-load hook, such as +``post_load_weights()``, on every module that defines it. Modules opt in by +defining the hook; no base class is needed. The post-load walks skip a module +whose ``_weights_removed`` flag is set, because its weights were freed on +purpose. +""" + +import inspect +from typing import TYPE_CHECKING, Any, Callable, Optional + +from torch import nn + +if TYPE_CHECKING: + from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper + + +def call_load_weights( + load_method: Callable[..., None], + weights: dict[str, Any], + weight_mapper: Optional["BaseWeightMapper"], + allow_partial_loading: bool = False, +) -> None: + """Call a model's weight-loading method with the arguments it accepts. + + Models differ in what ``load_weights()`` or ``load_draft_weights()`` take: + some accept ``weight_mapper``, some accept ``allow_partial_loading``. Only + the arguments the method declares are passed. + + Args: + load_method: bound method to call, e.g. ``model.load_weights`` + weights: checkpoint tensors by checkpoint name + weight_mapper: weight mapper of the model, passed only if accepted + allow_partial_loading: whether missing tensors are allowed + + Raises: + AssertionError: ``allow_partial_loading`` is True but the method + does not accept it. + """ + args = inspect.getfullargspec(load_method).args + kwargs = {} + if "weight_mapper" in args: + kwargs["weight_mapper"] = weight_mapper + if "allow_partial_loading" in args: + kwargs["allow_partial_loading"] = allow_partial_loading + else: + assert allow_partial_loading is False, ( + "allow_partial_loading is not supported for this model" + ) + load_method(weights, **kwargs) + + +def setup_aliases(model: nn.Module) -> None: + """Call ``setup_aliases()`` on every module, to create tensors that share storage.""" + for module in model.modules(): + module_setup_aliases: Optional[Callable[[], None]] = getattr(module, "setup_aliases", None) + if module_setup_aliases is not None and not getattr(module, "_weights_removed", False): + module_setup_aliases() + + +def mark_weights_transformed(model: nn.Module) -> None: + """Mark modules with a transform guard as already transformed. + + Post-transform sharing skips ``transform_weights()``, because the incoming + bytes already have the final runtime layout. The flag keeps later code + from treating them as raw checkpoint bytes. + """ + for module in model.modules(): + if hasattr(module, "_weights_transformed") and not getattr( + module, "_weights_removed", False + ): + module._weights_transformed = True + + +def walk_cache_state(model: nn.Module) -> None: + """Call ``cache_derived_state()`` on every module. + + This recomputes Python-side state from weights that already have their + final layout, without transforming them again. + """ + for module in model.modules(): + cache_derived_state: Optional[Callable[[], None]] = getattr( + module, "cache_derived_state", None + ) + if cache_derived_state is not None and not getattr(module, "_weights_removed", False): + cache_derived_state() + + +def walk_full_post_load(model: nn.Module) -> None: + """Call ``post_load_weights()``, the full post-load hook, on every module.""" + for module in model.modules(): + post_load_weights: Optional[Callable[[], None]] = getattr(module, "post_load_weights", None) + if post_load_weights is not None and not getattr(module, "_weights_removed", False): + post_load_weights() + + +def reset_weights_transformed(model: nn.Module) -> None: + """Mark transformed modules as needing a new transform pass. + + Called before fresh, untransformed weights replace the current ones. Only + modules that already carry the flag are touched. + """ + for module in model.modules(): + if hasattr(module, "_weights_transformed"): + module._weights_transformed = False diff --git a/tensorrt_llm/_torch/models/checkpoints/mx/model_loading.py b/tensorrt_llm/_torch/models/checkpoints/mx/model_loading.py new file mode 100644 index 000000000000..45fa2b48d037 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/mx/model_loading.py @@ -0,0 +1,624 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MX helpers that `ModelLoader` calls while it loads a model. + +This holds the post-transform profiles an MX receiver may accept, and the +checks and loader kwargs that the model loading code adds for MX checkpoints. +""" + +from dataclasses import replace +from typing import TYPE_CHECKING, Any, Optional + +import torch + +from tensorrt_llm._torch.models.checkpoints import model_hooks +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin +from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner +from tensorrt_llm._torch.weight_sharing import ( + LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + PHI3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_5_DENSE_VL_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_DENSE_FP8_BLOCK128_FP32SCALE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, + QWEN3_MOE_CUTLASS_POST_TRANSFORM_LAYOUT_ABI_V1, + LazyRootModelIdentity, + PostTransformConfigIdentity, + PostTransformFeature, + PostTransformProfile, + PostTransformProfileRegistry, + PostTransformQualificationDecision, + PostTransformRuntimeConfig, + PostTransformRuntimeConstraints, + PostTransformTransferScope, + SourceIdentity, + WeightManifestWriteResult, + maybe_write_weight_manifest, +) +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.checkpoint_format import CheckpointFormat +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantAlgo + +if TYPE_CHECKING: + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader + from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, LoadFormat, TorchLlmArgs + from tensorrt_llm.mapping import Mapping + +MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION = 1 + +# Dense models do not consume the MoE backend or MoE mapping dimensions. Their +# runtime topology remains bounded by the general and attention dimensions. +_MX_BF16_DENSE_RUNTIME_CONSTRAINTS = PostTransformRuntimeConstraints( + dtypes=frozenset({"bfloat16"}), + quant_algorithms=frozenset({"none"}), + kv_cache_quant_algorithms=frozenset({"none"}), + layerwise_quantization=frozenset({False}), + force_dynamic_quantization=frozenset({False}), + lora_enabled=frozenset({False}), + sparse_attention_enabled=frozenset({False}), + attention_backends=frozenset({"TRTLLM"}), + tp_sizes=frozenset({1, 2}), + pp_sizes=frozenset({1}), + cp_sizes=frozenset({1}), + attention_tp_sizes=frozenset({1, 2}), + attention_cp_sizes=frozenset({1}), + attention_dp=frozenset({False}), + multi_node=frozenset({False}), + tied_word_embeddings=frozenset({False}), + rope_types=frozenset({"default"}), + rope_fusion=frozenset({True}), +) +_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS = replace( + _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + rope_fusion=frozenset({False}), +) +# Tied Qwen3 checkpoints alias `lm_head.weight` to `model.embed_tokens.weight` +# at construction, so the shared table is transferred once and the tied profile +# has its own layout ABI. +_MX_QWEN3_BF16_DENSE_TIED_RUNTIME_CONSTRAINTS = replace( + _MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, + tied_word_embeddings=frozenset({True}), +) +# Mistral realizes a per-layer attention window from its checkpoint config. +# The dense profile is qualified only when every layer runs full attention. +_MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS = replace( + _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + sliding_windows=frozenset({"none"}), +) +# Qwen3.5 vision-language roots realize interleaved M-RoPE (`rope_type` +# `mrope`), applied by the fused QK-norm/RoPE kernel rather than the +# attention op. +_MX_QWEN3_5_VL_BF16_RUNTIME_CONSTRAINTS = replace( + _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + rope_types=frozenset({"mrope"}), + rope_fusion=frozenset({False}), +) +# Qwen3 MoE fuses RoPE into the QK-norm kernel, as Qwen3 does, and pins the +# configured MoE backend and the routed-expert partition. The routed-expert +# layout differs per MoE implementation, so the pinned backend and +# quantization must select exactly one implementation on every SM and admitted +# partition: `CUTLASS` with unquantized weights is the +# `trtllm.cutlass.grouped_gemm.none` leaf. With TP <= 2, MoE TP x MoE EP x MoE +# cluster == TP admits only the pure MoE-TP and MoE-EP partitions; wider TP +# needs a joint partition constraint. +_MX_QWEN3_MOE_BF16_CUTLASS_RUNTIME_CONSTRAINTS = replace( + _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + rope_fusion=frozenset({False}), + moe_backends=frozenset({"CUTLASS"}), + moe_tp_sizes=frozenset({1, 2}), + moe_ep_sizes=frozenset({1, 2}), + moe_cluster_sizes=frozenset({1}), +) + + +def _with_fp8_block_scale_weights( + constraints: PostTransformRuntimeConstraints, + sm_versions: frozenset[int] | set[int], +) -> PostTransformRuntimeConstraints: + """Admit FP8 block-scale (128x128) weights to an unquantized dense row. + + FP8BlockScalesLinearMethod realizes a different weight and scale layout + per GPU architecture: SM90 keeps checkpoint-layout E4M3 weights and FP32 + block scales, while other architectures can resmooth the weights and + repack the scales in transform_weights(). Each FP8 row therefore pins the + SM versions whose layout its transform-layout ABI describes. It also + requires the checkpoint to declare 128x128 weight blocks in its Hugging + Face quantization config, so a source that declares no block size or + another one falls back. The KV cache, activation dtype, and every other + dimension keep the base row's constraints, so `force_dynamic_quantization` + stays excluded. + + Args: + constraints: An unquantized dense row's runtime constraints. + sm_versions: SM versions that realize the row's FP8 layout. + + Returns: + The same constraints with declared 128x128 FP8 block-scale weights on + `sm_versions`. + """ + return replace( + constraints, + quant_algorithms=frozenset({QuantAlgo.FP8_BLOCK_SCALES.value}), + sm_versions=frozenset(sm_versions), + weight_block_sizes=frozenset({(128, 128)}), + ) + + +_MX_QWEN3_FP8_BLOCK_SM90_DENSE_RUNTIME_CONSTRAINTS = _with_fp8_block_scale_weights( + _MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, {90} +) + +# Root classes are named rather than held, so this table costs no model-zoo +# import; qualify() resolves only the profile whose class name matches the +# model being loaded. +MX_POST_TRANSFORM_PROFILE_REGISTRY = PostTransformProfileRegistry( + profiles=( + PostTransformProfile( + profile_id="llama-for-causal-lm-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_llama", "LlamaForCausalLM" + ), + architecture="LlamaForCausalLM", + model_type="llama", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + ), + PostTransformProfile( + profile_id="qwen2-for-causal-lm-bf16-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_qwen", "Qwen2ForCausalLM" + ), + architecture="Qwen2ForCausalLM", + model_type="qwen2", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + ), + # Ahead of the BF16 Qwen3 row: when both rows miss by one dimension, + # qualify() reports the first, so an FP8 load on another SM reports + # `sm_version` rather than `quant_algorithm`. + PostTransformProfile( + profile_id="qwen3-for-causal-lm-fp8-block-sm90-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_qwen3", "Qwen3ForCausalLM" + ), + architecture="Qwen3ForCausalLM", + model_type="qwen3", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=QWEN3_DENSE_FP8_BLOCK128_FP32SCALE_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_QWEN3_FP8_BLOCK_SM90_DENSE_RUNTIME_CONSTRAINTS, + ), + PostTransformProfile( + profile_id="qwen3-for-causal-lm-bf16-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_qwen3", "Qwen3ForCausalLM" + ), + architecture="Qwen3ForCausalLM", + model_type="qwen3", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, + ), + PostTransformProfile( + profile_id="qwen3-for-causal-lm-bf16-tied-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_qwen3", "Qwen3ForCausalLM" + ), + architecture="Qwen3ForCausalLM", + model_type="qwen3", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_QWEN3_BF16_DENSE_TIED_RUNTIME_CONSTRAINTS, + ), + PostTransformProfile( + profile_id="mistral-for-causal-lm-bf16-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_mistral", "MistralForCausalLM" + ), + architecture="MistralForCausalLM", + model_type="mistral", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS, + ), + PostTransformProfile( + profile_id="phi3-for-causal-lm-bf16-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_phi3", "Phi3ForCausalLM" + ), + architecture="Phi3ForCausalLM", + model_type="phi3", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=PHI3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_BF16_DENSE_RUNTIME_CONSTRAINTS, + ), + PostTransformProfile( + profile_id="qwen3_5-for-conditional-generation-bf16-complete-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_qwen3_5", "Qwen3_5VLModel" + ), + architecture="Qwen3_5ForConditionalGeneration", + model_type="qwen3_5", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=QWEN3_5_DENSE_VL_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.COMPLETE_MODEL, + runtime_constraints=_MX_QWEN3_5_VL_BF16_RUNTIME_CONSTRAINTS, + ), + PostTransformProfile( + profile_id="qwen3-moe-for-causal-lm-bf16-cutlass-target-v1", + root_model_class=LazyRootModelIdentity( + "tensorrt_llm._torch.models.modeling_qwen3_moe", "Qwen3MoeForCausalLM" + ), + architecture="Qwen3MoeForCausalLM", + model_type="qwen3_moe", + speculative_mode=None, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=QWEN3_MOE_CUTLASS_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + runtime_constraints=_MX_QWEN3_MOE_BF16_CUTLASS_RUNTIME_CONSTRAINTS, + ), + ) +) + + +def get_speculative_mode_name(spec_config: Optional["DecodingBaseConfig"]) -> Optional[str]: + """Return the lower-case speculative decoding mode used to match profiles.""" + if spec_config is None: + return None + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + mode_name = getattr(spec_dec_mode, "name", None) + if not isinstance(mode_name, str): + logger.warning( + "Unable to identify the speculative decoding mode from %s; " + "post-transform sharing is disabled for this load.", + type(spec_dec_mode).__name__, + ) + return "unknown" + return mode_name.lower() + + +def get_post_transform_transfer_scope(model: torch.nn.Module) -> PostTransformTransferScope: + """Return the model component that an MX transfer of `model` carries. + + MX transfers every parameter of the constructed root. A text-only root + is the target model. A registered multimodal root is the complete model + when it owns parameters outside its language model (a local multimodal + encoder), and only the language model otherwise, as with + `disable_mm_encoder` or multimodal encoder/prefill disaggregation. + Roots without a registered profile keep the target-model scope without + querying `language_model`; qualification rejects them by root class. + """ + if not isinstance(model, MultimodalModelMixin): + return PostTransformTransferScope.TARGET_MODEL + root_model_class = type(model) + if not any( + profile.root_model_class.matches(root_model_class) + for profile in MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles + ): + return PostTransformTransferScope.TARGET_MODEL + try: + language_model = getattr(model, "language_model", None) + except NotImplementedError: + language_model = None + if not isinstance(language_model, torch.nn.Module): + return PostTransformTransferScope.TARGET_MODEL + language_model_parameters = {id(parameter) for parameter in language_model.parameters()} + if any(id(parameter) not in language_model_parameters for parameter in model.parameters()): + return PostTransformTransferScope.COMPLETE_MODEL + return PostTransformTransferScope.LANGUAGE_MODEL + + +def qualify_post_transform_profile( + model: "DecoderModelForCausalLM", + *, + config_identity: Optional[PostTransformConfigIdentity] = None, + speculative_mode: Optional[str], + loads_draft_weights: bool, +) -> PostTransformQualificationDecision: + """Find whether this load matches a profile that may take post-transform MX weights.""" + if config_identity is None: + config_identity = PostTransformConfigIdentity.from_model_config(model.model_config) + enabled_features = set() + if loads_draft_weights: + enabled_features.add(PostTransformFeature.SEPARATE_DRAFT_MODEL) + locality_domain_policy = getattr(model.model_config, "locality_domain_policy", None) + if getattr(locality_domain_policy, "enabled", False): + enabled_features.add(PostTransformFeature.LOCALITY_DOMAIN_LOCALIZED_WEIGHTS) + # The realized balancer, not the configured one, decides the expert + # slot layout: MoE layers create it only when EPLB actually applies. + # Both gates read their state directly, so a renamed attribute fails + # the load instead of silently reporting the feature as off. + if any( + is_moe_weight_owner(module) and module._using_load_balancer() for module in model.modules() + ): + enabled_features.add(PostTransformFeature.MOE_LOAD_BALANCER) + if model.model_config.mapping.dwdp_enabled: + enabled_features.add(PostTransformFeature.DWDP) + return MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( + root_model_class=type(model), + architecture=config_identity.architecture, + model_type=config_identity.model_type, + speculative_mode=speculative_mode, + protocol_version=MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transfer_scope=get_post_transform_transfer_scope(model), + enabled_features=frozenset(enabled_features), + runtime_config=PostTransformRuntimeConfig.from_model_config( + model.model_config, model=model, sm_version=get_sm_version() + ), + ) + + +def get_mx_load_weights_kwargs( + *, + model_config: "ModelConfig", + llm_args: "TorchLlmArgs", + qualification: PostTransformQualificationDecision, +) -> dict[str, Any]: + """Return the extra `MXCheckpointLoader.load_weights` kwargs for the target model. + + Args: + model_config: final model config of the target model + llm_args: LLM args of this load + qualification: post-transform profile decision for this load + + Returns: + A kwargs dict to merge into the generic load kwargs. + """ + mx_kwargs: dict[str, Any] = { + "model_config": model_config, + "load_config": llm_args, + "post_transform_protocol_version": MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + # If a separate draft model still needs a raw disk load, do not accept + # post-transform bytes for only the target model. Enable this only + # after target and draft subgraphs have an explicit mixed-layout policy. + "allow_post_transform_weights": qualification.qualified, + } + if qualification.qualified: + # The receiver sets up aliases before post-transform bytes arrive. + mx_kwargs["prepare_post_transform_receiver"] = model_hooks.setup_aliases + return mx_kwargs + + +def get_mx_draft_load_weights_kwargs(checkpoint_loader: "BaseCheckpointLoader") -> dict[str, Any]: + """Return the extra `load_weights` kwargs for the separate draft checkpoint.""" + if checkpoint_loader.checkpoint_format != CheckpointFormat.MX: + return {} + # This loader instance still owns the target model's MX session. Preserve + # it while the native loader reads the auxiliary draft checkpoint so the + # finalized target can publish through it. + return {"_preserve_mx_session": True} + + +def _should_run_mx_staged_receiver_path( + checkpoint_loader: "BaseCheckpointLoader", + model: "DecoderModelForCausalLM", + *, + weights_preloaded: bool, + qualification: PostTransformQualificationDecision, +) -> bool: + """Whether an MX receiver can skip one-shot weight transforms. + + MXCheckpointLoader only accepts post-transform P2P bytes when this same + exact-profile check passes before transfer. It also refuses + target-plus-draft mixed layouts until there is an explicit policy for + that combination, so this post-load branch should never see an unsafe + post-transform receiver in normal use. + + Raises: + RuntimeError: the loader got post-transform weights but the load does + not match a qualified profile. + """ + if checkpoint_loader.checkpoint_format != CheckpointFormat.MX or not weights_preloaded: + return False + + method = getattr(type(checkpoint_loader), "is_post_transform_weights_preloaded", None) + if method is None or not checkpoint_loader.is_post_transform_weights_preloaded(): + return False + + profile = qualification.profile + if qualification.qualified and profile is not None: + logger.info( + "MX receiver using staged post-load profile %s for %s " + "(transform protocol v%d, layout ABI %s).", + profile.profile_id, + type(model).__name__, + MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + profile.transform_abi_id, + ) + return True + + raise RuntimeError( + f"MX receiver got post-transform weights for {type(model).__name__}, " + "but the load does not match a qualified staged post-load profile " + f"for protocol v{MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION}: " + f"{qualification.describe()}. " + "Refusing to run the full post-load path on already-transformed " + "weights." + ) + + +def run_post_load_walks( + checkpoint_loader: "BaseCheckpointLoader", + model: "DecoderModelForCausalLM", + *, + weights_preloaded: bool, + qualification: PostTransformQualificationDecision, +) -> None: + """Run the per-module post-load hooks after the weights are in place. + + An MX receiver that got post-transform weights only sets up aliases and + refreshes derived state, because the bytes already have the final layout. + Every other load runs the full post-load hook on each module. + + Args: + checkpoint_loader: checkpoint loader of this load + model: model whose weights are in place + weights_preloaded: whether the checkpoint loader filled the model directly + qualification: post-transform profile decision for this load + """ + if _should_run_mx_staged_receiver_path( + checkpoint_loader, + model, + weights_preloaded=weights_preloaded, + qualification=qualification, + ): + model_hooks.setup_aliases(model) + model_hooks.mark_weights_transformed(model) + model_hooks.walk_cache_state(model) + else: + model_hooks.walk_full_post_load(model) + + +def post_load_publish( + checkpoint_loader: "BaseCheckpointLoader", + model: "DecoderModelForCausalLM", + *, + checkpoint_dir: str, + weights_preloaded: bool, + qualification: PostTransformQualificationDecision, + source_identity: Optional[SourceIdentity], +) -> None: + """Run the checkpoint loader's publish step after post-load processing. + + Any checkpoint loader may publish. For MX, only qualified profiles + publish, and the publish carries this rank's source identity. + """ + kwargs = { + "checkpoint_dir": checkpoint_dir, + "weights_preloaded": weights_preloaded, + } + if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: + if not _is_mx_publish_allowed( + model, weights_preloaded=weights_preloaded, qualification=qualification + ): + return + kwargs["source_identity"] = source_identity + checkpoint_loader.post_load_publish(model, **kwargs) + + +def run_post_load_walks_and_publish( + checkpoint_loader: "BaseCheckpointLoader", + model: "DecoderModelForCausalLM", + *, + checkpoint_dir: str, + weights_preloaded: bool, + qualification: PostTransformQualificationDecision, + source_identity: Optional[SourceIdentity], +) -> None: + """Run the per-module post-load hooks, then the checkpoint loader's publish step. + + `run_post_load_walks` followed by `post_load_publish`, for loads with + nothing to do in between. The GMS writer calls the two separately, + because it moves stray parameters into its pool before publishing. + + Args: + checkpoint_loader: checkpoint loader of this load + model: model whose weights are in place + checkpoint_dir: checkpoint directory of the target model + weights_preloaded: whether the checkpoint loader filled the model directly + qualification: post-transform profile decision for this load + source_identity: identity of this rank's model for MX, or None + """ + run_post_load_walks( + checkpoint_loader, + model, + weights_preloaded=weights_preloaded, + qualification=qualification, + ) + post_load_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + qualification=qualification, + source_identity=source_identity, + ) + + +def _is_mx_publish_allowed( + model: "DecoderModelForCausalLM", + *, + weights_preloaded: bool, + qualification: PostTransformQualificationDecision, +) -> bool: + """Whether this rank may publish its post-transform weights through MX. + + Only qualified profiles publish. A rank that loaded from disk logs why it + skipped publishing. + """ + if qualification.qualified: + return True + if not weights_preloaded: + logger.info( + "Skipping MX post-transform publish for " + f"{type(model).__name__}: qualification " + f"{qualification.describe()}." + ) + return False + + +def dump_final_weight_manifest( + checkpoint_loader: "BaseCheckpointLoader", + model: "DecoderModelForCausalLM", + *, + weights_preloaded: bool, + mapping: "Mapping", + load_format: "LoadFormat", +) -> Optional[WeightManifestWriteResult]: + """Write the env-gated final-state weight manifest for this rank. + + This is a no-op unless `MX_WEIGHT_MANIFEST_DIR` is set. It runs after + every post-load hook and after MoE load-balancer finalization, and + before engine warmup, so the manifest describes the final post-load + state that all roles of the MX qualification harness share. + `ModelLoader.reload()` is deliberately not covered: incremental weight + updates own their own lifecycle. + """ + return maybe_write_weight_manifest( + model, + family="final", + rank=mapping.rank, + context={ + "boundary": "model_loader_load_end", + "checkpoint_format": checkpoint_loader.checkpoint_format, + "weights_preloaded": bool(weights_preloaded), + "load_format": str(load_format), + "tp_rank": mapping.tp_rank, + "pp_rank": mapping.pp_rank, + "world_size": mapping.world_size, + "model_class": type(model).__name__, + }, + ) diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 99d2a34d40b7..89cd2645241b 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -8,6 +8,7 @@ import os import time from dataclasses import dataclass +from enum import Enum from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Dict, Generic, Iterator, List, Literal, Optional, Tuple, Type, TypeVar, Union) @@ -60,6 +61,22 @@ def timing_metric(metric_name: str, metric_dict: dict[str, time.perf_counter() - start) +class ModelLoaderMetricNames(Enum): + """Names of the model loader's startup timing metrics.""" + + TOTAL_MODEL_LOADING_SECONDS = "total_model_loading_seconds" + CHECKPOINT_PREPARATION_SECONDS = "checkpoint_preparation_seconds" + WEIGHT_POPULATION_SECONDS = "weight_population_seconds" + CHECKPOINT_FINALIZATION_SECONDS = "checkpoint_finalization_seconds" + DRAFT_CHECKPOINT_PREPARATION_SECONDS = ( + "draft_checkpoint_preparation_seconds") + DRAFT_WEIGHT_POPULATION_SECONDS = "draft_weight_population_seconds" + DRAFT_CHECKPOINT_FINALIZATION_SECONDS = ( + "draft_checkpoint_finalization_seconds") + POST_LOAD_PROCESSING_SECONDS = "post_load_processing_seconds" + WEIGHT_MANIFEST_SECONDS = "weight_manifest_seconds" + + @dataclass class EagerFusionConfig: PRE_MOE_FUSION: bool = False diff --git a/tensorrt_llm/_torch/modules/engram/engram.py b/tensorrt_llm/_torch/modules/engram/engram.py index 7504b2301f10..2837e27cadaa 100644 --- a/tensorrt_llm/_torch/modules/engram/engram.py +++ b/tensorrt_llm/_torch/modules/engram/engram.py @@ -1569,8 +1569,9 @@ def __init__( self.dtype = dtype if dtype is not None else torch.bfloat16 self.tp_size = tp_size self.tp_rank = tp_rank - # The HF loader fills pinned allocations in place; other loaders may replace them. - self._requires_standard_hf_loading = True + # Regular checkpoint loading fills pinned allocations in place; other loaders may + # replace them. + self._requires_regular_checkpoint_loading = True # Rebuild UVA views when host addresses change to avoid stale pointers. self._views: Optional[Tuple[torch.Tensor, torch.Tensor]] = None self._views_src: Optional[Tuple[int, int]] = None diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/ple.py b/tensorrt_llm/_torch/modules/qwen4_exp/ple.py index 430a511cf163..4327eeb94654 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/ple.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/ple.py @@ -360,7 +360,7 @@ def _gather_pinned_embedding_kernel( class Qwen4ExpPinnedHostEmbedding(nn.Module): """A row-sharded PLE table held in pinned host memory and read through UVA.""" - _requires_standard_hf_loading = True + _requires_regular_checkpoint_loading = True def __init__( self, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index f7c054210ab6..40415f1e4ec6 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -2,38 +2,23 @@ # SPDX-License-Identifier: Apache-2.0 import copy -import inspect import os -import sys -import traceback -import warnings -from contextlib import AbstractContextManager, contextmanager from dataclasses import replace -from enum import Enum -from typing import Any, Callable, Iterator, Optional, Tuple +from typing import Any, Callable, Optional, Tuple import torch from tensorrt_llm._errors import ConfigInputError, UnsupportedError -from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import ( - AutoCheckpointMapper, BaseCheckpointLoader) +from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ + BaseCheckpointLoader +from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import ( + HfCheckpointLoader, checkpoint_io_policy_metadata) +from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import \ + HfWeightLoader from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.weight_sharing import ( - LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, - MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - PHI3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - QWEN3_5_DENSE_VL_POST_TRANSFORM_LAYOUT_ABI_V1, - QWEN3_DENSE_FP8_BLOCK128_FP32SCALE_POST_TRANSFORM_LAYOUT_ABI_V1, - QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, - QWEN3_MOE_CUTLASS_POST_TRANSFORM_LAYOUT_ABI_V1, ArtifactIdentity, - IdentityCheckPolicy, LazyRootModelIdentity, PostTransformConfigIdentity, - PostTransformFeature, PostTransformProfile, PostTransformProfileRegistry, - PostTransformQualificationDecision, PostTransformRuntimeConfig, - PostTransformRuntimeConstraints, PostTransformTransferScope, SourceIdentity, - WeightManifestWriteResult, check_weight_sharing_compatibility, - maybe_write_weight_manifest) + ArtifactIdentity, PostTransformConfigIdentity, + PostTransformQualificationDecision, SourceIdentity) from tensorrt_llm._utils import get_sm_version, str_dtype_to_torch from tensorrt_llm.checkpoint_format import CheckpointFormat from tensorrt_llm.llmapi.llm_args import (DecodingBaseConfig, @@ -53,19 +38,20 @@ from ..locality_domain.policy import LocalityDomainPolicy from ..model_config import ModelConfig from ..models import AutoModelForCausalLM -from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader +from ..models.checkpoints import memory_init, model_hooks +from ..models.checkpoints.base_weight_mapper import BaseWeightMapper from ..models.checkpoints.checkpoint_catalog import CheckpointCatalog +from ..models.checkpoints.gms import model_loading as gms_loading +from ..models.checkpoints.mx import model_loading as mx_loading from ..models.checkpoints.weight_load_plan import WeightLoadPlan -from ..models.modeling_multimodal_mixin import MultimodalModelMixin -from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode, +from ..models.modeling_utils import (DecoderModelForCausalLM, + ModelLoaderMetricNames, get_registered_model_class, timing_metric) from ..modules.embedding import Embedding, LMHead from ..modules.low_m_gemm import LOW_M_GEMM_ACTIVE, prepare_low_m_gemm from ..moe.fused_moe.moe_load_balancer import (BaseMoeLoadBalancer, maybe_create_moe_load_balancer) -from ..moe.fused_moe.weight_owner import is_moe_weight_owner from ..virtual_memory import RestoreMode -from ..virtual_memory import scope as virtual_memory_scope from .config_utils import (is_hybrid_linear, is_minimax_m3, is_mla, resolve_auto_ssm_cache_dtype, supports_fp4_mla_attention, uses_fp4_mla_attention, @@ -88,70 +74,6 @@ "mixed_nvfp4_bfloat16") _SHADOW_WEIGHT_LOAD_PLAN_ENV = "TRTLLM_SHADOW_WEIGHT_LOAD_PLAN" -# Dense models do not consume the MoE backend or MoE mapping dimensions. Their -# runtime topology remains bounded by the general and attention dimensions. -_MX_BF16_DENSE_RUNTIME_CONSTRAINTS = PostTransformRuntimeConstraints( - dtypes=frozenset({"bfloat16"}), - quant_algorithms=frozenset({"none"}), - kv_cache_quant_algorithms=frozenset({"none"}), - layerwise_quantization=frozenset({False}), - force_dynamic_quantization=frozenset({False}), - lora_enabled=frozenset({False}), - sparse_attention_enabled=frozenset({False}), - attention_backends=frozenset({"TRTLLM"}), - tp_sizes=frozenset({1, 2}), - pp_sizes=frozenset({1}), - cp_sizes=frozenset({1}), - attention_tp_sizes=frozenset({1, 2}), - attention_cp_sizes=frozenset({1}), - attention_dp=frozenset({False}), - multi_node=frozenset({False}), - tied_word_embeddings=frozenset({False}), - rope_types=frozenset({"default"}), - rope_fusion=frozenset({True}), -) -_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS = replace( - _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, - rope_fusion=frozenset({False}), -) -# Tied Qwen3 checkpoints alias `lm_head.weight` to `model.embed_tokens.weight` -# at construction, so the shared table is transferred once and the tied profile -# has its own layout ABI. -_MX_QWEN3_BF16_DENSE_TIED_RUNTIME_CONSTRAINTS = replace( - _MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, - tied_word_embeddings=frozenset({True}), -) -# Mistral realizes a per-layer attention window from its checkpoint config. -# The dense profile is qualified only when every layer runs full attention. -_MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS = replace( - _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, - sliding_windows=frozenset({"none"}), -) -# Qwen3.5 vision-language roots realize interleaved M-RoPE (`rope_type` -# `mrope`), applied by the fused QK-norm/RoPE kernel rather than the -# attention op. -_MX_QWEN3_5_VL_BF16_RUNTIME_CONSTRAINTS = replace( - _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, - rope_types=frozenset({"mrope"}), - rope_fusion=frozenset({False}), -) -# Qwen3 MoE fuses RoPE into the QK-norm kernel, as Qwen3 does, and pins the -# configured MoE backend and the routed-expert partition. The routed-expert -# layout differs per MoE implementation, so the pinned backend and -# quantization must select exactly one implementation on every SM and admitted -# partition: `CUTLASS` with unquantized weights is the -# `trtllm.cutlass.grouped_gemm.none` leaf. With TP <= 2, MoE TP x MoE EP x MoE -# cluster == TP admits only the pure MoE-TP and MoE-EP partitions; wider TP -# needs a joint partition constraint. -_MX_QWEN3_MOE_BF16_CUTLASS_RUNTIME_CONSTRAINTS = replace( - _MX_BF16_DENSE_RUNTIME_CONSTRAINTS, - rope_fusion=frozenset({False}), - moe_backends=frozenset({"CUTLASS"}), - moe_tp_sizes=frozenset({1, 2}), - moe_ep_sizes=frozenset({1, 2}), - moe_cluster_sizes=frozenset({1}), -) - # TorchLlmArgs fields that only the model classes listing them in # `model_specific_llm_args` implement. _MODEL_SPECIFIC_LLM_ARGS = ("replicate_embedding", @@ -166,44 +88,6 @@ def _reject_unsupported_model_llm_args(llm_args: TorchLlmArgs, raise ValueError(f"{name} is not supported by {model_cls.__name__}") -def _with_fp8_block_scale_weights( - constraints: PostTransformRuntimeConstraints, - sm_versions: frozenset[int] | set[int], -) -> PostTransformRuntimeConstraints: - """Admit FP8 block-scale (128x128) weights to an unquantized dense row. - - FP8BlockScalesLinearMethod realizes a different weight and scale layout - per GPU architecture: SM90 keeps checkpoint-layout E4M3 weights and FP32 - block scales, while other architectures can resmooth the weights and - repack the scales in transform_weights(). Each FP8 row therefore pins the - SM versions whose layout its transform-layout ABI describes. It also - requires the checkpoint to declare 128x128 weight blocks in its Hugging - Face quantization config, so a source that declares no block size or - another one falls back. The KV cache, activation dtype, and every other - dimension keep the base row's constraints, so `force_dynamic_quantization` - stays excluded. - - Args: - constraints: An unquantized dense row's runtime constraints. - sm_versions: SM versions that realize the row's FP8 layout. - - Returns: - The same constraints with declared 128x128 FP8 block-scale weights on - `sm_versions`. - """ - return replace( - constraints, - quant_algorithms=frozenset({QuantAlgo.FP8_BLOCK_SCALES.value}), - sm_versions=frozenset(sm_versions), - weight_block_sizes=frozenset({(128, 128)}), - ) - - -_MX_QWEN3_FP8_BLOCK_SM90_DENSE_RUNTIME_CONSTRAINTS = ( - _with_fp8_block_scale_weights(_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, - {90})) - - def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig, llm_args: TorchLlmArgs) -> None: """Validate snapshot reuse after the model and V2 setting are resolved.""" @@ -545,19 +429,6 @@ def _get_random_min_max(dtype: torch.dtype) -> Tuple[int, int]: initialize_dummy_state(low=low, high=high, seed=seed) -def get_rank_model_storage(model): - total_bytes = 0 - for _, param in model.named_parameters(): - if param.device.type == 'cuda' and param.device.index == torch.cuda.current_device( - ): - total_bytes += param.element_size() * param.nelement() - for _, buf in model.named_buffers(): - if buf.device.type == 'cuda' and buf.device.index == torch.cuda.current_device( - ): - total_bytes += buf.element_size() * buf.nelement() - return total_bytes - - def _checkpoint_io_incompatibility_reason( checkpoint_loader: Optional[BaseCheckpointLoader], checkpoint_format: CheckpointFormat | str | None, @@ -583,10 +454,6 @@ def _checkpoint_io_incompatibility_reason( if partial_model_loading: return "partial model loading was requested" - from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import \ - HfCheckpointLoader - from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import \ - HfWeightLoader from tensorrt_llm._torch.models.modeling_utils import ( CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING, get_checkpoint_weight_loader) @@ -653,10 +520,6 @@ def construct_checkpoint_loader( requested_checkpoint_io_policy, partial_model_loading, ) - from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \ - BaseCheckpointLoader - from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import \ - HfWeightLoader from tensorrt_llm._torch.models.modeling_utils import ( get_checkpoint_weight_loader, get_config_loader) @@ -699,48 +562,6 @@ def construct_checkpoint_loader( return checkpoint_loader -@contextmanager -def _legacy_weight_session(checkpoint_loader: Any, checkpoint_dir: str, - **kwargs: Any) -> Iterator[dict[str, Any]]: - yield checkpoint_loader.load_weights(checkpoint_dir, **kwargs) - - -def _open_checkpoint_weight_session( - checkpoint_loader: Any, checkpoint_dir: str, - **kwargs: Any) -> AbstractContextManager[dict[str, Any]]: - """Preserve duck-typed loaders that predate the session interface.""" - if getattr(type(checkpoint_loader), "open_weight_session", None) is None: - return _legacy_weight_session(checkpoint_loader, checkpoint_dir, - **kwargs) - return checkpoint_loader.open_weight_session(checkpoint_dir, **kwargs) - - -@contextmanager -def _timed_checkpoint_weight_session( - checkpoint_loader: Any, - checkpoint_dir: str, - metrics: dict[str, float], - checkpoint_preparation_metric_name: str, - checkpoint_finalization_metric_name: str, - **kwargs: Any, -) -> Iterator[dict[str, Any]]: - """Time session setup and finalization around materialization.""" - session = _open_checkpoint_weight_session(checkpoint_loader, checkpoint_dir, - **kwargs) - with timing_metric(checkpoint_preparation_metric_name, metrics): - weights = session.__enter__() - try: - yield weights - except BaseException: - with timing_metric(checkpoint_finalization_metric_name, metrics): - suppress = session.__exit__(*sys.exc_info()) - if not suppress: - raise - else: - with timing_metric(checkpoint_finalization_metric_name, metrics): - session.__exit__(None, None, None) - - def _inspect_shadow_weight_load_plan( checkpoint_loader: Any, checkpoint_dir: str, weight_mapper: Any, **kwargs: Any) -> tuple[CheckpointCatalog, WeightLoadPlan] | None: @@ -778,15 +599,6 @@ def _inspect_shadow_weight_load_plan( return catalog, plan -def _apply_to_buffers_only(model: torch.nn.Module, fn): - """Apply *fn* to every buffer in *model*, skipping parameters. - """ - for module in model.modules(): - for key, buf in module._buffers.items(): - if buf is not None: - module._buffers[key] = fn(buf) - - def _load_weights_strips_compile_wrapper(model: torch.nn.Module) -> bool: """Whether ``model.load_weights`` tolerates a torch.compile wrapper in the tree. @@ -865,20 +677,6 @@ def _reject_dense_dwdp_update(model: torch.nn.Module, action: str) -> None: "dwdp_config.enable_dense to update weights.") -class ModelLoaderMetricNames(Enum): - TOTAL_MODEL_LOADING_SECONDS = "total_model_loading_seconds" - CHECKPOINT_PREPARATION_SECONDS = "checkpoint_preparation_seconds" - WEIGHT_POPULATION_SECONDS = "weight_population_seconds" - CHECKPOINT_FINALIZATION_SECONDS = "checkpoint_finalization_seconds" - DRAFT_CHECKPOINT_PREPARATION_SECONDS = ( - "draft_checkpoint_preparation_seconds") - DRAFT_WEIGHT_POPULATION_SECONDS = "draft_weight_population_seconds" - DRAFT_CHECKPOINT_FINALIZATION_SECONDS = ( - "draft_checkpoint_finalization_seconds") - POST_LOAD_PROCESSING_SECONDS = "post_load_processing_seconds" - WEIGHT_MANIFEST_SECONDS = "weight_manifest_seconds" - - _CheckpointStartupMetadata = dict[str, str | dict[str, str | bool | None] | None] @@ -906,37 +704,222 @@ def _checkpoint_startup_metadata( } -def _checkpoint_io_policy_metadata( - checkpoint_loader: BaseCheckpointLoader +def _check_post_init_model_compat( + model: DecoderModelForCausalLM, load_format: LoadFormat, + checkpoint_loader: BaseCheckpointLoader) -> None: + """Raise error if the initialized LLM model cannot load the checkpoint. + """ + requires_regular_checkpoint_loading = any( + getattr(module, "_requires_regular_checkpoint_loading", False) + for module in model.modules()) + # Pinned-host parameters must be materialized and filled in place. The + # regular AUTO load copies the read weights into the existing parameters, + # which keeps their stable addresses. AUTO may still resolve to MX, so + # check the resolved format too. + if requires_regular_checkpoint_loading and ( + load_format != LoadFormat.AUTO + or checkpoint_loader.checkpoint_format == CheckpointFormat.MX): + raise UnsupportedError( + "Host-resident model weights currently require the regular " + "checkpoint loading of load_format=AUTO; DUMMY, VISION_ONLY, GMS, " + "and MX loaders cannot materialize them") + + +def _load_checkpoint_into_model( + checkpoint_loader: BaseCheckpointLoader, + checkpoint_dir: str, + *, + model: DecoderModelForCausalLM, + is_draft_model: bool, + mapping: Mapping, + weight_mapper: Optional[BaseWeightMapper], + extra_load_weights_kwargs: Optional[dict[str, Any]], + metrics: dict[str, float], ) -> dict[str, str | bool | None] | None: - """Snapshot observed policy for loads using the built-in HF weight loader. + """Load weights from one main or draft checkpoint, under auto load format. + + If enabled and compatible, use rank-striped read-ahead IO policy for a + faster load. Otherwise, use the native loading solution by calling + `checkpoint_loader.load_weights()`. Args: - checkpoint_loader: Loader whose completed weight load is observed. + checkpoint_loader: checkpoint loader of this load + checkpoint_dir: checkpoint to read, e.g. the target checkpoint + directory or ``spec_config.speculative_model`` + model: model that holds the target and, if any, the draft model + is_draft_model: whether this is the separate draft checkpoint + mapping: parallel mapping of this rank + weight_mapper: initialized weight mapper for this checkpoint + extra_load_weights_kwargs: extra kwargs for the read, e.g. + ``source_identity`` and the MX kwargs for the target; or None + metrics: model loader timing metrics, updated in place Returns: - A detached policy snapshot, or None for unsupported or unobserved loads. + The checkpoint I/O policy that read this checkpoint, or None when the + built-in HF weight loader did not read it from storage. """ - from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import \ - HfCheckpointLoader - from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import \ - HfWeightLoader - - if (not isinstance(checkpoint_loader, HfCheckpointLoader) - or type(checkpoint_loader.weight_loader) is not HfWeightLoader - or checkpoint_loader.is_weights_preloaded()): - return None - status = checkpoint_loader.weight_loader.last_checkpoint_io_status - if status.effective not in (CheckpointIOPolicy.NATIVE, - CheckpointIOPolicy.RANK_STRIPED_READ_AHEAD): - return None - return { - "requested": status.requested, - "selected": status.selected, - "activated": status.activated, - "effective": status.effective, - "fallback_reason": status.fallback_reason, + load_weights_kwargs: dict[str, Any] = {"mapping": mapping} + if is_draft_model: + preparation_metric = ModelLoaderMetricNames.DRAFT_CHECKPOINT_PREPARATION_SECONDS + population_metric = ModelLoaderMetricNames.DRAFT_WEIGHT_POPULATION_SECONDS + finalization_metric = ModelLoaderMetricNames.DRAFT_CHECKPOINT_FINALIZATION_SECONDS + model_load_weights = model.load_draft_weights + else: + preparation_metric = ModelLoaderMetricNames.CHECKPOINT_PREPARATION_SECONDS + population_metric = ModelLoaderMetricNames.WEIGHT_POPULATION_SECONDS + finalization_metric = ModelLoaderMetricNames.CHECKPOINT_FINALIZATION_SECONDS + model_load_weights = model.load_weights + # Only the target read gets the model, so a context-only worker reads + # the target checkpoint lazily. + load_weights_kwargs["model"] = model + load_weights_kwargs.update(extra_load_weights_kwargs or {}) + + # If possible, use rank-striped read-ahead IO policy to faster reading: + # - Must be HfCheckpointLoader, not its derived class + # - weight_loader must be HfWeightLoader, not its derived class + # - checkpoint_io_policy is RANK_STRIPED_READ_AHEAD + # - no context only lazy loading + if (type(checkpoint_loader) is HfCheckpointLoader + and type(checkpoint_loader.weight_loader) is HfWeightLoader + and checkpoint_loader.weight_loader.checkpoint_io_policy + == CheckpointIOPolicy.RANK_STRIPED_READ_AHEAD + and not checkpoint_loader.uses_context_only_lazy_loading( + load_weights_kwargs.get("model"))): + # read ahead, and copy the weights inside the weight loader's session + weight_loader = checkpoint_loader.weight_loader + with weight_loader.open_weight_session( + checkpoint_dir, **load_weights_kwargs) as weights: + _copy_weights_into_model(model_load_weights, weights, weight_mapper, + population_metric.value, metrics) + # The weight loader's session records its own timings. + io_status = weight_loader.last_checkpoint_io_status + metrics[preparation_metric.value] = io_status.preparation_seconds + metrics[finalization_metric.value] = io_status.finalization_seconds + else: + # Call ``checkpoint_loader.load_weights()`` to open weight tensors as + # a map (str -> tensor) first before calling LLM model class' + # ``load_weights()`` function to load them into GPU. + + # Situations that go into this branch: + # - Hf checkpoint loader with native IO policy or context-only lazy loading + # - MX or Mistral checkpoint loader + # - 3rd-party custom checkpoint loader or weight loader + with timing_metric(preparation_metric.value, metrics): + weights = checkpoint_loader.load_weights(checkpoint_dir, + **load_weights_kwargs) + # When using MX, it may have already loaded the weights in + # ``load_weights()``. So it can skip the weight copy below. + # MX indicate weights to be preloaded successfully by returning + # `weights` as empty. It does not yet support draft model. + if (weights or is_draft_model + or not checkpoint_loader.is_weights_preloaded()): + _copy_weights_into_model(model_load_weights, weights, weight_mapper, + population_metric.value, metrics) + metrics.setdefault(finalization_metric.value, 0.0) + + _inspect_shadow_weight_load_plan(checkpoint_loader, checkpoint_dir, + weight_mapper, **load_weights_kwargs) + # Take the policy after the weight session finished, which can turn + # read-ahead into native I/O, and before a draft read with the same loader + # resets it. + return checkpoint_io_policy_metadata(checkpoint_loader) + + +def _copy_weights_into_model(model_load_weights: Callable[..., None], + weights: dict[str, Any], + weight_mapper: Optional[BaseWeightMapper], + population_metric_name: str, + metrics: dict[str, float]) -> None: + """Copy read checkpoint weights into the model and time the copy.""" + with timing_metric(population_metric_name, metrics): + model_hooks.call_load_weights(model_load_weights, weights, + weight_mapper) + if torch.cuda.is_available(): # CPU guard + torch.cuda.synchronize() + + +def load_model_from_auto_load_format( + checkpoint_loader: BaseCheckpointLoader, + checkpoint_dir: str, + *, + model_config: ModelConfig, + model: DecoderModelForCausalLM, + mapping: Mapping, + llm_args: TorchLlmArgs, + spec_config: Optional[DecodingBaseConfig], + source_identity: Optional[SourceIdentity], + qualification: PostTransformQualificationDecision, + metrics: dict[str, float], + startup_metadata: _CheckpointStartupMetadata, +) -> tuple[BaseWeightMapper, bool]: + """Load the target and any separate draft checkpoint for ``LoadFormat.AUTO``. + + The model is already built and has CUDA memory; post-load hooks run + afterwards in `ModelLoader.load()`. + + Built-in and custom checkpoint loaders, including custom weight loaders, + read the checkpoint here. + + Args: + checkpoint_loader: checkpoint loader of this load + checkpoint_dir: checkpoint directory of the target model + model_config: config the model was built from + model: model to fill + mapping: parallel mapping of this rank + llm_args: LLM args of this load + spec_config: speculative decoding config, or None + source_identity: identity of this rank's model for MX, or None + qualification: MX post-transform profile decision for this load + metrics: model loader timing metrics, updated in place + startup_metadata: model loader startup metadata; the checkpoint I/O + policy of the target and draft reads is recorded in place + + Returns: + The target weight mapper, and whether the checkpoint loader wrote + the target weights into the model directly (MX P2P). + """ + # load target weights + weights_checkpoint_dir = (model.llm_checkpoint_dir if hasattr( + model, 'llm_checkpoint_dir') else checkpoint_dir) + # Generic loaders ignore source_identity; MXCheckpointLoader pops it. + target_load_weights_kwargs: dict[str, Any] = { + "source_identity": source_identity } + if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: + target_load_weights_kwargs.update( + mx_loading.get_mx_load_weights_kwargs(model_config=model_config, + llm_args=llm_args, + qualification=qualification)) + weight_mapper = checkpoint_loader.get_initialized_weight_mapper( + model, model_config) + startup_metadata["checkpoint_io_policy"] = _load_checkpoint_into_model( + checkpoint_loader, + weights_checkpoint_dir, + model=model, + is_draft_model=False, + mapping=mapping, + weight_mapper=weight_mapper, + extra_load_weights_kwargs=target_load_weights_kwargs, + metrics=metrics) + # Read before the draft load: an MX draft read without MX state resets + # this flag. + weights_preloaded = checkpoint_loader.is_weights_preloaded() + + # load separate draft weights + if spec_config is not None and spec_config.needs_separate_draft_weights: + draft_weight_mapper = checkpoint_loader.get_initialized_draft_weight_mapper( + model, weight_mapper) + startup_metadata[ + "draft_checkpoint_io_policy"] = _load_checkpoint_into_model( + checkpoint_loader, + spec_config.speculative_model, + model=model, + is_draft_model=True, + mapping=mapping, + weight_mapper=draft_weight_mapper, + extra_load_weights_kwargs=None, + metrics=metrics) + return weight_mapper, weights_preloaded class ModelLoader: @@ -944,137 +927,6 @@ class ModelLoader: Handles the loading, configuration, and weight initialization of a PyTorch model. This class isolates model loading logic from the main execution engine. """ - _MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION = 1 - - # Root classes are named rather than held, so this table costs no model-zoo - # import; qualify() resolves only the profile whose class name matches the - # model being loaded. - _POST_TRANSFORM_PROFILE_REGISTRY = PostTransformProfileRegistry(profiles=( - PostTransformProfile( - profile_id="llama-for-causal-lm-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_llama", - "LlamaForCausalLM"), - architecture="LlamaForCausalLM", - model_type="llama", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints=_MX_BF16_DENSE_RUNTIME_CONSTRAINTS, - ), - PostTransformProfile( - profile_id="qwen2-for-causal-lm-bf16-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_qwen", "Qwen2ForCausalLM"), - architecture="Qwen2ForCausalLM", - model_type="qwen2", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=QWEN2_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints=_MX_BF16_DENSE_RUNTIME_CONSTRAINTS, - ), - # Ahead of the BF16 Qwen3 row: when both rows miss by one dimension, - # qualify() reports the first, so an FP8 load on another SM reports - # `sm_version` rather than `quant_algorithm`. - PostTransformProfile( - profile_id="qwen3-for-causal-lm-fp8-block-sm90-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_qwen3", - "Qwen3ForCausalLM"), - architecture="Qwen3ForCausalLM", - model_type="qwen3", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id= - QWEN3_DENSE_FP8_BLOCK128_FP32SCALE_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints= - _MX_QWEN3_FP8_BLOCK_SM90_DENSE_RUNTIME_CONSTRAINTS, - ), - PostTransformProfile( - profile_id="qwen3-for-causal-lm-bf16-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_qwen3", - "Qwen3ForCausalLM"), - architecture="Qwen3ForCausalLM", - model_type="qwen3", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=QWEN3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints=_MX_QWEN3_BF16_DENSE_RUNTIME_CONSTRAINTS, - ), - PostTransformProfile( - profile_id="qwen3-for-causal-lm-bf16-tied-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_qwen3", - "Qwen3ForCausalLM"), - architecture="Qwen3ForCausalLM", - model_type="qwen3", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints=_MX_QWEN3_BF16_DENSE_TIED_RUNTIME_CONSTRAINTS, - ), - PostTransformProfile( - profile_id="mistral-for-causal-lm-bf16-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_mistral", - "MistralForCausalLM"), - architecture="MistralForCausalLM", - model_type="mistral", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=MISTRAL_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints=_MX_MISTRAL_BF16_DENSE_RUNTIME_CONSTRAINTS, - ), - PostTransformProfile( - profile_id="phi3-for-causal-lm-bf16-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_phi3", "Phi3ForCausalLM"), - architecture="Phi3ForCausalLM", - model_type="phi3", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=PHI3_DENSE_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints=_MX_BF16_DENSE_RUNTIME_CONSTRAINTS, - ), - PostTransformProfile( - profile_id="qwen3_5-for-conditional-generation-bf16-complete-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_qwen3_5", - "Qwen3_5VLModel"), - architecture="Qwen3_5ForConditionalGeneration", - model_type="qwen3_5", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=QWEN3_5_DENSE_VL_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.COMPLETE_MODEL, - runtime_constraints=_MX_QWEN3_5_VL_BF16_RUNTIME_CONSTRAINTS, - ), - PostTransformProfile( - profile_id="qwen3-moe-for-causal-lm-bf16-cutlass-target-v1", - root_model_class=LazyRootModelIdentity( - "tensorrt_llm._torch.models.modeling_qwen3_moe", - "Qwen3MoeForCausalLM"), - architecture="Qwen3MoeForCausalLM", - model_type="qwen3_moe", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=QWEN3_MOE_CUTLASS_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - runtime_constraints=_MX_QWEN3_MOE_BF16_CUTLASS_RUNTIME_CONSTRAINTS, - ), - )) - - @classmethod - def _post_transform_profile_registry(cls) -> PostTransformProfileRegistry: - return cls._POST_TRANSFORM_PROFILE_REGISTRY def __init__(self, llm_args: TorchLlmArgs, @@ -1236,17 +1088,6 @@ def load_config_and_apply_defaults( return llm_args - @staticmethod - def _needs_source_identity(checkpoint_loader: BaseCheckpointLoader, - load_format: LoadFormat) -> bool: - """Whether this load path can consume a SourceIdentity gate. - - Keep ordinary loading a strict no-op: SourceIdentity is currently used - only by MX (pre-transfer fallback gate) and GMS (pre-materialize - strict gate), so default HF/AUTO paths should not even build one. - """ - return load_format == LoadFormat.GMS or checkpoint_loader.checkpoint_format == CheckpointFormat.MX - @staticmethod def _build_source_identity( config: ModelConfig, @@ -1297,918 +1138,214 @@ def load( Returns: The loaded and initialized PyTorch model. """ + # Reset before the total timer starts: the timer writes into this dict. self._metrics = {} - self._checkpoint_loader = checkpoint_loader - config = self._load_and_validate_config(checkpoint_dir, - checkpoint_loader) - # Some model constructors normalize or rewrite config fields. Capture - # the registry identity from the resolved input before construction so - # publication and reception qualify the architecture the user asked - # to load, not a post-construction alias. - post_transform_config_identity = PostTransformConfigIdentity.from_model_config( - config) - load_format = self.llm_args.load_format - locality_domain_policy = getattr(config, "locality_domain_policy", None) - if (load_format == LoadFormat.GMS - and getattr(locality_domain_policy, "enabled", False)): - raise ConfigInputError( - "LoadFormat.GMS is incompatible with locality domain localized weights. " - "GMS shares registered parameters, while locality domain execution " - "releases those parameters after creating partition-local " - "weight shards.") - self._startup_metadata = _checkpoint_startup_metadata( - checkpoint_loader, load_format) - with timing_metric( ModelLoaderMetricNames.TOTAL_MODEL_LOADING_SECONDS.value, - self._metrics), maybe_create_moe_load_balancer( - config, self.mapping) as moe_load_balancer: - try: - # config will be modified in-place for some models, like Qwen2 - config_copy = copy.deepcopy(config) - with MetaInitMode(): - model = AutoModelForCausalLM.from_config(config_copy) - config = config_copy - is_meta_init = True - except Exception: - logger.info( - f"Fallback to regular model init: {traceback.format_exc(limit=10)}" - ) - model = AutoModelForCausalLM.from_config(config) - is_meta_init = False - - validate_lm_head_output_dtype(model, - self.llm_args.lm_head_output_dtype) - - requires_standard_hf_loading = any( - getattr(module, "_requires_standard_hf_loading", False) - for module in model.modules()) - # Pinned-host parameters must be materialized and filled in place. - # AUTO uses the HF mapper that preserves their stable addresses; - # AUTO may still resolve to MX, so check the resolved format too. - uses_standard_hf_loader = (load_format == LoadFormat.AUTO - and checkpoint_loader.checkpoint_format - != CheckpointFormat.MX) - if requires_standard_hf_loading and not uses_standard_hf_loader: - raise UnsupportedError( - "Host-resident model weights currently require the standard " - "Hugging Face AUTO loader; DUMMY, VISION_ONLY, GMS, and MX " - "loaders cannot materialize them") - - loads_draft_weights = ( - self.spec_config is not None - and self.spec_config.needs_separate_draft_weights) - speculative_mode = self._speculative_mode_name(self.spec_config) - post_transform_qualification = self._qualify_post_transform_profile( - model, - config_identity=post_transform_config_identity, - speculative_mode=speculative_mode, - loads_draft_weights=loads_draft_weights) - - self._source_identity: Optional[SourceIdentity] = None - if self._needs_source_identity(checkpoint_loader, load_format): - # Receiver's local SourceIdentity, built once from the final - # ModelConfig and the freshly constructed module. The module's - # realized parameter/buffer (shape, dtype) map is the layout - # ground truth; building it here (post-construction, - # pre-weight-load) gives producer and consumer a common, - # comparable lifecycle point. - self._source_identity = self._build_source_identity( - config, - model, - checkpoint_dir=checkpoint_dir, - model_name=str( - getattr(self.llm_args, "model", None) - or checkpoint_dir), - transform_abi_id=( - post_transform_qualification.transform_abi_id), - fallback_on_artifact_error=( - load_format != LoadFormat.GMS - and checkpoint_loader.checkpoint_format - == CheckpointFormat.MX), - ) - - memo: dict[torch.Tensor, torch.Tensor] = {} - - if (self.model_weights_memory_tag is not None - and load_format != LoadFormat.GMS): - # Allocate buffers to the outer virtual_memory_scope, - # but parameters (weights) to the dedicated inner virtual_memory_scope. - - def allocate_buffer_on_cuda(t: torch.Tensor): - if t not in memo: - if t.device == torch.device('meta'): - cuda_t = torch.empty_like(t, device='cuda') - else: - cuda_t = t.cuda() - memo[t] = cuda_t - memo[cuda_t] = cuda_t - return memo[t] - - _apply_to_buffers_only(model, allocate_buffer_on_cuda) - - need_initialized_weights = load_format not in (LoadFormat.AUTO, - LoadFormat.DUMMY) - - def allocate_weights_on_cuda(t: torch.Tensor): - if t not in memo: - cuda_t = torch.empty_like(t, device='cuda') - if t.device != torch.device('meta') and ( - need_initialized_weights or is_meta_init): - if t.is_cuda: - memory_type_map = { - ExecutorMemoryType.MODEL_WEIGHTS_MAIN: - ExecutorMemoryType.MODEL_ENGINE_MAIN, - ExecutorMemoryType.MODEL_WEIGHTS_DRAFT: - ExecutorMemoryType.MODEL_ENGINE_DRAFT, - } - - warnings.warn( - f"A weight tensor of shape {t.shape} is already allocated on CUDA device before " - f"the weight allocation stage. This will cause extra CUDA memory usage in the " - f"'{memory_type_map[self.model_weights_memory_tag]}' scope." - ) - cuda_t.copy_(t) - memo[t] = cuda_t - memo[cuda_t] = cuda_t - return memo[t] - - with virtual_memory_scope( - self.model_weights_memory_tag, - self.model_weights_restore_mode) as pool: - model._apply(allocate_weights_on_cuda) - self._weight_pool_proxy = pool - elif is_meta_init and load_format != LoadFormat.GMS: - - def init_meta_tensor(t: torch.Tensor): - if t.device != torch.device('meta'): - return t - - if t not in memo: - memo[t] = torch.empty_like(t, device='cuda') - return memo[t] - - model._apply(init_meta_tensor) - - # Ensure everything is at least on CUDA - # No-op if worked as expected - if load_format != LoadFormat.GMS: - model.to("cuda") - memo.clear() - - rank_model_storage = get_rank_model_storage(model) - logger.info( - f"Use {rank_model_storage / (1024**3):.2f} GB for model weights." - ) - weights_preloaded = False - # Set when either GMS RW or GMS RO branch has already run the - # post_load_* hooks itself, so the shared post-load block below - # must skip them. RW handles them inside `mem_pool_scope` so the - # committed pool reflects the post-post_load layout; RO runs - # `setup_aliases()` before `materialize_module` to wire aliases - # prior to zero-copy mapping, then refreshes derived state after - # real GMS tensors are bound. - gms_post_load_handled = False - if load_format == LoadFormat.AUTO: - # Pass model= so format-specific loaders (e.g. MX) can - # write weights directly into parameter buffers via P2P. - # Generic loaders ignore model=; loaders that can consume a - # live module reference (MX) use it for direct writes. - load_weights_kwargs: dict = { - "mapping": self.mapping, - "model": model, - # Generic loaders ignore it; MXCheckpointLoader pops it. - "source_identity": self._source_identity, - } - if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: - load_weights_kwargs["model_config"] = config - load_weights_kwargs["load_config"] = self.llm_args - load_weights_kwargs[ - "post_transform_protocol_version"] = self._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION - # If a separate draft model still needs a raw disk load, - # do not accept post-transform bytes for only the target - # model. Enable this only after target and draft subgraphs - # have an explicit mixed-layout policy. - load_weights_kwargs[ - "allow_post_transform_weights"] = post_transform_qualification.qualified - if post_transform_qualification.qualified: - load_weights_kwargs[ - "prepare_post_transform_receiver"] = self._setup_aliases - - weights_checkpoint_dir = (model.llm_checkpoint_dir if hasattr( - model, 'llm_checkpoint_dir') else checkpoint_dir) - weights_preloaded = self._materialize_checkpoint_weights( - checkpoint_loader, weights_checkpoint_dir, model, config, - load_weights_kwargs) - - if loads_draft_weights: - self._materialize_draft_checkpoint_weights( - checkpoint_loader, model) - - elif load_format == LoadFormat.GMS: - # GPU Memory Service path: weight tensors live in a - # node-shared GPU memory pool so multiple TRT-LLM instances - # zero-copy share the same weight bytes. Two roles: - # - RW (writer): opens `mem_pool_scope` BEFORE meta- - # materialization and to(cuda), runs the entire model - # bring-up (load + draft load + post_load_*) inside the - # pool, then commits via `finalize_write` so RO peers - # receive the post-post_load layout. - # - RO (reader): zero-copy materializes weights from the - # pool into the model via `materialize_module`; no - # disk I/O, no per-instance copies. - # Imported lazily so the optional `gpu_memory_service` - # dependency only loads when GMS is actually selected. - from tensorrt_llm._torch.memory import GMSBackend - - gms_backend = GMSBackend( - socket_path=self.llm_args.gms_config.socket_path, - mapping=self.mapping, - mode=self.llm_args.gms_config.mode, - tag=self.llm_args.gms_config.tag, - ) - - if not gms_backend.connect(): - raise RuntimeError( - "Failed to connect to GMS at " - f"{self.llm_args.gms_config.socket_path}") - - self._gms_backend = gms_backend - try: - # `is_rw` is `Optional[bool]` on the GPU memory backend - # protocol: True = RW lock granted, False = RO lock granted, - # None = unset. After a successful `connect()` it must be - # True or False; `None` is treated as an adapter bug and - # surfaces as a hard error rather than silently falling - # through to the RO path. - if gms_backend.is_rw is True: - # RW path: open the GMS pool BEFORE meta-tensor - # materialization and to(cuda), so every CUDA - # allocation needed to bring the model up — meta - # materialization, to(cuda), checkpoint load, draft - # load, post_load hooks — lands in the GMS pool. After - # the pool is closed, finalize_write commits the - # post-post_load layout for RO peers to map zero-copy. - # Mirrors the upstream Dynamo TRT-LLM shim pattern. - with gms_backend.mem_pool_scope(torch.device("cuda")): - # Materialize meta tensors inside the pool. Local - # memo because the outer one was deleted (and the - # outer `init_meta_tensor` branch is gated off - # for GMS so it never ran). - if is_meta_init: - gms_memo: dict = {} - - def init_meta_tensor_in_pool(t: torch.Tensor): - if t.device != torch.device('meta'): - return t - if t not in gms_memo: - gms_memo[t] = torch.empty_like( - t, device='cuda') - return gms_memo[t] - - model._apply(init_meta_tensor_in_pool) - - # Catch-all for any tensors not on CUDA yet. - model.to("cuda") - - weight_source = (model.llm_checkpoint_dir - if hasattr(model, - 'llm_checkpoint_dir') - else checkpoint_dir) - # Pass model= for symmetry with the LoadFormat.AUTO - # primary load above. Generic loaders (HF) ignore - # this kwarg; loaders that consume a live module - # reference (MX) use it for direct P2P writes into - # parameter buffers. Keeping the call shape - # consistent here avoids forgetting it when MX+GMS - # composition lands later. - load_weights_kwargs = { - "mapping": self.mapping, - "model": model, - "source_identity": self._source_identity, - } - if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: - load_weights_kwargs["model_config"] = config - load_weights_kwargs[ - "load_config"] = self.llm_args - load_weights_kwargs[ - "post_transform_protocol_version"] = self._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION - load_weights_kwargs[ - "allow_post_transform_weights"] = post_transform_qualification.qualified - if post_transform_qualification.qualified: - load_weights_kwargs[ - "prepare_post_transform_receiver"] = self._setup_aliases - with timing_metric( - ModelLoaderMetricNames. - CHECKPOINT_PREPARATION_SECONDS.value, - self._metrics): - weights = checkpoint_loader.load_weights( - weight_source, **load_weights_kwargs) - - # `weights` may be: - # - non-empty dict: standard mapping pipeline runs - # - empty/None + `is_weights_preloaded()=True`: - # the loader populated the model directly - # (e.g. MX P2P), so no mapping is needed - # - empty/None + `is_weights_preloaded()=False`: - # loader returned nothing AND didn't preload — - # the model would be committed in an - # uninitialized state, so abort. Per-tensor - # partial loads (missing keys in a non-empty - # dict) are caught downstream by - # `model.load_weights` strict checking. - weights_preloaded = checkpoint_loader.is_weights_preloaded( - ) - if weights: - self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( - model, config) - with timing_metric( - ModelLoaderMetricNames. - WEIGHT_POPULATION_SECONDS.value, - self._metrics): - self._call_load_weights( - model.load_weights, weights, - self.weight_mapper) - if torch.cuda.is_available(): # CPU guard - torch.cuda.synchronize() - elif not weights_preloaded: - raise RuntimeError( - f"GMS RW: checkpoint loader " - f"'{checkpoint_loader.checkpoint_format}' " - f"returned empty weights ({weights!r}) and " - "did not preload the model. Refusing to " - "commit an unpopulated model to the GMS " - "pool.") - - self._startup_metadata["checkpoint_io_policy"] = \ - _checkpoint_io_policy_metadata(checkpoint_loader) - - if loads_draft_weights: - self._load_separate_draft_weights( - model, checkpoint_loader) - - with timing_metric( - ModelLoaderMetricNames. - POST_LOAD_PROCESSING_SECONDS.value, - self._metrics): - # Run post_load hooks INSIDE the pool so any - # tensors they create or rebind (fused QKV, - # quantization scales, derived aliases) land in - # GMS and become part of the committed layout - # that RO peers receive. Closes hhzhang16's - # narrow-scope and commit-ordering concerns. - checkpoint_loader.post_load_apply( - model, weights_preloaded=weights_preloaded) - - mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( - checkpoint_loader, - model, - weights_preloaded=weights_preloaded, - qualification=post_transform_qualification) - if mx_staged_receiver_path: - self._setup_aliases(model) - self._mark_weights_transformed(model) - self._walk_cache_state(model) - else: - self._walk_full_post_load(model) - - # Defensive last-mile sweep: catches strays from - # C++ ops that bypassed the active torch - # allocator (e.g. native cudaMalloc). - gms_backend.move_untracked_params(model) - # Safe with active GMS mappings: GMSBackend.connect() - # installed upstream's patch_empty_cache(), which makes - # torch.cuda.empty_cache() skip GMS-backed VMM regions. - # Frees the non-GMS originals dropped by - # move_untracked_params (tensor.data was rebound to - # GMS-backed replacements above) before commit so the - # cached size doesn't show as live in memory accounting. - torch.cuda.empty_cache() - - self._post_load_publish( - checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=weights_preloaded, - qualification=post_transform_qualification) - - # Pool closed. Commit the post-post_load layout. - with timing_metric( - ModelLoaderMetricNames. - POST_LOAD_PROCESSING_SECONDS.value, - self._metrics): - gms_backend.finalize_write(model) - gms_post_load_handled = True - logger.info( - "LoadFormat.GMS (RW): loaded and committed weights via %s", - checkpoint_loader.checkpoint_format) - elif gms_backend.is_rw is False: - # RO path: weights are coming from a GMS donor that - # has already committed the post-post_load layout, so - # the receiver flags `weights_preloaded=True` on the - # checkpoint-loader hooks. This prevents duplicate - # transforms while leaving any receiver-side publication - # policy to the format-specific hook. - # - # Hook order: - # 1. `post_load_apply`: format-specific apply - # work (e.g., MX preshard markers). - # 2. Per-module `setup_aliases`: creates structural - # aliases BEFORE `materialize_module` walks the - # final module tree (including `draft_model` for - # spec dec). - # 3. SourceIdentity gate: STRICT pre-materialize - # compatibility check (GMS has no disk fallback). - # 4. `materialize_module`: zero-copy bind GMS - # pool storage onto the model parameters. - # 5. Mark materialized weights as transformed, then - # recompute Python-side derived state without - # re-running one-shot transforms. - # 6. `post_load_publish`: apply the format-specific - # receiver publication policy. - with timing_metric( - ModelLoaderMetricNames. - POST_LOAD_PROCESSING_SECONDS.value, - self._metrics): - checkpoint_loader.post_load_apply( - model, weights_preloaded=True) - - self._setup_aliases(model) - - # Pre-materialize compatibility gate. GMS has no - # disk-fallback path, so a mismatch raises under STRICT - # rather than falling back. - self._check_gms_source_identity(gms_backend) - - gms_backend.materialize_module(model) - self._mark_weights_transformed(model) - self._walk_cache_state(model) - - self._post_load_publish( - checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=True, - qualification=post_transform_qualification) - gms_post_load_handled = True - logger.info("LoadFormat.GMS (RO): materialized weights") - else: - raise RuntimeError( - f"GMS backend connected but lock state is unset " - f"(is_rw={gms_backend.is_rw!r}); expected True (RW) " - "or False (RO). This indicates a bug in the GMS " - "adapter or a protocol violation.") - except Exception: - gms_backend.cleanup() - self._gms_backend = None - raise - - elif load_format == LoadFormat.DUMMY: - self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( - model, config) - initialize_dummy_weights(model) - if loads_draft_weights: - model.draft_model.load_weights_from_target_model(model) - - elif load_format == LoadFormat.VISION_ONLY: - # Vision weights are already loaded within the model. - logger.info( - "LoadFormat.VISION_ONLY: skipping weight loading; using preloaded vision weights." - ) - - else: - raise NotImplementedError( - f"No load support for load format: {load_format}") - - with timing_metric( - ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, - self._metrics): - if not gms_post_load_handled: - checkpoint_loader.post_load_apply( - model, weights_preloaded=weights_preloaded) - mx_staged_receiver_path = self._should_run_mx_staged_receiver_path( - checkpoint_loader, - model, - weights_preloaded=weights_preloaded, - qualification=post_transform_qualification) - if mx_staged_receiver_path: - self._setup_aliases(model) - self._mark_weights_transformed(model) - self._walk_cache_state(model) - else: - self._walk_full_post_load(model) - self._post_load_publish( - checkpoint_loader, - model, - checkpoint_dir=checkpoint_dir, - weights_preloaded=weights_preloaded, - qualification=post_transform_qualification) - - # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled, - # `register_weight_slots_after_to_cuda` and `finalize_model` - # must run INSIDE `mem_pool_scope` and BEFORE `finalize_write` - # so MoE allocations become part of the committed layout that - # RO peers receive. Today they run outside the pool and after - # commit, which would silently produce a broken MoE routing - # state on RO peers — that combination is REJECTED at config - # time by `TorchLlmArgs.validate_gms_moe_compat` in - # `tensorrt_llm/llmapi/llm_args.py`. When implementing the - # follow-up, drop the validator gate AFTER moving these calls - # into the pool scope. - if isinstance(moe_load_balancer, BaseMoeLoadBalancer): - moe_load_balancer.register_weight_slots_after_to_cuda() - logger.info("moe_load_balancer finalizing model...") - moe_load_balancer.finalize_model() - logger.info("moe_load_balancer finalize model done") - - if torch.cuda.is_available(): # CPU guard - torch.cuda.synchronize() - # Reclaim segments freed during per-module post_load_weights (e.g. - # MegaMoE _transform_main_weights releases ~5-6 GiB of redundant - # weight Parameters via `.data = empty(0)` that PyTorch's caching - # allocator otherwise holds onto). Returning them to the driver - # gives downstream stages (KV cache estimation, attention workspace - # alloc, autotuner warmup symmetric-fabric setup) full visibility - # of free HBM. Single one-shot call after all modules are - # finalized; per-layer empty_cache here is unsafe because it can - # perturb NVLink barrier synchronization in multi-rank DG init. - torch.cuda.empty_cache() - - manifest_result = self._dump_final_weight_manifest( - checkpoint_loader, model, weights_preloaded=weights_preloaded) - if manifest_result is not None: - self._metrics[ModelLoaderMetricNames.WEIGHT_MANIFEST_SECONDS. - value] = (manifest_result.build_seconds + - manifest_result.write_seconds) + self._metrics): + model, moe_load_balancer = self._timed_load(checkpoint_dir, + checkpoint_loader) metrics = ", ".join(f"{name}={value:.4f}" for name, value in self._metrics.items()) logger.info( f"Model loading metrics for rank {self.mapping.rank}: {metrics}") - if LOW_M_GEMM_ACTIVE: - prepare_low_m_gemm(model) return model, moe_load_balancer - def _check_gms_source_identity(self, gms_backend) -> None: - """Pre-materialize SourceIdentity gate for the GMS read-only path. - - GMS has no disk-fallback path, so a missing or mismatched identity - raises under `STRICT` rather than falling back. - - Args: - gms_backend: The connected GMS backend (RO role) exposing - `get_source_identity()`. - - Raises: - SourceIdentityMismatchError: When the writer's identity is - available and incompatible with this receiver's identity. - """ - check_weight_sharing_compatibility( - self._source_identity, - gms_backend.get_source_identity(), - IdentityCheckPolicy.STRICT, - ) + def _timed_load( + self, + checkpoint_dir: str, + checkpoint_loader: BaseCheckpointLoader, + ) -> tuple[DecoderModelForCausalLM, BaseMoeLoadBalancer | None]: + """Do all the work of `load()`, which times it as the total loading time.""" + self._checkpoint_loader = checkpoint_loader + config: ModelConfig = self._load_and_validate_config( + checkpoint_dir, checkpoint_loader) + load_format = self.llm_args.load_format + gms_loading.check_gms_locality_domain_compat(load_format, config) + self._startup_metadata = _checkpoint_startup_metadata( + checkpoint_loader, load_format) - @classmethod - def _should_run_mx_staged_receiver_path( - cls, checkpoint_loader: BaseCheckpointLoader, - model: DecoderModelForCausalLM, *, weights_preloaded: bool, - qualification: PostTransformQualificationDecision) -> bool: - """Whether an MX receiver can skip one-shot weight transforms. - - MXCheckpointLoader only accepts post-transform P2P bytes when this same - exact-profile check passes before transfer. It also refuses - target-plus-draft mixed layouts until there is an explicit policy for - that combination, so this post-load branch should never see an unsafe - post-transform receiver in normal use. - """ - if checkpoint_loader.checkpoint_format != CheckpointFormat.MX or not weights_preloaded: - return False + # Some model constructors normalize or rewrite config fields. Capture the + # original model type and architecture specified in the config here + original_config_identity = PostTransformConfigIdentity.from_model_config( + config) - method = getattr(type(checkpoint_loader), - 'is_post_transform_weights_preloaded', None) - if method is None or not checkpoint_loader.is_post_transform_weights_preloaded( - ): - return False + # build the model; MoE layers register with the active load balancer + # while they are built, so only construction needs its block + with maybe_create_moe_load_balancer(config, + self.mapping) as moe_load_balancer: + # Some constructors adjust the config they are built from, so the + # model's config gets its own name. + model, model_config, is_meta_init = memory_init.create_model_with_meta_tensors( + config) + + validate_lm_head_output_dtype(model, self.llm_args.lm_head_output_dtype) + + # raise error when the built model cannot load the checkpoint as is + _check_post_init_model_compat(model, load_format, checkpoint_loader) + + loads_draft_weights = (self.spec_config is not None and + self.spec_config.needs_separate_draft_weights) + speculative_mode = mx_loading.get_speculative_mode_name( + self.spec_config) + post_transform_qualification = mx_loading.qualify_post_transform_profile( + model, + config_identity=original_config_identity, + speculative_mode=speculative_mode, + loads_draft_weights=loads_draft_weights) - profile = qualification.profile - if qualification.qualified and profile is not None: - logger.info( - "MX receiver using staged post-load profile %s for %s " - "(transform protocol v%d, layout ABI %s).", - profile.profile_id, - type(model).__name__, - cls._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - profile.transform_abi_id, + # GMS allocates its own memory later + if load_format != LoadFormat.GMS: + weight_pool = memory_init.allocate_model_on_cuda( + model, + is_meta_init=is_meta_init, + load_format=load_format, + model_weights_memory_tag=self.model_weights_memory_tag, + model_weights_restore_mode=self.model_weights_restore_mode) + if weight_pool is not None: + self._weight_pool_proxy = weight_pool + + source_identity: Optional[SourceIdentity] = None + if load_format == LoadFormat.GMS or checkpoint_loader.checkpoint_format == CheckpointFormat.MX: + # Receiver's local SourceIdentity, built once from the final + # ModelConfig and the freshly constructed module. The module's + # realized parameter/buffer (shape, dtype) map is the layout + # ground truth; building it here (post-construction, + # pre-weight-load) gives producer and consumer a common, + # comparable lifecycle point. + source_identity = self._build_source_identity( + model_config, + model, + checkpoint_dir=checkpoint_dir, + model_name=str( + getattr(self.llm_args, "model", None) or checkpoint_dir), + transform_abi_id=( + post_transform_qualification.transform_abi_id), + fallback_on_artifact_error=(load_format != LoadFormat.GMS and + checkpoint_loader.checkpoint_format + == CheckpointFormat.MX), ) - return True - raise RuntimeError( - f"MX receiver got post-transform weights for {type(model).__name__}, " - "but the load does not match a qualified staged post-load profile " - f"for protocol v{cls._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION}: " - f"{qualification.describe()}. " - "Refusing to run the full post-load path on already-transformed " - "weights.") + weights_preloaded = False + if load_format == LoadFormat.AUTO: + # read the checkpoint with the checkpoint loader, built-in or custom + self.weight_mapper, weights_preloaded = load_model_from_auto_load_format( + checkpoint_loader, + checkpoint_dir, + model_config=model_config, + model=model, + mapping=self.mapping, + llm_args=self.llm_args, + spec_config=self.spec_config, + source_identity=source_identity, + qualification=post_transform_qualification, + metrics=self._metrics, + startup_metadata=self._startup_metadata) + + elif load_format == LoadFormat.GMS: + gms_result = gms_loading.load_gms_weights( + model, + gms_config=self.llm_args.gms_config, + mapping=self.mapping, + llm_args=self.llm_args, + checkpoint_dir=checkpoint_dir, + checkpoint_loader=checkpoint_loader, + model_config=model_config, + is_meta_init=is_meta_init, + loads_draft_weights=loads_draft_weights, + speculative_model=(self.spec_config.speculative_model + if loads_draft_weights else None), + qualification=post_transform_qualification, + source_identity=source_identity, + metrics=self._metrics, + startup_metadata=self._startup_metadata) + self._gms_backend = gms_result.gms_backend + weights_preloaded = gms_result.weights_preloaded + if gms_result.weight_mapper is not None: + self.weight_mapper = gms_result.weight_mapper + + elif load_format == LoadFormat.DUMMY: + self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( + model, model_config) + initialize_dummy_weights(model) + if loads_draft_weights: + model.draft_model.load_weights_from_target_model(model) - @staticmethod - def _speculative_mode_name( - spec_config: Optional[DecodingBaseConfig]) -> Optional[str]: - if spec_config is None: - return None - spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) - mode_name = getattr(spec_dec_mode, "name", None) - if not isinstance(mode_name, str): - logger.warning( - "Unable to identify the speculative decoding mode from %s; " - "post-transform sharing is disabled for this load.", - type(spec_dec_mode).__name__, + elif load_format == LoadFormat.VISION_ONLY: + # Vision weights are already loaded within the model. + logger.info( + "LoadFormat.VISION_ONLY: skipping weight loading; using preloaded vision weights." ) - return "unknown" - return mode_name.lower() - - def _load_separate_draft_weights( - self, model: DecoderModelForCausalLM, - checkpoint_loader: BaseCheckpointLoader) -> None: - """Load draft/MTP weights from ``speculative_model`` into the one-engine model. - Each draft config initializes its own architecture mapper so replacement - MTP weights use the draft's quantization and MoE backend. - """ - draft_load_kwargs = {} - if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: - # This loader instance still owns the target model's MX session. - # Preserve it while the native loader reads the auxiliary draft - # checkpoint so the finalized target can publish through it. - draft_load_kwargs["_preserve_mx_session"] = True - with timing_metric( - ModelLoaderMetricNames.DRAFT_CHECKPOINT_PREPARATION_SECONDS. - value, self._metrics): - draft_weights = checkpoint_loader.load_weights( - self.spec_config.speculative_model, - mapping=self.mapping, - **draft_load_kwargs) - - if model.draft_config is not None: - draft_model_arch = model.draft_config.pretrained_config.architectures[ - 0] - draft_weight_mapper = AutoCheckpointMapper.get( - checkpoint_loader.checkpoint_format, draft_model_arch) - draft_weight_mapper.init_model_and_config(model.draft_model, - model.draft_config) else: - # Drafters without a separate config reuse the target mapper. - draft_weight_mapper = self.weight_mapper + raise NotImplementedError( + f"No load support for load format: {load_format}") with timing_metric( - ModelLoaderMetricNames.DRAFT_WEIGHT_POPULATION_SECONDS.value, + ModelLoaderMetricNames.POST_LOAD_PROCESSING_SECONDS.value, self._metrics): - self._call_load_weights(model.load_draft_weights, draft_weights, - draft_weight_mapper) + # GMS RW and RO have already run the post_load_* hooks + # themselves. RW handles them inside `mem_pool_scope` so the + # committed pool reflects the post-post_load layout; RO runs + # `setup_aliases()` before `materialize_module` to wire aliases + # prior to zero-copy mapping, then refreshes derived state after + # real GMS tensors are bound. + if load_format != LoadFormat.GMS: + checkpoint_loader.post_load_apply( + model, weights_preloaded=weights_preloaded) + mx_loading.run_post_load_walks_and_publish( + checkpoint_loader, + model, + checkpoint_dir=checkpoint_dir, + weights_preloaded=weights_preloaded, + qualification=post_transform_qualification, + source_identity=source_identity) + + # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled, + # `register_weight_slots_after_to_cuda` and `finalize_model` + # must run INSIDE `mem_pool_scope` and BEFORE `finalize_write` + # so MoE allocations become part of the committed layout that + # RO peers receive. Today they run outside the pool and after + # commit, which would silently produce a broken MoE routing + # state on RO peers — that combination is REJECTED at config + # time by `TorchLlmArgs.validate_gms_moe_compat` in + # `tensorrt_llm/llmapi/llm_args.py`. When implementing the + # follow-up, drop the validator gate AFTER moving these calls + # into the pool scope. + if isinstance(moe_load_balancer, BaseMoeLoadBalancer): + moe_load_balancer.register_weight_slots_after_to_cuda() + logger.info("moe_load_balancer finalizing model...") + moe_load_balancer.finalize_model() + logger.info("moe_load_balancer finalize model done") + if torch.cuda.is_available(): # CPU guard torch.cuda.synchronize() - - @classmethod - def _post_transform_transfer_scope( - cls, model: torch.nn.Module) -> PostTransformTransferScope: - """Return the model component that an MX transfer of `model` carries. - - MX transfers every parameter of the constructed root. A text-only root - is the target model. A registered multimodal root is the complete model - when it owns parameters outside its language model (a local multimodal - encoder), and only the language model otherwise, as with - `disable_mm_encoder` or multimodal encoder/prefill disaggregation. - Roots without a registered profile keep the target-model scope without - querying `language_model`; qualification rejects them by root class. - """ - if not isinstance(model, MultimodalModelMixin): - return PostTransformTransferScope.TARGET_MODEL - root_model_class = type(model) - if not any( - profile.root_model_class.matches(root_model_class) - for profile in cls._post_transform_profile_registry().profiles): - return PostTransformTransferScope.TARGET_MODEL - try: - language_model = getattr(model, "language_model", None) - except NotImplementedError: - language_model = None - if not isinstance(language_model, torch.nn.Module): - return PostTransformTransferScope.TARGET_MODEL - language_model_parameters = { - id(parameter) - for parameter in language_model.parameters() - } - if any( - id(parameter) not in language_model_parameters - for parameter in model.parameters()): - return PostTransformTransferScope.COMPLETE_MODEL - return PostTransformTransferScope.LANGUAGE_MODEL - - @classmethod - def _qualify_post_transform_profile( - cls, - model: DecoderModelForCausalLM, - *, - config_identity: Optional[PostTransformConfigIdentity] = None, - speculative_mode: Optional[str], - loads_draft_weights: bool) -> PostTransformQualificationDecision: - if config_identity is None: - config_identity = PostTransformConfigIdentity.from_model_config( - model.model_config) - enabled_features = set() - if loads_draft_weights: - enabled_features.add(PostTransformFeature.SEPARATE_DRAFT_MODEL) - locality_domain_policy = getattr(model.model_config, - "locality_domain_policy", None) - if getattr(locality_domain_policy, "enabled", False): - enabled_features.add( - PostTransformFeature.LOCALITY_DOMAIN_LOCALIZED_WEIGHTS) - # The realized balancer, not the configured one, decides the expert - # slot layout: MoE layers create it only when EPLB actually applies. - # Both gates read their state directly, so a renamed attribute fails - # the load instead of silently reporting the feature as off. - if any( - is_moe_weight_owner(module) and module._using_load_balancer() - for module in model.modules()): - enabled_features.add(PostTransformFeature.MOE_LOAD_BALANCER) - if model.model_config.mapping.dwdp_enabled: - enabled_features.add(PostTransformFeature.DWDP) - return cls._post_transform_profile_registry().qualify( - root_model_class=type(model), - architecture=config_identity.architecture, - model_type=config_identity.model_type, - speculative_mode=speculative_mode, - protocol_version=cls._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transfer_scope=cls._post_transform_transfer_scope(model), - enabled_features=frozenset(enabled_features), - runtime_config=PostTransformRuntimeConfig.from_model_config( - model.model_config, model=model, sm_version=get_sm_version()), - ) - - def _dump_final_weight_manifest( - self, checkpoint_loader: BaseCheckpointLoader, - model: DecoderModelForCausalLM, *, - weights_preloaded: bool) -> Optional[WeightManifestWriteResult]: - """Write the env-gated final-state weight manifest for this rank. - - This is a no-op unless `MX_WEIGHT_MANIFEST_DIR` is set. It runs after - every post-load hook and after MoE load-balancer finalization, and - before engine warmup, so the manifest describes the final post-load - state that all roles of the MX qualification harness share. - `reload()` is deliberately not covered: incremental weight updates own - their own lifecycle. - """ - return maybe_write_weight_manifest( - model, - family="final", - rank=self.mapping.rank, - context={ - "boundary": "model_loader_load_end", - "checkpoint_format": checkpoint_loader.checkpoint_format, - "weights_preloaded": bool(weights_preloaded), - "load_format": str(self.llm_args.load_format), - "tp_rank": self.mapping.tp_rank, - "pp_rank": self.mapping.pp_rank, - "world_size": self.mapping.world_size, - "model_class": type(model).__name__, - }) - - def _post_load_publish( - self, checkpoint_loader: BaseCheckpointLoader, - model: DecoderModelForCausalLM, *, checkpoint_dir: str, - weights_preloaded: bool, - qualification: PostTransformQualificationDecision) -> None: - kwargs = { - "checkpoint_dir": checkpoint_dir, - "weights_preloaded": weights_preloaded, - } - if checkpoint_loader.checkpoint_format == CheckpointFormat.MX: - if not qualification.qualified: - if not weights_preloaded: - logger.info("Skipping MX post-transform publish for " - f"{type(model).__name__}: qualification " - f"{qualification.describe()}.") - return - kwargs["source_identity"] = self._source_identity - checkpoint_loader.post_load_publish(model, **kwargs) - - @staticmethod - def _mark_weights_transformed(model: DecoderModelForCausalLM) -> None: - """Mark modules with transform guards as already transformed. - - Post-transform sharing paths skip transform_weights() because the - incoming bytes already use the final runtime layout. Preserve that - lifecycle state on modules that participate in the transform guard - protocol so a later orchestrator/refactor does not treat them as raw - checkpoint bytes. - """ - for module in model.modules(): - if hasattr(module, '_weights_transformed') and not getattr( - module, '_weights_removed', False): - module._weights_transformed = True - - @staticmethod - def _setup_aliases(model: DecoderModelForCausalLM) -> None: - """Run structural alias setup on eligible modules. - - The walk is duck-typed so modules can opt in without inheriting a - shared base class. Modules whose weights were removed are skipped, - matching the legacy full post-load walk. - - Args: - model: Root decoder model whose module tree should be visited. - - Returns: - None. - """ - for module in model.modules(): - setup_aliases: Optional[Callable[[], None]] = getattr( - module, 'setup_aliases', None) - if setup_aliases is not None and not getattr( - module, '_weights_removed', False): - setup_aliases() - - @staticmethod - def _walk_transform(model: DecoderModelForCausalLM) -> None: - """Run one-shot weight transforms on eligible modules. - - The walk is duck-typed so modules can opt in without inheriting a shared - base class. Modules whose weights were removed are skipped, and modules - already marked `_weights_transformed` are left untouched until an - orchestrator resets the flag after rebinding fresh weight bytes. - - Args: - model: Root decoder model whose module tree should be visited. - - Returns: - None. - """ - for module in model.modules(): - transform_weights: Optional[Callable[[], None]] = getattr( - module, 'transform_weights', None) - if transform_weights is not None and not getattr( - module, '_weights_removed', False) and not getattr( - module, '_weights_transformed', False): - transform_weights() - - @staticmethod - def _walk_cache_state(model: DecoderModelForCausalLM) -> None: - """Recompute derived Python-side state on eligible modules. - - This walk is separate from weight transforms so callers that already - have transformed weight bytes can refresh local Python state without - mutating tensor layouts again. - - Args: - model: Root decoder model whose module tree should be visited. - - Returns: - None. - """ - for module in model.modules(): - cache_derived_state: Optional[Callable[[], None]] = getattr( - module, 'cache_derived_state', None) - if cache_derived_state is not None and not getattr( - module, '_weights_removed', False): - cache_derived_state() - - @staticmethod - def _walk_full_post_load(model: DecoderModelForCausalLM) -> None: - """Run the backward-compatible post-load hook on eligible modules. - - This preserves the previous `ModelLoader` behavior for standard load - paths while staged-hook migration proceeds incrementally. - - Args: - model: Root decoder model whose module tree should be visited. - - Returns: - None. - """ - for module in model.modules(): - post_load_weights: Optional[Callable[[], None]] = getattr( - module, 'post_load_weights', None) - if post_load_weights is not None and not getattr( - module, '_weights_removed', False): - post_load_weights() - - @staticmethod - def _reset_weights_transformed(model: DecoderModelForCausalLM) -> None: - """Mark transformed modules as needing a new transform pass. - - Orchestrators call this before rebinding fresh, untransformed weights. - The reset only touches modules that already carry the flag so unrelated - modules do not grow staged-hook state eagerly. - - Args: - model: Root decoder model whose module tree should be visited. - - Returns: - None. - """ - for module in model.modules(): - if hasattr(module, '_weights_transformed'): - module._weights_transformed = False + # Reclaim segments freed during per-module post_load_weights (e.g. + # MegaMoE _transform_main_weights releases ~5-6 GiB of redundant + # weight Parameters via `.data = empty(0)` that PyTorch's caching + # allocator otherwise holds onto). Returning them to the driver + # gives downstream stages (KV cache estimation, attention workspace + # alloc, autotuner warmup symmetric-fabric setup) full visibility + # of free HBM. Single one-shot call after all modules are + # finalized; per-layer empty_cache here is unsafe because it can + # perturb NVLink barrier synchronization in multi-rank DG init. + torch.cuda.empty_cache() + + manifest_result = mx_loading.dump_final_weight_manifest( + checkpoint_loader, + model, + weights_preloaded=weights_preloaded, + mapping=self.mapping, + load_format=self.llm_args.load_format) + if manifest_result is not None: + self._metrics[ModelLoaderMetricNames.WEIGHT_MANIFEST_SECONDS. + value] = (manifest_result.build_seconds + + manifest_result.write_seconds) + + if LOW_M_GEMM_ACTIVE: + prepare_low_m_gemm(model) + return model, moe_load_balancer def reload(self, model: DecoderModelForCausalLM, @@ -2222,7 +1359,8 @@ def reload(self, guards intact because untouched modules may already contain transformed live weights. The owner of the update lifecycle is responsible for running post-load processing once all bytes are present. Weight - manifests are not written here; see `_dump_final_weight_manifest`. + manifests are not written here; see + `mx_loading.dump_final_weight_manifest`. A torch.compile wrapper left in the tree is fine for models that load through the generic loaders (they strip the wrapper's ``_orig_mod`` @@ -2251,11 +1389,12 @@ def reload(self, "VISION_ONLY, which bypass the standard weight mapping path.") _check_compile_wrapper_before_reload(model) if not allow_partial_loading: - self._reset_weights_transformed(model) - self._call_load_weights(model.load_weights, - weights, - self.weight_mapper, - allow_partial_loading=allow_partial_loading) + model_hooks.reset_weights_transformed(model) + model_hooks.call_load_weights( + model.load_weights, + weights, + self.weight_mapper, + allow_partial_loading=allow_partial_loading) if torch.cuda.is_available(): # CPU guard torch.cuda.current_stream().synchronize() @@ -2299,11 +1438,8 @@ def cleanup(self) -> None: if self._gms_backend is not None: gms_backend = self._gms_backend self._gms_backend = None - try: - gms_backend.cleanup() - except Exception as exc: # noqa: BLE001 - shutdown must remain best effort - logger.warning( - f"Failed to clean up GMS backend {gms_backend!r}: {exc!r}") + # GMSBackend.cleanup() logs and swallows its own errors. + gms_backend.cleanup() if self._checkpoint_loader is not None: try: self._checkpoint_loader.cleanup() @@ -2493,91 +1629,3 @@ def _load_and_validate_config( f"max_draft_len. Extra MTP layer instances will share " f"checkpoint weights via mod-indexing.") return config - - def _materialize_checkpoint_weights(self, checkpoint_loader, - checkpoint_dir: str, model, config, - load_weights_kwargs: dict) -> bool: - """Keep loader-specific work alive through weight materialization.""" - self._startup_metadata["checkpoint_io_policy"] = None - with _timed_checkpoint_weight_session( - checkpoint_loader, checkpoint_dir, self._metrics, - ModelLoaderMetricNames.CHECKPOINT_PREPARATION_SECONDS.value, - ModelLoaderMetricNames.CHECKPOINT_FINALIZATION_SECONDS.value, - **load_weights_kwargs) as weights: - weights_preloaded = checkpoint_loader.is_weights_preloaded() - self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( - model, config) - _inspect_shadow_weight_load_plan( - checkpoint_loader, - checkpoint_dir, - self.weight_mapper, - **load_weights_kwargs, - ) - if weights: - with timing_metric( - ModelLoaderMetricNames.WEIGHT_POPULATION_SECONDS.value, - self._metrics): - self._call_load_weights(model.load_weights, weights, - self.weight_mapper) - if torch.cuda.is_available(): # CPU guard - torch.cuda.synchronize() - # Session finalization can degrade activated read-ahead to native. - # Snapshot before a same-engine draft load resets the shared loader. - self._startup_metadata["checkpoint_io_policy"] = \ - _checkpoint_io_policy_metadata(checkpoint_loader) - return weights_preloaded - - def _materialize_draft_checkpoint_weights(self, checkpoint_loader, - model) -> None: - """Keep loader-specific work alive through draft materialization.""" - self._startup_metadata["draft_checkpoint_io_policy"] = None - with _timed_checkpoint_weight_session( - checkpoint_loader, - self.spec_config.speculative_model, - self._metrics, - ModelLoaderMetricNames.DRAFT_CHECKPOINT_PREPARATION_SECONDS. - value, - ModelLoaderMetricNames.DRAFT_CHECKPOINT_FINALIZATION_SECONDS. - value, - mapping=self.mapping) as weights: - if model.draft_config is not None: - draft_model_arch = model.draft_config.pretrained_config.architectures[ - 0] - draft_weight_mapper = AutoCheckpointMapper.get( - checkpoint_loader.checkpoint_format, draft_model_arch) - draft_weight_mapper.init_model_and_config( - model.draft_model, model.draft_config) - else: - # Drafters without a separate config reuse the target mapper. - draft_weight_mapper = self.weight_mapper - _inspect_shadow_weight_load_plan( - checkpoint_loader, - self.spec_config.speculative_model, - draft_weight_mapper, - mapping=self.mapping, - ) - with timing_metric( - ModelLoaderMetricNames.DRAFT_WEIGHT_POPULATION_SECONDS. - value, self._metrics): - self._call_load_weights(model.load_draft_weights, weights, - draft_weight_mapper) - if torch.cuda.is_available(): # CPU guard - torch.cuda.synchronize() - self._startup_metadata["draft_checkpoint_io_policy"] = \ - _checkpoint_io_policy_metadata(checkpoint_loader) - - def _call_load_weights(self, - load_method: Callable, - weights, - weight_mapper, - allow_partial_loading: bool = False): - """Calls the model's weight loading method with the correct arguments.""" - args = inspect.getfullargspec(load_method).args - kargs = {} - if "weight_mapper" in args: - kargs["weight_mapper"] = weight_mapper - if "allow_partial_loading" in args: - kargs["allow_partial_loading"] = allow_partial_loading - else: - assert allow_partial_loading is False, "allow_partial_loading is not supported for this model" - load_method(weights, **kargs) diff --git a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense_locality_domain_rubin.py b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense_locality_domain_rubin.py index 38b5f900509f..98789e5744dc 100644 --- a/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense_locality_domain_rubin.py +++ b/tests/unittest/_torch/kernels/parallel/test_cute_dsl_bf16_dense_locality_domain_rubin.py @@ -375,11 +375,11 @@ def test_bf16_linear_locality_domain_end_to_end_rubin(): locality_domain_policy=LocalityDomainPolicy(enabled=True), ).cuda() linear.load_weights([{"weight": weight, "bias": bias}]) - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader + from tensorrt_llm._torch.models.checkpoints import model_hooks - ModelLoader._walk_transform(linear) + linear.transform_weights() shards = linear._locality_domain_weight_shards - ModelLoader._walk_cache_state(linear) + model_hooks.walk_cache_state(linear) linear.post_load_weights() assert linear._locality_domain_weight_shards is shards diff --git a/tests/unittest/_torch/modeling/test_modeling_musespark_open.py b/tests/unittest/_torch/modeling/test_modeling_musespark_open.py index 30ae5b459980..f464b39beb33 100644 --- a/tests/unittest/_torch/modeling/test_modeling_musespark_open.py +++ b/tests/unittest/_torch/modeling/test_modeling_musespark_open.py @@ -241,12 +241,12 @@ def __init__(self): def test_frozen_constant_cache_fills_from_either_loader_walk(stage): """Both loader walks must reach both cache owners. - `ModelLoader` has two post-load walks and they look up *different* names - (`pyexecutor/model_loader.py`): + The model loader has two post-load walks and they look up *different* + names (`models/checkpoints/model_hooks.py`): - * `_walk_full_post_load` looks up `post_load_weights`. This is the + * `walk_full_post_load` looks up `post_load_weights`. This is the branch the ordinary load takes. - * `_walk_cache_state` looks up `cache_derived_state`. This is the branch + * `walk_cache_state` looks up `cache_derived_state`. This is the branch the GPU-memory-service read-only reader and the staged-receiver path take; an RO rank never calls the MuseSparkOpen weight mapper at all, it zero-copy imports weights another rank committed and then runs exactly this stage. diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py index 2ae878887807..a730826f46a2 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py @@ -16,6 +16,7 @@ from tensorrt_llm import LLM from tensorrt_llm._torch.models import modeling_nemotron_h_multimodal as nemotron_h_multimodal +from tensorrt_llm._torch.models.checkpoints import model_hooks from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.modeling_multimodal_utils import get_multimodal_embeddings from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM @@ -220,8 +221,6 @@ def _spec_forward_stub(): @pytest.mark.cpu_only def test_nemotron_nano_delegates_draft_loading(): - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader - model = object.__new__(NemotronHMultimodalModel) torch.nn.Module.__init__(model) model.llm = MagicMock(draft_config=SimpleNamespace(), draft_model=object()) @@ -229,7 +228,7 @@ def test_nemotron_nano_delegates_draft_loading(): assert model.draft_config is model.llm.draft_config assert model.draft_model is model.llm.draft_model - ModelLoader._call_load_weights(None, model.load_draft_weights, weights, mapper) + model_hooks.call_load_weights(model.load_draft_weights, weights, mapper) model.llm.load_draft_weights.assert_called_once_with(weights, weight_mapper=mapper) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_rank_striped_read_ahead.py b/tests/unittest/_torch/models/checkpoints/hf/test_rank_striped_read_ahead.py index 14f7e75724b0..6d67c60f74c3 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_rank_striped_read_ahead.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_rank_striped_read_ahead.py @@ -238,6 +238,9 @@ def allgather(value): assert status.selected == "native" assert status.effective == "native" assert "mapping.world_size" in status.fallback_reason + # The fallback load is preparation; without read-ahead there is nothing to finish. + assert status.preparation_seconds > 0 + assert status.finalization_seconds == 0.0 reader_start.assert_not_called() warning.assert_called_once() @@ -369,6 +372,9 @@ def test_ineligible_request_uses_exact_native_fallback( assert not status.activated assert status.effective == "native" assert status.fallback_reason + # The fallback load is preparation; without read-ahead there is nothing to finish. + assert status.preparation_seconds > 0 + assert status.finalization_seconds == 0.0 reader_start.assert_not_called() if requested_policy == "rank_striped_read_ahead": warning.assert_called_once() @@ -450,7 +456,11 @@ def wait_for_cancel(session, _extent): assert "model.norm.weight" in weights assert stopped.is_set() - assert loader.last_checkpoint_io_status.effective == "rank_striped_read_ahead" + status = loader.last_checkpoint_io_status + assert status.effective == "rank_striped_read_ahead" + # Preparation started the readers; finalization waited for them to stop. + assert status.preparation_seconds > 0 + assert status.finalization_seconds > 0 def test_node_major_page_faults_reads_vmstat(monkeypatch): @@ -926,6 +936,8 @@ def test_mapping_and_materialization_failures_never_retry(tmp_path, monkeypatch) assert status.fallback_reason == ( "model materialization failed: RuntimeError: materialization failed" ) + # Finishing read-ahead is timed even when the with-block fails. + assert status.finalization_seconds > 0 assert "model materialization failed" in error_log.call_args.args[0] assert "cleanup failure" not in error_log.call_args.args[0] diff --git a/tests/unittest/_torch/modules/test_engram.py b/tests/unittest/_torch/modules/test_engram.py index e0123dfd27cb..b7c121c3c644 100644 --- a/tests/unittest/_torch/modules/test_engram.py +++ b/tests/unittest/_torch/modules/test_engram.py @@ -104,7 +104,7 @@ def test_v41_engram_gate_parity(monkeypatch: pytest.MonkeyPatch, add_residual: b table_pointers = (table.weight.data_ptr(), table.scale.data_ptr()) module.cuda() table.to(device="cuda", dtype=torch.bfloat16) - assert table._requires_standard_hf_loading + assert table._requires_regular_checkpoint_loading assert table.weight.device.type == table.scale.device.type == "cpu" assert table.weight.is_pinned() and table.scale.is_pinned() assert table.weight.dtype == torch.float8_e4m3fn diff --git a/tests/unittest/_torch/modules/test_lm_head_output_dtype.py b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py index 16fefac631ff..bf6664cec384 100644 --- a/tests/unittest/_torch/modules/test_lm_head_output_dtype.py +++ b/tests/unittest/_torch/modules/test_lm_head_output_dtype.py @@ -29,6 +29,7 @@ from transformers import LlamaConfig from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints import memory_init from tensorrt_llm._torch.modules.embedding import Embedding, LMHead from tensorrt_llm._torch.modules.linear import TensorParallelMode from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_module @@ -280,7 +281,7 @@ def test_model_loader_rejects_ignored_lm_head_output_dtype( max_seq_len=8, ) monkeypatch.setattr(loader, "_load_and_validate_config", lambda *_args: config) - monkeypatch.setattr(model_loader_module, "MetaInitMode", nullcontext) + monkeypatch.setattr(memory_init, "MetaInitMode", nullcontext) monkeypatch.setattr( model_loader_module, "maybe_create_moe_load_balancer", lambda *_args: nullcontext() ) diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 09a79c5ebf88..50fbec102a19 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -1544,12 +1544,12 @@ def test_moe_weight_owner_declares_transform_guard_before_load() -> None: """A staged receiver can mark a never-loaded expert-weight owner transformed. Post-transform receivers install already transformed expert bytes without - calling `load_weights()`. `ModelLoader._mark_weights_transformed` only marks + calling `load_weights()`. `model_hooks.mark_weights_transformed` only marks modules that carry the guard, so it must exist from construction on the backend that sits behind the `_weights_removed` `ConfigurableMoE` wrapper. """ + from tensorrt_llm._torch.models.checkpoints import model_hooks from tensorrt_llm._torch.moe.fused_moe.impl_blocks import MoEWeightOwnerMixin - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader assert MoEWeightOwnerMixin._weights_transformed is False # `__new__` rather than the constructor: the guard is a class-level @@ -1566,15 +1566,13 @@ def test_moe_weight_owner_declares_transform_guard_before_load() -> None: model = torch.nn.Module() model.experts = wrapper - ModelLoader._mark_weights_transformed(model) + model_hooks.mark_weights_transformed(model) assert backend._weights_transformed is True assert wrapper._weights_transformed is False - # The walk itself must skip the marked backend; the backend's own guard - # would also stop the quantization transform, so spy on the hook it calls. - backend.transform_weights = MagicMock(wraps=backend.transform_weights) - ModelLoader._walk_transform(model) - backend.transform_weights.assert_not_called() + # The guard makes a later transform of the marked backend skip the + # quantization transform, so spy on the hook it would call. + backend.transform_weights() backend.quant_method.transform_weights.assert_not_called() @@ -3283,7 +3281,7 @@ def test_megamoe_cache_derived_state_survives_the_read_only_reader_walk(): surfaces as the assert in ``run_moe`` -- so the walk itself is asserted, not just the override in isolation. """ - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader + from tensorrt_llm._torch.models.checkpoints import model_hooks # ``__new__`` rather than the constructor: the real ``__init__`` would want # MPI, a process group, and a DG SymmBuffer, none of which this walk needs. @@ -3300,7 +3298,7 @@ def test_megamoe_cache_derived_state_survives_the_read_only_reader_walk(): model = torch.nn.Module() model.moe = wrapper - ModelLoader._walk_cache_state(model) + model_hooks.walk_cache_state(model) backend._alloc_symm_buffer.assert_called_once_with() backend.quant_method.cache_derived_state.assert_called_once_with(backend) @@ -3316,7 +3314,7 @@ def test_megamoe_cutedsl_cache_derived_state_survives_the_read_only_reader_walk( ``post_load_weights``, so losing the override here leaves the derived MegaMoE-format state silently never recomputed. """ - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader + from tensorrt_llm._torch.models.checkpoints import model_hooks # ``__new__`` rather than the constructor: the real ``__init__`` would want # a process group and the cuMem symmetric-memory rendezvous. @@ -3336,7 +3334,7 @@ def test_megamoe_cutedsl_cache_derived_state_survives_the_read_only_reader_walk( model = torch.nn.Module() model.moe = wrapper - ModelLoader._walk_cache_state(model) + model_hooks.walk_cache_state(model) backend.create_weights.assert_called_once_with() quant_method.cache_derived_state.assert_called_once_with(backend) diff --git a/tests/unittest/_torch/pyexecutor/test_model_loader_checkpoint_io.py b/tests/unittest/_torch/pyexecutor/test_model_loader_checkpoint_io.py index 631db120a0c8..a144f0b80f5a 100644 --- a/tests/unittest/_torch/pyexecutor/test_model_loader_checkpoint_io.py +++ b/tests/unittest/_torch/pyexecutor/test_model_loader_checkpoint_io.py @@ -9,11 +9,15 @@ import pytest from tensorrt_llm._torch.models import modeling_utils +from tensorrt_llm._torch.models.checkpoints import base_checkpoint_loader, model_hooks from tensorrt_llm._torch.models.checkpoints.checkpoint_catalog import ( CheckpointCatalog, CheckpointTensor, ) -from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import HfCheckpointLoader +from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import ( + HfCheckpointLoader, + checkpoint_io_policy_metadata, +) from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import HfWeightLoader from tensorrt_llm._torch.models.checkpoints.mistral.checkpoint_loader import MistralCheckpointLoader from tensorrt_llm._torch.models.checkpoints.mx.checkpoint_loader import MXCheckpointLoader @@ -25,13 +29,11 @@ ) from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_module from tensorrt_llm._torch.pyexecutor.model_loader import ( - ModelLoader, - _checkpoint_io_policy_metadata, _checkpoint_startup_metadata, - _open_checkpoint_weight_session, + _load_checkpoint_into_model, _resolve_checkpoint_io_policy, - _timed_checkpoint_weight_session, construct_checkpoint_loader, + load_model_from_auto_load_format, ) from tensorrt_llm.checkpoint_format import CheckpointFormat from tensorrt_llm.llmapi.llm_args import CheckpointIOPolicy, LoadFormat @@ -67,20 +69,14 @@ def test_model_loader_snapshots_final_policy_before_draft_reset( monkeypatch.setattr(weight_loader, "_active_communicator", lambda: None) monkeypatch.setattr(checkpoint_loader, "get_initialized_weight_mapper", MagicMock()) monkeypatch.setattr(model_loader_module.torch.cuda, "is_available", lambda: False) + _record_weight_copies(monkeypatch, []) - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = _checkpoint_startup_metadata(checkpoint_loader, LoadFormat.AUTO) - loader.mapping = Mapping() - loader.spec_config = SimpleNamespace(speculative_model="/draft") - loader._call_load_weights = MagicMock() - model = MagicMock() - model.draft_config = None + startup_metadata = _checkpoint_startup_metadata(checkpoint_loader, LoadFormat.AUTO) fallback_reason = "read-ahead failure" if outcome in ("fallback", "degraded") else None activated = outcome in ("activated", "degraded") def finish(_body_error: BaseException | None) -> RuntimeError | None: - assert loader.startup_metadata["checkpoint_io_policy"] is None + assert startup_metadata["checkpoint_io_policy"] is None return RuntimeError(fallback_reason) if fallback_reason else None session = SimpleNamespace(finish=MagicMock(side_effect=finish)) @@ -96,11 +92,6 @@ def start_read_ahead(*_args: Any, **_kwargs: Any) -> tuple[dict[str, object], An return weights, session monkeypatch.setattr(weight_loader, "_start_rank_striped_read_ahead", start_read_ahead) - - loader._materialize_checkpoint_weights( - checkpoint_loader, "/checkpoint", model, object(), {"mapping": loader.mapping} - ) - expected = { "requested": requested, "selected": "rank_striped_read_ahead" if activated else "native", @@ -108,14 +99,35 @@ def start_read_ahead(*_args: Any, **_kwargs: Any) -> tuple[dict[str, object], An "effective": "rank_striped_read_ahead" if outcome == "activated" else "native", "fallback_reason": fallback_reason, } - assert loader.startup_metadata["checkpoint_io_policy"] == expected - assert session.finish.call_count == int(activated) - weight_loader._configure_checkpoint_io_policy("native") - loader._materialize_draft_checkpoint_weights(checkpoint_loader, model) + def get_draft_weight_mapper(*_args: Any) -> MagicMock: + # Called between the target and the draft read: the target policy is + # already recorded, and the draft read is made native. + assert startup_metadata["checkpoint_io_policy"] == expected + weight_loader._configure_checkpoint_io_policy("native") + return MagicMock() + + monkeypatch.setattr( + checkpoint_loader, "get_initialized_draft_weight_mapper", get_draft_weight_mapper + ) + + load_model_from_auto_load_format( + checkpoint_loader, + "/checkpoint", + model_config=object(), + model=SimpleNamespace(load_weights=object(), load_draft_weights=object()), + mapping=Mapping(), + llm_args=None, + spec_config=SimpleNamespace(needs_separate_draft_weights=True, speculative_model="/draft"), + source_identity=None, + qualification=None, + metrics={}, + startup_metadata=startup_metadata, + ) - assert loader.startup_metadata["checkpoint_io_policy"] == expected - assert loader.startup_metadata["draft_checkpoint_io_policy"] == { + assert session.finish.call_count == int(activated) + assert startup_metadata["checkpoint_io_policy"] == expected + assert startup_metadata["draft_checkpoint_io_policy"] == { "requested": "native", "selected": "native", "activated": False, @@ -134,21 +146,24 @@ def test_hf_subclass_disk_loads_report_final_native_policy( ) weight_loader = checkpoint_loader.weight_loader native_load = MagicMock(return_value={}) - assert _checkpoint_io_policy_metadata(checkpoint_loader) is None + assert checkpoint_io_policy_metadata(checkpoint_loader) is None monkeypatch.setattr(weight_loader, "_load_weights_native", native_load) - monkeypatch.setattr(checkpoint_loader, "get_initialized_weight_mapper", MagicMock()) + _record_weight_copies(monkeypatch, []) fallback_reason = weight_loader.last_checkpoint_io_status.fallback_reason - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = _checkpoint_startup_metadata(checkpoint_loader, LoadFormat.AUTO) - weights_preloaded = loader._materialize_checkpoint_weights( - checkpoint_loader, "/checkpoint", MagicMock(), object(), {"mapping": Mapping()} + checkpoint_io_policy = _load_checkpoint_into_model( + checkpoint_loader, + "/checkpoint", + model=SimpleNamespace(load_weights=object()), + is_draft_model=False, + mapping=Mapping(), + weight_mapper=None, + extra_load_weights_kwargs=None, + metrics={}, ) native_load.assert_called_once() - assert weights_preloaded is False - assert loader.startup_metadata["checkpoint_io_policy"] == { + assert checkpoint_io_policy == { "requested": requested, "selected": "native", "activated": False, @@ -176,31 +191,51 @@ def test_unobserved_checkpoint_policy_is_not_reported_as_native( checkpoint_loader.weight_loader._last_checkpoint_io_status.effective = "native" monkeypatch.setattr(checkpoint_loader, "is_weights_preloaded", lambda: True) - assert _checkpoint_io_policy_metadata(checkpoint_loader) is None + assert checkpoint_io_policy_metadata(checkpoint_loader) is None @pytest.mark.parametrize("loader_cls", [MistralCheckpointLoader, MXCheckpointLoader]) def test_hf_subclasses_keep_polymorphic_load_weights( + monkeypatch: pytest.MonkeyPatch, loader_cls: type[MistralCheckpointLoader] | type[MXCheckpointLoader], ) -> None: checkpoint_loader = loader_cls.__new__(loader_cls) - checkpoint_loader.load_weights = MagicMock(return_value={"subclass.weight": object()}) + checkpoint_loader._weight_loader = HfWeightLoader( + checkpoint_io_policy="rank_striped_read_ahead" + ) + weights = {"subclass.weight": object()} + checkpoint_loader.load_weights = MagicMock(return_value=weights) + reader_session = MagicMock(side_effect=AssertionError("subclasses must not read ahead")) + monkeypatch.setattr(checkpoint_loader.weight_loader, "open_weight_session", reader_session) + call_load_weights = _record_weight_copies(monkeypatch, []) + model = SimpleNamespace(load_weights=object()) mapping = object() + metrics = {} - with checkpoint_loader.open_weight_session( - "/checkpoint", mapping=mapping, model="model" - ) as weights: - assert list(weights) == ["subclass.weight"] + _load_checkpoint_into_model( + checkpoint_loader, + "/checkpoint", + model=model, + is_draft_model=False, + mapping=mapping, + weight_mapper=None, + extra_load_weights_kwargs=None, + metrics=metrics, + ) checkpoint_loader.load_weights.assert_called_once_with( - "/checkpoint", mapping=mapping, model="model" + "/checkpoint", mapping=mapping, model=model ) + assert call_load_weights.call_args.args[1] is weights + reader_session.assert_not_called() + # The read is timed as preparation, with nothing to finalize. + assert metrics["checkpoint_finalization_seconds"] == 0.0 @pytest.mark.parametrize("policy", ["native", "rank_striped_read_ahead"]) -@pytest.mark.parametrize("session", [False, True]) +@pytest.mark.parametrize("via_model_loader", [False, True]) def test_context_only_loading_keeps_unused_checkpoint_tensors_lazy( - monkeypatch: pytest.MonkeyPatch, policy: str, session: bool + monkeypatch: pytest.MonkeyPatch, policy: str, via_model_loader: bool ) -> None: loader = HfCheckpointLoader(weight_loader=HfWeightLoader(checkpoint_io_policy=policy)) lazy_weights = {"layers.0.weight": object(), "layers.39.weight": object()} @@ -210,12 +245,23 @@ def test_context_only_loading_keeps_unused_checkpoint_tensors_lazy( monkeypatch.setattr(loader.weight_loader, "_load_lazy_safetensors", lazy_load) monkeypatch.setattr(loader.weight_loader, "_load_weights_native", eager_load) monkeypatch.setattr(loader.weight_loader, "open_weight_session", reader_session) - model = SimpleNamespace(model=SimpleNamespace(disagg_context_only=True)) + model = SimpleNamespace(model=SimpleNamespace(disagg_context_only=True), load_weights=object()) options = dict(mapping=Mapping(), model=model, use_consolidated=True) - if session: - with loader.open_weight_session("/checkpoint", **options) as weights: - assert weights is lazy_weights + if via_model_loader: + # A context-only worker reads natively even when read-ahead is selected. + call_load_weights = _record_weight_copies(monkeypatch, []) + _load_checkpoint_into_model( + loader, + "/checkpoint", + model=model, + is_draft_model=False, + mapping=options["mapping"], + weight_mapper=None, + extra_load_weights_kwargs={"use_consolidated": True}, + metrics={}, + ) + assert call_load_weights.call_args.args[1] is lazy_weights else: assert loader.load_weights("/checkpoint", **options) is lazy_weights @@ -243,14 +289,28 @@ def test_full_model_loading_retains_existing_checkpoint_reader( body = SimpleNamespace() if context_only is not None: body.disagg_context_only = context_only - model = SimpleNamespace(model=body) + model = SimpleNamespace(model=body, load_weights=object()) mapping = Mapping() + call_load_weights = _record_weight_copies(monkeypatch, []) + metrics = {} - with loader.open_weight_session("/checkpoint", mapping=mapping, model=model) as weights: - assert weights is expected + _load_checkpoint_into_model( + loader, + "/checkpoint", + model=model, + is_draft_model=False, + mapping=mapping, + weight_mapper=None, + extra_load_weights_kwargs=None, + metrics=metrics, + ) + assert call_load_weights.call_args.args[1] is expected eager_load.assert_called_once_with("/checkpoint", mapping, False, model=model) lazy_load.assert_not_called() + # A native read is timed as preparation, with nothing to finalize. + assert metrics["checkpoint_preparation_seconds"] > 0 + assert metrics["checkpoint_finalization_seconds"] == 0.0 @pytest.mark.parametrize("custom_wrapper", [False, True]) @@ -273,12 +333,22 @@ class _CustomWeightLoader(HfWeightLoader): lazy_load = MagicMock(side_effect=AssertionError("must not replace a custom loader")) monkeypatch.setattr(loader.weight_loader, "load_weights", custom_load) monkeypatch.setattr(loader.weight_loader, "_load_lazy_safetensors", lazy_load) - model = SimpleNamespace(model=SimpleNamespace(disagg_context_only=True)) + model = SimpleNamespace(model=SimpleNamespace(disagg_context_only=True), load_weights=object()) mapping = Mapping() + call_load_weights = _record_weight_copies(monkeypatch, []) - with loader.open_weight_session("/checkpoint", mapping=mapping, model=model) as weights: - assert weights is expected + _load_checkpoint_into_model( + loader, + "/checkpoint", + model=model, + is_draft_model=False, + mapping=mapping, + weight_mapper=None, + extra_load_weights_kwargs=None, + metrics={}, + ) + assert call_load_weights.call_args.args[1] is expected custom_load.assert_called_once_with("/checkpoint", mapping=mapping, model=model) lazy_load.assert_not_called() @@ -397,8 +467,18 @@ def test_construct_checkpoint_loader_preserves_explicit_loader_on_fallback( active_communicator = MagicMock(side_effect=AssertionError("rank-striped setup must not run")) monkeypatch.setattr(provided_loader.weight_loader, "_load_weights_native", native_load) monkeypatch.setattr(provided_loader.weight_loader, "_active_communicator", active_communicator) - with loader.open_weight_session("/checkpoint", mapping=Mapping()) as weights: - assert weights is native_weights + call_load_weights = _record_weight_copies(monkeypatch, []) + _load_checkpoint_into_model( + loader, + "/checkpoint", + model=SimpleNamespace(load_weights=object()), + is_draft_model=False, + mapping=Mapping(), + weight_mapper=None, + extra_load_weights_kwargs=None, + metrics={}, + ) + assert call_load_weights.call_args.args[1] is native_weights active_communicator.assert_not_called() @@ -459,8 +539,7 @@ def __init__(self) -> None: loader = construct_checkpoint_loader(None, "HF") assert isinstance(loader.weight_loader, _CustomWeightLoader) mapping = Mapping() - with loader.open_weight_session("/checkpoint", mapping=mapping) as weights: - assert list(weights) == ["custom.weight"] + assert list(loader.load_weights("/checkpoint", mapping=mapping)) == ["custom.weight"] loader.weight_loader.load_weights.assert_called_once_with("/checkpoint", mapping=mapping) warning = MagicMock() @@ -501,59 +580,60 @@ class _CustomCheckpointLoader(HfCheckpointLoader): assert "checkpoint loader is not the built-in" in warning.call_args.args[0] -class _SessionCheckpointLoader: - checkpoint_format = "HF" +class _RecordingCheckpointLoader(HfCheckpointLoader): + """HF checkpoint loader that records its reads instead of reading files.""" def __init__( self, events: list[str], - checkpoint_dir: str = "/checkpoint", weights: dict[str, object] | None = None, + *, + weights_preloaded: bool = False, + target_weight_mapper: object | None = None, ) -> None: + super().__init__() self.events = events - self.checkpoint_dir = checkpoint_dir self.weights = {"model.weight": object()} if weights is None else weights + self.weights_preloaded = weights_preloaded + self.target_weight_mapper = ( + object() if target_weight_mapper is None else target_weight_mapper + ) + # keyword arguments of each read, by checkpoint directory + self.read_kwargs: dict[str, dict[str, Any]] = {} - @contextmanager - def open_weight_session( - self, checkpoint_dir: str, **kwargs: Any - ) -> Iterator[dict[str, object]]: - assert checkpoint_dir == self.checkpoint_dir - assert "mapping" in kwargs - self.events.append("session_enter") - try: - yield self.weights - finally: - self.events.append("session_exit") + def load_weights(self, checkpoint_dir: str, mapping: Any, **kwargs: Any) -> dict[str, object]: + self.read_kwargs[checkpoint_dir] = dict(kwargs, mapping=mapping) + self.events.append(f"read:{checkpoint_dir}") + return self.weights def is_weights_preloaded(self) -> bool: - return False + return self.weights_preloaded def get_initialized_weight_mapper(self, model: object, config: object) -> object: del model, config self.events.append("mapper_init") - return object() - - -def test_legacy_weight_session_defers_load_until_enter() -> None: - weights = {"model.weight": object()} - mapping = object() - checkpoint_loader = SimpleNamespace(load_weights=MagicMock(return_value=weights)) + return self.target_weight_mapper - session = _open_checkpoint_weight_session(checkpoint_loader, "/checkpoint", mapping=mapping) - checkpoint_loader.load_weights.assert_not_called() - with session as loaded_weights: - assert loaded_weights is weights - checkpoint_loader.load_weights.assert_called_once_with("/checkpoint", mapping=mapping) +def _record_weight_copies(monkeypatch: pytest.MonkeyPatch, events: list[str]) -> MagicMock: + """Replace the copy of read weights into the model with a "materialize" event.""" + call_load_weights = MagicMock(side_effect=lambda *_args: events.append("materialize")) + monkeypatch.setattr(model_hooks, "call_load_weights", call_load_weights) + return call_load_weights -def test_checkpoint_session_times_preparation_and_finalization_separately( +def test_sessionless_loader_times_its_read_as_preparation( monkeypatch: pytest.MonkeyPatch, ) -> None: + """Loaders without an HF weight session read once, with nothing to finalize.""" events = [] metrics = {} - checkpoint_loader = _SessionCheckpointLoader(events) + weights = {"model.weight": object()} + checkpoint_loader = SimpleNamespace( + load_weights=MagicMock( + side_effect=lambda *_args, **_kwargs: events.append("read") or weights + ) + ) @contextmanager def record_timing(metric_name: str, metric_values: dict[str, float]) -> Iterator[None]: @@ -565,57 +645,139 @@ def record_timing(metric_name: str, metric_values: dict[str, float]) -> Iterator events.append(f"{metric_name}_end") monkeypatch.setattr(model_loader_module, "timing_metric", record_timing) + _record_weight_copies(monkeypatch, events) + model = SimpleNamespace(load_weights=object()) + mapping = object() - with _timed_checkpoint_weight_session( + _load_checkpoint_into_model( checkpoint_loader, "/checkpoint", - metrics, - "preparation", - "finalization", - mapping=object(), - ): - events.append("materialize") + model=model, + is_draft_model=False, + mapping=mapping, + weight_mapper=object(), + extra_load_weights_kwargs={"source_identity": None}, + metrics=metrics, + ) + checkpoint_loader.load_weights.assert_called_once_with( + "/checkpoint", mapping=mapping, model=model, source_identity=None + ) assert events == [ - "preparation_start", - "session_enter", - "preparation_end", + "checkpoint_preparation_seconds_start", + "read", + "checkpoint_preparation_seconds_end", + "weight_population_seconds_start", "materialize", - "finalization_start", - "session_exit", - "finalization_end", + "weight_population_seconds_end", ] - assert metrics == {"preparation": 1.0, "finalization": 1.0} + assert metrics == { + "checkpoint_preparation_seconds": 1.0, + "weight_population_seconds": 1.0, + "checkpoint_finalization_seconds": 0.0, + } -def test_model_loader_session_spans_mapper_and_materialization() -> None: +@pytest.mark.parametrize("weights_preloaded", [False, True]) +def test_target_mapper_is_initialized_before_checkpoint_read( + monkeypatch: pytest.MonkeyPatch, weights_preloaded: bool +) -> None: events = [] - model = MagicMock() - checkpoint_loader = _SessionCheckpointLoader(events) - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = {} - loader._call_load_weights = MagicMock( - side_effect=lambda *_args, **_kwargs: events.append("materialize") - ) - weights_preloaded = loader._materialize_checkpoint_weights( + # A preloading loader (MX P2P) writes the weights into the model itself + # and returns none to copy. + checkpoint_loader = _RecordingCheckpointLoader( + events, weights={} if weights_preloaded else None, weights_preloaded=weights_preloaded + ) + _record_weight_copies(monkeypatch, events) + model = SimpleNamespace(load_weights=object()) + mapping, source_identity = object(), object() + metrics = {} + + weight_mapper, loaded_weights_preloaded = load_model_from_auto_load_format( checkpoint_loader, "/checkpoint", - model, - object(), - {"mapping": object()}, + model_config=object(), + model=model, + mapping=mapping, + llm_args=None, + spec_config=None, + source_identity=source_identity, + qualification=None, + metrics=metrics, + startup_metadata={}, + ) + + assert weight_mapper is checkpoint_loader.target_weight_mapper + assert loaded_weights_preloaded is weights_preloaded + expected_events = ["mapper_init", "read:/checkpoint"] + if not weights_preloaded: + expected_events.append("materialize") + assert events == expected_events + # The target read gets the model, so a context-only worker reads lazily. + assert checkpoint_loader.read_kwargs == { + "/checkpoint": { + "mapping": mapping, + "model": model, + "source_identity": source_identity, + } + } + # The read is timed as preparation, with nothing to finalize. + assert "checkpoint_preparation_seconds" in metrics + assert metrics["checkpoint_finalization_seconds"] == 0.0 + assert ("weight_population_seconds" in metrics) is not weights_preloaded + + +def test_read_ahead_copies_inside_weight_session_and_reports_its_timings( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events = [] + checkpoint_loader = HfCheckpointLoader( + weight_loader=HfWeightLoader(checkpoint_io_policy="rank_striped_read_ahead") ) + weight_loader = checkpoint_loader.weight_loader + weights = {"model.weight": object()} + session_kwargs = {} - assert not weights_preloaded - assert events == [ - "session_enter", - "mapper_init", - "materialize", - "session_exit", - ] - assert "checkpoint_preparation_seconds" in loader.metrics - assert "weight_population_seconds" in loader.metrics - assert "checkpoint_finalization_seconds" in loader.metrics + @contextmanager + def open_weight_session(checkpoint_dir: str, **kwargs: Any) -> Iterator[dict[str, object]]: + session_kwargs[checkpoint_dir] = kwargs + events.append("session_enter") + weight_loader.record_session_timings(preparation_seconds=2.0) + try: + yield weights + finally: + weight_loader.record_session_timings(preparation_seconds=2.0, finalization_seconds=3.0) + events.append("session_exit") + + monkeypatch.setattr(weight_loader, "open_weight_session", open_weight_session) + checkpoint_loader.load_weights = MagicMock( + side_effect=AssertionError("read-ahead must read in the weight session") + ) + call_load_weights = _record_weight_copies(monkeypatch, events) + model = SimpleNamespace(load_weights=object()) + mapping = object() + metrics = {} + + _load_checkpoint_into_model( + checkpoint_loader, + "/checkpoint", + model=model, + is_draft_model=False, + mapping=mapping, + weight_mapper=None, + extra_load_weights_kwargs={"source_identity": None}, + metrics=metrics, + ) + + # The weights are copied before the session finishes read-ahead. + assert events == ["session_enter", "materialize", "session_exit"] + assert call_load_weights.call_args.args[1] is weights + assert session_kwargs == { + "/checkpoint": {"mapping": mapping, "model": model, "source_identity": None} + } + # The model loader reports the timings the session recorded. + assert metrics["checkpoint_preparation_seconds"] == 2.0 + assert metrics["checkpoint_finalization_seconds"] == 3.0 def test_shadow_plan_disabled_does_not_call_inspection_hooks( @@ -648,41 +810,34 @@ def test_shadow_plan_is_advisory_and_preserves_materialization_order( demands=(WeightDemand("all", ("model.weight",), (0,)),), ) events = [] - checkpoint_loader = _SessionCheckpointLoader(events) + checkpoint_loader = _RecordingCheckpointLoader(events) checkpoint_loader.build_checkpoint_catalog = MagicMock( side_effect=lambda *_args, **_kwargs: events.append("catalog") or catalog ) weight_mapper = SimpleNamespace( build_weight_load_plan=MagicMock(side_effect=lambda _catalog: events.append("plan") or plan) ) - checkpoint_loader.get_initialized_weight_mapper = MagicMock( - side_effect=lambda *_args: events.append("mapper_init") or weight_mapper - ) - model = MagicMock() - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = {} - loader._call_load_weights = MagicMock( - side_effect=lambda *_args, **_kwargs: events.append("materialize") - ) + call_load_weights = _record_weight_copies(monkeypatch, events) - loader._materialize_checkpoint_weights( + _load_checkpoint_into_model( checkpoint_loader, "/checkpoint", - model, - object(), - {"mapping": object()}, + model=SimpleNamespace(load_weights=object()), + is_draft_model=False, + mapping=object(), + weight_mapper=weight_mapper, + extra_load_weights_kwargs=None, + metrics={}, ) + # The plan is inspected after the copy. assert events == [ - "session_enter", - "mapper_init", + "read:/checkpoint", + "materialize", "catalog", "plan", - "materialize", - "session_exit", ] - loader._call_load_weights.assert_called_once() + call_load_weights.assert_called_once() @pytest.mark.parametrize("failure_site", ["catalog", "plan"]) @@ -720,12 +875,12 @@ def test_shadow_plan_failure_warns_and_continues( weight_mapper.build_weight_load_plan.assert_not_called() -def test_draft_session_spans_mapper_and_materialization( +def test_draft_checkpoint_loads_after_target_with_its_own_mapper( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setenv("TRTLLM_SHADOW_WEIGHT_LOAD_PLAN", "1") events = [] - checkpoint_loader = _SessionCheckpointLoader(events, "/draft") + checkpoint_loader = _RecordingCheckpointLoader(events) catalog = CheckpointCatalog(objects=(), tensors=(CheckpointTensor("draft.weight"),)) plan = WeightLoadPlan( catalog_id=catalog.catalog_id, @@ -736,48 +891,71 @@ def test_draft_session_spans_mapper_and_materialization( demands=(WeightDemand("all", ("draft.weight",), (0,)),), ) checkpoint_loader.build_checkpoint_catalog = MagicMock( - side_effect=lambda *_args, **_kwargs: events.append("catalog") or catalog - ) - model = MagicMock() - model.draft_config = SimpleNamespace( - pretrained_config=SimpleNamespace(architectures=["DraftForCausalLM"]) - ) - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = {} - loader.spec_config = SimpleNamespace(speculative_model="/draft") - loader.mapping = object() - loader._call_load_weights = MagicMock( - side_effect=lambda *_args, **_kwargs: events.append("materialize") + side_effect=lambda checkpoint_dir, **_kwargs: ( + events.append(f"catalog:{checkpoint_dir}") or catalog + ) ) + _record_weight_copies(monkeypatch, events) draft_mapper = MagicMock() - draft_mapper.init_model_and_config.side_effect = lambda *_args: events.append("mapper_init") + draft_mapper.init_model_and_config.side_effect = lambda *_args: events.append( + "draft_mapper_init" + ) draft_mapper.build_weight_load_plan.side_effect = lambda _catalog: events.append("plan") or plan - monkeypatch.setattr( - model_loader_module.AutoCheckpointMapper, "get", MagicMock(return_value=draft_mapper) + get_mapper = MagicMock(return_value=draft_mapper) + monkeypatch.setattr(base_checkpoint_loader.AutoCheckpointMapper, "get", get_mapper) + model = SimpleNamespace( + load_weights=object(), + load_draft_weights=object(), + draft_model=object(), + draft_config=SimpleNamespace( + pretrained_config=SimpleNamespace(architectures=["DraftForCausalLM"]) + ), ) + mapping = object() + metrics = {} - loader._materialize_draft_checkpoint_weights(checkpoint_loader, model) - + load_model_from_auto_load_format( + checkpoint_loader, + "/checkpoint", + model_config=object(), + model=model, + mapping=mapping, + llm_args=None, + spec_config=SimpleNamespace(needs_separate_draft_weights=True, speculative_model="/draft"), + source_identity=None, + qualification=None, + metrics=metrics, + startup_metadata={}, + ) + + # The target mapper has no plan builder, so only the draft read is inspected. assert events == [ - "session_enter", "mapper_init", - "catalog", - "plan", + "read:/checkpoint", + "materialize", + "draft_mapper_init", + "read:/draft", "materialize", - "session_exit", + "catalog:/draft", + "plan", ] - assert "draft_checkpoint_preparation_seconds" in loader.metrics - assert "draft_weight_population_seconds" in loader.metrics - assert "draft_checkpoint_finalization_seconds" in loader.metrics + get_mapper.assert_called_once_with("HF", "DraftForCausalLM") + draft_mapper.init_model_and_config.assert_called_once_with( + model.draft_model, model.draft_config + ) + # The draft read does not get the model. + assert checkpoint_loader.read_kwargs["/draft"] == {"mapping": mapping} + assert "draft_checkpoint_preparation_seconds" in metrics + assert "draft_weight_population_seconds" in metrics + assert metrics["draft_checkpoint_finalization_seconds"] == 0.0 -def test_mtp_draft_session_reuses_target_mapper_during_materialization( - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize("target_preloaded", [False, True]) +def test_mtp_draft_reuses_target_mapper( + monkeypatch: pytest.MonkeyPatch, target_preloaded: bool ) -> None: monkeypatch.setenv("TRTLLM_SHADOW_WEIGHT_LOAD_PLAN", "1") events = [] - checkpoint_loader = _SessionCheckpointLoader(events, "/draft") catalog = CheckpointCatalog(objects=(), tensors=(CheckpointTensor("draft.weight"),)) plan = WeightLoadPlan( catalog_id=catalog.catalog_id, @@ -787,33 +965,55 @@ def test_mtp_draft_session_reuses_target_mapper_during_materialization( ordering=WeightLoadOrderConfidence.OPAQUE, demands=(WeightDemand("all", ("draft.weight",), (0,)),), ) - checkpoint_loader.build_checkpoint_catalog = MagicMock( - side_effect=lambda *_args, **_kwargs: events.append("catalog") or catalog - ) - model = MagicMock() - model.draft_config = None - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = {} - loader.spec_config = SimpleNamespace(speculative_model="/draft") - loader.mapping = object() - loader.weight_mapper = SimpleNamespace( + # The draft read reuses the target mapper, so that mapper also builds the + # draft's shadow plan. + target_mapper = SimpleNamespace( build_weight_load_plan=MagicMock(side_effect=lambda _catalog: events.append("plan") or plan) ) - loader._call_load_weights = MagicMock( - side_effect=lambda *_args, **_kwargs: events.append("materialize") + # An MX P2P target read writes the weights into the model and returns none + # to copy. The draft read is still copied, even when it reads nothing, so + # the model's draft load can reject an empty draft checkpoint. + checkpoint_loader = _RecordingCheckpointLoader( + events, + weights={} if target_preloaded else None, + weights_preloaded=target_preloaded, + target_weight_mapper=target_mapper, ) + checkpoint_loader.build_checkpoint_catalog = MagicMock( + side_effect=lambda checkpoint_dir, **_kwargs: ( + events.append(f"catalog:{checkpoint_dir}") or catalog + ) + ) + call_load_weights = _record_weight_copies(monkeypatch, events) + model = SimpleNamespace(load_weights=object(), load_draft_weights=object(), draft_config=None) - loader._materialize_draft_checkpoint_weights(checkpoint_loader, model) - - assert loader._call_load_weights.call_args.args[2] is loader.weight_mapper - assert events == [ - "session_enter", - "catalog", + weight_mapper, _ = load_model_from_auto_load_format( + checkpoint_loader, + "/checkpoint", + model_config=object(), + model=model, + mapping=object(), + llm_args=None, + spec_config=SimpleNamespace(needs_separate_draft_weights=True, speculative_model="/draft"), + source_identity=None, + qualification=None, + metrics={}, + startup_metadata={}, + ) + + assert weight_mapper is target_mapper + draft_copy = call_load_weights.call_args + assert draft_copy.args[0] is model.load_draft_weights + assert draft_copy.args[2] is target_mapper + expected_events = ["mapper_init", "read:/checkpoint"] + if not target_preloaded: + expected_events.append("materialize") + expected_events += [ + "catalog:/checkpoint", "plan", + "read:/draft", "materialize", - "session_exit", + "catalog:/draft", + "plan", ] - assert "draft_checkpoint_preparation_seconds" in loader.metrics - assert "draft_weight_population_seconds" in loader.metrics - assert "draft_checkpoint_finalization_seconds" in loader.metrics + assert events == expected_events diff --git a/tests/unittest/_torch/pyexecutor/test_model_loader_gms.py b/tests/unittest/_torch/pyexecutor/test_model_loader_gms.py index c55dab737340..79c48f57073b 100644 --- a/tests/unittest/_torch/pyexecutor/test_model_loader_gms.py +++ b/tests/unittest/_torch/pyexecutor/test_model_loader_gms.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for GMS-specific branches in ``ModelLoader``.""" +"""Unit tests for ``LoadFormat.GMS`` loading, driven through ``ModelLoader.load``.""" from collections.abc import Callable from contextlib import contextmanager, nullcontext @@ -14,6 +14,9 @@ from tensorrt_llm._torch import memory as memory_mod from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy +from tensorrt_llm._torch.models.checkpoints import memory_init, model_hooks +from tensorrt_llm._torch.models.checkpoints.gms import model_loading as gms_loading +from tensorrt_llm._torch.models.checkpoints.mx import model_loading as mx_loading from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE from tensorrt_llm._torch.moe.fused_moe.cutlass import TrtllmCutlassUnquantizedImpl from tensorrt_llm._torch.pyexecutor import model_loader as model_loader_mod @@ -73,6 +76,9 @@ def to(self, *args, **kwargs): def load_weights(self, weights, mapper): self._events.append("load_weights") + def load_draft_weights(self, weights, mapper): + self._events.append("load_draft_weights") + def setup_aliases(self) -> None: self._events.append("setup_aliases") @@ -109,8 +115,18 @@ def _make_loader(monkeypatch, *, events, spec_config=None): max_num_tokens=128, max_seq_len=128, ) - loader._call_load_weights = MagicMock( - side_effect=lambda fn, weights, mapper, **kwargs: fn(weights, mapper) + # Record each weight copy; tests read the calls from `model_hooks.call_load_weights`. + monkeypatch.setattr( + model_hooks, + "call_load_weights", + MagicMock(side_effect=lambda fn, weights, mapper, **kwargs: fn(weights, mapper)), + ) + # Spy on the GMS load so tests can read the source identity that + # `ModelLoader` built and passed in (see `_gms_load_source_identity`). + monkeypatch.setattr( + gms_loading, + "load_gms_weights", + MagicMock(wraps=gms_loading.load_gms_weights), ) loader._load_and_validate_config = MagicMock( return_value=SimpleNamespace( @@ -125,7 +141,7 @@ def _make_loader(monkeypatch, *, events, spec_config=None): monkeypatch.setattr(model_loader_mod, "timing_metric", lambda *_args, **_kwargs: nullcontext()) monkeypatch.setattr(model_loader_mod, "maybe_create_moe_load_balancer", _moe_context) - monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext()) + monkeypatch.setattr(memory_init, "MetaInitMode", lambda: nullcontext()) # These tests stub ModelConfig, while SourceIdentity has dedicated # coverage. Keep this file focused on ModelLoader GMS branch behavior. @@ -152,16 +168,20 @@ def _build_source_identity(_cls, *_args, **kwargs): classmethod(_build_source_identity), ) monkeypatch.setattr( - model_loader_mod.AutoModelForCausalLM, + memory_init.AutoModelForCausalLM, "from_config", MagicMock(return_value=_TinyModel(events, include_draft=spec_config is not None)), ) - monkeypatch.setattr(model_loader_mod, "get_rank_model_storage", lambda _model: 0) monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) return loader +def _gms_load_source_identity() -> model_loader_mod.SourceIdentity: + """Return the source identity `ModelLoader` passed to the GMS load.""" + return gms_loading.load_gms_weights.call_args.kwargs["source_identity"] + + def _build_gms_backend(*, is_rw, events): backend = MagicMock() backend.connect.return_value = True @@ -221,7 +241,7 @@ def _tiny_profile_registry() -> PostTransformProfileRegistry: architecture="TinyForCausalLM", model_type="tiny", speculative_mode=None, - protocol_version=(ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION), + protocol_version=mx_loading.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_constraints=PostTransformRuntimeConstraints(), @@ -318,17 +338,21 @@ def test_gms_load_branch(monkeypatch, is_rw, expected_events): "/ckpt", mapping=loader.mapping, model=model, - source_identity=loader._source_identity, + source_identity=_gms_load_source_identity(), ) - loader._call_load_weights.assert_called_once() + model_hooks.call_load_weights.assert_called_once() backend.move_untracked_params.assert_called_once_with(model) backend.finalize_write.assert_called_once_with(model) + # The writer's weight mapper reaches `ModelLoader` for `reload()`. + assert loader.weight_mapper is checkpoint_loader.get_initialized_weight_mapper.return_value else: # RO: setup_aliases() must run before the GMS materialize step so # module aliases are wired up before zero-copy mapping. checkpoint_loader.load_weights.assert_not_called() - loader._call_load_weights.assert_not_called() + model_hooks.call_load_weights.assert_not_called() backend.materialize_module.assert_called_once_with(model) + # Readers map weights without a weight mapper. + assert loader.weight_mapper is None @pytest.mark.parametrize( @@ -357,7 +381,7 @@ def test_gms_checkpoint_policy_records_only_writer_disk_load( "fallback_reason": weight_loader.last_checkpoint_io_status.fallback_reason, } assert "load_format='auto'" in expected["fallback_reason"] - model = model_loader_mod.AutoModelForCausalLM.from_config.return_value + model = memory_init.AutoModelForCausalLM.from_config.return_value if include_draft: model.draft_config = None model.load_draft_weights = MagicMock() @@ -418,8 +442,8 @@ def _append(*_args, **_kwargs): # materialization; record it to pin the ordering without exercising the # comparison logic, which is covered in test_source_identity.py. monkeypatch.setattr( - model_loader_mod, - "check_weight_sharing_compatibility", + gms_loading, + "check_gms_source_identity", lambda *_args, **_kwargs: events.append("check_source_identity"), ) @@ -455,7 +479,7 @@ def cache_derived_state(self): raise AssertionError("GMS read-only weights must not be transformed again") self.cache_rebuilt = True - model = model_loader_mod.AutoModelForCausalLM.from_config.return_value + model = memory_init.AutoModelForCausalLM.from_config.return_value model.projection = _AlreadyTransformedProjection() checkpoint_loader = MagicMock(name="checkpoint_loader") checkpoint_loader.checkpoint_format = "HF" @@ -473,8 +497,8 @@ def test_gms_ro_marks_moe_backend_behind_removed_wrapper_transformed(monkeypatch The RW writer already transformed the expert bytes in the pool. A MoE backend behind the `_weights_removed` `ConfigurableMoE` wrapper never loads on the RO path, so only its guard declared before any load lets - `_mark_weights_transformed` mark it, and a later transform walk must then - leave it alone. + `model_hooks.mark_weights_transformed` mark it, and the load must leave its + weights untransformed. """ events = [] loader = _make_loader(monkeypatch, events=events) @@ -494,7 +518,7 @@ def test_gms_ro_marks_moe_backend_behind_removed_wrapper_transformed(monkeypatch wrapper.layer_load_balancer = None wrapper._weights_removed = True wrapper.backend = experts - model = model_loader_mod.AutoModelForCausalLM.from_config.return_value + model = memory_init.AutoModelForCausalLM.from_config.return_value model.experts = wrapper checkpoint_loader = MagicMock(name="checkpoint_loader") checkpoint_loader.checkpoint_format = "HF" @@ -506,7 +530,6 @@ def test_gms_ro_marks_moe_backend_behind_removed_wrapper_transformed(monkeypatch assert experts._weights_transformed is True assert wrapper._weights_transformed is False experts.quant_method.cache_derived_state.assert_called_once_with(experts) - ModelLoader._walk_transform(loaded_model) experts.transform_weights.assert_not_called() @@ -514,13 +537,19 @@ def test_gms_rw_post_load_runs_inside_pool_before_finalize(monkeypatch): """Every step that may allocate or rebind tensors must run inside the GMS pool. The widened ``mem_pool_scope`` covers meta-tensor materialization, - ``model.to('cuda')``, weight load, and the post_load_* hooks. Only - ``finalize_write`` runs after the pool closes, so the committed - layout reflects the post-post_load model. Asserts the exact ordering - so a future refactor cannot silently re-narrow the scope. + ``model.to('cuda')``, weight load, draft weight load, and the + post_load_* hooks. Only ``finalize_write`` runs after the pool closes, + so the committed layout reflects the post-post_load model. Asserts the + exact ordering so a future refactor cannot silently re-narrow the scope. """ events = [] - loader = _make_loader(monkeypatch, events=events) + # The spec config adds a separate draft checkpoint, so the writer also + # loads draft weights inside the pool. + loader = _make_loader( + monkeypatch, + events=events, + spec_config=_spec_config_needing_draft_weights(), + ) backend = _build_gms_backend(is_rw=True, events=events) backend.move_untracked_params.side_effect = lambda _model: events.append( @@ -539,13 +568,14 @@ def test_gms_rw_post_load_runs_inside_pool_before_finalize(monkeypatch): "post_load_publish" ) - loader.load("/ckpt", checkpoint_loader) + model, _ = loader.load("/ckpt", checkpoint_loader) assert events == [ "pool_enter", "_apply", "to", "load_weights", + "load_draft_weights", "post_load_apply", "post_load_weights", "move_untracked_params", @@ -565,6 +595,7 @@ def test_gms_rw_post_load_runs_inside_pool_before_finalize(monkeypatch): "_apply", "to", "load_weights", + "load_draft_weights", "post_load_apply", "post_load_publish", "post_load_weights", @@ -575,6 +606,18 @@ def test_gms_rw_post_load_runs_inside_pool_before_finalize(monkeypatch): f"{inside!r} (idx={idx}) must run inside the pool (enter={enter_idx}, exit={exit_idx})" ) + # The HF draft read gets only the mapping, and the draft weights go + # through the draft mapper built from the target mapper. + checkpoint_loader.load_weights.assert_called_with("/draft", mapping=loader.mapping) + checkpoint_loader.get_initialized_draft_weight_mapper.assert_called_once_with( + model, checkpoint_loader.get_initialized_weight_mapper.return_value + ) + model_hooks.call_load_weights.assert_called_with( + model.load_draft_weights, + checkpoint_loader.load_weights.return_value, + checkpoint_loader.get_initialized_draft_weight_mapper.return_value, + ) + def test_gms_rw_loader_preload_skips_mapping_pipeline(monkeypatch): """RW + empty ``weights`` + ``is_weights_preloaded()=True`` is legitimate. @@ -596,7 +639,8 @@ def test_gms_rw_loader_preload_skips_mapping_pipeline(monkeypatch): model, _ = loader.load("/ckpt", checkpoint_loader) - loader._call_load_weights.assert_not_called() + model_hooks.call_load_weights.assert_not_called() + assert loader.weight_mapper is None backend.move_untracked_params.assert_called_once_with(model) backend.finalize_write.assert_called_once_with(model) @@ -606,8 +650,8 @@ def test_gms_rw_mx_post_transform_preload_uses_staged_path(monkeypatch): events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) backend = _build_gms_backend(is_rw=True, events=events) @@ -645,13 +689,14 @@ def test_gms_rw_mx_post_transform_preload_uses_staged_path(monkeypatch): _args, kwargs = checkpoint_loader.load_weights.call_args assert kwargs["allow_post_transform_weights"] is True assert callable(kwargs["prepare_post_transform_receiver"]) - assert loader._source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 - loader._call_load_weights.assert_not_called() + source_identity = _gms_load_source_identity() + assert source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 + model_hooks.call_load_weights.assert_not_called() checkpoint_loader.post_load_publish.assert_called_once_with( model, checkpoint_dir="/ckpt", weights_preloaded=True, - source_identity=loader._source_identity, + source_identity=source_identity, ) backend.finalize_write.assert_called_once_with(model) @@ -686,7 +731,7 @@ def test_gms_rw_exception_during_weight_mapping_cleans_up(monkeypatch): """A writer error after pool population must release the GMS session.""" events = [] loader = _make_loader(monkeypatch, events=events) - loader._call_load_weights.side_effect = RuntimeError("mapping failed") + model_hooks.call_load_weights.side_effect = RuntimeError("mapping failed") backend = _build_gms_backend(is_rw=True, events=events) _install_gms_backend(monkeypatch, backend) @@ -781,5 +826,5 @@ def test_gms_rejects_localized_weights_before_model_init(monkeypatch): ): loader.load("/ckpt", checkpoint_loader) - model_loader_mod.AutoModelForCausalLM.from_config.assert_not_called() + memory_init.AutoModelForCausalLM.from_config.assert_not_called() checkpoint_loader.load_weights.assert_not_called() diff --git a/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py b/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py index 5d16eaf3823b..61039e66231e 100644 --- a/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py +++ b/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py @@ -48,7 +48,11 @@ from tensorrt_llm._torch.models import modeling_qwen as modeling_qwen_mod from tensorrt_llm._torch.models import modeling_qwen3 as modeling_qwen3_mod from tensorrt_llm._torch.models import modeling_qwen3_moe as modeling_qwen3_moe_mod +from tensorrt_llm._torch.models.checkpoints import memory_init as memory_init_mod +from tensorrt_llm._torch.models.checkpoints import model_hooks as model_hooks_mod +from tensorrt_llm._torch.models.checkpoints.gms import model_loading as gms_loading_mod from tensorrt_llm._torch.models.checkpoints.hf.qwen3_moe_weight_mapper import Qwen3MoeHfWeightMapper +from tensorrt_llm._torch.models.checkpoints.mx import model_loading as mx_loading_mod from tensorrt_llm._torch.models.checkpoints.mx.checkpoint_loader import MXCheckpointLoader from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.models.modeling_utils import MetaInitMode, get_registered_model_class @@ -433,7 +437,7 @@ def _simulate_sm(monkeypatch: pytest.MonkeyPatch, sm_version: int) -> None: monkeypatch.setattr(linear_mod, "get_sm_version", lambda: sm_version) monkeypatch.setattr(linear_mod, "is_sm_100f", lambda *_args, **_kwargs: 100 <= sm_version < 110) monkeypatch.setattr(linear_mod, "is_sm_120f", lambda *_args, **_kwargs: 120 <= sm_version < 130) - monkeypatch.setattr(model_loader_mod, "get_sm_version", lambda: sm_version) + monkeypatch.setattr(mx_loading_mod, "get_sm_version", lambda: sm_version) def _fp8_block_quant_config() -> QuantConfig: @@ -1075,7 +1079,7 @@ def _tiny_profile_registry(*, speculative_mode: str | None = None) -> PostTransf architecture="TinyForCausalLM", model_type="tiny", speculative_mode=speculative_mode, - protocol_version=(ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION), + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_constraints=PostTransformRuntimeConstraints(), @@ -1094,8 +1098,11 @@ def _make_loader(monkeypatch, *, events, spec_config=None): max_num_tokens=128, max_seq_len=128, ) - loader._call_load_weights = MagicMock( - side_effect=lambda fn, weights, mapper, **kwargs: fn(weights, mapper) + # Record each weight copy; tests read the calls from `model_hooks_mod.call_load_weights`. + monkeypatch.setattr( + model_hooks_mod, + "call_load_weights", + MagicMock(side_effect=lambda fn, weights, mapper, **kwargs: fn(weights, mapper)), ) loader._load_and_validate_config = MagicMock( return_value=SimpleNamespace( @@ -1110,7 +1117,7 @@ def _make_loader(monkeypatch, *, events, spec_config=None): monkeypatch.setattr(model_loader_mod, "timing_metric", lambda *_args, **_kwargs: nullcontext()) monkeypatch.setattr(model_loader_mod, "maybe_create_moe_load_balancer", _moe_context) - monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext()) + monkeypatch.setattr(memory_init_mod, "MetaInitMode", lambda: nullcontext()) # These tests stub ModelConfig, while SourceIdentity has dedicated # coverage. Keep this file focused on ModelLoader MX branch behavior. @@ -1138,14 +1145,21 @@ def _build_source_identity(_cls, *_args, **kwargs): classmethod(_build_source_identity), ) monkeypatch.setattr( - model_loader_mod.AutoModelForCausalLM, + memory_init_mod.AutoModelForCausalLM, "from_config", MagicMock(return_value=_TinyModel(events)), ) - monkeypatch.setattr(model_loader_mod, "get_rank_model_storage", lambda _model: 0) + monkeypatch.setattr(memory_init_mod, "get_rank_model_storage", lambda _model: 0) return loader +def _mock_mx_checkpoint_loader() -> MagicMock: + """Return a mocked MX checkpoint loader; `ModelLoader.load()` calls its `load_weights`.""" + checkpoint_loader = MagicMock(spec=MXCheckpointLoader, name="checkpoint_loader") + checkpoint_loader.checkpoint_format = "MX" + return checkpoint_loader + + @pytest.mark.cpu_only def test_construct_checkpoint_loader_passes_mx_config(): mx_config = SimpleNamespace( @@ -1286,7 +1300,7 @@ def _documented_constraints(profile: PostTransformProfile) -> str: @pytest.mark.cpu_only def test_public_support_table_matches_qualified_profile_registry() -> None: - profiles = ModelLoader._post_transform_profile_registry().profiles + profiles = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles documentation = (Path(__file__).parents[4] / "docs/source/features/model-express.md").read_text( encoding="utf-8" ) @@ -1478,7 +1492,7 @@ def test_mx_e2e_rows_and_ci_list_match_qualified_profile_registry() -> None: scheme, and `l0_model_express.yml` must list each row exactly once, in a block whose GPU count matches the row's two TP groups. """ - profiles = ModelLoader._post_transform_profile_registry().profiles + profiles = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles assert set(_MX_E2E_ROWS) == {profile.profile_id for profile in profiles}, ( "Every qualified profile needs exactly one `_MX_E2E_ROWS` entry" ) @@ -1538,7 +1552,7 @@ def test_documentation_prose_matches_qualified_profile_registry() -> None: limitations list, and the case IDs and `TRTLLM_MX_*_MODEL` variables named in the "Qualification Test" section. """ - profiles = ModelLoader._post_transform_profile_registry().profiles + profiles = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles families = {_documented_family(profile) for profile in profiles} text_only_families = { _documented_family(profile) @@ -1590,12 +1604,11 @@ def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works( events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.is_weights_preloaded.return_value = True checkpoint_loader.load_weights.return_value = {} @@ -1605,34 +1618,34 @@ def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works( _args, kwargs = checkpoint_loader.load_weights.call_args assert kwargs["mapping"] is loader.mapping assert kwargs["model"] is model - assert kwargs["source_identity"] is loader._source_identity + source_identity = kwargs["source_identity"] assert ( - kwargs["model_config"] - is model_loader_mod.AutoModelForCausalLM.from_config.call_args.args[0] + kwargs["model_config"] is memory_init_mod.AutoModelForCausalLM.from_config.call_args.args[0] ) assert kwargs["load_config"] is loader.llm_args assert ( kwargs["post_transform_protocol_version"] - == ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION + == mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION ) assert kwargs["allow_post_transform_weights"] is True - assert loader._source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 - assert loader._call_load_weights.call_count == 0 + assert source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 + assert model_hooks_mod.call_load_weights.call_count == 0 checkpoint_loader.get_initialized_weight_mapper.assert_called_once() assert loader.weight_mapper is checkpoint_loader.get_initialized_weight_mapper.return_value + # The publish carries the same identity that the read got. checkpoint_loader.post_load_publish.assert_called_once_with( model, checkpoint_dir="/ckpt", weights_preloaded=True, - source_identity=loader._source_identity, + source_identity=source_identity, ) # reload() uses self.weight_mapper unconditionally; MX success must - # initialize it even though the initial load skipped _call_load_weights. + # initialize it even though the initial load skipped call_load_weights. model._weights_transformed = True model.linear._weights_transformed = True loader.reload(model, {"reloaded": MagicMock()}) - assert loader._call_load_weights.call_count == 1 + assert model_hooks_mod.call_load_weights.call_count == 1 assert model._weights_transformed is False assert model.linear._weights_transformed is False assert events == ["post_load_weights", "load_weights"] @@ -1640,8 +1653,7 @@ def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works( def test_cleanup_releases_active_checkpoint_loader(monkeypatch): loader = _make_loader(monkeypatch, events=[]) - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.is_weights_preloaded.return_value = False checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} @@ -1670,16 +1682,12 @@ def test_cleanup_swallows_checkpoint_loader_failure(monkeypatch): ) -def test_cleanup_continues_after_gms_failure(monkeypatch): +def test_cleanup_releases_gms_backend_and_checkpoint_loader(monkeypatch): loader = _make_loader(monkeypatch, events=[]) gms_backend = MagicMock(name="gms_backend") - cleanup_error = RuntimeError("gms cleanup failed") - gms_backend.cleanup.side_effect = cleanup_error checkpoint_loader = MagicMock(name="checkpoint_loader") loader._gms_backend = gms_backend loader._checkpoint_loader = checkpoint_loader - warning = MagicMock() - monkeypatch.setattr(model_loader_mod.logger, "warning", warning) loader.cleanup() @@ -1687,9 +1695,6 @@ def test_cleanup_continues_after_gms_failure(monkeypatch): checkpoint_loader.cleanup.assert_called_once_with() assert loader._gms_backend is None assert loader._checkpoint_loader is None - warning.assert_called_once_with( - f"Failed to clean up GMS backend {gms_backend!r}: {cleanup_error!r}" - ) def test_config_failure_retains_checkpoint_loader_for_cleanup(monkeypatch): @@ -1716,8 +1721,8 @@ def test_reload_partial_loading_preserves_weights_transformed_flags(monkeypatch) loader.reload(model, {"reloaded": MagicMock()}, allow_partial_loading=True) - assert loader._call_load_weights.call_count == 1 - assert loader._call_load_weights.call_args.kwargs["allow_partial_loading"] is True + assert model_hooks_mod.call_load_weights.call_count == 1 + assert model_hooks_mod.call_load_weights.call_args.kwargs["allow_partial_loading"] is True assert model._weights_transformed is True assert model.linear._weights_transformed is True assert events == ["load_weights"] @@ -1730,20 +1735,19 @@ def test_mx_partial_fallback_merges_returned_weights( events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.is_weights_preloaded.return_value = True fallback_weights = {"mismatched.weight": MagicMock()} checkpoint_loader.load_weights.return_value = fallback_weights model, _ = loader.load("/ckpt", checkpoint_loader) - assert loader._call_load_weights.call_count == 1 - load_fn, weights, mapper = loader._call_load_weights.call_args.args + assert model_hooks_mod.call_load_weights.call_count == 1 + load_fn, weights, mapper = model_hooks_mod.call_load_weights.call_args.args assert load_fn == model.load_weights assert weights is fallback_weights assert mapper is loader.weight_mapper @@ -1751,14 +1755,13 @@ def test_mx_partial_fallback_merges_returned_weights( model, checkpoint_dir="/ckpt", weights_preloaded=True, - source_identity=loader._source_identity, + source_identity=checkpoint_loader.load_weights.call_args.kwargs["source_identity"], ) -class _PostTransformMxLoader: - checkpoint_format = "MX" - +class _PostTransformMxLoader(MXCheckpointLoader): def __init__(self, *, post_transform: bool) -> None: + super().__init__() self._post_transform = post_transform self._weights_preloaded = True self._disk_weight = MagicMock() @@ -1796,32 +1799,31 @@ def test_mx_post_transform_receiver_uses_staged_path_when_qualified( events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) checkpoint_loader = _PostTransformMxLoader(post_transform=True) model, _ = loader.load("/ckpt", checkpoint_loader) - loader._call_load_weights.assert_not_called() + model_hooks_mod.call_load_weights.assert_not_called() _args, kwargs = checkpoint_loader.load_weights.call_args assert kwargs["allow_post_transform_weights"] is True assert ( - kwargs["model_config"] - is model_loader_mod.AutoModelForCausalLM.from_config.call_args.args[0] + kwargs["model_config"] is memory_init_mod.AutoModelForCausalLM.from_config.call_args.args[0] ) assert kwargs["load_config"] is loader.llm_args assert ( kwargs["post_transform_protocol_version"] - == ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION + == mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION ) assert callable(kwargs["prepare_post_transform_receiver"]) checkpoint_loader.post_load_publish.assert_called_once_with( model, checkpoint_dir="/ckpt", weights_preloaded=True, - source_identity=loader._source_identity, + source_identity=kwargs["source_identity"], ) # Post-transform receivers skip transform_weights(), but the accepted # tensors are already in final layout. Keep the transform guard in sync so @@ -1879,7 +1881,7 @@ def _tiny_multimodal_profile( architecture="TinyForCausalLM", model_type="tiny", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, transfer_scope=PostTransformTransferScope.COMPLETE_MODEL, runtime_constraints=PostTransformRuntimeConstraints(), @@ -1898,8 +1900,8 @@ def test_post_transform_transfer_scope_is_derived_from_the_root( plain load. """ monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", PostTransformProfileRegistry( profiles=( _tiny_multimodal_profile(_TinyMultimodalModel), @@ -1907,7 +1909,7 @@ def test_post_transform_transfer_scope_is_derived_from_the_root( ) ), ) - derive_scope = ModelLoader._post_transform_transfer_scope + derive_scope = mx_loading_mod.get_post_transform_transfer_scope assert derive_scope(_TinyModel([])) is PostTransformTransferScope.TARGET_MODEL assert ( @@ -1946,12 +1948,12 @@ def test_mx_load_offers_post_transform_weights_only_for_the_registered_scope( events: list[str] = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", PostTransformProfileRegistry(profiles=(_tiny_multimodal_profile(_TinyMultimodalModel),)), ) monkeypatch.setattr( - model_loader_mod.AutoModelForCausalLM, + memory_init_mod.AutoModelForCausalLM, "from_config", MagicMock(return_value=_TinyMultimodalModel(events, with_encoder=with_encoder)), ) @@ -1961,18 +1963,19 @@ def test_mx_load_offers_post_transform_weights_only_for_the_registered_scope( _args, kwargs = checkpoint_loader.load_weights.call_args assert kwargs["allow_post_transform_weights"] is with_encoder + source_identity = kwargs["source_identity"] if with_encoder: - assert loader._source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 - loader._call_load_weights.assert_not_called() + assert source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 + model_hooks_mod.call_load_weights.assert_not_called() checkpoint_loader.post_load_publish.assert_called_once_with( model, checkpoint_dir="/ckpt", weights_preloaded=True, - source_identity=loader._source_identity, + source_identity=source_identity, ) else: - assert loader._source_identity.transform_abi_id is None - assert loader._call_load_weights.call_count == 1 + assert source_identity.transform_abi_id is None + assert model_hooks_mod.call_load_weights.call_count == 1 checkpoint_loader.post_load_publish.assert_not_called() @@ -1988,7 +1991,7 @@ def test_qualified_profiles_bind_published_transform_abis() -> None: for name in weight_sharing_mod.__all__ if re.fullmatch(r".*_POST_TRANSFORM_LAYOUT_ABI_V\d+", name) } - profiles = ModelLoader._post_transform_profile_registry().profiles + profiles = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles assert {profile.transform_abi_id for profile in profiles} <= published @@ -2003,7 +2006,7 @@ def test_default_profile_qualifies_real_tiny_llama_lifecycle( monkeypatch, model_class=_UnqualifiedLlamaForCausalLM, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2041,7 +2044,7 @@ def test_default_profile_qualifies_real_tiny_llama_tp2_rank_lifecycle( tp_size=2, rank=rank, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2068,7 +2071,7 @@ def test_qwen2_dense_profile_qualifies_full_staged_lifecycle() -> None: unqualified_model_factory=lambda: _tiny_qwen2_model( model_class=_UnqualifiedQwen2ForCausalLM ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2105,7 +2108,7 @@ def test_qwen2_dense_profile_rejects_effective_unfused_rope() -> None: kv_cache_compression_config=SimpleNamespace(changes_physical_kv_length=True) ) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2132,7 +2135,7 @@ def test_qwen2_dense_profile_qualifies_tp2_rank_lifecycle( tp_size=2, rank=rank, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2167,7 +2170,7 @@ def test_qwen3_dense_profile_qualifies_full_staged_lifecycle( monkeypatch, model_class=_UnqualifiedQwen3ForCausalLM, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2221,7 +2224,7 @@ def test_qwen3_dense_profile_qualifies_tp2_rank_lifecycle( tp_size=2, rank=rank, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2297,7 +2300,7 @@ def test_qwen3_dense_tied_profile_qualifies_full_staged_lifecycle( model_class=_UnqualifiedQwen3ForCausalLM, tie_word_embeddings=True, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2316,7 +2319,7 @@ def test_qwen3_dense_tied_profile_qualifies_full_staged_lifecycle( producer, receiver = assert_post_transform_lifecycle_equivalent(case) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( receiver, speculative_mode=None, loads_draft_weights=False, @@ -2355,7 +2358,7 @@ def test_qwen3_dense_tied_profile_qualifies_tp2_rank_lifecycle( rank=rank, tie_word_embeddings=True, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2436,13 +2439,13 @@ def test_staged_qwen3_tied_finalization_preserves_mx_tensor_catalog( model = _tiny_qwen3_model(monkeypatch, tie_word_embeddings=True) # MX discovers and registers tensors after receiver alias preparation. - ModelLoader._setup_aliases(model) + model_hooks_mod.setup_aliases(model) registered_catalog = _mx_canonical_parameter_catalog(model) # This is the receiver-side finalization sequence after RDMA completes. - ModelLoader._setup_aliases(model) - ModelLoader._mark_weights_transformed(model) - ModelLoader._walk_cache_state(model) + model_hooks_mod.setup_aliases(model) + model_hooks_mod.mark_weights_transformed(model) + model_hooks_mod.walk_cache_state(model) published_catalog = _mx_canonical_parameter_catalog(model) assert published_catalog == registered_catalog @@ -2467,7 +2470,7 @@ def test_qwen3_dense_profiles_reject_tie_the_model_does_not_realize( else: model.lm_head.weight = model.model.embed_tokens.weight - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2487,7 +2490,7 @@ def test_qwen3_tied_and_untied_source_identities_never_pair( (True, QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1), ): model = _tiny_qwen3_model(monkeypatch, tie_word_embeddings=tie_word_embeddings) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2534,7 +2537,7 @@ def test_mx_staged_receiver_load_keeps_tied_qwen3_table_in_one_allocation( monkeypatch: pytest.MonkeyPatch, ) -> None: producer = _tiny_qwen3_model(monkeypatch, tie_word_embeddings=True).cuda() - ModelLoader._walk_full_post_load(producer) + model_hooks_mod.walk_full_post_load(producer) real_source_identity = model_loader_mod.SourceIdentity.__dict__["from_model_config"] loader = _make_loader(monkeypatch, events=[]) @@ -2556,16 +2559,17 @@ def _construct(model_config: ModelConfig) -> nn.Module: return model # Run the loader's real meta-init construction and memoized materialization. - monkeypatch.setattr(model_loader_mod, "MetaInitMode", MetaInitMode) - monkeypatch.setattr(model_loader_mod.AutoModelForCausalLM, "from_config", _construct) + monkeypatch.setattr(memory_init_mod, "MetaInitMode", MetaInitMode) + monkeypatch.setattr(memory_init_mod.AutoModelForCausalLM, "from_config", _construct) checkpoint_loader = _CatalogCopyingMxLoader(producer) model, _ = loader.load("/ckpt", checkpoint_loader) assert construction_devices == ["meta"] - loader._call_load_weights.assert_not_called() + model_hooks_mod.call_load_weights.assert_not_called() assert checkpoint_loader.load_weights.call_args.kwargs["allow_post_transform_weights"] is True - assert loader._source_identity.transform_abi_id == QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1 + source_identity = checkpoint_loader.load_weights.call_args.kwargs["source_identity"] + assert source_identity.transform_abi_id == QWEN3_DENSE_TIED_POST_TRANSFORM_LAYOUT_ABI_V1 # Materialization may give each module its own Parameter object, so the tie # is checked through memory: one view of one allocation. @@ -2617,7 +2621,7 @@ def _qwen3_fp8_block_case( tp_size=tp_size, rank=rank, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2683,7 +2687,7 @@ def test_qwen3_fp8_block_profile_qualifies_full_staged_lifecycle( _qwen3_fp8_block_case(monkeypatch) ) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( receiver, speculative_mode=None, loads_draft_weights=False, @@ -2738,7 +2742,7 @@ def test_mistral_dense_profile_qualifies_full_staged_lifecycle() -> None: unqualified_model_factory=lambda: _tiny_mistral_model( model_class=_UnqualifiedMistralForCausalLM ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2785,7 +2789,7 @@ def test_mistral_dense_profile_qualifies_tp2_rank_lifecycle( tp_size=2, rank=rank, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2830,7 +2834,7 @@ def test_mistral_dense_profile_rejects_sliding_window_models( ) -> None: model = _tiny_mistral_model(sliding_window=sliding_window, layer_types=layer_types) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2851,7 +2855,7 @@ def test_mistral_dense_profile_qualifies_full_attention_layer_types() -> None: layer_types=("full_attention", "full_attention"), ) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -2866,14 +2870,14 @@ def test_mistral_dense_profile_qualifies_full_attention_layer_types() -> None: def _qwen3_moe_profile() -> PostTransformProfile: (profile,) = ( candidate - for candidate in ModelLoader._post_transform_profile_registry().profiles + for candidate in mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles if candidate.profile_id == _QWEN3_MOE_PROFILE_ID ) return profile def _qualify_qwen3_moe_model(model: nn.Module) -> PostTransformQualificationDecision: - return ModelLoader._qualify_post_transform_profile( + return mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -3123,12 +3127,12 @@ def test_staged_qwen3_moe_finalization_preserves_mx_tensor_catalog( """ model = _tiny_qwen3_moe_model(monkeypatch) - ModelLoader._setup_aliases(model) + model_hooks_mod.setup_aliases(model) registered_catalog = _mx_canonical_parameter_catalog(model) - ModelLoader._setup_aliases(model) - ModelLoader._mark_weights_transformed(model) - ModelLoader._walk_cache_state(model) + model_hooks_mod.setup_aliases(model) + model_hooks_mod.mark_weights_transformed(model) + model_hooks_mod.walk_cache_state(model) published_catalog = _mx_canonical_parameter_catalog(model) assert published_catalog == registered_catalog @@ -3188,12 +3192,12 @@ def test_qwen3_moe_profile_rejects_unqualified_moe_variants( `test_bf16_profiles_reject_unqualified_runtime_variants` covers the variants that every BF16 profile rejects. """ - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=modeling_qwen3_moe_mod.Qwen3MoeForCausalLM, architecture="Qwen3MoeForCausalLM", model_type="qwen3_moe", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_qwen3_moe_runtime_config(**overrides), ) @@ -3229,12 +3233,12 @@ def test_qwen3_moe_profile_rejects_unsupported_moe_features( def qualify( enabled_features: frozenset[PostTransformFeature], ) -> PostTransformQualificationDecision: - return ModelLoader._post_transform_profile_registry().qualify( + return mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=modeling_qwen3_moe_mod.Qwen3MoeForCausalLM, architecture="Qwen3MoeForCausalLM", model_type="qwen3_moe", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, enabled_features=enabled_features, runtime_config=_qwen3_moe_runtime_config(**overrides), @@ -3298,7 +3302,7 @@ def test_moe_profile_backend_literals_realize_one_implementation( """ moe_profiles = [ profile - for profile in ModelLoader._post_transform_profile_registry().profiles + for profile in mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles if profile.runtime_constraints.moe_backends is not None ] assert [profile.profile_id for profile in moe_profiles] == [_QWEN3_MOE_PROFILE_ID], ( @@ -3339,12 +3343,12 @@ def test_qwen3_vl_moe_root_does_not_inherit_qwen3_moe_profile() -> None: is modeling_qwen3_moe_mod.Qwen3MoeForCausalLM ) - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=vl_root, architecture="Qwen3MoeForCausalLM", model_type="qwen3_moe", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_qwen3_moe_runtime_config(), ) @@ -3369,7 +3373,7 @@ def test_moe_profiles_pin_partition_and_one_backend_per_abi() -> None: """ abi_owners: dict[str, tuple[str, str]] = {} moe_profiles = 0 - for profile in ModelLoader._post_transform_profile_registry().profiles: + for profile in mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles: constraints = profile.runtime_constraints if constraints.moe_backends is None: assert constraints.moe_tp_sizes is None @@ -3397,12 +3401,12 @@ def test_legacy_llama_file_mistral_root_is_not_qualified() -> None: get_registered_model_class("MistralForCausalLM") is modeling_mistral_mod.MistralForCausalLM ) - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=legacy_root, architecture="MistralForCausalLM", model_type="mistral", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config(sliding_window="none"), ) @@ -3413,12 +3417,12 @@ def test_legacy_llama_file_mistral_root_is_not_qualified() -> None: @pytest.mark.cpu_only def test_mistral_dense_profile_rejects_native_format_model_type() -> None: - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=modeling_mistral_mod.MistralForCausalLM, architecture="MistralForCausalLM", model_type="mistral_common", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config(sliding_window="none"), ) @@ -3433,7 +3437,7 @@ def test_phi3_dense_profile_qualifies_full_staged_lifecycle() -> None: profile_id="phi3-for-causal-lm-bf16-target-v1", model_factory=_tiny_phi3_model, unqualified_model_factory=lambda: _tiny_phi3_model(model_class=_UnqualifiedPhi3ForCausalLM), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -3448,7 +3452,7 @@ def test_phi3_dense_profile_qualifies_full_staged_lifecycle() -> None: producer, _receiver = assert_post_transform_lifecycle_equivalent(case) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( producer, speculative_mode=None, loads_draft_weights=False, @@ -3488,7 +3492,7 @@ def test_phi3_dense_profile_qualifies_tp2_rank_lifecycle( tp_size=2, rank=rank, ), - qualify_model=lambda model: ModelLoader._qualify_post_transform_profile( + qualify_model=lambda model: mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -3516,7 +3520,7 @@ def test_phi3_dense_profile_qualifies_tp2_rank_lifecycle( def test_phi3_dense_profile_rejects_unregistered_subclass_root() -> None: """An otherwise identical subclass must not inherit the exact-root profile.""" - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( _tiny_phi3_model(model_class=_UnqualifiedPhi3ForCausalLM), speculative_mode=None, loads_draft_weights=False, @@ -3534,7 +3538,7 @@ def test_phi3_dense_profile_realizes_configured_window_as_full_attention() -> No """ model = _tiny_phi3_model(sliding_window=8) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -3575,7 +3579,7 @@ def test_phi3_dense_profile_rejects_longrope_and_tied_embedding_variants( """LongRoPE and tied-embedding Phi3 variants fall outside the profile's constraints.""" model = _tiny_phi3_model(**variant) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -3626,12 +3630,12 @@ def test_phi4mm_root_does_not_inherit_phi3_profile() -> None: phi4mm_root = get_registered_model_class("Phi4MMForCausalLM") assert phi4mm_root is not None - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=phi4mm_root, architecture="Phi3ForCausalLM", model_type="phi3", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config(), ) @@ -3767,12 +3771,12 @@ def test_bf16_profiles_reject_unqualified_runtime_variants( The MoE-only variants of the Qwen3 MoE profile are in `test_qwen3_moe_profile_rejects_unqualified_moe_variants`. """ - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=root_model_class, architecture=architecture, model_type=model_type, speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config( rope_fusion=supported_rope_fusion, @@ -3853,12 +3857,12 @@ def test_bf16_dense_profiles_scope_tied_embeddings( """Only roots with a tied sibling profile qualify a realized tie; no root accepts an unknown tie.""" def qualify(tied_word_embeddings: bool | None): - return ModelLoader._post_transform_profile_registry().qualify( + return mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=root_model_class, architecture=architecture, model_type=model_type, speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config( rope_fusion=supported_rope_fusion, @@ -3900,12 +3904,12 @@ def test_qwen3_tied_profile_rejects_unqualified_runtime_variants( overrides: dict[str, object], expected_dimensions: set[str], ) -> None: - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=modeling_qwen3_mod.Qwen3ForCausalLM, architecture="Qwen3ForCausalLM", model_type="qwen3", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config( rope_fusion=False, @@ -3999,12 +4003,12 @@ def test_bf16_profiles_reject_wrong_realized_dimension( expected_dimension: str, ) -> None: """Each BF16 profile rejects a realized RoPE fusion or window it was not qualified for.""" - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=root_model_class, architecture=architecture, model_type=model_type, speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config(**realized_overrides), ) @@ -4094,12 +4098,12 @@ def test_bf16_dense_profiles_ignore_moe_only_runtime_dimensions( supported_sliding_window: str | None, ) -> None: """Dense profiles still qualify when only MoE-specific runtime dimensions differ.""" - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=root_model_class, architecture=architecture, model_type=model_type, speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config( rope_fusion=supported_rope_fusion, @@ -4166,12 +4170,12 @@ def test_bf16_dense_profiles_qualify_on_any_sm_version( supported_rope_fusion: bool, supported_sliding_window: str | None, ) -> None: - decision = ModelLoader._post_transform_profile_registry().qualify( + decision = mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=root_model_class, architecture=architecture, model_type=model_type, speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=_bf16_dense_runtime_config( rope_fusion=supported_rope_fusion, @@ -4201,8 +4205,8 @@ def test_qualify_post_transform_profile_uses_device_sm_version( ) -> None: (tiny_profile,) = _tiny_profile_registry().profiles monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", PostTransformProfileRegistry( profiles=( replace( @@ -4214,9 +4218,9 @@ def test_qualify_post_transform_profile_uses_device_sm_version( ) ), ) - monkeypatch.setattr(model_loader_mod, "get_sm_version", lambda: device_sm_version) + monkeypatch.setattr(mx_loading_mod, "get_sm_version", lambda: device_sm_version) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( _TinyModel([]), speculative_mode=None, loads_draft_weights=False, @@ -4227,12 +4231,12 @@ def test_qualify_post_transform_profile_uses_device_sm_version( def _qualify_qwen3_runtime(runtime_config: PostTransformRuntimeConfig): - return ModelLoader._post_transform_profile_registry().qualify( + return mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify( root_model_class=modeling_qwen3_mod.Qwen3ForCausalLM, architecture="Qwen3ForCausalLM", model_type="qwen3", speculative_mode=None, - protocol_version=ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + protocol_version=mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, transfer_scope=PostTransformTransferScope.TARGET_MODEL, runtime_config=runtime_config, ) @@ -4242,7 +4246,7 @@ def _qualify_qwen3_runtime(runtime_config: PostTransformRuntimeConfig): def test_qwen3_fp8_block_profile_is_the_sm90_variant_of_the_bf16_row() -> None: profiles = { profile.profile_id: profile - for profile in ModelLoader._post_transform_profile_registry().profiles + for profile in mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles } fp8_profile = profiles[_QWEN3_FP8_BLOCK_PROFILE_ID] bf16_profile = profiles["qwen3-for-causal-lm-bf16-target-v1"] @@ -4252,7 +4256,7 @@ def test_qwen3_fp8_block_profile_is_the_sm90_variant_of_the_bf16_row() -> None: == QWEN3_DENSE_FP8_BLOCK128_FP32SCALE_POST_TRANSFORM_LAYOUT_ABI_V1 ) assert fp8_profile.transform_abi_id != bf16_profile.transform_abi_id - assert fp8_profile.runtime_constraints == model_loader_mod._with_fp8_block_scale_weights( + assert fp8_profile.runtime_constraints == mx_loading_mod._with_fp8_block_scale_weights( bf16_profile.runtime_constraints, {90} ) assert fp8_profile.runtime_constraints == replace( @@ -4268,7 +4272,7 @@ def test_qwen3_fp8_block_profile_is_the_sm90_variant_of_the_bf16_row() -> None: @pytest.mark.cpu_only def test_qwen3_fp8_block_row_precedes_the_bf16_qwen3_row() -> None: profile_ids = [ - profile.profile_id for profile in ModelLoader._post_transform_profile_registry().profiles + profile.profile_id for profile in mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.profiles ] assert profile_ids.index(_QWEN3_FP8_BLOCK_PROFILE_ID) + 1 == profile_ids.index( @@ -4441,8 +4445,8 @@ def test_mx_donor_publishes_only_on_a_qualified_sm_version( loader = _make_loader(monkeypatch, events=[]) (tiny_profile,) = _tiny_profile_registry().profiles monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", PostTransformProfileRegistry( profiles=( replace( @@ -4454,9 +4458,8 @@ def test_mx_donor_publishes_only_on_a_qualified_sm_version( ) ), ) - monkeypatch.setattr(model_loader_mod, "get_sm_version", lambda: device_sm_version) - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + monkeypatch.setattr(mx_loading_mod, "get_sm_version", lambda: device_sm_version) + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.is_weights_preloaded.return_value = False checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} checkpoint_loader.get_initialized_weight_mapper.return_value = MagicMock() @@ -4474,13 +4477,13 @@ def test_staged_llama_finalization_preserves_mx_tensor_catalog( model = _tiny_llama_model(monkeypatch) # MX discovers and registers tensors after receiver alias preparation. - ModelLoader._setup_aliases(model) + model_hooks_mod.setup_aliases(model) registered_catalog = _mx_canonical_parameter_catalog(model) # This is the receiver-side finalization sequence after RDMA completes. - ModelLoader._setup_aliases(model) - ModelLoader._mark_weights_transformed(model) - ModelLoader._walk_cache_state(model) + model_hooks_mod.setup_aliases(model) + model_hooks_mod.mark_weights_transformed(model) + model_hooks_mod.walk_cache_state(model) published_catalog = _mx_canonical_parameter_catalog(model) assert registered_catalog @@ -4492,12 +4495,12 @@ def test_separate_draft_model_is_not_qualified_by_target_only_profile( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(speculative_mode="eagle3_one_model"), ) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( _TinyModel([]), speculative_mode="eagle3_one_model", loads_draft_weights=True, @@ -4544,8 +4547,8 @@ def test_qualification_reports_realized_moe_lifecycle_features( the loader reads fails it. """ monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) model = _TinyModel([]) @@ -4559,7 +4562,7 @@ def test_qualification_reports_realized_moe_lifecycle_features( world_size=dwdp_size, tp_size=dwdp_size, dwdp_size=dwdp_size ) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, speculative_mode=None, loads_draft_weights=False, @@ -4613,15 +4616,16 @@ def test_skipped_mx_publish_logs_only_reported_details( empty feature set there does not mean that no unsupported feature is active. """ messages: list[str] = [] - monkeypatch.setattr(model_loader_mod.logger, "info", messages.append) + monkeypatch.setattr(mx_loading_mod.logger, "info", messages.append) checkpoint_loader = SimpleNamespace(checkpoint_format="MX", post_load_publish=MagicMock()) - ModelLoader.__new__(ModelLoader)._post_load_publish( + mx_loading_mod.post_load_publish( checkpoint_loader, _TinyModel([]), checkpoint_dir="/checkpoint", weights_preloaded=False, qualification=decision, + source_identity=None, ) assert messages == [ @@ -4632,42 +4636,43 @@ def test_skipped_mx_publish_logs_only_reported_details( @pytest.mark.cpu_only -def test_separate_draft_load_preserves_mx_session( +def test_gms_writer_draft_load_preserves_mx_session( monkeypatch: pytest.MonkeyPatch, ) -> None: mapping = MagicMock() - spec_config = SimpleNamespace(speculative_model="draft-checkpoint") - loader = ModelLoader( - llm_args=MagicMock(), - mapping=mapping, - spec_config=spec_config, - sparse_attention_config=None, - max_num_tokens=1, - max_seq_len=1, - ) checkpoint_loader = MagicMock() checkpoint_loader.checkpoint_format = "MX" - checkpoint_loader.load_weights.return_value = {"draft": object()} - draft_mapper = MagicMock() - monkeypatch.setattr( - model_loader_mod.AutoCheckpointMapper, - "get", - MagicMock(return_value=draft_mapper), - ) + draft_weights = {"draft": object()} + checkpoint_loader.load_weights.return_value = draft_weights + target_weight_mapper = MagicMock() + call_load_weights = MagicMock() + monkeypatch.setattr(model_hooks_mod, "call_load_weights", call_load_weights) monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - model = SimpleNamespace( - draft_config=_make_draft_model_config(), - draft_model=MagicMock(), - load_draft_weights=lambda _weights: None, - ) + model = SimpleNamespace(load_draft_weights=lambda _weights: None) - loader._load_separate_draft_weights(model, checkpoint_loader) + gms_loading_mod._load_gms_writer_draft_weights( + model, + checkpoint_loader, + speculative_model="draft-checkpoint", + mapping=mapping, + target_weight_mapper=target_weight_mapper, + metrics={}, + ) checkpoint_loader.load_weights.assert_called_once_with( "draft-checkpoint", mapping=mapping, _preserve_mx_session=True, ) + # The draft weights go through the draft mapper that the checkpoint loader builds. + checkpoint_loader.get_initialized_draft_weight_mapper.assert_called_once_with( + model, target_weight_mapper + ) + call_load_weights.assert_called_once_with( + model.load_draft_weights, + draft_weights, + checkpoint_loader.get_initialized_draft_weight_mapper.return_value, + ) @pytest.mark.cpu_only @@ -4675,12 +4680,12 @@ def test_one_engine_speculative_mode_is_not_qualified_by_target_only_profile( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( _TinyModel([]), speculative_mode="mtp", loads_draft_weights=False, @@ -4696,18 +4701,18 @@ def test_speculative_mode_name_is_canonical_and_fails_closed( monkeypatch: pytest.MonkeyPatch, ) -> None: warning = MagicMock() - monkeypatch.setattr(model_loader_mod.logger, "warning", warning) - assert ModelLoader._speculative_mode_name(None) is None + monkeypatch.setattr(mx_loading_mod.logger, "warning", warning) + assert mx_loading_mod.get_speculative_mode_name(None) is None warning.assert_not_called() assert ( - ModelLoader._speculative_mode_name( + mx_loading_mod.get_speculative_mode_name( SimpleNamespace(spec_dec_mode=SimpleNamespace(name="MTP")) ) == "mtp" ) warning.assert_not_called() assert ( - ModelLoader._speculative_mode_name(SimpleNamespace(spec_dec_mode=SimpleNamespace())) + mx_loading_mod.get_speculative_mode_name(SimpleNamespace(spec_dec_mode=SimpleNamespace())) == "unknown" ) warning.assert_called_once_with( @@ -4730,14 +4735,14 @@ def test_mx_post_transform_receiver_falls_back_for_unqualified_model( loader = _make_loader(monkeypatch, events=events) if qualification_case == "locality_domain-feature": monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) model = _TinyModel(events) model.model_config.locality_domain_policy = LocalityDomainPolicy(enabled=True) monkeypatch.setattr( - model_loader_mod.AutoModelForCausalLM, + memory_init_mod.AutoModelForCausalLM, "from_config", MagicMock(return_value=model), ) @@ -4749,12 +4754,12 @@ def test_mx_post_transform_receiver_falls_back_for_unqualified_model( assert kwargs["allow_post_transform_weights"] is False assert "prepare_post_transform_receiver" not in kwargs assert checkpoint_loader.is_weights_preloaded() is False - assert loader._call_load_weights.call_count == 1 - load_fn, weights, mapper = loader._call_load_weights.call_args.args + assert model_hooks_mod.call_load_weights.call_count == 1 + load_fn, weights, mapper = model_hooks_mod.call_load_weights.call_args.args assert load_fn == model.load_weights assert weights == {"disk.weight": checkpoint_loader._disk_weight} assert mapper is loader.weight_mapper - assert loader._source_identity.transform_abi_id is None + assert kwargs["source_identity"].transform_abi_id is None assert events == ["load_weights", "post_load_weights"] checkpoint_loader.post_load_publish.assert_not_called() @@ -4766,20 +4771,19 @@ def test_load_qualifies_with_preconstruction_identity_after_model_normalization( loader = _make_loader(monkeypatch, events=events) registry = MagicMock(wraps=_tiny_profile_registry()) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", registry, ) normalized_model = _TinyModel(events) normalized_model.model_config.pretrained_config.architectures = ["NormalizedForCausalLM"] normalized_model.model_config.pretrained_config.model_type = "normalized" monkeypatch.setattr( - model_loader_mod.AutoModelForCausalLM, + memory_init_mod.AutoModelForCausalLM, "from_config", MagicMock(return_value=normalized_model), ) - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} checkpoint_loader.is_weights_preloaded.return_value = False @@ -4788,7 +4792,7 @@ def test_load_qualifies_with_preconstruction_identity_after_model_normalization( _args, kwargs = checkpoint_loader.load_weights.call_args registry.qualify.assert_called_once() assert kwargs["allow_post_transform_weights"] is True - assert loader._source_identity.transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 + assert kwargs["source_identity"].transform_abi_id == LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1 @pytest.mark.cpu_only @@ -4816,26 +4820,25 @@ def test_mx_fallback_runs_standard_weight_mapping( events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.is_weights_preloaded.return_value = False checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} checkpoint_loader.get_initialized_weight_mapper.return_value = MagicMock() model, _ = loader.load("/ckpt", checkpoint_loader) - assert loader._call_load_weights.call_count == 1 + assert model_hooks_mod.call_load_weights.call_count == 1 assert events[0] == "load_weights" assert "post_load_weights" in events checkpoint_loader.post_load_publish.assert_called_once_with( model, checkpoint_dir="/ckpt", weights_preloaded=False, - source_identity=loader._source_identity, + source_identity=checkpoint_loader.load_weights.call_args.kwargs["source_identity"], ) @@ -4843,8 +4846,8 @@ def test_mx_artifact_identity_failure_falls_back_to_disk(monkeypatch): events = [] loader = _make_loader(monkeypatch, events=events) monkeypatch.setattr( - ModelLoader, - "_POST_TRANSFORM_PROFILE_REGISTRY", + mx_loading_mod, + "MX_POST_TRANSFORM_PROFILE_REGISTRY", _tiny_profile_registry(), ) artifact_error = ValueError( @@ -4864,15 +4867,13 @@ def test_mx_artifact_identity_failure_falls_back_to_disk(monkeypatch): warning = MagicMock() monkeypatch.setattr(model_loader_mod.logger, "warning", warning) - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.is_weights_preloaded.return_value = False checkpoint_loader.load_weights.return_value = {"weight": MagicMock()} checkpoint_loader.get_initialized_weight_mapper.return_value = MagicMock() model, _ = loader.load("/ckpt", checkpoint_loader) - assert loader._source_identity is None source_identity_factory.assert_not_called() warning.assert_called_once() assert "falling back to regular checkpoint loading" in warning.call_args.args[0] @@ -4906,10 +4907,6 @@ def __init__( def setup_aliases(self) -> None: self.events.append((self.name, "setup_aliases")) - def transform_weights(self) -> None: - self.events.append((self.name, "transform_weights")) - self._weights_transformed = True - def cache_derived_state(self) -> None: self.events.append((self.name, "cache_derived_state")) @@ -4930,7 +4927,7 @@ def test_staged_hook_setup_aliases_walks_skip_removed_modules(): events = [] model = _HookModel(events) - ModelLoader._setup_aliases(model) + model_hooks_mod.setup_aliases(model) assert events == [ ("model", "setup_aliases"), @@ -4940,17 +4937,14 @@ def test_staged_hook_setup_aliases_walks_skip_removed_modules(): @pytest.mark.cpu_only -def test_staged_hook_walks_skip_removed_and_transformed_modules(): +def test_cache_state_and_full_post_load_walks_skip_removed_modules(): events = [] model = _HookModel(events) - ModelLoader._walk_transform(model) - ModelLoader._walk_cache_state(model) - ModelLoader._walk_full_post_load(model) + model_hooks_mod.walk_cache_state(model) + model_hooks_mod.walk_full_post_load(model) assert events == [ - ("model", "transform_weights"), - ("child", "transform_weights"), ("model", "cache_derived_state"), ("child", "cache_derived_state"), ("transformed_child", "cache_derived_state"), @@ -4967,7 +4961,7 @@ def test_reset_weights_transformed_only_resets_existing_flags(): model._weights_transformed = True model.child._weights_transformed = True - ModelLoader._reset_weights_transformed(model) + model_hooks_mod.reset_weights_transformed(model) assert model._weights_transformed is False assert model.child._weights_transformed is False @@ -4982,7 +4976,7 @@ def test_mark_weights_transformed_only_sets_existing_flags(): model._weights_transformed = False model.child._weights_transformed = False - ModelLoader._mark_weights_transformed(model) + model_hooks_mod.mark_weights_transformed(model) assert model._weights_transformed is True assert model.child._weights_transformed is True @@ -5067,7 +5061,7 @@ def _make_manifest_loader( loader = _make_loader(monkeypatch, events=events) loader.mapping = mapping or mapping_mod.Mapping(world_size=1, rank=0, tp_size=1) monkeypatch.setattr( - model_loader_mod.AutoModelForCausalLM, + memory_init_mod.AutoModelForCausalLM, "from_config", MagicMock(return_value=_ManifestTinyModel(events)), ) @@ -5077,13 +5071,13 @@ def _make_manifest_loader( else: monkeypatch.setenv(WEIGHT_MANIFEST_DIR_ENV, str(tmp_path)) monkeypatch.setenv(WEIGHT_MANIFEST_ROLE_ENV, role) - real_write = model_loader_mod.maybe_write_weight_manifest + real_write = mx_loading_mod.maybe_write_weight_manifest def _recording_write(*args, **kwargs): events.append("manifest") return real_write(*args, **kwargs) - monkeypatch.setattr(model_loader_mod, "maybe_write_weight_manifest", _recording_write) + monkeypatch.setattr(mx_loading_mod, "maybe_write_weight_manifest", _recording_write) return loader @@ -5101,9 +5095,13 @@ def test_load_writes_final_manifest_at_end_of_load( ) -> None: events = [] loader = _make_manifest_loader(monkeypatch, tmp_path, events=events, role="baseline") + from_checkpoint = MagicMock() + monkeypatch.setattr(model_loader_mod.ArtifactIdentity, "from_checkpoint", from_checkpoint) model, _ = loader.load("/ckpt", _hf_checkpoint_loader()) + # HF/AUTO loading needs no source identity, so it never reads or hashes the checkpoint. + from_checkpoint.assert_not_called() assert events == ["load_weights", "post_load_weights", "manifest"] files = sorted(path.name for path in tmp_path.iterdir()) assert files == ["manifest.final.baseline.rank0.json"] @@ -5124,8 +5122,7 @@ def test_load_final_manifest_records_mx_receiver_facts( ) -> None: events = [] loader = _make_manifest_loader(monkeypatch, tmp_path, events=events, role="receiver") - checkpoint_loader = MagicMock(name="checkpoint_loader") - checkpoint_loader.checkpoint_format = "MX" + checkpoint_loader = _mock_mx_checkpoint_loader() checkpoint_loader.is_weights_preloaded.return_value = True # `_ManifestTinyModel` is not a registered root, so the post-transform # staged path must stay off; the full post-load path still ends in a manifest. diff --git a/tests/unittest/_torch/pyexecutor/test_model_loader_mx_qwen3_5_vl.py b/tests/unittest/_torch/pyexecutor/test_model_loader_mx_qwen3_5_vl.py index 128073a935b0..aa0ec60e730a 100644 --- a/tests/unittest/_torch/pyexecutor/test_model_loader_mx_qwen3_5_vl.py +++ b/tests/unittest/_torch/pyexecutor/test_model_loader_mx_qwen3_5_vl.py @@ -38,11 +38,12 @@ from tensorrt_llm._torch.models import modeling_qwen3_5 as modeling_qwen3_5_mod from tensorrt_llm._torch.models import modeling_qwen3_next as modeling_qwen3_next_mod from tensorrt_llm._torch.models import modeling_qwen3vl as modeling_qwen3vl_mod +from tensorrt_llm._torch.models.checkpoints import model_hooks as model_hooks_mod from tensorrt_llm._torch.models.checkpoints.hf.qwen3_5_weight_mapper import Qwen3_5MoeHfWeightMapper +from tensorrt_llm._torch.models.checkpoints.mx import model_loading as mx_loading_mod from tensorrt_llm._torch.models.modeling_utils import MetaInitMode from tensorrt_llm._torch.modules.linear import Linear from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config -from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( ARTIFACT_IDENTITY_FORMAT_VERSION, QWEN3_5_DENSE_VL_POST_TRANSFORM_LAYOUT_ABI_V1, @@ -201,7 +202,7 @@ def _load(model: nn.Module) -> None: def _qualify(model: nn.Module): """Qualify `model` with the identity that `ModelLoader` captures before construction.""" - return ModelLoader._qualify_post_transform_profile( + return mx_loading_mod.qualify_post_transform_profile( model, config_identity=_CONFIG_IDENTITY, speculative_mode=None, @@ -381,7 +382,7 @@ def test_qwen3_5_vl_profile_qualifies_full_staged_lifecycle(tmp_path: Path) -> N decision = _qualify(receiver) assert decision.transform_abi_id == QWEN3_5_DENSE_VL_POST_TRANSFORM_LAYOUT_ABI_V1 assert ( - ModelLoader._post_transform_transfer_scope(receiver) + mx_loading_mod.get_post_transform_transfer_scope(receiver) is PostTransformTransferScope.COMPLETE_MODEL ) assert _runtime_config(producer) == _qwen3_5_vl_runtime_config() @@ -420,14 +421,14 @@ def test_qwen3_5_vl_mx_tensor_catalog_is_complete_and_stable(tmp_path: Path) -> checkpoint_dir = _write_tiny_checkpoint_config(tmp_path / "ckpt") producer = _tiny_qwen3_5_vl_model(checkpoint_dir) _hf_checkpoint_loader(_tiny_hf_state_dict())(producer) - ModelLoader._walk_full_post_load(producer) + model_hooks_mod.walk_full_post_load(producer) receiver = _tiny_qwen3_5_vl_model(checkpoint_dir) - ModelLoader._setup_aliases(receiver) + model_hooks_mod.setup_aliases(receiver) registered_catalog = mx_canonical_parameter_catalog(receiver) - ModelLoader._setup_aliases(receiver) - ModelLoader._mark_weights_transformed(receiver) - ModelLoader._walk_cache_state(receiver) + model_hooks_mod.setup_aliases(receiver) + model_hooks_mod.mark_weights_transformed(receiver) + model_hooks_mod.walk_cache_state(receiver) published_catalog = mx_canonical_parameter_catalog(receiver) assert published_catalog == registered_catalog @@ -470,7 +471,7 @@ def test_qwen3_5_vl_qualification_requires_preconstruction_identity(tmp_path: Pa model = _tiny_qwen3_5_vl_model(_write_tiny_checkpoint_config(tmp_path / "ckpt")) post_construction_identity = PostTransformConfigIdentity.from_model_config(model.model_config) - decision = ModelLoader._qualify_post_transform_profile( + decision = mx_loading_mod.qualify_post_transform_profile( model, config_identity=post_construction_identity, speculative_mode=None, @@ -560,7 +561,7 @@ def test_qwen3_5_vl_without_local_encoder_falls_back( assert model.mm_encoder is None assert ( - ModelLoader._post_transform_transfer_scope(model) + mx_loading_mod.get_post_transform_transfer_scope(model) is PostTransformTransferScope.LANGUAGE_MODEL ) assert _qualify(model).reason is PostTransformQualificationReason.TRANSFER_SCOPE_NOT_REGISTERED @@ -598,12 +599,12 @@ def _qualify_registry(**overrides: object): "architecture": "Qwen3_5ForConditionalGeneration", "model_type": "qwen3_5", "speculative_mode": None, - "protocol_version": ModelLoader._MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + "protocol_version": mx_loading_mod.MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, "transfer_scope": PostTransformTransferScope.COMPLETE_MODEL, "runtime_config": _qwen3_5_vl_runtime_config(), } arguments.update(overrides) - return ModelLoader._post_transform_profile_registry().qualify(**arguments) + return mx_loading_mod.MX_POST_TRANSFORM_PROFILE_REGISTRY.qualify(**arguments) @pytest.mark.cpu_only diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py index a8e1db22c708..934451a818bf 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_mtp_separate_checkpoint.py @@ -511,6 +511,47 @@ def fake_load_weights( return captured +def _load_separate_draft_checkpoint( + model, weights: dict, *, load_format: str, target_config: ModelConfig +) -> object: + """Load ``weights`` as the separate draft checkpoint of ``model``. + + ``load_format`` "auto" reads like `load_model_from_auto_load_format`; + "gms" reads like the GMS writer. Returns the target weight mapper. + """ + from tensorrt_llm._torch.models.checkpoints.gms import model_loading as gms_loading + from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import HfCheckpointLoader + from tensorrt_llm._torch.pyexecutor import model_loader + + checkpoint_loader = HfCheckpointLoader() + checkpoint_loader.load_weights = lambda *args, **kwargs: weights + target_weight_mapper = object() + speculative_model = model.spec_config.speculative_model + if load_format == "auto": + model_loader._load_checkpoint_into_model( + checkpoint_loader, + speculative_model, + model=model, + is_draft_model=True, + mapping=target_config.mapping, + weight_mapper=checkpoint_loader.get_initialized_draft_weight_mapper( + model, target_weight_mapper + ), + extra_load_weights_kwargs=None, + metrics={}, + ) + else: + gms_loading._load_gms_writer_draft_weights( + model, + checkpoint_loader, + speculative_model=speculative_model, + mapping=target_config.mapping, + target_weight_mapper=target_weight_mapper, + metrics={}, + ) + return target_weight_mapper + + def test_mtp_head_module_names_covers_both_aliases(): model = _make_one_engine_stub( MTPDecodingConfig(max_draft_len=3, speculative_model="/path/to/mtp") @@ -552,12 +593,12 @@ def test_separate_mtp_target_load_skips_heads_without_partial_loading(monkeypatc @pytest.mark.cpu_only -@pytest.mark.parametrize("materialize", [False, True]) +@pytest.mark.parametrize("load_format", ["gms", "auto"]) @pytest.mark.parametrize("pp_size,owns_layer", [(1, True), (2, True), (2, False)]) def test_qwen4_replacement_loads_mtp_and_hyper_connection( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, - materialize: bool, + load_format: str, pp_size: int, owns_layer: bool, ) -> None: @@ -565,7 +606,6 @@ def test_qwen4_replacement_loads_mtp_and_hyper_connection( from tensorrt_llm._torch.models.checkpoints.hf.qwen4_exp_weight_mapper import ( Qwen4ExpHfWeightMapper, ) - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader captured = {} monkeypatch.setattr(torch.cuda, "is_available", lambda: False) @@ -617,20 +657,10 @@ def capture_load(model, weights, mapper, **kwargs) -> None: "mtp.layers.0.attn_hyper_connection.block_inject_weight.weight": inject, "mtp.hyper_connection_mixer.input_mix_weight_down.weight": final_down, } - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = {} - loader.spec_config = model.spec_config - loader.mapping = config.mapping - loader.weight_mapper = object() - checkpoint_loader = SimpleNamespace( - checkpoint_format="HF", load_weights=lambda *args, **kwargs: weights - ) - - if materialize: - loader._materialize_draft_checkpoint_weights(checkpoint_loader, model) - else: - loader._load_separate_draft_weights(model, checkpoint_loader) + + _load_separate_draft_checkpoint( + model, weights, load_format=load_format, target_config=target_config + ) assert isinstance(captured["mapper"], Qwen4ExpHfWeightMapper) assert captured["mapper"].config is config @@ -938,10 +968,9 @@ def mtp_layer(config, layer_idx, aux): @pytest.mark.cpu_only -@pytest.mark.parametrize("materialize", [False, True]) -def test_replacement_loader_initializes_independent_architecture_mapper(monkeypatch, materialize): +@pytest.mark.parametrize("load_format", ["gms", "auto"]) +def test_replacement_loader_initializes_independent_architecture_mapper(monkeypatch, load_format): from tensorrt_llm._torch.models import modeling_utils - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader captured = {} monkeypatch.setattr(torch.cuda, "is_available", lambda: False) @@ -967,26 +996,20 @@ def capture_load(model, weights, mapper, **kwargs): model.draft_model.model_config = config model.draft_model.config = config.pretrained_config model.draft_model.owns_lm_head = False - loader = ModelLoader.__new__(ModelLoader) - loader._metrics = {} - loader._startup_metadata = {} - loader.spec_config = model.spec_config - loader.mapping = config.mapping - loader.weight_mapper = object() - checkpoint_loader = SimpleNamespace( - checkpoint_format="HF", - load_weights=lambda *args, **kwargs: _nemotron_style_mtp_weights(include_shared_head=False), - ) - - if materialize: - loader._materialize_draft_checkpoint_weights(checkpoint_loader, model) - else: - loader._load_separate_draft_weights(model, checkpoint_loader) + + # A replacement head uses the target architecture, so this config also + # serves as the target config. + target_weight_mapper = _load_separate_draft_checkpoint( + model, + _nemotron_style_mtp_weights(include_shared_head=False), + load_format=load_format, + target_config=config, + ) assert isinstance(captured["mapper"], NemotronHHfWeightMapper) assert captured["mapper"].config is config assert captured["mapper"].model is model.draft_model - assert captured["mapper"] is not loader.weight_mapper + assert captured["mapper"] is not target_weight_mapper assert captured["model"] is model.draft_model assert "mtp_layers.0.layers.0.enorm.weight" in captured["weights"] assert captured["allow_partial_loading"] is False diff --git a/tests/unittest/_torch/speculative/test_vocab_parallel.py b/tests/unittest/_torch/speculative/test_vocab_parallel.py index b59250725670..235ddf567ce0 100644 --- a/tests/unittest/_torch/speculative/test_vocab_parallel.py +++ b/tests/unittest/_torch/speculative/test_vocab_parallel.py @@ -285,7 +285,7 @@ def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits=F assert worker.target_vocab_parallel_head is None -@pytest.mark.parametrize("walk", ["_setup_aliases", "_walk_full_post_load"]) +@pytest.mark.parametrize("walk", ["setup_aliases", "walk_full_post_load"]) def test_loader_post_load_walk_binds_the_draft_side(monkeypatch, walk): """The worker is a module of the engine, registered after the drafter, so the model loader's post-load walks bind its target side and its draft side @@ -294,7 +294,7 @@ def test_loader_post_load_walk_binds_the_draft_side(monkeypatch, walk): from torch import nn - from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader + from tensorrt_llm._torch.models.checkpoints import model_hooks from tensorrt_llm._torch.speculative.dflash import DFlashWorker lm_head = nn.Linear(4, 4) @@ -310,7 +310,7 @@ def test_loader_post_load_walk_binds_the_draft_side(monkeypatch, walk): engine.spec_worker = worker # The drafter takes over the target's head while the weights load. engine.draft_model.lm_head = lm_head - getattr(ModelLoader, walk)(engine) + getattr(model_hooks, walk)(engine) assert worker.target_vocab_parallel_head.lm_head is lm_head assert worker._vocab_parallel.draft_tokens == worker._argmax_block_tokens diff --git a/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py b/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py index e67260daa657..7d9f45351090 100644 --- a/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py +++ b/tests/unittest/_torch/weight_sharing/test_gms_source_identity_gate.py @@ -14,31 +14,22 @@ # limitations under the License. """Exercise the real GMS read-only pre-materialize SourceIdentity gate. -These tests drive `ModelLoader._check_gms_source_identity` directly -- the -single decision point the GMS RO path consults before `materialize_module`. -Unlike MX, GMS has no disk fallback, so a verified mismatch must *raise* under -`STRICT` rather than fall back. The `ModelLoader` is constructed via -`__new__` to bypass its heavy `__init__`; the GMS backend is a stub -exposing only `get_source_identity` so no GMS pool, GPU, or model is touched. +These tests drive `check_gms_source_identity` directly -- the single decision +point the GMS RO path consults before `materialize_module`. Unlike MX, GMS has +no disk fallback, so a verified mismatch must *raise* under `STRICT` rather +than fall back. The GMS backend is a stub exposing only `get_source_identity` +so no GMS pool, GPU, or model is touched. """ import pytest from _source_identity_fakes import make_identity as _identity -from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader +from tensorrt_llm._torch.models.checkpoints.gms.model_loading import check_gms_source_identity from tensorrt_llm._torch.weight_sharing import SourceIdentityMismatchError -from tensorrt_llm.llmapi.llm_args import LoadFormat pytestmark = pytest.mark.cpu_only -class _FakeCheckpointLoader: - """Minimal checkpoint-loader stub exposing `checkpoint_format`.""" - - def __init__(self, checkpoint_format): - self.checkpoint_format = checkpoint_format - - class _FakeGMSBackend: """Minimal GMS backend stub exposing only the gate's fetch seam.""" @@ -49,50 +40,31 @@ def get_source_identity(self): return self._source_identity -def _new_loader(local_identity): - """Construct a loader bypassing heavy base __init__, wire the local id.""" - loader = ModelLoader.__new__(ModelLoader) - loader._source_identity = local_identity - return loader - - -def test_source_identity_is_only_needed_for_mx_or_gms(): - # Default HF/AUTO loading should be a strict no-op: no SourceIdentity - # construction, no tensor-layout traversal, no behavior change. - assert not ModelLoader._needs_source_identity(_FakeCheckpointLoader("HF"), LoadFormat.AUTO) - assert ModelLoader._needs_source_identity(_FakeCheckpointLoader("MX"), LoadFormat.AUTO) - assert ModelLoader._needs_source_identity(_FakeCheckpointLoader("HF"), LoadFormat.GMS) - - def test_gate_passes_on_matching_identity(): local = _identity(attn_backend="TRTLLM") writer = _identity(attn_backend="TRTLLM") - loader = _new_loader(local) # Compatible writer: the gate is a no-op (no raise). - loader._check_gms_source_identity(_FakeGMSBackend(writer)) + check_gms_source_identity(local, _FakeGMSBackend(writer)) def test_gate_raises_on_mismatch(): local = _identity(attn_backend="TRTLLM") writer = _identity(attn_backend="FLASHINFER") - loader = _new_loader(local) # No disk fallback for GMS: an incompatible writer must raise. with pytest.raises(SourceIdentityMismatchError): - loader._check_gms_source_identity(_FakeGMSBackend(writer)) + check_gms_source_identity(local, _FakeGMSBackend(writer)) def test_gate_raises_on_checkpoint_artifact_mismatch(): local = _identity(artifact_key="fine-tune-a") writer = _identity(artifact_key="fine-tune-b") - loader = _new_loader(local) with pytest.raises(SourceIdentityMismatchError): - loader._check_gms_source_identity(_FakeGMSBackend(writer)) + check_gms_source_identity(local, _FakeGMSBackend(writer)) def test_gate_raises_when_writer_identity_unavailable(): # Publisher metadata not wired yet (get_source_identity returns None); # GMS has no disk fallback, so unverified sharing must raise. local = _identity() - loader = _new_loader(local) with pytest.raises(SourceIdentityMismatchError): - loader._check_gms_source_identity(_FakeGMSBackend(None)) + check_gms_source_identity(local, _FakeGMSBackend(None)) diff --git a/tests/unittest/utils/post_transform_qualification.py b/tests/unittest/utils/post_transform_qualification.py index 2b5dcbde7b7f..8c652e4c478e 100644 --- a/tests/unittest/utils/post_transform_qualification.py +++ b/tests/unittest/utils/post_transform_qualification.py @@ -24,10 +24,10 @@ from torch import nn from transformers import PretrainedConfig +from tensorrt_llm._torch.models.checkpoints import model_hooks from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM from tensorrt_llm._torch.moe.fused_moe.fused_moe_vanilla import VanillaMoE from tensorrt_llm._torch.moe.fused_moe.weight_owner import is_moe_weight_owner -from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader from tensorrt_llm._torch.weight_sharing import ( PostTransformQualificationDecision, PostTransformQualificationReason, @@ -292,7 +292,7 @@ def _holds_routed_experts(model: nn.Module) -> bool: def _runs_real_transform(module: nn.Module) -> bool: - """Whether `ModelLoader._walk_transform` would run a transform on `module`. + """Whether the full post-load walk would run a real transform on `module`. Reads the hook from the class, because the call recorders replace it on the instance; the no-op default of `DecoderModelForCausalLM` transforms nothing. @@ -374,15 +374,15 @@ def assert_post_transform_lifecycle_equivalent( transform_calls = _install_transform_call_recorders(receiver) if case.prepare_producer is not None: case.prepare_producer(producer) - ModelLoader._walk_full_post_load(producer) - ModelLoader._setup_aliases(receiver) + model_hooks.walk_full_post_load(producer) + model_hooks.setup_aliases(receiver) case.prepare_receiver(producer, receiver) # Runtime finalization repeats setup_aliases(), whose contract is # idempotent, after the checkpoint loader's pre-transfer preparation. - ModelLoader._setup_aliases(receiver) - ModelLoader._mark_weights_transformed(receiver) - ModelLoader._walk_cache_state(receiver) + model_hooks.setup_aliases(receiver) + model_hooks.mark_weights_transformed(receiver) + model_hooks.walk_cache_state(receiver) assert transform_calls == [], f"Staged receiver invoked transform_weights(): {transform_calls}" # Staged finalization marks only modules that already carry the guard, and