[TRTLLM-13246][feat] Wave 1: migrate aliases to setup_aliases and stage GMS RO load

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Author: chienchunhung

tensorrt_llm/_torch/memory/gpu_memory_backend.py

@@ -36,8 +36,14 @@
   CUDA memory pool. After loading, weights are committed for read-only
   access by other workers and the client transitions to RO mode in place.
 - **RO (Read-Only)**: Subsequent workers zero-copy import already-committed
-  weights from the GMS pool. `post_load_weights()` must run BEFORE
-  materialization so that module aliases are set up correctly.
+  weights from the GMS pool. `setup_aliases()` must run BEFORE
+  materialization so that module aliases are set up correctly, while derived
+  state is refreshed after real tensors are bound. RO is validated for models
+  whose `post_load_weights()` is pure alias wiring; models that additionally
+  rely on plain Python attributes set inside `post_load_weights()` (rather
+  than registered `nn.Buffer` / `nn.Parameter` assignments) need to migrate
+  those side effects to `transform_weights()` or `cache_derived_state()`
+  before they are safe on the RO path.
 """
 
 from contextlib import contextmanager
@@ -477,7 +483,7 @@ def materialize_module(self, model: nn.Module) -> None:
         by GPU pointers from the shared memory region — no data copies,
         no disk I/O, just CUDA VMM remapping. The model's submodule
         layout must already match the writer's at commit time, including
-        any aliases / derived buffers introduced by `post_load_weights`.
+        any aliases introduced by `setup_aliases`.
 
         Args:
             model: The `nn.Module` to materialize. Walks the full
@@ -489,11 +495,10 @@ def materialize_module(self, model: nn.Module) -> None:
             RuntimeError: If `connect()` has not been called yet.
 
         Note:
-            `post_load_weights()` must be called on the model BEFORE
-            this method. The order ensures that any aliases / derived
-            parameters created by post-load hooks are present on the
-            module tree at materialization time, so they are bound to
-            the same GMS storage as their primary tensor.
+            `setup_aliases()` must be called on the model BEFORE this method.
+            The order ensures that any structural aliases created by post-load
+            hooks are present on the module tree at materialization time, so
+            they are bound to the same GMS storage as their primary tensor.
         """
         if self._client is None:
             raise RuntimeError("GMS client not connected. Call connect() first.")

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -1920,7 +1920,7 @@ def load_weights(self, weights: ConsumableWeightsDict):
         weight_loader = DeepseekV3WeightLoader(self)
         weight_loader.load_weights(weights)
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_exaone_moe.py

