Fix gpt-oss router_indices in EP
Author: jiqing-feng
Pull request: yes

The indices has been masked by `0`, but the `0` will be recognized as `experts[0]`. We need a new class specific for masking

I run gpt-oss with EP=2, and found both rank0 and rank1 computed expert 0
<img width="800" height="250" alt="image" src="https://github.com/user-attachments/assets/3c2f40b7-9852-4118-9356-39c246605582" />

After this PR, we can see the masking expert is num_expert (16) here, and 16 will be skipped.
<img width="928" height="438" alt="image" src="https://github.com/user-attachments/assets/d719cd98-1a4d-4257-a593-103dc32012dd" />




Comments:
cc @arthurzucker I think?

run-slow: gpt_oss

I also need to change all the other MOE, like mixtral, after this is verified.

You should not need for now because only `gpt_oss` supports EP for now!@

To reproduce the error:
Please run this command and script on a CPU-only device.
`mpirun -np 2 --map-by ppr:1:numa --bind-to numa -genv MASTER_ADDR=127.0.0.1 -genv MASTER_PORT=29500 -genv OMP_NUM_THREADS=32 python tp_hf.py`

```python
import os
import torch
import torch.distributed as dist
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
from transformers.distributed import DistributedConfig


model_id = "lmsys/gpt-oss-20b-bf16"
os.environ['RANK'] = str(os.environ.get('PMI_RANK', 0))
os.environ['LOCAL_RANK'] = str(os.environ.get('PMI_RANK', 0))
os.environ['WORLD_SIZE'] = str(os.environ.get('PMI_SIZE', 1))
rank = int(os.environ['LOCAL_RANK'])
world_size = int(os.environ['WORLD_SIZE'])

def main(rank, world_size) -> None:
    is_tp = world_size > 1
    model_kwargs = dict(torch_dtype=torch.bfloat16)
    if is_tp:
        model_kwargs["tp_plan"] = "auto"
        # model_kwargs["distributed_config"] = DistributedConfig(enable_expert_parallel=1)
    else:
        model_kwargs["device_map"] = "cpu"

    # Retrieve tensor parallel model
    model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
    if dist.is_initialized():
        print("Backend:", dist.get_backend())
    else:
        print("Distributed process group is not initialized.")


    # Retrieve tensor parallel model
    config = AutoConfig.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id, config=config,  **model_kwargs)
    if dist.is_initialized():
        print("Backend:", dist.get_backend())
    else:
        print("Distributed process group is not initialized.")

    # Prepare input tokens
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    prompt = "Once upon a time, there existed a little girl, who liked to have adventures. She wanted to go to places and meet new people, and have fun."
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, do_sample=False, max_new_tokens=32)

    if rank == 0:
        print(tokenizer.batch_decode(outputs, skip_special_tokens=True))


if __name__ == "__main__":
    rank = int(os.environ["RANK"]) if "RANK" in os.environ else 0
    world_size = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    main(rank, world_size)
```

Output before this PR:
```
["Once upon a time, there existed a little girl, who liked to have adventures. She wanted to go to places and meet new people, and have fun. She was a very good friend, and she was a very good\n\nIt sounds like you're sharing a story or a prompt! If you'd like to continue the story"]
```

Output after this PR:
```
['Once upon a time, there existed a little girl, who liked to have adventures. She wanted to go to places and meet new people, and have fun. She was a very good friend, and she was a good friend. She was a good friend. She was a good\n\nIt seems like your text got cut']
```

Output without EP: `python tp_hf.py`
```
['Once upon a time, there existed a little girl, who liked to have adventures. She wanted to go to places and meet new people, and have fun. She was a very good friend, and she was a good friend. She was a good friend. She was a good\n\nIt seems like your message got cut']
```

You can see that the output is almost the same as without EP after this PR.

cc @SunMarc @ArthurZucker 

Hi @ArthurZucker @SunMarc @Rocketknight1 . Would you please review this PR? I've copied the codes to reproduce the issue. You can run it under the cpu-only torch without transformers customized kernels. Installing a cpu-only torch `pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu`

I also added `base_model_ep_plan` in `configuration_gpt_oss.py` so the `distributed_config = DistributedConfig(enable_expert_parallel=1)` can work. Otherwise it has no parallelism implemented as `ep_plan` is None.

Hi @SunMarc . Would you please review this PR? We have other tasks on gpt-oss model which blocked by this PR. Waiting for your review! Thanks!

I read this [blog](https://cookbook.openai.com/articles/gpt-oss/run-transformers#multi-gpu--distributed-inference) and followed the instruction of `distributed_config=DistributedConfig(enable_expert_parallel=1)`, but the `ep_plan` is `None` in gpt-oss model.

I don't know how cuda performs the EP, maybe because cuda uses kernels. But CPU definitely needs to pass ep_plan if we want to enable EP. So I added `base_model_ep_plan` in `configuration_gpt_oss.py`



Sorry I was off for a week!

> Sorry I was off for a week!

@ArthurZucker . No worries. This is a bug fix and the case can be easily reproduced on CPU. Please review this PR as we have a blog pending on this PR being merged, so we can release.

The failed CI is not related to my changes.

Don't worry I'll do the last review today and merge!

**[For maintainers]** Suggested jobs to run (before merge)

run-slow: gpt_oss, mxfp4

Hi @ArthurZucker . I have fixed your comments. Please review it. Thanks!

Thanks 🤗