Fix XGLMModelLanguageGenerationTest.test_batched_nan_fp16

[ydshieh]

What does this PR do?

#18057 added this test to test running with fp16.

However, from_pretrained(model_name, torch_dtype=torch.float16 seems not able to change the dtype for weights registered below:

transformers/src/transformers/models/xglm/modeling_xglm.py

Lines 168 to 176 in a7bc422

	self.make_weights(num_positions + self.offset, embedding_dim, padding_idx)
	def make_weights(self, num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None):
	emb_weights = self.get_embedding(num_embeddings, embedding_dim, padding_idx)
	if hasattr(self, "weights"):
	# in forward put the weights on the correct dtype and device of the param
	emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)
	self.register_buffer("weights", emb_weights)

and hidden_states becomes again float32 (because position is) at

transformers/src/transformers/models/xglm/modeling_xglm.py

Line 715 in a7bc422

hidden_states = inputs_embeds + positions

and finally failed at hidden_states = self.self_attn_layer_norm(hidden_states) with

RuntimeError: expected scalar type Float but found Half

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[ydshieh]

position (and embed_positions's weight) are in float32 even if we load the model in float16. Need this cast to make later LayerNorm layer work.

[sgugger]

I think the right fix would be to make sure the weights have the correct dtype, as the embedding layer is the biggest one, so the memory save is very important.

[sgugger]

Thanks for looking into this. While this fixes the issue, I'm not sure if it's the right fix.

[sgugger]

I think the right fix would be to make sure the weights have the correct dtype, as the embedding layer is the biggest one, so the memory save is very important.

[ydshieh]

@sgugger OK, let me check if I can do something for (not a real) weights defined by

self.register_buffer("weights", emb_weights)

[ydshieh]

The involved test in this PR uses

from_pretrained(model_name, torch_dtype=torch.float16, ...)

but at init. time, it uses float to create some tensors that is then registered to the buffer (and keep float32)

transformers/src/transformers/models/xglm/modeling_xglm.py

Line 189 in 9ed80b0

emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)

[sgugger]

Thanks, this feels like a much better fix!