[fsdp, model] feat: support glm_moe_dsa FSDP training with DSA attention

[Kite0011]

What does this PR do?

Add FSDP training support for the glm_moe_dsa model, including:

• Custom attention forward with DSA (Dynamic Sparse Attention) and Ulysses sequence parallelism support
• Monkey patch registration for glm_moe_dsa model type
• FLOPs estimation function accounting for MLA + DSA indexer + sparse attention

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, veomni, sglang, vllm, rollout, trainer, ci, training_utils, recipe, hardware, deployment, ray, worker, single_controller, misc, perf, model, algo, env, tool, ckpt, doc, data, cfg, reward, fully_async, one_step_off
  • 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

loss converge using a random glm_moe_dsa model 10 step
task: open_reasoning_math sft

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 (飞书群).)
• If your PR is related to the recipe submodule, please also update the reference to the submodule commit via git submodule update --remote or cd recipe && git pull origin main.

[gemini-code-assist]

Code Review

This pull request introduces support for the glm_moe_dsa model type, including a custom attention forward implementation (glm_moe_dsa_attn_forward_with_dsa) that supports Ulysses sequence parallelism and padding removal, corresponding monkey patching, and FLOPs estimation. The review feedback highlights several potential runtime issues: position_ids needs to be unsqueezed to 2D if it is 1D to avoid shape mismatches; 2D or 3D attention_mask inputs should be unsqueezed to 4D before combining with index_mask to prevent crashes; the FLOPs counter must safely handle configurations where q_lora_rank is None; and use_fused_kernels should be explicitly disabled for glm_moe_dsa since they are not supported.