执行摘要
重构 vLLM 权重同步逻辑
PR body 说明:"Restructure update_weights_from_ipc to make vllm weigt sync more clear." 目的是使权重同步逻辑更清晰易读,降低维护成本。
建议合并,但应补充测试覆盖 FP8 权重同步路径,以确保重构后行为一致。
无 review 评论,仅有一条 Gemini Code Assist 的自动评论(已停用)。
PR body 说明:"Restructure update_weights_from_ipc to make vllm weigt sync more clear." 目的是使权重同步逻辑更清晰易读,降低维护成本。
建议合并,但应补充测试覆盖 FP8 权重同步路径,以确保重构后行为一致。
无 review 评论,仅有一条 Gemini Code Assist 的自动评论(已停用)。
重构 update_weights_from_ipc 控制流:在 verl/workers/rollout/vllm_rollout/utils.py 中,将原有嵌套的 if/elif 条件重写为按步骤组织的结构,分为准备、接收、后处理三个阶段,每个阶段内部使用清晰的 if/elif/else 链,去除了辅助变量 use_standard_weight_load,并将备选路径(QAT/ModelOpt/LoRA/FP8/普通)配置在同一层级。
简化 vllm_fp8_utils.py 中的 API:将 prepare_quanted_weights_for_loading 和 process_quanted_weights_after_loading 改为直接接收 model 对象而非 model_runner。prepare_quanted_weights_for_loading 现在内部调用 restore_mxfp8_weights_for_loading(原在 load_quanted_weights 中),process_quanted_weights_after_loading 内部调用 apply_mxfp8_transformation_after_loading(原在 load_quanted_weights 中)。load_quanted_weights 移除了 prepare_model 和 process_model 参数,使得调用者必须在外部处理 prepare/process,职责更清晰。
调用方适配:在 update_weights_from_ipc 中,FP8 路径改为对每个模型调用 prepare_quanted_weights_for_loading 并存储结果到列表,后处理时再逐个调用 process_quanted_weights_after_loading。同时移除了 _update_weights 调用中的 quant_prepared 参数。
| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
verl/utils/vllm/vllm_fp8_utils.py |
FP8 工具 | modified | 7.74 |
verl/workers/rollout/vllm_rollout/utils.py |
Rollout | modified | 6.75 |
verl/utils/vllm/vllm_fp8_utils.py
core-logic
核心重构文件,修改了 prepare_quanted_weights_for_loading、process_quanted_weights_after_loading 和 load_quanted_weights 三个函数的签名和内部逻辑,分离了 MXFP8 相关操作,使职责更单一。
def prepare_quanted_weights_for_loading(model):
"""Restore quantized params to the layout their ``weight_loader`` expects.
Must run once before the first bucket, and pairs with
``process_quanted_weights_after_loading``, which re-applies the inference
layout once the last bucket has landed. Both helpers select work by
inspecting the model, so they are no-ops for quantization schemes that need
no restore. The returned value is opaque reload state for the paired call.
"""
# 新:prepare 内部先处理 MXFP8 restore
restore_mxfp8_weights_for_loading(model)
if not is_deepseek_v4_model(model):
return False
return prepare_deepseek_v4_weights_for_loading(model, _copy_param_subclass_attrs)
def process_quanted_weights_after_loading(model, reload_state):
"""Re-apply the inference layout undone by ``prepare_quanted_weights_for_loading``."""
# 新:process 内部先处理 MXFP8 apply
apply_mxfp8_transformation_after_loading(model)
process_deepseek_v4_weights_after_loading(model, reload_state)
def load_quanted_weights(weights, model_runner, is_drafter=False):
# 简化:移除了 prepare_model 和 process_model 参数
if is_drafter:
drafter = getattr(model_runner, "drafter", None)
model = drafter.model if drafter is not None and hasattr(drafter, "model") else None
assert model is not None, ...
else:
model = model_runner.model
quant_config = model_runner.vllm_config.quant_config
vllm_dtype = model_runner.vllm_config.model_config.dtype
weights = list(weights)
cache_deepseek_v4_dense_fp8_scales(model, weights)
weights_quantized = quant_weights(weights, model, quant_config, dtype=vllm_dtype)
# ... monkey patch and load_weights ...
return loaded_params
verl/workers/rollout/vllm_rollout/utils.py
core-logic
主要修改入口,update_weights_from_ipc 重构为核心的三步结构,分支逻辑扁平化。
def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False, use_shm: bool = False):
"""Update the weights of the rollout model."""
from vllm.platforms import current_platform
from verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import BucketedWeightReceiver
if current_platform.device_type == "npu" and self.device is None:
self.device = torch.device(f"npu:{self.local_rank}")
assert self.device is not None
# ===== step 1: prepare for weight loading =====
quant_reload_states = None
if self._is_qat_model:
from verl.utils.qat import prepare_qat_for_load_weights
for model in self._iter_all_models():
prepare_qat_for_load_weights(model, device=self.device)
elif self._is_modelopt_qat:
from verl.utils.modelopt.vllm_modelopt_patch import prepare_modelopt_for_weight_reload
prepare_modelopt_for_weight_reload(self.model_runner.model, device=self.device)
elif peft_config and base_sync_done:
self.remove_lora(VLLM_LORA_INT_ID)
elif is_fp8_model(self.model_runner.vllm_config):
from verl.utils.vllm.vllm_fp8_utils import prepare_quanted_weights_for_loading
quant_reload_states = [
(model, prepare_quanted_weights_for_loading(model)) for model in self._iter_all_models()
]
else:
for model in self._iter_all_models():
patch_vllm_moe_model_weight_loader(model)
# ===== step 2: receive weights and update =====
receiver = BucketedWeightReceiver(...)
receiver.receive_weights(
on_bucket_received=lambda weights: self._update_weights(weights, peft_config=peft_config, base_sync_done=base_sync_done)
)
# ===== step 3: process weights after loading =====
if self._is_qat_model:
for model in self._iter_all_models():
manual_process_weights_after_loading(model)
elif self._is_modelopt_qat:
modelopt_process_weights_after_loading(self.model_runner.model)
elif peft_config and base_sync_done:
pass # no post-process needed
elif is_fp8_model(self.model_runner.vllm_config):
from verl.utils.vllm.vllm_fp8_utils import process_quanted_weights_after_loading
for model, reload_state in quant_reload_states:
process_quanted_weights_after_loading(model, reload_state)
else:
# ... process_weights_after_loading for standard models
当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。
prepare_quanted_weights_for_loading 和 process_quanted_weights_after_loading 的签名变更影响所有调用者,但仅有两处调用(vllm_fp8_utils.py 内的 load_quanted_weights 和 vllm_rollout/utils.py 内的 update_weights_from_ipc),均在本次修改中更新。prepare_model/process_model 参数,并重新分配了 MXFP8 restore/apply 逻辑的归属。如果其他未发现的代码路径(如 drafter 场景)调用了 load_quanted_weights 并依赖这些参数,可能行为不一致。但从 patch 看,drafter 路径也通过 is_drafter 参数得到了处理。test_vllm_fp8_utils_moe_on_cpu.py)未更新,可能导致回归未被捕获。当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。
参与讨论