modelling_llama -> spda_attention; ValueError: too many values to unpack (expected 4)
Author: hpcpony
Pull request: no
Labels: name: bug

### System Info

torch = 2.6.0
transformers = 4.51.0 and 4.51.2




### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

I'm poking around just trying to do a simple "hello world" (being a newbie and all).  I've managed to fill in the blanks and actually get to the point where the Trainer will run.  I'm sure there are plenty of problems with what I'm doing and how I'm doing it, but I'm stuck at this error.  It's not obvious what I'm doing to cause it.

With the following code...

```
# Util.py
from transformers import LlamaConfig, LlamaForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import Dataset
import os

##
## Load an empty model
##
def load_model(model_dir, TRACE=False) :
    rank = os.environ.get('RANK')

    # get just the config from the pretrained directory
    if TRACE : print("[{}] Loading Configuration".format(rank))
    config = LlamaConfig.from_pretrained(model_dir)
    if TRACE : print("Config:\n", config)

    # create an empty model based on just the confic (i.e. no weights)
    if TRACE : print("[{}] Loading Model".format(rank))
    model = LlamaForCausalLM(config)
    if TRACE : print("Model:\n", model)

    # get the tokenizer
    print("[{}] Loading Tokenizer".format(rank))
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    tokenizer.pad_token = "
"        # for PreTrainedTokenizerFast and Llama 3.1
    if TRACE : print("Tokenizer:\n", tokenizer)

    # return model and tokenizer
    return model, tokenizer

##
## Create some training data
##
def load_dataset(TRACE=False) :
    rank = os.environ.get('RANK')
    if TRACE : print("[{}] Loading Dataset".format(rank))
    # make up data with a couple of entries
#    dataset = Dataset.from_dict({
#        "input_ids":["The quick brown fox jumped over the lazy dog's back"]*20
#        })
    dataset = Dataset.from_dict({
        "text":[["The quick brown fox jumped over the lazy dog's back"*20]]
        })
    # synthesize a bogus dataset with 'train' and 'test' sets
    datasets = {
        'train': dataset, 
        'test' : dataset
        }
    if TRACE : 
      print("Dataset\n")
      print("train\n")
      pprint.pp(datasets['train'][0:5])
      print("test\n")
      pprint.pp(datasets['test'][0:5])
    return datasets

```
and...

