diff --git a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py
index cf9b055c99f1..4d78b3dcb19b 100644
--- a/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py
+++ b/tensorrt_llm/_torch/models/checkpoints/base_weight_mapper.py
@@ -114,9 +114,16 @@ 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
@@ -124,9 +131,12 @@ def does_require_special_handling(self, module_name: str) -> bool:
     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
diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/gemma3_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/gemma3_weight_mapper.py
index a8d31d6526d9..8382588dc243 100644
--- a/tensorrt_llm/_torch/models/checkpoints/hf/gemma3_weight_mapper.py
+++ b/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)
diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py
index 41b40042c762..35234afe96fb 100644
--- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen2_moe_weight_mapper.py
+++ b/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)
diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py
index 63f812a139c1..9afcfd2842af 100644
--- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py
+++ b/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[:],
diff --git a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py b/tensorrt_llm/_torch/models/modeling_llama_min_latency.py
index c4cc71bc0bfa..9af2d99f0b83 100644
--- a/tensorrt_llm/_torch/models/modeling_llama_min_latency.py
+++ b/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(
diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py b/tensorrt_llm/_torch/models/modeling_qwen3_moe.py
index 4bf631279415..afaf31138a2f 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen3_moe.py
+++ b/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:
diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py
index a812f54ef4b0..fc39df56f534 100755
--- a/tensorrt_llm/_torch/models/modeling_speculative.py
+++ b/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,
diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py
index c31e7e70ee62..67f5f9223706 100755
--- a/tensorrt_llm/_torch/models/modeling_utils.py
+++ b/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"]:
diff --git a/tensorrt_llm/_torch/modules/attention.py b/tensorrt_llm/_torch/modules/attention.py
index 05e7924b7203..dfe84bdecc5f 100644
--- a/tensorrt_llm/_torch/modules/attention.py
+++ b/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])
diff --git a/tensorrt_llm/_torch/modules/embedding.py b/tensorrt_llm/_torch/modules/embedding.py
index badce20b44ca..046954da46d3 100644
--- a/tensorrt_llm/_torch/modules/embedding.py
+++ b/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
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
index 80a4475e3c6c..94bc3f17b2fe 100755
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py
@@ -13,6 +13,7 @@
 from ...model_config import ModelConfig
 from ...utils import AuxStreamType, EventType, Fp4QuantizedTensor, ceil_div
 from .interface import AlltoallMethodType, MoE
+from .quantization import UnquantizedFusedMoEMethod
 
 # isort: off
 from .quantization import (
@@ -767,12 +768,23 @@ def forward_fake(
             **kwargs,
         )
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
         assert self._weights_created
         assert len(weights) == 1
         weights = weights[0]
 
-        self.quant_method.load_weights(self, weights, self.weight_loading_mode)
+        if not isinstance(self.quant_method, UnquantizedFusedMoEMethod):
+            assert not allow_partial_loading, "Partial loading is not supported for quantized MoE now"
+            self.quant_method.load_weights(self, weights,
+                                           self.weight_loading_mode)
+        else:
+            self.quant_method.load_weights(
+                self,
+                weights,
+                self.weight_loading_mode,
+                allow_partial_loading=allow_partial_loading)
 
     def post_load_weights(self):
         self.quant_method.post_load_weights(self)
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py
index 3a08ffc87438..23450d264f4a 100755
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py
@@ -1389,12 +1389,19 @@ def forward_impl(
 
         return final_hidden_states
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
+        assert not allow_partial_loading, "Partial loading is not supported for TritonFusedMoE now"
         assert self._weights_created
         assert len(weights) == 1
         weights = weights[0]
 
-        self.quant_method.load_weights(self, weights, self.weight_loading_mode)
+        self.quant_method.load_weights(
+            self,
+            weights,
+            self.weight_loading_mode,
+            allow_partial_loading=allow_partial_loading)
 
     def post_load_weights(self):
         self.quant_method.post_load_weights(self)
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
index fbfd7808e3c0..2c2268b6dbe5 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
@@ -17,12 +17,14 @@
 from ...model_config import ModelConfig
 from ...utils import AuxStreamType, Fp4QuantizedTensor, ceil_div
 from .interface import AlltoallMethodType, MoE, MoEWeightLoadingMode
-from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethod,
-                           NVFP4TRTLLMGenFusedMoEMethod,
-                           W4A8MXFP4FP8TRTLLMGenFusedMoEMethod,
-                           W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod,
-                           W4A8NVFP4FP8TRTLLMGenFusedMoEMethod,
-                           W4A16MXFP4TRTLLMGenFusedMoEMethod)
+
+# isort: off
+from .quantization import (
+    DeepSeekFP8BlockScalesFusedMoEMethod, NVFP4TRTLLMGenFusedMoEMethod,
+    UnquantizedFusedMoEMethod, W4A8MXFP4FP8TRTLLMGenFusedMoEMethod,
+    W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod, W4A8NVFP4FP8TRTLLMGenFusedMoEMethod,
+    W4A16MXFP4TRTLLMGenFusedMoEMethod)
+# isort: on
 from .routing import BaseMoeRoutingMethod, DeepSeekV3MoeRoutingMethod
 
 
@@ -256,13 +258,24 @@ def create_weights(self):
                                         requires_grad=False)
             self.register_parameter("w2_bias", self.w2_bias)
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
         assert self._weights_created
 
         assert len(weights) == 1
         weights = weights[0]
 
-        self.quant_method.load_weights(self, weights, self.weight_loading_mode)
+        if not isinstance(self.quant_method, UnquantizedFusedMoEMethod):
+            assert not allow_partial_loading, "Partial loading is not supported for quantized MoE now"
+            self.quant_method.load_weights(self, weights,
+                                           self.weight_loading_mode)
+        else:
+            self.quant_method.load_weights(
+                self,
+                weights,
+                self.weight_loading_mode,
+                allow_partial_loading=allow_partial_loading)
 
     def post_load_weights(self):
         self.quant_method.post_load_weights(self)
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py
index ed6f11993b20..a13ff07bad8a 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py
@@ -416,9 +416,12 @@ def pack_params(self, experts, module_name: str, weight_name: str):
         packed_weight = packed_weight.view(len(weights), *weights_data[0].shape)
         getattr(self, f"{module_name}_{weight_name}").data = packed_weight
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
         from ...models.modeling_utils import filter_weights
 
+        assert not allow_partial_loading, "Partial loading is not supported for vanilla MoE now"
         assert self._weights_created
         assert len(weights) == 1
         weights = weights[0]
diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
index 7f1e819484a8..b46e96ddf7db 100755
--- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
@@ -940,12 +940,23 @@ def _supports_load_balancer(self) -> bool:
         """WideEPMoE supports load balancer."""
         return True
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
         assert self._weights_created
         assert len(weights) == 1
         weights = weights[0]
 
