KeyError: 'hidden_states'
Author: officialsahyaboutorabi
Pull request: no
Labels: name: bug

### System Info

```
- `transformers` version: 4.57.1
- Platform: Linux-6.6.105+-x86_64-with-glibc2.35
- Python version: 3.12.12
- Huggingface_hub version: 0.36.0
- Safetensors version: 0.6.2
- Accelerate version: 1.11.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.5.1+cu121 (CUDA)
- Tensorflow version (GPU?): 2.19.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.10.7 (gpu)
- Jax version: 0.7.2
- JaxLib version: 0.7.2
- Using distributed or parallel set-up in script?: False
- Using GPU in script?: False
- GPU type: NVIDIA A100-SXM4-80GB
```


### Who can help?

Anybody.

### Information

- [ ] The official example scripts
- [x] My own modified scripts

Hello everyone,

I am running a DPO training script on google colab and when I run the cell dpo_trainer.train()

I get the following error below:

```
[/usr/local/lib/python3.12/dist-packages/transformers/utils/generic.py](https://localhost:8080/#) in wrapped_forward(*args, **kwargs)
   1013             @wraps(orig_forward)
   1014             def wrapped_forward(*args, **kwargs):
-> 1015                 if key == "hidden_states" and len(collected_outputs[key]) == 0:
   1016                     collected_outputs[key] += (args[0],)
   1017                 if kwargs.get("debug_io", False):

KeyError: 'hidden_states'
```

For context. 

Here is my google colab file:

[dpo_integration_for_POLARISANDOTHERSTrainingAndTesting_notebook_triggered_by_n8n(5).ipynb](https://github.com/user-attachments/files/23469233/dpo_integration_for_POLARISANDOTHERSTrainingAndTesting_notebook_triggered_by_n8n.5.ipynb)

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

Steps to reproduce:

All instructions are in the google colab :)

Run all the cells up to the cell with the code: `dpo_trainer.train()`

[dpo_integration_for_POLARISANDOTHERSTrainingAndTesting_notebook_triggered_by_n8n(5).ipynb](https://github.com/user-attachments/files/23469290/dpo_integration_for_POLARISANDOTHERSTrainingAndTesting_notebook_triggered_by_n8n.5.ipynb)

### Expected behavior

I want it to perform DPO training with the google colab on my custom dataset.

Comments:
Hi, we can't debug big notebooks like this for you, unfortunately, because it's hard to tell if the bug is in our code or in yours without a lot of work! Closing for now, but feel free to reopen if you can make a **minimal** reproducer of the issue so that we can triage it.

Alright here is a shorter version: 


```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# --- 1. Configuration ---

max_seq_length = 4096 # Choose any! We auto support RoPE Scaling internally!
dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+

full_model_name = huggingface_username + "/" + model_name

print(f"Loading tokenizer: {full_model_name}")
tokenizer = AutoTokenizer.from_pretrained(full_model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print(f"Loading model in bfloat16 (recommended & stable)")
model = AutoModelForCausalLM.from_pretrained(
    full_model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
    attn_implementation="flash_attention_2",   # Works perfectly in bf16
)

# === CRITICAL QWEN3 DPO FIXES ===
model.config.use_cache = False
model.config.pad_token_id = tokenizer.pad_token_id

# This is the ONLY thing that reliably kills the hidden_states KeyError
for layer in model.model.layers:
    layer.self_attn.collect_hidden_states = False
    if hasattr(layer.self_attn, "q_norm"):
        layer.self_attn.q_norm = None
    if hasattr(layer.self_attn, "k_norm"):
        layer.self_attn.k_norm = None

# Final nuke of the broken hook (belt-and-suspenders)
from transformers.utils import generic
generic.ModelOutput.__call__ = lambda *a, **k: {}

print("Qwen3 DPO fixes applied — training will now work.")
```

Another cell after:

```
#Run the Trainer (DPO version)
from trl import DPOTrainer, DPOConfig
from peft import LoraConfig
import os

# --- Safely unload adapters if the cell is re-run ---
try:
    if hasattr(model, 'unload'):
        model = model.unload()
        print("✅ Successfully unloaded existing PEFT adapters.")
except Exception as e:
    print(f"No PEFT adapters to unload or model is already in base state: {e}")

# Configuration
final_output_dir = "outputs"

# --- LoRA Configuration (same as before) ---
peft_config = LoraConfig(
    r=64,
    lora_alpha=128,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
)

# --- DPO Config ---
dpo_config = DPOConfig(
    output_dir=final_output_dir,
    beta=0.1,
    loss_type="sigmoid",
    max_length=8192,
    max_prompt_length=6144,           # ↑ use more context
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,    # ↑ effective batch 16
    num_train_epochs=3,
    learning_rate=3e-6,               # ↓ lower LR = stable DPO
    warmup_ratio=0.03,                # ↓ shorter warmup
    logging_steps=5,
    save_steps=200,
    optim="paged_adamw_8bit",         # ↑ VRAM-efficient
    weight_decay=0.01,
    lr_scheduler_type="cosine",       # ↑ smoother decay
    seed=42,
    bf16=True,                        # ↑ A100 native
    gradient_checkpointing=True,     # ↑ ~50% VRAM saved
    torch_compile=True,               # ↑ 15–25% faster
    report_to="none",
    remove_unused_columns=False,
)


# --- DPO Trainer ---
dpo_trainer = DPOTrainer(
    model=model,
    ref_model=None,                # Use implicit reference (no separate ref model)
    args=dpo_config,
    train_dataset=dataset,         # must have "prompt", "chosen", "rejected"
    peft_config=peft_config,
)

print("✅ DPOTrainer initialized with LoRA config.")
```

Then the next cell

`dpo_trainer.train()`

```
I get the following error:


The tokenizer has new PAD/BOS/EOS tokens that differ from the model config and generation config. The model config and generation config were aligned accordingly, being updated with the tokenizer's values. Updated tokens: {'bos_token_id': None, 'pad_token_id': 151643}.

---------------------------------------------------------------------------

KeyError                                  Traceback (most recent call last)

/tmp/ipython-input-414979759.py in <cell line: 0>()
----> 1 dpo_trainer.train()

34 frames

/usr/local/lib/python3.12/dist-packages/transformers/utils/generic.py in wrapped_forward(*args, **kwargs)
   1013             @wraps(orig_forward)
   1014             def wrapped_forward(*args, **kwargs):
-> 1015                 if key == "hidden_states" and len(collected_outputs[key]) == 0:
   1016                     collected_outputs[key] += (args[0],)
   1017                 if kwargs.get("debug_io", False):

KeyError: 'hidden_states'
```