feat: general fsdp2 on non-MoE models + HF TP plan

[yuki-666]

What does this PR do ?

1. Support FSDP2 on non-MoE models.
2. Support Hugging Face TP plan.
3. The priority for using which parallel plan is custom-parallel-plan > opt-parallel-plan (which we implemented for certain models in FSDP2) > hf-tp-plan (HF's _tp_plan).

Convergence test on LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM, Gemma2ForCausalLM, Gemma3ForCausalLM, Phi3ForCausalLM run well.

Convergence Test Detail

Llama-3.1-8B-Instruct (LlamaForCausalLM)
FSDP2-tp8-opt_plan vs FSDP2-tp8-hf_tp_plan

Qwen2ForCausalLM / Qwen3ForCausalLM

Qwen2.5-7B-Instruct (Qwen2ForCausalLM) FSDP2-tp4-opt_plan vs FSDP2-tp4-hf_tp_plan	Qwen3-0.6B (Qwen3ForCausalLM) FSDP1 vs FSDP2-tp1

Gemma2ForCausalLM / Gemma3ForCausalLM

gemma-2-9b-it (Gemma2ForCausalLM) FSDP1 vs FSDP2-tp1 vs FSDP2-tp4-hf_tp_plan	gemma-3-1b-it (Gemma3ForCausalLM) FSDP1 vs FSDP2-tp1

Issues

Closes #156

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[yuki-666]

File another issue #413 to trace FSDP2 for MoE models.

1. Qwen3-30B-A3B is obviously slower than Qwen3-32B, especially on the refit process or when using hf-tp-plan with dtensor tp > 1.
2. DeepseekV2ForCausalLM using fsdp2 will fail on the following error on model.layers.0.self_attn.rotary_emb.cos_cached, said v.shape=torch.Size([2048, 64]) and self.reference_model_buffers[k].shape=torch.Size([163840, 64]) shape mismatch in self.use_reference_model().

• DeepseekV2ForCausalLM should use vllm==0.8.5 and add this line https://github.com/vllm-project/vllm/blob/v0.9.0/vllm/v1/attention/backends/mla/common.py#L494 to vllm for test.

[yuki-666]

@yuki-666 what were the parameters of your gemma2 run? I can't seem to get it to run correctly:

uv run examples/run_grpo_math.py policy.model_name=google/gemma-2-2b-it logger.wandb_enabled=True cluster.gpus_per_node=8 +policy.generation.vllm_cfg.load_format=auto

@terrykong Thanks very much for pointing out this!

I tested with almost the same script as yours before this commit fdb565c.
After this commit, load_format of vllm during training is default set to dummy, only specific models will change it through nemo_rl/models/huggingface/common.py. The param policy.generation.vllm_cfg.load_format is removed from yaml and has no effect even if we pass it.

It is fixed now, and other models won't be affect since they don't need special handle on load_format.

[yuki-666]

Thanks @jgerh , have updated from your suggestions.

[terrykong]

Thanks for the quick fix @yuki-666 . Gemma2 seems to be okay now from a quick run: