[Bugfix] Re-sync parameter tp_rank after process_weights_after_loading (fix replicated / disable_tp weight reload)

vllm/model_executor/layers/linear.py

@@ -281,6 +281,15 @@ def __init__(
         self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
 
     def update_param_tp_status(self):
+        # Single source of truth for a parameter's TP state. BasevLLMParameter
+        # stamps self.tp_rank with the *global* rank in __init__; this reconciles
+        # every child parameter to the *layer's* tp_rank/tp_size (which correctly
+        # accounts for disable_tp -> replicated weights with tp_rank == 0).
+        #
+        # Must be re-run whenever parameters are (re-)created after construction,
+        # e.g. after quant_method.process_weights_after_loading() swaps in fresh
+        # Parameters. Otherwise a later load_weights()/weight-refit would narrow a
+        # replicated weight at global_rank * shard_size and overflow.
         for param in self.parameters():
             if isinstance(param, BasevLLMParameter):
                 param.tp_rank = self.tp_rank
@@ -910,7 +919,6 @@ def weight_loader_v2(
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
-            tp_rank=self.tp_rank,
         )
 
     def load_weights(
@@ -1110,12 +1118,10 @@ def weight_loader_v2(
                 # to ensure that any subsequent reduction (like .max())
                 # works correctly while preserving the parameter shape.
                 for idx in range(param.data.shape[0]):
-                    param.load_qkv_weight(
-                        loaded_weight=loaded_weight, shard_id=idx, tp_rank=self.tp_rank
-                    )
+                    param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=idx)
                 return
             elif type(param) in (RowvLLMParameter, BasevLLMParameter):
-                param.load_qkv_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank)
+                param.load_qkv_weight(loaded_weight=loaded_weight)
                 return
             # TODO: @dsikka - move to parameter.py
             self._load_fused_module_from_checkpoint(param, loaded_weight)
@@ -1139,7 +1145,6 @@ def weight_loader_v2(
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
-            tp_rank=self.tp_rank,
         )
 
     def weight_loader(
@@ -1493,7 +1498,6 @@ def weight_loader_v2(
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
-            tp_rank=self.tp_rank,
         )
 
     def weight_loader(

vllm/model_executor/model_loader/reload/layerwise.py

@@ -364,6 +364,11 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
     quant_method = getattr(layer, "quant_method", None)
     if isinstance(quant_method, QuantizeMethodBase):
         quant_method.process_weights_after_loading(layer)
+        # Re-reconcile parameter TP state: process_weights_after_loading may
+        # have re-created Parameters (stamped with the global rank), which would
+        # otherwise break replicated (disable_tp) weights on a subsequent reload.
+        if hasattr(layer, "update_param_tp_status"):
+            layer.update_param_tp_status()
 
     # Copy processed values into original tensor storage (preserves cudagraph refs)
     # this code is a no-op if not reloading (because kernel tensors is empty)

vllm/model_executor/model_loader/utils.py

@@ -111,6 +111,13 @@ def process_weights_after_loading(
             # parameters onto device for processing and back off after.
             with device_loading_context(module, target_device):
                 quant_method.process_weights_after_loading(module)
+            # process_weights_after_loading may swap in freshly-created
+            # Parameters (e.g. FP8 requantization), which are stamped with the
+            # global rank in BasevLLMParameter.__init__. Re-reconcile their TP
+            # state to the layer so a later weight reload / RL weight-refit
+            # narrows replicated (disable_tp) weights at the correct offset.
+            if hasattr(module, "update_param_tp_status"):
+                module.update_param_tp_status()
             # Repacking transients above can leave large amounts of memory in
             # the caching allocator, which starves the OS on UMA devices.
             release_device_memory_under_pressure(target_device)