Mixed precision training + TransformerDecoder with a boolean mask

[alvaropp]

Hi there,

I'm have a custom model that uses a TransformerDecoder module and a boolean causal mask, as follows:

class MyModel(L.LightningModule):
    [...]

    def forward(self, data):
        [...]
        mask = torch.triu(
            torch.ones(size, size),
            diagonal=1,
        ).bool()
        transformed = self.transformer_decoder(
            tgt=data,
            memory=memory,
            tgt_mask=mask.bool(),
            memory_mask=mask.bool(),
        )

This works fine on both CPU and GPU, however, I'm trying to use mixed precision for training, as follows:

trainer = L.Trainer(
    accelerator="cpu",
    precision="16-mixed",
)

and I'm running into the following error: RuntimeError: Expected attn_mask dtype to be bool or to match query dtype, but got attn_mask.dtype: float and query.dtype: c10::BFloat16 instead..

As can be seen in the code above, I've tried casting the mask to bool everywhere, as I think that's compatible with bfloat16 (I may be mistaken).

I've also tried using:

from torch.nn.modules.transformer import _generate_square_subsequent_mask

mask = _generate_square_subsequent_mask(num_patches, device=self.device).bool()

but getting the same error.

Any ideas appreciated, thanks!

[alvaropp]

Ahh, it seems to work on GPU, so it could be that 16-mixed is not supported on CPU?