-        self.quant_method.load_weights(self, weights, self.weight_loading_mode)
+        if not isinstance(self.quant_method, UnquantizedFusedMoEMethod):
+            assert not allow_partial_loading, "Partial loading is not supported for quantized MoE now"
+            self.quant_method.load_weights(self, weights,
+                                           self.weight_loading_mode)
+        else:
+            self.quant_method.load_weights(
+                self,
+                weights,
+                self.weight_loading_mode,
+                allow_partial_loading=allow_partial_loading)
 
     def post_load_weights(self):
         self.quant_method.post_load_weights(self)
diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py
index 4c6a07333d33..388b1a51d6f7 100644
--- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py
+++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py
@@ -1,3 +1,4 @@
+import inspect
 import math
 from abc import ABC, abstractmethod
 from typing import Dict, List, NamedTuple, Optional, Union
@@ -237,12 +238,24 @@ def create_weights(
             module.w2_bias = None
 
     def load_expert_weights_to_dst(
-            self, module: torch.nn.Module, weights: List[Dict],
+            self,
+            module: torch.nn.Module,
+            weights: List[Dict],
             weight_loading_mode: MoEWeightLoadingMode,
-            load_expert_ids: List[int], dst_w3_w1_weights_tensor: torch.Tensor,
+            load_expert_ids: List[int],
+            dst_w3_w1_weights_tensor: torch.Tensor,
             dst_w2_weights_tensor: torch.Tensor,
             dst_w3_w1_bias_tensor: Optional[torch.Tensor],
-            dst_w2_bias_tensor: Optional[torch.Tensor]):
+            dst_w2_bias_tensor: Optional[torch.Tensor],
+            allow_partial_loading: bool = False):
+        w3_w1_kargs = {}
+        w2_kargs = {}
+        w3_w1_args = inspect.getfullargspec(self.load_expert_w3_w1_weight).args
+        w2_args = inspect.getfullargspec(self.load_expert_w2_weight).args
+        if "allow_partial_loading" in w3_w1_args:
+            w3_w1_kargs["allow_partial_loading"] = allow_partial_loading
+        if "allow_partial_loading" in w2_args:
+            w2_kargs["allow_partial_loading"] = allow_partial_loading
         # Multithread weight load is superseded by prefetch_files() in model_engine.py
         # Also, threading adds overhead in order to protect shuffle index cache with critical section.
         for local_slot_id, expert_id in enumerate(load_expert_ids):
@@ -253,107 +266,170 @@ def load_expert_weights_to_dst(
                     MoEWeightLoadingMode.VANILLA,
                     MoEWeightLoadingMode.W4A8_CUSTOM
             ]:
-                w1_weight = weights[f"{expert_id}.w1.weight"]
-                w3_weight = weights[f"{expert_id}.w3.weight"]
-                w2_weight = weights[f"{expert_id}.w2.weight"]
+                w1_weight = weights[
+                    f"{expert_id}.w1.weight"] if f"{expert_id}.w1.weight" in weights else None
+                w3_weight = weights[
+                    f"{expert_id}.w3.weight"] if f"{expert_id}.w3.weight" in weights else None
+                w2_weight = weights[
+                    f"{expert_id}.w2.weight"] if f"{expert_id}.w2.weight" in weights else None
                 if module.bias:
-                    w1_bias = weights[f"{expert_id}.w1.bias"]
-                    w3_bias = weights[f"{expert_id}.w3.bias"]
-                    w2_bias = weights[f"{expert_id}.w2.bias"]
+                    w1_bias = weights[
+                        f"{expert_id}.w1.bias"] if f"{expert_id}.w1.bias" in weights else None
+                    w3_bias = weights[
+                        f"{expert_id}.w3.bias"] if f"{expert_id}.w3.bias" in weights else None
+                    w2_bias = weights[
+                        f"{expert_id}.w2.bias"] if f"{expert_id}.w2.bias" in weights else None
             elif weight_loading_mode == MoEWeightLoadingMode.FUSED_GATE_UP_PROJ:
-                w1_w3_weight = weights["gate_up_proj"][expert_id].transpose(
-                    0, 1)
-                w1_weight, w3_weight = w1_w3_weight.chunk(2, dim=0)
+                w1_weight, w3_weight = None, None
+                if "gate_up_proj" in weights:
+                    w1_w3_weight = weights["gate_up_proj"][expert_id].transpose(
+                        0, 1)
+                    w1_weight, w3_weight = w1_w3_weight.chunk(2, dim=0)
                 w2_weight = weights["down_proj"][expert_id].transpose(
-                    0, 1).contiguous()
+                    0, 1).contiguous() if "down_proj" in weights else None
                 if module.bias:
-                    w1_w3_bias = weights["gate_up_proj.bias"][expert_id]
-                    w1_bias, w3_bias = w1_w3_bias.chunk(2, dim=0)
-                    w2_bias = weights["down_proj.bias"][expert_id]
+                    w1_bias, w3_bias = None, None
+                    if "gate_up_proj.bias" in weights:
+                        w1_w3_bias = weights["gate_up_proj.bias"][expert_id]
+                        w1_bias, w3_bias = w1_w3_bias.chunk(2, dim=0)
+                    if "down_proj.bias" in weights:
+                        w2_bias = weights["down_proj.bias"][expert_id]
             else:
                 raise NotImplementedError(
                     f"Unknown weight loading mode in MoE: {weight_loading_mode}"
                 )
 
             self.load_expert_w3_w1_weight(module, w1_weight, w3_weight,
-                                          dst_w3_w1_weights_tensor[expert_idx])
+                                          dst_w3_w1_weights_tensor[expert_idx],
+                                          **w3_w1_kargs)
 
             self.load_expert_w2_weight(module, w2_weight,
-                                       dst_w2_weights_tensor[expert_idx])
-            module._add_raw_shared_weights_for_unmap(
-                [w1_weight, w3_weight, w2_weight])
+                                       dst_w2_weights_tensor[expert_idx],
+                                       **w2_kargs)
+            unmap_weights = [
+                weight for weight in [w1_weight, w3_weight, w2_weight]
+                if weight is not None
+            ]
+            module._add_raw_shared_weights_for_unmap(unmap_weights)
 
             if module.bias:
                 self.load_expert_w3_w1_weight(
                     module, w1_bias, w3_bias,
-                    dst_w3_w1_bias_tensor.data[expert_idx])
+                    dst_w3_w1_bias_tensor.data[expert_idx], **w3_w1_kargs)
 
                 self.load_expert_w2_weight(module, w2_bias,
-                                           dst_w2_bias_tensor.data[expert_idx])
-                module._add_raw_shared_weights_for_unmap(
-                    [w1_bias, w3_bias, w2_bias])
-
-    def load_weights(self, module: torch.nn.Module, weights: List[Dict],
-                     weight_loading_mode: MoEWeightLoadingMode):
+                                           dst_w2_bias_tensor.data[expert_idx],
+                                           **w2_kargs)
+                unmap_weights = [
+                    weight for weight in [w1_bias, w3_bias, w2_bias]
+                    if weight is not None
+                ]
+                module._add_raw_shared_weights_for_unmap(unmap_weights)
+
+    def load_weights(self,
+                     module: torch.nn.Module,
+                     weights: List[Dict],
+                     weight_loading_mode: MoEWeightLoadingMode,
+                     allow_partial_loading: bool = False):
 
         self.load_expert_weights_to_dst(
-            module, weights, weight_loading_mode,
-            module.initial_local_expert_ids, module.w3_w1_weight.data,
+            module,
+            weights,
+            weight_loading_mode,
+            module.initial_local_expert_ids,
+            module.w3_w1_weight.data,
             module.w2_weight.data,
             module.w3_w1_bias.data if module.bias else None,
-            module.w2_bias.data if module.bias else None)
+            module.w2_bias.data if module.bias else None,
+            allow_partial_loading=allow_partial_loading)
 
         self.load_quant_scales(module, weights)
 
         if self.need_load_shared_weights(module):
             local_shared_load_expert_ids = module.layer_load_balancer.get_load_expert_ids(
             )
