diff --git a/modelexpress_client/python/modelexpress/trtllm_live_transfer.py b/modelexpress_client/python/modelexpress/trtllm_live_transfer.py index 75d1962ea..7b3746b72 100644 --- a/modelexpress_client/python/modelexpress/trtllm_live_transfer.py +++ b/modelexpress_client/python/modelexpress/trtllm_live_transfer.py @@ -32,6 +32,62 @@ logger = logging.getLogger("modelexpress.trtllm_live_transfer") +_TRTLLM_RUNTIME_ALIAS_COMPONENTS = frozenset( + {"next_attn", "next_layer_layernorm"} +) + + +def _is_trtllm_runtime_alias_name(name: str) -> bool: + """Whether a parameter path exists only after TRT-LLM runtime alias setup.""" + return bool( + _TRTLLM_RUNTIME_ALIAS_COMPONENTS.intersection(name.split(".")) + ) + + +def _canonical_named_parameters(torch_model: Any) -> list[tuple[str, torch.Tensor]]: + """Return one stable, non-runtime-alias name for each parameter storage.""" + canonical = [] + canonical_ptrs = set() + alias_names_by_ptr = {} + for name, param in torch_model.named_parameters(remove_duplicate=False): + ptr = param.data.data_ptr() + storage_key = (param.device.type, param.device.index, ptr) + if _is_trtllm_runtime_alias_name(name): + alias_names_by_ptr.setdefault(storage_key, []).append(name) + continue + if storage_key in canonical_ptrs: + logger.debug("Skipping duplicate canonical param: %s (ptr=%x)", name, ptr) + continue + canonical_ptrs.add(storage_key) + canonical.append((name, param)) + + alias_only_ptrs = set(alias_names_by_ptr).difference(canonical_ptrs) + if alias_only_ptrs: + examples = [alias_names_by_ptr[ptr][0] for ptr in list(alias_only_ptrs)[:3]] + raise RuntimeError( + "TRT-LLM runtime aliases have no canonical parameter path: " + f"{len(alias_only_ptrs)} storages; examples: {examples}" + ) + return canonical + + +def _require_exact_catalog_match( + source_descs: dict[str, Any], target_params: dict[str, torch.Tensor] +) -> None: + """Reject P2P unless source and target expose exactly the same names.""" + source_names = set(source_descs) + target_names = set(target_params) + absent_from_target = sorted(source_names.difference(target_names)) + absent_from_source = sorted(target_names.difference(source_names)) + if absent_from_target or absent_from_source: + raise RuntimeError( + "MX P2P source/target parameter catalogs do not match: " + f"{len(absent_from_target)} source tensors are absent from the target " + f"(examples: {absent_from_target[:3]}); " + f"{len(absent_from_source)} target tensors are absent from the source " + f"(examples: {absent_from_source[:3]})" + ) + def _build_trtllm_identity( model_name: str, @@ -62,8 +118,9 @@ def _build_trtllm_identity( def publish_model_params(torch_model: Any) -> None: """Publish this rank's model params to ModelExpress directly from a torch model. - Called from ModelLoader.load() BEFORE post_load_weights() so that targets - receive pre-processed weights and can run their own post_load_weights(). + TensorRT-LLM may call this after post-load transformations. Runtime-only + alias paths are excluded so receivers can match the final bytes against + their canonical, pre-alias parameter tree. Each rank publishes independently via MxClient (per-worker API). """ @@ -84,15 +141,9 @@ def publish_model_params(torch_model: Any) -> None: mx_server = envs.MODEL_EXPRESS_URL or "modelexpress-server:8001" param_tensors = {} - seen_data_ptrs = set() total_bytes = 0 - for name, param in torch_model.named_parameters(): + for name, param in _canonical_named_parameters(torch_model): if param.device.type == "cuda" and param.device.index == device_id: - ptr = param.data.data_ptr() - if ptr in seen_data_ptrs: - logger.debug("Skipping aliased param: %s (ptr=%x)", name, ptr) - continue - seen_data_ptrs.add(ptr) param_tensors[name] = param.data total_bytes += param.numel() * param.element_size() @@ -127,14 +178,10 @@ def publish_model_params(torch_model: Any) -> None: for name, tensor in param_tensors.items() ] - # Dual-write legacy `tensors` alongside `tensor_source` for servers that - # predate the tensor_source oneof (see publish.py for the full rationale). worker = p2p_pb2.WorkerMetadata( worker_rank=mpi_rank, nixl_metadata=nixl_mgr.nixl_metadata, - tensors=tensor_protos, tensor_source=tensor_source_metadata(tensor_protos), - accelerator="cuda", ) identity = _build_trtllm_identity(model_name=model_name) @@ -195,15 +242,9 @@ def publish_from_worker(worker: Any) -> None: mx_server = envs.MODEL_EXPRESS_URL or "modelexpress-server:8001" param_tensors = {} - seen_data_ptrs = set() total_bytes = 0 - for name, param in torch_model.named_parameters(): + for name, param in _canonical_named_parameters(torch_model): if param.device.type == "cuda" and param.device.index == device_id: - ptr = param.data.data_ptr() - if ptr in seen_data_ptrs: - logger.debug("Skipping aliased param: %s (ptr=%x)", name, ptr) - continue - seen_data_ptrs.add(ptr) param_tensors[name] = param.data total_bytes += param.numel() * param.element_size() @@ -247,14 +288,10 @@ def publish_from_worker(worker: Any) -> None: for name, tensor in param_tensors.items() ] - # Dual-write legacy `tensors` alongside `tensor_source` for servers that - # predate the tensor_source oneof (see publish.py for the full rationale). my_worker = p2p_pb2.WorkerMetadata( worker_rank=mpi_rank, nixl_metadata=nixl_mgr.nixl_metadata, - tensors=tensor_protos, tensor_source=tensor_source_metadata(tensor_protos), - accelerator="cuda", ) identity = _build_trtllm_identity(model_name=model_name) @@ -290,6 +327,41 @@ def load_weights( mapping: Any = None, model: Any = None, **kwargs, + ) -> dict[str, Any]: + device_id = torch.cuda.current_device() + try: + from mpi4py import MPI + mpi_rank = MPI.COMM_WORLD.Get_rank() + except Exception: + mpi_rank = device_id + + log_dir = envs.MX_TRANSFER_LOG_DIR + os.makedirs(log_dir, exist_ok=True) + rank_log = os.path.join(log_dir, f"rank{mpi_rank}.log") + fh = logging.FileHandler(rank_log, mode="w") + fh.setLevel(logging.INFO) + fh.setFormatter( + logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s") + ) + mx_logger = logging.getLogger("modelexpress") + mx_logger.addHandler(fh) + try: + return self._load_weights( + checkpoint_dir=checkpoint_dir, + mapping=mapping, + model=model, + **kwargs, + ) + finally: + mx_logger.removeHandler(fh) + fh.close() + + def _load_weights( + self, + checkpoint_dir: str, + mapping: Any = None, + model: Any = None, + **kwargs, ) -> dict[str, Any]: from .nixl_transfer import NixlTransferManager from .types import TensorDescriptor @@ -315,15 +387,6 @@ def load_weights( except Exception: mpi_rank = device_id - # MPI workers' stdout is swallowed by TRT-LLM — write to per-rank file - log_dir = envs.MX_TRANSFER_LOG_DIR - os.makedirs(log_dir, exist_ok=True) - rank_log = os.path.join(log_dir, f"rank{mpi_rank}.log") - fh = logging.FileHandler(rank_log, mode="w") - fh.setLevel(logging.INFO) - fh.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s")) - logging.getLogger("modelexpress").addHandler(fh) - logger.info( "Live transfer: loading '%s' rank %d (GPU %d)", model_name, mpi_rank, device_id ) @@ -343,7 +406,7 @@ def load_weights( # 2. Build name→param map from target model target_params = {} - for name, param in model.named_parameters(): + for name, param in _canonical_named_parameters(model): if param.device.index == device_id: target_params[name] = param.data @@ -353,11 +416,11 @@ def load_weights( # 3. Build source name→descriptor map source_descs = {t.name: t for t in worker_tensor_descriptors(source_worker)} + _require_exact_catalog_match(source_descs, target_params) # 4. Match source and target by name matched = [] dtype_cast_needed = [] - unmatched_source = [] for src_name, src_desc in source_descs.items(): if src_name in target_params: dst_param = target_params[src_name] @@ -381,15 +444,6 @@ def load_weights( "Size mismatch for %s: source=%d target=%d (numel src=%d dst=%d)", src_name, src_size, dst_size, src_numel, dst_param.numel(), ) - else: - unmatched_source.append(src_name) - - if unmatched_source: - logger.warning( - "%d source tensors not found in target: %s...", - len(unmatched_source), unmatched_source[:3], - ) - # For dtype-mismatched tensors, allocate temp buffers at source dtype dtype_map = {"torch.bfloat16": torch.bfloat16, "torch.float16": torch.float16, "torch.float32": torch.float32, "torch.uint8": torch.uint8, diff --git a/modelexpress_client/python/tests/test_trtllm_live_transfer.py b/modelexpress_client/python/tests/test_trtllm_live_transfer.py new file mode 100644 index 000000000..f98f6558b --- /dev/null +++ b/modelexpress_client/python/tests/test_trtllm_live_transfer.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for TensorRT-LLM live-transfer catalog validation.""" + +import logging + +import pytest +import torch +from torch import nn + +from modelexpress.trtllm_live_transfer import ( + MxLiveWeightLoader, + _canonical_named_parameters, + _require_exact_catalog_match, +) + + +class _AliasLayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.next_attn = None + self.self_attn = None + + +class _AliasedModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.layers = nn.ModuleList([_AliasLayer(), _AliasLayer()]) + self.layers[1].self_attn = nn.Linear(2, 2, bias=False) + self.layers[0].next_attn = self.layers[1].self_attn + + +def test_runtime_aliases_are_excluded_from_canonical_catalog(): + model = _AliasedModel() + + default_names = dict(model.named_parameters()) + canonical_names = dict(_canonical_named_parameters(model)) + + assert "layers.0.next_attn.weight" in default_names + assert "layers.1.self_attn.weight" not in default_names + assert "layers.0.next_attn.weight" not in canonical_names + assert canonical_names["layers.1.self_attn.weight"] is model.layers[1].self_attn.weight + + +def test_exact_catalog_match_is_accepted(): + tensor = torch.zeros(1) + _require_exact_catalog_match({"model.weight": object()}, {"model.weight": tensor}) + + +@pytest.mark.parametrize( + ("source", "target", "message"), + [ + ({"source.only": object()}, {}, "1 source tensors are absent"), + ({}, {"target.only": torch.zeros(1)}, "1 target tensors are absent"), + ], +) +def test_incomplete_catalogs_fail_closed(source, target, message): + with pytest.raises(RuntimeError, match=message): + _require_exact_catalog_match(source, target) + + +def test_rank_log_handler_is_closed_on_failure(monkeypatch, tmp_path): + monkeypatch.setenv("MX_TRANSFER_LOG_DIR", str(tmp_path)) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + loader = MxLiveWeightLoader() + monkeypatch.setattr( + loader, + "_load_weights", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("catalog mismatch")), + ) + + with pytest.raises(RuntimeError, match="catalog mismatch"): + loader.load_weights("checkpoint", model=object()) + + rank_log = tmp_path / "rank0.log" + assert all( + getattr(handler, "baseFilename", None) != str(rank_log) + for handler in logging.getLogger("modelexpress").handlers + )