fix: filter None router logits in Qwen3 MoE and handle empty router logits (#39203)
Author: SwiftAkira
Pull request: yes

## What does this PR do?

This PR fixes issue #39203 where Qwen3 MoE models crash when mlp_only_layers is non-empty and output_router_logits=True. The issue occurs because MLP-only layers return None router logits, which are incorrectly collected and passed to load_balancing_loss_func, causing a TypeError.

## Root Cause Analysis

The problem was in the router logits collection logic in Qwen3MoeModel.forward(). Unlike Qwen2 MoE which properly filters None values, Qwen3 MoE was collecting all layer outputs without null checks:

- MLP-only layers (specified in mlp_only_layers) return None for router logits since they don't use expert routing
- The original code collected these None values into the router_logits tuple
- When load_balancing_loss_func processes this tuple, it fails on None entries

## Solution

This PR implements two complementary fixes:

1. **Router logits null check**: Added proper filtering during collection to match Qwen2 MoE pattern:
   ```python
   # Before (broken):
   if output_router_logits:
       all_router_logits += (layer_outputs[-1],)
   
   # After (fixed):
   if output_router_logits and layer_outputs[-1] is not None:
       all_router_logits += (layer_outputs[-1],)
   ```

2. **Empty tuple handling**: Added a custom load_balancing_loss_func that gracefully handles the edge case where all layers are MLP-only (resulting in an empty router_logits tuple):
   ```python
   if len(gate_logits) == 0:
       return 0
   ```

## Implementation Details

All changes were made in the modular architecture:
- **Source file**: src/transformers/models/qwen3_moe/modular_qwen3_moe.py (hand-edited)
- **Generated file**: src/transformers/models/qwen3_moe/modeling_qwen3_moe.py (auto-generated)

The fix follows the established pattern from Qwen2 MoE, ensuring consistency across the codebase.

## Testing

Comprehensive testing was performed with various configurations:

1. **Mixed configuration** (mlp_only_layers=[1,3]): 
   - Correctly collects 2 router logits from MoE layers
   - Successfully computes auxiliary loss
   
2. **All MoE configuration** (mlp_only_layers=[]):
   - Collects router logits from all layers
   - Standard auxiliary loss computation
   
3. **All MLP configuration** (mlp_only_layers=[0,1,2,3]):
   - Results in empty router logits tuple
   - Auxiliary loss returns 0 (no routing needed)

All test cases pass without errors.

## Backward Compatibility

This fix is fully backward compatible:
- Existing models continue to work unchanged
- Only adds null checks with minimal performance overhead
- Maintains the same API and behavior for valid configurations

## Fixes

Closes #39203

## How was this patch tested?

- Manual testing with Qwen3 MoE models using different mlp_only_layers configurations
- Verified proper router logits collection and auxiliary loss computation
- Tested edge cases including all-MLP and all-MoE scenarios
- Validated that no None values appear in the final router_logits tuple

cc @ArthurZucker @ntenenz 

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

run-slow: qwen3_moe

# PR Update: Qwen3 MoE Router Logits Fix - Ready for Review

## 🔄 Response to @ArthurZucker's Review

Hi @ArthurZucker! Thank you for the review. I've successfully rebased onto the latest main and thoroughly tested the scenario. While PR #39120 was indeed a comprehensive refactor affecting output handling, **the specific issue with Qwen3 MoE router logits collection still exists** and requires our targeted fix.

## ✅ Testing Confirms the Fix is Still Necessary

After rebasing and extensive testing, I can confirm:

- **✅ Original issue persists**: MLP layers (from `mlp_only_layers`) return `None` router logits when `output_router_logits=True`
- **✅ Fix works correctly**: Our null check prevents crashes and filters `None` values properly
- **✅ Load balancing handles edge cases**: Empty tuples are gracefully handled
- **✅ Backward compatibility**: No impact on normal MoE operation
- **✅ CI compatibility**: Fixed modular conversion issues

## 🔧 Technical Details

**The Root Cause:**
The issue occurs when `mlp_only_layers` is non-empty AND `output_router_logits=True`:

```python
# In Qwen3MoeDecoderLayer.forward()
if isinstance(hidden_states, tuple):
    hidden_states, router_logits = hidden_states  # MoE layers
else:
    router_logits = None  # ← MLP layers return None

# Later in Qwen3MoeModel.forward() - ORIGINAL CODE (problematic)
if output_router_logits:
    all_router_logits += (layer_outputs[-1],)  # ← Crashes with None values!
```

**Why PR #39120 Didn't Fix This:**
- PR #39120 focused on output handling refactoring across models
- The Qwen3 MoE specific issue with `mlp_only_layers` architecture wasn't addressed
- The core problem is architectural: MLP layers correctly return `None` but collection logic didn't handle it

## 🛠️ Our Solution

### 1. Null Check in Router Logits Collection
```python
# 🔧 FIX: Add null check to prevent None router logits from being collected
if output_router_logits and layer_outputs[-1] is not None:
    all_router_logits += (layer_outputs[-1],)
```

### 2. Empty Tuple Handling in Load Balancing Loss
```python
# 🔧 FIX: Handle empty tuple case (when all layers are MLP-only)
if len(gate_logits) == 0:
    return 0
```

### 3. CI Compatibility Fixes
- **Fixed modular/main file sync**: Updated kwargs signatures to match
- **Added missing parameters**: Fixed mask function call consistency
- **Ensured proper imports**: Aligned with transformers conventions

## 🧪 Comprehensive Test Results

I created extensive tests that confirm:

| Test Scenario | Status | Result |
|---------------|---------|---------|
| **Mixed MLP/MoE layers** | ✅ PASS | Forward pass works correctly |
| **Router logits filtering** | ✅ PASS | Only non-None values collected |
| **Empty tuple handling** | ✅ PASS | Load balancing returns 0 gracefully |
| **Normal MoE operation** | ✅ PASS | No regression in standard use |
| **Edge case: All MLP layers** | ✅ PASS | Handles `mlp_only_layers=[0,1,2,3]` |
| **CI modular conversion** | ✅ PASS | Generated/main files match |

**Example Test Output:**
```
Testing original Qwen3 MoE router logits issue...
Config: mlp_only_layers = [0, 2]
Config: decoder_sparse_step = 2
Config: num_hidden_layers = 4
Config: output_router_logits = True

✅ SUCCESS: Forward pass completed without errors!
Router logits shape: 2
Router logits per layer:
  Layer 0: torch.Size([10, 4])  # MoE layer
  Layer 1: torch.Size([10, 4])  # MoE layer
```

## 🎯 Why This Fix is Architecturally Correct

The Qwen3 MoE model intentionally supports mixed architectures:
- **`mlp_only_layers`**: Specific layers use regular MLP instead of MoE
- **Design intention**: Allows fine-grained control over which layers are MoE vs MLP
- **Our fix**: Respects this design by handling `None` router logits from MLP layers

The fix doesn't change model behavior - it just prevents crashes when using the intended architectural feature.

## 📈 Impact Assessment

- **✅ Minimal and targeted**: Only affects router logits collection logic
- **✅ Zero performance impact**: No additional computation in normal paths
- **✅ Maintains full compatibility**: All existing functionality preserved
- **✅ Enables intended use cases**: `mlp_only_layers` now works with `output_router_logits=True`

## 🚀 Ready for Merge

**Current Status:**
- ✅ Rebased onto latest main (includes PR #39120 changes)
- ✅ All tests pass locally
- ✅ CI compatibility verified
- ✅ Code style compliant (ruff formatted)
- ✅ Comprehensive testing completed

**Files Changed:**
- `src/transformers/models/qwen3_moe/modeling_qwen3_moe.py`: Main fix + load balancing improvement
- `src/transformers/models/qwen3_moe/modular_qwen3_moe.py`: Modular version + CI compatibility

The fix is **minimal, necessary, and ready for production**. It solves a real crash scenario while maintaining full backward compatibility and following Hugging Face's coding standards.

## 🙏 Thank You

Thank you for the thorough review process! The rebase and additional testing have confirmed that this fix is still essential even after the comprehensive changes in PR #39120. Ready for final approval! 🎉