Cannot fine-tune T5Gemma with Seq2SeqTrainer
Author: Crissium
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.56.1
- Platform: Linux-6.8.0-60-generic-x86_64-with-glibc2.39
- Python version: 3.12.11
- Huggingface_hub version: 0.34.4
- Safetensors version: 0.5.3
- Accelerate version: 1.8.0
- Accelerate config: 	not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.7.1 (CUDA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: Parallel
- Using GPU in script?: Yes
- GPU type: NVIDIA H100 NVL

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

Sample code:

```python
import datasets
import datetime
import jiwer
import os
import torch
import tracto
import transformers


data_date = '0901'
training_date = datetime.datetime.now().strftime('%m%d')
os.environ['WANDB_PROJECT'] = 'ASR Error Simulation'
os.environ['WANDB_NOTES'] = f'Model: T5Gemma Small UL2, Data: Splits{data_date}'
dataset_root = f'/mnt/users_home/cpii.local/yxing/Workspace/ASR/Data/Generated/Transcripts/Splits{data_date}'

ds_train = datasets.load_dataset(os.path.join(dataset_root, 'Train'), split='train')
ds_val = datasets.load_dataset(os.path.join(dataset_root, 'Val'), split='validation')

model = transformers.T5GemmaForConditionalGeneration.from_pretrained('google/t5gemma-s-s-ul2', attn_implementation='eager')
model.generation_config.max_new_tokens = 90
tokeniser = transformers.GemmaTokenizerFast.from_pretrained('google/t5gemma-s-s-ul2')
tokeniser.model_max_length = 90


def preprocess_data(examples):
	return tokeniser(
		examples['gt'],
		text_target=examples['predicted'],
		padding='max_length',
		truncation=True
	)


def normalise_text(text: str) -> str:
	text = text.replace('-', ' ')
	return ''.join(c for c in text.lower() if not tracto.is_punct(c) or c == "'")


def compute_wer_difference(eval_pred: transformers.EvalPrediction) -> dict[str, float]:
	gt_text = tokeniser.batch_decode(eval_pred.inputs, skip_special_tokens=True)
	original_pred = tokeniser.batch_decode(eval_pred.label_ids, skip_special_tokens=True)
	model_pred = tokeniser.batch_decode(eval_pred.predictions, skip_special_tokens=True)
	original_wer = jiwer.wer(list(map(normalise_text, gt_text)), list(map(normalise_text, original_pred)))
	model_wer = jiwer.wer(list(map(normalise_text, gt_text)), list(map(normalise_text, model_pred)))
	return {'original_wer': original_wer, 'model_wer': model_wer, 'wer_difference': original_wer - model_wer}


data_collator = transformers.DataCollatorForSeq2Seq(tokeniser, model=model)

training_args = transformers.Seq2SeqTrainingArguments(
	run_name=f'{training_date} Small T5Gemma UL2',
	report_to='wandb',
	output_dir=f'../Checkpoints/{training_date} Small T5Gemma UL2',
	eval_strategy='epoch',
	save_strategy='epoch',
	logging_strategy='epoch',
	include_for_metrics=['inputs'],
	num_train_epochs=500,
	learning_rate=1e-5,
	weight_decay=0.01,
	predict_with_generate=True,
	per_device_train_batch_size=150,
	per_device_eval_batch_size=150,
	ddp_find_unused_parameters=False
)

with training_args.main_process_first(desc='dataset map pre-processing'):
	ds_train = ds_train.map(preprocess_data, batched=True)
	ds_val = ds_val.map(preprocess_data, batched=True)

	ds_train = ds_train.remove_columns([col for col in ds_train.column_names if col not in ['input_ids', 'attention_mask', 'labels']])
	ds_val = ds_val.remove_columns([col for col in ds_val.column_names if col not in ['input_ids', 'attention_mask', 'labels']])

trainer = transformers.Seq2SeqTrainer(
	model=model,
	args=training_args,
	data_collator=data_collator,
	train_dataset=ds_train,
	eval_dataset=ds_val,
	compute_metrics=compute_wer_difference
)

trainer.train()

torch.distributed.destroy_process_group()
```

Sample data:
```json
{"source":"Common Voice","filename":"common_voice_en_41645384.mp3","gt":"Can all this be happening?","predicted":"Can all this be happening?"}
{"source":"Common Voice","filename":"common_voice_en_39702948.mp3","gt":"Whether he represented winter or summer is not quite clear.","predicted":"Whether he represented winter or summer is not quite clear."}
{"source":"Common Voice","filename":"common_voice_en_38381818.mp3","gt":"They are believed to have spoken an Iberian language.","predicted":"They are believed to have spoken an Iberian language."}
{"source":"Common Voice","filename":"common_voice_en_32260000.mp3","gt":"Glas Bheinn lies between Ullapool and Durness.","predicted":"The last binin lies between Olapu and Durnes."}
{"source":"Common Voice","filename":"common_voice_en_27647177.mp3","gt":"Cuba is in the fourth stage of demographic transition.","predicted":"Cuba is in the fourth stage of demographic transition."}
{"source":"Common Voice","filename":"common_voice_en_22373557.mp3","gt":"The terminology is made confusing by the etymology of these words.","predicted":"The terminology is made confusing by the etymology of these words."}
{"source":"Common Voice","filename":"common_voice_en_37213325.mp3","gt":"There is also the City Island Seaside Trolley run by the Bronx Tourism Council.","predicted":"There is also the City Island Seaside Trolley run by the Bronx Tourism Council."}
{"source":"Common Voice","filename":"common_voice_en_19687576.mp3","gt":"These went from Mauchline to the Isle of Wight.","predicted":"These went from McElwain to the Italy of Wight."}
{"source":"Common Voice","filename":"common_voice_en_30574664.mp3","gt":"The crater is very old and is crisscrossed by chains of secondary impact craters.","predicted":"The crater is very old and is crisscrossed by chains of secondary impact craters."}
{"source":"Common Voice","filename":"common_voice_en_19794969.mp3","gt":"Its motto, on the nameplate below the title, is One of America's Great Newspapers.","predicted":"Its motto on the nameplate below the title is one of America's great newspapers."}
```

Error:
```
[rank1]: Traceback (most recent call last):
[rank1]:   File "/mnt/users_home/cpii.local/yxing/Workspace/ASR/Noise/src/train.py", line 86, in <module>
[rank1]:     trainer.train()
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer.py", line 2328, in train
[rank1]:     return inner_training_loop(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer.py", line 2788, in _inner_training_loop
[rank1]:     self._maybe_log_save_evaluate(
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer.py", line 3227, in _maybe_log_save_evaluate
[rank1]:     metrics = self._evaluate(trial, ignore_keys_for_eval)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer.py", line 3176, in _evaluate
[rank1]:     metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
[rank1]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer_seq2seq.py", line 191, in evaluate
[rank1]:     return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer.py", line 4469, in evaluate
[rank1]:     output = eval_loop(
[rank1]:              ^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer.py", line 4665, in evaluation_loop
[rank1]:     losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
[rank1]:                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/trainer_seq2seq.py", line 327, in prediction_step
[rank1]:     generated_tokens = self.model.generate(**generation_inputs, **gen_kwargs)
[rank1]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank1]:     return func(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/generation/utils.py", line 2399, in generate
[rank1]:     self._prepare_cache_for_generation(
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/generation/utils.py", line 2007, in _prepare_cache_for_generation
[rank1]:     else EncoderDecoderCache(DynamicCache(**dynamic_cache_kwargs), DynamicCache(**dynamic_cache_kwargs))
[rank1]:                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/cache_utils.py", line 1018, in __init__
[rank1]:     for _ in range(config.num_hidden_layers)
[rank1]:                    ^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/mnt/users_home/cpii.local/yxing/miniconda3/envs/g/lib/python3.12/site-packages/transformers/configuration_utils.py", line 207, in __getattribute__
[rank1]:     return super().__getattribute__(key)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: AttributeError: 'T5GemmaConfig' object has no attribute 'num_hidden_layers'
```


### Expected behavior

`Seq2SeqTrainer` works with `T5GemmaForConditionalGeneration`.

Comments:
cc @SunMarc 

It seems related to how the trainer handles parallel/distributed training with this model. A potential workaround could involve reviewing the training arguments or adjusting optimizer setting probab. It would be great if the maintainers could provide guidance or a fix for better GPU parallel support.

The issue will be fixed here https://github.com/huggingface/transformers/pull/40939

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.