diff --git a/.github/kernel-maintainers.json b/.github/kernel-maintainers.json index 16eacacd4..3215df6f9 100644 --- a/.github/kernel-maintainers.json +++ b/.github/kernel-maintainers.json @@ -51,6 +51,7 @@ "scattermoe": ["U08Q42K8MGW"], "sgl-flash-attn3": [], "sonic-moe": ["U08Q42K8MGW", "U054KJK89C2"], + "tdt-loss": ["U07U53UJECT"], "tinygrad-rms": [], "topk": ["U0586KAH2E9"], "trimul-gpumode": [], diff --git a/scripts/check_kernel_freshness.py b/scripts/check_kernel_freshness.py index 47b59c5cb..20302a861 100644 --- a/scripts/check_kernel_freshness.py +++ b/scripts/check_kernel_freshness.py @@ -48,6 +48,7 @@ "scattermoe": "https://github.com/shawntan/scattermoe", "sgl-flash-attn3": "https://github.com/sgl-project/sgl-flash-attn", "sonic-moe": "https://github.com/Dao-AILab/sonic-moe", + "tdt-loss": "", "tinygrad-rms": "https://github.com/tinygrad/tinygrad", "trimul-gpumode": "https://github.com/davidberard98/gpumode-trimul", "triton-kernels": "https://github.com/triton-lang/triton.git", diff --git a/tdt-loss/CARD.md b/tdt-loss/CARD.md new file mode 100644 index 000000000..1b3832451 --- /dev/null +++ b/tdt-loss/CARD.md @@ -0,0 +1,42 @@ +--- +library_name: kernels +{% if license %}license: {{ license }} +{% endif %}--- + +This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated. + +## How to use +{% if functions %} + +```python +# make sure `kernels` is installed: `pip install -U kernels` +from kernels import get_kernel + +kernel_module = get_kernel("{{ repo_id }}") +{{ functions[0] }} = kernel_module.{{ functions[0] }} + +{{ functions[0] }}(...) +``` +{% else %} + +Usage example not available. +{% endif %} + +## Available functions +{% if functions %} +{% for func in functions %} +- `{{ func }}` +{% endfor %} +{% else %} + +Function list not available. +{% endif %} + +## Benchmarks +{% if has_benchmark %} + +Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }}`. +{% else %} + +Benchmarking script is not available. +{% endif %} diff --git a/tdt-loss/benchmarks/benchmark.py b/tdt-loss/benchmarks/benchmark.py new file mode 100644 index 000000000..42fffaa34 --- /dev/null +++ b/tdt-loss/benchmarks/benchmark.py @@ -0,0 +1,81 @@ +import sys +from pathlib import Path + +import torch + +from kernels.benchmark import Benchmark + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from tests.reference import tdt_loss_reference # noqa: E402 + +DURATIONS = [0, 1, 2, 3, 4] + + +class TDTLossBenchmark(Benchmark): + seed: int = 42 + + def _make_inputs(self, B, T, U, V, dtype): + D = len(DURATIONS) + # Slices of a joint output, as produced by ParakeetForTDT. + self.logits = torch.randn(B, T, U + 1, V + D, device=self.device, dtype=dtype, requires_grad=True) + self.targets = torch.randint(1, V, (B, U), device=self.device, dtype=torch.int32) + self.logit_lengths = torch.full((B,), T, device=self.device, dtype=torch.int32) + self.target_lengths = torch.full((B,), U, device=self.device, dtype=torch.int32) + self.V = V + + def _run(self): + self.logits.grad = None + losses = self.kernel.tdt_loss( + self.logits[..., : self.V], + self.logits[..., self.V :], + self.targets, + self.logit_lengths, + self.target_lengths, + 0, + DURATIONS, + reduction="none", + ) + losses.sum().backward() + self.out = losses.detach() + + def _reference(self): + with torch.no_grad(): + return tdt_loss_reference( + self.logits[..., : self.V], + self.logits[..., self.V :], + self.targets, + self.logit_lengths, + self.target_lengths, + 0, + DURATIONS, + ) + + # Small vocabulary, short utterances. + def setup(self): + self._make_inputs(B=8, T=100, U=20, V=1025, dtype=torch.float32) + + def benchmark_base(self): + self._run() + + def verify_base(self) -> torch.Tensor: + return self._reference() + + # Parakeet TDT v3 sizes (8192 tokens + blank), float32. + def setup_large(self): + self._make_inputs(B=4, T=200, U=40, V=8193, dtype=torch.float32) + + def benchmark_large(self): + self._run() + + def verify_large(self) -> torch.Tensor: + return self._reference() + + # Same sizes with bfloat16 logits. + def setup_large_bf16(self): + self._make_inputs(B=4, T=200, U=40, V=8193, dtype=torch.bfloat16) + + def benchmark_large_bf16(self): + self._run() + + def verify_large_bf16(self) -> torch.Tensor: + return self._reference() diff --git a/tdt-loss/build.toml b/tdt-loss/build.toml new file mode 100644 index 000000000..7028ad9cd --- /dev/null +++ b/tdt-loss/build.toml @@ -0,0 +1,28 @@ +[general] +name = "tdt-loss" +license = "Apache-2.0" +version = 1 +edition = 5 +backends = ["cuda"] + +[general.hub] +repo-id = "kernels-community/tdt-loss" + +[torch] +src = [ + "torch-ext/torch_binding.cpp", + "torch-ext/torch_binding.h", +] + +[torch.stable-abi] +cuda = "2.10" + +[kernel.tdt_loss] +backend = "cuda" +depends = ["torch"] +src = [ + "tdt_loss/common.cuh", + "tdt_loss/lattice.cu", + "tdt_loss/logprobs.cu", + "tdt_loss/stable.h", +] diff --git a/tdt-loss/flake.lock b/tdt-loss/flake.lock new file mode 100644 index 000000000..f80d937fa --- /dev/null +++ b/tdt-loss/flake.lock @@ -0,0 +1,117 @@ +{ + "nodes": { + "flake-compat": { + "locked": { + "lastModified": 1767039857, + "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "kernel-builder": { + "inputs": { + "flake-compat": "flake-compat", + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs", + "rust-overlay": "rust-overlay" + }, + "locked": { + "lastModified": 1790766967, + "narHash": "sha256-P4QLsZU4S/ZOaxIgc3KXTdaZ3uHJ+R7ko0ACKZG31vA=", + "owner": "huggingface", + "repo": "kernels", + "rev": "a56da98e235a3f52b4b766e4d93fac6a9e599e3d", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "kernels", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1789524109, + "narHash": "sha256-GZDvT/tbSVQB2KjI86l7VCM5D/vCrJo2sE3PCaaD0oQ=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "be3dab5b1a3c41ed2db86093b9f4e048305630b7", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable-small", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "kernel-builder": "kernel-builder" + } + }, + "rust-overlay": { + "inputs": { + "nixpkgs": [ + "kernel-builder", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1789543870, + "narHash": "sha256-45ikFmtfgBZgMePIm4GevJ9VjDSWO3PGQgOuuwHDpes=", + "owner": "oxalica", + "repo": "rust-overlay", + "rev": "35ca0490d13a3d38c4602d0eb9600a30fa63a367", + "type": "github" + }, + "original": { + "owner": "oxalica", + "repo": "rust-overlay", + "type": "github" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/tdt-loss/flake.nix b/tdt-loss/flake.nix new file mode 100644 index 000000000..ef8a89a25 --- /dev/null +++ b/tdt-loss/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for tdt-loss kernel"; + + inputs = { + kernel-builder.url = "github:huggingface/kernels"; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/tdt-loss/tdt_loss/common.cuh b/tdt-loss/tdt_loss/common.cuh new file mode 100644 index 000000000..f3dbb4b19 --- /dev/null +++ b/tdt-loss/tdt_loss/common.cuh @@ -0,0 +1,115 @@ +#pragma once + +#include +#include + +#include +#include + +namespace tdt_loss { + +constexpr int kWarpSize = 32; +// Upper bound on the number of durations, only used to size shared memory. +constexpr int kMaxDurations = 64; +// Block sizes of the per-node kernels and of the lattice recursions. The +// kernels are compiled with matching __launch_bounds__ so that register +// allocation always allows launching the largest block. +constexpr int kMaxRowThreads = 256; +constexpr int kMaxLatticeThreads = 1024; + +__device__ __forceinline__ float neg_inf() { return -INFINITY; } + +// log(exp(a) + exp(b)), safe when either argument is -inf. The arguments are +// float64 to keep large log-likelihoods exact, but the correction term only +// depends on |a - b| and is computed in float32. +__device__ __forceinline__ double log_add(double a, double b) { + if (a == -INFINITY) + return b; + if (b == -INFINITY) + return a; + return fmax(a, b) + log1pf(expf(-fabsf(static_cast(a - b)))); +} + +// log(sum(exp(v))) over the values produced by `for_each(f)`, which calls f(v) +// for every value. Two passes (max, then sum) keep the exponentials +// independent of each other instead of chaining log_add calls. +template +__device__ __forceinline__ double log_sum_exp(ForEach for_each) { + double m = -INFINITY; + for_each([&](double v) { m = fmax(m, v); }); + if (m == -INFINITY) + return m; + float s = 0.f; + for_each([&](double v) { s += expf(static_cast(v - m)); }); + return m + logf(s); +} + +// Merge two partial (max, sum(exp(x - max))) pairs of an online softmax. +__device__ __forceinline__ void merge_max_sum(float &m, float &s, float m_other, + float s_other) { + const float m_new = fmaxf(m, m_other); + if (m_new == -INFINITY) { + m = m_new; + s = 0.f; + return; + } + s = s * expf(m - m_new) + s_other * expf(m_other - m_new); + m = m_new; +} + +// Block-wide reduction of (max, sum) pairs. The result is valid in all threads. +// `shm_m` and `shm_s` must hold at least blockDim.x / kWarpSize floats. +__device__ __forceinline__ void +block_reduce_max_sum(float &m, float &s, float *shm_m, float *shm_s) { + const int lane = threadIdx.x % kWarpSize; + const int warp = threadIdx.x / kWarpSize; + const int num_warps = (blockDim.x + kWarpSize - 1) / kWarpSize; + +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset /= 2) { + const float m_other = __shfl_xor_sync(0xffffffff, m, offset); + const float s_other = __shfl_xor_sync(0xffffffff, s, offset); + merge_max_sum(m, s, m_other, s_other); + } + if (lane == 0) { + shm_m[warp] = m; + shm_s[warp] = s; + } + __syncthreads(); + if (warp == 0) { + m = lane < num_warps ? shm_m[lane] : -INFINITY; + s = lane < num_warps ? shm_s[lane] : 0.f; +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset /= 2) { + const float m_other = __shfl_xor_sync(0xffffffff, m, offset); + const float s_other = __shfl_xor_sync(0xffffffff, s, offset); + merge_max_sum(m, s, m_other, s_other); + } + if (lane == 0) { + shm_m[0] = m; + shm_s[0] = s; + } + } + __syncthreads(); + m = shm_m[0]; + s = shm_s[0]; +} + +// Per-sample lengths: T (frames) and U (labels); the lattice for the sample +// spans t in [0, T) and u in [0, U]. Lengths outside of the padded lattice +// trigger a device-side assert, like out-of-range indexing does in PyTorch; +// they are also clamped so that no out-of-bounds access happens when asserts +// are compiled out. +__device__ __forceinline__ void sample_lengths(const int *logit_lengths, + const int *target_lengths, int b, + int max_T, int max_U, int &T, + int &U) { + T = logit_lengths[b]; + U = target_lengths[b]; + CUDA_KERNEL_ASSERT(T >= 0 && T <= max_T && "logit_lengths out of range"); + CUDA_KERNEL_ASSERT(U >= 0 && U < max_U && "target_lengths out of range"); + T = min(max(T, 0), max_T); + U = min(max(U, 0), max_U - 1); +} + +} // namespace tdt_loss diff --git a/tdt-loss/tdt_loss/lattice.cu b/tdt-loss/tdt_loss/lattice.cu new file mode 100644 index 000000000..76608df73 --- /dev/null +++ b/tdt-loss/tdt_loss/lattice.cu @@ -0,0 +1,238 @@ +// Forward (alpha) and backward (beta) recursions over the TDT lattice. +// +// For a sample with T frames and U labels the lattice nodes are (t, u) with +// t in [0, T) and u in [0, U]. From node (t, u), for every duration d: +// - a blank arc (only d > 0) goes to (t + d, u), +// - a label arc goes to (t + d, u + 1). +// The alignment ends with a blank arc from (T - d, U) that exits the lattice. +// +// Every arc strictly increases t + u, so all nodes of an anti-diagonal +// t + u = n are independent. One thread block handles one sample and sweeps +// the anti-diagonals, with the threads striding over u. +// +// The recursions accumulate in float64: they are a small fraction of the total +// cost, and the gradient combines alpha + beta - log_likelihood, which cancels +// catastrophically in float32 when the log-likelihood is large (long targets). + +#include "common.cuh" +#include "stable.h" + +namespace tdt_loss { + +namespace { + +__device__ __forceinline__ void load_durations(const int *durations, int D, + int *shm_durations) { + for (int i = threadIdx.x; i < D; i += blockDim.x) { + shm_durations[i] = durations[i]; + } + __syncthreads(); +} + +__global__ void __launch_bounds__(kMaxLatticeThreads) tdt_alpha_kernel( + const float *__restrict__ blank_lp, const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, const int *__restrict__ logit_lengths, + const int *__restrict__ target_lengths, const int *__restrict__ durations, + int max_T, int max_U, int D, double *__restrict__ alphas, + double *__restrict__ log_ll) { + __shared__ int shm_durations[kMaxDurations]; + + const int b = blockIdx.x; + int T, U; + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); + load_durations(durations, D, shm_durations); + + const int64_t sample = static_cast(b) * max_T * max_U; + const float *blank = blank_lp + sample; + const float *label = label_lp + sample; + const float *dur = dur_lp + sample * D; + double *alpha = alphas + sample; + + for (int n = 0; n < T + U; ++n) { + const int u_min = max(0, n - T + 1); + const int u_max = min(n, U); + for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { + const int t = n - u; + const double a = n == 0 ? 0.0 : log_sum_exp([&](auto &&f) { + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_prev = t - d; + if (t_prev < 0) + continue; + if (d > 0) { + const int64_t src = static_cast(t_prev) * max_U + u; + f(alpha[src] + blank[src] + dur[src * D + i]); + } + if (u > 0) { + const int64_t src = static_cast(t_prev) * max_U + u - 1; + f(alpha[src] + label[src] + dur[src * D + i]); + } + } + }); + alpha[static_cast(t) * max_U + u] = a; + } + __syncthreads(); + } + + if (threadIdx.x == 0) { + double ll = -INFINITY; + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_prev = T - d; + if (d == 0 || t_prev < 0) + continue; + const int64_t src = static_cast(t_prev) * max_U + U; + ll = log_add(ll, alpha[src] + blank[src] + dur[src * D + i]); + } + log_ll[b] = ll; + } +} + +__global__ void __launch_bounds__(kMaxLatticeThreads) tdt_beta_kernel( + const float *__restrict__ blank_lp, const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, const int *__restrict__ logit_lengths, + const int *__restrict__ target_lengths, const int *__restrict__ durations, + int max_T, int max_U, int D, double *__restrict__ betas) { + __shared__ int shm_durations[kMaxDurations]; + + const int b = blockIdx.x; + int T, U; + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); + load_durations(durations, D, shm_durations); + + const int64_t sample = static_cast(b) * max_T * max_U; + const float *blank = blank_lp + sample; + const float *label = label_lp + sample; + const float *dur = dur_lp + sample * D; + double *beta = betas + sample; + + for (int n = T + U - 1; n >= 0; --n) { + const int u_min = max(0, n - T + 1); + const int u_max = min(n, U); + for (int u = u_min + threadIdx.x; u <= u_max; u += blockDim.x) { + const int t = n - u; + const int64_t node = static_cast(t) * max_U + u; + const double bt = log_sum_exp([&](auto &&f) { + for (int i = 0; i < D; ++i) { + const int d = shm_durations[i]; + const int t_next = t + d; + const double lp_dur = dur[node * D + i]; + if (d > 0) { + if (t_next < T) { + const int64_t dst = static_cast(t_next) * max_U + u; + f(beta[dst] + blank[node] + lp_dur); + } else if (t_next == T && u == U) { + f(blank[node] + lp_dur); // terminal arc + } + } + if (u < U && t_next < T) { + const int64_t dst = static_cast(t_next) * max_U + u + 1; + f(beta[dst] + label[node] + lp_dur); + } + } + }); + beta[node] = bt; + } + __syncthreads(); + } +} + +int lattice_block_size(int64_t max_U) { + int threads = 32; + while (threads < kMaxLatticeThreads && threads < max_U) + threads *= 2; + return threads; +} + +void check_lattice_args(Tensor const &blank_lp, Tensor const &label_lp, + Tensor const &dur_lp, Tensor const &logit_lengths, + Tensor const &target_lengths, Tensor const &durations, + Tensor const &out, Tensor const *log_ll) { + STD_TORCH_CHECK(blank_lp.is_cuda(), "blank_lp must be a CUDA tensor"); + STD_TORCH_CHECK(blank_lp.dim() == 3, + "blank_lp must have shape (batch, T, U+1)"); + for (auto const *x : {&blank_lp, &label_lp, &dur_lp, &out}) { + STD_TORCH_CHECK(same_device(*x, blank_lp), + "all tensors must be on the same device"); + STD_TORCH_CHECK(x->is_contiguous(), "lattice tensors must be contiguous"); + } + for (auto const *x : {&blank_lp, &label_lp, &dur_lp}) { + STD_TORCH_CHECK(x->scalar_type() == ScalarType::Float, + "log-prob tensors must be float32"); + } + STD_TORCH_CHECK(out.scalar_type() == ScalarType::Double, + "alphas/betas must be float64"); + for (auto const *x : {&logit_lengths, &target_lengths, &durations}) { + STD_TORCH_CHECK(same_device(*x, blank_lp), + "all tensors must be on the same device"); + STD_TORCH_CHECK(x->scalar_type() == ScalarType::Int, + "lengths and durations must be int32"); + STD_TORCH_CHECK(x->is_contiguous(), + "lengths and durations must be contiguous"); + } + const int64_t B = blank_lp.size(0); + const int64_t nodes = blank_lp.numel(); + STD_TORCH_CHECK(label_lp.numel() == nodes && out.numel() == nodes, + "lattice tensors must have shape (batch, T, U+1)"); + STD_TORCH_CHECK(dur_lp.numel() == nodes * durations.numel(), + "dur_lp must have shape (batch, T, U+1, num_durations)"); + STD_TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, + "lengths must have shape (batch,)"); + if (log_ll != nullptr) { + STD_TORCH_CHECK( + same_device(*log_ll, blank_lp) && + log_ll->scalar_type() == ScalarType::Double && + log_ll->is_contiguous() && log_ll->numel() == B, + "log_ll must be a contiguous float64 tensor of shape (batch,)"); + } + STD_TORCH_CHECK(durations.numel() <= kMaxDurations, "at most ", kMaxDurations, + " durations are supported"); +} + +} // namespace + +} // namespace tdt_loss + +void tdt_loss_fwd(Tensor const &blank_lp, Tensor const &label_lp, + Tensor const &dur_lp, Tensor const &logit_lengths, + Tensor const &target_lengths, Tensor const &durations, + Tensor &alphas, Tensor &log_ll) { + using namespace tdt_loss; + check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, + durations, alphas, &log_ll); + const int64_t B = blank_lp.size(0); + if (B == 0) + return; + + const torch::stable::accelerator::DeviceGuard device_guard( + blank_lp.get_device_index()); + const cudaStream_t stream = current_stream(blank_lp); + tdt_alpha_kernel<<>>( + ptr(blank_lp), ptr(label_lp), ptr(dur_lp), + ptr(logit_lengths), ptr(target_lengths), ptr(durations), + blank_lp.size(1), blank_lp.size(2), durations.numel(), + ptr(alphas), ptr(log_ll)); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void tdt_loss_bwd(Tensor const &blank_lp, Tensor const &label_lp, + Tensor const &dur_lp, Tensor const &logit_lengths, + Tensor const &target_lengths, Tensor const &durations, + Tensor &betas) { + using namespace tdt_loss; + check_lattice_args(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, + durations, betas, nullptr); + const int64_t B = blank_lp.size(0); + if (B == 0) + return; + + const torch::stable::accelerator::DeviceGuard device_guard( + blank_lp.get_device_index()); + const cudaStream_t stream = current_stream(blank_lp); + tdt_beta_kernel<<>>( + ptr(blank_lp), ptr(label_lp), ptr(dur_lp), + ptr(logit_lengths), ptr(target_lengths), ptr(durations), + blank_lp.size(1), blank_lp.size(2), durations.numel(), + ptr(betas)); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/tdt-loss/tdt_loss/logprobs.cu b/tdt-loss/tdt_loss/logprobs.cu new file mode 100644 index 000000000..b74a64e34 --- /dev/null +++ b/tdt-loss/tdt_loss/logprobs.cu @@ -0,0 +1,408 @@ +// Per-lattice-node kernels of the TDT loss: the log-softmax gather that feeds +// the alpha/beta recursions, and the fused gradient w.r.t. the logits. +// +// Both kernels use one thread block per lattice node (b, t, u) and stream the +// vocabulary dimension. Nodes outside of a sample's lattice (t >= T or u > U) +// never read the logits. + +#include + +#include "common.cuh" +#include "stable.h" + +namespace tdt_loss { + +namespace { + +template +__global__ void __launch_bounds__(kMaxRowThreads) tdt_logprobs_fwd_kernel( + const scalar_t *__restrict__ token_logits, int64_t tok_sb, int64_t tok_st, + int64_t tok_su, const scalar_t *__restrict__ duration_logits, + int64_t dur_sb, int64_t dur_st, int64_t dur_su, + const int *__restrict__ targets, int64_t targets_stride, + const int *__restrict__ logit_lengths, + const int *__restrict__ target_lengths, int max_T, int max_U, int V, int D, + int blank_id, float sigma, float *__restrict__ blank_lp, + float *__restrict__ label_lp, float *__restrict__ dur_lp, + float *__restrict__ token_lse) { + __shared__ float shm_m[kWarpSize]; + __shared__ float shm_s[kWarpSize]; + + const int64_t row = blockIdx.x; + const int u = row % max_U; + const int t = (row / max_U) % max_T; + const int b = row / (static_cast(max_U) * max_T); + + int T, U; + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); + + if (t >= T || u > U) { + if (threadIdx.x == 0) { + blank_lp[row] = neg_inf(); + label_lp[row] = neg_inf(); + token_lse[row] = 0.f; + } + for (int i = threadIdx.x; i < D; i += blockDim.x) { + dur_lp[row * D + i] = neg_inf(); + } + return; + } + + const scalar_t *x = token_logits + b * tok_sb + t * tok_st + u * tok_su; + + // Online softmax normalizer over the vocabulary. + float m = neg_inf(); + float s = 0.f; + for (int v = threadIdx.x; v < V; v += blockDim.x) { + const float xv = static_cast(x[v]); + if (xv > m) { + s = s * expf(m - xv) + 1.f; + m = xv; + } else if (xv != -INFINITY) { + s += expf(xv - m); + } + } + block_reduce_max_sum(m, s, shm_m, shm_s); + + if (threadIdx.x == 0) { + const float lse = m + logf(s); + token_lse[row] = lse; + blank_lp[row] = static_cast(x[blank_id]) - lse - sigma; + + float label = neg_inf(); + if (u < U) { + const int target = targets[b * targets_stride + u]; + CUDA_KERNEL_ASSERT(target >= 0 && target < V && "target out of range"); + if (target >= 0 && target < V) { + label = static_cast(x[target]) - lse - sigma; + } + } + label_lp[row] = label; + + // The duration head is tiny (a handful of classes), a serial + // log-softmax is cheaper than another block reduction. + const scalar_t *y = duration_logits + b * dur_sb + t * dur_st + u * dur_su; + float dm = neg_inf(); + for (int i = 0; i < D; ++i) + dm = fmaxf(dm, static_cast(y[i])); + float ds = 0.f; + for (int i = 0; i < D; ++i) + ds += expf(static_cast(y[i]) - dm); + const float dlse = dm + logf(ds); + for (int i = 0; i < D; ++i) { + dur_lp[row * D + i] = static_cast(y[i]) - dlse; + } + } +} + +template +__global__ void __launch_bounds__(kMaxRowThreads) tdt_logits_grad_kernel( + const scalar_t *__restrict__ token_logits, int64_t tok_sb, int64_t tok_st, + int64_t tok_su, const int *__restrict__ targets, int64_t targets_stride, + const int *__restrict__ logit_lengths, + const int *__restrict__ target_lengths, const int *__restrict__ durations, + const float *__restrict__ blank_lp, const float *__restrict__ label_lp, + const float *__restrict__ dur_lp, const float *__restrict__ token_lse, + const double *__restrict__ alphas, const double *__restrict__ betas, + const double *__restrict__ log_ll, const float *__restrict__ grad_loss, + int max_T, int max_U, int V, int D, int blank_id, + scalar_t *__restrict__ grad_token_logits, + scalar_t *__restrict__ grad_duration_logits) { + // Posterior occupancy of the blank / label arc leaving this node, for each + // duration. + __shared__ float shm_blank[kMaxDurations]; + __shared__ float shm_label[kMaxDurations]; + + const int64_t row = blockIdx.x; + const int u = row % max_U; + const int t = (row / max_U) % max_T; + const int b = row / (static_cast(max_U) * max_T); + + int T, U; + sample_lengths(logit_lengths, target_lengths, b, max_T, max_U, T, U); + + scalar_t *gx = grad_token_logits + row * V; + scalar_t *gd = grad_duration_logits + row * D; + const scalar_t zero = static_cast(0.f); + + const double ll = log_ll[b]; + // Infeasible samples (infinite loss) get a zero gradient instead of NaNs. + if (t >= T || u > U || !isfinite(ll)) { + for (int v = threadIdx.x; v < V; v += blockDim.x) + gx[v] = zero; + for (int i = threadIdx.x; i < D; i += blockDim.x) + gd[i] = zero; + return; + } + + const int64_t sample = static_cast(b) * max_T * max_U; + const double alpha = alphas[row]; + const double blank = blank_lp[row]; + const double label = label_lp[row]; + + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const int d = durations[i]; + const int t_next = t + d; + const double lp_dur = dur_lp[row * D + i]; + + float p_blank = 0.f; + if (d > 0) { + double beta_next = -INFINITY; + if (t_next < T) { + beta_next = betas[sample + static_cast(t_next) * max_U + u]; + } else if (t_next == T && u == U) { + beta_next = 0.0; // terminal arc + } + if (beta_next != -INFINITY) { + p_blank = + expf(static_cast(alpha + blank + lp_dur + beta_next - ll)); + } + } + + float p_label = 0.f; + if (u < U && t_next < T) { + const double beta_next = + betas[sample + static_cast(t_next) * max_U + u + 1]; + p_label = + expf(static_cast(alpha + label + lp_dur + beta_next - ll)); + } + + shm_blank[i] = p_blank; + shm_label[i] = p_label; + } + __syncthreads(); + + float occ_blank = 0.f; + float occ_label = 0.f; + for (int i = 0; i < D; ++i) { + occ_blank += shm_blank[i]; + occ_label += shm_label[i]; + } + + // loss = -log_ll, so d loss / d log_prob(arc) = -occupancy(arc). + const float go = grad_loss[b]; + const float g_blank = -go * occ_blank; + const float g_label = -go * occ_label; + const float g_total = g_blank + g_label; + + // Through log_softmax: d/dx_j = g_j - softmax_j * sum_k g_k. + for (int i = threadIdx.x; i < D; i += blockDim.x) { + const float g_dur = -go * (shm_blank[i] + shm_label[i]); + gd[i] = static_cast(g_dur - expf(dur_lp[row * D + i]) * g_total); + } + + if (g_blank == 0.f && g_label == 0.f) { + for (int v = threadIdx.x; v < V; v += blockDim.x) + gx[v] = zero; + return; + } + + const scalar_t *x = token_logits + b * tok_sb + t * tok_st + u * tok_su; + const float lse = token_lse[row]; + const int target = u < U ? targets[b * targets_stride + u] : -1; + for (int v = threadIdx.x; v < V; v += blockDim.x) { + float g = -expf(static_cast(x[v]) - lse) * g_total; + if (v == blank_id) + g += g_blank; + if (v == target) + g += g_label; + gx[v] = static_cast(g); + } +} + +int row_block_size(int64_t V) { + int threads = 32; + while (threads < kMaxRowThreads && threads < V) + threads *= 2; + return threads; +} + +void check_logits(Tensor const &token_logits, Tensor const &duration_logits) { + STD_TORCH_CHECK(token_logits.is_cuda(), "token_logits must be a CUDA tensor"); + STD_TORCH_CHECK(token_logits.dim() == 4, + "token_logits must have shape (batch, T, U+1, vocab_size+1)"); + STD_TORCH_CHECK( + duration_logits.dim() == 4, + "duration_logits must have shape (batch, T, U+1, num_durations)"); + STD_TORCH_CHECK( + same_device(duration_logits, token_logits), + "token_logits and duration_logits must be on the same device"); + STD_TORCH_CHECK(duration_logits.scalar_type() == token_logits.scalar_type(), + "token_logits and duration_logits must have the same dtype"); + for (int i = 0; i < 3; ++i) { + STD_TORCH_CHECK(duration_logits.size(i) == token_logits.size(i), + "token_logits and duration_logits must agree on (batch, T, " + "U+1)"); + } + STD_TORCH_CHECK(token_logits.stride(3) == 1 && duration_logits.stride(3) == 1, + "the last dimension of the logits must be contiguous"); + STD_TORCH_CHECK(duration_logits.size(3) <= kMaxDurations, "at most ", + kMaxDurations, " durations are supported"); +} + +void check_int_tensor(Tensor const &x, char const *name, Tensor const &like) { + STD_TORCH_CHECK(same_device(x, like), name, + " must be on the same device as the logits"); + STD_TORCH_CHECK(x.scalar_type() == ScalarType::Int, name, " must be int32"); + STD_TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); +} + +void check_float_out(Tensor const &x, char const *name, Tensor const &like) { + STD_TORCH_CHECK(same_device(x, like), name, + " must be on the same device as the logits"); + STD_TORCH_CHECK(x.scalar_type() == ScalarType::Float, name, + " must be float32"); + STD_TORCH_CHECK(x.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +} // namespace tdt_loss + +void tdt_logprobs_fwd(Tensor const &token_logits, Tensor const &duration_logits, + Tensor const &targets, Tensor const &logit_lengths, + Tensor const &target_lengths, int64_t blank_id, + double sigma, Tensor &blank_lp, Tensor &label_lp, + Tensor &dur_lp, Tensor &token_lse) { + using namespace tdt_loss; + + check_logits(token_logits, duration_logits); + const int64_t B = token_logits.size(0); + const int64_t max_T = token_logits.size(1); + const int64_t max_U = token_logits.size(2); + const int64_t V = token_logits.size(3); + const int64_t D = duration_logits.size(3); + + STD_TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); + STD_TORCH_CHECK(targets.dim() == 2 && targets.size(0) == B && + targets.size(1) >= max_U - 1, + "targets must have shape (batch, U)"); + check_int_tensor(targets, "targets", token_logits); + check_int_tensor(logit_lengths, "logit_lengths", token_logits); + check_int_tensor(target_lengths, "target_lengths", token_logits); + STD_TORCH_CHECK(logit_lengths.numel() == B && target_lengths.numel() == B, + "lengths must have shape (batch,)"); + for (auto *out : {&blank_lp, &label_lp, &token_lse}) { + check_float_out(*out, "log-prob output", token_logits); + STD_TORCH_CHECK(out->numel() == B * max_T * max_U, + "log-prob outputs must have shape (batch, T, U+1)"); + } + check_float_out(dur_lp, "dur_lp", token_logits); + STD_TORCH_CHECK(dur_lp.numel() == B * max_T * max_U * D, + "dur_lp must have shape (batch, T, U+1, num_durations)"); + + const int64_t rows = B * max_T * max_U; + if (rows == 0) + return; + STD_TORCH_CHECK(rows <= std::numeric_limits::max(), + "batch * T * (U+1) is too large"); + + const torch::stable::accelerator::DeviceGuard device_guard( + token_logits.get_device_index()); + const cudaStream_t stream = current_stream(token_logits); + const int threads = row_block_size(V); + + TDT_DISPATCH_FLOATING_TYPES( + token_logits.scalar_type(), "tdt_logprobs_fwd", [&] { + tdt_logprobs_fwd_kernel + <<(rows), threads, 0, stream>>>( + ptr(token_logits), token_logits.stride(0), + token_logits.stride(1), token_logits.stride(2), + ptr(duration_logits), duration_logits.stride(0), + duration_logits.stride(1), duration_logits.stride(2), + ptr(targets), targets.stride(0), ptr(logit_lengths), + ptr(target_lengths), max_T, max_U, V, D, blank_id, + static_cast(sigma), ptr(blank_lp), + ptr(label_lp), ptr(dur_lp), + ptr(token_lse)); + }); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void tdt_logits_grad(Tensor const &token_logits, Tensor const &targets, + Tensor const &logit_lengths, Tensor const &target_lengths, + Tensor const &durations, Tensor const &blank_lp, + Tensor const &label_lp, Tensor const &dur_lp, + Tensor const &token_lse, Tensor const &alphas, + Tensor const &betas, Tensor const &log_ll, + Tensor const &grad_loss, int64_t blank_id, + Tensor &grad_token_logits, Tensor &grad_duration_logits) { + using namespace tdt_loss; + + STD_TORCH_CHECK(token_logits.is_cuda() && token_logits.dim() == 4, + "token_logits must be a 4D CUDA tensor"); + STD_TORCH_CHECK(token_logits.stride(3) == 1, + "the last dimension of token_logits must be contiguous"); + const int64_t B = token_logits.size(0); + const int64_t max_T = token_logits.size(1); + const int64_t max_U = token_logits.size(2); + const int64_t V = token_logits.size(3); + const int64_t D = durations.numel(); + + STD_TORCH_CHECK(D <= kMaxDurations, "at most ", kMaxDurations, + " durations are supported"); + STD_TORCH_CHECK(blank_id >= 0 && blank_id < V, "blank_id out of range"); + check_int_tensor(targets, "targets", token_logits); + check_int_tensor(logit_lengths, "logit_lengths", token_logits); + check_int_tensor(target_lengths, "target_lengths", token_logits); + check_int_tensor(durations, "durations", token_logits); + for (auto *in : {&blank_lp, &label_lp, &token_lse}) { + check_float_out(*in, "log-prob tensor", token_logits); + STD_TORCH_CHECK(in->numel() == B * max_T * max_U, + "log-prob tensors must have shape (batch, T, U+1)"); + } + for (auto *in : {&alphas, &betas}) { + STD_TORCH_CHECK( + same_device(*in, token_logits) && + in->scalar_type() == ScalarType::Double && in->is_contiguous() && + in->numel() == B * max_T * max_U, + "alphas and betas must be contiguous float64 tensors of shape " + "(batch, T, U+1)"); + } + check_float_out(dur_lp, "dur_lp", token_logits); + STD_TORCH_CHECK(same_device(log_ll, token_logits) && + log_ll.scalar_type() == ScalarType::Double && + log_ll.is_contiguous(), + "log_ll must be a contiguous float64 tensor"); + check_float_out(grad_loss, "grad_loss", token_logits); + STD_TORCH_CHECK(log_ll.numel() == B && grad_loss.numel() == B, + "log_ll and grad_loss must have shape (batch,)"); + STD_TORCH_CHECK( + grad_token_logits.is_contiguous() && + same_sizes(grad_token_logits, token_logits) && + grad_token_logits.scalar_type() == token_logits.scalar_type(), + "grad_token_logits must be contiguous and match token_logits"); + STD_TORCH_CHECK(grad_duration_logits.is_contiguous() && + grad_duration_logits.numel() == B * max_T * max_U * D && + grad_duration_logits.scalar_type() == + token_logits.scalar_type(), + "grad_duration_logits must be contiguous and match " + "duration_logits"); + + const int64_t rows = B * max_T * max_U; + if (rows == 0) + return; + STD_TORCH_CHECK(rows <= std::numeric_limits::max(), + "batch * T * (U+1) is too large"); + + const torch::stable::accelerator::DeviceGuard device_guard( + token_logits.get_device_index()); + const cudaStream_t stream = current_stream(token_logits); + const int threads = row_block_size(V); + + TDT_DISPATCH_FLOATING_TYPES( + token_logits.scalar_type(), "tdt_logits_grad", [&] { + tdt_logits_grad_kernel + <<(rows), threads, 0, stream>>>( + ptr(token_logits), token_logits.stride(0), + token_logits.stride(1), token_logits.stride(2), + ptr(targets), targets.stride(0), ptr(logit_lengths), + ptr(target_lengths), ptr(durations), + ptr(blank_lp), ptr(label_lp), ptr(dur_lp), + ptr(token_lse), ptr(alphas), ptr(betas), + ptr(log_ll), ptr(grad_loss), max_T, max_U, V, D, + blank_id, ptr(grad_token_logits), + ptr(grad_duration_logits)); + }); + TDT_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/tdt-loss/tdt_loss/stable.h b/tdt-loss/tdt_loss/stable.h new file mode 100644 index 000000000..20b6b65cf --- /dev/null +++ b/tdt-loss/tdt_loss/stable.h @@ -0,0 +1,73 @@ +#pragma once + +// Helpers for the host code on top of the Torch stable ABI. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// The shim's stream accessor is guarded by USE_CUDA, so declare it here. +extern "C" AOTITorchError +aoti_torch_get_current_cuda_stream(int32_t device_index, void **ret_stream); + +using torch::headeronly::ScalarType; +using torch::stable::Tensor; + +// Dispatches on the logits dtype, with the type available as `scalar_t`. +#define TDT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + THO_DISPATCH_SWITCH( \ + TYPE, NAME, \ + THO_PRIVATE_CASE_TYPE_USING_HINT(ScalarType::Float, scalar_t, \ + __VA_ARGS__) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(ScalarType::Half, scalar_t, \ + __VA_ARGS__) \ + THO_PRIVATE_CASE_TYPE_USING_HINT(ScalarType::BFloat16, scalar_t, \ + __VA_ARGS__)) + +// Like STD_CUDA_KERNEL_LAUNCH_CHECK, which needs shim functions that are only +// declared with USE_CUDA. +#define TDT_CUDA_KERNEL_LAUNCH_CHECK() \ + do { \ + const cudaError_t err = cudaGetLastError(); \ + STD_TORCH_CHECK(err == cudaSuccess, \ + "CUDA error: ", cudaGetErrorString(err)); \ + } while (0) + +namespace tdt_loss { + +template T *ptr(Tensor const &x) { + return static_cast(x.data_ptr()); +} + +inline bool same_device(Tensor const &a, Tensor const &b) { + return a.is_cuda() == b.is_cuda() && + a.get_device_index() == b.get_device_index(); +} + +inline bool same_sizes(Tensor const &a, Tensor const &b) { + if (a.dim() != b.dim()) + return false; + for (int64_t i = 0; i < a.dim(); ++i) { + if (a.size(i) != b.size(i)) + return false; + } + return true; +} + +inline cudaStream_t current_stream(Tensor const &x) { + void *stream = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(x.get_device_index(), &stream)); + return static_cast(stream); +} + +} // namespace tdt_loss diff --git a/tdt-loss/tests/__init__.py b/tdt-loss/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tdt-loss/tests/reference.py b/tdt-loss/tests/reference.py new file mode 100644 index 000000000..62d6f470c --- /dev/null +++ b/tdt-loss/tests/reference.py @@ -0,0 +1,81 @@ +"""Pure PyTorch reference of the TDT loss, shared by the tests and the benchmark.""" + +import torch + + +def tdt_loss_reference( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_id, + durations, + sigma=0.0, + dtype=torch.float32, +): + """Per-sample TDT loss, as implemented in `transformers.loss.loss_tdt.tdt_loss` (computed in `dtype`).""" + device = token_logits.device + batch_size, max_t, max_u, _ = token_logits.shape + + token_log_probs = torch.log_softmax(token_logits.to(dtype), dim=-1) - sigma + duration_log_probs = torch.log_softmax(duration_logits.to(dtype), dim=-1) + + log_alpha = torch.full((batch_size, max_t, max_u), float("-inf"), device=device, dtype=dtype) + log_alpha[:, 0, 0] = 0.0 + blank_log_probs = token_log_probs[:, :, :, blank_id] + if max_u > 1: + targets_expanded = targets.long().unsqueeze(1).expand(-1, max_t, -1) + label_log_probs = torch.gather( + token_log_probs[:, :, : max_u - 1, :], dim=3, index=targets_expanded.unsqueeze(-1) + ).squeeze(-1) + + neg_inf = torch.tensor(float("-inf"), device=device, dtype=dtype) + for n in range(1, max_t + max_u - 1): + u_indices = torch.arange(max(0, n - max_t + 1), min(n + 1, max_u), device=device) + t_indices = n - u_indices + candidates = [] + for i, dur in enumerate(durations): + t_prev = t_indices - dur + valid_t = t_prev >= 0 + if not valid_t.any(): + continue + t_src = t_prev.clamp(min=0) + if dur > 0: + contrib = ( + log_alpha[:, t_src, u_indices] + + blank_log_probs[:, t_src, u_indices] + + duration_log_probs[:, t_src, u_indices, i] + ) + candidates.append(torch.where(valid_t.unsqueeze(0), contrib, neg_inf)) + valid_both = valid_t & (u_indices > 0) + if valid_both.any(): + u_src = (u_indices - 1).clamp(min=0) + u_src_label = u_src.clamp(max=max_u - 2) if max_u > 1 else u_src + contrib = ( + log_alpha[:, t_src, u_src] + + label_log_probs[:, t_src, u_src_label] + + duration_log_probs[:, t_src, u_src, i] + ) + candidates.append(torch.where(valid_both.unsqueeze(0), contrib, neg_inf)) + if candidates: + log_alpha[:, t_indices, u_indices] = torch.logsumexp(torch.stack(candidates), dim=0) + + batch_idx = torch.arange(batch_size, device=device) + target_lengths = target_lengths.long() + log_probs = torch.full((batch_size,), float("-inf"), device=device, dtype=dtype) + for i, dur in enumerate(durations): + if dur == 0: + continue + t_final = logit_lengths.long() - dur + valid = t_final >= 0 + if not valid.any(): + continue + t_clamped = t_final.clamp(min=0) + terminal = ( + log_alpha[batch_idx, t_clamped, target_lengths] + + token_log_probs[batch_idx, t_clamped, target_lengths, blank_id] + + duration_log_probs[batch_idx, t_clamped, target_lengths, i] + ) + log_probs = torch.where(valid, torch.logsumexp(torch.stack([log_probs, terminal]), dim=0), log_probs) + return -log_probs diff --git a/tdt-loss/tests/test_tdt_loss.py b/tdt-loss/tests/test_tdt_loss.py new file mode 100644 index 000000000..7c060ff5c --- /dev/null +++ b/tdt-loss/tests/test_tdt_loss.py @@ -0,0 +1,395 @@ +"""Tests for the TDT loss kernel. + +The kernel is checked against two pure PyTorch references: +- `tdt_loss_reference` (tests/reference.py): the vectorized implementation used in transformers, whose autograd + gradients are used to validate the kernel backward; +- `tdt_loss_naive`: a loop-based implementation that is easy to verify by hand. +""" + +import pytest +import torch + +import kernels + +from tests.reference import tdt_loss_reference + +if not torch.cuda.is_available(): + pytest.skip("CUDA not available", allow_module_level=True) + +tdt_loss_kernel = kernels.get_kernel("kernels-community/tdt-loss", version=1) + +DEVICE = torch.device("cuda") +DURATIONS = [0, 1, 2, 3, 4] +REDUCTIONS = ["mean_volume", "mean_batch", "mean", "sum", "none"] +# Gradient tolerance against a float64 reference for ~1000 labels (float32 log-probs, float64 lattice). +GRAD_ATOL_LONG = 1e-3 + + +def reduce(losses, target_lengths, reduction): + target_lengths = target_lengths.float() + if reduction == "mean_volume": + return losses.sum() / target_lengths.sum() + if reduction == "mean_batch": + return losses.mean() + if reduction == "mean": + return (losses / target_lengths).mean() + if reduction == "sum": + return losses.sum() + return losses + + +def tdt_loss_naive(token_logits, duration_logits, targets, logit_lengths, target_lengths, blank_id, durations, sigma=0.0): + """Per-sample TDT loss with explicit loops over the lattice, in float64 on the CPU.""" + token_lp = (torch.log_softmax(token_logits.double().cpu(), dim=-1) - sigma).tolist() + dur_lp = torch.log_softmax(duration_logits.double().cpu(), dim=-1).tolist() + targets = targets.tolist() + neg_inf = float("-inf") + + def log_add(a, b): + return torch.logaddexp(torch.tensor(a, dtype=torch.float64), torch.tensor(b, dtype=torch.float64)).item() + + losses = [] + for b in range(len(targets)): + T, U = int(logit_lengths[b]), int(target_lengths[b]) + alpha = [[neg_inf] * (U + 1) for _ in range(T)] + alpha[0][0] = 0.0 + for t in range(T): + for u in range(U + 1): + if t == 0 and u == 0: + continue + for i, d in enumerate(durations): + t_src = t - d + if t_src < 0: + continue + if d > 0: + arc = alpha[t_src][u] + token_lp[b][t_src][u][blank_id] + dur_lp[b][t_src][u][i] + alpha[t][u] = log_add(alpha[t][u], arc) + if u > 0: + label = targets[b][u - 1] + arc = alpha[t_src][u - 1] + token_lp[b][t_src][u - 1][label] + dur_lp[b][t_src][u - 1][i] + alpha[t][u] = log_add(alpha[t][u], arc) + log_ll = neg_inf + for i, d in enumerate(durations): + if d > 0 and T - d >= 0: + t_src = T - d + log_ll = log_add(log_ll, alpha[t_src][U] + token_lp[b][t_src][U][blank_id] + dur_lp[b][t_src][U][i]) + losses.append(-log_ll) + return torch.tensor(losses, dtype=torch.float64) + + +def assert_grad_close(actual, expected, **kwargs): + # The reference gets NaN gradients on lattice nodes that cannot be reached (e.g. u > t when there is no + # zero duration), from the backward of a logsumexp over only -inf. The true gradient there is zero. + torch.testing.assert_close(actual.to(expected.dtype), expected.nan_to_num(nan=0.0), **kwargs) + + +def make_inputs(B, T, U, V, durations, dtype=torch.float32, seed=0, joint=False): + """Random inputs with variable lengths. `joint=True` returns slices of a single joint tensor, + like `ParakeetForTDT` does.""" + generator = torch.Generator(device="cpu").manual_seed(seed) + D = len(durations) + logits = torch.randn(B, T, U + 1, V + D, generator=generator).to(DEVICE, dtype) + if joint: + token_logits, duration_logits = logits[..., :V], logits[..., V:] + else: + token_logits, duration_logits = logits[..., :V].contiguous(), logits[..., V:].contiguous() + targets = torch.randint(1, V, (B, U), generator=generator).to(DEVICE, torch.int32) + logit_lengths = torch.tensor([T - (b % 3) for b in range(B)], device=DEVICE, dtype=torch.int32) + target_lengths = torch.tensor([U - (b % 2) for b in range(B)], device=DEVICE, dtype=torch.int32) + return token_logits, duration_logits, targets, logit_lengths, target_lengths + + +def run_kernel(token_logits, duration_logits, targets, logit_lengths, target_lengths, durations, sigma, reduction): + return tdt_loss_kernel.tdt_loss( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + 0, + durations, + sigma=sigma, + reduction=reduction, + ) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("sigma", [0.0, 0.05]) +def test_matches_naive(sigma): + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(3, 12, 5, 16, DURATIONS) + expected = tdt_loss_naive( + token_logits, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS, sigma=sigma + ) + losses = run_kernel( + token_logits, duration_logits, targets, logit_lengths, target_lengths, DURATIONS, sigma, "none" + ) + torch.testing.assert_close(losses.double().cpu(), expected, atol=1e-4, rtol=1e-5) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("reduction", REDUCTIONS) +@pytest.mark.parametrize("sigma", [0.0, 0.05]) +@pytest.mark.parametrize("durations", [DURATIONS, [1, 2], [0, 1, 3, 5]]) +def test_forward_backward_match_reference(durations, sigma, reduction): + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(4, 30, 8, 64, durations) + + ref_tok = token_logits.clone().requires_grad_(True) + ref_dur = duration_logits.clone().requires_grad_(True) + ref_losses = tdt_loss_reference( + ref_tok, ref_dur, targets, logit_lengths, target_lengths, 0, durations, sigma=sigma + ) + expected = reduce(ref_losses, target_lengths, reduction) + expected.sum().backward() + + tok = token_logits.clone().requires_grad_(True) + dur = duration_logits.clone().requires_grad_(True) + loss = run_kernel(tok, dur, targets, logit_lengths, target_lengths, durations, sigma, reduction) + loss.sum().backward() + + torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) + assert_grad_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + assert_grad_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) + + +@pytest.mark.kernels_ci +def test_joint_logits_slices(): + """Non-contiguous slices of a joint output (as in ParakeetForTDT) are handled without a copy.""" + V = 128 + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs( + 3, 25, 6, V, DURATIONS, joint=True + ) + assert not token_logits.is_contiguous() and not duration_logits.is_contiguous() + + joint = token_logits._base.clone().requires_grad_(True) + loss = run_kernel( + joint[..., :V], joint[..., V:], targets, logit_lengths, target_lengths, DURATIONS, 0.0, "mean" + ) + loss.backward() + + ref_joint = token_logits._base.clone().requires_grad_(True) + expected = reduce( + tdt_loss_reference( + ref_joint[..., :V], ref_joint[..., V:], targets, logit_lengths, target_lengths, 0, DURATIONS + ), + target_lengths, + "mean", + ) + expected.backward() + + torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) + assert_grad_close(joint.grad, ref_joint.grad, atol=1e-5, rtol=1e-4) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_half_precision(dtype): + """Half-precision logits are read directly and gradients come back in the input dtype.""" + if dtype == torch.bfloat16 and torch.cuda.get_device_capability() < (8, 0): + pytest.skip("bfloat16 requires compute capability >= 8.0") + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs( + 4, 30, 8, 256, DURATIONS, dtype=dtype + ) + + tok = token_logits.clone().requires_grad_(True) + dur = duration_logits.clone().requires_grad_(True) + loss = run_kernel(tok, dur, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "mean") + loss.backward() + assert tok.grad.dtype == dtype and dur.grad.dtype == dtype + + # The reference upcasts the same half-precision values to float32. + ref_tok = token_logits.float().requires_grad_(True) + ref_dur = duration_logits.float().requires_grad_(True) + expected = reduce( + tdt_loss_reference(ref_tok, ref_dur, targets, logit_lengths, target_lengths, 0, DURATIONS), + target_lengths, + "mean", + ) + expected.backward() + + torch.testing.assert_close(loss, expected, atol=1e-4, rtol=1e-4) + assert_grad_close(tok.grad, ref_tok.grad, atol=1e-3, rtol=1e-2) + assert_grad_close(dur.grad, ref_dur.grad, atol=1e-3, rtol=1e-2) + + +def test_parakeet_like_shapes(): + """Vocabulary and label lengths of the size used by Parakeet TDT.""" + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs( + 2, 60, 20, 8193, DURATIONS, joint=True + ) + tok = token_logits.clone().requires_grad_(True) + dur = duration_logits.clone().requires_grad_(True) + losses = run_kernel(tok, dur, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "none") + losses.sum().backward() + + ref_tok = token_logits.clone().requires_grad_(True) + ref_dur = duration_logits.clone().requires_grad_(True) + expected = tdt_loss_reference(ref_tok, ref_dur, targets, logit_lengths, target_lengths, 0, DURATIONS) + expected.sum().backward() + + torch.testing.assert_close(losses, expected, atol=1e-3, rtol=1e-4) + assert_grad_close(tok.grad, ref_tok.grad, atol=1e-5, rtol=1e-4) + assert_grad_close(dur.grad, ref_dur.grad, atol=1e-5, rtol=1e-4) + + +def test_long_targets(): + """More labels than threads in a block (the lattice kernels stride over u).""" + U = 1100 + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(1, 6, U, 8, DURATIONS) + tok = token_logits.clone().requires_grad_(True) + loss = run_kernel(tok, duration_logits, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "sum") + loss.backward() + + # The log-likelihood is in the thousands, so compare against a float64 reference. + ref_tok = token_logits.double().requires_grad_(True) + expected = tdt_loss_reference( + ref_tok, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS, dtype=torch.float64 + ).sum() + expected.backward() + + torch.testing.assert_close(loss.double(), expected, atol=1e-2, rtol=1e-5) + assert_grad_close(tok.grad, ref_tok.grad, atol=GRAD_ATOL_LONG, rtol=0) + + +@pytest.mark.kernels_ci +def test_empty_targets(): + """Samples without labels: the only path is a sequence of blanks.""" + token_logits, duration_logits, _, logit_lengths, _ = make_inputs(2, 10, 0, 16, DURATIONS) + targets = torch.empty(2, 0, device=DEVICE, dtype=torch.int32) + target_lengths = torch.zeros(2, device=DEVICE, dtype=torch.int32) + + losses = run_kernel(token_logits, duration_logits, targets, logit_lengths, target_lengths, DURATIONS, 0.0, "none") + expected = tdt_loss_naive(token_logits, duration_logits, targets, logit_lengths, target_lengths, 0, DURATIONS) + torch.testing.assert_close(losses.double().cpu(), expected, atol=1e-4, rtol=1e-5) + + +@pytest.mark.kernels_ci +def test_infeasible_alignment(): + """No valid path: the loss is infinite and the gradient is zero rather than NaN.""" + durations = [0, 4] + token_logits, duration_logits, targets, _, target_lengths = make_inputs(2, 6, 3, 16, durations) + # Sample 0 emits all labels at t=0 and exits with a blank of duration 4; sample 1 has T=3 < 4. + logit_lengths = torch.tensor([4, 3], device=DEVICE, dtype=torch.int32) + + tok = token_logits.clone().requires_grad_(True) + losses = run_kernel(tok, duration_logits, targets, logit_lengths, target_lengths, durations, 0.0, "none") + assert torch.isfinite(losses[0]) and torch.isinf(losses[1]) + losses.backward(torch.ones_like(losses)) + assert torch.isfinite(tok.grad).all() + assert tok.grad[0].abs().sum() > 0 + assert (tok.grad[1] == 0).all() + + +@pytest.mark.kernels_ci +def test_deterministic(): + inputs = make_inputs(4, 30, 8, 64, DURATIONS) + grads = [] + for _ in range(2): + tok = inputs[0].clone().requires_grad_(True) + run_kernel(tok, *inputs[1:], DURATIONS, 0.0, "mean").backward() + grads.append(tok.grad) + assert torch.equal(grads[0], grads[1]) + + +@pytest.mark.kernels_ci +def test_invalid_reduction(): + inputs = make_inputs(1, 5, 2, 4, [0, 1]) + with pytest.raises(ValueError, match="Invalid reduction"): + run_kernel(*inputs, [0, 1], 0.0, "invalid") + + +@pytest.mark.kernels_ci +def test_negative_durations(): + inputs = make_inputs(1, 5, 2, 4, [0, 1]) + with pytest.raises(ValueError, match="non-negative"): + run_kernel(*inputs, [-1, 1], 0.0, "mean") + + +@pytest.mark.kernels_ci +def test_layer(): + """The `TDTLoss` layer matches the functional API and mirrors the `transformers` `tdt_loss` signature.""" + import inspect + + params = list(inspect.signature(tdt_loss_kernel.layers.TDTLoss.forward).parameters) + assert params == [ + "self", + "token_logits", + "duration_logits", + "targets", + "logit_lengths", + "target_lengths", + "blank_token_id", + "durations", + "sigma", + "reduction", + ] + + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(2, 12, 4, 16, DURATIONS) + kwargs = dict( + token_logits=token_logits, + duration_logits=duration_logits, + targets=targets, + logit_lengths=logit_lengths, + target_lengths=target_lengths, + blank_token_id=0, + durations=DURATIONS, + sigma=0.05, + reduction="none", + ) + torch.testing.assert_close(tdt_loss_kernel.layers.TDTLoss()(**kwargs), tdt_loss_kernel.tdt_loss(**kwargs)) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_torch_compile(dtype): + """The loss and its gradients are the same with `torch.compile(fullgraph=True)`.""" + if dtype == torch.bfloat16 and torch.cuda.get_device_capability() < (8, 0): + pytest.skip("bfloat16 requires compute capability >= 8.0") + V = 64 + token_logits, _, targets, logit_lengths, target_lengths = make_inputs(3, 20, 6, V, DURATIONS, dtype=dtype, joint=True) + + def loss_fn(logits, targets, logit_lengths, target_lengths): + return tdt_loss_kernel.tdt_loss( + logits[..., :V], logits[..., V:], targets, logit_lengths, target_lengths, 0, DURATIONS, sigma=0.05 + ) + + torch._dynamo.reset() + results = [] + for fn in (loss_fn, torch.compile(loss_fn, fullgraph=True)): + logits = token_logits._base.detach().clone().requires_grad_(True) + loss = fn(logits, targets, logit_lengths, target_lengths) + loss.backward() + results.append((loss.detach(), logits.grad)) + + torch.testing.assert_close(results[1][0], results[0][0]) + torch.testing.assert_close(results[1][1], results[0][1]) + + +@pytest.mark.kernels_ci +def test_opcheck(): + """The schemas and fake implementations of the custom ops are consistent (`torch.library.opcheck`).""" + ops = tdt_loss_kernel.loss.ops + B, T, U, V = 2, 8, 4, 16 + token_logits, duration_logits, targets, logit_lengths, target_lengths = make_inputs(B, T, U, V, DURATIONS) + durations = torch.tensor(DURATIONS, device=DEVICE, dtype=torch.int32) + f32 = dict(device=DEVICE, dtype=torch.float32) + f64 = dict(device=DEVICE, dtype=torch.float64) + shape = (B, T, U + 1) + + blank_lp, label_lp, token_lse = (torch.empty(shape, **f32) for _ in range(3)) + dur_lp = torch.empty(*shape, len(DURATIONS), **f32) + args = (token_logits, duration_logits, targets, logit_lengths, target_lengths, 0, 0.0) + args += (blank_lp, label_lp, dur_lp, token_lse) + torch.library.opcheck(ops.tdt_logprobs_fwd, args) + ops.tdt_logprobs_fwd(*args) + + lattice = (blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations) + alphas, betas, log_ll = torch.empty(shape, **f64), torch.empty(shape, **f64), torch.empty(B, **f64) + torch.library.opcheck(ops.tdt_loss_fwd, lattice + (alphas, log_ll)) + ops.tdt_loss_fwd(*lattice, alphas, log_ll) + torch.library.opcheck(ops.tdt_loss_bwd, lattice + (betas,)) + ops.tdt_loss_bwd(*lattice, betas) + + args = (token_logits, targets, logit_lengths, target_lengths, durations, blank_lp, label_lp, dur_lp, token_lse) + args += (alphas, betas, log_ll, torch.ones(B, **f32), 0) + args += (torch.empty_like(token_logits), torch.empty_like(duration_logits)) + torch.library.opcheck(ops.tdt_logits_grad, args) diff --git a/tdt-loss/torch-ext/tdt_loss/__init__.py b/tdt-loss/torch-ext/tdt_loss/__init__.py new file mode 100644 index 000000000..e5f7d3cb9 --- /dev/null +++ b/tdt-loss/torch-ext/tdt_loss/__init__.py @@ -0,0 +1,6 @@ +"""TDT (Token-and-Duration Transducer) loss CUDA kernel.""" + +from . import layers +from .loss import tdt_loss + +__all__ = ["layers", "tdt_loss"] diff --git a/tdt-loss/torch-ext/tdt_loss/layers.py b/tdt-loss/torch-ext/tdt_loss/layers.py new file mode 100644 index 000000000..fe75c1eec --- /dev/null +++ b/tdt-loss/torch-ext/tdt_loss/layers.py @@ -0,0 +1,40 @@ +from typing import Sequence, Union + +import torch +import torch.nn as nn + +from .loss import tdt_loss + + +class TDTLoss(nn.Module): + """ + Stateless layer computing the TDT loss with the CUDA kernels. Its forward has the same signature as + `transformers.loss.loss_tdt.tdt_loss`, so it can replace it through `kernels.kernelize`. + """ + + has_backward = True + can_torch_compile = True + + def forward( + self, + token_logits: torch.Tensor, + duration_logits: torch.Tensor, + targets: torch.Tensor, + logit_lengths: torch.Tensor, + target_lengths: torch.Tensor, + blank_token_id: int, + durations: Union[Sequence[int], torch.Tensor], + sigma: float = 0.0, + reduction: str = "mean", + ) -> torch.Tensor: + return tdt_loss( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_token_id, + durations, + sigma=sigma, + reduction=reduction, + ) diff --git a/tdt-loss/torch-ext/tdt_loss/loss.py b/tdt-loss/torch-ext/tdt_loss/loss.py new file mode 100644 index 000000000..3beb979b1 --- /dev/null +++ b/tdt-loss/torch-ext/tdt_loss/loss.py @@ -0,0 +1,254 @@ +"""TDT (Token-and-Duration Transducer) loss: autograd function and functional API.""" + +from typing import Dict, Sequence, Tuple, Union + +import torch + +from ._ops import add_op_namespace_prefix, ops + +_REDUCTIONS = ("mean_volume", "mean_batch", "mean", "sum", "none") + + +_DURATIONS_CACHE: Dict[Tuple[Tuple[int, ...], torch.device], torch.Tensor] = {} + + +def _durations_tensor(durations: Union[Sequence[int], torch.Tensor], device: torch.device) -> torch.Tensor: + """Validated int32 tensor of durations on `device`, cached in eager mode since they are fixed for a model.""" + if isinstance(durations, torch.Tensor): + durations = durations.tolist() + key = (tuple(int(d) for d in durations), device) + if any(d < 0 for d in key[0]): + raise ValueError(f"Durations must be non-negative, got {list(key[0])}.") + if torch.compiler.is_compiling(): + return torch.tensor(key[0], device=device, dtype=torch.int32) + if key not in _DURATIONS_CACHE: + _DURATIONS_CACHE[key] = torch.tensor(key[0], device=device, dtype=torch.int32) + return _DURATIONS_CACHE[key] + + +def _last_dim_contiguous(x: torch.Tensor) -> torch.Tensor: + # The kernels accept arbitrary strides for the (batch, T, U) dims, so slices of a + # joint `(..., vocab_size + num_durations)` output can be passed without a copy. + return x if x.stride(-1) == 1 else x.contiguous() + + +# The ops only write to preallocated outputs, so their fake (meta) implementations have nothing to do. They let +# `torch.compile` trace through the ops. +@torch.library.register_fake(add_op_namespace_prefix("tdt_logprobs_fwd")) +def _( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_id, + sigma, + blank_lp, + label_lp, + dur_lp, + token_lse, +): + return None + + +@torch.library.register_fake(add_op_namespace_prefix("tdt_loss_fwd")) +def _(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, alphas, log_ll): + return None + + +@torch.library.register_fake(add_op_namespace_prefix("tdt_loss_bwd")) +def _(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, betas): + return None + + +@torch.library.register_fake(add_op_namespace_prefix("tdt_logits_grad")) +def _( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + betas, + log_ll, + grad_loss, + blank_id, + grad_token_logits, + grad_duration_logits, +): + return None + + +class _TDTLossFunction(torch.autograd.Function): + """Per-sample TDT loss (negative log-likelihood) with a CUDA forward and backward.""" + + @staticmethod + def forward( + ctx, + token_logits: torch.Tensor, + duration_logits: torch.Tensor, + targets: torch.Tensor, + logit_lengths: torch.Tensor, + target_lengths: torch.Tensor, + durations: torch.Tensor, + blank_id: int, + sigma: float, + ) -> torch.Tensor: + token_logits = _last_dim_contiguous(token_logits) + duration_logits = _last_dim_contiguous(duration_logits) + + B, max_T, max_U, _ = token_logits.shape + D = duration_logits.shape[-1] + f32 = dict(device=token_logits.device, dtype=torch.float32) + + blank_lp = torch.empty(B, max_T, max_U, **f32) + label_lp = torch.empty(B, max_T, max_U, **f32) + dur_lp = torch.empty(B, max_T, max_U, D, **f32) + token_lse = torch.empty(B, max_T, max_U, **f32) + ops.tdt_logprobs_fwd( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + blank_id, + sigma, + blank_lp, + label_lp, + dur_lp, + token_lse, + ) + + # The lattice recursions run in float64, see lattice.cu. + f64 = dict(device=token_logits.device, dtype=torch.float64) + alphas = torch.empty(B, max_T, max_U, **f64) + log_ll = torch.empty(B, **f64) + ops.tdt_loss_fwd(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, alphas, log_ll) + + ctx.save_for_backward( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + log_ll, + ) + ctx.blank_id = blank_id + ctx.duration_shape = duration_logits.shape + ctx.duration_dtype = duration_logits.dtype + return -log_ll.float() + + @staticmethod + def backward(ctx, grad_loss: torch.Tensor): + ( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + log_ll, + ) = ctx.saved_tensors + if not (ctx.needs_input_grad[0] or ctx.needs_input_grad[1]): + return (None,) * 8 + + betas = torch.empty_like(alphas) + ops.tdt_loss_bwd(blank_lp, label_lp, dur_lp, logit_lengths, target_lengths, durations, betas) + + grad_token_logits = torch.empty(token_logits.shape, device=token_logits.device, dtype=token_logits.dtype) + grad_duration_logits = torch.empty(ctx.duration_shape, device=token_logits.device, dtype=ctx.duration_dtype) + ops.tdt_logits_grad( + token_logits, + targets, + logit_lengths, + target_lengths, + durations, + blank_lp, + label_lp, + dur_lp, + token_lse, + alphas, + betas, + log_ll, + grad_loss.float().contiguous(), + ctx.blank_id, + grad_token_logits, + grad_duration_logits, + ) + return grad_token_logits, grad_duration_logits, None, None, None, None, None, None + + +def tdt_loss( + token_logits: torch.Tensor, + duration_logits: torch.Tensor, + targets: torch.Tensor, + logit_lengths: torch.Tensor, + target_lengths: torch.Tensor, + blank_token_id: int, + durations: Union[Sequence[int], torch.Tensor], + sigma: float = 0.0, + reduction: str = "mean", +) -> torch.Tensor: + """Compute the TDT (Token-and-Duration Transducer) loss (https://arxiv.org/abs/2304.06795). + + Args: + token_logits: Token logits of shape `(batch, T, U+1, vocab_size+1)`, in float32, float16 or bfloat16. + Only the last dimension needs to be contiguous. + duration_logits: Duration logits of shape `(batch, T, U+1, num_durations)`, same dtype as `token_logits`. + targets: Target labels of shape `(batch, U)`. + logit_lengths: Encoder output lengths of shape `(batch,)`. + target_lengths: Target lengths of shape `(batch,)`. + blank_token_id: Blank token id. + durations: Non-negative duration values, e.g. `[0, 1, 2, 3, 4]`. + sigma: Logit undernormalization constant (see TDT paper). Defaults to `0.0`. + reduction: One of `"mean_volume"`, `"mean_batch"`, `"mean"`, `"sum"` or `"none"`, mirroring NeMo's + `RNNTLoss`. `"mean"` divides each loss by its target length before averaging over the batch. + + Returns: + Scalar loss tensor (or per-example losses of shape `(batch,)` if `reduction="none"`). Samples without a + valid alignment get an infinite loss and a zero gradient. + """ + if reduction not in _REDUCTIONS: + raise ValueError( + f'Invalid reduction mode "{reduction}". Expected one of {", ".join(repr(r) for r in _REDUCTIONS)}.' + ) + + device = token_logits.device + durations = _durations_tensor(durations, device) + targets = targets.to(device=device, dtype=torch.int32).contiguous() + logit_lengths = logit_lengths.to(device=device, dtype=torch.int32).contiguous() + target_lengths = target_lengths.to(device=device, dtype=torch.int32).contiguous() + + losses = _TDTLossFunction.apply( + token_logits, + duration_logits, + targets, + logit_lengths, + target_lengths, + durations, + int(blank_token_id), + float(sigma), + ) + + if reduction == "mean_volume": + return losses.sum() / target_lengths.sum().float() + if reduction == "mean_batch": + return losses.mean() + if reduction == "mean": + return (losses / target_lengths.float()).mean() + if reduction == "sum": + return losses.sum() + return losses diff --git a/tdt-loss/torch-ext/torch_binding.cpp b/tdt-loss/torch-ext/torch_binding.cpp new file mode 100644 index 000000000..f41039291 --- /dev/null +++ b/tdt-loss/torch-ext/torch_binding.cpp @@ -0,0 +1,35 @@ +#include + +#include "registration.h" +#include "torch_binding.h" + +STABLE_TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { + ops.def("tdt_logprobs_fwd(Tensor token_logits, Tensor duration_logits, " + "Tensor targets, Tensor logit_lengths, Tensor target_lengths, " + "int blank_id, float sigma, Tensor! blank_lp, Tensor! label_lp, " + "Tensor! dur_lp, Tensor! token_lse) -> ()"); + + ops.def("tdt_loss_fwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, " + "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " + "Tensor! alphas, Tensor! log_ll) -> ()"); + + ops.def("tdt_loss_bwd(Tensor blank_lp, Tensor label_lp, Tensor dur_lp, " + "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " + "Tensor! betas) -> ()"); + + ops.def("tdt_logits_grad(Tensor token_logits, Tensor targets, " + "Tensor logit_lengths, Tensor target_lengths, Tensor durations, " + "Tensor blank_lp, Tensor label_lp, Tensor dur_lp, Tensor token_lse, " + "Tensor alphas, Tensor betas, Tensor log_ll, Tensor grad_loss, " + "int blank_id, Tensor! grad_token_logits, " + "Tensor! grad_duration_logits) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, ops) { + ops.impl("tdt_logprobs_fwd", TORCH_BOX(&tdt_logprobs_fwd)); + ops.impl("tdt_loss_fwd", TORCH_BOX(&tdt_loss_fwd)); + ops.impl("tdt_loss_bwd", TORCH_BOX(&tdt_loss_bwd)); + ops.impl("tdt_logits_grad", TORCH_BOX(&tdt_logits_grad)); +} + +REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/tdt-loss/torch-ext/torch_binding.h b/tdt-loss/torch-ext/torch_binding.h new file mode 100644 index 000000000..3acc065cc --- /dev/null +++ b/tdt-loss/torch-ext/torch_binding.h @@ -0,0 +1,53 @@ +#pragma once + +#include + +// Log-softmax gather over the vocabulary and duration heads. Writes the blank +// and label log-probabilities of every lattice node, the duration +// log-probabilities and the token log-normalizer (reused by the backward). +void tdt_logprobs_fwd(torch::stable::Tensor const &token_logits, + torch::stable::Tensor const &duration_logits, + torch::stable::Tensor const &targets, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + int64_t blank_id, double sigma, + torch::stable::Tensor &blank_lp, + torch::stable::Tensor &label_lp, + torch::stable::Tensor &dur_lp, + torch::stable::Tensor &token_lse); + +// Forward recursion over the lattice (alphas) and per-sample log-likelihood. +void tdt_loss_fwd(torch::stable::Tensor const &blank_lp, + torch::stable::Tensor const &label_lp, + torch::stable::Tensor const &dur_lp, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + torch::stable::Tensor const &durations, + torch::stable::Tensor &alphas, torch::stable::Tensor &log_ll); + +// Backward recursion over the lattice (betas). +void tdt_loss_bwd(torch::stable::Tensor const &blank_lp, + torch::stable::Tensor const &label_lp, + torch::stable::Tensor const &dur_lp, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + torch::stable::Tensor const &durations, + torch::stable::Tensor &betas); + +// Gradient of the per-sample losses (scaled by grad_loss) w.r.t. the token +// and duration logits. +void tdt_logits_grad(torch::stable::Tensor const &token_logits, + torch::stable::Tensor const &targets, + torch::stable::Tensor const &logit_lengths, + torch::stable::Tensor const &target_lengths, + torch::stable::Tensor const &durations, + torch::stable::Tensor const &blank_lp, + torch::stable::Tensor const &label_lp, + torch::stable::Tensor const &dur_lp, + torch::stable::Tensor const &token_lse, + torch::stable::Tensor const &alphas, + torch::stable::Tensor const &betas, + torch::stable::Tensor const &log_ll, + torch::stable::Tensor const &grad_loss, int64_t blank_id, + torch::stable::Tensor &grad_token_logits, + torch::stable::Tensor &grad_duration_logits);