KeyError in Llama-4-Maverick-17B-128E-Instruct-FP8 Inference with Offloading
Author: pchu2025
Pull request: no

### Issue Description
Loading `meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8` succeeds with `transformers==4.51.0`, but inference fails with `KeyError: 'model.layers.37.feed_forward.experts.gate_up_proj'` during `model.generate`. This occurs on 4x NVIDIA RTX A6000 (~196GB VRAM, CUDA 12.4, Python 3.12.3, Ubuntu 24.04.2) with offloading, critical for sentiment analysis (~100–150GB/day, ~85–90% accuracy). Disabling MoE (`num_experts=0`) didn’t resolve it.

### Steps to Reproduce
1. Install dependencies:
   ```bash
   pip install torch==2.4.1 accelerate==1.7.0 compressed-tensors==0.9.4 transformers==4.51.0

2. Confirm model files (~389GB, 84 .safetensors) at /mnt/data/ai_super_palace/models/llama4/.

3. Run:
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
os.environ["TORCHVISION_DISABLE_NMS"] = "1"
model = AutoModelForCausalLM.from_pretrained(
    '/mnt/data/ai_super_palace/models/llama4',
    torch_dtype=torch.float16,
    device_map="auto",
    low_cpu_mem_usage=True,
    offload_folder="/mnt/data/ai_super_palace/models/llama4/offload",
    config={"parallel_style": "none"}
)
tokenizer = AutoTokenizer.from_pretrained('/mnt/data/ai_super_palace/models/llama4')
prompt = "What is the sentiment of this text: 'I love this product, it's amazing!'"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

4. Error:
KeyError: 'model.layers.37.feed_forward.experts.gate_up_proj'

**Environment**
Transformers: 4.51.0
Python: 3.12.3
PyTorch: 2.4.1
CUDA: 12.4
Accelerate: 1.7.0
Compressed-tensors: 0.9.4
OS: Ubuntu 24.04.2 LTS
Hardware: 4x NVIDIA RTX A6000 (~196GB VRAM)
Model: meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8

**Additional Details**
Model card requires transformers>=4.51.0, supports FP8 via compressed-tensors.
Warnings: Uninitialized MoE weights (feed_forward.experts.*), offloaded parameters (VRAM limit).
Prior errors (TypeError: NoneType not iterable) resolved with config={"parallel_style": "none"}.
Suspect bug in accelerate offloading or MoE weight initialization.

**Request**
Is this a known llama4 MoE offloading issue?
Can MoE weights be initialized or offloading fixed?
Workaround for inference without re-downloading (~389GB)?
Urgent for sentiment analysis.

**Logs**
See traceback above. config.json (40KB) available.

Thank you!














Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi @pchu2025, I'm sorry, I missed this issue the first time! cc @mekkcyber for quantization

Hi @pchu2025, can you share the device_map of the model : `model.hf_device_map`

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.