KeyError: 'hidden_states'
Author: officialsahyaboutorabi
Pull request: no

### 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'
```

### 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:

Alright here is a shorter version: 

libraries:

```
%%capture
# THIS IS THE NEW, CORRECTED SCRIPT

# Step 1: Uninstall conflicting default libraries from the Colab environment.
# This prevents version clashes that cause the errors you are seeing.
!pip uninstall -y torch torchvision torchaudio transformers accelerate datasets tokenizers

!pip install -q torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

!pip install -q -U bitsandbytes
!pip install -q -U datasets trl

#!pip uninstall -y transformers tokenizers accelerate
!pip install transformers tokenizers accelerate peft -U

# Step 3: Install your application-specific libraries.
!pip install -q streamlit pyngrok logging huggingface_hub

print("✅ Installation and upgrade complete.")
```

in another cell run:

```
!pip install -q flash_attn
```


```
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'
```



### Expected behavior

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

Comments:
I shortened it so you don't have to go through the entire google colab to find the problem. But you can see the entire file if you want to. 

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

Sorry, this is still too messy for us to deal with. The reproducer should just be a single block that we can paste, with all lines that don't need to be there to see the bug removed, but at least 50% of this is just code agent slop! For example, we can't tell if gradient checkpointing is the cause, or if it's the DPOTrainer, or LoRA/peft, or the `paged_adamw_8bit` optimizer, or the interaction between multiple of those things.

The goal of a minimal reproducer is that you should do the work to figure out exactly where the bug is, and what conditions trigger it. Ideally, you should have code that works, with as few lines as possible, and a single change that makes the bug appear, so we can see exactly what's going on. Pasting us a huge chunk of code with lots of unnecessary variables just makes us do the work of figuring that for you, and we don't have time!

> Sorry, this is still too messy for us to deal with. The reproducer should just be a single block that we can paste, with all lines that don't need to be there to see the bug removed, but at least 50% of this is just code agent slop! For example, we can't tell if gradient checkpointing is the cause, or if it's the DPOTrainer, or LoRA/peft, or the `paged_adamw_8bit` optimizer, or the interaction between multiple of those things.
> 
> The goal of a minimal reproducer is that you should do the work to figure out exactly where the bug is, and what conditions trigger it. Ideally, you should have code that works, with as few lines as possible, and a single change that makes the bug appear, so we can see exactly what's going on. Pasting us a huge chunk of code with lots of unnecessary variables just makes us do the work of figuring that for you, and we don't have time!


Well here is the extremely short version. Below.

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'
```



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.