[TRTLLM-13077][feat] Decompose post_load_weights()

tensorrt_llm/_torch/models/modeling_utils.py

@@ -579,6 +579,63 @@ def config(self):
     def vocab_size_padded(self) -> int:
         return self.lm_head.vocab_size_padded
 
+    def setup_aliases(self) -> None:
+        """Wire structural Python references between modules.
+
+        This stage is for module-tree structure only, such as assigning
+        cross-layer module references or shared module aliases. It must not
+        read or mutate tensor values, so callers may run it before weight bytes
+        are available, materialized, or transformed.
+
+        The method is intentionally idempotent. Reassigning the same module
+        reference should preserve the same module graph, matching
+        ``torch.nn.Module.__setattr__`` semantics.
+
+        Returns:
+            None.
+        """
+
+    def transform_weights(self) -> None:
+        """Apply one-shot post-load transformations to weight tensors.
+
+        This stage is for irreversible or layout-changing tensor operations,
+        such as fusing weights or converting quantized weight representations.
+        Subclasses that migrate transform logic here should return early when
+        ``_weights_transformed`` is already true, and set it only after the
+        transform succeeds. Orchestrators that replace the underlying tensors
+        with fresh, untransformed bytes are responsible for resetting that flag.
+
+        Returns:
+            None.
+        """
+
+    def cache_derived_state(self) -> None:
+        """Recompute Python-side state derived from currently loaded weights.
+
+        This stage is reserved for idempotent recomputation from real tensors,
+        such as cached scalars, validation results, or fingerprints. It should
+        not perform one-shot weight transforms. Callers may run it after weight
+        bytes arrive from any loading or sharing mechanism.
+
+        Returns:
+            None.
+        """
+
+    def post_load_weights(self) -> None:
+        """Run the default staged post-load hook sequence.
+
+        Existing model-loading paths continue to call this method for backward
+        compatibility. More specialized loaders can call individual stages when
+        they need a subset of alias setup, tensor transformation, or derived
+        state recomputation.
+
+        Returns:
+            None.
+        """
+        self.setup_aliases()
+        self.transform_weights()
+        self.cache_derived_state()
+
     def forward(
         self,
         attn_metadata: AttentionMetadata,

tensorrt_llm/_torch/pyexecutor/model_loader.py

@@ -717,10 +717,7 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
                     checkpoint_dir=checkpoint_dir,
                     weights_preloaded=weights_preloaded)
 
-                for module in model.modules():
-                    if hasattr(module, 'post_load_weights') and not getattr(
-                            module, '_weights_removed', False):
-                        module.post_load_weights()
+                self._walk_full_post_load(model)
 
             # TODO(GMS-MOE-LB): when the (MoE, GMS) combination is enabled,
             # ``register_weight_slots_after_to_cuda`` and ``finalize_model``
@@ -743,10 +740,126 @@ def init_meta_tensor_in_pool(t: torch.Tensor):
 
         return model, moe_load_balancer
 
+    @staticmethod
+    def _setup_aliases(model: DecoderModelForCausalLM) -> None:
+        """Run top-level structural alias setup if the model defines it.
+
+        Alias wiring is a model-level concern. It is intentionally not a
+        recursive module walk, because migrated aliases are expected to be set
+        by the root model that owns the layer graph.
+
+        Args:
+            model: Root decoder model whose top-level alias hook should run.
+
+        Returns:
+            None.
+        """
+        setup_aliases: Optional[Callable[[], None]] = getattr(
+            model, 'setup_aliases', None)
+        if setup_aliases is not None:
+            setup_aliases()
+
+    @staticmethod
+    def _walk_transform(model: DecoderModelForCausalLM) -> None:
+        """Run one-shot weight transforms on eligible modules.
+
+        The walk is duck-typed so modules can opt in without inheriting a shared
+        base class. Modules whose weights were removed are skipped, and modules
+        already marked ``_weights_transformed`` are left untouched until an
+        orchestrator resets the flag after rebinding fresh weight bytes.
+
+        Args:
+            model: Root decoder model whose module tree should be visited.
+
+        Returns:
+            None.
+        """
+        for module in model.modules():
+            transform_weights: Optional[Callable[[], None]] = getattr(
+                module, 'transform_weights', None)
+            if transform_weights is not None and not getattr(
+                    module, '_weights_removed', False) and not getattr(
+                        module, '_weights_transformed', False):
+                transform_weights()
+
+    @staticmethod
+    def _walk_cache_state(model: DecoderModelForCausalLM) -> None:
+        """Recompute derived Python-side state on eligible modules.
+
+        This walk is separate from weight transforms so callers that already
+        have transformed weight bytes can refresh local Python state without
+        mutating tensor layouts again.
+
+        Args:
+            model: Root decoder model whose module tree should be visited.
+
+        Returns:
+            None.
+        """
+        for module in model.modules():
+            cache_derived_state: Optional[Callable[[], None]] = getattr(
+                module, 'cache_derived_state', None)
+            if cache_derived_state is not None and not getattr(
+                    module, '_weights_removed', False):
+                cache_derived_state()
+
+    @staticmethod
+    def _walk_full_post_load(model: DecoderModelForCausalLM) -> None:
+        """Run the backward-compatible post-load hook on eligible modules.
+
+        This preserves the previous ``ModelLoader`` behavior for standard load
+        paths while staged-hook migration proceeds incrementally.
+
+        Args:
+            model: Root decoder model whose module tree should be visited.
+
+        Returns:
+            None.
+        """
+        for module in model.modules():
+            post_load_weights: Optional[Callable[[], None]] = getattr(
+                module, 'post_load_weights', None)
+            if post_load_weights is not None and not getattr(
+                    module, '_weights_removed', False):
+                post_load_weights()
+
+    @staticmethod
+    def _reset_weights_transformed(model: DecoderModelForCausalLM) -> None:
+        """Mark transformed modules as needing a new transform pass.
+
+        Orchestrators call this before rebinding fresh, untransformed weights.
+        The reset only touches modules that already carry the flag so unrelated
+        modules do not grow staged-hook state eagerly.
+
+        Args:
+            model: Root decoder model whose module tree should be visited.
+
+        Returns:
+            None.
+        """
+        for module in model.modules():
+            if hasattr(module, '_weights_transformed'):
+                module._weights_transformed = False
+
     def reload(self,
                model: DecoderModelForCausalLM,
                weights: dict,
-               allow_partial_loading: bool = False):
+               allow_partial_loading: bool = False) -> None:
+        """Reload model weights without running post-load hooks.
+
+        Reload is used by incremental update paths that may provide only a
+        partial set of replacement weights. The owner of the update lifecycle is
+        responsible for running post-load processing once all bytes are present.
+
+        Args:
+            model: Model instance receiving the replacement weights.
+            weights: Checkpoint weights to pass to ``model.load_weights``.
+            allow_partial_loading: Whether missing replacement weights are
+                allowed by models that support partial loading.
+
+        Returns:
+            None.
+        """
         if self.weight_mapper is None:
             raise RuntimeError(
                 "Cannot reload weights: weight_mapper was not initialized. "

tests/integration/test_lists/test-db/l0_sanity_check.yml

@@ -34,6 +34,7 @@ l0_sanity_check:
     - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[UCX-mha-ctx_fp16_gen_fp16]
     - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mha]
     - unittest/others/test_kv_cache_transceiver.py::test_cancel_request_in_transmission[mla]
