LongT5 Models Are Not Initialized With Pretrained Weights
Author: reelmath
Pull request: no
Labels: name: bug

### System Info

 `transformers` version: 4.20.1
- Platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.13
- Huggingface_hub version: 0.8.1
- PyTorch version (GPU?): 1.11.0+cu113 (True)
- Tensorflow version (GPU?): 2.8.2 (True)
- 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?: No

### Who can help?

@LysandreJik @stancld @patrickvonplaten 

### 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 have tried using LongT5 fine-tuning on a long range summarization task on a custom dataset (consider it like CNN/DM in that it is highly extractive). While long-t5-tglobal-base works well (I am able to converge on a validation loss of ~1.25 and ROUGE-2 of ~21), the long-t5-local-base, long-t5-local-large, and long-t5-tglobal-large all end up getting me training/validation losses of 200+ with ROUGE scores of exactly 0, making me believe that these models haven't actually been initialized with Google's weights. Here are the json outputs associated with trainer.evaluate() after 1 epoch of training:

**google/long-t5-local-base**
{'epoch': 1.0,
 'eval_gen_len': 1023.0,
 'eval_loss': 366.21673583984375,
 'eval_rouge1': 0.0,
 'eval_rouge2': 0.0,
 'eval_rougeL': 0.0,
 'eval_rougeLsum': 0.0,
 'eval_runtime': 37.9896,
 'eval_samples_per_second': 0.132,
 'eval_steps_per_second': 0.053}

**google/long-t5-tglobal-base (This one works correctly)**
{'epoch': 1.0,
 'eval_gen_len': 708.2,
 'eval_loss': 1.6017440557479858,
 'eval_rouge1': 35.7791,
 'eval_rouge2': 11.5732,
 'eval_rougeL': 19.1541,
 'eval_rougeLsum': 31.8491,
 'eval_runtime': 34.8695,
 'eval_samples_per_second': 0.143,
 'eval_steps_per_second': 0.057}

**google/long-t5-local-large**
{'epoch': 0.77,
 'eval_gen_len': 1023.0,
 'eval_loss': 252.44662475585938,
 'eval_rouge1': 0.0,
 'eval_rouge2': 0.0,
 'eval_rougeL': 0.0,
 'eval_rougeLsum': 0.0,
 'eval_runtime': 89.2506,
 'eval_samples_per_second': 0.056,
 'eval_steps_per_second': 0.034}

**google/long-t5-tglobal-large**
{'epoch': 0.77,
 'eval_gen_len': 1023.0,
 'eval_loss': 241.6276397705078,
 'eval_rouge1': 0.0,
 'eval_rouge2': 0.0,
 'eval_rougeL': 0.0,
 'eval_rougeLsum': 0.0,
 'eval_runtime': 89.9801,
 'eval_samples_per_second': 0.056,
 'eval_steps_per_second': 0.033}

For reproduction, just run the standard Huggingface PyTorch training script for summarization on any official dataset (CNN/DM, XSum, etc.).

Note that I haven't tried the 3B parameter versions so cannot speak to whether this problem affects them as well.


### Expected behavior

All four models should have a low validation loss when fine tuning on summarization (as opposed to three of them having 300+ validation losses as if they are randomly initialized).

Comments:
Hey @reelmath, hmm - we've verified that `google/long-t5-tglobal-large` can be correctly fine-tuned - see: https://huggingface.co/Stancld/longt5-tglobal-large-16384-pubmed-3k_steps where @stancld fine-tuned the model on the PubMed summarization dataset. 

Could you maybe copy-paste the command that you used to fine-tune those models here? 

Also cc @patil-suraj - can you double check quickly that the uploaded weights are the correct ones?

from transformers import AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer
import torch

model_checkpoint = google/long-t5-tglobal-base # substitute with alternate model checkpoints
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

from datasets import load_metric
metric = load_metric("rouge")

batch_size = 4
model_name = model_checkpoint.split("/")[-1]
args = Seq2SeqTrainingArguments(
f"{model_name}-summaries",
evaluation_strategy = "steps",
save_strategy = "steps",
logging_strategy = "epoch",
learning_rate=5e-4,
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
weight_decay=0.01,
gradient_accumulation_steps=4,
warmup_steps=3000,
num_train_epochs=1,
eval_steps=100,
save_steps=100,
load_best_model_at_end=True,
predict_with_generate=True,
gradient_checkpointing=True,
adafactor=True,
fp16=False
)

data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

import nltk
import numpy as np

def compute_metrics(eval_pred):
predictions, labels = eval_pred
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)
decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

