[chenhaiq]

What does this PR do?

Enable expandable_segments to avoid the increasing memory fragmentation caused by temporary variables during the training process of fsdp2, which may trigger probabilistic out-of-memory (OOM) errors.

Since both sglang and vllm can not work with expandable_segments:True, it has to be turn off during rollout.

Checklist Before Starting

• Search for similar PRs. Paste at least one query link here: https://github.com/volcengine/verl/pulls?q=is%3Apr+is%3Aopen+fsdp2+
• 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

Without this fix, memory reserved could be very high after compute_log_prob or update_actor.

(WorkerDict pid=339320) [2025-08-11 17:43:01] dp actor After compute_log_prob, memory allocated (GB): 5.53, memory reserved (GB): 73.59, device memory used/total (GB): 77.47/79.15

With this fix, it stays low during training.

(WorkerDict pid=396879) [2025-08-12 07:39:42] dp actor After compute_log_prob, memory allocated (GB): 4.95, memory reserved (GB): 14.20, device memory used/total (GB): 17.72/79.15

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 (飞书群).)

---

Co-authored-by: Lu Hongyu luhongyu.4869@bytedance.com"

[gemini-code-assist]

Code Review

This pull request addresses a probabilistic out-of-memory (OOM) error with FSDP2 by dynamically toggling the expandable_segments setting in PyTorch's CUDA memory allocator. The changes correctly disable this setting for sglang and vllm during rollout and re-enable it for FSDP2 training.

My review focuses on the use of the private PyTorch function _set_allocator_settings. While this seems to be the only way to achieve the desired dynamic behavior, it introduces a maintainability risk as private APIs can change without notice. I've suggested wrapping these calls in try-except blocks to make the code more robust against future PyTorch updates and improve the comments to clarify why this approach is taken.

[zhaochenyang20]

@chenhaiq could you add @narutolhy as a co-author? He contributed the OOM scripts.

[weifengpy]

thanks for the fix! for fsdp2, becuase we allocate many small tensors (model.parameters() sharded by dim-i), we almost should always use torch.cuda.memory._set_allocator_settings(f"expandable_segments:true")

with this PR, the outstanding fragmentation for fsdp2 is only inter-stream fragmentation: 1) all-gather stream has its own memory pool, 2) default stream has its own memory pool. This fragmentation is much less pronouced than expandable_segments. But in case we want to eliminate inter-stream fragmentation, it's possible to call model._set_unshard_async_op(unshard_async_op) and call layer.unshard() manually to to allocate all-gather buffer in the default stream at the cost of less overlapping:

https://github.com/pytorch/pytorch/blob/8d1cf529229dce7cd5ea04abb0faac83b87ca6d1/test/distributed/_composable/fsdp/test_fully_shard_training.py#L447-L448

[chenhaiq]

@chenhaiq could you add @narutolhy as a co-author? He contributed the OOM scripts.

sure.

[vermouth1992]

@chenhaiq could you add @narutolhy as a co-author? He contributed the OOM scripts.

sure.

I guess we can skip CI. How can we add @narutolhy as a co-author?

[chenhaiq]

I guess we can skip CI. How can we add @narutolhy as a co-author?

I have updated commit message with all co-authors. Please merge it.

[Edenzzzz]

Wonder why SGLang can't use expandable segments?

[zhaochenyang20]

Wonder why SGLang can't use expandable segments?

Torch memory savor can not use it. So we turn on in several stages.