+    - unittest/_torch/pyexecutor/test_model_loader_mx.py
 - condition:
     ranges:
       system_gpu_count:

tests/unittest/_torch/pyexecutor/test_model_loader_mx.py

@@ -114,6 +114,7 @@ def test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works(mon
     # initialize it even though the initial load skipped _call_load_weights.
     loader.reload(model, {"reloaded": MagicMock()})
     assert loader._call_load_weights.call_count == 1
+    assert events == ["post_load_weights", "load_weights"]
 
 
 def test_mx_partial_fallback_merges_returned_weights(monkeypatch):
@@ -154,3 +155,85 @@ def test_mx_fallback_runs_standard_weight_mapping(monkeypatch):
     checkpoint_loader.post_load_publish.assert_called_once_with(
         model, checkpoint_dir="/ckpt", weights_preloaded=False
     )
+
+
+class _HookRecorder(nn.Module):
+    def __init__(
+        self,
+        name: str,
+        events: list[tuple[str, str]],
+        *,
+        removed: bool | None = None,
+        transformed: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.name = name
+        self.events = events
+        if removed is not None:
+            self._weights_removed = removed
+        if transformed is not None:
+            self._weights_transformed = transformed
+
+    def setup_aliases(self):
+        self.events.append((self.name, "setup_aliases"))
+
+    def transform_weights(self):
+        self.events.append((self.name, "transform_weights"))
+        self._weights_transformed = True
+
+    def cache_derived_state(self):
+        self.events.append((self.name, "cache_derived_state"))
+
+    def post_load_weights(self):
+        self.events.append((self.name, "post_load_weights"))
+
+
+class _HookModel(_HookRecorder):
+    def __init__(self, events):
+        super().__init__("model", events)
+        self.child = _HookRecorder("child", events)
+        self.transformed_child = _HookRecorder("transformed_child", events, transformed=True)
+        self.removed_child = _HookRecorder("removed_child", events, removed=True)
+
+
+def test_staged_hook_setup_aliases_is_top_level_only():
+    events = []
+    model = _HookModel(events)
+
+    ModelLoader._setup_aliases(model)
+
+    assert events == [("model", "setup_aliases")]
+
+
+def test_staged_hook_walks_skip_removed_and_transformed_modules():
+    events = []
+    model = _HookModel(events)
+
+    ModelLoader._walk_transform(model)
+    ModelLoader._walk_cache_state(model)
+    ModelLoader._walk_full_post_load(model)
+
+    assert events == [
+        ("model", "transform_weights"),
+        ("child", "transform_weights"),
+        ("model", "cache_derived_state"),
+        ("child", "cache_derived_state"),
+        ("transformed_child", "cache_derived_state"),
+        ("model", "post_load_weights"),
+        ("child", "post_load_weights"),
+        ("transformed_child", "post_load_weights"),
+    ]
+
+
+def test_reset_weights_transformed_only_resets_existing_flags():
+    events = []
+    model = _HookModel(events)
+    model._weights_transformed = True
+    model.child._weights_transformed = True
+
+    ModelLoader._reset_weights_transformed(model)
+
+    assert model._weights_transformed is False
+    assert model.child._weights_transformed is False
+    assert model.transformed_child._weights_transformed is False
+    assert not hasattr(model.removed_child, "_weights_transformed")