# Rouge expects a newline after each sentence
decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]

result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
# Extract a few results
result = {key: value.mid.fmeasure * 100 for key, value in result.items()}

# Add mean generated length
prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
result["gen_len"] = np.mean(prediction_lens)
result["eval_rougeLsum"] = result["rougeLsum"]

return {k: round(v, 4) for k, v in result.items()}
trainer = Seq2SeqTrainer(
model,
args,
train_dataset=training_dataset,
eval_dataset=validation_dataset,
data_collator=data_collator,
tokenizer=tokenizer,
compute_metrics=compute_metrics
)

import nltk
nltk.download('punkt')

trainer.train()

trainer.evaluate()

@patrickvonplaten

Hello @reelmath, don't you have please available any loss logs? Not sure if `lr=5e-4` with a batch size of 16 is not too high, so there might be a danger of the model diverging once the training starts?

@stancld Here are the results I got from fine tuning LongT5 TGlobal Large. I dropped the learning rate to 1e-5 and ran bsz=2 with 8 accumulation steps (so effective bsz of 16)

Step	Training Loss	Validation Loss	Rougelsum	Rouge1	Rouge2	Rougel	Gen Len
20	246.385800	241.613937	0.000000	0.000000	0.000000	0.000000	2047.000000
40	244.759000	241.159866	0.000000	0.000000	0.000000	0.000000	2047.000000
60	245.665300	240.319992	0.000000	0.000000	0.000000	0.000000	2047.000000
80	246.063400	239.445801	0.000000	0.000000	0.000000	0.000000	2047.000000

Logs attached below:
[runs-20220708T155211Z-001.zip](https://github.com/huggingface/transformers/files/9073383/runs-20220708T155211Z-001.zip)


Update: it appears that changing:

model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)

to

model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint, from_flax=True)

helps resolve the issue. This is true for longt5-tglobal-large, at least.

Update 2: Loading from flax works for longt5-tglobal-large and longt5-local-base, but does not work for longt5-local-large (which starts and flatlines at a training and validation loss of around 10).

@reelmath Thanks for all the information. Could you please also run the first 20 training steps and log loss every step?
It's really strange, as mentioned @patrickvonplaten, there are already a few successfully fine-tuned checkpoints of TGlobal-Large model which was initialized directly from the PyTorch checkpoint.

@patil-suraj Is there by chance any possibility that checkpoints became corrupted during the transfer from my repo to Google's?

LongT5 Local Large initialized from Flax into pt
[events.out.tfevents.1657582716.9bccc2c81a7b.71.zip](https://github.com/huggingface/transformers/files/9088386/events.out.tfevents.1657582716.9bccc2c81a7b.71.zip)

LongT5 Local Large initialized directly with pt
[events.out.tfevents.1657582716.9bccc2c81a7b.71.zip](https://github.com/huggingface/transformers/files/9088387/events.out.tfevents.1657582716.9bccc2c81a7b.71.zip)

LongT5 TGlobal Large initialized from Flax into pt
[events.out.tfevents.1657586027.abf516f8c8c8.73.zip](https://github.com/huggingface/transformers/files/9088616/events.out.tfevents.1657586027.abf516f8c8c8.73.zip)

LongT5 TGlobal Large initialized directly with pt
[events.out.tfevents.1657587671.abf516f8c8c8.73.zip](https://github.com/huggingface/transformers/files/9088622/events.out.tfevents.1657587671.abf516f8c8c8.73.zip)

My guess is that Local Base is going to look quite similar to TGlobal Large in patterns, though with higher training/validation losses (at least based on previous experiments).

@reelmath Thanks, try to will have a look at those at the weekend :]

@stancld Any progress? I've been hoping to try out Local Large sometime.

Hi @reelmath, sorry for my late reply, I've been really busy lately. When I load weights from the pre-trained checkpoints everything looks good to me, however, I have currently no access to computational resources to verify the aforementioned problem. Though, I think there's one way how to verify whether the weights have been broken during the transition to Google repo.

You can try to convert the original Google checkpoint to HF checkpoint using the script available in this repo: https://github.com/huggingface/transformers/blob/main/src/transformers/models/longt5/convert_longt5x_checkpoint_to_flax.py. The links to all the publicly available checkpoints can be found in this repo https://github.com/google-research/longt5. Then, it's need to download checkpoints from GDrive and subsequently convert these checkpoints with the script attached.

Please, let me know if this works for you :] Also, definitely stay in touch if you have any issue with converting those checkpoints.

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.