-            local_shared_w3_w1_tensors = torch.empty(
-                (len(local_shared_load_expert_ids), ) +
-                module.w3_w1_weight.data.shape[1:],
-                dtype=module.w3_w1_weight.data.dtype,
-                device='cpu')
-            local_shared_w2_tensors = torch.empty(
-                (len(local_shared_load_expert_ids), ) +
-                module.w2_weight.data.shape[1:],
-                dtype=module.w2_weight.data.dtype,
-                device='cpu')
-            if module.bias:
-                local_shared_w3_w1_bias_tensors = torch.empty(
+            if getattr(module, 'local_shared_w3_w1_tensors', None) is not None:
+                local_shared_w3_w1_tensors = getattr(
+                    module, 'local_shared_w3_w1_tensors')
+            else:
+                local_shared_w3_w1_tensors = torch.empty(
                     (len(local_shared_load_expert_ids), ) +
-                    module.w3_w1_bias.data.shape[1:],
-                    dtype=module.w3_w1_bias.data.dtype,
+                    module.w3_w1_weight.data.shape[1:],
+                    dtype=module.w3_w1_weight.data.dtype,
                     device='cpu')
-                local_shared_w2_bias_tensors = torch.empty(
+                setattr(module, 'local_shared_w3_w1_tensors',
+                        local_shared_w3_w1_tensors)
+            if getattr(module, 'local_shared_w2_tensors', None) is not None:
+                local_shared_w2_tensors = getattr(module,
+                                                  'local_shared_w2_tensors')
+            else:
+                local_shared_w2_tensors = torch.empty(
                     (len(local_shared_load_expert_ids), ) +
-                    module.w2_bias.data.shape[1:],
-                    dtype=module.w2_bias.data.dtype,
+                    module.w2_weight.data.shape[1:],
+                    dtype=module.w2_weight.data.dtype,
                     device='cpu')
+                setattr(module, 'local_shared_w2_tensors',
+                        local_shared_w2_tensors)
+            if module.bias:
+                if getattr(module, 'local_shared_w3_w1_bias_tensors',
+                           None) is not None:
+                    local_shared_w3_w1_bias_tensors = getattr(
+                        module, 'local_shared_w3_w1_bias_tensors')
+                else:
+                    local_shared_w3_w1_bias_tensors = torch.empty(
+                        (len(local_shared_load_expert_ids), ) +
+                        module.w3_w1_bias.data.shape[1:],
+                        dtype=module.w3_w1_bias.data.dtype,
+                        device='cpu')
+                    setattr(module, 'local_shared_w3_w1_bias_tensors',
+                            local_shared_w3_w1_bias_tensors)
+                if getattr(module, 'local_shared_w2_bias_tensors',
+                           None) is not None:
+                    local_shared_w2_bias_tensors = getattr(
+                        module, 'local_shared_w2_bias_tensors')
+                else:
+                    local_shared_w2_bias_tensors = torch.empty(
+                        (len(local_shared_load_expert_ids), ) +
+                        module.w2_bias.data.shape[1:],
+                        dtype=module.w2_bias.data.dtype,
+                        device='cpu')
+                    setattr(module, 'local_shared_w2_bias_tensors',
+                            local_shared_w2_bias_tensors)
             self.load_expert_weights_to_dst(
-                module, weights, weight_loading_mode,
-                local_shared_load_expert_ids, local_shared_w3_w1_tensors,
+                module,
+                weights,
+                weight_loading_mode,
+                local_shared_load_expert_ids,
+                local_shared_w3_w1_tensors,
                 local_shared_w2_tensors,
                 local_shared_w3_w1_bias_tensors if module.bias else None,
-                local_shared_w2_bias_tensors if module.bias else None)
+                local_shared_w2_bias_tensors if module.bias else None,
+                allow_partial_loading=allow_partial_loading)
+
+    def post_load_weights(self, module: torch.nn.Module):
+        if self.need_load_shared_weights(module):
             weight_fns = {
-                'w3_w1_weight': local_shared_w3_w1_tensors,
-                'w2_weight': local_shared_w2_tensors
+                'w3_w1_weight': getattr(module, 'local_shared_w3_w1_tensors'),
+                'w2_weight': getattr(module, 'local_shared_w2_tensors')
             }
+            delattr(module, 'local_shared_w3_w1_tensors')
+            delattr(module, 'local_shared_w2_tensors')
             if module.bias:
                 weight_fns.update({
-                    'w3_w1_bias': local_shared_w3_w1_bias_tensors,
-                    'w2_bias': local_shared_w2_bias_tensors
+                    'w3_w1_bias':
+                    getattr(module, 'local_shared_w3_w1_bias_tensors'),
+                    'w2_bias':
+                    getattr(module, 'local_shared_w2_bias_tensors')
                 })
+                delattr(module, 'local_shared_w3_w1_bias_tensors')
+                delattr(module, 'local_shared_w2_bias_tensors')
             module.register_all_parameter_slot_and_to_fix_weight_fns(weight_fns)
             module.layer_load_balancer.host_tensor_sharer.finalize_layer_weights(
             )
-
         if hasattr(module,
                    "layer_load_balancer") and module.layer_load_balancer:
             module.layer_load_balancer.set_initial_weight_assignments(
                 module.initial_global_assignments)
-
-    def post_load_weights(self, module: torch.nn.Module):
         # Re-setup quant scales after loading weights as the tensors may have been modified.
         self.setup_quant_scales(module)
 
@@ -374,48 +450,64 @@ def apply(self, module: torch.nn.Module, input: torch.Tensor, *args,
         raise NotImplementedError
 
     # Helper function
-    def load_expert_w3_w1_weight(self, module: torch.nn.Module,
+    def load_expert_w3_w1_weight(self,
+                                 module: torch.nn.Module,
                                  w1_weight: torch.Tensor,
                                  w3_weight: torch.Tensor,
-                                 dst_w3_w1_weight: torch.Tensor):
+                                 dst_w3_w1_weight: torch.Tensor,
+                                 allow_partial_loading: bool = False):
         """
         Load w1 and w3 weights for each expert.
         Override this method if you need to preprocess the weights differently.
         """
         # device don't have to be 'cuda', e.g. 'cpu' for online EPLB
         device = dst_w3_w1_weight.device
-        w1_weight_shard = load_weight_shard(w1_weight,
-                                            module.tp_size,
-                                            module.tp_rank,
-                                            TensorParallelMode.COLUMN,
-                                            device=device)
-        w3_weight_shard = load_weight_shard(w3_weight,
-                                            module.tp_size,
-                                            module.tp_rank,
-                                            TensorParallelMode.COLUMN,
-                                            device=device)
+        if not allow_partial_loading:
+            assert w1_weight is not None and w3_weight is not None
+        w1_weight_shard = load_weight_shard(
+            w1_weight,
+            module.tp_size,
+            module.tp_rank,
+            TensorParallelMode.COLUMN,
+            device=device) if w1_weight is not None else None
+        w3_weight_shard = load_weight_shard(
+            w3_weight,
+            module.tp_size,
+            module.tp_rank,
+            TensorParallelMode.COLUMN,
+            device=device) if w3_weight is not None else None
 
-        w31_weight_shard = torch.cat([w3_weight_shard, w1_weight_shard], dim=0)
-        dst_w3_w1_weight.copy_(w31_weight_shard.view(dst_w3_w1_weight.dtype),
-                               non_blocking=True)
+        dst_w3_weight, dst_w1_weight = dst_w3_w1_weight.chunk(2, dim=0)
+        if w1_weight is not None:
+            dst_w1_weight.copy_(w1_weight_shard.view(dst_w3_w1_weight.dtype),
+                                non_blocking=True)
+        if w3_weight is not None:
+            dst_w3_weight.copy_(w3_weight_shard.view(dst_w3_w1_weight.dtype),
+                                non_blocking=True)
 
     # Helper function
-    def load_expert_w2_weight(self, module: torch.nn.Module,
+    def load_expert_w2_weight(self,
+                              module: torch.nn.Module,
                               w2_weight: torch.Tensor,
-                              dst_w2_weight: torch.Tensor):
+                              dst_w2_weight: torch.Tensor,
+                              allow_partial_loading: bool = False):
         """
         Load w2 weight for each expert.
         Override this method if you need to preprocess the weights differently.
         """
         # device don't have to be 'cuda', e.g. 'cpu' for online EPLB
         device = dst_w2_weight.device
-        w2_weight_shard = load_weight_shard(w2_weight,
-                                            module.tp_size,
-                                            module.tp_rank,
-                                            TensorParallelMode.ROW,
-                                            device=device)
-        dst_w2_weight.copy_(w2_weight_shard.view(dst_w2_weight.dtype),
-                            non_blocking=True)
+        if not allow_partial_loading:
+            assert w2_weight is not None
+        w2_weight_shard = load_weight_shard(
+            w2_weight,
+            module.tp_size,
+            module.tp_rank,
+            TensorParallelMode.ROW,
+            device=device) if w2_weight is not None else None
+        if w2_weight is not None:
+            dst_w2_weight.copy_(w2_weight_shard.view(dst_w2_weight.dtype),
+                                non_blocking=True)
 
 
 class UnquantizedFusedMoEMethod(FusedMoEMethodBase):
diff --git a/tensorrt_llm/_torch/modules/gated_mlp.py b/tensorrt_llm/_torch/modules/gated_mlp.py
index 90af4440c36e..c1200c7d75ee 100644
--- a/tensorrt_llm/_torch/modules/gated_mlp.py
+++ b/tensorrt_llm/_torch/modules/gated_mlp.py
@@ -57,6 +57,15 @@ def __init__(
         else:
             mapping = config.mapping
 
+        # Calculate local intermediate size after tensor parallel sharding
+        tp_size = mapping.tp_size
+        local_intermediate_size = self.intermediate_size // tp_size
+
+        gateup_shard_indices_mapping = {
+            'gate': (0, local_intermediate_size),
+            'up': (local_intermediate_size, local_intermediate_size),
+        }
+
         self.gate_up_proj = Linear(
             self.hidden_size,
             self.intermediate_size * 2,
@@ -73,6 +82,7 @@ def __init__(
             force_dynamic_quantization=config.force_dynamic_quantization,
             use_cute_dsl_blockscaling_mm=use_cute_dsl_blockscaling_mm,
             disable_deep_gemm=disable_deep_gemm,
+            fused_weight_shard_indices_mapping=gateup_shard_indices_mapping,
         )
 
         self.down_lora = LoraLayer([LoraModuleType.MLP_4H_TO_H],
diff --git a/tensorrt_llm/_torch/modules/linear.py b/tensorrt_llm/_torch/modules/linear.py
index bba66f0385be..e24cf5f583cb 100644
--- a/tensorrt_llm/_torch/modules/linear.py
+++ b/tensorrt_llm/_torch/modules/linear.py
@@ -130,58 +130,99 @@ def copy_weight(dst: Parameter, src: torch.Tensor):
     dst.data.copy_(src)
 
 
+def copy_weight_shard(dst: Parameter, src: torch.Tensor, shard_offset: int,
+                      shard_size: int):
+    if dst.dtype != src.dtype:
+        src = src.to(dst.dtype)
+    assert dst.dtype == src.dtype, f"Incompatible dtype. dst: {dst.dtype}, src: {src.dtype}"
+    dst[shard_offset:shard_offset + shard_size].data.copy_(src)
+
+
 def load_weights_vanilla_helper(module: Linear,
                                 weights: List[Dict],
                                 weight_transform=lambda x: x,
-                                bias_transform=lambda x: x):
+                                bias_transform=lambda x: x,
+                                allow_partial_loading: bool = False):
     assert len(weights) == 1
+    if not allow_partial_loading:
+        assert "weight" in weights[0]
+        if module.bias is not None:
+            assert "bias" in weights[0]
     device = torch.device('cuda')
 
     weight = load_weight_shard(weights[0]['weight'], module.tp_size,
-                               module.tp_rank, module.tp_mode, device)
+                               module.tp_rank, module.tp_mode,
+                               device) if "weight" in weights[0] else None
 
-    if module.has_weight_only_quant:
-        # NOTE: without the preprocess during the runtime, the gemm output nan's. in order to use the preprocess_weights_for_mixed_gemm
-        # we need to cast the weight to int8 first.
-        activation_dtype = torch.float8_e4m3fn if module.has_w4a8_awq else torch.float16
-        weight_dtype, _ = get_weight_dtype_and_id(module)
-        weight = preprocess_weights_for_mixed_gemm(
-            weight.T.to(torch.int8).contiguous().cpu(), weight_dtype,
-            activation_dtype).cuda().contiguous()
+    if weight is not None:
+        if module.has_weight_only_quant:
+            # NOTE: without the preprocess during the runtime, the gemm output nan's. in order to use the preprocess_weights_for_mixed_gemm
+            # we need to cast the weight to int8 first.
+            activation_dtype = torch.float8_e4m3fn if module.has_w4a8_awq else torch.float16
+            weight_dtype, _ = get_weight_dtype_and_id(module)
+            weight = preprocess_weights_for_mixed_gemm(
+                weight.T.to(torch.int8).contiguous().cpu(), weight_dtype,
+                activation_dtype).cuda().contiguous()
 
-    copy_weight(module.weight, weight_transform(weight))
+        copy_weight(module.weight, weight_transform(weight))
 
     if module.bias is not None:
         bias = load_weight_shard(weights[0]['bias'], module.tp_size,
-                                 module.tp_rank, module.tp_mode, device)
-        copy_weight(module.bias, bias_transform(bias))
+                                 module.tp_rank, module.tp_mode,
+                                 device) if "bias" in weights[0] else None
+        if bias is not None:
+            copy_weight(module.bias, bias_transform(bias))
 
 
 def load_weights_fused_qkv_helper(
     module: Linear,
     weights: List[Dict],
     weight_transform=lambda x: x,
-    bias_transform=lambda x: x
+    bias_transform=lambda x: x,
+    allow_partial_loading: bool = False
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    assert len(weights) == 3
+    if not allow_partial_loading:
+        assert all('weight' in weights[i] for i in range(3))
+        if module.bias is not None:
+            assert all('bias' in weights[i] for i in range(3))
+    else:
+        assert getattr(
+            module, "fused_weight_shard_indices_mapping", None
+        ) is not None, "Fused weight shard indices mapping is required in partial loading"
     device = torch.device('cuda')
 
     q_weight = load_weight_shard(weights[0]['weight'], module.tp_size,
-                                 module.tp_rank, module.tp_mode, device)
+                                 module.tp_rank, module.tp_mode,
+                                 device) if "weight" in weights[0] else None
     k_weight = load_weight_shard(weights[1]['weight'], module.tp_size,
-                                 module.tp_rank, module.tp_mode, device)
+                                 module.tp_rank, module.tp_mode,
+                                 device) if "weight" in weights[1] else None
     v_weight = load_weight_shard(weights[2]['weight'], module.tp_size,
-                                 module.tp_rank, module.tp_mode, device)
+                                 module.tp_rank, module.tp_mode,
+                                 device) if "weight" in weights[2] else None
 
     if module.bias is not None:
         q_bias = load_weight_shard(weights[0]['bias'], module.tp_size,
-                                   module.tp_rank, module.tp_mode, device)
+                                   module.tp_rank, module.tp_mode,
+                                   device) if "bias" in weights[0] else None
         k_bias = load_weight_shard(weights[1]['bias'], module.tp_size,
-                                   module.tp_rank, module.tp_mode, device)
+                                   module.tp_rank, module.tp_mode,
+                                   device) if "bias" in weights[1] else None
         v_bias = load_weight_shard(weights[2]['bias'], module.tp_size,
-                                   module.tp_rank, module.tp_mode, device)
-        copy_weight(module.bias,
-                    bias_transform(torch.cat((q_bias, k_bias, v_bias))))
+                                   module.tp_rank, module.tp_mode,
+                                   device) if "bias" in weights[2] else None
+        if not allow_partial_loading:
+            copy_weight(module.bias,
+                        bias_transform(torch.cat((q_bias, k_bias, v_bias))))
+        else:
+            for shard_key, bias in zip(('q', 'k', 'v'),
+                                       (q_bias, k_bias, v_bias)):
+                if bias is not None:
+                    assert shard_key in module.fused_weight_shard_indices_mapping, f"Shard key {shard_key} not found in fused weight shard indices mapping"
+                    shard_offset, shard_size = module.fused_weight_shard_indices_mapping[
+                        shard_key]
+                    copy_weight_shard(module.bias, bias_transform(bias),
+                                      shard_offset, shard_size)
 
     return tuple(map(weight_transform, (q_weight, k_weight, v_weight)))
 
