Fix gpt-oss router_indices in EP

[jiqing-feng]

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

After this PR, we can see the masking expert is num_expert (16) here, and 16 will be skipped.

[Rocketknight1]

cc @ArthurZucker I think?

[jiqing-feng]

run-slow: gpt_oss

[jiqing-feng]

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

[ArthurZucker]

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

[jiqing-feng]

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

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

[jiqing-feng]

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

[jiqing-feng]

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.

[jiqing-feng]

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!

[jiqing-feng]

I read this blog 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

[ArthurZucker]

Sorry I was off for a week!

[jiqing-feng]

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.

[jiqing-feng]

The failed CI is not related to my changes.

[ArthurZucker]

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

[ArthurZucker]

Very nice, let's just add a small explanation int he doc and good to go!

[github-actions]

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

run-slow: gpt_oss, mxfp4

[jiqing-feng]

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

[ArthurZucker]

Thanks 🤗