TypeError: 'MambaCache' object is not iterable
Author: rangehow
Pull request: no

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

.

Comments:
This bug also appears in my partner's training（Although our code is written completely independently and separately）

Ah yes, when evaluating you should probably set use_cache = False if you are not doing generation because the cache_param will be returned otherwise. That is a good issue!
There should be a better way to fix this

The issue was fixed after editing `use_cache` in `config.json`.  Hope to have a solution that is imperceptible to user in future. 

Hi @ArthurZucker, I met another question that seems mamba implementation not compatible with zero-offload.
I train mamba-2.8b-hf on 4 3090 using zero-2 + optimizer offload and get error as follow:

> trainer.train(resume_from_checkpoint=True)
      File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 1626, in train
self._optimizer_step(i)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1805, in _optimizer_step
    self.optimizer.step()
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/optim/optimizer.py", line 385, in wrapper
    return inner_training_loop(
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 1963, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/transformers/trainer.py", line 2886, in training_step
    out = func(*args, **kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    self.accelerator.backward(loss)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/accelerate/accelerator.py", line 1960, in backward
    return func(*args, **kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py", line 179, in step
    self.deepspeed_engine_wrapped.backward(loss, **kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/accelerate/utils/deepspeed.py", line 176, in backward
    self.engine.step()
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2169, in step
    group['weight_decay'], group['bias_correction'], p.data, p.grad.data,
KeyError: 'bias_correction'
    self._take_model_step(lr_kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 2075, in _take_model_step
    self.optimizer.step()
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1856, in step
    self._optimizer_step(i)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1805, in _optimizer_step
    self.optimizer.step()
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/optim/optimizer.py", line 385, in wrapper
    out = func(*args, **kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/data/ruanjh/miniconda3/envs/mamba/lib/python3.10/site-packages/deepspeed/ops/adam/cpu_adam.py", line 179, in step
    group['weight_decay'], group['bias_correction'], p.data, p.grad.data,
KeyError: 'bias_correction'

while using zero2+gradient checkpointing is ok.  : )

Hi @rangehow, if the first problem was resolved and you're encountering a new one, could you open a new issue? This helps us better track what has and has not been resolved