[fsdp, sglang] fix: Using Agreesive Empty Cache instead

[zhaochenyang20]

What does this PR do?

Add concise overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review.

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 aims to fix an Out-of-Memory (OOM) issue during Qwen VLM training by replacing the standard PyTorch empty_cache() with a more aggressive_empty_cache function. The change is applied consistently in wake_up and sleep methods of FSDPSGLangShardingManager.

While the intention is good, I've identified a potential issue with the implementation of the new aggressive_empty_cache utility function. Its loop termination logic seems arbitrary and could lead to the fix being either ineffective in some scenarios or causing unnecessary performance overhead. I've added a detailed comment on this. I recommend revisiting the implementation of aggressive_empty_cache to make it more robust.

[gemini-code-assist]

While using a more aggressive cache clearing function is a good approach to mitigate OOM errors, the implementation of aggressive_empty_cache in verl/utils/memory_utils.py appears to have a potential flaw in its loop termination logic.

The loop terminates if the amount of freed reserved memory in an iteration is less than 1GB (reserved_freed < 1024**3). This has two potential drawbacks:

1. Ineffectiveness: The loop might terminate prematurely if it frees just under 1GB of memory, even if more could be freed in subsequent iterations. This could lead to the OOM error persisting in some cases.
2. Performance: If more than 1GB is freed, the loop continues, potentially for up to max_retries times. This could introduce unnecessary latency if a single pass was sufficient.

A more robust approach would be to continue looping as long as a significant amount of memory is being freed, and stop only when an iteration frees little to no memory. This would make the function both more effective at preventing OOMs and more efficient.

Since aggressive_empty_cache is a new core utility for this fix, it would be beneficial to refine its logic to be more reliable before merging.

[wuxibin89]

Should we also modify megatron_sglang.py?

[zhaochenyang20]

Let me do it

[zhaochenyang20]

Without this PR, I profiled Qwen VLM's memory leakage with #3099

We see that each step, there are 0.5GB leaked.

[zhaochenyang20]

Back to march 2025, we can steadily train VLM.