Fix transformers MoE compatibility with VLLM 
Author: 3outeille
Pull request: yes
Labels: name: for_v5?

- Since v5 pushes, it introduces MoE refactor which break compatibility with VLLM. I dont think `torch.zeros_like(router_logits).scatter_(1, router_indices, router_top_value)` is needed (argument was it is needed during EP but if EP is enabled, the whole function is ignore anyway)

NOTE: EP was broken already so need to check following PR https://github.com/huggingface/transformers/pull/42183

Take away:

- we don't want to use `zeros_like + scatter` so find another solution and fix breaking tests (cf https://github.com/huggingface/transformers/actions/runs/19738722845/job/56557321765) => Follow up PR https://github.com/huggingface/transformers/pull/42456
- unrelated to current change, vLLM needs fixing to adjust to transformers weight loading cc: @hmellor 


# Test results
<img width="2402" height="989" alt="image" src="https://github.com/user-attachments/assets/0e549049-e7b0-4414-a743-7b7da67bda4c" />

# Testing methodology

`mkdir tests_results && ./test_moe_inference.sh tests_results/`
```python
#!/bin/bash

# Script to test vLLM inference with transformers backend for MoE models
# Each model is tested in parallel on separate GPUs

# Define colors for output
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
DIM='\033[0;90m'
NC='\033[0m' # No Color

# Define models to test (architecture -> HuggingFace checkpoint)
declare -A MODELS=(
    ["flex_olmo"]="hf-internal-testing/tiny-random-FlexOlmoForCausalLM"
    ["gpt_oss"]="hf-internal-testing/tiny-random-GptOssForCausalLM"
    ["minimax"]="hf-internal-testing/tiny-random-MiniMaxForCausalLM"
    ["mixtral"]="hf-internal-testing/tiny-random-MixtralForCausalLM"
    ["olmoe"]="hf-internal-testing/tiny-random-OlmoeForCausalLM"
    ["qwen2_moe"]="hf-internal-testing/tiny-random-Qwen2MoeForCausalLM"
    ["qwen3_moe"]="hf-internal-testing/tiny-random-Qwen3MoeForCausalLM"
    ["qwen3_next"]="hf-internal-testing/tiny-random-Qwen3NextForCausalLM"
    ["qwen3_vl_moe"]="hf-internal-testing/tiny-random-Qwen3VLMoeForConditionalGeneration"
)

# Get number of available GPUs
NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
if [ "$NUM_GPUS" -eq 0 ]; then
    echo "No GPUs detected!"
    exit 1
fi
echo "Detected $NUM_GPUS GPUs"

# Handle results directory - use provided path or create temp directory
if [ -n "$1" ]; then
    RESULTS_DIR="$1"
    mkdir -p "$RESULTS_DIR"
    CLEANUP_RESULTS=false
elif [ -n "$RESULTS_DIR" ]; then
    # RESULTS_DIR already set via environment variable
    mkdir -p "$RESULTS_DIR"
    CLEANUP_RESULTS=false
else
    RESULTS_DIR=$(mktemp -d)
    CLEANUP_RESULTS=true
fi

# Only cleanup if we created a temp directory
if [ "$CLEANUP_RESULTS" = true ]; then
    trap "rm -rf $RESULTS_DIR" EXIT
fi

echo "Results directory: $RESULTS_DIR"

echo "=========================================="
echo "  MoE Models Inference Test Script"
echo "  (Parallel execution on $NUM_GPUS GPUs)"
echo "=========================================="
echo ""

# Function to run inference test on a specific GPU
run_test() {
    local model_name=$1
    local model_checkpoint=$2
    local gpu_id=$3
    local result_file="$RESULTS_DIR/${model_name}.result"
    
    echo -e "${YELLOW}[GPU $gpu_id] Starting: ${model_name} (${model_checkpoint})${NC}"
    
    # Set environment and run the command on specific GPU
    CUDA_VISIBLE_DEVICES=$gpu_id VLLM_ENABLE_V1_MULTIPROCESSING=0 \
        python vllm/examples/offline_inference/basic/generate.py \
        --model "$model_checkpoint" \
        --model-impl transformers \
        --enforce-eager \
        --no-enable-prefix-caching \
        > "$RESULTS_DIR/${model_name}.log" 2>&1
    
    local exit_code=$?
    
    # Write result to file (for collection later)
    if [ $exit_code -eq 0 ]; then
        echo "SUCCESS" > "$result_file"
        echo -e "${GREEN}✓ [GPU $gpu_id] ${model_name}: SUCCESS${NC}"
    else
        echo "FAILED (exit code: $exit_code)" > "$result_file"
        echo -e "${RED}✗ [GPU $gpu_id] ${model_name}: FAILED (exit code: $exit_code)${NC}"
    fi
}

# Convert associative array keys to indexed array for scheduling
MODEL_NAMES=(${!MODELS[@]})
NUM_MODELS=${#MODEL_NAMES[@]}

# Track PIDs for waiting
declare -a PIDS=()

# Launch tests in parallel, cycling through available GPUs
for i in "${!MODEL_NAMES[@]}"; do
    model_name="${MODEL_NAMES[$i]}"
    model_checkpoint="${MODELS[$model_name]}"
    gpu_id=$((i % NUM_GPUS))
    
    # If we've used all GPUs, wait for a slot to free up
    if [ ${#PIDS[@]} -ge $NUM_GPUS ]; then
        # Wait for any one process to complete
        wait -n 2>/dev/null || wait "${PIDS[0]}"
        # Remove completed PIDs (simplified: just clear and rebuild)
        NEW_PIDS=()
        for pid in "${PIDS[@]}"; do
            if kill -0 "$pid" 2>/dev/null; then
                NEW_PIDS+=("$pid")
            fi
        done
        PIDS=("${NEW_PIDS[@]}")
    fi
    
    run_test "$model_name" "$model_checkpoint" "$gpu_id" &
    PIDS+=($!)
done

# Wait for all remaining background jobs to complete
echo ""
echo "Waiting for all tests to complete..."
wait

# Print summary
echo ""
echo "=========================================="
echo "  SUMMARY"
echo "=========================================="
echo ""

success_count=0
fail_count=0

for model_name in "${MODEL_NAMES[@]}"; do
    result_file="$RESULTS_DIR/${model_name}.result"
    if [ -f "$result_file" ]; then
        result=$(cat "$result_file")
        if [[ "$result" == "SUCCESS" ]]; then
            echo -e "${GREEN}✓ ${model_name}: ${result}${NC}"
            ((success_count++))
        else
            echo -e "${RED}✗ ${model_name}: ${result}${NC}"
            # Show last few lines of error
            echo -e "${DIM}  Error snippet:"
            tail -n 3 "$RESULTS_DIR/${model_name}.log" | while read -r line; do echo -e "    ${DIM}${line}${NC}"; done
            ((fail_count++))
        fi
    else
        echo -e "${RED}✗ ${model_name}: NO RESULT (test may have crashed)${NC}"
        ((fail_count++))
    fi
done

echo ""
echo "-------------------------------------------"
echo -e "Total: ${GREEN}${success_count} passed${NC}, ${RED}${fail_count} failed${NC}"
echo "=========================================="

# Show logs for failed tests
if [ $fail_count -gt 0 ]; then
    echo ""
    echo "Failed test logs available in: $RESULTS_DIR"
    echo "To view: cat $RESULTS_DIR/<model_name>.log"
fi
```

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_42429). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

run-slow: flex_olmo, gpt_oss, minimax, mixtral, olmoe, qwen2_moe, qwen3_moe, qwen3_next, qwen3_omni_moe, qwen3_vl_moe

This comment contains `run-slow`, running the specified jobs: 

models: ["models/flex_olmo", "models/gpt_oss", "models/minimax", "models/mixtral", "models/olmoe", "models/qwen2_moe", "models/qwen3_moe", "models/qwen3_next", "models/qwen3_omni_moe", "models/qwen3_vl_moe"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/19738722845)

✅ No failing test specific to this PR 🎉 !


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

run-slow: flex_olmo, gpt_oss, minimax, mixtral, olmoe, qwen2_moe, qwen3_moe, qwen3_next, qwen3_omni_moe, qwen3_vl_moe