@@ -725,7 +725,7 @@ def load_weights(
             allow_partial_loading=allow_partial_loading,
         )
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         # For the cross-layer residual+LN fusion.
         for idx, layer in enumerate(self.model.layers[: self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_glm.py

@@ -1074,7 +1074,7 @@ def load_weights(self, weights: ConsumableWeightsDict, allow_partial_loading: bo
         weight_loader = Glm4WeightLoader(self)
         weight_loader.load_weights(weights, allow_partial_loading=allow_partial_loading)
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         for idx, layer in enumerate(self.model.layers[: self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:
                 layer.next_layer_layernorm = self.model.norm

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -631,7 +631,7 @@ def load_weights(self, weights: Dict):
         else:
             self.load_hf_weights(weights)
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         for idx, layer in enumerate(
                 self.model.block[:self.config.num_hidden_layers]):
             if idx == 0:

tensorrt_llm/_torch/models/modeling_llama.py

@@ -484,7 +484,7 @@ def __init__(
         self.input_layernorm = RMSNorm(hidden_size=config.hidden_size,
                                        eps=config.rms_norm_eps,
                                        dtype=config.torch_dtype)
-        # When post_load_weights() chains layernorms across layers,
+        # When setup_aliases() chains layernorms across layers,
         # this flag is set to True to skip the input layernorm in
         # forward() since it is handled by the previous layer.
         self.skip_input_layernorm = False
@@ -709,7 +709,7 @@ def __init__(
             quantize_type="nvfp4"
             if not self.disable_nvfp4_layernorm_fusion and self.is_nvfp4
             and not (differ_pp_stage_with_previous_layer) else None)
-        # When post_load_weights() chains layernorms across layers,
+        # When setup_aliases() chains layernorms across layers,
         # this flag is set to True to skip the input layernorm in
         # forward() since it is handled by the previous layer.
         self.skip_input_layernorm = False
@@ -983,7 +983,7 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]):
         self.norm = RMSNorm(hidden_size=config.hidden_size,
                             eps=config.rms_norm_eps,
                             dtype=config.torch_dtype)
-        # When post_load_weights() chains the final norm into the
+        # When setup_aliases() chains the final norm into the
         # last decoder layer, this flag is set to True to skip
         # applying it again in forward().
         self.skip_norm = False
@@ -1088,7 +1088,7 @@ def __init__(self, model_config: ModelConfig[LlamaConfig]):
         self.norm = RMSNorm(hidden_size=config.hidden_size,
                             eps=config.rms_norm_eps,
                             dtype=config.torch_dtype)
-        # When post_load_weights() chains the final norm into the
+        # When setup_aliases() chains the final norm into the
         # last decoder layer, this flag is set to True to skip
         # applying it again in forward().
         self.skip_norm = False
@@ -1140,7 +1140,7 @@ def __init__(
     ):
         super().__init__(LlamaModel(model_config), model_config)
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:
@@ -1564,7 +1564,7 @@ def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
             if had_mm_encoder:
                 self.mm_encoder = saved_mm_encoder
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_qwen3_moe.py

@@ -417,7 +417,7 @@ def __init__(
         )
         self.preload_weight_modules = self.model.preload_weight_modules
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_qwen3_next.py

@@ -980,7 +980,7 @@ def load_weights(self, weights: dict, weight_mapper: BaseWeightMapper):
         new_weights = weight_mapper.preprocess_weights(weights)
         super().load_weights(new_weights, weight_mapper)
 
-    def post_load_weights(self):
+    def setup_aliases(self):
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -462,8 +462,9 @@ def init_meta_tensor(t: torch.Tensor):
             # post_load_* hooks itself, so the shared post-load block below
             # must skip them. RW handles them inside `mem_pool_scope` so the
             # committed pool reflects the post-post_load layout; RO runs
-            # `module.post_load_weights()` before `materialize_module` to
-            # wire aliases prior to zero-copy mapping.
+            # ``setup_aliases()`` before ``materialize_module`` to wire aliases
+            # prior to zero-copy mapping, then refreshes derived state after
+            # real GMS tensors are bound.
             gms_post_load_handled = False
             if load_format == LoadFormat.AUTO:
                 # Pass model= so format-specific loaders (e.g. MX) can
@@ -692,31 +693,33 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                         # presharded modules).
                         #
                         # Hook order:
-                        #   1. `post_load_apply`: format-specific apply
+                        #   1. ``post_load_apply``: format-specific apply
                         #      work (e.g., MX preshard markers).
-                        #   2. Per-module `post_load_weights`: creates
-                        #      aliases/derived parameter attributes BEFORE
-                        #      `materialize_module` walks the final module
-                        #      tree (including `draft_model` for spec dec).
-                        #   3. `materialize_module`: zero-copy bind GMS
+                        #   2. Top-level ``setup_aliases``: creates structural
+                        #      aliases BEFORE ``materialize_module`` walks the
+                        #      final module tree (including ``draft_model`` for
+                        #      spec dec).
+                        #   3. SourceIdentity gate: STRICT pre-materialize
+                        #      compatibility check (GMS has no disk fallback).
+                        #   4. ``materialize_module``: zero-copy bind GMS
                         #      pool storage onto the model parameters.
-                        #   4. `post_load_publish`: any receiver-side
+                        #   5. Per-module ``cache_derived_state``: recompute
+                        #      Python-side state from real, materialized
+                        #      tensors without re-running one-shot transforms.
+                        #   6. ``post_load_publish``: any receiver-side
                         #      publish (no-op via the receiver guard).
                         checkpoint_loader.post_load_apply(
                             model, weights_preloaded=True)
 
-                        for module in model.modules():
-                            if hasattr(module,
-                                       'post_load_weights') and not getattr(
-                                           module, '_weights_removed', False):
-                                module.post_load_weights()
+                        self._setup_aliases(model)
 
                         # Pre-materialize compatibility gate. GMS has no
                         # disk-fallback path, so a mismatch raises under STRICT
                         # rather than falling back.
                         self._check_gms_source_identity(gms_backend)
 
                         gms_backend.materialize_module(model)
+                        self._walk_cache_state(model)
 
                         checkpoint_loader.post_load_publish(
                             model,
@@ -929,6 +932,7 @@ def reload(self,
                 "Cannot reload weights: weight_mapper was not initialized. "
                 "This can happen when the initial load used GMS, MX P2P, or "
                 "VISION_ONLY, which bypass the standard weight mapping path.")
+        self._reset_weights_transformed(model)
         self._call_load_weights(model.load_weights,
                                 weights,
                                 self.weight_mapper,

tests/unittest/_torch/pyexecutor/test_model_loader_mx.py

@@ -48,6 +48,12 @@ def load_weights(self, weights, mapper):
     def load_draft_weights(self, weights, mapper):
         self._events.append("load_draft_weights")
 
+    def setup_aliases(self):
+        self._events.append("setup_aliases")
+
+    def cache_derived_state(self):
+        self._events.append("cache_derived_state")
+
     def post_load_weights(self):
         self._events.append("post_load_weights")
 
@@ -57,8 +63,15 @@ def _moe_context(config, mapping):
     yield None
 
 
-def _make_loader(monkeypatch, *, events, spec_config=None):
-    llm_args = SimpleNamespace(load_format=LoadFormat.AUTO)
+def _make_loader(monkeypatch, *, events, spec_config=None, load_format=LoadFormat.AUTO):
+    llm_args = SimpleNamespace(
+        load_format=load_format,
+        gms_config=SimpleNamespace(
+            socket_path="/tmp/gms.sock",
+            mode="test",
+            tag="test",
+        ),
+    )
     loader = ModelLoader(
         llm_args=llm_args,
         mapping=MagicMock(name="mapping"),
@@ -75,6 +88,15 @@ def _make_loader(monkeypatch, *, events, spec_config=None):
     monkeypatch.setattr(model_loader_mod, "timing", lambda *_args, **_kwargs: nullcontext())
     monkeypatch.setattr(model_loader_mod, "maybe_create_moe_load_balancer", _moe_context)
     monkeypatch.setattr(model_loader_mod, "MetaInitMode", lambda: nullcontext())
+    # The MX and GMS load paths build a receiver-side SourceIdentity from the
+    # resolved ModelConfig. These tests stub the config, so short-circuit the
+    # fingerprint construction to a sentinel; identity-comparison logic is
+    # covered separately in test_source_identity.py.
+    monkeypatch.setattr(
+        model_loader_mod.SourceIdentity,
+        "from_model_config",
+        classmethod(lambda cls, *_args, **_kwargs: SimpleNamespace(name="local-identity")),
+    )
     monkeypatch.setattr(
         model_loader_mod.AutoModelForCausalLM,
         "from_config",
@@ -112,8 +134,12 @@ def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works(mon
 
     # reload() uses self.weight_mapper unconditionally; MX success must
     # initialize it even though the initial load skipped _call_load_weights.
+    model._weights_transformed = True
+    model.linear._weights_transformed = True
     loader.reload(model, {"reloaded": MagicMock()})
     assert loader._call_load_weights.call_count == 1
+    assert model._weights_transformed is False
+    assert model.linear._weights_transformed is False
     assert events == ["post_load_weights", "load_weights"]
 
 
@@ -157,6 +183,62 @@ def test_mx_fallback_runs_standard_weight_mapping(monkeypatch):
     )
 
 
+def test_gms_ro_materializes_between_alias_setup_and_cache_state(monkeypatch):
+    events = []
+    loader = _make_loader(monkeypatch, events=events, load_format=LoadFormat.GMS)
+    checkpoint_loader = MagicMock(name="checkpoint_loader")
+    checkpoint_loader.checkpoint_format = "GMS"
+
+    def record(event):
+        def _append(*_args, **_kwargs):
+            events.append(event)
+
+        return _append
+
+    checkpoint_loader.post_load_apply.side_effect = record("post_load_apply")
+    checkpoint_loader.post_load_publish.side_effect = record("post_load_publish")
+
+    # The STRICT pre-materialize identity gate runs between alias setup and
+    # materialization; record it to pin the ordering, without exercising the
+    # comparison logic (covered in test_source_identity.py).
+    monkeypatch.setattr(
+        model_loader_mod,
+        "check_weight_sharing_compatibility",
+        lambda *_args, **_kwargs: events.append("check_source_identity"),
+    )
+
+    class _GmsBackend:
+        def __init__(self, *args, **kwargs):
+            self.is_rw = False
+
+        def connect(self):
+            return True
+
+        def get_source_identity(self):
+            return SimpleNamespace(name="remote-identity")
+
+        def materialize_module(self, model):
+            events.append("materialize_module")
+
+        def cleanup(self):
+            events.append("cleanup")
+
+    monkeypatch.setattr("tensorrt_llm._torch.memory.GMSBackend", _GmsBackend)
+
+    loader.load("/ckpt", checkpoint_loader)
+
+    assert events == [
+        "post_load_apply",
+        "setup_aliases",
+        "check_source_identity",
+        "materialize_module",
+        "cache_derived_state",
+        "post_load_publish",
+    ]
+    assert "post_load_weights" not in events
+    checkpoint_loader.load_weights.assert_not_called()
+
+
 class _HookRecorder(nn.Module):
     def __init__(
         self,