[trainer,rollout] fix: ensure LoRA weights are loaded when vllm_sleep_level=2 and without using layerd_summon

[piood]

What does this PR do?

Fix issue where VLLM would only load base model parameters and not LoRA parameters when VLLM_SLEEP_LEVEL == 2 and not using layered_summon.

This fixes the LoRA trainer error where the first rollout would only use base model parameters, and subsequent rollouts would correctly load LoRA parameters.

Fixes: #3516
Related PR: #3461

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: ...
• Format the PR title as [{modules}] {type}: {description} (This will be checked by the CI)
  • {modules} include fsdp, megatron, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data
  • If this PR involves multiple modules, separate them with , like [megatron, fsdp, doc]
  • {type} is in feat, fix, refactor, chore, test
  • If this PR breaks any API (CLI arguments, config, function signature, etc.), add [BREAKING] to the beginning of the title.
  • Example: [BREAKING][fsdp, megatron] feat: dynamic batching

Test

For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.

API and Usage Example

Demonstrate how the API changes if any, and provide usage example(s) if possible.

# Add code snippet or script demonstrating how to use this

Design & Code Changes

Demonstrate the high-level design if this PR is complex, and list the specific changes.

Checklist Before Submitting

Important

Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.

• Read the Contribute Guide.
• Apply pre-commit checks: pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
• Add / Update the documentation.
• Add unit or end-to-end test(s) to the CI workflow to cover all the code. If not feasible, explain why: ...
• Once your PR is ready for CI, send a message in the ci-request channel in the verl Slack workspace. (If not accessible, please try the Feishu group (飞书群).)

[gemini-code-assist]

Code Review

This pull request addresses a bug where LoRA weights were not correctly loaded during the first rollout when using vLLM with VLLM_SLEEP_LEVEL=2 and without layered_summon. The fix correctly identifies this scenario and forces a full model weight synchronization on the initial rollout by setting base_sync_done to False. Subsequent rollouts will then correctly use the more efficient LoRA-only weight updates. The change also simplifies the logic by removing the now-redundant force_reload flag. The implementation is clean, targeted, and effectively resolves the described issue.

[techkang]

Have you tried to train on your patch with layer_summon=False? When VLLM_SLEEP_LEVEL=2, the base model's weights are destroyed every iteration. After this patch, I think it cannot load weights again. A simple way to fix this is let VLLM_SLEEP_LEVEL=1 whenever lora is enabled.

[piood]

Have you tried to train on your patch with layer_summon=False? When VLLM_SLEEP_LEVEL=2, the base model's weights are destroyed every iteration. After this patch, I think it cannot load weights again. A simple way to fix this is let VLLM_SLEEP_LEVEL=1 。

You are right.There are some problem, i will continue fix and test it.

[WncFht]

@techkang I think the simple way you mentioned maybe not that great. In that way, all the lora situation should set VLLM_SLEEP_LEVEL to be 1, which will make the lora mode cannot offload the base model and potentially when the rollout.load_format equals to dummy, it will never load base model.

So I think we may implement a loading strategy to first load base model and then load the lora adapter in every loading turn. If you agree with me, I will comlete the PR recently, since it may bring more change in the params collecting and model updating logic.

[techkang]

@WncFht Thanks for you discussion. But I still think when rollout.load_format equals to dummy, the base model will be loaded at the first iteration with these code:

verl/verl/utils/fsdp_utils.py

Lines 638 to 651 in d33c85e

	else:
	model = peft_model.base_model.model
	orig_dev = "cpu" if "cpu" in str(next(model.parameters()).device) else get_device_name()
	model = model.to("cpu")
	for name, param in model.state_dict().items():
	if any(x in name for x in ["_flat_param", "lora_"]):
	continue
	name = name.replace("_fsdp_wrapped_module.", "").replace(".base_layer", "")
	lora_params[name] = (
	param.full_tensor().detach().cpu()
	if hasattr(param, "full_tensor")
	else param.detach().cpu()
	)
	model = model.to(orig_dev)

[piood]

@techkang @WncFht I have polished this PR and tested it. It now correctly loads the base model parameters before the LoRA parameters in each iteration when vllm_sleep_level==2 and layer_summon=False.

[techkang]

@piood Thanks, LGTM! Left some minor suggestions.

[techkang]

cc @wuxibin89

[piood]

@techkang I found an issue when LoRA is enabled with vllm_sleep_level=2: when using dummy load format (base_sync_done=False), both params and base_model_params in rollout_mode contain the same base model parameters, causing duplication.

[techkang]

Yes, the duplication will only occur at the first iteration. And the duplication shares the same memory so I don't think it's a serious problem.

[piood]

@techkang Got it, thanks!