```
# 00_trainer_8B.py
import os
import time
import Util

from transformers import Trainer, TrainingArguments, DataCollatorWithPadding

model_dir = '/work/AI/Models/NousResearch/Meta-Llama-3.1-8B'
TRACE = True

from torch.distributed.elastic.multiprocessing.errors import record
import torch.distributed
from datetime import timedelta

@record
def main(): 
    # setup process group to not time out 'cause our slow systems
    torch.distributed.init_process_group(
     timeout=timedelta(minutes=20)
    )
    rank = os.environ.get('RANK')
    # get empty model and tokenizer
    model, tokenizer = Util.load_model(model_dir, TRACE=TRACE)
    # get bogus dataset
    dataset = Util.load_dataset(TRACE=False)
    # setup the model to the GPU
    if TRACE : print("[{}] Moving model to GPU".format(rank))
    model.to('cuda')
    ##
    ## Do the training
    ## from: https://huggingface.co/docs/transformers/en/trainer
    ##
    if TRACE : print("[{}] Setting up training".format(rank))
    # setup the parameters for the training
    training_args = TrainingArguments(
      output_dir="outputs/"+model_dir.split('/')[-1],
      learning_rate=2e-5,
      per_device_train_batch_size=4,
      per_device_eval_batch_size=4,
      num_train_epochs=4,
      weight_decay=0.01,
      eval_strategy="epoch",
      save_strategy="no",         # no, epoch, steps, best
      bf16=True,
      #label_names=['text'],    # huh?
      #local_rank= ,
      #load_best_model_at_end=True,
      push_to_hub=False,
    )
    # Tokenize the datasets
    dataset['train'] = dataset['train'].map(lambda x: tokenizer(x['text']))  # tokenize explicitly
    dataset['test'] = dataset['train'].map(lambda x: tokenizer(x['text']))
    print("Train> ", dataset['train'])
    print("Test> ", dataset['test'])
    # instantiate the Trainer
    data_collator = DataCollatorWithPadding(tokenizer)   # default
    trainer = Trainer(
      model=model,
      args=training_args,
      train_dataset=dataset["train"],
      eval_dataset=dataset["test"],
      processing_class=tokenizer,
      data_collator=data_collator,
      #compute_metrics=compute_metrics,
    )
    # Train...
    if TRACE : print("[{}] Training Start...".format(rank))
    trainer.train()
    if TRACE : print("[{}] ...Training Done".format(rank))

import torch
if __name__ == "__main__" :
    print("CUDA is available ",torch.cuda.is_available())
    print("Devices ", torch.cuda.device_count())
    main()
```
Here's execution....
```
[gpu:trl] torchrun --standalone --nnodes 1 00_trainer_8B.py 
CUDA is available  True
Devices  8
[0] Loading Configuration
Config:
 LlamaConfig {
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 128000,
  "eos_token_id": 128001,
  "head_dim": 128,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 14336,
  "max_position_embeddings": 131072,
  "mlp_bias": false,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 8,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-05,
  "rope_scaling": {
    "factor": 8.0,
    "high_freq_factor": 4.0,
    "low_freq_factor": 1.0,
    "original_max_position_embeddings": 8192,
    "rope_type": "llama3"
  },
  "rope_theta": 500000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.51.2",
  "use_cache": true,
  "vocab_size": 128256
}

[0] Loading Model
Model:
 LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 4096)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        )
        (mlp): LlamaMLP(
          (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
          (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
          (act_fn): SiLU()
        )
        (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((4096,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=4096, out_features=128256, bias=False)
)
[0] Loading Tokenizer
Tokenizer:
 PreTrainedTokenizerFast(name_or_path='/work/AI/Models/NousResearch/Meta-Llama-3.1-8B', vocab_size=128000, model_max_length=131072, is_fast=True, padding_side='right', truncation_side='right', special_tokens={'bos_token': '
', 'eos_token': '
', 'pad_token': '
'}, clean_up_tokenization_spaces=True, added_tokens_decoder={
	128000: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128001: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128002: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128003: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128004: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128005: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128006: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128007: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128008: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128009: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128010: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128011: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128012: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128013: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128014: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128015: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128016: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128017: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128018: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128019: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128020: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128021: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128022: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128023: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128024: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128025: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128026: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128027: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128028: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128029: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128030: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128031: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128032: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128033: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128034: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128035: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128036: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128037: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128038: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128039: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128040: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128041: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128042: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128043: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128044: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128045: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128046: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128047: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128048: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128049: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128050: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128051: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128052: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128053: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128054: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128055: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128056: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128057: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128058: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128059: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128060: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128061: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128062: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128063: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128064: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128065: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128066: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128067: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128068: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128069: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128070: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128071: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128072: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128073: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128074: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128075: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128076: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128077: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128078: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128079: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128080: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128081: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128082: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128083: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128084: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128085: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128086: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128087: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128088: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128089: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128090: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128091: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128092: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128093: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128094: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128095: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128096: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128097: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128098: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128099: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128100: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128101: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128102: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128103: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128104: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128105: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128106: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128107: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128108: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128109: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128110: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128111: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128112: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128113: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128114: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128115: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128116: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128117: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128118: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128119: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128120: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128121: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128122: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128123: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128124: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128125: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128126: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128127: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128128: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128129: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128130: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128131: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128132: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128133: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128134: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128135: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128136: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128137: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128138: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128139: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128140: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128141: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128142: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128143: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128144: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128145: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128146: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128147: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128148: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128149: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128150: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128151: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128152: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128153: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128154: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128155: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128156: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128157: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128158: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128159: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128160: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128161: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128162: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128163: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128164: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128165: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128166: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128167: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128168: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128169: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128170: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128171: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128172: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128173: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128174: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128175: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128176: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128177: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128178: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128179: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128180: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128181: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128182: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128183: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128184: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128185: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128186: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128187: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128188: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128189: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128190: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128191: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128192: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128193: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128194: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128195: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128196: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128197: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128198: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128199: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128200: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128201: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128202: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128203: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128204: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128205: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128206: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128207: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128208: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128209: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128210: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128211: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128212: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128213: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128214: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128215: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128216: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128217: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128218: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128219: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128220: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128221: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128222: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128223: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128224: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128225: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128226: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128227: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128228: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128229: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128230: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128231: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128232: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128233: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128234: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128235: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128236: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128237: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128238: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128239: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128240: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128241: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128242: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128243: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128244: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128245: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128246: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128247: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128248: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128249: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128250: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128251: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128252: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128253: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128254: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
	128255: AddedToken("
", rstrip=False, lstrip=False, single_word=False, normalized=False, special=True),
}
)
[0] Moving model to GPU
[0] Setting up training
Map: 100%|█████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 20.59 examples/s]
Map: 100%|████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 261.38 examples/s]
Train>  Dataset({
    features: ['text', 'input_ids', 'attention_mask'],
    num_rows: 1
})
Test>  Dataset({
    features: ['text', 'input_ids', 'attention_mask'],
    num_rows: 1
})
[0] Training Start...
  0%|                                                                                             | 0/4 [00:00<?, ?it/s][rank0]: Traceback (most recent call last):
[rank0]:   File "/work/trl/00_trainer_8B.py", line 75, in <module>
[rank0]:     main()
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
[rank0]:     return f(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^
[rank0]:   File "/work/trl/00_trainer_8B.py", line 68, in main
[rank0]:     trainer.train()
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 2245, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 2560, in _inner_training_loop
[rank0]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 3736, in training_step
[rank0]:     loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 3801, in compute_loss
[rank0]:     outputs = model(**inputs)
[rank0]:               ^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/parallel/distributed.py", line 1643, in forward
[rank0]:     else self._run_ddp_forward(*inputs, **kwargs)
[rank0]:          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/parallel/distributed.py", line 1459, in _run_ddp_forward
[rank0]:     return self.module(*inputs, **kwargs)  # type: ignore[index]
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/accelerate/utils/operations.py", line 814, in forward
[rank0]:     return model_forward(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/accelerate/utils/operations.py", line 802, in __call__
[rank0]:     return convert_to_fp32(self.model_forward(*args, **kwargs))
[rank0]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/utils/generic.py", line 965, in wrapper
[rank0]:     output = func(self, *args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 821, in forward
[rank0]:     outputs: BaseModelOutputWithPast = self.model(
[rank0]:                                        ^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/utils/generic.py", line 965, in wrapper
[rank0]:     output = func(self, *args, **kwargs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 571, in forward
[rank0]:     layer_outputs = decoder_layer(
[rank0]:                     ^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 318, in forward
[rank0]:     hidden_states, self_attn_weights = self.self_attn(
[rank0]:                                        ^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 274, in forward
[rank0]:     attn_output, attn_weights = attention_interface(
[rank0]:                                 ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 30, in sdpa_attention_forward
[rank0]:     key = repeat_kv(key, module.num_key_value_groups)
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 11, in repeat_kv
[rank0]:     batch, num_key_value_heads, slen, head_dim = hidden_states.shape
[rank0]:     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: ValueError: too many values to unpack (expected 4)
  0%|                                                                                             | 0/4 [00:00<?, ?it/s]
[rank0]:[W412 16:57:36.231382427 ProcessGroupNCCL.cpp:1496] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
E0412 16:57:44.945000 263707 /opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/api.py:869] failed (exitcode: 1) local_rank: 0 (pid: 263778) of binary: /opt/AI/trl-0.16.1/bin/python
E0412 16:57:44.963000 263707 /opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/error_handler.py:141] no error file defined for parent, to copy child error file (/tmp/torchelastic_4cfdgcr3/2cd61db4-dfff-458a-9ec7-954de88c8652_ms05csxo/attempt_0/0/error.json)
Traceback (most recent call last):
  File "/opt/AI/trl-0.16.1/bin/torchrun", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/run.py", line 918, in main
    run(args)
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/run.py", line 909, in run
    elastic_launch(
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 138, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 269, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
00_trainer_8B.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2025-04-12_16:57:35
  host      : gpu.super.org
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 263778)
  error_file: /tmp/torchelastic_4cfdgcr3/2cd61db4-dfff-458a-9ec7-954de88c8652_ms05csxo/attempt_0/0/error.json
  traceback : Traceback (most recent call last):
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 355, in wrapper
      return f(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^
    File "/work/trl/00_trainer_8B.py", line 68, in main
      trainer.train()
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 2245, in train
      return inner_training_loop(
             ^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 2560, in _inner_training_loop
      tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 3736, in training_step
      loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/trainer.py", line 3801, in compute_loss
      outputs = model(**inputs)
                ^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
      return self._call_impl(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
      return forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/parallel/distributed.py", line 1643, in forward
      else self._run_ddp_forward(*inputs, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/parallel/distributed.py", line 1459, in _run_ddp_forward
      return self.module(*inputs, **kwargs)  # type: ignore[index]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
      return self._call_impl(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
      return forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/accelerate/utils/operations.py", line 814, in forward
      return model_forward(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/accelerate/utils/operations.py", line 802, in __call__
      return convert_to_fp32(self.model_forward(*args, **kwargs))
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
      return func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/utils/generic.py", line 965, in wrapper
      output = func(self, *args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/utils/deprecation.py", line 172, in wrapped_func
      return func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 821, in forward
      outputs: BaseModelOutputWithPast = self.model(
                                         ^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
      return self._call_impl(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
      return forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/utils/generic.py", line 965, in wrapper
      output = func(self, *args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 571, in forward
      layer_outputs = decoder_layer(
                      ^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
      return self._call_impl(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
      return forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 318, in forward
      hidden_states, self_attn_weights = self.self_attn(
                                         ^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
      return self._call_impl(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
      return forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 274, in forward
      attn_output, attn_weights = attention_interface(
                                  ^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 30, in sdpa_attention_forward
      key = repeat_kv(key, module.num_key_value_groups)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 11, in repeat_kv
      batch, num_key_value_heads, slen, head_dim = hidden_states.shape
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  ValueError: too many values to unpack (expected 4)
  
============================================================
```

