diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index d92b9fc7d008..c0dbc776acb1 100644 --- a/vllm/model_executor/layers/linear.py +++ b/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( diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index eb609d63c50d..92f454f9a5f4 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/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) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index 6be057bff082..3367f4833e6c 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/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)