diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 1b79818d4288..69256521d364 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/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, diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 5e98e730b6b8..c0c6c9f28c58 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/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. " diff --git a/tests/integration/test_lists/test-db/l0_sanity_check.yml b/tests/integration/test_lists/test-db/l0_sanity_check.yml index 231e545f50d4..792387bc6db3 100644 --- a/tests/integration/test_lists/test-db/l0_sanity_check.yml +++ b/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: diff --git a/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py b/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py index ae790dd04255..d887e5e0f0d6 100644 --- a/tests/unittest/_torch/pyexecutor/test_model_loader_mx.py +++ b/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")