I poked around and the value of hidden_states.shape is: torch.Size([1, 221, 1, 8, 128])

### Expected behavior

Informative/useful error message, or no value error

Comments:
Hi @hpcpony! /) 

It's kind of hard to piece through all of that code to figure out the source of the problem! Rather than just pasting your training scripts, can you make a **minimal** reproducer script for us, ideally the shortest code possible that still triggers the bug? It'll make it a lot easier for us to track down any issue!

Is this better for you?


```
# bug2.py
from transformers import LlamaConfig, LlamaForCausalLM, AutoTokenizer
from datasets import Dataset
from transformers import TrainingArguments, Trainer

model_dir = '/work/AI/Models/NousResearch/Meta-Llama-3.1-8B'

if __name__ == "__main__" :

    config = LlamaConfig.from_pretrained(model_dir)
    model = LlamaForCausalLM(config)
    model.to('cuda')
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    tokenizer.pad_token = "
"
    datasets = {
        'train': Dataset.from_dict({
           "text":[["The quick brown fox jumped over the lazy dog's back"*20]]
           }),
        'test' : Dataset.from_dict({
           "text":[["The quick brown fox jumped over the lazy dog's back"*20]]
          })
        }
    datasets['train'] = datasets['train'].map(lambda x: tokenizer(x['text']))
    datasets['test'] = datasets['train'].map(lambda x: tokenizer(x['text']))
    training_args = TrainingArguments(
      output_dir="outputs/"+model_dir.split('/')[-1],
      per_device_train_batch_size=4,
      per_device_eval_batch_size=4,
      num_train_epochs=4,
      eval_strategy="epoch",
      save_strategy="no",         # no, epoch, steps, best
      push_to_hub=False,
    )
    trainer = Trainer(
      model=model,
      args=training_args,
      train_dataset=datasets["train"],
      eval_dataset=datasets["test"],
      processing_class=tokenizer,
    )
    trainer.train()
```

