Skip to content

[Bug]: Gemma 4 MoE (26B-A4B) crashes with --data-parallel-size > 1 — AssertionError in cuda_communicator all_gather #38999

Description

@leuski

Your current environment

  • vLLM version: v0.19.0 (also reproduced on vllm/vllm-openai:gemma4-x86_64-cu130)
  • GPU: 8x NVIDIA RTX PRO 6000 Blackwell Server Edition (96GB GDDR7, SM 12.0)
  • OS: Linux 5.15.0-171-generic
  • CUDA: 13.0
  • Python: 3.12
  • Docker image: vllm/vllm-openai:v0.19.0-x86_64-cu130-ubuntu2404 (with transformers>=5.5.0 installed for Gemma 4 support)

Model

google/gemma-4-26B-A4B-it (Mixture-of-Experts, 25.2B total params, 3.8B active, 128 experts with 8 active per token)

Command

vllm serve google/gemma-4-26B-A4B-it \
  --data-parallel-size 2 \
  --quantization fp8 \
  --max-model-len 25600 \
  --max-num-seqs 32 \
  --kv-cache-dtype fp8 \
  --gpu-memory-utilization 0.92 \
  --trust-remote-code

GPUs assigned: 2 GPUs (device IDs 6 and 7).

🐛 Describe the bug

Description

Gemma 4 MoE model crashes at runtime when --data-parallel-size is set to anything greater than 1. The model loads weights successfully, captures CUDA graphs, and starts API servers — but fails on the first inference request with an AssertionError inside the expert-parallel all-gather communication path.

This was also reproduced with --quantization mxfp4 (which crashes earlier during weight loading) and --quantization fp8 (crashes on first request after successful load).

DP=1 works perfectly. The issue is specific to DP > 1 with MoE models.

Dense models (e.g., google/gemma-4-E4B-it, google/gemma-4-31B-it) work fine with DP > 1.

Error traceback

Phase 1: Weight loading (with --quantization mxfp4)

When using --quantization mxfp4, the crash occurs during weight loading:

File "/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/fused_moe/layer.py", line 1073, in weight_loader
    dim2 = loaded_weight.shape[2]
IndexError: tuple index out of range

This is covered by existing issues #35329 and #35324.

Phase 2: Runtime crash (with --quantization fp8)

With fp8, weights load successfully and CUDA graphs are captured. The crash occurs on the first inference request:

File "/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/cuda_communicator.py", in _all_gather_single
AssertionError: 1 != 36

The all-gather operation in the MoE expert dispatch expects a tensor dimension matching sizes[rank] (36 tokens) but receives 1. This suggests the MoE dispatch layer is incorrectly routing through expert-parallelism communication paths when running under data parallelism, causing a tensor shape mismatch between DP workers.

Root cause analysis

The MoE fused expert layer appears to assume that multiple GPUs = expert parallelism (EP), triggering inter-GPU all-gather operations for expert routing. Under data parallelism, each GPU should run an independent full copy of the MoE model with no cross-GPU expert communication. The DP workers seem to be initializing with EP-style communication, leading to the shape assertion failure.

Workaround

Run separate container instances, each with --data-parallel-size 1 on its own GPU, with an external load balancer (e.g., Caddy, nginx) distributing requests round-robin across them. This provides effective data parallelism without triggering the MoE dispatch bug.

Related issues

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't workingunstaleRecieved activity after being labelled stale

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions