FP-16 training producing nans on t5-large/flan-t5-xl

[cassianlewis]

System Info

This was an issue a while back but seems to have resurfaced - https://discuss.huggingface.co/t/t5-fp16-issue-is-fixed/3139

I have tested the exact following code on t5-small and t5-base and they work fine. However, when using t5-large and/or flan-t5-xl, the model produces nan outputs. This is solely a result of using half precision (ignore the multiple GPUs, strategy etc, I have tested with every other variation):

trainer = pl.Trainer(
    precision="16",
    accelerator='gpu',
    strategy='auto',
    devices=4,)

I am using transformers == 4.28.1 and lightning == 2.0.0

Any ideas/help appreciated
Thanks!

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

trainer = pl.Trainer(
    precision="16",
    accelerator='gpu',
    strategy='auto',
    devices=4,)

Expected behavior

Nans!!!

[github-actions]

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 are likely to be ignored.

[amyeroberts]

Hi @cassianlewis, thanks for reporting this issue.

There's been a few updates recently regarding handling of loading of model weights and their precision. Could you update your transformers version to the latest release and let us know if the issue still persists?

Could you also provide some more details about the other strategies tested (when this behaviour doesn't occur)?

cc @younesbelkada

[cassianlewis]

Hi @amyeroberts, I think the issue is more with the model itself. See https://www.philschmid.de/fine-tune-flan-t5-deepspeed

When fine-tuning T5 models we cannot use fp16 since it leads to overflow issues, see: [#4586](https://github.com/huggingface/transformers/issues/4586), [#10830](https://github.com/huggingface/transformers/issues/10830), [#10956](https://github.com/huggingface/transformers/pull/10956)

(t5 was trained with bf16)

FYI this behaviour doesn't occur when I remove the precision="16" line

[amyeroberts]

@cassianlewis Thanks for providing these details and links.

Looking at the linked issues and Philip's blog, unfortunately I don't think there's much that can be done to resolve this. We can leave this issue open for now for people to share any updates or solutions they may have.

[younesbelkada]

Hi @cassianlewis
I think it is expected that pure fp16 training leads to unstable behaviour for some models and in some configuration; have you tried to train the models in bf16 precision?

[cassianlewis]

@younesbelkada My Tesla T4 doesn't currently support bf16 (you need Ampere architecture), so I will get back to you on that

[sjrl]

I was running into the same issue, and I found a few things that helped (but did make training time much slower).

1. Specifically, I found that I needed to add the line @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) to the forward method of this module

   transformers/src/transformers/models/t5/modeling_t5.py

   Lines 301 to 314 in 2ab75ad

   	class T5DenseGatedActDense(nn.Module):
   	def __init__(self, config: T5Config):
   	super().__init__()
   	self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False)
   	self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False)
   	self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)
   	self.dropout = nn.Dropout(config.dropout_rate)
   	self.act = ACT2FN[config.dense_act_fn]
   	def forward(self, hidden_states):
   	hidden_gelu = self.act(self.wi_0(hidden_states))
   	hidden_linear = self.wi_1(hidden_states)
   	hidden_states = hidden_gelu * hidden_linear
   	hidden_states = self.dropout(hidden_states)

Of course, this greatly slows things down since fp16 won't be utilized at all in this module anymore (and this would slow down other methods like bf16 training). I think if there is a way to enforce that hidden_linear should always stay in fp32 then the torch.cuda.amp... statement would not be needed, but I was unable to figure out how to get that to work.

2. Another thing I changed were these blocks in the code

   transformers/src/transformers/models/t5/modeling_t5.py

   Lines 707 to 713 in 2ab75ad

   	if hidden_states.dtype == torch.float16:
   	clamp_value = torch.where(
   	torch.isinf(hidden_states).any(),
   	torch.finfo(hidden_states.dtype).max - 1000,
   	torch.finfo(hidden_states.dtype).max,
   	)
   	hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)

   
   to use the alternate if statement

if torch.isinf(hidden_states).any():
    ...

since I found it was not guaranteed that hidden_states would be in torch.float16 during fp16 training, but infinity values could still be present.

[github-actions]

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 are likely to be ignored.