Reproduce with:

```
   torchrun --standalone --nnodes 1 bug2.py

```


Got it! cc @sunmarc, another trainer/torchrun issue

I had also tried accelerate and it fails in the same manner.
(don't know much about all the accelerate settings so they're mostly "good guesses")

```
accelerate launch --config-file 1p.yaml ./bug2.py
```

```
# 1p.yaml
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: 'NO'
downcast_bf16: 'no'
enable_cpu_affinity: false
gpu_ids: all
machine_rank: 0
main_training_function: main
mixed_precision: 'no'
num_machines: 1
num_processes: 1
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false
```

```
  .
  .
  .
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/llama/modeling_llama.py", line 274, in forward
    attn_output, attn_weights = attention_interface(
                                ^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 30, in sdpa_attention_forward
    key = repeat_kv(key, module.num_key_value_groups)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/integrations/sdpa_attention.py", line 11, in repeat_kv
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: too many values to unpack (expected 4)
```



I don't know if this helps, but I tried collecting some more information.  

I thought if tried another model and/or stopped using torchrun/accelerate it might make a difference.

Essentially the same "simple" program as before with a couple small changes:

* rearraged order so long model load comes last
* generalized to handle either llama or bloom model

```
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from datasets import Dataset
from datasets.utils.logging import disable_progress_bar
from transformers import TrainingArguments, Trainer

model_dir = '/work/AI/Models/NousResearch/Meta-Llama-3.1-8B'
model_dir = '/work/AI/Models/bigscience/bloom-3b'

if __name__ == "__main__" :
    disable_progress_bar()
    ## Create some training data
    datasets = {
        'train': Dataset.from_dict({
           "text":[["The quick brown fox jumped over the lazy dog's back"*20]]
           }),
        'test' : Dataset.from_dict({
           "text":[["The quick brown fox jumped over the lazy dog's back"*20]]
          })
        }

    ## Create the tokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    if not tokenizer.pad_token : tokenizer.pad_token = "
"   # for Llama
    ## Tokenize datasets
    datasets['train'] = datasets['train'].map(lambda x: tokenizer(x['text'], padding=True))
    datasets['test'] = datasets['test'].map(lambda x: tokenizer(x['text'], padding=True))

    ## Load the empty model
    config = AutoConfig.from_pretrained(model_dir)
    model = AutoModelForCausalLM.from_config(config)
    model.to('cuda')

    ## Set up training
    training_args = TrainingArguments(
      output_dir="outputs/"+model_dir.split('/')[-1],
      per_device_train_batch_size=4,
      per_device_eval_batch_size=4,
      num_train_epochs=4,
      eval_strategy="epoch",
      save_strategy="no",         # no, epoch, steps, best
      push_to_hub=False,
      disable_tqdm=True,
    )
    
    trainer = Trainer(
      model=model,
      args=training_args,
      train_dataset=datasets["train"],
      eval_dataset=datasets["test"],
      processing_class=tokenizer,
    )

    ## Train
    trainer.train()
```
Running using neither torchrun nor accelerate:
```
python ./bug3.py
```
The error is different, but similar in that it indicates a mismatch of shapes.  I don't understand the code well enough to tell if this is really the same problem or not.
```
  .
  .
  .
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/bloom/modeling_bloom.py", line 1001, in forward
    transformer_outputs = self.transformer(
                          ^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/AI/trl-0.16.1/lib/python3.12/site-packages/transformers/models/bloom/modeling_bloom.py", line 658, in forward
    batch_size, seq_length, _ = inputs_embeds.shape
    ^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: too many values to unpack (expected 3)
```
poking around, inputs_embeds.shape = torch.Size([1, 1, 220, 2560])


