Error when running eval on Mamba LORA with PEFT
I am running the following codefrom datasets import load_dataset
from trl import SFTTrainer
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf")
model = AutoModelForCausalLM.from_pretrained("state-spaces/mamba-130m-hf")
train_dataset = load_dataset("spider", split="train").select(range(50))
val_dataset = load_dataset("spider", split="validation").select(range(50))
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    logging_dir='./logs',
    logging_steps=10,
    learning_rate=2e-3
)
lora_config =  LoraConfig(
        r=8,
        target_modules=["x_proj", "embeddings", "in_proj", "out_proj"],
        task_type="CAUSAL_LM",
        bias="none"
)
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    peft_config=lora_config,
    eval_dataset=val_dataset,
    train_dataset=train_dataset,
    dataset_text_field="question",
)
trainer.train()
trainer.evaluate()Based on this nice gist from@ArthurZgist linkto do LORA on Mamba.When I try to run evaluate I stumble upon the following error:warnings.warn(
Traceback (most recent call last):
  File "/home/ppol/MaskMamba/src/test.py", line 33, in <module>
    trainer.evaluate()
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/transformers/trainer.py", line 3964, in evaluate
    output = eval_loop(
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/transformers/trainer.py", line 4158, in evaluation_loop
    losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/transformers/trainer.py", line 4374, in prediction_step
    loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/transformers/trainer.py", line 3625, in compute_loss
    outputs = model(**inputs)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py", line 194, in forward
    return self.gather(outputs, self.output_device)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py", line 217, in gather
    return gather(outputs, output_device, dim=self.dim)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 135, in gather
    res = gather_map(outputs)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 127, in gather_map
    return type(out)((k, gather_map([d[k] for d in outputs])) for k in out)
  File "<string>", line 7, in __init__
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/transformers/utils/generic.py", line 390, in __post_init__
    for idx, element in enumerate(iterator):
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 127, in <genexpr>
    return type(out)((k, gather_map([d[k] for d in outputs])) for k in out)
  File "/home/ppol/.conda/envs/mask_mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 130, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
TypeError: 'MambaCache' object is not iterableIf you could please point out what I am doing wrong it would be very helpful.
2024-10-29T23:36:46Z
Maybe a bug.github.com/huggingface/transformersTypeError: 'MambaCache' object is not iterableopened08:35AM - 07 Mar 24 UTCclosed02:30AM - 08 Mar 24 UTCrangehow### System Info

Hi @ArthurZucker, I met some problems in training mamba. In the…training phase, loss decreased as expected, while in eval phase, the error occurred.




logs

```
{'loss': 18.0794, 'grad_norm': 1740.3226318359375, 'learning_rate': 0.0019999999702625887, 'epoch': 0.0}                                                                        
{'loss': 17.8266, 'grad_norm': 4142.9677734375, 'learning_rate': 0.001999999881050357, 'epoch': 0.0}                                                                            
{'loss': 15.1686, 'grad_norm': 2723.932373046875, 'learning_rate': 0.00199999973236331, 'epoch': 0.01}                                                                          
{'loss': 13.9637, 'grad_norm': 121.74279022216797, 'learning_rate': 0.0019999995242014562, 'epoch': 0.01}                                                                       
{'loss': 12.0022, 'grad_norm': 76.66282653808594, 'learning_rate': 0.001999999256564809, 'epoch': 0.01}                                                                         
{'loss': 8.4667, 'grad_norm': 93.78446960449219, 'learning_rate': 0.001999998929453383, 'epoch': 0.01}                                                                          
{'loss': 7.4828, 'grad_norm': 14.276519775390625, 'learning_rate': 0.001999998542867199, 'epoch': 0.02}                                                                         
{'loss': 7.2502, 'grad_norm': 39.938140869140625, 'learning_rate': 0.001999998096806279, 'epoch': 0.02}                                                                         
{'loss': 7.1519, 'grad_norm': 17.629228591918945, 'learning_rate': 0.0019999975912706495, 'epoch': 0.02}                                                                        
{'loss': 6.9367, 'grad_norm': 11.50907039642334, 'learning_rate': 0.001999997026260341, 'epoch': 0.02}  
Traceback (most recent call last):
  File "/data/ruanjh/mamba-training/train.py", line 65, in <module>
    trainer.train()
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 1626, in train
    return inner_training_loop(
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 2034, in _inner_training_loop
    self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 2418, in _maybe_log_save_evaluate
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 3204, in evaluate
    output = eval_loop(
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 3393, in evaluation_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 3610, in prediction_step
    loss, outputs = self.compute_loss(model, inputs, return_outputs=True)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 2900, in compute_loss
    outputs = model(**inputs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py", line 186, in forward
    return self.gather(outputs, self.output_device)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py", line 203, in gather
    return gather(outputs, output_device, dim=self.dim)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 104, in gather
    res = gather_map(outputs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 95, in gather_map
    return type(out)((k, gather_map([d[k] for d in outputs]))
  File "<string>", line 7, in __init__
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/utils/generic.py", line 393, in __post_init__
    for idx, element in enumerate(iterator):
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 95, in <genexpr>
    return type(out)((k, gather_map([d[k] for d in outputs]))
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/nn/parallel/scatter_gather.py", line 99, in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
TypeError: 'MambaCache' object is not iterable
```

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

dataset is iwslt17 available on hf but cleaned a little by myself(you can just replace this part using dataset.load_dataset())
```
from transformers import MambaForCausalLM,AutoTokenizer,Seq2SeqTrainer,DataCollatorForSeq2Seq,Seq2SeqTrainingArguments,Trainer,TrainingArguments
import torch
from dataset import MyDataset
import json
model_dir='/data/ruanjh/mamba-370m-hf'
output_dir='./mamba-translate'
tokenizer=AutoTokenizer.from_pretrained(model_dir)

model=MambaForCausalLM.from_pretrained(model_dir,torch_dtype=torch.bfloat16)

collator=DataCollatorForSeq2Seq(tokenizer,model)

with open('/iwslt17/train.json') as f:
    train_data=json.load(f)
train_dataset=MyDataset(train_data,tokenizer)

with open('/iwslt17/validation.json') as f:
    eval_data=json.load(f)
eval_dataset=MyDataset(eval_data,tokenizer)


trainer = Trainer(
        model=model,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
        args=TrainingArguments(
            overwrite_output_dir =True,
            remove_unused_columns =False,
            gradient_accumulation_steps=8,
            #------------------------------
            evaluation_strategy='steps',
            # eval_delay=100,
            eval_steps =50,
            #-------------------------------
            save_strategy ='steps',
            save_steps = 50,
            save_total_limit =3,
            load_best_model_at_end=True,
            #--------------------------------
            dataloader_num_workers =10,
            learning_rate=2e-3,
            num_train_epochs=30,
            # auto_find_batch_size=True,
            per_device_train_batch_size=4,
            per_device_eval_batch_size =4,
            output_dir="./mamba-translate",
            logging_steps=5,
            bf16=True,
            prediction_loss_only=True,
            lr_scheduler_type="cosine",
            # save_safetensors =False,
        ),
        data_collator=collator,
    )

trainer.train()

```

### Expected behavior

.
2024-10-30T10:25:54Z
nice thanks!
2024-10-30T22:25:58Z
This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.