fix(client): use canonical TRT-LLM transfer catalogs

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,

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
+    )