Evaluation Padding_idx Error when using BERTScore and Deepspeed Zero3 but not with Zero2
Author: KMFODA
Pull request: no

## Environment info
- `transformers` version: 4.16.0.dev0
- Platform: Linux-5.11.0-37-generic-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.6.9
- PyTorch version (GPU?): 1.8.0 (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
- Using distributed or parallel set-up in script?: Yes

### Who can help
- Deepspeed: @stas00

## Information

Model I am using (google/pegasus-large):

The problem arises when using:
* [x] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [x] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Replace compute_metrics in transformers/examples/pytorch/summarization/run_summarization.py with the following:

```
from datasets import load_dataset, load_metric  
# Metric
metric = load_metric("rouge")

# load rouge
rouge = load_metric("rouge")

# load bert_score
bert_score_metric = load_metric("bertscore")

def compute_metrics(eval_preds):
  preds, labels = eval_preds
  if isinstance(preds, tuple):
      preds = preds[0]
  decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
  if data_args.ignore_pad_token_for_loss:
      # Replace -100 in the labels as we can't decode them.
      labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
  decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

  # Some simple post-processing
  decoded_preds, decoded_labels = postprocess_text(decoded_preds, decoded_labels)

  rouge1_output = rouge.compute(predictions=decoded_preds, references=decoded_labels, rouge_types=["rouge1"])["rouge1"].mid
  rouge2_output = rouge.compute(predictions=decoded_preds, references=decoded_labels, rouge_types=["rouge2"])["rouge2"].mid
  rougel_output = rouge.compute(predictions=decoded_preds, references=decoded_labels, rouge_types=["rougeL"])["rougeL"].mid
  
  result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
  # Extract a few results from ROUGE
  result = {key: value.mid.fmeasure * 100 for key, value in result.items()}

  prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
  result["gen_len"] = np.mean(prediction_lens)
  result = {k: round(v, 4) for k, v in result.items()}
  
  bert_results = bert_score_metric.compute(predictions=decoded_preds, references=decoded_labels, lang="en", )
  bert_f1 = sum(bert_results["f1"]) / len(bert_results["f1"])
  bert_precision = sum(bert_results["precision"]) / len(bert_results["precision"])
  bert_recall = sum(bert_results["recall"]) / len(bert_results["recall"])
  
  bert_rogue = ((bert_f1*100) + result["rougeL"]) / 2
  
  return {
      "result":result,
      "rogue1": {"rouge1_precision":round(rouge1_output.precision, 4)*100, "rouge1_recall":round(rouge1_output.recall, 4)*100, "rouge1_f1":round(rouge1_output.fmeasure, 4)*100},
      "rogue2": {"rouge2_precision":round(rouge2_output.precision, 4)*100, "rouge2_recall":round(rouge2_output.recall, 4)*100, "rouge2_f1":round(rouge2_output.fmeasure, 4)*100},
      "roguel": {"rougel_precision":round(rougel_output.precision, 4)*100, "rougel_recall":round(rougel_output.recall, 4)*100, "rougel_f1":round(rougel_output.fmeasure, 4)*100},
      "bertscore": {"bert_precision":round(bert_precision, 4)*100, "bert_recall":round(bert_recall, 4)*100, "bert_f1":round(bert_f1, 4)*100},
      "bert_rogue": round(bert_rogue, 4)
  }
```

2. Run the following command:

```
deepspeed transformers/examples/pytorch/summarization/run_summarization.py \
    --deepspeed transformers/tests/deepspeed/ds_config_zero3.json \
    --model_name_or_path google/pegasus-large \
    --per_device_train_batch_size 2 \
    --output_dir output_dir \
    --overwrite_output_dir \
    --fp16 \
    --do_train \
    --predict_with_generate \
    --report_to wandb \
    --load_best_model_at_end True \
    --greater_is_better True \
    --evaluation_strategy steps \
    --save_steps 1200 \
    --eval_steps 50 \
    --logging_steps 400 \
    --max_train_samples 100 \
    --max_eval_samples 10 \
    --dataset_name samsum
```

3. You'll get the following error message:

```
Traceback (most recent call last):
  File "run_summarization.py", line 769, in <module>
    main()
  File "run_summarization.py", line 713, in main
    metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer_seq2seq.py", line 70, in evaluate
Traceback (most recent call last):
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer.py", line 2164, in evaluate
  File "run_summarization.py", line 769, in <module>
    main()
  File "run_summarization.py", line 713, in main
    metric_key_prefix=metric_key_prefix,
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer.py", line 2400, in evaluation_loop
    metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer_seq2seq.py", line 70, in evaluate
    return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix)
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer.py", line 2164, in evaluate
    metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
  File "run_summarization.py", line 605, in compute_metrics
    metric_key_prefix=metric_key_prefix,
  File "/usr/local/lib/python3.6/dist-packages/transformers/trainer.py", line 2400, in evaluation_loop
    bert_results = bert_score_metric.compute(predictions=decoded_preds, references=decoded_labels, lang="en", )
  File "/usr/local/lib/python3.6/dist-packages/datasets/metric.py", line 405, in compute
    output = self._compute(predictions=predictions, references=references, **kwargs)
  File "/workspace/.cache/huggingface/modules/datasets_modules/metrics/bertscore/5075ace4c0abac8c79f820cc5e8dabc1d8263bf873cf93296a8f34d2aad42eb6/bertscore.py", line 184, in _compute
    batch_size=batch_size,
  File "/usr/local/lib/python3.6/dist-packages/bert_score/scorer.py", line 222, in score
    all_layers=self.all_layers,
  File "/usr/local/lib/python3.6/dist-packages/bert_score/utils.py", line 529, in bert_cos_score_idf
    metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
  File "run_summarization.py", line 605, in compute_metrics
    sen_batch, model, tokenizer, idf_dict, device=device, all_layers=all_layers
  File "/usr/local/lib/python3.6/dist-packages/bert_score/utils.py", line 408, in get_bert_embedding
    model, padded_sens[i : i + batch_size], attention_mask=mask[i : i + batch_size], all_layers=all_layers,
    bert_results = bert_score_metric.compute(predictions=decoded_preds, references=decoded_labels, lang="en", )
  File "/usr/local/lib/python3.6/dist-packages/bert_score/utils.py", line 318, in bert_encode
  File "/usr/local/lib/python3.6/dist-packages/datasets/metric.py", line 405, in compute
    out = model(x, attention_mask=attention_mask, output_hidden_states=all_layers)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    output = self._compute(predictions=predictions, references=references, **kwargs)
  File "/workspace/.cache/huggingface/modules/datasets_modules/metrics/bertscore/5075ace4c0abac8c79f820cc5e8dabc1d8263bf873cf93296a8f34d2aad42eb6/bertscore.py", line 184, in _compute
    batch_size=batch_size,
  File "/usr/local/lib/python3.6/dist-packages/bert_score/scorer.py", line 222, in score
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/models/roberta/modeling_roberta.py", line 848, in forward
    all_layers=self.all_layers,
  File "/usr/local/lib/python3.6/dist-packages/bert_score/utils.py", line 529, in bert_cos_score_idf
    sen_batch, model, tokenizer, idf_dict, device=device, all_layers=all_layers
  File "/usr/local/lib/python3.6/dist-packages/bert_score/utils.py", line 408, in get_bert_embedding
    past_key_values_length=past_key_values_length,
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    model, padded_sens[i : i + batch_size], attention_mask=mask[i : i + batch_size], all_layers=all_layers,
  File "/usr/local/lib/python3.6/dist-packages/bert_score/utils.py", line 318, in bert_encode
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/models/roberta/modeling_roberta.py", line 131, in forward
    out = model(x, attention_mask=attention_mask, output_hidden_states=all_layers)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    inputs_embeds = self.word_embeddings(input_ids)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/models/roberta/modeling_roberta.py", line 848, in forward
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/sparse.py", line 147, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 1897, in embedding
    past_key_values_length=past_key_values_length,
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/transformers/models/roberta/modeling_roberta.py", line 131, in forward
    assert padding_idx < weight.size(0), "Padding_idx must be within num_embeddings"
    inputs_embeds = self.word_embeddings(input_ids)
AssertionError: Padding_idx must be within num_embeddings
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/sparse.py", line 147, in forward
    self.norm_type, self.scale_grad_by_freq, self.sparse)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py", line 1897, in embedding
    assert padding_idx < weight.size(0), "Padding_idx must be within num_embeddings"
AssertionError: Padding_idx must be within num_embeddings
```

## Expected behavior

If this command is run instead training and evaluation completes without an AssertationError:

```
deepspeed transformers/examples/pytorch/summarization/run_summarization.py \
    --deepspeed transformers/tests/deepspeed/ds_config_zero2.json \
    --model_name_or_path google/pegasus-large \
    --per_device_train_batch_size 2 \
    --output_dir output_dir \
    --overwrite_output_dir \
    --fp16 \
    --do_train \
    --predict_with_generate \
    --report_to wandb \
    --load_best_model_at_end True \
    --greater_is_better True \
    --evaluation_strategy steps \
    --save_steps 1200 \
    --eval_steps 50 \
    --logging_steps 400 \
    --max_train_samples 100 \
    --max_eval_samples 10 \
    --dataset_name samsum
```

I'm not sure what Zero3 is doing vs Zero2 that could cause this padding error.

Comments:
Hi @KMFODA,

zero3 is different from zero2 in that it shards the params. zero3 magically gathers them on each gpu before the model's forward calls, so the model has no idea its weights were manipulated. But if you do other forward's that aren't part of the model and which are called before the model, deepspeed doesn't know about it and most likely the error comes from that.

Typically you then need to do your own gathering before accessing the params that aren't part of the model's forward, for example:

https://github.com/huggingface/transformers/blob/1fc0fa46174511a68889eacef72693c93dc00373/src/transformers/modeling_utils.py#L783-L787

Please let me know if that helps with your process. And of course if that helps please do share the outcome.

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/master/CONTRIBUTING.md) are likely to be ignored.

Thanks @stas00 for the response. Sorry for my delayed one. I ended up just stick to a ROUGE score which didn't need these changes in the interest of time. Will try this out next time I'm using zero3 and BertScore and report back her. Thank you!