[TRTLLM-9293][feat] Enable partial weight loading to support streaming update weights

tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py

@@ -114,19 +114,29 @@ def preprocess_weights(self, weights: dict) -> dict:
         """
         ...
 
-    def handle_manual_copy(self, module_name: str, module_weights: dict, n: str,
-                           p: nn.Parameter) -> None:
-        p.data.copy_(module_weights[n][:])
+    def handle_manual_copy(self,
+                           module_name: str,
+                           module_weights: dict,
+                           n: str,
+                           p: nn.Parameter,
+                           allow_partial_loading: bool = False) -> None:
+        if not allow_partial_loading:
+            assert n in module_weights
+        if n in module_weights:
+            p.data.copy_(module_weights[n][:])
 
     def does_require_special_handling(self, module_name: str) -> bool:
         return module_name in self.mapping
 
     def is_special_instance_module(self, module: nn.Module) -> bool:
         return False
 
-    def handle_special_instance_module(self, module: nn.Module,
-                                       module_name: str,
-                                       module_weights: dict) -> None:
+    def handle_special_instance_module(
+            self,
+            module: nn.Module,
+            module_name: str,
+            module_weights: dict,
+            allow_partial_loading: bool = False) -> None:
         raise NotImplementedError()
 
     @property

tensorrt_llm/_torch/models/checkpoints/hf/gemma3_weight_mapper.py

@@ -27,9 +27,21 @@ def should_skip_module(self, module_name: str) -> bool:
         return any(skip_module in module_name
                    for skip_module in self._skip_modules)
 
-    def handle_manual_copy(self, module_name: str, module_weights: dict, n: str,
-                           p: nn.Parameter) -> None:
+    def handle_manual_copy(self,
+                           module_name: str,
+                           module_weights: dict,
+                           n: str,
+                           p: nn.Parameter,
+                           allow_partial_loading: bool = False) -> None:
         if 'norm' in module_name:
-            p.data.copy_(module_weights[n][:] + 1)
+            if not allow_partial_loading:
+                assert n in module_weights
+            if n in module_weights:
+                p.data.copy_(module_weights[n][:] + 1)
         else:
-            super().handle_manual_copy(module_name, module_weights, n, p)
+            super().handle_manual_copy(
+                module_name,
+                module_weights,
+                n,
+                p,
+                allow_partial_loading=allow_partial_loading)

tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py

@@ -12,9 +12,12 @@ class Qwen2MoeHfWeightMapper(HfWeightMapper):
     def is_special_instance_module(self, module: nn.Module) -> bool:
         return isinstance(module, MoE)
 
-    def handle_special_instance_module(self, module: nn.Module,
-                                       module_name: str,
-                                       module_weights: dict) -> None:
+    def handle_special_instance_module(
+            self,
+            module: nn.Module,
+            module_name: str,
+            module_weights: dict,
+            allow_partial_loading: bool = False) -> None:
         if isinstance(module, MoE):
             updated_module_weights = {}
             for weight_name, weight_value in module_weights.items():
@@ -23,4 +26,5 @@ def handle_special_instance_module(self, module: nn.Module,
                                                "w3").replace("down_proj", "w2")
                 updated_module_weights[new_weight_name] = weight_value
             del module_weights
-            module.load_weights(weights=[updated_module_weights])
+            module.load_weights(weights=[updated_module_weights],
+                                allow_partial_loading=allow_partial_loading)

tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py

@@ -60,11 +60,15 @@ def should_skip_module(self, module_name: str) -> bool:
     def _duplicate_kv_weights(self, module: nn.Module, new_name: str,
                               weights: dict):
         if new_name in ['k_proj', 'v_proj']:
+            if "weight" not in weights and "bias" not in weights:
+                return weights
             # k_proj and v_proj shape is [num_kv_heads*head_dim, hidden_dim]
+            kv_shape = weights['weight'].shape[
+                0] if "weight" in weights else weights['bias'].shape[0]
             if isinstance(module.quant_method, W4A16_AWQ_LinearMethod):
-                num_kv_heads = weights['weight'].shape[0] * 2 // self._head_dim
+                num_kv_heads = kv_shape * 2 // self._head_dim
             else:
-                num_kv_heads = weights['weight'].shape[0] // self._head_dim
+                num_kv_heads = kv_shape // self._head_dim
             processed_weights = {
                 k:
                 self._duplicate_kv(weight=v[:],

tensorrt_llm/_torch/models/modeling_llama_min_latency.py

@@ -384,6 +384,11 @@ def __init__(
                 and self.floor_scale == 8192.0 \
                 and self.attn_scale == 0.1
 
+            qkv_shard_indices_mapping = {
+                "q": (0, self.q_size),
+                "k": (self.q_size, self.kv_size),
+                "v": (self.q_size + self.kv_size, self.kv_size),
+            }
             # When min-latency QKV gemm is enabled, override qkv_proj.
             self.qkv_proj = Llama4MinLatencyLinear(
                 self.hidden_size,
@@ -400,6 +405,7 @@ def __init__(
                 enable_fused_gemm_attn_scaling=self.
                 enable_fused_gemm_attn_scaling,
                 enable_trtllm_gen=True,
+                fused_weight_shard_indices_mapping=qkv_shard_indices_mapping,
             )
 
     def _forward_nope(

tensorrt_llm/_torch/models/modeling_qwen3_moe.py

@@ -58,10 +58,15 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             hidden_states, self.weight.t(), bias=None, out_dtype=self.out_dtype)
         return logits
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
         assert len(weights) == 1
-
-        self.weight.copy_(weights[0]["weight"][:])
+        w = weights[0].get("weight")
+        if not allow_partial_loading:
+            assert w is not None, "Qwen3Gate expects weight when partial loading is disabled"
+        if w is not None:
+            self.weight.copy_(w[:])
 
     @property
     def routing_method(self) -> BaseMoeRoutingMethod:

tensorrt_llm/_torch/models/modeling_speculative.py

@@ -55,6 +55,11 @@ def __init__(
         # Override the QKV projection. The number of input features
         # is twice as big for EAGLE3 draft models.
         if not self._next_layer_regular:
+            qkv_shard_indices_mapping = {
+                "q": (0, self.q_size),
+                "k": (self.q_size, self.kv_size),
+                "v": (self.q_size + self.kv_size, self.kv_size),
+            }
             self.qkv_proj = Linear(
                 2 * self.hidden_size,
                 tp_size * self.q_size + 2 * tp_size * self.kv_size,
@@ -67,6 +72,7 @@ def __init__(
                 quant_config=model_config.get_quant_config(),
                 skip_create_weights_in_init=model_config.
                 skip_create_weights_in_init,
+                fused_weight_shard_indices_mapping=qkv_shard_indices_mapping,
             )
 
 
@@ -642,10 +648,12 @@ def forward(
 
     def load_weights(self,
                      weights: Dict,
-                     weight_mapper: Optional[BaseWeightMapper] = None):
+                     weight_mapper: Optional[BaseWeightMapper] = None,
+                     allow_partial_loading: bool = False):
         super().load_weights(weights=weights,
                              weight_mapper=weight_mapper,
-                             skip_modules=["draft_model"])
+                             skip_modules=["draft_model"],
+                             allow_partial_loading=allow_partial_loading)
 
     def load_draft_weights(self,
                            weights: Dict,

tensorrt_llm/_torch/models/modeling_utils.py

@@ -1,4 +1,5 @@
 import contextlib
+import inspect
 import math
 import os
 import time
@@ -557,7 +558,8 @@ def forward(
     def load_weights(self,
                      weights: Dict,
                      weight_mapper: Optional["BaseWeightMapper"] = None,
-                     skip_modules: List[str] = []):
+                     skip_modules: List[str] = [],
+                     allow_partial_loading: bool = False):
         # TODO smor- this solution is a temporary solution to load weights while we are still using
         # the old checkpoint format loading process. Once checkpoint format is unified
         # this method will be removed.
@@ -566,13 +568,15 @@ def load_weights(self,
             _load_weights_impl(self,
                                weights,
                                skip_modules,
-                               preload_weight_modules=preload_weight_modules)
+                               preload_weight_modules=preload_weight_modules,
+                               allow_partial_loading=allow_partial_loading)
         else:
             _load_weights_impl_v2(self,
                                   weights,
                                   weight_mapper,
                                   skip_modules,
-                                  preload_weight_modules=preload_weight_modules)
+                                  preload_weight_modules=preload_weight_modules,
+                                  allow_partial_loading=allow_partial_loading)
 
     def infer_max_seq_len(self) -> int:
         # Modified from tensorrt_llm/builder.py _init_max_seq_len
@@ -817,7 +821,8 @@ def _load_weights_impl(model: Union[nn.Module, DecoderModelForCausalLM],
                        weights: Dict,
                        skip_modules: List[str] = [],
                        params_map: Optional[Dict[str, str]] = None,
-                       preload_weight_modules: Optional[List[str]] = None):
+                       preload_weight_modules: Optional[List[str]] = None,
+                       allow_partial_loading: bool = False):
     # TODO: remove preload_weight_modules - it is a workaround for min-latency llama4 model loading where
     # we need some order in the module loading. Once this is resolved, we can remove this workaround.
     # TODO smor- this method is here as a temporary solution to load weights.
@@ -871,8 +876,6 @@ def load_single_module(name, module):
                 for new_name in params_map[names[-1]]:
                     fw = filter_weights('.'.join(names[:-1] + [new_name]),
                                         weights)
-                    if not fw:
-                        continue
                     if new_name in ['k_proj', 'v_proj']:
                         num_kv_heads_list = [num_kv_heads
                                              ] * len(fw) if isinstance(
@@ -887,21 +890,29 @@ def load_single_module(name, module):
                             if k in ["weight", "bias"] else v
                             for i, (k, v) in enumerate(fw.items())
                         }
-
                     module_weights.append(fw)
-                # Note: module_weights may be empty after filtering (e.g., in streaming weight updates)
-                if module_weights:
-                    module.load_weights(weights=module_weights)
+                module.load_weights(weights=module_weights,
+                                    allow_partial_loading=allow_partial_loading)
 
             else:
                 module_weights = filter_weights(name, weights)
                 # Note: module_weights may be empty after filtering (e.g., in streaming weight updates)
                 if module_weights:
                     if hasattr(module, 'load_weights'):
-                        module.load_weights(weights=[module_weights])
+                        args = inspect.getfullargspec(module.load_weights).args
+                        if "allow_partial_loading" not in args:
+                            assert not allow_partial_loading, "allow_partial_loading is not supported for this model"
+                            module.load_weights(weights=[module_weights])
+                        else:
+                            module.load_weights(
+                                weights=[module_weights],
+                                allow_partial_loading=allow_partial_loading)
                     else:
                         for n, p in module.named_parameters(recurse=False):
-                            p.data.copy_(module_weights[n][:])
+                            if not allow_partial_loading:
+                                assert n in module_weights
+                            if n in module_weights:
+                                p.data.copy_(module_weights[n][:])
 
     if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL",
                       "True") in ["True", "true", "1", "yes", "y"]:
@@ -942,7 +953,8 @@ def _load_weights_impl_v2(model: Union[nn.Module, DecoderModelForCausalLM],
                           weight_mapper: "BaseWeightMapper",
                           skip_modules: List[str] = [],
                           params_map: Optional[Dict[str, str]] = None,
-                          preload_weight_modules: Optional[List[str]] = None):
+                          preload_weight_modules: Optional[List[str]] = None,
+                          allow_partial_loading: bool = False):
     # TODO: remove preload_weight_modules - it is a workaround for min-latency llama4 and Qwen3 model loading where
     # we need some order in the module loading. Once this is resolved, we can remove this workaround.
     weight_mapper.add_skip_modules(skip_modules)
@@ -961,26 +973,38 @@ def load_single_module(name, module):
             if weight_mapper.does_require_special_handling(module_name):
                 module_weights = weight_mapper.apply_callbacks(
                     module, module_name, module_names_breakdown, weights)
-                # Note: module_weights may be empty after filtering (e.g., in streaming weight updates)
-                if module_weights:
-                    module.load_weights(weights=module_weights)
+                module.load_weights(weights=module_weights,
+                                    allow_partial_loading=allow_partial_loading)
             else:
                 module_weights = weight_mapper.filter_weights(name, weights)
                 # Note: module_weights may be empty after filtering (e.g., in streaming weight updates)
                 if module_weights:
                     if weight_mapper.is_special_instance_module(module):
                         weight_mapper.handle_special_instance_module(
-                            module, module_name, module_weights)
+                            module,
+                            module_name,
+                            module_weights,
+                            allow_partial_loading=allow_partial_loading)
                     elif hasattr(module, 'load_weights'):
-                        if module_weights:
-                            if "linear_attn.conv1d" in name:
-                                module_weights['weight'] = module_weights[
-                                    'weight'].squeeze(dim=1)
+                        if "linear_attn.conv1d" in name:
+                            module_weights['weight'] = module_weights[
+                                'weight'].squeeze(dim=1)
+                        args = inspect.getfullargspec(module.load_weights).args
+                        if "allow_partial_loading" not in args:
+                            assert not allow_partial_loading, "allow_partial_loading is not supported for this model"
                             module.load_weights(weights=[module_weights])
+                        else:
+                            module.load_weights(
+                                weights=[module_weights],
+                                allow_partial_loading=allow_partial_loading)
                     else:
                         for n, p in module.named_parameters(recurse=False):
                             weight_mapper.handle_manual_copy(
-                                module_name, module_weights, n, p)
+                                module_name,
+                                module_weights,
+                                n,
+                                p,
+                                allow_partial_loading=allow_partial_loading)
 
     if os.environ.get("TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL",
                       "True") in ["True", "true", "1", "yes", "y"]:

tensorrt_llm/_torch/modules/attention.py

@@ -234,6 +234,15 @@ def __init__(
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_key_value_heads * self.head_dim
 
+        qkv_shard_indices_mapping = {
+            "q": (0, self.q_size * (2 if self.attn_output_gate else 1)),
+            "k":
+            (self.q_size * (2 if self.attn_output_gate else 1), self.kv_size),
+            "v":
+            (self.q_size * (2 if self.attn_output_gate else 1) + self.kv_size,
+             self.kv_size),
+        }
+
         self.qkv_proj = Linear(
             self.hidden_size,
             tp_size * self.q_size * (2 if self.attn_output_gate else 1) +
@@ -249,7 +258,8 @@ def __init__(
             allreduce_strategy=config.allreduce_strategy,
             force_dynamic_quantization=config.force_dynamic_quantization,
             disable_deep_gemm=disable_deep_gemm,
-            use_custom_cublas_mm=use_custom_cublas_mm)
+            use_custom_cublas_mm=use_custom_cublas_mm,
+            fused_weight_shard_indices_mapping=qkv_shard_indices_mapping)
 
         self.o_lora = LoraLayer([LoraModuleType.ATTENTION_DENSE],
                                 [self.hidden_size])

tensorrt_llm/_torch/modules/embedding.py

@@ -120,14 +120,17 @@ def skip_forward(
         output = input.new_empty(output_shape)
         return output
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
         original_weight = None
         if self.tp_mode == TensorParallelMode.COLUMN:
             if self.tp_rank == self.tp_size - 1 and self.padding_size > 0:
                 original_weight = self.weight.data.zero_()
                 self.weight.data = self.weight[:-self.padding_size, :]
 
-        super().load_weights(weights)
+        super().load_weights(weights,
+                             allow_partial_loading=allow_partial_loading)
 
         if original_weight is not None:
             self.weight.data = original_weight