Possible issue with learning rate scheduler while using Fairscale in Seq2Seq Trainer
Author: arijitthegame
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.21.1
- Platform: Linux-5.4.0-1078-aws-x86_64-with-glibc2.27
- Python version: 3.8.13
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.11.0+cu102 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes, 8 GPUs
- Using distributed or parallel set-up in script?: Yes

### Who can help?

@patrickvonplaten, @Narsil  I am running a summarization task on Mimic Dataset using T5 v1-large. I am using p3.16x EC2 instance which has 8 V100 GPUs with 16G each. I am using the seq2seq trainer with the config below. However, I am getting the following warning as in the screenshot pasted below. However, turning off fairscale removes this warning. Thank you once again for your help.
<img width="1026" alt="Screen Shot 2022-08-25 at 4 13 40 PM" src="https://user-images.githubusercontent.com/25961440/186759804-312d725a-6910-48e1-ab71-ec531557a8d1.png">


### Information

- [X] The official example scripts
- [ ] 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 am using the MIMIC-III dataset for summarization. Other than some light data preprocessing, everything is unchanged and I follow the exact same functions as in this tutorial https://huggingface.co/docs/transformers/tasks/summarization 
'''
# arguments below
seed=42
max_input_length = 512
max_target_length = 128
batch_size = 1
num_train_epochs = 5
gradient_accumulation_steps = 8
logging_steps = 91544 // batch_size #hardcoded to the length of mimic training set
learning_rate=3e-4
per_device_train_batch_size=batch_size
per_device_eval_batch_size=batch_size
weight_decay=0.01
save_total_limit=3
optim=Adafactor
warmup_steps = 1000

model_checkpoint = "google/t5-v1_1-large"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)
args = Seq2SeqTrainingArguments(
    output_dir="/home/ubuntu/summarization_models/t5_large/",
    evaluation_strategy="epoch",
    learning_rate=learning_rate,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=weight_decay,
    save_total_limit=save_total_limit,
    num_train_epochs=num_train_epochs,
    predict_with_generate=True,
    logging_steps=logging_steps,
    push_to_hub=False,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim='adafactor',
    max_grad_norm=1.0,
    ddp_find_unused_parameters=True, 
    warmup_steps=warmup_steps,
    fp16 = True, 
    sharded_ddp = 'zero_dp_2'
)
trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

trainer.train()

'''

### Expected behavior

Expected behavior: Normal training. No warning message. 

Comments:
Hi @arijitthegame ,

I have no idea what is fairscale and the stacktrace you provide is quite incomplete.
Do you mind trying to setup a small script that is reproduceable on a simple desktop/laptop so we can reproduce  (not t5-large) ?
Otherwise, reading the warning I think it shouldn't matter that much, at worse you're missing 1 step of the scheduler. It shouldn't change the end result for your model, only if you want reproduceable results with fixed seeds and so on should it matter.

Thank you 

Hi @Narsil, 

 Thank you so much for your quick reply. Fairscale is a libarary developed by META for DDP + sharding and is integrated in the HF Trainer. Let me attach a MWE. The same warning appears for T5-base and T5-small. 

```py
from datasets import load_dataset
import numpy as np

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer
from transformers import DataCollatorForSeq2Seq
from transformers import Adafactor
import evaluate
import nltk
from nltk.tokenize import sent_tokenize

seed=42
max_input_length = 1024
max_target_length = 128
batch_size = 4
num_train_epochs = 20
gradient_accumulation_steps = 1
logging_steps = 989 // batch_size #hardcoded to the length of training set
learning_rate=3e-4
per_device_train_batch_size=batch_size
per_device_eval_batch_size=batch_size
weight_decay=0.01
save_total_limit=3
optim=Adafactor

billsum = load_dataset("billsum", split="ca_test")
billsum = billsum.train_test_split(test_size=0.2)


rouge_score = evaluate.load("rouge")
model_checkpoint = "google/t5-v1_1-base"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

prefix = "summarize: " # not important for a single tasking model

def preprocess_function(examples):
    inputs = [prefix + doc for doc in examples["text"]]
    model_inputs = tokenizer(inputs, max_length=1024, truncation=True)

    with tokenizer.as_target_tokenizer():
        labels = tokenizer(examples["summary"], max_length=128, truncation=True)

    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    # Decode generated summaries into text
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    # Replace -100 in the labels as we can't decode them
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    # Decode reference summaries into text
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    # ROUGE expects a newline after each sentence
    decoded_preds = ["\n".join(sent_tokenize(pred.strip())) for pred in decoded_preds]
    decoded_labels = ["\n".join(sent_tokenize(label.strip())) for label in decoded_labels]
    # Compute ROUGE scores
    result = rouge_score.compute(
        predictions=decoded_preds, references=decoded_labels, use_stemmer=True
    )
    # Extract the scores
    result = {key: value * 100 for key, value in result.items()}
    return result

def postprocess_text(preds, labels):
    preds = [pred.strip() for pred in preds]
    labels = [label.strip() for label in labels]

    # ROUGE expects a newline after each sentence
    preds = ["\n".join(nltk.sent_tokenize(pred)) for pred in preds]
    labels = ["\n".join(nltk.sent_tokenize(label)) for label in labels]

    return preds, labels

tokenized_dataset = billsum.map(preprocess_function, batched=True)
tokenized_dataset = tokenized_dataset.remove_columns(
    billsum["train"].column_names
)

model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

args = Seq2SeqTrainingArguments(
    output_dir="/home/ubuntu/summarization_models/test/",
    evaluation_strategy="epoch",
    learning_rate=learning_rate,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=weight_decay,
    save_total_limit=save_total_limit,
    num_train_epochs=num_train_epochs,
    predict_with_generate=True,
    logging_steps=logging_steps,
    push_to_hub=False,
    gradient_accumulation_steps=gradient_accumulation_steps,
    optim='adafactor',
    max_grad_norm=1.0,
    ddp_find_unused_parameters=True, 
   # warmup_steps=warmup_steps,
    fp16 = True, 
    sharded_ddp = 'zero_dp_2'
)


trainer = Seq2SeqTrainer(
    model,
    args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
    data_collator=data_collator,
    tokenizer=tokenizer,
  #  compute_metrics=compute_metrics, # get different issue when uncommenting this line HALP!!
)

trainer.train()
```

Turning on compute metrics gives a different error. I can open a separate issue for that. Thank you once again. 

Hey @arijitthegame, 

Sorry could you also attach the error message of the script you've posted above? Also does the error happen when you provide the `compute_metrics` function or not?



