Model training works on GPU but not on Pytorch/XLA

[ahmedo42]

Issue description

I'm finetuning a BERT model , the model trains just fine on GPU but when I use TPU on Kaggle kernels or Colab it produces the following stack trace

Exception in device=TPU:5: torch_xla/csrc/aten_xla_type.cpp:163 : Check failed: at::canCast( resultType, out.scalar_type()) 
*** Begin stack trace ***
	tensorflow::CurrentStackTrace()
	
	torch_xla::AtenXlaType::mul_(at::Tensor&, at::Tensor const&)
	at::Tensor::mul_(at::Tensor const&) const
	at::native::binary_cross_entropy_with_logits(at::Tensor const&, at::Tensor const&, at::Tensor const&, at::Tensor const&, long)
	torch_xla::AtenXlaType::binary_cross_entropy_with_logits(at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, long)
	c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoRuntimeFunctor_<at::Tensor (*)(at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, long), at::Tensor, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, long> >, at::Tensor (at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, long)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, long)
	at::binary_cross_entropy_with_logits(at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, long)
	
	at::binary_cross_entropy_with_logits(at::Tensor const&, at::Tensor const&, c10::optional<at::Tensor> const&, c10::optional<at::Tensor> const&, long)

Code example

I believe the error happens when I'm calculating the loss using nn.BCEWithLogitsLoss() due to some type error?

this is how I'm doing it in training:

def train_step(engine, batch):

        input_ids = batch["input_ids"]
        attention_mask = batch["attention_mask"]
        token_type_ids = batch["token_type_ids"]
        labels = batch["label"].unsqueeze(-1)

        if input_ids.device != device:
            input_ids = input_ids.to(device, non_blocking=True,dtype=torch.long)
            attention_mask = attention_mask.to(device, non_blocking=True,dtype=torch.long)
            token_type_ids = token_type_ids.to(device, non_blocking=True,dtype=torch.long)
            labels = labels.to(device, non_blocking=True,dtype=torch.float)

        model.train()

        with autocast(enabled=with_amp):
            output = model(input_ids, attention_mask,token_type_ids).float()
            loss = criterion(output,labels)

[jysohn23]

Yeah it looks like there's some type mismatch where one isn't able to be promoted to the other. I see you're using autocast could you try disabling that if you haven't already?

[ahmedo42]

I set with_amp=False by default , even though if it is True it should work according to this https://pytorch.org/docs/stable/amp.html#prefer-binary-cross-entropy-with-logits-over-binary-cross-entropy .

[jysohn23]

Yeah just that for AMP we have our own flavor just literally introduced in #2654.

Would you try log the types the output and label tensors? Also do you have a small repro you could share with us?

[ahmedo42]

so I removed the whole amp decorator and ran the code on GPU , both are of type torch.float32 before being called by the loss function any idea why TPU might behave differently ?

[ahmedo42]

Hi! , The issue has been resolved , I formatted the dataset columns to the appropriate dtypes beforehand instead of converting them on the fly.