[model] refactor: polishing FSDP model engine

[vermouth1992]

What does this PR do?

• Extract a separate prepare_micro_batches
• Fix prepare_dynamic_batch
• Make forward_step more modular

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 refactors the FSDP model engine for better modularity. It extracts prepare_micro_batches into a utility function, fixes an issue in prepare_dynamic_batch by deep-copying meta_info, and makes the forward_step method more modular by splitting it into input preparation, model execution, and output preparation steps. The refactoring is well-executed and improves code structure. However, I've found a critical issue in the VLM input preparation logic that was present before but should be fixed as part of this polishing effort.

[gemini-code-assist]

The check for is_vlm_model incorrectly uses micro_batch_tensor to look for multi_modal_inputs. However, multi_modal_inputs are handled as non-tensor data and are present in micro_batch_non_tensor. This will cause is_vlm_model to always be False, leading to incorrect padding logic for VLM models when Ulysses sequence parallelism is enabled.

Suggested change

	is_vlm_model = "multi_modal_inputs" in micro_batch_tensor.keys()
	is_vlm_model = "multi_modal_inputs" in micro_batch_non_tensor.keys()

[vermouth1992]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces several valuable refactorings to the FSDP model engine, significantly improving code modularity and maintainability. Key changes include extracting prepare_micro_batches into a utility function, and breaking down the monolithic forward_step method in FSDPEngineWithLMHead into smaller, more manageable prepare_model_inputs and prepare_model_outputs methods. The PR also includes important correctness fixes, such as ensuring each micro-batch gets a deep copy of meta_info to prevent side-effects, and correctly sourcing multi_modal_inputs from the non-tensor part of the data batch. Overall, these changes enhance code quality and fix subtle bugs. The implementation is solid and I have no further suggestions.