gpt-oss dtensor worker v2 does not support tensor parallel
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**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**
```bash
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.
Contributor guide
Assessment
This issue has not been assessed yet.