You should use a dataset from HF instead of passing your own custom dataset. The dataset that you passed is incompatible as you passed a nested list, hence after tokenization, we have an additional dimension. Also, you to pas `labels` with you want to train the model. You can have a look at this course if you want to get started with transformers ! https://huggingface.co/learn/llm-course/en/chapter3/1?fw=pt

Fix : remove the list 

```python
    dataset = Dataset.from_dict({
        "text":["The quick brown fox jumped over the lazy dog's back"*20]
        })
```

I need to use my own dataset, but i will certainly check out the course to
see how to do that properly.

Thanks for the help.

On Thu, Apr 17, 2025, 07:02 Marc Sun ***@***.***> wrote:

> You should use a dataset from HF instead of passing your own custom
> dataset. The dataset that you passed is incompatible as you passed a nested
> list, hence after tokenization, we have an additional dimension. Also, you
> to pas labels with you want to train the model. You can have a look at
> this course if you want to get started with transformers !
> https://huggingface.co/learn/llm-course/en/chapter3/1?fw=pt
>
>     dataset = Dataset.from_dict({
>         "text":[["The quick brown fox jumped over the lazy dog's back"*20]]
>         })
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/37470#issuecomment-2812516716>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ALWTX42EEPOC4XHOFNLJTGL2Z6C4VAVCNFSM6AAAAAB3AMSHAWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMJSGUYTMNZRGY>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
> *SunMarc* left a comment (huggingface/transformers#37470)
> <https://github.com/huggingface/transformers/issues/37470#issuecomment-2812516716>
>
> You should use a dataset from HF instead of passing your own custom
> dataset. The dataset that you passed is incompatible as you passed a nested
> list, hence after tokenization, we have an additional dimension. Also, you
> to pas labels with you want to train the model. You can have a look at
> this course if you want to get started with transformers !
> https://huggingface.co/learn/llm-course/en/chapter3/1?fw=pt
>
>     dataset = Dataset.from_dict({
>         "text":[["The quick brown fox jumped over the lazy dog's back"*20]]
>         })
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/issues/37470#issuecomment-2812516716>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/ALWTX42EEPOC4XHOFNLJTGL2Z6C4VAVCNFSM6AAAAAB3AMSHAWVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDQMJSGUYTMNZRGY>
> .
> You are receiving this because you were mentioned.Message ID:
> ***@***.***>
>