@@ -190,21 +231,44 @@ def load_weights_fused_gate_up_helper(
         module: Linear,
         weights: List[Dict],
         weight_transform=lambda x: x,
-        bias_transform=lambda x: x) -> tuple[torch.Tensor, torch.Tensor]:
-    assert len(weights) == 2
+        bias_transform=lambda x: x,
+        allow_partial_loading: bool = False
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not allow_partial_loading:
+        assert all('weight' in weights[i] for i in range(2))
+        if module.bias is not None:
+            assert all('bias' in weights[i] for i in range(2))
+    else:
+        assert getattr(
+            module, "fused_weight_shard_indices_mapping", None
+        ) is not None, "Fused weight shard indices mapping is required in partial loading"
     device = torch.device('cuda')
 
     gate_weight = load_weight_shard(weights[0]['weight'], module.tp_size,
-                                    module.tp_rank, module.tp_mode, device)
+                                    module.tp_rank, module.tp_mode,
+                                    device) if "weight" in weights[0] else None
     up_weight = load_weight_shard(weights[1]['weight'], module.tp_size,
-                                  module.tp_rank, module.tp_mode, device)
+                                  module.tp_rank, module.tp_mode,
+                                  device) if "weight" in weights[1] else None
     if module.bias is not None:
         gate_bias = load_weight_shard(weights[0]['bias'], module.tp_size,
-                                      module.tp_rank, module.tp_mode, device)
+                                      module.tp_rank, module.tp_mode,
+                                      device) if "bias" in weights[0] else None
         up_bias = load_weight_shard(weights[1]['bias'], module.tp_size,
-                                    module.tp_rank, module.tp_mode, device)
-        copy_weight(module.bias, bias_transform(torch.cat(
-            (gate_bias, up_bias))))
+                                    module.tp_rank, module.tp_mode,
+                                    device) if "bias" in weights[1] else None
+        if not allow_partial_loading:
+            copy_weight(module.bias,
+                        bias_transform(torch.cat((gate_bias, up_bias))))
+        else:
+            for shard_key, bias in zip(('gate', 'up'), (gate_bias, up_bias)):
+                if bias is not None:
+                    assert shard_key in module.fused_weight_shard_indices_mapping, f"Shard key {shard_key} not found in fused weight shard indices mapping"
+                    shard_offset, shard_size = module.fused_weight_shard_indices_mapping[
+                        shard_key]
+                    copy_weight_shard(module.bias, bias_transform(bias),
+                                      shard_offset, shard_size)
+
     return tuple(map(weight_transform, (gate_weight, up_weight)))
 
 
