Paligemma- fix devices and dtype assignments #31008
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
cc @ArthurZucker wdyt? |
| causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( | ||
| token_type_ids[:, None, None, :] == 0, 0 | ||
| token_type_ids[:, None, None, :].to(causal_mask.device) == 0, 0 | ||
| ) |
There was a problem hiding this comment.
this one does not make sense to me
There was a problem hiding this comment.
with the masked_fill, you need both tensors to be on the same device, right?
There was a problem hiding this comment.
Ah, sorry, I read to fast.
So token type ids's device is not correctly inferred ? Or are we not creating the causal mask on the correct device? It should be created on the input or attention mask's device for consistency, since when it's used accelerate will transfer it accordingly I think.
There was a problem hiding this comment.
you're right! will update to move to device at creation time
There was a problem hiding this comment.
hmm, I think we are setting the causal_mask to the correct device. It's the token_type_id device that is indeed not correctly inferred
There was a problem hiding this comment.
But from the comment of @SunMarc I would suspect both devices to be the same no
There was a problem hiding this comment.
hmm, doesn't seem like it, tried on a multi-gpu env with device_map to auto and got
There was a problem hiding this comment.
@SunMarc if you have an idea here - token_type_ids is created by the processor along with input_ids and passed to the forward normally
There was a problem hiding this comment.
So, from the code and the image you shared, I see that token_type_ids is indeed on the same device as input_ids. However, since you created the causal_mask to be on the same device as inputs_embeds.device, token_type_ids and input_ids might not be on the same device.
causal_mask = torch.full(
(sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
)where
dtype, device = inputs_embeds.dtype, inputs_embeds.device
There was a problem hiding this comment.
alright, thanks! in that case, we can keep it as it is? The other way is to create the causal mask on the input_ids.device, I'm not sure if one is better than the other - inputs_embeds is much larger in general
|
fwiw most of the lines in here are nearly identical to the changes i have done locally as well besides the |
|
@grahamannett , good to know. For |
* fix devices and dtype assignments * [run-slow]paligemma
What does this PR do?
Moves tensors to correct devices in case of multi-gpu training on accelerate and device_map = auto.
Additionally ensures bf16 training works as well.
Fixes #30997