[Bugfix] Re-sync parameter tp_rank after process_weights_after_loading (fix replicated / disable_tp weight reload) (#48025)
Signed-off-by: Alex Xu <alexxu@roblox.com> Co-authored-by: YQ-Wang <yiqingwang@roblox.com> Co-authored-by: alexhxu <alex.xu1015@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
alexxu-roblox committed
Jul 18, 2026 at 01:40 UTC
d96aee09518a57fdfe1e83369d2f4ab515f78c30
3 files changed
+23
-7
vllm/model_executor/layers/linear.py
+11
-7
@@ -281,6 +281,15 @@ class LinearBase(PluggableLayer):
281
self.tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
282
283
def update_param_tp_status(self):
284
+ # Single source of truth for a parameter's TP state. BasevLLMParameter
285
+ # stamps self.tp_rank with the *global* rank in __init__; this reconciles
286
+ # every child parameter to the *layer's* tp_rank/tp_size (which correctly
287
+ # accounts for disable_tp -> replicated weights with tp_rank == 0).
288
+ #
289
+ # Must be re-run whenever parameters are (re-)created after construction,
290
+ # e.g. after quant_method.process_weights_after_loading() swaps in fresh
291
+ # Parameters. Otherwise a later load_weights()/weight-refit would narrow a
292
+ # replicated weight at global_rank * shard_size and overflow.
293
for param in self.parameters():
294
if isinstance(param, BasevLLMParameter):
295
param.tp_rank = self.tp_rank
@@ -910,7 +919,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
919
shard_id=loaded_shard_id,
920
shard_offset=shard_offset,
921
shard_size=shard_size,
913
- tp_rank=self.tp_rank,
922
)
923
924
def load_weights(
@@ -1110,12 +1118,10 @@ class QKVParallelLinear(ColumnParallelLinear):
1118
# to ensure that any subsequent reduction (like .max())
1119
# works correctly while preserving the parameter shape.
1120
for idx in range(param.data.shape[0]):
1113
- param.load_qkv_weight(
1114
- loaded_weight=loaded_weight, shard_id=idx, tp_rank=self.tp_rank
1115
- )
1121
+ param.load_qkv_weight(loaded_weight=loaded_weight, shard_id=idx)
1122
return
1123
elif type(param) in (RowvLLMParameter, BasevLLMParameter):
1118
- param.load_qkv_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank)
1124
+ param.load_qkv_weight(loaded_weight=loaded_weight)
1125
return
1126
# TODO: @dsikka - move to parameter.py
1127
self._load_fused_module_from_checkpoint(param, loaded_weight)
@@ -1139,7 +1145,6 @@ class QKVParallelLinear(ColumnParallelLinear):
1145
shard_id=loaded_shard_id,
1146
shard_offset=shard_offset,
1147
shard_size=shard_size,
1142
- tp_rank=self.tp_rank,
1148
)
1149
1150
def weight_loader(
@@ -1493,7 +1498,6 @@ class MinimaxM3QKVParallelLinearWithIndexer(QKVParallelLinear):
1498
shard_id=loaded_shard_id,
1499
shard_offset=shard_offset,
1500
shard_size=shard_size,
1496
- tp_rank=self.tp_rank,
1501
)
1502
1503
def weight_loader(
vllm/model_executor/model_loader/reload/layerwise.py
+5
@@ -364,6 +364,11 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
364
quant_method = getattr(layer, "quant_method", None)
365
if isinstance(quant_method, QuantizeMethodBase):
366
quant_method.process_weights_after_loading(layer)
367
+ # Re-reconcile parameter TP state: process_weights_after_loading may
368
+ # have re-created Parameters (stamped with the global rank), which would
369
+ # otherwise break replicated (disable_tp) weights on a subsequent reload.
370
+ if hasattr(layer, "update_param_tp_status"):
371
+ layer.update_param_tp_status()
372
373
# Copy processed values into original tensor storage (preserves cudagraph refs)
374
# this code is a no-op if not reloading (because kernel tensors is empty)
vllm/model_executor/model_loader/utils.py
+7
@@ -111,6 +111,13 @@ def process_weights_after_loading(
111
# parameters onto device for processing and back off after.
112
with device_loading_context(module, target_device):
113
quant_method.process_weights_after_loading(module)
114
+ # process_weights_after_loading may swap in freshly-created
115
+ # Parameters (e.g. FP8 requantization), which are stamped with the
116
+ # global rank in BasevLLMParameter.__init__. Re-reconcile their TP
117
+ # state to the layer so a later weight reload / RL weight-refit
118
+ # narrows replicated (disable_tp) weights at the correct offset.
119
+ if hasattr(module, "update_param_tp_status"):
120
+ module.update_param_tp_status()
121
# Repacking transients above can leave large amounts of memory in
122
# the caching allocator, which starves the OS on UMA devices.
123
release_device_memory_under_pressure(target_device)