@@ -245,17 +309,23 @@ def apply(self, module: Linear, input: torch.Tensor,
               bias: Optional[torch.Tensor], *args, **kwargs):
         raise NotImplementedError
 
-    def load_weights(self, module: Linear, weights: List[Dict],
-                     weight_mode: WeightMode):
+    def load_weights(self,
+                     module: Linear,
+                     weights: List[Dict],
+                     weight_mode: WeightMode,
+                     allow_partial_loading: bool = False):
         """
         Load weights from the checkpoint.
         """
+        kargs = {}
+        if isinstance(self, UnquantizedLinearMethod):
+            kargs['allow_partial_loading'] = allow_partial_loading
         if weight_mode == WeightMode.VANILLA:
-            self.load_weights_vanilla(module, weights)
+            self.load_weights_vanilla(module, weights, **kargs)
         elif weight_mode == WeightMode.FUSED_QKV_LINEAR:
-            self.load_weights_fused_qkv_linear(module, weights)
+            self.load_weights_fused_qkv_linear(module, weights, **kargs)
         elif weight_mode == WeightMode.FUSED_GATE_UP_LINEAR:
-            self.load_weights_fused_gate_up_linear(module, weights)
+            self.load_weights_fused_gate_up_linear(module, weights, **kargs)
         else:
             raise ValueError(f'unsupported weight mode: {weight_mode}')
 
