diff --git a/docs/source/features/checkpoint-loading.md b/docs/source/features/checkpoint-loading.md index c6ff47447c26..6c6701151bcc 100644 --- a/docs/source/features/checkpoint-loading.md +++ b/docs/source/features/checkpoint-loading.md @@ -110,6 +110,12 @@ checkpoint bytes reach host memory: 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. +- `demand_ordered_rank_striped_read_ahead` is an explicit experimental variant + that uses mapper demand hints to prioritize complete read-ahead chunks across + node-local ranks. It retains every chunk, existing full-checkpoint memory + admission, and native materialization. It does not add consumption-based + pacing, destination-selective reads, or a streaming materializer. `auto` + continues to select the original rank-striped policy. ```yaml checkpoint_format: HF @@ -117,6 +123,14 @@ load_format: auto checkpoint_io_policy: rank_striped_read_ahead ``` +Policy IDs and display metadata are centralized in +`tensorrt_llm/_checkpoint_io_policy.py`. `auto` is a selector, not a separate +loading implementation. Display names may evolve without changing the stable +configuration IDs; future modes become selectable only when implemented. +The generated LLM-args telemetry manifest records the accepted IDs for tools +that must not import TRT-LLM. Eligibility, fallback, and CI allocation remain +separate from the catalog. + 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 @@ -157,6 +171,50 @@ 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. +### Demand-ordering qualification + +Select `checkpoint_io_policy: demand_ordered_rank_striped_read_ahead` in the +normal HF configuration. The same eligibility restrictions apply to both +rank-striped policies. Demand planning runs after mapper initialization and +before the new readers are released. Opaque plans use physical order; qualified +hints prioritize earlier demands while preserving physical locality among +equal priorities and retaining an unmatched physical-order tail. Advisory +ordering is not proof that source data can be omitted. + +The generic mapper prioritizes explicit native preload groups; concurrently +loaded modules share a priority. It does not infer an exact tensor-consumption +timeline or impose a serial layer order. Unsupported ordering stays opaque. + +For a planning-cost ablation, set +`TRTLLM_DEMAND_ORDERED_READ_AHEAD_PHYSICAL_ORDER=1` on every rank. This still +builds and validates the demand plan but applies physical read order. Leave it +unset for demand ordering. The switch affects only the experimental policy. + +Compare four cases using the same TRT-LLM revision, runtime image, checkpoint, +GPU/node, parallel configuration, and controlled source-cache conditions: + +| Case | Policy | Physical-order ablation | +| --- | --- | --- | +| Native control | `native` | Unset | +| Original rank-striped | `rank_striped_read_ahead` | Unset | +| Planning overhead | `demand_ordered_rank_striped_read_ahead` | `1` | +| Demand ordering | `demand_ordered_rank_striped_read_ahead` | Unset | + +Repeat cold and warm trials with counterbalanced order. Start with Qwen, +Llama Maverick, and a preprocessing-heavy model; include TP and PP/EP cases. +Check generation/weight parity, memory peaks, effective policy and applied +order, and net model-loading latency. Do not infer a speedup from read bandwidth +or weight population alone. The experimental policy includes its additional +planning/activation in preparation; existing native/v1 phase boundaries are +unchanged. Compare total model-loading time as well as the checkpoint phases, +so unclassified work such as mapper initialization is not hidden. These metrics +do not measure process launch to first token. + +Explicit policy configurations bypass the regular CI `auto`/`native` +assignment. Keep demand-ordered observations separate from the original +rank-striped treatment and retain the physical-order ablation setting with +benchmark results. + ### CI startup experiment Pre-merge perf-sanity runs that upload telemetry explicitly assign eligible diff --git a/tensorrt_llm/_checkpoint_io_policy.py b/tensorrt_llm/_checkpoint_io_policy.py new file mode 100644 index 000000000000..30e6a7995953 --- /dev/null +++ b/tensorrt_llm/_checkpoint_io_policy.py @@ -0,0 +1,51 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint I/O names and display metadata, independent of loader execution.""" + +from typing import Final, Literal, NamedTuple, TypeAlias, get_args + +CheckpointIoPolicy: TypeAlias = Literal[ + "auto", "native", "rank_striped_read_ahead", "demand_ordered_rank_striped_read_ahead" +] + +AUTO_IO_POLICY: Final = "auto" +NATIVE_IO_POLICY: Final = "native" +RANK_STRIPED_IO_POLICY: Final = "rank_striped_read_ahead" +DEMAND_ORDERED_IO_POLICY: Final = "demand_ordered_rank_striped_read_ahead" + +CHECKPOINT_IO_POLICIES: tuple[CheckpointIoPolicy, ...] = get_args(CheckpointIoPolicy) +RANK_STRIPED_IO_POLICIES = (RANK_STRIPED_IO_POLICY, DEMAND_ORDERED_IO_POLICY) +EXECUTABLE_CHECKPOINT_IO_POLICIES = (NATIVE_IO_POLICY,) + RANK_STRIPED_IO_POLICIES + + +class CheckpointIoPolicyInfo(NamedTuple): + """Human-readable metadata; eligibility and dispatch remain in the loader.""" + + display_name: str + description: str + kind: Literal["selector", "implementation"] + + +CHECKPOINT_IO_POLICY_INFO: dict[CheckpointIoPolicy, CheckpointIoPolicyInfo] = { + AUTO_IO_POLICY: CheckpointIoPolicyInfo( + "Automatic selection", + "Select basic rank-striped read-ahead where eligible, otherwise native I/O.", + "selector", + ), + NATIVE_IO_POLICY: CheckpointIoPolicyInfo( + "Native checkpoint loading", + "Preserve the existing checkpoint I/O and materialization path.", + "implementation", + ), + RANK_STRIPED_IO_POLICY: CheckpointIoPolicyInfo( + "Rank-cooperative read-ahead", + "Read complete checkpoint chunks across node-local ranks while native materialization runs.", + "implementation", + ), + DEMAND_ORDERED_IO_POLICY: CheckpointIoPolicyInfo( + "Demand-ordered rank-cooperative read-ahead", + "Prioritize complete checkpoint chunks using native materialization demand hints.", + "implementation", + ), +} diff --git a/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py index 6b4817a41f22..bd6db59e5bc6 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py @@ -82,6 +82,16 @@ def is_weights_preloaded(self) -> bool: """Whether the last load wrote weights directly into the model.""" return False + def activate_weight_session( + self, + readiness_error: BaseException | None = None, + *, + weight_mapper: BaseWeightMapper | None = None) -> None: + """Activate deferred advisory work before model materialization.""" + del weight_mapper + if readiness_error is not None: + raise readiness_error + def is_post_transform_weights_preloaded(self) -> bool: """Whether the last direct preload delivered post-transform weights. diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py index 31fc0d0470f7..a533fb3a0a7f 100644 --- a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py @@ -1,7 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os from abc import ABC, abstractmethod +from bisect import bisect_left from collections.abc import Callable, Mapping import torch @@ -13,7 +15,8 @@ from tensorrt_llm._torch.models.checkpoints.weight_load_plan import ( WeightDemand, WeightLoadOrderConfidence, WeightLoadPlan, WeightLoadPlanCoverage) -from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM +from tensorrt_llm._torch.models.modeling_utils import (DecoderModelForCausalLM, + is_moe_weight_owner) class BaseWeightMapper(ABC): @@ -86,18 +89,202 @@ def cleanup(self) -> None: def build_weight_load_plan(self, catalog: CheckpointCatalog) -> WeightLoadPlan: - """Build a conservative shadow plan for this rank. - - The generic mapper cannot infer model-specific fusion, transformation, - or ordering dependencies. Treat the entire catalog as one indivisible - demand so consumers may use the plan for all-source read-ahead, but not - destination-aware selective I/O. Specialized mappers may override this - method to return an exact plan. + """Build a conservative, source-neutral load plan for this rank. + + This method mirrors the native module walk without invoking any weight + transformation or loading callback. It uses module structure, explicit + fusion mappings, and preload hints to identify likely consumer groups. + The inference is advisory: every source tensor is retained exactly once, + unmatched tensors remain in physical catalog order at the tail, and the + plan never permits selective I/O. Specialized mappers may override this + method when they can provide stronger guarantees. """ if self._model is None: raise ValueError( "weight mapper must be initialized before building a load plan") mapping = self._model.model_config.mapping + + # Catalog order is the source adapter's physical order. Preserve it + # within each inferred group and in the unmatched tail. + physical_names = tuple(tensor.name for tensor in catalog.tensors) + if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", + "False") in ("True", "true", "1", "yes", "y"): + # The serial override bypasses preload grouping in the native + # materializer; do not apply the concurrent-path ordering hints. + return self._build_opaque_weight_load_plan(catalog, physical_names) + physical_index = { + name: index + for index, name in enumerate(physical_names) + } + sorted_names = tuple(sorted(physical_names)) + + named_modules = getattr(self._model, "named_modules", None) + if named_modules is None: + return self._build_opaque_weight_load_plan(catalog, physical_names) + + modules = tuple(named_modules(remove_duplicate=False)) + remaining_names = set(physical_names) + demands: list[WeightDemand] = [] + + def take_prefixes(prefixes: tuple[str, ...]) -> tuple[str, ...]: + """Take unassigned tensor names below structural module prefixes.""" + matched_indexes: set[int] = set() + for prefix in prefixes: + if not prefix: + continue + if prefix in remaining_names: + matched_indexes.add(physical_index[prefix]) + + descendant_prefix = f"{prefix}." + index = bisect_left(sorted_names, descendant_prefix) + while (index < len(sorted_names) + and sorted_names[index].startswith(descendant_prefix)): + name = sorted_names[index] + if name in remaining_names: + matched_indexes.add(physical_index[name]) + index += 1 + + result = tuple(physical_names[index] + for index in sorted(matched_indexes)) + remaining_names.difference_update(result) + return result + + def take_exact(names: tuple[str, ...]) -> tuple[str, ...]: + matched_indexes = { + physical_index[name] + for name in names + if name in remaining_names and name in physical_index + } + result = tuple(physical_names[index] + for index in sorted(matched_indexes)) + remaining_names.difference_update(result) + return result + + def source_module_path(module_path: str, module: nn.Module) -> str: + """Mirror native path normalization without touching payloads.""" + parts = module_path.split('.') if module_path else [] + if (parts and parts[-1] == "backend" + and is_moe_weight_owner(module)): + return '.'.join(parts[:-1]) + return module_path + + def infer_source_names(module_path: str, + module: nn.Module) -> tuple[str, ...]: + # Match the materializer's early exit: only modules that own at + # least one parameter slot can be consumers in the native walk. + if not module._parameters or self.should_skip_module(module_path): + return () + + normalized_path = source_module_path(module_path, module) + module_parts = normalized_path.split('.') if normalized_path else [] + module_name = module_parts[-1] if module_parts else "" + parent_parts = module_parts[:-1] + + # A fused destination consumes raw tensors from explicitly named + # sibling subtrees, for example q_proj/k_proj/v_proj -> qkv_proj. + if module_name and self.does_require_special_handling(module_name): + prefixes = tuple('.'.join(parent_parts + [source_name]) + for source_name in self.mapping[module_name]) + return take_prefixes(prefixes) + + # Custom module loaders receive the full structurally matched + # subtree. Do not treat the root as a subtree consumer: an empty + # prefix contains the whole checkpoint and conveys no useful order. + # Other model-specific path rewrites that are not expressed by the + # mapper (for example a checkpoint-only namespace prefix) + # intentionally miss here and remain in the safe, all-source + # unmatched tail. + if normalized_path and (self.is_special_instance_module(module) + or hasattr(module, 'load_weights')): + return take_prefixes((normalized_path, )) + + # The native fallback copies only direct parameters. Descendant + # parameters belong to their own named-module consumers. + parameter_names = tuple( + f"{normalized_path}.{name}" if normalized_path else name + for name, _ in module.named_parameters(recurse=False)) + return take_exact(parameter_names) + + preload_suffixes = tuple( + getattr(self._model, "preload_weight_modules", None) or ()) + preload_modules: list[tuple[str, nn.Module]] = [] + preload_paths: set[str] = set() + for suffix in preload_suffixes: + for module_path, module in modules: + if (module_path.endswith(suffix) + and module_path not in preload_paths): + preload_paths.add(module_path) + preload_modules.append((module_path, module)) + + previous_preload_group: str | None = None + for module_path, module in preload_modules: + source_names = infer_source_names(module_path, module) + if not source_names: + continue + group_id = f"preload:{len(demands):06d}:{module_path}" + predecessors = ((previous_preload_group, ) + if previous_preload_group is not None else ()) + demands.append( + WeightDemand( + group_id=group_id, + source_names=source_names, + destination_ranks=(mapping.rank, ), + priority=len(demands), + predecessors=predecessors, + )) + previous_preload_group = group_id + + # Native materialization schedules non-preloaded modules concurrently. + # Equal priorities let read-ahead preserve physical locality without + # inventing dependencies between these consumers. + nonpreload_priority = len(demands) + for module_path, module in modules: + if module_path in preload_paths: + continue + source_names = infer_source_names(module_path, module) + if not source_names: + continue + demands.append( + WeightDemand( + group_id= + f"module:{len(demands):06d}:{module_path or ''}", + source_names=source_names, + destination_ranks=(mapping.rank, ), + priority=nonpreload_priority, + )) + + # If structural inspection found nothing useful, expose no ordering + # claim. Otherwise retain every unresolved tensor as a physical tail. + if not demands: + return self._build_opaque_weight_load_plan(catalog, physical_names) + + unmatched_names = tuple(name for name in physical_names + if name in remaining_names) + if unmatched_names: + demands.append( + WeightDemand( + group_id="unmatched_checkpoint_tensors", + source_names=unmatched_names, + destination_ranks=(mapping.rank, ), + priority=nonpreload_priority + 1, + )) + + plan = WeightLoadPlan( + catalog_id=catalog.catalog_id, + rank=mapping.rank, + world_size=mapping.world_size, + coverage=WeightLoadPlanCoverage.CONSERVATIVE, + ordering=WeightLoadOrderConfidence.ADVISORY, + demands=tuple(demands), + ) + plan.validate_against(catalog) + return plan + + def _build_opaque_weight_load_plan( + self, catalog: CheckpointCatalog, + physical_names: tuple[str, ...]) -> WeightLoadPlan: + """Build an all-source plan when no useful ordering can be inferred.""" + mapping = self.model.model_config.mapping plan = WeightLoadPlan( catalog_id=catalog.catalog_id, rank=mapping.rank, @@ -106,7 +293,7 @@ def build_weight_load_plan(self, ordering=WeightLoadOrderConfidence.OPAQUE, demands=(WeightDemand( group_id="all_checkpoint_tensors", - source_names=tuple(sorted(catalog.tensor_names)), + source_names=physical_names, destination_ranks=(mapping.rank, ), ), ), ) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py index c495230779cc..aaf5da7ed545 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/checkpoint_loader.py @@ -93,6 +93,20 @@ def build_checkpoint_catalog(self, checkpoint_dir: str, use_consolidated=kwargs.get("use_consolidated", False), ) + def activate_weight_session( + self, + readiness_error: BaseException | None = None, + *, + weight_mapper: BaseWeightMapper | None = None) -> None: + """Start deferred read-ahead only on the built-in HF path.""" + if (type(self) is HfCheckpointLoader + and type(self.weight_loader) is HfWeightLoader): + self.weight_loader.activate_weight_session( + readiness_error, weight_mapper=weight_mapper) + return + super().activate_weight_session(readiness_error, + weight_mapper=weight_mapper) + @property def weight_loader(self) -> BaseWeightLoader: return self._weight_loader diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/rank_striped_read_ahead.py b/tensorrt_llm/_torch/models/checkpoints/hf/rank_striped_read_ahead.py index 6ee41ba14ac6..fd99e201c19f 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/rank_striped_read_ahead.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/rank_striped_read_ahead.py @@ -5,13 +5,22 @@ import os import threading +from collections import Counter from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass +from heapq import heappop, heappush from pathlib import Path from typing import Sequence import psutil +from tensorrt_llm._torch.models.checkpoints.checkpoint_catalog import CheckpointCatalog +from tensorrt_llm._torch.models.checkpoints.weight_load_plan import ( + WeightDemand, + WeightLoadOrderConfidence, + WeightLoadPlan, +) + # Work-assignment granularity across ranks; one extent stays with one issuer. _CHUNK_SIZE = 256 * 1024 * 1024 # I/O and cancellation granularity within an extent; one pread reads this much. @@ -64,21 +73,188 @@ def distribute_worker_budget(local_size: int) -> tuple[int, ...]: return tuple(workers + (rank < remainder) for rank in range(local_size)) +def _build_source_extents( + files: Sequence[tuple[str, int]], +) -> list[ReadAheadExtent]: + extents = [] + for path, file_size in sorted(files): + if file_size < 0: + raise ValueError("checkpoint file sizes must be nonnegative") + for offset in range(0, file_size, _CHUNK_SIZE): + extents.append( + ReadAheadExtent( + path, + offset, + min(_CHUNK_SIZE, file_size - offset), + ) + ) + return extents + + +def _topological_demands(plan: WeightLoadPlan) -> tuple[WeightDemand, ...]: + """Return a stable priority-aware topological order for one rank plan.""" + demand_by_id = {demand.group_id: demand for demand in plan.demands} + successors = {group_id: [] for group_id in demand_by_id} + remaining_predecessors = {} + for demand in plan.demands: + remaining_predecessors[demand.group_id] = len(demand.predecessors) + for predecessor in demand.predecessors: + successors[predecessor].append(demand.group_id) + + ready = [ + ( + demand.priority, + demand.group_id, + ) + for demand in plan.demands + if not demand.predecessors + ] + ready.sort() + ordered = [] + while ready: + _, group_id = heappop(ready) + demand = demand_by_id[group_id] + ordered.append(demand) + for successor in successors[group_id]: + remaining_predecessors[successor] -= 1 + if remaining_predecessors[successor] == 0: + successor_demand = demand_by_id[successor] + heappush( + ready, + ( + successor_demand.priority, + successor, + ), + ) + + if len(ordered) != len(plan.demands): + raise ValueError("weight demand predecessor graph must be acyclic") + return tuple(ordered) + + +def compile_weight_plan_extent_order( + catalog: CheckpointCatalog, + plans: Sequence[WeightLoadPlan], + files: Sequence[tuple[str, int]], +) -> tuple[ReadAheadExtent, ...]: + """Prioritize whole source chunks without changing read coverage. + + The earliest ordering-qualified demand for a tensor across node-local rank + plans wins. Tensor ranges only reprioritize the existing fixed-size source + chunks; chunks are never split or omitted, and physically ordered unmatched + chunks form the tail. Conservative plans are consumed through + ``described_tensor_names`` rather than the selective-I/O-only + ``selected_tensor_names`` property. If every plan has opaque ordering, the + source's physical chunk order is preserved. + """ + if not catalog.has_complete_byte_ranges: + raise ValueError("checkpoint catalog must expose complete byte ranges") + if not plans: + raise ValueError("at least one weight load plan is required") + + paths_by_object_id = {} + sizes_by_object_id = {} + for path, size_bytes in files: + object_id = os.path.basename(path) + if object_id in paths_by_object_id: + raise ValueError(f"checkpoint paths have duplicate basename object ID {object_id!r}") + paths_by_object_id[object_id] = path + sizes_by_object_id[object_id] = size_bytes + + catalog_sizes = {obj.object_id: obj.size_bytes for obj in catalog.objects} + missing_objects = catalog_sizes.keys() - paths_by_object_id.keys() + unexpected_objects = paths_by_object_id.keys() - catalog_sizes.keys() + if missing_objects or unexpected_objects: + raise ValueError( + "checkpoint catalog and source files resolve different objects: " + f"missing={sorted(missing_objects)}, " + f"unexpected={sorted(unexpected_objects)}" + ) + mismatched_sizes = { + object_id: (catalog_sizes[object_id], sizes_by_object_id[object_id]) + for object_id in catalog_sizes + if catalog_sizes[object_id] != sizes_by_object_id[object_id] + } + if mismatched_sizes: + raise ValueError( + f"checkpoint catalog object sizes do not match source files: {mismatched_sizes}" + ) + + ordered_plans = tuple(sorted(plans, key=lambda plan: plan.rank)) + if len({plan.rank for plan in ordered_plans}) != len(ordered_plans): + raise ValueError("node-local weight load plans must have unique ranks") + world_sizes = {plan.world_size for plan in ordered_plans} + if len(world_sizes) != 1: + raise ValueError("node-local weight load plans must agree on world_size") + for plan in ordered_plans: + plan.validate_against(catalog) + + # Dependencies establish an earliest read-ahead stage. Equal-priority, + # independent groups share a stage: their native consumers may execute + # concurrently, so preserve physical locality instead of inventing a + # tensor-name ordering. + tensor_priorities: dict[str, int] = {} + for plan in ordered_plans: + if plan.ordering is WeightLoadOrderConfidence.OPAQUE: + continue + group_priorities: dict[str, int] = {} + for demand in _topological_demands(plan): + priority = max( + demand.priority, + max( + (group_priorities[p] + 1 for p in demand.predecessors), + default=demand.priority, + ), + ) + group_priorities[demand.group_id] = priority + for tensor_name in demand.source_names: + tensor_priorities[tensor_name] = min( + priority, tensor_priorities.get(tensor_name, priority) + ) + + source_extents = _build_source_extents(files) + chunk_priorities: dict[ReadAheadExtent, int] = {} + chunks_by_object_id = {object_id: [] for object_id in paths_by_object_id} + for extent in source_extents: + chunks_by_object_id[os.path.basename(extent.path)].append(extent) + + for tensor_name, priority in tensor_priorities.items(): + for tensor_extent in catalog.get_tensor(tensor_name).extents: + if tensor_extent.length_bytes == 0: + continue + chunks = chunks_by_object_id[tensor_extent.object_id] + first_chunk = tensor_extent.offset_bytes // _CHUNK_SIZE + final_chunk = (tensor_extent.end_offset_bytes - 1) // _CHUNK_SIZE + for chunk in chunks[first_chunk : final_chunk + 1]: + chunk_priorities[chunk] = min(priority, chunk_priorities.get(chunk, priority)) + + tail_priority = max(chunk_priorities.values(), default=0) + 1 + # Python's stable sort preserves physical order within each priority and + # among all unmatched chunks. Every original chunk remains present once. + return tuple( + sorted( + source_extents, + key=lambda extent: chunk_priorities.get(extent, tail_priority), + ) + ) + + def build_local_plan( files: Sequence[tuple[str, int]], local_rank: int, local_size: int, + *, + ordered_extents: Sequence[ReadAheadExtent] | None = None, ) -> ReadAheadPlan: """Assign every checkpoint extent to exactly one node-local issuer.""" if not 0 <= local_rank < local_size: raise ValueError("local rank and size must describe a valid group") - extents = [] - for path, file_size in sorted(files): - if file_size < 0: - raise ValueError("checkpoint file sizes must be nonnegative") - for offset in range(0, file_size, _CHUNK_SIZE): - extents.append(ReadAheadExtent(path, offset, min(_CHUNK_SIZE, file_size - offset))) + extents = _build_source_extents(files) + if ordered_extents is not None: + if Counter(ordered_extents) != Counter(extents): + raise ValueError("ordered_extents must be a complete permutation of checkpoint extents") + extents = list(ordered_extents) worker_counts = distribute_worker_budget(local_size) issuer_ranks = [rank for rank, workers in enumerate(worker_counts) if workers > 0] @@ -207,19 +383,22 @@ def __init__(self, active_communicator, node_communicator, plan: ReadAheadPlan) self._node_communicator = node_communicator self._plan = plan self._cancel = threading.Event() + self._read_release = threading.Event() self._thread: threading.Thread | None = None self._file_descriptors: dict[str, int] = {} self._read_error: BaseException | None = None self._closed = False try: - for path in {extent.path for extent in plan.extents}: + for path in dict.fromkeys(extent.path for extent in plan.extents): self._file_descriptors[path] = os.open(path, os.O_RDONLY) except Exception: self._close_file_descriptors() raise - def start(self) -> "RankStripedReadAheadSession": + def start(self, *, defer_reads: bool = False) -> "RankStripedReadAheadSession": + if not defer_reads: + self._read_release.set() if not self._plan.extents: return self self._thread = threading.Thread( @@ -233,6 +412,10 @@ def start(self) -> "RankStripedReadAheadSession": raise return self + def release_reads(self) -> None: + """Release prepared workers only after activation consensus.""" + self._read_release.set() + def _read_extent(self, extent: ReadAheadExtent) -> int: file_descriptor = self._file_descriptors[extent.path] offset = extent.offset @@ -251,8 +434,17 @@ def _read_extent(self, extent: ReadAheadExtent) -> int: completed += bytes_read return completed + def _request_cancel(self) -> None: + self._cancel.set() + # Wake a prepared worker that is still behind the activation gate. It + # observes cancellation before issuing any payload I/O. + self._read_release.set() + def _run(self) -> None: try: + self._read_release.wait() + if self._cancel.is_set(): + return with ThreadPoolExecutor(max_workers=self._plan.workers) as executor: futures = [ executor.submit(self._read_extent, extent) for extent in self._plan.extents @@ -265,7 +457,7 @@ def _run(self) -> None: break future.result() except Exception: - self._cancel.set() + self._request_cancel() for future in futures: future.cancel() raise @@ -309,7 +501,7 @@ def cancel_and_close(self) -> BaseException | None: def cancel_reads(self) -> BaseException | None: """Stop background work while leaving the node communicator usable.""" - self._cancel.set() + self._request_cancel() if self._thread is not None: self._thread.join() return self._close_file_descriptors() @@ -324,7 +516,7 @@ def finish(self, body_error: BaseException | None) -> RuntimeError | None: # cannot improve first-token latency. Bound the remaining I/O volume # to the currently executing small pread on each worker; synchronous # storage latency itself is not cancellable here. - self._cancel.set() + self._request_cancel() if self._thread is not None: self._thread.join() diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index 3eeae6ceb74a..2fc48ece3e30 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -31,6 +31,18 @@ import tqdm from mpi4py import MPI as _MPI +from tensorrt_llm._checkpoint_io_policy import \ + CHECKPOINT_IO_POLICIES as _SUPPORTED_REQUESTED_IO_POLICIES +from tensorrt_llm._checkpoint_io_policy import \ + DEMAND_ORDERED_IO_POLICY as _DEMAND_ORDERED_IO_POLICY +from tensorrt_llm._checkpoint_io_policy import \ + EXECUTABLE_CHECKPOINT_IO_POLICIES as _SUPPORTED_IO_POLICIES +from tensorrt_llm._checkpoint_io_policy import \ + NATIVE_IO_POLICY as _NATIVE_IO_POLICY +from tensorrt_llm._checkpoint_io_policy import \ + RANK_STRIPED_IO_POLICIES as _RANK_STRIPED_IO_POLICIES +from tensorrt_llm._checkpoint_io_policy import \ + RANK_STRIPED_IO_POLICY as _RANK_STRIPED_IO_POLICY from tensorrt_llm._torch.mmap_utils import populate_file_pages from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ( BaseWeightLoader, ConsumableWeightsDict) @@ -40,7 +52,10 @@ build_safetensors_checkpoint_catalog from tensorrt_llm._torch.models.checkpoints.hf.rank_striped_read_ahead import ( RankStripedReadAheadSession, build_local_plan, close_node_communicator, - coordinate_error, effective_available_host_memory, memory_admission) + compile_weight_plan_extent_order, coordinate_error, + effective_available_host_memory, memory_admission) +from tensorrt_llm._torch.models.checkpoints.weight_load_plan import \ + WeightLoadOrderConfidence from tensorrt_llm._torch.models.modeling_utils import ( register_checkpoint_weight_loader, run_concurrently) from tensorrt_llm._utils import (ENABLE_MULTI_DEVICE, local_mpi_barrier, @@ -51,11 +66,7 @@ _WEIGHT_CACHE_ENV = "TRTLLM_HF_WEIGHT_CACHE" _WEIGHT_CACHE_MAX_ENTRIES_ENV = "TRTLLM_HF_WEIGHT_CACHE_MAX_ENTRIES" -_AUTO_IO_POLICY = "auto" -_NATIVE_IO_POLICY = "native" -_RANK_STRIPED_IO_POLICY = "rank_striped_read_ahead" -_SUPPORTED_IO_POLICIES = (_NATIVE_IO_POLICY, _RANK_STRIPED_IO_POLICY) -_SUPPORTED_REQUESTED_IO_POLICIES = (_AUTO_IO_POLICY, ) + _SUPPORTED_IO_POLICIES +_PHYSICAL_ORDER_ENV = "TRTLLM_DEMAND_ORDERED_READ_AHEAD_PHYSICAL_ORDER" # Model families whose checkpoints are too large to materialize in host RAM; # their models stream rank-local slices out of the lazy mmapped handles. _LAZY_SAFETENSORS_MODEL_TYPES = ("kimi_k3", "kimi_linear") @@ -101,6 +112,19 @@ class _CheckpointIOStatus: fallback_reason: str | None = None +@dataclass +class _DemandReadAheadContext: + """Keep preflight metadata alive until mapper-informed reader setup.""" + + active_communicator: Any + node_communicator: Any + catalog: CheckpointCatalog + files: list[tuple[str, int]] + local_rank: int + local_size: int + reader: RankStripedReadAheadSession | None = None + + @register_checkpoint_weight_loader("MX") @register_checkpoint_weight_loader("mistral") @register_checkpoint_weight_loader("mistral_large_3") @@ -116,6 +140,7 @@ class HfWeightLoader(BaseWeightLoader): _requested_checkpoint_io_policy = _NATIVE_IO_POLICY _selection_fallback_reason: str | None = None _partial_model_loading = False + _pending_read_ahead_session: _DemandReadAheadContext | None = None _last_checkpoint_io_status = _CheckpointIOStatus( requested=_NATIVE_IO_POLICY, selected=_NATIVE_IO_POLICY, @@ -128,6 +153,7 @@ def __init__(self, selection_fallback_reason: str | None = None, partial_model_loading: bool = False) -> None: self._partial_model_loading = partial_model_loading + self._pending_read_ahead_session = None self._configure_checkpoint_io_policy( checkpoint_io_policy, requested_checkpoint_io_policy=requested_checkpoint_io_policy, @@ -164,6 +190,150 @@ def _configure_checkpoint_io_policy( def checkpoint_io_policy(self) -> str: return self._checkpoint_io_policy + def activate_weight_session(self, + readiness_error: BaseException | None = None, + *, + weight_mapper: Any = None) -> None: + """Collectively validate readiness, then activate deferred reads.""" + context = self._pending_read_ahead_session + if (self._checkpoint_io_policy != _DEMAND_ORDERED_IO_POLICY + or context is None): + if readiness_error is not None: + raise readiness_error + return + + coordinated_readiness_error = coordinate_error( + context.active_communicator, + "rank-striped mapper readiness", + readiness_error, + ) + if coordinated_readiness_error is not None: + # Every rank rejects activation before materialization. The open + # context still owns the node communicator for coordinated cleanup. + status = self._last_checkpoint_io_status + status.activated = False + status.fallback_reason = str(coordinated_readiness_error) + if readiness_error is not None: + raise readiness_error + raise coordinated_readiness_error + + local_plan = None + plan_error = None + try: + if weight_mapper is None: + raise ValueError( + "demand-ordered activation requires an initialized " + "weight mapper") + catalog = context.catalog + local_plan = weight_mapper.build_weight_load_plan(catalog) + local_plan.validate_against(catalog) + except Exception as error: + plan_error = error + coordinated_plan_error = coordinate_error( + context.active_communicator, + "demand-ordered weight plan build", + plan_error, + ) + if coordinated_plan_error is not None: + self._degrade_demand_ordered_activation("weight plan build", + coordinated_plan_error) + return + assert local_plan is not None + + session = None + setup_error = None + try: + node_communicator = context.node_communicator + physical_order = os.environ.get(_PHYSICAL_ORDER_ENV, + "0").lower() in ("1", "true", "yes", + "on") + node_inputs = ([(local_plan, physical_order)] if node_communicator + is None else node_communicator.allgather( + (local_plan, physical_order))) + if any(item[1] != physical_order for item in node_inputs): + raise ValueError( + "physical-order ablation must agree across node-local ranks" + ) + node_plans = [item[0] for item in node_inputs] + ordered_extents = compile_weight_plan_extent_order( + context.catalog, node_plans, context.files) + compiled_plan = build_local_plan( + context.files, + context.local_rank, + context.local_size, + ordered_extents=None if physical_order else ordered_extents) + session = RankStripedReadAheadSession(context.active_communicator, + node_communicator, + compiled_plan) + except Exception as error: + setup_error = error + coordinated_setup_error = coordinate_error( + context.active_communicator, + "demand-ordered reader setup", + setup_error, + ) + if coordinated_setup_error is not None: + self._close_unactivated_session(session, + context.node_communicator, + context.active_communicator, + "demand-ordered setup cleanup", + preserve_node_communicator=True) + self._degrade_demand_ordered_activation("reader setup", + coordinated_setup_error) + return + assert session is not None + context.reader = session + ordering = ("physical_ablation" if physical_order else "demand" if any( + plan.ordering is not WeightLoadOrderConfidence.OPAQUE + for plan in node_plans) else "physical_opaque") + logger.info( + "Demand-ordered checkpoint read-ahead configured: " + f"ordering={ordering}, local_plans={len(node_plans)}, " + f"assigned_extents={len(compiled_plan.extents)}. " + "All checkpoint chunks are retained; materialization is unchanged.") + + activation_error = None + try: + session.start(defer_reads=True) + except Exception as error: + activation_error = error + coordinated_error = coordinate_error( + context.active_communicator, + "rank-striped reader activation", + activation_error, + ) + if coordinated_error is not None: + # Every participating rank reaches this decision before model + # mutation. Cancel any readers that started successfully and keep + # using the already-mapped native tensors. + cleanup_error = session.cancel_reads() + coordinated_cleanup_error = coordinate_error( + context.active_communicator, + "rank-striped activation cleanup", + cleanup_error, + ) + status = self._last_checkpoint_io_status + status.activated = False + status.fallback_reason = str(coordinated_error) + logger.warning( + "Rank-striped activation degraded collectively before " + f"model materialization: {coordinated_error}") + if coordinated_cleanup_error is not None: + raise coordinated_cleanup_error + return + session.release_reads() + self._last_checkpoint_io_status.activated = True + + def _degrade_demand_ordered_activation(self, phase: str, + error: BaseException) -> None: + """Disable unstarted advisory work after an all-rank decision.""" + status = self._last_checkpoint_io_status + status.activated = False + status.fallback_reason = str(error) + logger.warning( + "Demand-ordered read-ahead degraded collectively during " + f"{phase}; keeping the already-mapped native tensors: {error}") + @property def last_checkpoint_io_status(self) -> _CheckpointIOStatus: return replace(self._last_checkpoint_io_status) @@ -522,6 +692,8 @@ def open_weight_session(self, yield weights return + if isinstance(session, _DemandReadAheadContext): + self._pending_read_ahead_session = session body_error = None try: yield weights @@ -530,38 +702,63 @@ def open_weight_session(self, 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) + try: + if isinstance(session, _DemandReadAheadContext): + reader = session.reader + if reader is None: + coordinated_body_error = coordinate_error( + session.active_communicator, + "rank-striped model materialization", + body_error) + self._close_unactivated_session( + None, session.node_communicator, + session.active_communicator, + "rank-striped cleanup") + if coordinated_body_error is not None: + raise coordinated_body_error + read_error = None + else: + read_error = reader.finish(body_error) + else: + 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() - 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 = _RANK_STRIPED_IO_POLICY + # 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: - # Read-ahead is advisory. The native mmap/materialization - # path succeeded, so never retry a partially mutated model. - status.effective = _NATIVE_IO_POLICY - 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() + status = self._last_checkpoint_io_status + if read_error is None and status.activated: + status.effective = status.selected + else: + # Read-ahead is advisory. The native mmap/materialization + # path succeeded, so never retry a partially mutated model. + status.effective = _NATIVE_IO_POLICY + if read_error is not None: + status.fallback_reason = str(read_error) + logger.warning( + "Rank-striped read-ahead degraded after " + "activation; keeping the successfully " + f"materialized model: {read_error}") + elif status.fallback_reason is None: + status.fallback_reason = "read-ahead session was not activated" + self._log_checkpoint_io_status() + finally: + if self._pending_read_ahead_session is session: + self._pending_read_ahead_session = None @staticmethod def _active_communicator(): @@ -636,7 +833,7 @@ def _fallback_to_native(self, message = ("Checkpoint I/O policy is falling back before model " f"materialization: requested={status.requested}, " f"selected={status.selected}, reason={reason}.") - if status.requested == _RANK_STRIPED_IO_POLICY: + if status.requested in _RANK_STRIPED_IO_POLICIES: logger.warning(message) else: logger.info(message) @@ -723,7 +920,8 @@ def _start_rank_striped_read_ahead( use_consolidated: bool, active_communicator, **kwargs, - ) -> tuple[dict[str, Any], RankStripedReadAheadSession | None]: + ) -> tuple[dict[str, Any], RankStripedReadAheadSession + | _DemandReadAheadContext | None]: """Load weights via rank-striped read-ahead, or fall back to native.""" node_communicator = None split_error = None @@ -749,6 +947,7 @@ def _start_rank_striped_read_ahead( weight_files = [] stats = [] + checkpoint_catalog = None available_memory = 0 eligibility_reason = None preflight_error = None @@ -768,6 +967,10 @@ def _start_rank_striped_read_ahead( elif (self._partial_model_loading or int(os.environ.get("TLLM_OVERRIDE_LAYER_NUM", "0")) != 0): eligibility_reason = "partial model loading was requested" + if (self._checkpoint_io_policy == _DEMAND_ORDERED_IO_POLICY + and eligibility_reason is None): + checkpoint_catalog = build_safetensors_checkpoint_catalog( + weight_files) available_memory = effective_available_host_memory() except Exception as error: preflight_error = error @@ -827,10 +1030,12 @@ def _start_rank_striped_read_ahead( discovery_signature = tuple( (os.path.basename(path), size) for path, size in file_sizes) + catalog_id = (None if checkpoint_catalog is None else + checkpoint_catalog.catalog_id) selections = ([ - (discovery_signature, eligibility_reason) + (discovery_signature, catalog_id, eligibility_reason) ] if active_communicator is None else active_communicator.allgather( - (discovery_signature, eligibility_reason))) + (discovery_signature, catalog_id, eligibility_reason))) if any(selection[0] != selections[0][0] for selection in selections[1:]): self._close_unactivated_session( @@ -839,9 +1044,9 @@ def _start_rank_striped_read_ahead( raise RuntimeError( "SafeTensors checkpoint discovery must match " f"across model-load ranks; received {selections}.") - fallback_reasons = [(rank, selection[1]) + fallback_reasons = [(rank, selection[2]) for rank, selection in enumerate(selections) - if selection[1] is not None] + if selection[2] is not None] if fallback_reasons: rank, reason = fallback_reasons[0] self._last_checkpoint_io_status.selected = _NATIVE_IO_POLICY @@ -850,36 +1055,56 @@ def _start_rank_striped_read_ahead( f"rank {rank}: {reason}", active_communicator, node_communicator, **kwargs), None - - self._last_checkpoint_io_status.selected = _RANK_STRIPED_IO_POLICY - session = None - setup_error = None - try: - plan = build_local_plan(file_sizes, local_rank, local_size) - session = RankStripedReadAheadSession(active_communicator, - node_communicator, - plan).start() - except Exception as error: - setup_error = error - coordinated_setup_error = coordinate_error(active_communicator, - "rank-striped reader setup", - setup_error) - if coordinated_setup_error is not None: - return self._fallback_to_native(checkpoint_dir, mapping, - use_consolidated, - str(coordinated_setup_error), - active_communicator, - node_communicator, session, - **kwargs), None - assert session is not None + if any(selection[1] != selections[0][1] + for selection in selections[1:]): + self._close_unactivated_session( + None, node_communicator, active_communicator, + "rank-striped catalog mismatch cleanup") + raise RuntimeError( + "SafeTensors checkpoint catalogs must match across " + f"model-load ranks; received {selections}.") + + self._last_checkpoint_io_status.selected = self._checkpoint_io_policy + session: RankStripedReadAheadSession | _DemandReadAheadContext | None + if self._checkpoint_io_policy == _DEMAND_ORDERED_IO_POLICY: + assert checkpoint_catalog is not None + session = _DemandReadAheadContext( + active_communicator, + node_communicator, + checkpoint_catalog, + file_sizes, + local_rank, + local_size, + ) + else: + session = None + setup_error = None + try: + plan = build_local_plan(file_sizes, local_rank, local_size) + # Preserve v1 overlap with SafeTensors mapping and mapper init. + session = RankStripedReadAheadSession(active_communicator, + node_communicator, plan) + session.start() + except Exception as error: + setup_error = error + coordinated_setup_error = coordinate_error( + active_communicator, "rank-striped reader setup", setup_error) + if coordinated_setup_error is not None: + return self._fallback_to_native(checkpoint_dir, mapping, + use_consolidated, + str(coordinated_setup_error), + active_communicator, + node_communicator, session, + **kwargs), None + assert session is not None + logger.info("Rank-striped checkpoint read-ahead activated: " + f"local_rank={local_rank}, local_size={local_size}, " + f"workers={plan.workers}, assigned_bytes=" + f"{plan.assigned_bytes}, checkpoint_bytes=" + f"{checkpoint_bytes}.") status = self._last_checkpoint_io_status - status.activated = True - logger.info("Rank-striped checkpoint read-ahead activated: " - f"local_rank={local_rank}, local_size={local_size}, " - f"workers={plan.workers}, assigned_bytes=" - f"{plan.assigned_bytes}, checkpoint_bytes=" - f"{checkpoint_bytes}.") + status.activated = self._checkpoint_io_policy == _RANK_STRIPED_IO_POLICY weights = None mapping_error = None @@ -893,13 +1118,14 @@ def _start_rank_striped_read_ahead( active_communicator, "rank-striped SafeTensors mapping", mapping_error) if coordinated_mapping_error is not None: - cleanup_error = session.cancel_and_close() - coordinated_cleanup_error = coordinate_error( - active_communicator, "rank-striped mapping cleanup", - cleanup_error) - if coordinated_cleanup_error is not None: + try: + self._close_unactivated_session( + None if isinstance(session, _DemandReadAheadContext) else + session, node_communicator, active_communicator, + "rank-striped mapping cleanup") + except Exception as cleanup_error: logger.error("Rank-striped cleanup also failed after mapping " - f"failure: {coordinated_cleanup_error}") + f"failure: {cleanup_error}") status.effective = "none" status.fallback_reason = str(coordinated_mapping_error) self._log_checkpoint_io_status() diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index e1ee1d5ee08f..bac5643f5e88 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -7,13 +7,25 @@ import sys import traceback import warnings -from contextlib import AbstractContextManager, contextmanager +from contextlib import AbstractContextManager, contextmanager, nullcontext from dataclasses import replace from enum import Enum from typing import Any, Callable, Iterator, Optional, Tuple import torch +from tensorrt_llm._checkpoint_io_policy import \ + AUTO_IO_POLICY as _AUTO_CHECKPOINT_IO_POLICY +from tensorrt_llm._checkpoint_io_policy import \ + CHECKPOINT_IO_POLICIES as _SUPPORTED_CHECKPOINT_IO_POLICIES +from tensorrt_llm._checkpoint_io_policy import \ + DEMAND_ORDERED_IO_POLICY as _DEMAND_ORDERED_CHECKPOINT_IO_POLICY +from tensorrt_llm._checkpoint_io_policy import \ + NATIVE_IO_POLICY as _NATIVE_CHECKPOINT_IO_POLICY +from tensorrt_llm._checkpoint_io_policy import \ + RANK_STRIPED_IO_POLICIES as _RANK_STRIPED_CHECKPOINT_IO_POLICIES +from tensorrt_llm._checkpoint_io_policy import \ + RANK_STRIPED_IO_POLICY as _RANK_STRIPED_CHECKPOINT_IO_POLICY from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import ( AutoCheckpointMapper, BaseCheckpointLoader) from tensorrt_llm._torch.peft.lora.config import LoraConfig @@ -344,14 +356,6 @@ def get_rank_model_storage(model): return total_bytes -_AUTO_CHECKPOINT_IO_POLICY = "auto" -_NATIVE_CHECKPOINT_IO_POLICY = "native" -_RANK_STRIPED_CHECKPOINT_IO_POLICY = "rank_striped_read_ahead" -_SUPPORTED_CHECKPOINT_IO_POLICIES = ( - _AUTO_CHECKPOINT_IO_POLICY, - _NATIVE_CHECKPOINT_IO_POLICY, - _RANK_STRIPED_CHECKPOINT_IO_POLICY, -) _SHADOW_WEIGHT_LOAD_PLAN_ENV = "TRTLLM_SHADOW_WEIGHT_LOAD_PLAN" @@ -407,13 +411,15 @@ def _resolve_checkpoint_io_policy( f"requested={requested_policy}, " f"selected={_NATIVE_CHECKPOINT_IO_POLICY}, " f"reason={reason}.") - if requested_policy == _RANK_STRIPED_CHECKPOINT_IO_POLICY: + if requested_policy in _RANK_STRIPED_CHECKPOINT_IO_POLICIES: logger.warning(message) else: logger.info(message) return _NATIVE_CHECKPOINT_IO_POLICY, reason - return _RANK_STRIPED_CHECKPOINT_IO_POLICY, None + if requested_policy == _AUTO_CHECKPOINT_IO_POLICY: + return _RANK_STRIPED_CHECKPOINT_IO_POLICY, None + return requested_policy, None def _construct_checkpoint_loader( @@ -501,6 +507,27 @@ def _open_checkpoint_weight_session( return checkpoint_loader.open_weight_session(checkpoint_dir, **kwargs) +def _activate_checkpoint_weight_session( + checkpoint_loader: Any, + metrics: dict[str, float], + preparation_metric_name: str, + readiness_error: BaseException | None = None, + *, + weight_mapper: Any = None, +) -> None: + """Coordinate mapper readiness before activating deferred loader work.""" + activate = getattr(checkpoint_loader, "activate_weight_session", None) + if activate is not None: + weight_loader = getattr(checkpoint_loader, "weight_loader", None) + timer = (timing_metric(preparation_metric_name, metrics) if getattr( + weight_loader, "checkpoint_io_policy", None) + == _DEMAND_ORDERED_CHECKPOINT_IO_POLICY else nullcontext()) + with timer: + activate(readiness_error, weight_mapper=weight_mapper) + if readiness_error is not None: + raise readiness_error + + @contextmanager def _timed_checkpoint_weight_session( checkpoint_loader: Any, @@ -1946,20 +1973,33 @@ 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.""" + preparation_metric_name = ModelLoaderMetricNames.CHECKPOINT_PREPARATION_SECONDS.value with _timed_checkpoint_weight_session( checkpoint_loader, checkpoint_dir, self._metrics, - ModelLoaderMetricNames.CHECKPOINT_PREPARATION_SECONDS.value, + preparation_metric_name, 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) + try: + weights_preloaded = checkpoint_loader.is_weights_preloaded() + self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper( + model, config) + except BaseException as error: + _activate_checkpoint_weight_session(checkpoint_loader, + self._metrics, + preparation_metric_name, + error) + raise _inspect_shadow_weight_load_plan( checkpoint_loader, checkpoint_dir, self.weight_mapper, **load_weights_kwargs, ) + _activate_checkpoint_weight_session( + checkpoint_loader, + self._metrics, + preparation_metric_name, + weight_mapper=self.weight_mapper) if weights: with timing_metric( ModelLoaderMetricNames.WEIGHT_POPULATION_SECONDS.value, @@ -1973,31 +2013,43 @@ def _materialize_checkpoint_weights(self, checkpoint_loader, def _materialize_draft_checkpoint_weights(self, checkpoint_loader, model) -> None: """Keep loader-specific work alive through draft materialization.""" + preparation_metric_name = ModelLoaderMetricNames.DRAFT_CHECKPOINT_PREPARATION_SECONDS.value with _timed_checkpoint_weight_session( checkpoint_loader, self.spec_config.speculative_model, self._metrics, - ModelLoaderMetricNames.DRAFT_CHECKPOINT_PREPARATION_SECONDS. - value, + preparation_metric_name, 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: - # MTP one-model + separate MTP checkpoint: no draft HF architecture. - draft_weight_mapper = self.weight_mapper + try: + 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: + # MTP one-model + separate MTP checkpoint: no draft HF architecture. + draft_weight_mapper = self.weight_mapper + except BaseException as error: + _activate_checkpoint_weight_session(checkpoint_loader, + self._metrics, + preparation_metric_name, + error) + raise _inspect_shadow_weight_load_plan( checkpoint_loader, self.spec_config.speculative_model, draft_weight_mapper, mapping=self.mapping, ) + _activate_checkpoint_weight_session( + checkpoint_loader, + self._metrics, + preparation_metric_name, + weight_mapper=draft_weight_mapper) with timing_metric( ModelLoaderMetricNames.DRAFT_WEIGHT_POPULATION_SECONDS. value, self._metrics): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 59c5030de54c..d8c3c756651c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -47,6 +47,7 @@ LoraConfig, get_default_trtllm_modules_to_hf_modules) from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType +from .._checkpoint_io_policy import RANK_STRIPED_IO_POLICIES, CheckpointIoPolicy from .._utils import (_str_to_torch_dtype_dict, is_sm_100f, mpi_rank, prefer_pinned) @@ -5885,24 +5886,28 @@ def validate_mla_skip_correction_config(self) -> 'TorchLlmArgs': status="prototype", ) - checkpoint_io_policy: Literal[ - "auto", "native", "rank_striped_read_ahead"] = Field( - default="auto", - description= - "Controls checkpoint storage I/O independently of checkpoint format. " - "'auto' selects rank-striped read-ahead for compatible built-in " - "PyTorch/HF loads and selects native I/O otherwise. " - "'native' preserves the existing loader. " - "'rank_striped_read_ahead' lets node-local ranks read disjoint " - "SafeTensors extents while native mapping, materialization, and H2D " - "continue. Incompatible configurations select native I/O before " - "optimized reader or collective setup. Runtime-ineligible loads fall " - "back to native before model mutation.", - status="prototype", - json_schema_extra={ - "type": "Literal['auto', 'native', 'rank_striped_read_ahead']" - }, - ) + checkpoint_io_policy: CheckpointIoPolicy = Field( + default="auto", + description= + "Controls checkpoint storage I/O independently of checkpoint format. " + "'auto' selects rank-striped read-ahead for compatible built-in " + "PyTorch/HF loads and selects native I/O otherwise. " + "'native' preserves the existing loader. " + "'rank_striped_read_ahead' lets node-local ranks read disjoint " + "SafeTensors extents while native mapping, materialization, and H2D " + "continue. Incompatible configurations select native I/O before " + "optimized reader or collective setup. Runtime-ineligible loads fall " + "back to native before model mutation. " + "'demand_ordered_rank_striped_read_ahead' is an explicit experimental " + "variant that prioritizes complete checkpoint chunks using " + "node-local demand hints. It preserves all chunks, existing memory " + "admission, and native materialization; it does not enable " + "consumption-based pacing or destination-selective reads.", + status="prototype", + json_schema_extra={ + "type": str(CheckpointIoPolicy).replace("typing.", "") + }, + ) @property def is_partial_model_loading(self) -> bool: @@ -6541,11 +6546,11 @@ def warn_non_pytorch_checkpoint_io_policy_fallback(self) -> 'TorchLlmArgs': # requested policy for telemetry while reporting its native selection. # PyTorch requests are resolved at loader construction, where the actual # format and registered loader implementations are known. - if (self.checkpoint_io_policy == "rank_striped_read_ahead" + if (self.checkpoint_io_policy in RANK_STRIPED_IO_POLICIES and self.backend != "pytorch"): logger.warning( "Checkpoint I/O policy resolved before loading: " - "requested=rank_striped_read_ahead, selected=native, " + f"requested={self.checkpoint_io_policy}, selected=native, " "reason=rank-striped read-ahead requires the PyTorch backend.") return self diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e4f0bcc10f7e..0afea03dad89 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -169,7 +169,8 @@ "allowed_values": [ "auto", "native", - "rank_striped_read_ahead" + "rank_striped_read_ahead", + "demand_ordered_rank_striped_read_ahead" ], "capture_policy": "literal", "kind": "categorical", diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 85c16ae8d2bd..8a044f2f9164 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -152,6 +152,7 @@ l0_cpu: - unittest/quantization/test_mode.py - unittest/scaffolding - unittest/scripts + - unittest/test_checkpoint_io_policy.py - unittest/tools - unittest/usage - unittest/utils/test_logger.py diff --git a/tests/unittest/_torch/executor/test_model_loader_checkpoint_io.py b/tests/unittest/_torch/executor/test_model_loader_checkpoint_io.py index bdff42d60ce6..1387c314f548 100644 --- a/tests/unittest/_torch/executor/test_model_loader_checkpoint_io.py +++ b/tests/unittest/_torch/executor/test_model_loader_checkpoint_io.py @@ -35,6 +35,11 @@ pytestmark = pytest.mark.cpu_only +_RANK_STRIPED_POLICIES = ( + "rank_striped_read_ahead", + "demand_ordered_rank_striped_read_ahead", +) + def test_checkpoint_startup_metadata_describes_actual_loader() -> None: checkpoint_loader = HfCheckpointLoader() @@ -85,6 +90,24 @@ def test_construct_checkpoint_loader_selects_rank_striped_for_builtin_hf( assert status.selected == "rank_striped_read_ahead" +def test_demand_ordered_rank_striped_is_explicit_only() -> None: + demand_ordered = _construct_checkpoint_loader( + "pytorch", + None, + "HF", + checkpoint_io_policy="demand_ordered_rank_striped_read_ahead", + ) + automatic = _construct_checkpoint_loader("pytorch", None, "HF", checkpoint_io_policy="auto") + + assert demand_ordered.weight_loader.checkpoint_io_policy == ( + "demand_ordered_rank_striped_read_ahead" + ) + assert demand_ordered.weight_loader.last_checkpoint_io_status.selected == ( + "demand_ordered_rank_striped_read_ahead" + ) + assert automatic.weight_loader.checkpoint_io_policy == ("rank_striped_read_ahead") + + @pytest.mark.parametrize( ("kwargs", "reason"), [ @@ -93,10 +116,12 @@ def test_construct_checkpoint_loader_selects_rank_striped_for_builtin_hf( ({"partial_model_loading": True}, "partial model loading"), ], ) +@pytest.mark.parametrize("checkpoint_io_policy", _RANK_STRIPED_POLICIES) def test_construct_checkpoint_loader_falls_back_before_incompatible_request( monkeypatch: pytest.MonkeyPatch, kwargs: dict[str, Any], reason: str, + checkpoint_io_policy: str, ) -> None: warning = MagicMock() monkeypatch.setattr(model_loader_module.logger, "warning", warning) @@ -107,24 +132,26 @@ def test_construct_checkpoint_loader_falls_back_before_incompatible_request( "pytorch", None, checkpoint_format, - checkpoint_io_policy="rank_striped_read_ahead", + checkpoint_io_policy=checkpoint_io_policy, **options, ) assert loader.weight_loader.checkpoint_io_policy == "native" status = loader.weight_loader.last_checkpoint_io_status - assert status.requested == "rank_striped_read_ahead" + assert status.requested == checkpoint_io_policy assert status.selected == "native" assert reason in status.fallback_reason warning.assert_called_once() assert "selected=native" in warning.call_args.args[0] +@pytest.mark.parametrize("checkpoint_io_policy", _RANK_STRIPED_POLICIES) def test_construct_checkpoint_loader_preserves_explicit_loader_on_fallback( monkeypatch: pytest.MonkeyPatch, + checkpoint_io_policy: str, ) -> None: provided_loader = HfCheckpointLoader( - weight_loader=HfWeightLoader(checkpoint_io_policy="rank_striped_read_ahead") + weight_loader=HfWeightLoader(checkpoint_io_policy=checkpoint_io_policy) ) warning = MagicMock() monkeypatch.setattr(model_loader_module.logger, "warning", warning) @@ -133,13 +160,13 @@ def test_construct_checkpoint_loader_preserves_explicit_loader_on_fallback( "pytorch", provided_loader, "HF", - checkpoint_io_policy="rank_striped_read_ahead", + checkpoint_io_policy=checkpoint_io_policy, ) assert loader is provided_loader assert provided_loader.weight_loader.checkpoint_io_policy == "native" status = provided_loader.weight_loader.last_checkpoint_io_status - assert status.requested == "rank_striped_read_ahead" + assert status.requested == checkpoint_io_policy assert status.selected == "native" assert "explicit checkpoint loader" in status.fallback_reason warning.assert_called_once() @@ -177,14 +204,16 @@ def test_auto_selects_native_for_incompatible_config_without_warning( assert any("selected=native" in call.args[0] for call in info.call_args_list) +@pytest.mark.parametrize("checkpoint_io_policy", _RANK_STRIPED_POLICIES) def test_static_native_selection_skips_rank_striped_setup( monkeypatch: pytest.MonkeyPatch, + checkpoint_io_policy: str, ) -> None: loader = _construct_checkpoint_loader( "pytorch", None, "HF", - checkpoint_io_policy="rank_striped_read_ahead", + checkpoint_io_policy=checkpoint_io_policy, partial_model_loading=True, ) native_weights = {"native": object()} @@ -200,8 +229,10 @@ def test_static_native_selection_skips_rank_striped_setup( active_communicator.assert_not_called() +@pytest.mark.parametrize("checkpoint_io_policy", _RANK_STRIPED_POLICIES) def test_construct_checkpoint_loader_preserves_custom_hf_weight_loader( monkeypatch: pytest.MonkeyPatch, + checkpoint_io_policy: str, ) -> None: class _CustomWeightLoader: def __init__(self) -> None: @@ -224,15 +255,17 @@ def __init__(self) -> None: "pytorch", None, "HF", - checkpoint_io_policy="rank_striped_read_ahead", + checkpoint_io_policy=checkpoint_io_policy, ) assert isinstance(fallback_loader.weight_loader, _CustomWeightLoader) warning.assert_called_once() assert "selected=native" in warning.call_args.args[0] +@pytest.mark.parametrize("checkpoint_io_policy", _RANK_STRIPED_POLICIES) def test_construct_checkpoint_loader_detects_custom_hf_wrapper( monkeypatch: pytest.MonkeyPatch, + checkpoint_io_policy: str, ) -> None: class _CustomCheckpointLoader(HfCheckpointLoader): pass @@ -249,7 +282,7 @@ class _CustomCheckpointLoader(HfCheckpointLoader): "pytorch", None, "HF", - checkpoint_io_policy="rank_striped_read_ahead", + checkpoint_io_policy=checkpoint_io_policy, ) assert isinstance(loader, _CustomCheckpointLoader) @@ -270,6 +303,8 @@ def __init__( self.events = events self.checkpoint_dir = checkpoint_dir self.weights = {"model.weight": object()} if weights is None else weights + self.initialized_weight_mapper = object() + self.activation_mappers: list[object | None] = [] @contextmanager def open_weight_session( @@ -289,7 +324,18 @@ def is_weights_preloaded(self) -> bool: def get_initialized_weight_mapper(self, model: object, config: object) -> object: del model, config self.events.append("mapper_init") - return object() + return self.initialized_weight_mapper + + def activate_weight_session( + self, + readiness_error: BaseException | None = None, + *, + weight_mapper: object | None = None, + ) -> None: + self.activation_mappers.append(weight_mapper) + self.events.append("session_activate") + if readiness_error is not None: + raise readiness_error def test_legacy_weight_session_defers_load_until_enter() -> None: @@ -366,12 +412,70 @@ def test_model_loader_session_spans_mapper_and_materialization() -> None: assert events == [ "session_enter", "mapper_init", + "session_activate", "materialize", "session_exit", ] assert "checkpoint_preparation_seconds" in loader.metrics assert "weight_population_seconds" in loader.metrics assert "checkpoint_finalization_seconds" in loader.metrics + assert checkpoint_loader.activation_mappers == [checkpoint_loader.initialized_weight_mapper] + + +@pytest.mark.parametrize("draft", [False, True]) +@pytest.mark.parametrize("policy", ("native", *_RANK_STRIPED_POLICIES)) +def test_checkpoint_preparation_times_only_demand_ordered_activation( + monkeypatch: pytest.MonkeyPatch, + draft: bool, + policy: str, +) -> None: + clock = [0.0] + monkeypatch.setattr(modeling_utils.time, "perf_counter", lambda: clock[0]) + monkeypatch.setattr(model_loader_module.torch.cuda, "is_available", lambda: False) + checkpoint_loader = _SessionCheckpointLoader([], "/checkpoint") + checkpoint_loader.weight_loader = SimpleNamespace(checkpoint_io_policy=policy) + loader = ModelLoader.__new__(ModelLoader) + loader._metrics = {} + loader.mapping = object() + loader.spec_config = SimpleNamespace(speculative_model="/checkpoint") + loader.weight_mapper = object() + + def init_mapper(*_args: object) -> object: + clock[0] += 2.0 + return loader.weight_mapper + + def activate(*_args: object, **_kwargs: object) -> None: + clock[0] += 3.0 + + def inspect_shadow(*_args: object, **_kwargs: object) -> None: + clock[0] += 5.0 + + def materialize(*_args: object, **_kwargs: object) -> None: + clock[0] += 7.0 + + checkpoint_loader.get_initialized_weight_mapper = init_mapper + checkpoint_loader.activate_weight_session = activate + monkeypatch.setattr(model_loader_module, "_inspect_shadow_weight_load_plan", inspect_shadow) + loader._call_load_weights = materialize + model = MagicMock() + if draft: + model.draft_config = SimpleNamespace( + pretrained_config=SimpleNamespace(architectures=["DraftForCausalLM"]) + ) + mapper = SimpleNamespace(init_model_and_config=init_mapper) + monkeypatch.setattr(model_loader_module.AutoCheckpointMapper, "get", lambda *_args: mapper) + loader._materialize_draft_checkpoint_weights(checkpoint_loader, model) + else: + loader._materialize_checkpoint_weights( + checkpoint_loader, "/checkpoint", model, object(), {"mapping": object()} + ) + + prefix = "draft_" if draft else "" + expected_preparation = 3.0 if policy == "demand_ordered_rank_striped_read_ahead" else 0.0 + assert loader.metrics[f"{prefix}checkpoint_preparation_seconds"] == expected_preparation + assert loader.metrics[f"{prefix}weight_population_seconds"] == 7.0 + assert loader.metrics[f"{prefix}checkpoint_finalization_seconds"] == 0.0 + assert clock[0] == 17.0 def test_shadow_plan_disabled_does_not_call_inspection_hooks( @@ -390,6 +494,21 @@ def test_shadow_plan_disabled_does_not_call_inspection_hooks( weight_mapper.build_weight_load_plan.assert_not_called() +def test_preloaded_custom_loader_does_not_require_weight_mapper() -> None: + checkpoint_loader = _SessionCheckpointLoader([], weights={}) + checkpoint_loader.is_weights_preloaded = lambda: True + checkpoint_loader.get_initialized_weight_mapper = lambda *_args: None + loader = ModelLoader.__new__(ModelLoader) + loader._metrics = {} + loader._call_load_weights = MagicMock() + + assert loader._materialize_checkpoint_weights( + checkpoint_loader, "/checkpoint", MagicMock(), object(), {"mapping": object()} + ) + loader._call_load_weights.assert_not_called() + assert checkpoint_loader.activation_mappers == [None] + + def test_shadow_plan_is_advisory_and_preserves_materialization_order( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -406,15 +525,13 @@ def test_shadow_plan_is_advisory_and_preserves_materialization_order( events = [] checkpoint_loader = _SessionCheckpointLoader(events) checkpoint_loader.build_checkpoint_catalog = MagicMock( - side_effect=lambda *_args, **_kwargs: (events.append("catalog") or catalog) + 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) - ) + 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) + side_effect=lambda *_args: events.append("mapper_init") or weight_mapper ) model = MagicMock() loader = ModelLoader.__new__(ModelLoader) @@ -436,6 +553,7 @@ def test_shadow_plan_is_advisory_and_preserves_materialization_order( "mapper_init", "catalog", "plan", + "session_activate", "materialize", "session_exit", ] @@ -477,6 +595,99 @@ def test_shadow_plan_failure_warns_and_continues( weight_mapper.build_weight_load_plan.assert_not_called() +@pytest.mark.parametrize("draft", [False, True]) +def test_mapper_failure_reaches_activation_before_session_finish( + monkeypatch: pytest.MonkeyPatch, + draft: bool, +) -> None: + events = [] + checkpoint_loader = _SessionCheckpointLoader(events) + loader = ModelLoader.__new__(ModelLoader) + loader._metrics = {} + loader.mapping = object() + loader.spec_config = SimpleNamespace(speculative_model="/checkpoint") + loader._call_load_weights = MagicMock() + mapper_error = ValueError("mapper failed") + + def fail_mapper_init(*_args: object) -> None: + events.append("mapper_init") + raise mapper_error + + checkpoint_loader.get_initialized_weight_mapper = fail_mapper_init + monkeypatch.setattr(model_loader_module.AutoCheckpointMapper, "get", fail_mapper_init) + + with pytest.raises(ValueError, match="mapper failed") as raised: + if draft: + loader._materialize_draft_checkpoint_weights(checkpoint_loader, MagicMock()) + else: + loader._materialize_checkpoint_weights( + checkpoint_loader, "/checkpoint", MagicMock(), object(), {"mapping": object()} + ) + + assert raised.value is mapper_error + assert events == [ + "session_enter", + "mapper_init", + "session_activate", + "session_exit", + ] + loader._call_load_weights.assert_not_called() + + +@pytest.mark.parametrize("local_failure", [False, True]) +def test_preload_readiness_failure_is_coordinated_before_activation_start( + local_failure: bool, +) -> None: + communicator = MagicMock() + communicator.Get_size.return_value = 2 + communicator.allgather.return_value = [None, "OSError: peer preload status failed"] + + peer_error = OSError("peer preload status failed") + events = [] + checkpoint_loader = _SessionCheckpointLoader(events) + loader = ModelLoader.__new__(ModelLoader) + loader._metrics = {} + loader._call_load_weights = MagicMock() + demand_ordered_loader = HfWeightLoader( + checkpoint_io_policy="demand_ordered_rank_striped_read_ahead" + ) + session = MagicMock(active_communicator=communicator) + demand_ordered_loader._pending_read_ahead_session = session + + def check_preloaded() -> bool: + events.append("preload_check") + if local_failure: + raise peer_error + return False + + def activate( + readiness_error: BaseException | None = None, + *, + weight_mapper: object | None = None, + ) -> None: + events.append("session_activate") + demand_ordered_loader.activate_weight_session(readiness_error, weight_mapper=weight_mapper) + + checkpoint_loader.is_weights_preloaded = check_preloaded + checkpoint_loader.activate_weight_session = activate + + with pytest.raises((RuntimeError, OSError)) as raised: + loader._materialize_checkpoint_weights( + checkpoint_loader, "/checkpoint", MagicMock(), object(), {"mapping": object()} + ) + + expected_events = ["session_enter", "preload_check"] + if not local_failure: + expected_events.append("mapper_init") + assert events == expected_events + ["session_activate", "session_exit"] + assert not demand_ordered_loader.last_checkpoint_io_status.activated + assert (raised.value is peer_error) if local_failure else isinstance(raised.value, RuntimeError) + communicator.allgather.assert_called_once_with( + "OSError: peer preload status failed" if local_failure else None + ) + loader._call_load_weights.assert_not_called() + + def test_draft_session_spans_mapper_and_materialization( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -493,7 +704,7 @@ 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) + side_effect=lambda *_args, **_kwargs: events.append("catalog") or catalog ) model = MagicMock() model.draft_config = SimpleNamespace( @@ -508,9 +719,7 @@ def test_draft_session_spans_mapper_and_materialization( ) draft_mapper = MagicMock() draft_mapper.init_model_and_config.side_effect = lambda *_args: events.append("mapper_init") - draft_mapper.build_weight_load_plan.side_effect = lambda _catalog: ( - events.append("plan") or plan - ) + 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) ) @@ -522,12 +731,14 @@ def test_draft_session_spans_mapper_and_materialization( "mapper_init", "catalog", "plan", + "session_activate", "materialize", "session_exit", ] 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 checkpoint_loader.activation_mappers == [draft_mapper] def test_mtp_draft_session_reuses_target_mapper_during_materialization( @@ -546,7 +757,7 @@ def test_mtp_draft_session_reuses_target_mapper_during_materialization( demands=(WeightDemand("all", ("draft.weight",), (0,)),), ) checkpoint_loader.build_checkpoint_catalog = MagicMock( - side_effect=lambda *_args, **_kwargs: (events.append("catalog") or catalog) + side_effect=lambda *_args, **_kwargs: events.append("catalog") or catalog ) model = MagicMock() model.draft_config = None @@ -555,9 +766,7 @@ def test_mtp_draft_session_reuses_target_mapper_during_materialization( loader.spec_config = SimpleNamespace(speculative_model="/draft") loader.mapping = object() loader.weight_mapper = SimpleNamespace( - build_weight_load_plan=MagicMock( - side_effect=lambda _catalog: (events.append("plan") or plan) - ) + 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") @@ -570,6 +779,7 @@ def test_mtp_draft_session_reuses_target_mapper_during_materialization( "session_enter", "catalog", "plan", + "session_activate", "materialize", "session_exit", ] 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 58fcdcdf2455..6b22b86d8ad0 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 @@ -3,7 +3,9 @@ import pickle import sys +import tempfile import threading +import time from pathlib import Path from unittest import mock @@ -13,9 +15,21 @@ from mpi4py import MPI from safetensors.torch import save_file +from tensorrt_llm._torch.models.checkpoints.checkpoint_catalog import ( + CheckpointCatalog, + CheckpointExtent, + CheckpointObject, + CheckpointTensor, +) from tensorrt_llm._torch.models.checkpoints.hf import rank_striped_read_ahead as read_ahead from tensorrt_llm._torch.models.checkpoints.hf import weight_loader as weight_loader_module from tensorrt_llm._torch.models.checkpoints.hf.weight_loader import HfWeightLoader +from tensorrt_llm._torch.models.checkpoints.weight_load_plan import ( + WeightDemand, + WeightLoadOrderConfidence, + WeightLoadPlan, + WeightLoadPlanCoverage, +) from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.cpu_only @@ -32,6 +46,7 @@ def _isolated_loader_environment(monkeypatch): monkeypatch.delenv("TRTLLM_HF_WEIGHT_CACHE", raising=False) monkeypatch.delenv("TLLM_OVERRIDE_LAYER_NUM", raising=False) + monkeypatch.delenv("TRTLLM_DEMAND_ORDERED_READ_AHEAD_PHYSICAL_ORDER", raising=False) monkeypatch.setattr(weight_loader_module, "ENABLE_MULTI_DEVICE", False) monkeypatch.setattr(weight_loader_module, "effective_available_host_memory", lambda: 1 << 40) @@ -54,6 +69,43 @@ def _rank_striped_loader(*, partial_model_loading: bool = False) -> HfWeightLoad ) +def _demand_ordered_rank_striped_loader(*, partial_model_loading: bool = False) -> HfWeightLoader: + return HfWeightLoader( + checkpoint_io_policy="demand_ordered_rank_striped_read_ahead", + partial_model_loading=partial_model_loading, + ) + + +class _TestWeightMapper: + def __init__(self, mapping: Mapping, *, fail: bool = False) -> None: + self._mapping = mapping + self._fail = fail + + def build_weight_load_plan(self, catalog: CheckpointCatalog) -> WeightLoadPlan: + if self._fail: + raise RuntimeError("injected weight-plan failure") + plan = WeightLoadPlan( + catalog_id=catalog.catalog_id, + rank=self._mapping.rank, + world_size=self._mapping.world_size, + coverage=WeightLoadPlanCoverage.CONSERVATIVE, + ordering=WeightLoadOrderConfidence.ADVISORY, + demands=( + WeightDemand( + group_id="all_checkpoint_tensors", + source_names=tuple(tensor.name for tensor in catalog.tensors), + destination_ranks=(self._mapping.rank,), + ), + ), + ) + plan.validate_against(catalog) + return plan + + +def _test_weight_mapper(mapping: Mapping | None = None, *, fail: bool = False) -> _TestWeightMapper: + return _TestWeightMapper(mapping if mapping is not None else Mapping(), fail=fail) + + def test_subclass_without_super_init_uses_native_compatibility_defaults( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -163,6 +215,316 @@ def test_extent_plan_is_complete_disjoint_and_fair(tmp_path, monkeypatch): assert max(workers) - min(workers) == 1 +def test_ordering_preserves_complete_fair_stripes(tmp_path, monkeypatch): + monkeypatch.setattr(read_ahead, "_CHUNK_SIZE", 4) + monkeypatch.setattr(read_ahead, "_READ_SIZE", 1) + monkeypatch.setattr(read_ahead, "_WORKERS_PER_LOAD_GROUP", 3) + monkeypatch.setattr(read_ahead, "_WORKERS_PER_RANK", 2) + files = [(str(tmp_path / "a"), 10), (str(tmp_path / "b"), 7)] + source_order = [ + read_ahead.ReadAheadExtent(str(tmp_path / "a"), 0, 4), + read_ahead.ReadAheadExtent(str(tmp_path / "a"), 4, 4), + read_ahead.ReadAheadExtent(str(tmp_path / "a"), 8, 2), + read_ahead.ReadAheadExtent(str(tmp_path / "b"), 0, 4), + read_ahead.ReadAheadExtent(str(tmp_path / "b"), 4, 3), + ] + plans = [ + read_ahead.build_local_plan( + files, + rank, + 4, + ordered_extents=list(reversed(source_order)), + ) + for rank in range(4) + ] + + assert [plan.workers for plan in plans] == [1, 1, 1, 0] + assert [extent for plan in plans for extent in plan.extents] + assert set(extent for plan in plans for extent in plan.extents) == set(source_order) + assert plans[0].extents[0] == source_order[-1] + + with pytest.raises(ValueError, match="complete permutation"): + read_ahead.build_local_plan( + files, + 0, + 4, + ordered_extents=source_order[:-1], + ) + + +def test_weight_plans_compile_to_a_complete_node_demand_order(tmp_path, monkeypatch): + monkeypatch.setattr(read_ahead, "_CHUNK_SIZE", 4) + paths = (tmp_path / "a.safetensors", tmp_path / "b.safetensors") + files = ((str(paths[0]), 12), (str(paths[1]), 8)) + catalog = CheckpointCatalog( + objects=( + CheckpointObject("a.safetensors", 12), + CheckpointObject("b.safetensors", 8), + ), + tensors=( + CheckpointTensor("late", extents=(CheckpointExtent("a.safetensors", 0, 4),)), + CheckpointTensor("middle", extents=(CheckpointExtent("a.safetensors", 8, 4),)), + CheckpointTensor("early", extents=(CheckpointExtent("b.safetensors", 4, 4),)), + ), + ) + + def make_plan(rank, names, ordering=WeightLoadOrderConfidence.ADVISORY): + return WeightLoadPlan( + catalog_id=catalog.catalog_id, + rank=rank, + world_size=2, + coverage=WeightLoadPlanCoverage.CONSERVATIVE, + ordering=ordering, + demands=tuple( + WeightDemand( + # Lexical ordering intentionally differs from tuple order. + group_id=f"layer.{10 - index}", + source_names=(name,), + destination_ranks=(rank,), + priority=index, + ) + for index, name in enumerate(names) + ), + ) + + ordered = read_ahead.compile_weight_plan_extent_order( + catalog, + ( + make_plan(0, ("late", "middle", "early")), + make_plan(1, ("early", "middle", "late")), + ), + files, + ) + + assert ordered == ( + read_ahead.ReadAheadExtent(str(paths[0]), 0, 4), + read_ahead.ReadAheadExtent(str(paths[1]), 4, 4), + read_ahead.ReadAheadExtent(str(paths[0]), 8, 4), + read_ahead.ReadAheadExtent(str(paths[0]), 4, 4), + read_ahead.ReadAheadExtent(str(paths[1]), 0, 4), + ) + + physical = read_ahead.compile_weight_plan_extent_order( + catalog, + ( + make_plan( + 0, + ("early", "middle", "late"), + WeightLoadOrderConfidence.OPAQUE, + ), + make_plan( + 1, + ("late", "middle", "early"), + WeightLoadOrderConfidence.OPAQUE, + ), + ), + files, + ) + assert physical == tuple(read_ahead._build_source_extents(files)) + + dependency_plan = WeightLoadPlan( + catalog_id=catalog.catalog_id, + rank=0, + world_size=1, + coverage=WeightLoadPlanCoverage.CONSERVATIVE, + ordering=WeightLoadOrderConfidence.ADVISORY, + demands=( + WeightDemand( + group_id="dependency", + source_names=("late", "middle"), + destination_ranks=(0,), + priority=100, + ), + WeightDemand( + group_id="urgent-dependent", + source_names=("early",), + destination_ranks=(0,), + priority=-100, + predecessors=("dependency",), + ), + ), + ) + predecessor_order = read_ahead.compile_weight_plan_extent_order( + catalog, (dependency_plan,), files + ) + assert predecessor_order[:3] == ( + read_ahead.ReadAheadExtent(str(paths[0]), 0, 4), + read_ahead.ReadAheadExtent(str(paths[0]), 8, 4), + read_ahead.ReadAheadExtent(str(paths[1]), 4, 4), + ) + + +def test_weight_plan_compilation_is_canonical_for_equal_plan_ids(tmp_path, monkeypatch): + monkeypatch.setattr(read_ahead, "_CHUNK_SIZE", 4) + path = tmp_path / "model.safetensors" + files = ((str(path), 12),) + catalog = CheckpointCatalog( + objects=(CheckpointObject(path.name, 12),), + tensors=( + CheckpointTensor("a", extents=(CheckpointExtent(path.name, 0, 4),)), + CheckpointTensor("b", extents=(CheckpointExtent(path.name, 4, 4),)), + CheckpointTensor("c", extents=(CheckpointExtent(path.name, 8, 4),)), + ), + ) + first = WeightDemand("first", ("b", "a"), (0,), priority=0) + second = WeightDemand("second", ("c",), (0,), priority=0) + + def make_plan(demands): + return WeightLoadPlan( + catalog_id=catalog.catalog_id, + rank=0, + world_size=1, + coverage=WeightLoadPlanCoverage.CONSERVATIVE, + ordering=WeightLoadOrderConfidence.ADVISORY, + demands=demands, + ) + + forward = make_plan((first, second)) + reordered = make_plan( + ( + second, + WeightDemand("first", ("a", "b"), (0,), priority=0), + ) + ) + + assert forward.plan_id == reordered.plan_id + assert read_ahead.compile_weight_plan_extent_order( + catalog, (forward,), files + ) == read_ahead.compile_weight_plan_extent_order(catalog, (reordered,), files) + + # Equal-priority independent groups retain physical locality, regardless + # of their group IDs, input order, or source-name ordering. + assert read_ahead.compile_weight_plan_extent_order(catalog, (forward,), files) == tuple( + read_ahead._build_source_extents(files) + ) + + +def test_physical_order_ablation_still_compiles_demand_plan(tmp_path, monkeypatch): + monkeypatch.setattr(read_ahead, "_CHUNK_SIZE", 4) + path = tmp_path / "model.safetensors" + path.write_bytes(b"abcdefgh") + files = ((str(path), 8),) + catalog = CheckpointCatalog( + objects=(CheckpointObject(path.name, 8),), + tensors=(CheckpointTensor("early", extents=(CheckpointExtent(path.name, 4, 4),)),), + ) + demand_plan = WeightLoadPlan( + catalog_id=catalog.catalog_id, + rank=0, + world_size=1, + coverage=WeightLoadPlanCoverage.CONSERVATIVE, + ordering=WeightLoadOrderConfidence.ADVISORY, + demands=(WeightDemand("early", ("early",), (0,)),), + ) + original_plan = read_ahead.build_local_plan(files, 0, 1) + compile_plan = mock.Mock(wraps=read_ahead.compile_weight_plan_extent_order) + monkeypatch.setattr(weight_loader_module, "compile_weight_plan_extent_order", compile_plan) + for physical_order in (False, True): + monkeypatch.setenv( + "TRTLLM_DEMAND_ORDERED_READ_AHEAD_PHYSICAL_ORDER", str(int(physical_order)) + ) + loader = _demand_ordered_rank_striped_loader() + context = weight_loader_module._DemandReadAheadContext( + None, None, catalog, list(files), 0, 1 + ) + loader._pending_read_ahead_session = context + loader.activate_weight_session( + weight_mapper=mock.Mock(build_weight_load_plan=mock.Mock(return_value=demand_plan)) + ) + assert context.reader is not None + assert context.reader._plan.extents == ( + original_plan.extents if physical_order else tuple(reversed(original_plan.extents)) + ) + assert context.reader.cancel_and_close() is None + assert compile_plan.call_count == 2 + + +def test_read_ahead_handles_short_reads(tmp_path, monkeypatch): + path = tmp_path / "weights" + path.write_bytes(b"0123456789" * 20) + monkeypatch.setattr(read_ahead, "_READ_SIZE", 40) + extent = read_ahead.ReadAheadExtent(str(path), 0, 200) + real_pread = read_ahead.os.pread + + def short_pread(file_descriptor, size, offset): + return real_pread(file_descriptor, max(1, size // 2), offset) + + monkeypatch.setattr(read_ahead.os, "pread", short_pread) + session = read_ahead.RankStripedReadAheadSession( + None, None, read_ahead.ReadAheadPlan((extent,), 1) + ) + try: + assert session._read_extent(extent) == extent.length + finally: + assert session.cancel_and_close() is None + + +def test_deferred_session_issues_no_payload_io_before_release(tmp_path, monkeypatch): + path = tmp_path / "weights" + path.write_bytes(b"x") + plan = read_ahead.ReadAheadPlan( + (read_ahead.ReadAheadExtent(str(path), 0, 1),), + workers=1, + ) + pread = mock.Mock(return_value=b"x") + monkeypatch.setattr(read_ahead.os, "pread", pread) + session = read_ahead.RankStripedReadAheadSession(None, None, plan).start(defer_reads=True) + + time.sleep(0.01) + pread.assert_not_called() + session.release_reads() + assert session._thread is not None + session._thread.join(timeout=5) + pread.assert_called_once() + assert session.cancel_and_close() is None + + +def test_reader_open_failure_closes_already_opened_descriptors( + tmp_path, + monkeypatch, +): + opened_path = tmp_path / "opened" + failing_path = tmp_path / "failing" + for path in (opened_path, failing_path): + path.write_bytes(b"x") + failed_plan = read_ahead.ReadAheadPlan( + ( + read_ahead.ReadAheadExtent(str(opened_path), 0, 1), + read_ahead.ReadAheadExtent(str(failing_path), 0, 1), + ), + workers=1, + ) + real_open = read_ahead.os.open + real_close = read_ahead.os.close + newly_opened_descriptors = [] + closed_descriptors = [] + + def fail_second_open(path, flags, *args, **kwargs): + if path == str(failing_path): + raise OSError("injected open failure") + descriptor = real_open(path, flags, *args, **kwargs) + if path == str(opened_path): + newly_opened_descriptors.append(descriptor) + return descriptor + + def record_close(descriptor): + closed_descriptors.append(descriptor) + return real_close(descriptor) + + # os is process-wide: preserve unrelated opens and restore it before pytest + # resumes capture/fixture teardown. + with monkeypatch.context() as patch: + patch.setattr(read_ahead.os, "open", fail_second_open) + patch.setattr(read_ahead.os, "close", record_close) + + with pytest.raises(OSError, match="injected open failure"): + read_ahead.RankStripedReadAheadSession(None, None, failed_plan) + assert newly_opened_descriptors + assert newly_opened_descriptors[0] in closed_descriptors + with tempfile.TemporaryFile() as capture_file: + capture_file.write(b"unrelated capture I/O") + + @pytest.mark.parametrize( ("cgroup", "expected"), [ @@ -327,7 +689,10 @@ def test_collective_native_fallback_failure_escapes_without_cleanup( close_node_communicator.assert_not_called() -@pytest.mark.parametrize("requested_policy", ["auto", "rank_striped_read_ahead"]) +@pytest.mark.parametrize( + "requested_policy", + ["auto", "rank_striped_read_ahead", "demand_ordered_rank_striped_read_ahead"], +) @pytest.mark.parametrize("partial,available", [(True, 1 << 40), (False, 0)]) def test_ineligible_request_uses_exact_native_fallback( tmp_path: Path, @@ -338,7 +703,9 @@ def test_ineligible_request_uses_exact_native_fallback( ) -> None: checkpoint_dir, _ = _write_checkpoint(tmp_path) loader = HfWeightLoader( - checkpoint_io_policy="rank_striped_read_ahead", + checkpoint_io_policy=( + "rank_striped_read_ahead" if requested_policy == "auto" else requested_policy + ), requested_checkpoint_io_policy=requested_policy, partial_model_loading=partial, ) @@ -363,7 +730,7 @@ def test_ineligible_request_uses_exact_native_fallback( assert status.effective == "native" assert status.fallback_reason reader_start.assert_not_called() - if requested_policy == "rank_striped_read_ahead": + if requested_policy != "auto": warning.assert_called_once() assert "falling back before model materialization" in warning.call_args.args[0] else: @@ -386,11 +753,27 @@ def test_disabled_weight_cache_does_not_force_fallback(tmp_path, monkeypatch): assert loader.last_checkpoint_io_status.activated -def test_reader_start_failure_cleans_up_and_falls_back(tmp_path, monkeypatch): +def test_demand_ordered_policy_requires_full_checkpoint_memory(tmp_path, monkeypatch): checkpoint_dir, _ = _write_checkpoint(tmp_path) + monkeypatch.setattr(read_ahead, "_HOST_MEMORY_HEADROOM_BYTES", 0) + monkeypatch.setattr(read_ahead, "_HOST_MEMORY_HEADROOM_FRACTION", 0) + monkeypatch.setattr(weight_loader_module, "effective_available_host_memory", lambda: 1) + loader = _demand_ordered_rank_striped_loader() + start = mock.Mock(side_effect=AssertionError("must not start")) + monkeypatch.setattr(read_ahead.RankStripedReadAheadSession, "start", start) + + with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()): + pass + + assert loader.last_checkpoint_io_status.effective == "native" + assert not loader.last_checkpoint_io_status.activated + start.assert_not_called() + + +def test_reader_start_failure_cleans_up_and_degrades_without_reload(tmp_path, monkeypatch): + checkpoint_dir, expected = _write_checkpoint(tmp_path) loader = _rank_striped_loader() - native_weights = {"native": object()} - native_load = mock.Mock(return_value=native_weights) + native_load = mock.Mock(side_effect=AssertionError("must not reload")) monkeypatch.setattr(loader, "_load_weights_native", native_load) close_file_descriptors = [] original_close = read_ahead.RankStripedReadAheadSession._close_file_descriptors @@ -413,9 +796,9 @@ def record_close(session): with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()) as weights: pass - assert weights is native_weights + assert torch.equal(weights["model.norm.weight"], expected["model.norm.weight"]) assert close_file_descriptors - native_load.assert_called_once() + native_load.assert_not_called() status = loader.last_checkpoint_io_status assert status.selected == "rank_striped_read_ahead" assert not status.activated @@ -423,6 +806,240 @@ def record_close(session): assert "start failed" in status.fallback_reason +def test_demand_ordered_reader_defers_activation_until_mapper_boundary(tmp_path, monkeypatch): + checkpoint_dir, _ = _write_checkpoint(tmp_path) + loader = _demand_ordered_rank_striped_loader() + started = threading.Event() + + def wait_for_cancel(session, _extent): + started.set() + assert session._cancel.wait(timeout=5) + return 0 + + monkeypatch.setattr(read_ahead.RankStripedReadAheadSession, "_read_extent", wait_for_cancel) + reader_factory = mock.Mock(wraps=read_ahead.RankStripedReadAheadSession) + monkeypatch.setattr(weight_loader_module, "RankStripedReadAheadSession", reader_factory) + + with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()): + reader_factory.assert_not_called() + assert not started.is_set() + assert not loader.last_checkpoint_io_status.activated + loader.activate_weight_session(weight_mapper=_test_weight_mapper()) + assert started.wait(timeout=5) + assert loader.last_checkpoint_io_status.activated + reader_factory.assert_called_once() + + assert loader.last_checkpoint_io_status.effective == ("demand_ordered_rank_striped_read_ahead") + + +@pytest.mark.parametrize("policy", ["native", "rank_striped_read_ahead"]) +def test_existing_policies_do_not_plan_or_restart_on_activation(policy, monkeypatch): + loader = HfWeightLoader(checkpoint_io_policy=policy) + session = mock.Mock() + mapper = mock.Mock() + loader._pending_read_ahead_session = session + consensus = mock.Mock(side_effect=AssertionError("must not coordinate")) + monkeypatch.setattr(weight_loader_module, "coordinate_error", consensus) + + loader.activate_weight_session(weight_mapper=mapper) + + mapper.build_weight_load_plan.assert_not_called() + session.start.assert_not_called() + consensus.assert_not_called() + + +def test_mapper_readiness_failure_is_rank_coherent_before_start(): + communicator = mock.Mock() + communicator.Get_size.return_value = 2 + communicator.allgather.return_value = [None, "ValueError: peer mapper failed"] + local_errors = (None, ValueError("peer mapper failed")) + raised_errors = [] + + for local_error in local_errors: + loader = _demand_ordered_rank_striped_loader() + session = mock.Mock() + session.active_communicator = communicator + loader._pending_read_ahead_session = session + + with pytest.raises((RuntimeError, ValueError)) as raised: + loader.activate_weight_session(local_error) + raised_errors.append(raised.value) + session.start.assert_not_called() + assert not loader.last_checkpoint_io_status.activated + assert "peer mapper failed" in loader.last_checkpoint_io_status.fallback_reason + + assert isinstance(raised_errors[0], RuntimeError) + assert raised_errors[1] is local_errors[1] + assert communicator.allgather.call_args_list == [ + mock.call(None), + mock.call("ValueError: peer mapper failed"), + ] + + +def test_demand_ordered_activation_orders_consensus_before_node_plan_exchange(monkeypatch): + events = [] + catalog = CheckpointCatalog( + objects=(CheckpointObject("model.safetensors", 1),), + tensors=( + CheckpointTensor("weight", extents=(CheckpointExtent("model.safetensors", 0, 1),)), + ), + ) + plan = WeightLoadPlan( + catalog_id=catalog.catalog_id, + rank=0, + world_size=1, + coverage=WeightLoadPlanCoverage.CONSERVATIVE, + ordering=WeightLoadOrderConfidence.ADVISORY, + demands=(WeightDemand("weight", ("weight",), (0,)),), + ) + compiled_plan = read_ahead.ReadAheadPlan((), 0) + mapper = mock.Mock() + mapper.build_weight_load_plan.side_effect = lambda _catalog: (events.append("plan") or plan) + node_communicator = mock.Mock() + node_communicator.allgather.side_effect = lambda local_input: ( + events.append("node-allgather") or [local_input] + ) + context = weight_loader_module._DemandReadAheadContext( + None, node_communicator, catalog, [], 0, 1 + ) + compile_plan = mock.Mock(side_effect=lambda *_args: (events.append("compile") or compiled_plan)) + session = mock.Mock() + reader_factory = mock.Mock(side_effect=lambda *_args: (events.append("construct") or session)) + session.start.side_effect = lambda **_kwargs: events.append("start") + session.release_reads.side_effect = lambda: events.append("release") + + def record_consensus(_communicator, phase, error): + assert error is None + events.append(phase) + return None + + monkeypatch.setattr(weight_loader_module, "coordinate_error", record_consensus) + monkeypatch.setattr(weight_loader_module, "compile_weight_plan_extent_order", compile_plan) + monkeypatch.setattr( + weight_loader_module, "build_local_plan", mock.Mock(return_value=compiled_plan) + ) + monkeypatch.setattr(weight_loader_module, "RankStripedReadAheadSession", reader_factory) + loader = _demand_ordered_rank_striped_loader() + loader._pending_read_ahead_session = context + + loader.activate_weight_session(weight_mapper=mapper) + + assert events == [ + "rank-striped mapper readiness", + "plan", + "demand-ordered weight plan build", + "node-allgather", + "compile", + "construct", + "demand-ordered reader setup", + "start", + "rank-striped reader activation", + "release", + ] + + +def test_physical_order_ablation_mismatch_disables_all_readers(monkeypatch): + monkeypatch.delenv("TRTLLM_DEMAND_ORDERED_READ_AHEAD_PHYSICAL_ORDER", raising=False) + catalog = CheckpointCatalog(objects=(), tensors=(CheckpointTensor("weight"),)) + mapper = _test_weight_mapper() + plan = mapper.build_weight_load_plan(catalog) + node_communicator = mock.Mock() + node_communicator.allgather.return_value = [(plan, False), (plan, True)] + context = weight_loader_module._DemandReadAheadContext( + None, node_communicator, catalog, [], 0, 1 + ) + loader = _demand_ordered_rank_striped_loader() + loader._pending_read_ahead_session = context + + loader.activate_weight_session(weight_mapper=mapper) + + assert context.reader is None + assert not loader.last_checkpoint_io_status.activated + assert "must agree across node-local ranks" in loader.last_checkpoint_io_status.fallback_reason + + +def test_unactivated_demand_session_reports_native(tmp_path): + checkpoint_dir, expected = _write_checkpoint(tmp_path) + loader = _demand_ordered_rank_striped_loader() + + with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()) as weights: + assert torch.equal(weights["model.norm.weight"], expected["model.norm.weight"]) + + status = loader.last_checkpoint_io_status + assert not status.activated + assert status.effective == "native" + assert "not activated" in status.fallback_reason + + +def test_demand_ordered_activation_failure_degrades_collectively_without_reload( + tmp_path, monkeypatch +): + checkpoint_dir, _ = _write_checkpoint(tmp_path) + loader = _demand_ordered_rank_striped_loader() + native_load = mock.Mock(side_effect=AssertionError("must not reload")) + monkeypatch.setattr(loader, "_load_weights_native", native_load) + monkeypatch.setattr( + read_ahead.RankStripedReadAheadSession, + "start", + mock.Mock(side_effect=RuntimeError("activation failed")), + ) + + with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()): + loader.activate_weight_session(weight_mapper=_test_weight_mapper()) + assert not loader.last_checkpoint_io_status.activated + + native_load.assert_not_called() + status = loader.last_checkpoint_io_status + assert status.effective == "native" + assert "activation failed" in status.fallback_reason + + +@pytest.mark.parametrize( + ("failure_phase", "failure_message"), + [ + ("plan", "injected weight-plan failure"), + ("compiler", "injected compiler failure"), + ("open", "injected descriptor-open failure"), + ], +) +def test_demand_ordered_planning_failure_degrades_before_start_without_reload( + tmp_path, + monkeypatch, + failure_phase, + failure_message, +): + checkpoint_dir, _ = _write_checkpoint(tmp_path) + loader = _demand_ordered_rank_striped_loader() + native_load = mock.Mock(side_effect=AssertionError("must not reload")) + reader_start = mock.Mock(side_effect=AssertionError("must not start")) + monkeypatch.setattr(loader, "_load_weights_native", native_load) + monkeypatch.setattr(read_ahead.RankStripedReadAheadSession, "start", reader_start) + if failure_phase == "compiler": + monkeypatch.setattr( + weight_loader_module, + "compile_weight_plan_extent_order", + mock.Mock(side_effect=RuntimeError(failure_message)), + ) + elif failure_phase == "open": + monkeypatch.setattr( + weight_loader_module, + "RankStripedReadAheadSession", + mock.Mock(side_effect=OSError(failure_message)), + ) + + mapper = _test_weight_mapper(fail=failure_phase == "plan") + with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()) as weights: + assert weights + loader.activate_weight_session(weight_mapper=mapper) + assert not loader.last_checkpoint_io_status.activated + + reader_start.assert_not_called() + native_load.assert_not_called() + status = loader.last_checkpoint_io_status + assert status.effective == "native" + assert failure_message in status.fallback_reason + + def test_session_overlaps_materialization_and_cancels_tail(tmp_path, monkeypatch): checkpoint_dir, _ = _write_checkpoint(tmp_path) loader = _rank_striped_loader() @@ -438,6 +1055,7 @@ def wait_for_cancel(session, _extent): monkeypatch.setattr(read_ahead.RankStripedReadAheadSession, "_read_extent", wait_for_cancel) with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()) as weights: + loader.activate_weight_session() assert started.wait(timeout=5) assert not stopped.is_set() assert "model.norm.weight" in weights @@ -468,9 +1086,14 @@ def test_advisory_read_failure_keeps_materialized_weights(tmp_path, monkeypatch) assert "injected read failure" in status.fallback_reason -def test_mapping_and_materialization_failures_never_retry(tmp_path, monkeypatch): +@pytest.mark.parametrize( + "loader_factory", [_rank_striped_loader, _demand_ordered_rank_striped_loader] +) +def test_mapping_and_materialization_failures_never_retry(tmp_path, monkeypatch, loader_factory): checkpoint_dir, _ = _write_checkpoint(tmp_path) - loader = _rank_striped_loader() + loader = loader_factory() + close_communicator = mock.Mock(wraps=weight_loader_module.close_node_communicator) + monkeypatch.setattr(weight_loader_module, "close_node_communicator", close_communicator) native_load = mock.Mock(side_effect=AssertionError("must not reload")) monkeypatch.setattr(loader, "_load_weights_native", native_load) monkeypatch.setattr( @@ -481,8 +1104,9 @@ def test_mapping_and_materialization_failures_never_retry(tmp_path, monkeypatch) with loader.open_weight_session(str(checkpoint_dir), mapping=Mapping()): pass native_load.assert_not_called() + close_communicator.assert_called_once() - loader = _rank_striped_loader() + loader = loader_factory() native_load = mock.Mock(side_effect=AssertionError("must not reload")) error_log = mock.Mock() monkeypatch.setattr(loader, "_load_weights_native", native_load) @@ -500,6 +1124,27 @@ def test_mapping_and_materialization_failures_never_retry(tmp_path, monkeypatch) assert "cleanup failure" not in error_log.call_args.args[0] +def test_mapper_failure_closes_context_without_constructing_reader(monkeypatch): + context = weight_loader_module._DemandReadAheadContext( + None, + mock.Mock(), + CheckpointCatalog(objects=(), tensors=(CheckpointTensor("weight"),)), + [], + 0, + 1, + ) + loader = _demand_ordered_rank_striped_loader() + monkeypatch.setattr( + loader, "_start_rank_striped_read_ahead", lambda *_args, **_kwargs: ({}, context) + ) + with pytest.raises(ValueError, match="mapper failed"): + with loader.open_weight_session("checkpoint", mapping=Mapping()): + loader.activate_weight_session(ValueError("mapper failed")) + assert context.reader is None + assert loader._pending_read_ahead_session is None + context.node_communicator.Free.assert_called_once() + + def test_native_and_rank_striped_outputs_match(tmp_path, monkeypatch): # CPU-only CI has no CUDA device from which local_mpi_rank can derive a # rank. Model the single-rank environment explicitly; this test validates @@ -540,6 +1185,80 @@ def _run_real_mpi_scenarios(checkpoint_dir: str) -> dict: success = success_loader.last_checkpoint_io_status.effective communicator.Barrier() + demand_ordered_success_loader = _demand_ordered_rank_striped_loader() + with demand_ordered_success_loader.open_weight_session( + checkpoint_dir, mapping=mapping + ) as demand_ordered_success_weights: + demand_ordered_success_loader.activate_weight_session( + weight_mapper=_test_weight_mapper(mapping) + ) + assert demand_ordered_success_weights + demand_ordered_success = demand_ordered_success_loader.last_checkpoint_io_status.effective + communicator.Barrier() + + plan_failure_loader = _demand_ordered_rank_striped_loader() + with plan_failure_loader.open_weight_session( + checkpoint_dir, mapping=mapping + ) as plan_failure_weights: + plan_failure_loader.activate_weight_session( + weight_mapper=_test_weight_mapper(mapping, fail=rank == world_size - 1) + ) + assert plan_failure_weights + plan_failure_status = plan_failure_loader.last_checkpoint_io_status + communicator.Barrier() + + original_reader = weight_loader.RankStripedReadAheadSession + if rank == world_size - 1: + + def fail_reader_open(*_args, **_kwargs): + raise OSError("rank-local reader-open failure") + + weight_loader.RankStripedReadAheadSession = fail_reader_open + open_failure_loader = _demand_ordered_rank_striped_loader() + try: + with open_failure_loader.open_weight_session(checkpoint_dir, mapping=mapping): + open_failure_loader.activate_weight_session(weight_mapper=_test_weight_mapper(mapping)) + open_failure_status = open_failure_loader.last_checkpoint_io_status + finally: + weight_loader.RankStripedReadAheadSession = original_reader + communicator.Barrier() + + readiness_loader = _demand_ordered_rank_striped_loader() + local_readiness_error = ( + RuntimeError("rank-local readiness failure") if rank == world_size - 1 else None + ) + readiness_error = None + try: + with readiness_loader.open_weight_session( + checkpoint_dir, mapping=mapping + ) as readiness_weights: + assert readiness_weights + readiness_loader.activate_weight_session(local_readiness_error) + except BaseException as error: + readiness_error = f"{type(error).__name__}: {error}" + readiness_status = readiness_loader.last_checkpoint_io_status + communicator.Barrier() + + original_start = rank_striped_read_ahead.RankStripedReadAheadSession.start + if rank == world_size - 1: + + def fail_reader_start(_session, *, defer_reads=False): + del defer_reads + raise RuntimeError("rank-local demand-ordered reader-start failure") + + rank_striped_read_ahead.RankStripedReadAheadSession.start = fail_reader_start + reader_start_loader = _demand_ordered_rank_striped_loader() + try: + with reader_start_loader.open_weight_session( + checkpoint_dir, mapping=mapping + ) as reader_start_weights: + assert reader_start_weights + reader_start_loader.activate_weight_session(weight_mapper=_test_weight_mapper(mapping)) + reader_start_status = reader_start_loader.last_checkpoint_io_status + finally: + rank_striped_read_ahead.RankStripedReadAheadSession.start = original_start + communicator.Barrier() + fallback_failure_loader = _rank_striped_loader() if rank == world_size - 1: original_native_load = fallback_failure_loader._load_weights_native @@ -583,6 +1302,7 @@ def fail_after_native_load(*args, **kwargs): try: with body_loader.open_weight_session(checkpoint_dir, mapping=mapping) as weights: assert weights + body_loader.activate_weight_session() if rank == world_size - 1: raise RuntimeError("rank-local body failure") except RuntimeError as error: @@ -607,6 +1327,7 @@ def run_and_record_error(session): read_loader = _rank_striped_loader() try: with read_loader.open_weight_session(checkpoint_dir, mapping=mapping) as read_weights: + read_loader.activate_weight_session() local_error_recorded = rank != world_size - 1 or read_error_recorded.wait(timeout=5) error_recorded = communicator.allgather(local_error_recorded) if not all(error_recorded): @@ -621,6 +1342,19 @@ def run_and_record_error(session): "rank": rank, "success": success, "sum": float(success_weights["model.norm.weight"].sum()), + "demand_ordered_success": demand_ordered_success, + "demand_ordered_sum": float(demand_ordered_success_weights["model.norm.weight"].sum()), + "plan_failure_effective": plan_failure_status.effective, + "plan_failure_reason": plan_failure_status.fallback_reason, + "plan_failure_sum": float(plan_failure_weights["model.norm.weight"].sum()), + "open_failure_effective": open_failure_status.effective, + "open_failure_reason": open_failure_status.fallback_reason, + "readiness_error": readiness_error, + "readiness_effective": readiness_status.effective, + "readiness_reason": readiness_status.fallback_reason, + "reader_start_effective": reader_start_status.effective, + "reader_start_reason": reader_start_status.fallback_reason, + "reader_start_sum": float(reader_start_weights["model.norm.weight"].sum()), "fallback_failure_error": fallback_failure_error, "subgroup_effective": subgroup_effective, "subgroup_sum": float(subgroup_weights["model.norm.weight"].sum()), @@ -649,6 +1383,29 @@ def test_real_mpi_success_and_rank_local_failures(tmp_path, mpi_pool_executor): assert {result["rank"] for result in results} == {0, 1} assert all(result["success"] == "rank_striped_read_ahead" for result in results) assert all(result["sum"] == expected_sum for result in results) + assert all( + result["demand_ordered_success"] == "demand_ordered_rank_striped_read_ahead" + for result in results + ) + assert all(result["demand_ordered_sum"] == expected_sum for result in results) + assert all(result["plan_failure_effective"] == "native" for result in results) + assert all( + "injected weight-plan failure" in result["plan_failure_reason"] for result in results + ) + assert all(result["plan_failure_sum"] == expected_sum for result in results) + assert all(result["open_failure_effective"] == "native" for result in results) + assert all( + "rank-local reader-open failure" in result["open_failure_reason"] for result in results + ) + assert all("rank-local readiness failure" in result["readiness_error"] for result in results) + assert all(result["readiness_effective"] == "none" for result in results) + assert all("rank-local readiness failure" in result["readiness_reason"] for result in results) + assert all(result["reader_start_effective"] == "native" for result in results) + assert all( + "rank-local demand-ordered reader-start failure" in result["reader_start_reason"] + for result in results + ) + assert all(result["reader_start_sum"] == expected_sum for result in results) assert all( "rank-local native fallback failure" in result["fallback_failure_error"] for result in results diff --git a/tests/unittest/_torch/models/checkpoints/test_weight_load_plan.py b/tests/unittest/_torch/models/checkpoints/test_weight_load_plan.py index c7842f0336d0..ada8c744c929 100644 --- a/tests/unittest/_torch/models/checkpoints/test_weight_load_plan.py +++ b/tests/unittest/_torch/models/checkpoints/test_weight_load_plan.py @@ -4,6 +4,8 @@ from types import SimpleNamespace import pytest +import torch +from torch import nn from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper @@ -112,6 +114,83 @@ def apply_callbacks(self, module, module_name, module_names_breakdown, weights): raise AssertionError("not used") +class _ManualModule(nn.Module): + def __init__(self, value: float = 1.0, *, with_child: bool = False): + super().__init__() + self.weight = nn.Parameter(torch.tensor([value])) + if with_child: + self.child = _ManualModule(value + 1.0) + + +class _CustomLoadModule(_ManualModule): + def __init__(self): + super().__init__(3.0, with_child=True) + + def load_weights(self, *args, **kwargs): + raise AssertionError("planning must not invoke custom weight loading") + + +class _SpecialInstanceModule(_ManualModule): + def __init__(self): + super().__init__(5.0, with_child=True) + + +class _MoeBackendModule(_CustomLoadModule): + pass + + +class _StructuredModel(nn.Module): + def __init__(self, model_config: ModelConfig): + super().__init__() + self.model_config = model_config + self.config = SimpleNamespace() + self.preload_weight_modules = ["special", "qkv_proj"] + self.layer = nn.Module() + self.layer.qkv_proj = _ManualModule() + self.layer.special = _SpecialInstanceModule() + self.layer.custom = _CustomLoadModule() + self.layer.manual = _ManualModule(with_child=True) + + +class _MoePathModel(nn.Module): + def __init__(self, model_config: ModelConfig): + super().__init__() + self.model_config = model_config + self.config = SimpleNamespace() + self.preload_weight_modules = ["experts"] + self.layer = nn.Module() + self.layer.experts = nn.Module() + self.layer.experts.backend = _MoeBackendModule() + + +class _ManyCustomModulesModel(nn.Module): + def __init__(self, model_config: ModelConfig, module_count: int): + super().__init__() + self.model_config = model_config + self.config = SimpleNamespace() + self.blocks = nn.ModuleList(_CustomLoadModule() for _ in range(module_count)) + + +class _CountingName(str): + startswith_calls = 0 + + def startswith(self, prefix, *args): + type(self).startswith_calls += 1 + return super().startswith(prefix, *args) + + +class _StructuredMapper(_Mapper): + def map_weights(self) -> None: + self.mapping.update({"qkv_proj": ["q_proj", "k_proj", "v_proj"]}) + + def is_special_instance_module(self, module: nn.Module) -> bool: + return isinstance(module, _SpecialInstanceModule) + + +def _catalog_with_names(*names: str) -> CheckpointCatalog: + return CheckpointCatalog(objects=(), tensors=tuple(CheckpointTensor(name) for name in names)) + + def test_base_mapper_emits_valid_conservative_plan() -> None: catalog = _catalog() model_config = ModelConfig(mapping=Mapping(world_size=2, rank=1, tp_size=2)) @@ -126,3 +205,197 @@ def test_base_mapper_emits_valid_conservative_plan() -> None: assert plan.selected_tensor_names is None assert plan.demands[0].source_names == ("a", "b", "c") assert plan.demands[0].destination_ranks == (1,) + + +def test_base_mapper_builds_deterministic_advisory_consumer_groups() -> None: + catalog = _catalog_with_names( + "unmatched.head", + "layer.q_proj.bias", + "layer.custom.child.weight", + "layer.special.child.weight", + "layer.manual.child.weight", + "layer.q_proj.weight", + "layer.special.weight", + "layer.k_proj.weight", + "layer.custom.weight", + "layer.manual.weight", + "layer.v_proj.weight", + # These look similar to inferred names but are not structural prefix or + # exact-parameter matches and must remain in the physical tail. + "layer.q_proj_extra.weight", + "layer.manual.weight.extra", + ) + model_config = ModelConfig(mapping=Mapping(world_size=2, rank=1)) + model = _StructuredModel(model_config) + mapper = _StructuredMapper() + mapper.init_model_and_config(model, model_config) + + original_state = {name: value.detach().clone() for name, value in model.state_dict().items()} + original_mapping = {name: tuple(source_names) for name, source_names in mapper.mapping.items()} + original_preloads = tuple(model.preload_weight_modules) + + plan = mapper.build_weight_load_plan(catalog) + repeated_plan = mapper.build_weight_load_plan(catalog) + + assert plan.coverage is WeightLoadPlanCoverage.CONSERVATIVE + assert plan.ordering is WeightLoadOrderConfidence.ADVISORY + assert plan.selected_tensor_names is None + assert plan.plan_id == repeated_plan.plan_id + assert plan.demands == repeated_plan.demands + + special, fusion = plan.demands[:2] + assert special.source_names == ( + "layer.special.child.weight", + "layer.special.weight", + ) + assert fusion.source_names == ( + "layer.q_proj.bias", + "layer.q_proj.weight", + "layer.k_proj.weight", + "layer.v_proj.weight", + ) + assert special.priority == 0 + assert special.predecessors == () + assert fusion.priority == 1 + assert fusion.predecessors == (special.group_id,) + + nonpreload = [demand for demand in plan.demands if demand.group_id.startswith("module:")] + assert [demand.source_names for demand in nonpreload] == [ + ( + "layer.custom.child.weight", + "layer.custom.weight", + ), + ("layer.manual.weight",), + ("layer.manual.child.weight",), + ] + assert {demand.priority for demand in nonpreload} == {2} + assert all(not demand.predecessors for demand in nonpreload) + + assert plan.demands[-1].group_id == "unmatched_checkpoint_tensors" + assert plan.demands[-1].source_names == ( + "unmatched.head", + "layer.q_proj_extra.weight", + "layer.manual.weight.extra", + ) + + described_names = [name for demand in plan.demands for name in demand.source_names] + assert len(described_names) == len(set(described_names)) + assert set(described_names) == catalog.tensor_names + assert all(demand.destination_ranks == (1,) for demand in plan.demands) + + assert mapper.mapping == { + name: list(source_names) for name, source_names in original_mapping.items() + } + assert tuple(model.preload_weight_modules) == original_preloads + assert model.state_dict().keys() == original_state.keys() + assert all( + torch.equal(value, original_state[name]) for name, value in model.state_dict().items() + ) + + +def test_base_mapper_uses_opaque_order_for_unmodeled_source_aliases() -> None: + catalog = _catalog_with_names( + # Some mappers strip source-only namespaces before materialization. + # Generic planning must degrade conservatively until such aliases are + # expressed through a pure mapper planning hook. + "language_model.layer.manual.weight", + "language_model.layer.q_proj.weight", + ) + model_config = ModelConfig(mapping=Mapping(world_size=2, rank=1)) + model = _StructuredModel(model_config) + mapper = _StructuredMapper() + mapper.init_model_and_config(model, model_config) + + plan = mapper.build_weight_load_plan(catalog) + + assert plan.coverage is WeightLoadPlanCoverage.CONSERVATIVE + assert plan.ordering is WeightLoadOrderConfidence.OPAQUE + assert plan.selected_tensor_names is None + assert len(plan.demands) == 1 + assert plan.demands[0].group_id == "all_checkpoint_tensors" + assert plan.demands[0].source_names == ( + "language_model.layer.manual.weight", + "language_model.layer.q_proj.weight", + ) + + +def test_base_mapper_normalizes_moe_sources_without_inventing_preload_membership( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + "tensorrt_llm._torch.models.checkpoints.base_weight_mapper.is_moe_weight_owner", + lambda module: isinstance(module, _MoeBackendModule), + ) + catalog = _catalog_with_names( + "tail.first", + "layer.experts.weight_scale", + "layer.experts.weight", + "tail.second", + ) + model_config = ModelConfig(mapping=Mapping(world_size=2, rank=1)) + model = _MoePathModel(model_config) + mapper = _Mapper() + mapper.init_model_and_config(model, model_config) + original_state = {name: value.detach().clone() for name, value in model.state_dict().items()} + + plan = mapper.build_weight_load_plan(catalog) + + assert plan.coverage is WeightLoadPlanCoverage.CONSERVATIVE + assert plan.ordering is WeightLoadOrderConfidence.ADVISORY + assert plan.selected_tensor_names is None + assert plan.demands[0].group_id.endswith("layer.experts.backend") + assert plan.demands[0].group_id.startswith("module:") + assert plan.demands[0].priority == 0 + assert plan.demands[0].source_names == ( + "layer.experts.weight_scale", + "layer.experts.weight", + ) + assert plan.demands[-1].source_names == ( + "tail.first", + "tail.second", + ) + described_names = [name for demand in plan.demands for name in demand.source_names] + assert len(described_names) == len(set(described_names)) + assert set(described_names) == catalog.tensor_names + assert all( + torch.equal(value, original_state[name]) for name, value in model.state_dict().items() + ) + + +def test_base_mapper_uses_opaque_order_for_serial_loading_override(monkeypatch) -> None: + monkeypatch.setenv("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL", "true") + catalog = _catalog_with_names("layer.q_proj.weight", "layer.special.weight") + model_config = ModelConfig(mapping=Mapping()) + model = _StructuredModel(model_config) + mapper = _StructuredMapper() + mapper.init_model_and_config(model, model_config) + + plan = mapper.build_weight_load_plan(catalog) + + assert plan.ordering is WeightLoadOrderConfidence.OPAQUE + assert plan.described_tensor_names == catalog.tensor_names + + +def test_base_mapper_uses_indexed_structural_prefix_matching() -> None: + module_count = 128 + catalog = _catalog_with_names( + *(_CountingName(f"blocks.{index}.weight") for index in range(module_count)), + _CountingName("unmatched.tail"), + ) + model_config = ModelConfig(mapping=Mapping(world_size=2, rank=1)) + model = _ManyCustomModulesModel(model_config, module_count) + mapper = _Mapper() + mapper.init_model_and_config(model, model_config) + _CountingName.startswith_calls = 0 + + plan = mapper.build_weight_load_plan(catalog) + + # A full physical-name scan per custom module would make roughly + # module_count**2 startswith calls. Bisected ranges inspect only matching + # names plus each range boundary. + assert _CountingName.startswith_calls < module_count * 4 + assert plan.coverage is WeightLoadPlanCoverage.CONSERVATIVE + assert plan.ordering is WeightLoadOrderConfidence.ADVISORY + assert plan.selected_tensor_names is None + assert plan.demands[-1].source_names == ("unmatched.tail",) + assert plan.described_tensor_names == catalog.tensor_names diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 48ece4abb49b..b682ec4d6ebe 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -128,7 +128,7 @@ methods: default: null status: prototype checkpoint_io_policy: - annotation: Literal['auto', 'native', 'rank_striped_read_ahead'] + annotation: Literal['auto', 'native', 'rank_striped_read_ahead', 'demand_ordered_rank_striped_read_ahead'] default: auto status: prototype mx_config: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index e9c9925dfb37..fcf63e59a3f0 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -91,8 +91,15 @@ def test_generation_config_auto_rejects_autodeploy() -> None: @pytest.mark.cpu_only -@pytest.mark.parametrize("policy", - ["auto", "native", "rank_striped_read_ahead"]) +@pytest.mark.parametrize( + "policy", + [ + "auto", + "native", + "rank_striped_read_ahead", + "demand_ordered_rank_striped_read_ahead", + ], +) def test_checkpoint_io_policy_defaults_to_auto_and_accepts_supported_values( policy: str) -> None: assert TorchLlmArgs(model=llama_model_path).checkpoint_io_policy == "auto" diff --git a/tests/unittest/test_checkpoint_io_policy.py b/tests/unittest/test_checkpoint_io_policy.py new file mode 100644 index 000000000000..051e31ce52a0 --- /dev/null +++ b/tests/unittest/test_checkpoint_io_policy.py @@ -0,0 +1,74 @@ +# 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. + +import json +from pathlib import Path +from typing import get_args + +import pytest +import yaml +from pydantic import BaseModel + +from tensorrt_llm._checkpoint_io_policy import ( + AUTO_IO_POLICY, + CHECKPOINT_IO_POLICIES, + CHECKPOINT_IO_POLICY_INFO, + DEMAND_ORDERED_IO_POLICY, + EXECUTABLE_CHECKPOINT_IO_POLICIES, + NATIVE_IO_POLICY, + RANK_STRIPED_IO_POLICIES, + RANK_STRIPED_IO_POLICY, + CheckpointIoPolicy, +) +from tensorrt_llm.usage.llmapi_config import manifest_rows + +pytestmark = pytest.mark.cpu_only + + +def test_policy_catalog_is_complete() -> None: + assert CHECKPOINT_IO_POLICIES == get_args(CheckpointIoPolicy) + assert tuple(CHECKPOINT_IO_POLICY_INFO) == CHECKPOINT_IO_POLICIES + assert RANK_STRIPED_IO_POLICIES == (RANK_STRIPED_IO_POLICY, DEMAND_ORDERED_IO_POLICY) + assert EXECUTABLE_CHECKPOINT_IO_POLICIES == (NATIVE_IO_POLICY, *RANK_STRIPED_IO_POLICIES) + assert CHECKPOINT_IO_POLICIES == (AUTO_IO_POLICY, *EXECUTABLE_CHECKPOINT_IO_POLICIES) + for policy, info in CHECKPOINT_IO_POLICY_INFO.items(): + assert info.display_name.strip() and info.description.strip() + assert info.kind == ("selector" if policy == AUTO_IO_POLICY else "implementation") + + +def test_policy_field_preserves_string_api() -> None: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + field = TorchLlmArgs.model_fields["checkpoint_io_policy"] + assert field.annotation == CheckpointIoPolicy + assert field.default == AUTO_IO_POLICY + assert field.json_schema_extra["type"] == str(CheckpointIoPolicy).replace("typing.", "") + + +def test_policy_catalog_matches_committed_artifacts() -> None: + class PolicyArgs(BaseModel): + checkpoint_io_policy: CheckpointIoPolicy = AUTO_IO_POLICY + + root = Path(__file__).resolve().parents[2] + manifest = json.loads((root / "tensorrt_llm/usage/llm_args_golden_manifest.json").read_text()) + entry = next( + item for item in manifest["TorchLlmArgs"] if item["path"] == "checkpoint_io_policy" + ) + assert entry == manifest_rows(PolicyArgs)[0] + reference = yaml.safe_load( + (root / "tests/unittest/api_stability/references/llm.yaml").read_text() + )["methods"]["__init__"]["parameters"]["checkpoint_io_policy"] + assert reference["annotation"] == str(CheckpointIoPolicy).replace("typing.", "") + assert reference["default"] == AUTO_IO_POLICY diff --git a/tests/unittest/usage/test_llmapi_config_capture.py b/tests/unittest/usage/test_llmapi_config_capture.py index 9798459343dd..217e26cd8302 100644 --- a/tests/unittest/usage/test_llmapi_config_capture.py +++ b/tests/unittest/usage/test_llmapi_config_capture.py @@ -697,7 +697,7 @@ def test_collect_llm_api_config_captures_gms_load_format(): def test_collect_llm_api_config_captures_checkpoint_io_policy() -> None: - for policy in ("auto", "rank_striped_read_ahead"): + for policy in ("auto", "rank_striped_read_ahead", "demand_ordered_rank_striped_read_ahead"): args = TorchLlmArgs( model="/customer/private/Llama", skip_tokenizer_init=True,