[megatron] feat: set_expandable_segments for megatron

[vermouth1992]

What does this PR do?

• As title
• We use set_expandable_segments to resolve memory fragmentation

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 introduces the use of set_expandable_segments to manage memory fragmentation in Megatron-based sharding managers for both SGLang and vLLM. The approach of disabling expandable segments before memory-intensive operations and re-enabling them after is sound. However, this relies on an internal, undocumented PyTorch API, which poses a future compatibility risk. My feedback highlights this concern and suggests making the implementation more robust.

[gemini-code-assist]

The set_expandable_segments function uses torch.cuda.memory._set_allocator_settings, which is an internal PyTorch API. Relying on such private APIs is risky as they are not guaranteed to be stable and can be changed or removed without notice in future PyTorch releases. This could lead to unexpected failures. To make this more robust, I recommend that the set_expandable_segments function itself is modified to include error handling, such as a try-except AttributeError block, to gracefully handle cases where this internal API is no longer available.

[gemini-code-assist]

This call to set_expandable_segments depends on an internal PyTorch function (_set_allocator_settings). This creates a dependency on an unstable, private API that may break in future PyTorch versions. While it addresses memory fragmentation, it introduces a maintainability risk. It would be safer if the set_expandable_segments implementation were updated to catch potential AttributeError and log a warning if the API is not found. This would prevent the application from crashing due to changes in PyTorch internals.