NVIDIA-NeMo / NVIDIA-NeMo/RL

gpt-oss dtensor worker v2 does not support tensor parallel

Open
#1,684 3 comments 1 reaction 1 assignee Claimed by @hemildesai View on GitHub
bug community-request waiting-on-customer
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.