Skip to content

trunk/aab577b236447c87bb4970bb78c9c01b99b7ba62

@kurtamohler kurtamohler tagged this 29 Jun 05:22
Closes #160828

The backward pass is broken up into two kernel calls, one to calculate `log_beta` and one to perform logsumexp-reduction of `log_beta + log_alpha` and write the input grad.

Avoids some code duplication by refactoring the `log_alpha` code to be able to conditionally calculate `log_beta`.

The reduction step is parallelized by creating one thread per `T * N`, where `log_probs.shape = (T, N, C)`, `T` is the input length and `N` is the batch size. Each thread iterates over the target sequence, logsumexp-reducing the corresponding elements of `log_alpha` and `log_beta` into whichever element of the input gradient matches the target token. Then the thread iterates over its assigned row of the output grad and applies the formula $\nabla y = (y - \exp(R + L - \log y)) \nabla L$, where $\log y$ is `log_probs`, $R$ is the result of the reduction, $L$ is the loss calculated by the forward call.

[Notes on how ctc_loss works](https://github.com/kurtamohler/notes/blob/a56ce01e65ac591c6d851d2fa0d7438cbd523aa9/pytorch/ctc_loss.ipynb)

Co-authored with Claude
Pull Request resolved: https://github.com/pytorch/pytorch/pull/188187
Approved by: https://github.com/jhavukainen, https://github.com/malfet
Assets 2
Loading