@@ -268,23 +338,32 @@ def load_weight_scales(self, weights: List[Dict], *args, **kwargs):
         """
 
     @abstractmethod
-    def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None:
+    def load_weights_vanilla(self,
+                             module: Linear,
+                             weights: List[Dict],
+                             allow_partial_loading: bool = False) -> None:
         """
         Load weights for the VANILLA weight mode.
         """
         raise NotImplementedError
 
     @abstractmethod
-    def load_weights_fused_qkv_linear(self, module: Linear,
-                                      weights: List[Dict]) -> None:
+    def load_weights_fused_qkv_linear(
+            self,
+            module: Linear,
+            weights: List[Dict],
+            allow_partial_loading: bool = False) -> None:
         """
         Load weights for the FUSED_QKV_LINEAR weight mode.
         """
         raise NotImplementedError
 
     @abstractmethod
-    def load_weights_fused_gate_up_linear(self, module: Linear,
-                                          weights: List[Dict]) -> None:
+    def load_weights_fused_gate_up_linear(
+            self,
+            module: Linear,
+            weights: List[Dict],
+            allow_partial_loading: bool = False) -> None:
         """
         Load weights for the FUSED_GATE_UP_LINEAR weight mode.
         """
@@ -316,22 +395,52 @@ def apply(self, module: Linear, input: torch.Tensor,
             output = F.linear(input, module.weight, bias)
         return output
 
-    def load_weights_vanilla(self, module: Linear, weights: List[Dict]) -> None:
-        load_weights_vanilla_helper(module, weights)
+    def load_weights_vanilla(self,
+                             module: Linear,
+                             weights: List[Dict],
+                             allow_partial_loading: bool = False) -> None:
+        load_weights_vanilla_helper(module,
+                                    weights,
+                                    allow_partial_loading=allow_partial_loading)
 
-    def load_weights_fused_qkv_linear(self, module: Linear,
-                                      weights: List[Dict]) -> None:
+    def load_weights_fused_qkv_linear(
+            self,
+            module: Linear,
+            weights: List[Dict],
+            allow_partial_loading: bool = False) -> None:
         q_weight, k_weight, v_weight = load_weights_fused_qkv_helper(
-            module, weights)
-        fused_weight = torch.cat((q_weight, k_weight, v_weight))
-        copy_weight(module.weight, fused_weight)
-
-    def load_weights_fused_gate_up_linear(self, module: Linear,
-                                          weights: List[Dict]) -> None:
+            module, weights, allow_partial_loading=allow_partial_loading)
+        if not allow_partial_loading:
+            copy_weight(module.weight, torch.cat(
+                (q_weight, k_weight, v_weight)))
+        else:
+            for shard_key, weight in zip(('q', 'k', 'v'),
+                                         (q_weight, k_weight, v_weight)):
+                if weight is not None:
+                    assert shard_key in module.fused_weight_shard_indices_mapping, f"Shard key {shard_key} not found in fused weight shard indices mapping"
+                    shard_offset, shard_size = module.fused_weight_shard_indices_mapping[
+                        shard_key]
+                    copy_weight_shard(module.weight, weight, shard_offset,
+                                      shard_size)
+
+    def load_weights_fused_gate_up_linear(
+            self,
+            module: Linear,
+            weights: List[Dict],
+            allow_partial_loading: bool = False) -> None:
         gate_weight, up_weight = load_weights_fused_gate_up_helper(
-            module, weights)
-        fused_weight = torch.cat((gate_weight, up_weight))
-        copy_weight(module.weight, fused_weight)
+            module, weights, allow_partial_loading=allow_partial_loading)
+        if not allow_partial_loading:
+            copy_weight(module.weight, torch.cat((gate_weight, up_weight)))
+        else:
+            for shard_key, weight in zip(('gate', 'up'),
+                                         (gate_weight, up_weight)):
+                if weight is not None:
+                    assert shard_key in module.fused_weight_shard_indices_mapping, f"Shard key {shard_key} not found in fused weight shard indices mapping"
+                    shard_offset, shard_size = module.fused_weight_shard_indices_mapping[
+                        shard_key]
+                    copy_weight_shard(module.weight, weight, shard_offset,
+                                      shard_size)
 
 
 class FP8QDQLinearMethod(LinearMethodBase):
@@ -1906,6 +2015,7 @@ def __init__(
         use_cute_dsl_nvfp4_blockscaling_mm: bool = False,
         use_cublaslt_nvfp4_blockscaling_mm: bool = False,
         disable_deep_gemm: bool = False,
+        fused_weight_shard_indices_mapping: Optional[dict] = None,
     ):
         from ..distributed import AllReduce
 
@@ -1926,6 +2036,7 @@ def __init__(
         self.use_cute_dsl_nvfp4_blockscaling_mm = use_cute_dsl_nvfp4_blockscaling_mm
         self.use_cublaslt_nvfp4_blockscaling_mm = use_cublaslt_nvfp4_blockscaling_mm
         self.disable_deep_gemm = disable_deep_gemm
+        self.fused_weight_shard_indices_mapping = fused_weight_shard_indices_mapping
 
         local_in_features = in_features
         local_out_features = out_features
@@ -2101,11 +2212,19 @@ def forward(
 
         return output
 
-    def load_weights(self, weights: List[Dict]):
+    def load_weights(self,
+                     weights: List[Dict],
+                     allow_partial_loading: bool = False):
         assert self._weights_created
 
         weight_mode = self.weights_loading_config.weight_mode
-        self.quant_method.load_weights(self, weights, weight_mode)
+        if not isinstance(self.quant_method, UnquantizedLinearMethod):
+            assert allow_partial_loading is False, "allow_partial_loading is only supported for non-unquantized linear methods now"
+        self.quant_method.load_weights(
+            self,
+            weights,
+            weight_mode,
+            allow_partial_loading=allow_partial_loading)
 
     def post_load_weights(self):
         self.quant_method.post_load_weights(self)
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 9aefc0761b48..36c8b6ed2588 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -187,7 +187,8 @@ def __init__(
         if model is None:
             lora_config: Optional[
                 LoraConfig] = None if is_draft_model else llm_args.lora_config
-            loader = ModelLoader(
+            # Keep the model_loader to support reloading the model weights later
+            self.model_loader = ModelLoader(
                 llm_args=llm_args,
                 mapping=self.mapping,
                 spec_config=self.spec_config,
@@ -196,7 +197,7 @@ def __init__(
                 max_seq_len=self.max_seq_len,
                 lora_config=lora_config,
             )
-            self.model, moe_load_balancer = loader.load(
+            self.model, moe_load_balancer = self.model_loader.load(
                 checkpoint_dir=model_path, checkpoint_loader=checkpoint_loader)
             if isinstance(moe_load_balancer, MoeLoadBalancer):
                 setattr(self, "moe_load_balancer", moe_load_balancer)
diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py
index b9c1377cd98c..c99e01f8082d 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_loader.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py
@@ -20,7 +20,8 @@
 from ..model_config import ModelConfig
 from ..models import AutoModelForCausalLM
 from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader
-from ..models.modeling_utils import MetaInitMode, timing
+from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode,
+                                     timing)
 from ..modules.fused_moe.moe_load_balancer import (
     MoeLoadBalancer, maybe_create_moe_load_balancer)
 
@@ -268,19 +269,21 @@ def init_meta_tensor(t: torch.Tensor):
                 else:
                     weights = checkpoint_loader.load_weights(checkpoint_dir)
 
-                weight_mapper = checkpoint_loader.get_initialized_weight_mapper(
+                self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper(
                     model, config)
                 self._call_load_weights(model.load_weights, weights,
-                                        weight_mapper)
+                                        self.weight_mapper)
 
                 if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
                 ):
                     weights = checkpoint_loader.load_weights(
                         self.spec_config.speculative_model_dir)
                     self._call_load_weights(model.load_draft_weights, weights,
-                                            weight_mapper)
+                                            self.weight_mapper)
 
             elif load_format == LoadFormat.DUMMY:
+                self.weight_mapper = checkpoint_loader.get_initialized_weight_mapper(
+                    model, config)
                 initialize_dummy_weights(model)
                 if self.spec_config is not None and self.spec_config.spec_dec_mode.need_load_draft_weights(
                 ):
@@ -311,6 +314,16 @@ def init_meta_tensor(t: torch.Tensor):
 
         return model, moe_load_balancer
 
+    def reload(self,
+               model: DecoderModelForCausalLM,
+               weights: dict,
+               allow_partial_loading: bool = False):
+        self._call_load_weights(model.load_weights,
+                                weights,
+                                self.weight_mapper,
+                                allow_partial_loading=allow_partial_loading)
+        torch.cuda.current_stream().synchronize()
+
     def _load_and_validate_config(
             self, checkpoint_dir: str,
             checkpoint_loader: BaseCheckpointLoader) -> ModelConfig:
@@ -354,10 +367,18 @@ def _load_and_validate_config(
                             sub_config).num_hidden_layers = num_layers_override
         return config
 
-    def _call_load_weights(self, load_method: Callable, weights, weight_mapper):
+    def _call_load_weights(self,
+                           load_method: Callable,
+                           weights,
+                           weight_mapper,
+                           allow_partial_loading: bool = False):
         """Calls the model's weight loading method with the correct arguments."""
         args = inspect.getfullargspec(load_method).args