Hi @patrickvonplaten, 

 Attached is the complete error message when turning on computing metrics. It crashes and does not train after a few steps.
 
 ```
Using cuda_amp half precision backend
***** Running training *****
  Num examples = 989
  Num Epochs = 20
  Instantaneous batch size per device = 4
  Total train batch size (w. parallel, distributed & accumulation) = 32
  Gradient Accumulation steps = 1
  Total optimization steps = 620
  5%|███▎                                                             | 31/620 [00:12<02:59,  3.27it/s]***** Running Evaluation *****
  Num examples = 248
  Batch size = 4
Traceback (most recent call last):
Traceback (most recent call last):
Traceback (most recent call last):
  File "test_trainer.py", line 117, in <module>
  File "test_trainer.py", line 117, in <module>
  File "test_trainer.py", line 117, in <module>
        trainer.train()trainer.train()

    Traceback (most recent call last):
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
trainer.train()  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
  File "test_trainer.py", line 117, in <module>

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
Traceback (most recent call last):
Traceback (most recent call last):
Traceback (most recent call last):
  File "test_trainer.py", line 117, in <module>
  File "test_trainer.py", line 117, in <module>
  File "test_trainer.py", line 117, in <module>
Traceback (most recent call last):
  File "test_trainer.py", line 117, in <module>
    trainer.train()
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
                trainer.train()trainer.train()trainer.train()trainer.train()



  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1500, in train
        return inner_training_loop(return inner_training_loop(

      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
return inner_training_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
    return inner_training_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
                return inner_training_loop(return inner_training_loop(return inner_training_loop(return inner_training_loop(



  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 1834, in _inner_training_loop
        self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate
      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate
self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate
        self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
    
self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate
    self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2040, in _maybe_log_save_evaluate
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate

      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate
metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate
    metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate
        metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)


  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 79, in evaluate
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)    
return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate

return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2760, in evaluate
    output = eval_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
    output = eval_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
    output = eval_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
    output = eval_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
    output = eval_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
    output = eval_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
    output = eval_loop(
      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
output = eval_loop(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer.py", line 2938, in evaluation_loop
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
    generated_tokens = self.model.generate(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
        loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)generated_tokens = self.model.generate(

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)    
loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
    return func(*args, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
    generated_tokens = self.model.generate(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    generated_tokens = self.model.generate(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
        loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)return func(*args, **kwargs)

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)    
loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
    loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/trainer_seq2seq.py", line 201, in prediction_step
        model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(generated_tokens = self.model.generate(

      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
    generated_tokens = self.model.generate(  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    generated_tokens = self.model.generate(
    model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
generated_tokens = self.model.generate(  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context


  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
    return func(*args, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
    return func(*args, **kwargs)
          File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
return func(*args, **kwargs)model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 1182, in generate
          File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
    model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
    model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
        return forward_call(*input, **kwargs)model_kwargs = self._prepare_encoder_decoder_kwargs_for_generation(

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/generation_utils.py", line 525, in _prepare_encoder_decoder_kwargs_for_generation
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward
    model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
            model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)return forward_call(*input, **kwargs)    return forward_call(*input, **kwargs)

model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward
      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
model_kwargs["encoder_outputs"]: ModelOutput = encoder(**encoder_kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    inputs_embeds = self.embed_tokens(input_ids)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    inputs_embeds = self.embed_tokens(input_ids)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward
    inputs_embeds = self.embed_tokens(input_ids)
      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    inputs_embeds = self.embed_tokens(input_ids)return forward_call(*input, **kwargs)    

return forward_call(*input, **kwargs)  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
      File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward

return forward_call(*input, **kwargs)  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/transformers/models/t5/modeling_t5.py", line 936, in forward
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
    return F.embedding(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
    inputs_embeds = self.embed_tokens(input_ids)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return F.embedding(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
            inputs_embeds = self.embed_tokens(input_ids)    return forward_call(*input, **kwargs)inputs_embeds = self.embed_tokens(input_ids)
inputs_embeds = self.embed_tokens(input_ids)


  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return F.embedding(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
    return F.embedding(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
    return F.embedding(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
    return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
        return forward_call(*input, **kwargs)
return forward_call(*input, **kwargs)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
    return F.embedding(RuntimeError    
: return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)    The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
    
return F.embedding(
return F.embedding(

  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/nn/functional.py", line 2183, in embedding
RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)    RuntimeError
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse): 
The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.RuntimeError
: The tensor has a non-zero number of elements, but its data is not allocated yet. Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() to actually allocate memory.
  5%|███▎                                                             | 31/620 [00:12<04:01,  2.44it/s]
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 42297) of binary: /home/ubuntu/nlp_prompting/env/bin/python
Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/distributed/launch.py", line 193, in <module>
    main()
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/distributed/launch.py", line 189, in main
    launch(args)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/distributed/launch.py", line 174, in launch
    run(args)
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/distributed/run.py", line 715, in run
    elastic_launch(
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/ubuntu/nlp_prompting/env/lib/python3.8/site-packages/torch/distributed/launcher/api.py", line 245, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
test_trainer.py FAILED
------------------------------------------------------------
```

Without the computing metrics, the model trains but of course I can not compute the rouge metrics with the code I posted above but I get a warning about scheduler. 
<img width="1026" alt="Screen Shot 2022-08-25 at 4 13 40 PM" src="https://user-images.githubusercontent.com/25961440/186977962-62c53340-0270-4214-aedf-489e5a6e2f5d.png">


Uff I don't really know what's going on here. The error seems to come from:

```
inputs_embeds = self.embed_tokens(input_ids)
```

But don't know `fairscale` well enough here. Maybe cc @sgugger 

Thank you so much. I believe it may be due to this. `This feature is incompatible with --predict_with_generate in the run_translation.py script. in ` https://huggingface.co/docs/transformers/main_classes/trainer#fairscale

This bug does not appear if I use `--fsdp full_shard` instead. 

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.