gpt-oss dtensor worker v2 does not support tensor parallel

[jordane95]

Describe the bug

When fine-tuning gpt-oss using dtensor worker v2 and tensor parallel, some tp plan is not recognized by fsdp, for example, attention sink use local_rowwise, moe use gather or grouped_gemm. After adding support for those plans in translate function, it still raise type error when callling parallelize_module since moe parallel strategy imported from transformers lib has different type to ParallelStyle.

Steps/Code to reproduce bug

python examples/run_sft.py cluster.gpus_per_node=4 policy.model_name=openai/gpt-oss-20b policy.dtensor_cfg.tensor_parallel_size=4

Expected behavior

A clear and concise description of what you expected to happen.

Additional context

Add any other context about the problem here.

[coderabbitai]

📝 CodeRabbit Plan Mode

Generate an implementation plan and prompts that you can use with your favorite coding agent.

• Create Plan

Examples

• Example 1
• Example 2

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord or schedule a call!

[terrykong]

@hemildesai to take a look