+        kargs = {}
         if "weight_mapper" in args:
-            load_method(weights, weight_mapper=weight_mapper)
+            kargs["weight_mapper"] = weight_mapper
+        if "allow_partial_loading" in args:
+            kargs["allow_partial_loading"] = allow_partial_loading
         else:
-            load_method(weights)
+            assert allow_partial_loading is False, "allow_partial_loading is not supported for this model"
+        load_method(weights, **kargs)
diff --git a/tensorrt_llm/llmapi/rlhf_utils.py b/tensorrt_llm/llmapi/rlhf_utils.py
index b3d63ec236ef..4934d40e9791 100644
--- a/tensorrt_llm/llmapi/rlhf_utils.py
+++ b/tensorrt_llm/llmapi/rlhf_utils.py
@@ -1,6 +1,9 @@
+from typing import Optional
+
 import torch
 
 from tensorrt_llm._ray_utils import control_action_decorator
+from tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import MoeLoadBalancer
 from tensorrt_llm._torch.utils import get_device_uuid
 from tensorrt_llm.logger import logger
 
@@ -28,7 +31,7 @@ class WorkerExtension:
     """
 
     @control_action_decorator
-    def update_weights(self, ipc_handles: dict):
+    def update_weights(self, ipc_handles: Optional[dict] = None):
         """Update model weights from IPC (Inter-Process Communication) handles.
 
         This method receives shared memory handles from another process (typically FSDP training),
