Hi! I think that instead of adding gradient_checkpointing as an argument to training arguments, I used this line when I defined the model:
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
model.gradient_checkpointing_enable()
Anyway we ended up training this model in a GCP, it was too big.
Hope this helps!