I'm skeptical about this without some repro. My intuition is that if it produces invalid NaNs occasonally, it will likely produce wrong non-NaN results. Could you please capture the inputs to CTC loss that produces NaNs (i.e. keep them around and when you get NaN gradients save them)?
If you use 2d targets: is this something that might be helped by #20971 ?

Oh, and thanks for investigating this. It's awesome to see that you tracked this down!

@t-vi yes, creating a consistent repo is the most important thing here. I was able to repo NaN occasionally by injecting gradient check code during training and pdb once there is any NaN in gradient . But unfortunately, supplying the exact same input to CTCLoss again did not produce NaN. So I am not able to give you a consistent repo yet. I am thinking other approach to get a repo.

Btw, we use 1-d target so #20971 won’t help. Also we notice that once we disable large batch handling, NaN also disappears.

Since a lot of people’s work pending a numerical stable CTC loss, may I suggest we can merge this PR as a stopgap solution, similar as the purpose of zero_infinity ?

OK, so it's something different. I would indeed want to get us to the situation of zero_infinity - where we clearly understand why we're getting NaNs and so can target avoiding them.
I would expect that there would be something about the inputs when the grads turn NaN, but I'd really like to know what it is.

The other thing I'm wondering about: Do you get the NaN in the forward or the backward? Because you also delete them in the forward, even though you describe observing them in the gradient only.

@t-vi

The other thing I'm wondering about: Do you get the NaN in the forward or the backward? Because you also delete them in the forward, even though you describe observing them in the gradient only.

sorry that I miss this question. I observe NaN appears in gradient first( no NaN in the input) and loss is a normal number in that case. After gradient has NaN, it propagates to model parameters, and then input (to CTCLoss) in the next mini-batch and then everywhere.

@t-vi @yqwangustc and I discussed this offline and decided that we should land this patch to stem the bleeding (but we should continue to investigate what the real problem is.)

I have one more request, which is that zero_infinities documentation be updated to say that it also zeros NaNs, as a quirk of the current implementation (link to the bug or PR). (I know eventually we will probably want to change it back, but if we forget and this ends up in the release, I'd like the docs to be accurate.)

Related: #14335 (no code was provided in the reports but it looks similar)

A bit more information on the input causing NaN.

• lprobs.shape
  torch.Size([984, 16, 42])
• targets_flat.shape
  torch.Size([1316])
• input_length
  tensor([984, 956, 943, 924, 921, 904, 901, 897, 864, 824, 717, 710, 695, 666, 578, 552], device='cuda:0')
• target_lengths
  tensor([ 96, 78, 85, 140, 107, 88, 126, 95, 70, 52, 86, 38, 88, 69, 60, 38], device='cuda:0')
• blank_id = 41

and the loss is a norm floating point, so forward pass is fine. grad[983, 1:10, 41] has all the NaN value, i.e., all the NaN appears at BLANK posterior but they are just padded position in time sequence. Switching to small batch size condition does not yield any NaN, so I am checking grad before/after this line https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/LossCTC.cu#L560

@t-vi will these information help you to identify something ?

update:

by inserting the following code, I confirm that https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/LossCTC.cu#L560 is the offending line.

auto grad_blank = grad.narrow(2, BLANK, 1);
    if (at::isnan(grad_blank).any().item().to<uint8_t>() == 1) {
      printf("Detecting NAN in grad_blank before\n");
    }

    grad_blank -= (at::logsumexp(log_alpha.as_strided({batch_size, log_alpha.size(1), max_target_length+1},
                                                      {log_alpha.stride(0), log_alpha.stride(1), log_alpha.stride(2)*2})
                                 + log_beta.as_strided({batch_size, log_beta.size(1), max_target_length+1},
                                                       {log_beta.stride(0), log_beta.stride(1), log_beta.stride(2)*2}),
                                 2, true)
                   .permute({1, 0, 2})
                   .add_(neg_log_likelihood.view({1, batch_size, 1}))
                   .sub_(log_probs.narrow(2, BLANK, 1))
                   .exp_()
                   );

     if (at::isnan(grad_blank).any().item().to<uint8_t>() == 1) {
       printf("Detecting NAN in grad_blank after\n");
     }

@t-vi if I understand it correctly, using as_strided needs extreme care. Would you mind looking deep into this line to see whether there is any chance of error ? In particular, I suspect the stride*2 may cause issue at some edge cases ? Maybe we can rewrite it using gather ?

update again

Additional data examination reveal that log_beta has nan in it. Looking into why it will have nan.

Thank you for the analysis! This is extremely useful.
So it occurs in the padded part. In this log_alpha/log_beta will be -infinity. I don't think the as_strided is particularly problematic, but we should really replace it with .slice.
Now if log_probs is infinite or NaN (which should not have an effect in the padded part), it'll give NaN (either from NaN or from inf - inf).

Thanks again for tracking this down!

So I'm happy with your fix, maybe we can add a comment that this is in the padding and that when the blank calculation is moved to a kernel, properly ignoring those bits will be better.
I'd have a preference to not change the forward (the two other bits), but only zero out the NaNs in the backward.

@t-vi

Actually, after many iterations, I think I found the real root cause now !

I start to notice after log_beta = at::empty_like(log_alpha); there is already some nan in log_beta. Initially I thought this might be fine, since all the (b, t, s) cases will be covered in the GPU kernel ctc_loss_backward_log_beta_gpu_kernel. However, after carefully reviewing the code, I found two exceptions:

• if s=2*target_length + 1, since the condition in https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/LossCTC.cu#L342 is s > 2*target_length+1 not s>=2*target_length+1 (really tricky here !), any element in log_beta[b, :, 2*target_lengths[b] + 1] is possibly not initialized.

• For t=max_input_length - 1, I didn't see anywhere in the code to assign value, so log_beta[b, max_input_length - 1, :] is also possibly not initialized.

These are consistent with the pattern of either NaN or some widely large value (e.g., -1.03845089e+34) I saw in log_beta. This also perfectly explain why I cannot reliably reproduce this issue before.

So the suggest change:

• at https://github.com/pytorch/pytorch/blame/2ee2d78a29a583e3f40cd27828f02f6ef627a01c/aten/src/ATen/native/cuda/LossCTC.cu#L520, I am not sure why we use empty_like, changing it to full_like(log_alpha, neginf) solve all the problems!

@t-vi Please kindly correct me if I am wrong. I am running more experiments, but I am quite confident this will permanent fix the problem, :-) If this looks good to you, I can abandon this PR, and create a new one (which only has 1 line change !)

He, shaking that function really has the bugs falling out... Awesome work!
I'd probably stick with empty (saves calling a kernel to initialize the memory) if we can make sure we write all of them.

• So in your first observation > should be >= clearly.
• The t = max_input_length - 1 case should be handled by extending the s-loop. I was thinking that t = input_length - 1 was already covered and did not consider that we would end up using the uninitialized value.
  Then t >= input_length would cover that (unless input_length == max_input_length).

But I think we might still have to zero "out of sequence" gradients due to the use of log_probs.

Do you want help to update the PR? I'd be glad to help code something up, but you did all the hard work already...

@t-vi new PR is on the way and will send it out soon. sorry that just wake up, :-)

Supercool! Thanks!

Abandon this PR in favor of pytorch/pytorch #21392