@@ -45,25 +48,41 @@ def update_weights(self, ipc_handles: dict):
             Exception: Re-raises any exception encountered during weight update.
         """
         try:
-            logger.info("Update weights from IPC handles")
-            device_uuid = get_device_uuid(self.device_id)
-
-            if device_uuid not in ipc_handles:
-                raise ValueError(f"Device UUID {device_uuid} not found in ipc_handles")
-
-            weights = {}
-            all_handles = ipc_handles[device_uuid]
-
-            for param_name, tensor_handle in all_handles:
-                func, args = tensor_handle
-                list_args = list(args)
-                list_args[6] = self.device_id  # Set target device
-                tensor = func(*list_args)
-                weights[param_name] = tensor
-
-            self.engine.model_engine.model.load_weights(weights)
-            torch.cuda.synchronize()
-            self.engine.reset_prefix_cache()
+            if ipc_handles is not None:
+                logger.info("Update weights from IPC handles")
+                device_uuid = get_device_uuid(self.device_id)
+
+                if device_uuid not in ipc_handles:
+                    raise ValueError(f"Device UUID {device_uuid} not found in ipc_handles")
+
+                weights = {}
+                all_handles = ipc_handles[device_uuid]
+
+                for param_name, tensor_handle in all_handles:
+                    func, args = tensor_handle
+                    list_args = list(args)
+                    list_args[6] = self.device_id  # Set target device
+                    tensor = func(*list_args)
+                    weights[param_name] = tensor
+
+                logger.info(f"weights key size: {len(weights.keys())}")
+                self.engine.model_engine.model_loader.reload(
+                    self.engine.model_engine.model, weights, allow_partial_loading=True
+                )
+            else:
+                logger.info("Finalize update weights")
+                for module in self.engine.model_engine.model.modules():
+                    if hasattr(module, "post_load_weights") and not getattr(
+                        module, "_weights_removed", False
+                    ):
+                        module.post_load_weights()
+                moe_load_balancer = getattr(self.engine.model_engine, "moe_load_balancer", None)
+                if isinstance(moe_load_balancer, MoeLoadBalancer):
+                    moe_load_balancer.register_weight_slots_after_to_cuda()
+                    logger.info("moe_load_balancer finalizing model...")
+                    moe_load_balancer.finalize_model()
+                    logger.info("moe_load_balancer finalize model done")
+                self.engine.reset_prefix_cache()
 
         except Exception as e:
             logger.error("Encountered an error in update_weights")
diff --git a/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py b/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
index 5ce554d7dc29..96e882261222 100644
--- a/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
+++ b/tests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_update_weights.py
@@ -1,5 +1,6 @@
-from typing import List
+from typing import Callable, List, Optional
 
+import pytest
 import torch
 from torch.multiprocessing.reductions import reduce_tensor
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -18,6 +19,7 @@ def __init__(self, model_name: str):
         self.cuda_device = torch.cuda.current_device()
         self.all_weights = {}
         self.device_uuid = [HFModel.get_device_uuid(i) for i in range(torch.cuda.device_count())]
+        self._replicate_weights()
 
     @staticmethod
     def get_device_uuid(cuda_device: int):
@@ -25,12 +27,6 @@ def get_device_uuid(cuda_device: int):
 
         return get_device_uuid(cuda_device)
 
-    def flip_weights(self):
-        for _, p in self.model.named_parameters():
-            p.data = -p.data
-
-        self._replicate_weights()
-
     def _replicate_weights(self):
         model_weights = []
         for n, p in self.model.named_parameters():
@@ -44,18 +40,35 @@ def _replicate_weights(self):
                     cur_weights.append((n, p.to("cuda:" + str(i))))
                 self.all_weights[i] = cur_weights
 
-    def get_weight_ipc_handles(self, cuda_device: int = None):
+    def get_weight_ipc_handles(
+        self,
+        cuda_device: Optional[List[int]] = None,
+        weight_filter: Optional[Callable[[str], bool]] = None,
+    ):
+        """
+        Get IPC handles for model weights with flexible filtering.
+
+        Args:
+            cuda_device: List of CUDA device indices to get weights from
+            weight_filter: Optional function that takes weight name and returns True if weight should be included
+
+        Returns:
+            ret: Dictionary containing weight handles
+        """
         ret = {}
-        device_list = (
-            list(range(torch.cuda.device_count())) if cuda_device is None else [cuda_device]
-        )
+        device_list = list(range(torch.cuda.device_count())) if cuda_device is None else cuda_device
+
         for device in device_list:
             all_handles = []
             for item in self.all_weights[device]:
                 name, p = item
+                # Apply filter if provided
+                if weight_filter is not None and not weight_filter(name):
+                    continue
                 handle = reduce_tensor(p)
                 all_handles.append((name, handle))
             ret[self.device_uuid[device]] = all_handles
+
         return ret
 
     def generate_batch_incremental(
@@ -81,6 +94,7 @@ def generate_batch_incremental(
                 ret = self.model.generate(
                     input_ids=cur_token_ids.unsqueeze(0).cuda(),
                     max_new_tokens=1,
+                    do_sample=False,
                     return_dict_in_generate=True,
                     output_scores=True,
                 )
@@ -106,7 +120,7 @@ def compare_logits(
     logits_list: List[torch.Tensor],
     ref_logits_list: List[torch.Tensor],
     topk: int = 20,
-    threshold: float = 0.85,
+    threshold: float = 0.9,
 ):
     assert len(logits_list) == len(ref_logits_list)
 
@@ -139,16 +153,21 @@ def run_generate(llm, hf_model, prompts, sampling_params):
     return llm_logits, ref_logits
 
 
-def test_llm_update_weights():
-    llama_model_path = str(llm_models_root() / "llama-models-v2/TinyLlama-1.1B-Chat-v1.0")
+@pytest.mark.parametrize(
+    "model_dir",
+    ["Qwen2.5-0.5B-Instruct", "Qwen3/Qwen3-8B", "llama-models-v2/TinyLlama-1.1B-Chat-v1.0"],
+)
+def test_llm_update_weights(model_dir):
+    model_dir = str(llm_models_root() / model_dir)
     kv_cache_config = KvCacheConfig(enable_block_reuse=True, free_gpu_memory_fraction=0.1)
 
-    hf_model = HFModel(llama_model_path)
+    hf_model = HFModel(model_dir)
 
     llm = LLM(
-        model=llama_model_path,
+        model=model_dir,
         ray_worker_extension_cls="tensorrt_llm.llmapi.rlhf_utils.WorkerExtension",
         tensor_parallel_size=1,
+        load_format="dummy",
         pipeline_parallel_size=1,
         kv_cache_config=kv_cache_config,
     )
@@ -163,20 +182,79 @@ def test_llm_update_weights():
 
     sampling_params = SamplingParams(temperature=0, return_generation_logits=True)
 
-    results = []
+    ipc_handles = hf_model.get_weight_ipc_handles([0])
+
+    llm._collective_rpc("update_weights", (ipc_handles,))
+    # Finalize the update weights
+    llm._collective_rpc("update_weights", (None,))
 
-    # Stage 1: Generate with original model
-    results.append(run_generate(llm, hf_model, prompts, sampling_params))
-    llm_logits, ref_logits = results[0]
+    llm_logits, ref_logits = run_generate(llm, hf_model, prompts, sampling_params)
     compare_logits(llm_logits, ref_logits)
 
-    # Stage 2: Test update with flipped weights
-    hf_model.flip_weights()
-    ipc_handles = hf_model.get_weight_ipc_handles()
 
-    llm._collective_rpc("update_weights", (ipc_handles,))
-    results.append(run_generate(llm, hf_model, prompts, sampling_params))
-    llm_logits, ref_logits = results[1]
+@pytest.mark.parametrize(
+    "model_dir",
+    ["Qwen2.5-0.5B-Instruct", "Qwen3/Qwen3-8B", "llama-models-v2/TinyLlama-1.1B-Chat-v1.0"],
+)
+def test_llm_partial_update_weights(model_dir):
+    model_dir = str(llm_models_root() / model_dir)
+    kv_cache_config = KvCacheConfig(enable_block_reuse=True, free_gpu_memory_fraction=0.1)
+
+    hf_model = HFModel(model_dir)
+
+    llm = LLM(
+        model=model_dir,
+        ray_worker_extension_cls="tensorrt_llm.llmapi.rlhf_utils.WorkerExtension",
+        tensor_parallel_size=1,
+        load_format="dummy",
+        pipeline_parallel_size=1,
+        kv_cache_config=kv_cache_config,
+    )
+
+    # Generate texts from the prompts.
+    prompts = [
+        "Hello, my name is",
+        "The president of the United States is",
+        "The capital of France is",
+        "The future of AI is",
+    ]
 
-    # Compare the logits for this phase since output should be random
+    sampling_params = SamplingParams(temperature=0, return_generation_logits=True)
+
+    ipc_handles = hf_model.get_weight_ipc_handles([0])
+
+    def common_filter(filter_name: str) -> Callable[[str], bool]:
+        def filter_fn(name: str) -> bool:
+            return filter_name in name
+
+        return filter_fn
+
+    filter_list = [
+        "q_proj.weight",
+        "k_proj.weight",
+        "v_proj.weight",
+        "o_proj.weight",
+        "gate_proj.weight",
+        "up_proj.weight",
+        "down_proj.weight",
+        "norm.weight",
+        "embed_tokens.weight",
+        "lm_head.weight",
+    ]
+    if "Qwen2.5" in model_dir or "Qwen2" in model_dir:
+        filter_list.extend(
+            [
+                "q_proj.bias",
+                "k_proj.bias",
+                "v_proj.bias",
+            ]
+        )
+    for filter_name in filter_list:
+        weight_filter = common_filter(filter_name=filter_name)
+        ipc_handles = hf_model.get_weight_ipc_handles([0], weight_filter=weight_filter)
+        llm._collective_rpc("update_weights", (ipc_handles,), non_block=True)
+    # Finalize the update weights
+    llm._collective_rpc("update_weights", (None,))
+
+    llm_logits, ref_logits = run_generate(llm, hf_model, prompts, sampling_params)
     compare_logits(llm_logits, ref_logits)