NVIDIA-NeMo/Speech · #16437
feat(asr): native CTC loss kernel over sequence-packed log probabilities
nemo/collections/asr/losses/ctc.py40 + / 4 −
@@ -14,11 +14,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import torch from torch import Tensor, nn from nemo.collections.asr.parts.packed_sequence import PackedEncoderActivations, unpack_encoder_output from nemo.core.classes import Serialization, Typing, typecheck from nemo.core.neural_types import LabelsType, LengthsType, LogprobsType, LossType, NeuralType+from nemo.core.utils.optional_libs import TRITON_AVAILABLE++if TRITON_AVAILABLE:+ from nemo.collections.asr.parts.triton.ctc_loss import PACKED_CTC_LOSS_DTYPES, packed_ctc_loss __all__ = ['CTCLoss'] @@ -68,10 +73,15 @@ def reduce(self, losses, target_lengths): def forward_packed(self, log_probs: PackedEncoderActivations, targets: Tensor, target_lengths: Tensor) -> Tensor: """Compute the CTC loss from token-flat log probabilities with explicit sequence boundaries. - This is a compatibility boundary, not a padding-free CTC kernel:- it materializes the padded log probabilities for the existing kernel and calls ``forward``,- so the encoder and the vocabulary projection can stay packed up to this call.- ``reduction`` and ``zero_infinity`` apply unchanged.+ On CUDA, for float32, bfloat16 and float16 log probabilities,+ the loss runs as a kernel over the token-flat log probabilities+ (:func:`nemo.collections.asr.parts.triton.ctc_loss.packed_ctc_loss`),+ so neither the padded ``(B, T, D)`` log probabilities nor their gradient are formed.+ Elsewhere, and for float64 (the kernel computes in float32),+ the padded log probabilities are materialized for the padded kernel of ``forward``.+ Both give the loss and the gradient of ``forward``;+ ``reduction`` and ``zero_infinity`` apply unchanged,+ and ``targets`` and ``target_lengths`` may live on the CPU, as for ``forward``. Args: log_probs: concatenated per-utterance log probabilities, shape ``(sum(T_i), D)`` including the blank.@@ -86,6 +96,32 @@ def forward_packed(self, log_probs: PackedEncoderActivations, targets: Tensor, t f"forward_packed expects PackedEncoderActivations as log_probs, got {type(log_probs).__name__}; " "padded log probabilities go to forward." )+ if log_probs.data.is_cuda and TRITON_AVAILABLE and log_probs.data.dtype in PACKED_CTC_LOSS_DTYPES:+ # the kernel reads the targets and their lengths on the device of the log probabilities,+ # and the reductions below divide by the lengths, so both move there+ targets = targets.to(device=log_probs.data.device, dtype=torch.long)+ target_lengths = target_lengths.to(device=log_probs.data.device, dtype=torch.long)+ losses = packed_ctc_loss(+ log_probs.data,+ log_probs.cu_seqlens,+ log_probs.lengths,+ targets,+ target_lengths,+ blank=self._blank,+ zero_infinity=self.zero_infinity,+ )+ if self.reduction == 'sum':+ loss = losses.sum()+ elif self.reduction == 'mean':+ # the padded kernel's mean: each loss over its target length, then the batch mean+ loss = (losses / target_lengths.clamp(min=1)).mean()+ elif self.reduction == 'none':+ loss = losses+ else:+ raise ValueError(f"unexpected reduction {self.reduction!r}")+ if self._apply_reduction:+ loss = self.reduce(loss, target_lengths)+ return loss return self.forward( log_probs=unpack_encoder_output(log_probs), targets=targets,nemo/collections/asr/parts/triton/ctc_loss.pyadded337 + / 0 −
@@ -0,0 +1,337 @@+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+# http://www.apache.org/licenses/LICENSE-2.0+#+# Unless required by applicable law or agreed to in writing, software+# distributed under the License is distributed on an "AS IS" BASIS,+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.+# See the License for the specific language governing permissions and+# limitations under the License.++"""CTC loss over sequence-packed log probabilities.++The log probabilities of a packed batch are one token-flat ``(T_total, D)`` buffer,+utterance ``b`` at the rows ``[cu_seqlens[b], cu_seqlens[b + 1])``.+The kernels read that buffer and write the gradient in the same layout,+so the loss needs neither the padded ``(B, T, D)`` copy that ``torch.nn.functional.ctc_loss`` takes+nor its padded gradient.++One program per utterance walks the utterance's frames.+The forward program runs the alpha recursion over the extended label sequence+``blank, l_1, blank, l_2, ..., l_U, blank`` (``S = 2 U + 1`` states, one lane per state)+and stores every alpha row into ``(T_total, S_max)``, as the padded kernel stores its ``(B, T, S)`` alphas.+The backward program runs the beta recursion from the last frame,+forms the state posteriors from the stored alphas,+and writes each frame's gradient row into the token-flat gradient:+``exp(log_probs) - posterior`` summed per class,+which is the gradient of the padded kernel+(the gradient with respect to the logits when ``log_probs`` is a ``log_softmax`` output).+The posteriors of the blank states are reduced in registers;+the label states are scattered with atomic adds,+as the padded CUDA kernel does for its non-blank labels.++The design follows RETURNN's packed full-sum loss+(``ctc_loss_packed`` and ``FastBaumWelchPackedOp`` in ``returnn/torch/util/native_op.py`` and ``returnn/native_op.py``+of https://github.com/rwth-i6/returnn):+the scores are addressed in the token-flat buffer through per-sequence start offsets,+the posteriors are scattered back into a buffer of the same layout,+the gradient is ``exp(log_probs) - posterior`` masked per frame,+the shapes are static (no host synchronization in the kernels, so a CUDA graph can capture them),+and the state count follows the longest target of the batch rather than the width of the targets buffer,+so an over-allocated buffer costs nothing (RETURNN's content-sized FSA edge layout).+It differs in three points:+the CTC topology is built into the recursion instead of a generic FSA edge list,+each program walks the frames of its own utterance instead of frame-wise kernels over all edges of the batch,+and the alphas are stored while the betas are recomputed in the backward,+instead of the full ``(T_total, D)`` posterior tensor.+"""++import torch+import triton+import triton.language as tl+from torch import Tensor++__all__ = ["PACKED_CTC_LOSS_DTYPES", "packed_ctc_loss"]++# the dtypes the kernels read; the recursion runs in float32 either way+PACKED_CTC_LOSS_DTYPES = (torch.float32, torch.bfloat16, torch.float16)+++def packed_ctc_loss(+ log_probs: Tensor,+ cu_seqlens: Tensor,+ lengths: Tensor,+ targets: Tensor,+ target_lengths: Tensor,+ *,+ blank: int,+ zero_infinity: bool = False,+) -> Tensor:+ """CTC loss per utterance over token-flat log probabilities.++ Same semantics as ``torch.nn.functional.ctc_loss`` with ``reduction="none"``:+ the loss of an utterance that cannot be aligned is ``inf``,+ or zero with ``zero_infinity``, which also zeroes its gradient,+ and the gradient is the one with respect to the logits of a ``log_softmax``.+ The label scatter of the backward uses atomic adds,+ so the summation order of repeated labels varies between runs,+ as it does in the padded CUDA kernel.++ Args:+ log_probs: log probabilities of all utterances concatenated along time, shape ``(T_total, D)``+ with the blank among the ``D`` classes; float32, bfloat16 or float16, computed in float32.+ cu_seqlens: cumulative frame counts, shape ``(B + 1,)``, integer.+ lengths: frames per utterance, shape ``(B,)``, integer.+ targets: padded target ids, shape ``(B, U)``, integer.+ target_lengths: target ids per utterance, shape ``(B,)``, integer, between 0 and ``U``+ (checked outside a CUDA graph capture, where no host read is possible).+ blank: index of the blank class.+ zero_infinity: replace an infinite loss and its gradient by zero.++ Returns:+ The loss per utterance, shape ``(B,)``, float32.+ """+ if log_probs.ndim != 2:+ raise ValueError(f"log_probs must have shape (T_total, D), got {tuple(log_probs.shape)}.")+ if not log_probs.is_cuda:+ raise RuntimeError("The packed CTC loss kernel requires CUDA tensors.")+ if log_probs.dtype not in PACKED_CTC_LOSS_DTYPES:+ raise TypeError(f"log_probs must be float32, bfloat16 or float16, got {log_probs.dtype}.")+ batch_size = lengths.numel()+ if cu_seqlens.shape != (batch_size + 1,):+ raise ValueError(f"cu_seqlens must have shape ({batch_size + 1},), got {tuple(cu_seqlens.shape)}.")+ if targets.ndim != 2 or targets.shape[0] != batch_size:+ raise ValueError(f"targets must have shape ({batch_size}, U), got {tuple(targets.shape)}.")+ if target_lengths.shape != (batch_size,):+ raise ValueError(f"target_lengths must have shape ({batch_size},), got {tuple(target_lengths.shape)}.")+ if not 0 <= blank < log_probs.shape[1]:+ raise ValueError(f"blank must be a class index below {log_probs.shape[1]}, got {blank}.")+ # The host checks of the padded kernel (no negative length, none beyond the targets width),+ # and the longest target sizes the states, so that an over-allocated targets buffer costs nothing.+ # A graph capture cannot sync: there the width stands in, and the sampler guarantees the bounds.+ longest = targets.shape[1]+ if batch_size and not torch.cuda.is_current_stream_capturing():+ shortest, longest = (int(bound) for bound in torch.aminmax(target_lengths))+ if shortest < 0:+ raise ValueError(f"target_lengths must not be negative, got {shortest}.")+ if longest > targets.shape[1]:+ raise ValueError(f"target_lengths must not exceed the targets width {targets.shape[1]}, got {longest}.")+ return _PackedCTCLoss.apply(+ log_probs.contiguous(),+ cu_seqlens.contiguous(),+ lengths.contiguous(),+ targets.contiguous(),+ target_lengths.contiguous(),+ blank,+ zero_infinity,+ 2 * longest + 1,+ )+++class _PackedCTCLoss(torch.autograd.Function):+ @staticmethod+ def forward(ctx, log_probs, cu_seqlens, lengths, targets, target_lengths, blank, zero_infinity, num_states_max):+ batch_size, target_width = targets.shape+ alpha = torch.empty((log_probs.shape[0], num_states_max), device=log_probs.device, dtype=torch.float32)+ loss = torch.empty((batch_size,), device=log_probs.device, dtype=torch.float32)+ if batch_size:+ _alpha_kernel[(batch_size,)](+ log_probs,+ cu_seqlens,+ lengths,+ targets,+ target_lengths,+ alpha,+ loss,+ log_probs.shape[1],+ target_width,+ num_states_max,+ blank,+ BLOCK_S=_state_block(num_states_max),+ num_warps=_num_warps(num_states_max),+ )+ ctx.save_for_backward(log_probs, cu_seqlens, lengths, targets, target_lengths, alpha, loss)+ ctx.blank = blank+ ctx.zero_infinity = zero_infinity+ ctx.num_states_max = num_states_max+ if zero_infinity:+ return torch.where(torch.isposinf(loss), 0.0, loss)+ return loss++ @staticmethod+ def backward(ctx, grad_output):+ log_probs, cu_seqlens, lengths, targets, target_lengths, alpha, loss = ctx.saved_tensors+ batch_size, target_width = targets.shape+ num_states_max = ctx.num_states_max+ # zeros, not empty: a frame outside every utterance (a gapped layout) has no gradient+ grad = torch.zeros(log_probs.shape, device=log_probs.device, dtype=torch.float32)+ if batch_size:+ _beta_grad_kernel[(batch_size,)](+ log_probs,+ cu_seqlens,+ lengths,+ targets,+ target_lengths,+ alpha,+ loss,+ grad_output.contiguous().to(torch.float32),+ grad,+ log_probs.shape[1],+ target_width,+ num_states_max,+ ctx.blank,+ ZERO_INFINITY=ctx.zero_infinity,+ BLOCK_S=_state_block(num_states_max),+ BLOCK_D=1024,+ num_warps=_num_warps(num_states_max),+ )+ return grad.to(log_probs.dtype), None, None, None, None, None, None, None+++def _state_block(num_states_max: int) -> int:+ return max(16, triton.next_power_of_2(num_states_max))+++def _num_warps(num_states_max: int) -> int:+ return 4 if num_states_max <= 512 else 8+++@triton.jit+def _logsumexp3(a, b, c):+ peak = tl.maximum(tl.maximum(a, b), c)+ finite_peak = tl.where(peak == float("-inf"), 0.0, peak)+ total = tl.exp(a - finite_peak) + tl.exp(b - finite_peak) + tl.exp(c - finite_peak)+ return tl.where(peak == float("-inf"), float("-inf"), finite_peak + tl.log(total))+++@triton.jit+def _extended_labels(targets_ptr, b, target_width, state, valid, blank):+ """``l'_s``: the blank at even states, ``targets[b, (s - 1) // 2]`` at odd states; blank beyond the utterance."""+ is_label = valid & (state % 2 == 1)+ return tl.load(targets_ptr + b * target_width + (state - 1) // 2, mask=is_label, other=blank).to(tl.int32)+++@triton.jit+def _frame_log_probs(log_probs_ptr, row, num_classes, label, valid):+ return tl.load(log_probs_ptr + row * num_classes + label, mask=valid, other=float("-inf")).to(tl.float32)+++@triton.jit+def _alpha_kernel(+ log_probs_ptr,+ cu_seqlens_ptr,+ lengths_ptr,+ targets_ptr,+ target_lengths_ptr,+ alpha_ptr,+ loss_ptr,+ num_classes,+ target_width,+ num_states_max,+ blank,+ BLOCK_S: tl.constexpr,+):+ b = tl.program_id(0)+ start = tl.load(cu_seqlens_ptr + b).to(tl.int64)+ num_frames = tl.load(lengths_ptr + b).to(tl.int32)+ num_labels = tl.load(target_lengths_ptr + b).to(tl.int32)+ num_states = 2 * num_labels + 1+ state = tl.arange(0, BLOCK_S)+ valid = state < num_states+ label = _extended_labels(targets_ptr, b, target_width, state, valid, blank)+ # the skip over the blank between two different labels; a target id equal to the blank is a label like any other,+ # as in the padded kernel+ label_before = tl.gather(label, tl.maximum(state - 2, 0), 0)+ can_skip = valid & (state >= 2) & (state % 2 == 1) & (label != label_before)++ loss = tl.where(num_labels == 0, 0.0, float("inf"))+ if num_frames > 0:+ # the paths start in the first blank or the first label+ alpha = tl.where(+ valid & (state < 2), _frame_log_probs(log_probs_ptr, start, num_classes, label, valid), float("-inf")+ )+ tl.store(alpha_ptr + start * num_states_max + state, alpha, mask=valid)+ for t in range(1, num_frames):+ row = start + t+ stay = alpha+ advance = tl.where(state >= 1, tl.gather(alpha, tl.maximum(state - 1, 0), 0), float("-inf"))+ skip = tl.where(can_skip, tl.gather(alpha, tl.maximum(state - 2, 0), 0), float("-inf"))+ alpha = _logsumexp3(stay, advance, skip) + _frame_log_probs(log_probs_ptr, row, num_classes, label, valid)+ alpha = tl.where(valid, alpha, float("-inf"))+ tl.store(alpha_ptr + row * num_states_max + state, alpha, mask=valid)+ # the paths end in the last blank or the last label+ last = tl.max(tl.where(state == num_states - 1, alpha, float("-inf")), 0)+ before_last = tl.max(tl.where(state == num_states - 2, alpha, float("-inf")), 0)+ loss = -_logsumexp3(last, before_last, float("-inf"))+ tl.store(loss_ptr + b, loss)+++@triton.jit+def _beta_grad_kernel(+ log_probs_ptr,+ cu_seqlens_ptr,+ lengths_ptr,+ targets_ptr,+ target_lengths_ptr,+ alpha_ptr,+ loss_ptr,+ grad_output_ptr,+ grad_ptr,+ num_classes,+ target_width,+ num_states_max,+ blank,+ ZERO_INFINITY: tl.constexpr,+ BLOCK_S: tl.constexpr,+ BLOCK_D: tl.constexpr,+):+ b = tl.program_id(0)+ start = tl.load(cu_seqlens_ptr + b).to(tl.int64)+ num_frames = tl.load(lengths_ptr + b).to(tl.int32)+ num_labels = tl.load(target_lengths_ptr + b).to(tl.int32)+ num_states = 2 * num_labels + 1+ state = tl.arange(0, BLOCK_S)+ valid = state < num_states+ label = _extended_labels(targets_ptr, b, target_width, state, valid, blank)+ label_after = tl.gather(label, tl.minimum(state + 2, BLOCK_S - 1), 0)+ can_skip = valid & (state + 2 < num_states) & (state % 2 == 1) & (label != label_after)+ is_label = valid & (state % 2 == 1)+ is_blank = valid & (state % 2 == 0)++ nll = tl.load(loss_ptr + b)+ scale = tl.load(grad_output_ptr + b)+ if ZERO_INFINITY:+ scale = tl.where(nll == float("inf"), 0.0, scale)+ classes = tl.arange(0, BLOCK_D)++ beta = tl.full([BLOCK_S], float("-inf"), tl.float32)+ for i in range(0, num_frames):+ t = num_frames - 1 - i+ row = start + t+ lp = _frame_log_probs(log_probs_ptr, row, num_classes, label, valid)+ if i == 0:+ beta = tl.where(valid & (state >= num_states - 2), lp, float("-inf"))+ else:+ stay = beta+ advance = tl.where(+ valid & (state + 1 < num_states), tl.gather(beta, tl.minimum(state + 1, BLOCK_S - 1), 0), float("-inf")+ )+ skip = tl.where(can_skip, tl.gather(beta, tl.minimum(state + 2, BLOCK_S - 1), 0), float("-inf"))+ beta = tl.where(valid, _logsumexp3(stay, advance, skip) + lp, float("-inf"))+ alpha = tl.load(alpha_ptr + row * num_states_max + state, mask=valid, other=float("-inf"))+ # the posterior of state s at frame t; nothing when the gradient is zero, so that inf - inf never enters+ posterior = tl.where(valid & (scale != 0.0), tl.exp(alpha + beta + nll - lp) * scale, 0.0)+ blank_posterior = tl.sum(tl.where(is_blank, posterior, 0.0), 0)+ for class_start in range(0, num_classes, BLOCK_D):+ c = class_start + classes+ class_valid = c < num_classes+ probs = tl.exp(tl.load(log_probs_ptr + row * num_classes + c, mask=class_valid, other=0.0).to(tl.float32))+ grad_row = probs * scale - tl.where(c == blank, blank_posterior, 0.0)+ tl.store(grad_ptr + row * num_classes + c, grad_row, mask=class_valid)+ # the label posteriors go into the stored row, so the stores must be complete first+ tl.debug_barrier()+ tl.atomic_add(grad_ptr + row * num_classes + label, -posterior, mask=is_label)tests/collections/asr/test_ctc_loss.py288 + / 0 −
@@ -16,8 +16,12 @@ import pytest import torch +from nemo.collections.asr.losses import ctc as ctc_module from nemo.collections.asr.losses.ctc import CTCLoss from nemo.collections.asr.parts.packed_sequence import pack_encoder_output+from nemo.core.utils.optional_libs import TRITON_AVAILABLE++CUDA_TRITON_AVAILABLE = TRITON_AVAILABLE and torch.cuda.is_available() class TestPackedCTCLoss:@@ -68,3 +72,287 @@ def test_rejects_padded_log_probs(self): criterion = CTCLoss(num_classes=2) with pytest.raises(TypeError, match="PackedEncoderActivations"): criterion.forward_packed(torch.zeros(1, 3, 3), torch.zeros(1, 1, dtype=torch.long), torch.tensor([1]))+++@pytest.mark.skipif(not CUDA_TRITON_AVAILABLE, reason="the packed CTC loss kernel needs CUDA and Triton")+class TestPackedCTCLossKernel:+ """On CUDA `forward_packed` runs the Triton kernel over the token-flat log probabilities:+ loss and gradient parity with the padded `forward`, the edge cases of the recursion, dtypes."""++ @staticmethod+ def _compare(criterion, logits, lengths, targets, target_lengths, *, atol=1e-5, rtol=1e-5, autocast=False):+ """Loss and gradient of `forward_packed` (the kernel) against `forward` (padded) on the same logits."""+ dense = logits.detach().clone().requires_grad_()+ flat = logits.detach().clone().requires_grad_()+ with torch.autocast("cuda", dtype=torch.bfloat16, enabled=autocast):+ expected = criterion(+ log_probs=dense.log_softmax(-1), targets=targets, input_lengths=lengths, target_lengths=target_lengths+ )+ packed = pack_encoder_output(flat.log_softmax(-1), lengths)+ actual = criterion.forward_packed(packed, targets, target_lengths)+ torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol)+ expected.sum().backward()+ actual.sum().backward()+ torch.testing.assert_close(flat.grad, dense.grad, atol=atol, rtol=rtol)+ return actual++ @staticmethod+ def _small_case(device):+ torch.manual_seed(7)+ lengths = torch.tensor([8, 5, 2], device=device)+ target_lengths = torch.tensor([3, 1, 0], device=device)+ targets = torch.tensor([[1, 1, 2], [2, 0, 0], [0, 0, 0]], device=device)+ logits = torch.randn(3, 11, 5, device=device)+ return logits, lengths, targets, target_lengths++ @pytest.mark.unit+ @pytest.mark.parametrize("reduction", ["none", "mean", "sum", "mean_batch", "mean_volume"])+ def test_loss_and_gradient_match_padded(self, reduction):+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ criterion = CTCLoss(num_classes=4, reduction=reduction)+ self._compare(criterion, logits, lengths, targets, target_lengths)++ @pytest.mark.unit+ @pytest.mark.parametrize("reduction", ["none", "mean", "sum", "mean_batch", "mean_volume"])+ @pytest.mark.parametrize("on_cpu", ["targets", "target_lengths", "both"])+ def test_cpu_target_metadata(self, reduction, on_cpu):+ # the targets and their lengths may stay on the CPU, as forward accepts them+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ if on_cpu != "target_lengths":+ targets = targets.cpu()+ if on_cpu != "targets":+ target_lengths = target_lengths.cpu()+ self._compare(CTCLoss(num_classes=4, reduction=reduction), logits, lengths, targets, target_lengths)++ @pytest.mark.unit+ def test_kernel_path_is_taken(self, monkeypatch):+ def no_unpack(*args, **kwargs):+ raise AssertionError("forward_packed formed the padded log probabilities on CUDA")++ monkeypatch.setattr(ctc_module, "unpack_encoder_output", no_unpack)+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ self._compare(CTCLoss(num_classes=4), logits, lengths, targets, target_lengths)++ @pytest.mark.unit+ def test_packed_ctc_loss_function(self):+ from nemo.collections.asr.parts.triton.ctc_loss import packed_ctc_loss++ logits, lengths, targets, target_lengths = self._small_case("cuda")+ log_probs = logits.log_softmax(-1)+ packed = pack_encoder_output(log_probs, lengths)+ losses = packed_ctc_loss(+ packed.data, packed.cu_seqlens, packed.lengths, targets, target_lengths, blank=4, zero_infinity=False+ )+ expected = torch.nn.functional.ctc_loss(+ log_probs.transpose(0, 1), targets, lengths, target_lengths, blank=4, reduction="none"+ )+ assert losses.shape == (3,) and losses.dtype == torch.float32+ torch.testing.assert_close(losses, expected, atol=1e-5, rtol=1e-5)++ @pytest.mark.unit+ def test_repeated_and_blank_labels(self):+ # a repeated label needs the blank between its two copies, and a blank id inside the targets is a label+ # that may not skip, as in the padded kernel+ torch.manual_seed(3)+ lengths = torch.tensor([9, 9, 7], device="cuda")+ targets = torch.tensor([[1, 1, 2, 2], [4, 4, 1, 0], [2, 4, 2, 0]], device="cuda")+ target_lengths = torch.tensor([4, 3, 3], device="cuda")+ logits = torch.randn(3, 9, 5, device="cuda")+ self._compare(CTCLoss(num_classes=4, reduction="none"), logits, lengths, targets, target_lengths)++ @pytest.mark.unit+ def test_empty_targets(self):+ torch.manual_seed(5)+ lengths = torch.tensor([6, 3], device="cuda")+ targets = torch.zeros(2, 0, dtype=torch.long, device="cuda")+ target_lengths = torch.tensor([0, 0], device="cuda")+ logits = torch.randn(2, 6, 3, device="cuda")+ self._compare(CTCLoss(num_classes=2, reduction="none"), logits, lengths, targets, target_lengths)++ @pytest.mark.unit+ @pytest.mark.parametrize("zero_infinity", [False, True])+ def test_zero_length_utterance(self, zero_infinity):+ # an utterance without frames: loss zero without targets, infinite (or zero) with targets, no gradient+ torch.manual_seed(11)+ lengths = torch.tensor([5, 0, 3], device="cuda")+ targets = torch.tensor([[1, 2], [1, 0], [2, 0]], device="cuda")+ target_lengths = torch.tensor([2, 1, 1], device="cuda")+ logits = torch.randn(3, 5, 4, device="cuda")+ criterion = CTCLoss(num_classes=3, reduction="none", zero_infinity=zero_infinity)+ losses = self._compare(criterion, logits, lengths, targets, target_lengths)+ assert losses[1].item() == (0.0 if zero_infinity else float("inf"))+ target_lengths[1] = 0+ losses = self._compare(criterion, logits, lengths, targets, target_lengths)+ assert losses[1].item() == 0.0++ @pytest.mark.unit+ @pytest.mark.parametrize("zero_infinity", [False, True])+ def test_impossible_alignment(self, zero_infinity):+ logits = torch.tensor([[[1.0, 0.5, -0.5]], [[0.2, 0.1, 0.0]]], device="cuda", requires_grad=True)+ lengths = torch.tensor([1, 1], device="cuda")+ targets = torch.tensor([[0, 0], [1, 0]], device="cuda")+ target_lengths = torch.tensor([2, 1], device="cuda")+ criterion = CTCLoss(num_classes=2, reduction="none", zero_infinity=zero_infinity)+ packed = pack_encoder_output(logits.log_softmax(-1), lengths)+ actual = criterion.forward_packed(packed, targets, target_lengths)+ if zero_infinity:+ assert actual[0].item() == 0.0+ actual.sum().backward()+ torch.testing.assert_close(logits.grad[0], torch.zeros_like(logits[0]))+ assert logits.grad[1].abs().sum() > 0+ else:+ assert torch.isposinf(actual[0])+ assert torch.isfinite(actual[1])++ @pytest.mark.unit+ def test_production_shape(self):+ # long utterances of different lengths, a vocabulary of 1,027 classes, one target per six frames+ torch.manual_seed(13)+ lengths = torch.tensor([667, 401, 123, 50, 599, 300], device="cuda")+ target_lengths = lengths // 6+ targets = torch.randint(0, 1026, (6, int(target_lengths.max())), device="cuda")+ logits = torch.randn(6, 667, 1027, device="cuda")+ # both kernels form the posteriors from alpha + beta + nll in float32, which at magnitude 1e3 carries 1e-4+ criterion = CTCLoss(num_classes=1026, reduction="mean_batch")+ self._compare(criterion, logits, lengths, targets, target_lengths, atol=1e-4, rtol=1e-4)++ @pytest.mark.unit+ def test_float32_under_autocast(self):+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ self._compare(CTCLoss(num_classes=4, reduction="sum"), logits, lengths, targets, target_lengths, autocast=True)++ @pytest.mark.unit+ def test_bfloat16_log_probs(self):+ # bfloat16 log probabilities are read as they are; the padded kernel takes them through autocast's float32 cast+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ dense = logits.to(torch.bfloat16).requires_grad_()+ flat = dense.detach().clone().requires_grad_()+ criterion = CTCLoss(num_classes=4, reduction="sum")+ with torch.autocast("cuda", dtype=torch.bfloat16):+ expected = criterion(+ log_probs=dense.log_softmax(-1), targets=targets, input_lengths=lengths, target_lengths=target_lengths+ )+ packed = pack_encoder_output(flat.log_softmax(-1), lengths)+ assert packed.data.dtype == torch.bfloat16+ actual = criterion.forward_packed(packed, targets, target_lengths)+ assert actual.dtype == torch.float32+ torch.testing.assert_close(actual, expected, atol=1e-4, rtol=1e-4)+ expected.backward()+ actual.backward()+ assert flat.grad.dtype == torch.bfloat16+ torch.testing.assert_close(flat.grad, dense.grad)++ @pytest.mark.unit+ def test_float64_takes_padded_path(self, monkeypatch):+ # the kernel computes in float32; float64 log probabilities keep the padded kernel and its precision+ def no_kernel(*args, **kwargs):+ raise AssertionError("forward_packed ran the kernel on float64 log probabilities")++ monkeypatch.setattr(ctc_module, "packed_ctc_loss", no_kernel)+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ criterion = CTCLoss(num_classes=4, reduction="none")+ losses = self._compare(criterion, logits.double(), lengths, targets, target_lengths, atol=1e-10, rtol=1e-10)+ assert losses.dtype == torch.float64++ @pytest.mark.unit+ def test_kernel_rejects_float64(self):+ from nemo.collections.asr.parts.triton.ctc_loss import packed_ctc_loss++ logits, lengths, targets, target_lengths = self._small_case("cuda")+ packed = pack_encoder_output(logits.double().log_softmax(-1), lengths)+ with pytest.raises(TypeError, match="float32"):+ packed_ctc_loss(packed.data, packed.cu_seqlens, packed.lengths, targets, target_lengths, blank=4)++ @pytest.mark.unit+ def test_over_allocated_targets_and_filler_utterances(self):+ # the static-shape regime of a captured step: the targets buffer is far wider than the longest target+ # and holds junk beyond each length, and the batch is padded with zero-length utterances at the end;+ # loss and gradient must not depend on any of it+ from nemo.collections.asr.parts.triton.ctc_loss import packed_ctc_loss++ torch.manual_seed(17)+ lengths = torch.tensor([9, 7, 4], device="cuda")+ target_lengths = torch.tensor([4, 3, 2], device="cuda")+ targets = torch.randint(0, 5, (3, 4), device="cuda")+ log_probs = torch.randn(20, 6, device="cuda").log_softmax(-1)+ cu_seqlens = torch.tensor([0, 9, 16, 20], dtype=torch.int32, device="cuda")+ tight = log_probs.detach().clone().requires_grad_()+ loss_tight = packed_ctc_loss(tight, cu_seqlens, lengths, targets, target_lengths, blank=5)+ loss_tight.sum().backward()++ wide_targets = torch.randint(0, 5, (5, 32), device="cuda")+ wide_targets[:3, :4] = targets+ wide_lengths = torch.cat([lengths, lengths.new_zeros(2)])+ wide_target_lengths = torch.cat([target_lengths, target_lengths.new_zeros(2)])+ wide_cu_seqlens = torch.tensor([0, 9, 16, 20, 20, 20], dtype=torch.int32, device="cuda")+ wide = log_probs.detach().clone().requires_grad_()+ loss_wide = packed_ctc_loss(wide, wide_cu_seqlens, wide_lengths, wide_targets, wide_target_lengths, blank=5)+ loss_wide.sum().backward()+ torch.testing.assert_close(loss_wide[:3], loss_tight)+ assert loss_wide[3:].abs().sum().item() == 0.0+ torch.testing.assert_close(wide.grad, tight.grad)++ @pytest.mark.unit+ def test_cuda_graph_replay(self):+ # under capture no host read is possible, so the width of the targets buffer sizes the states;+ # replays with other targets, lengths and junk beyond the lengths in the same buffers+ # must give the eager loss and gradient, which the longest target of each batch sizes+ from nemo.collections.asr.parts.triton.ctc_loss import packed_ctc_loss++ torch.manual_seed(19)+ lengths = torch.tensor([70, 7, 4, 0], device="cuda")+ cu_seqlens = torch.tensor([0, 70, 77, 81, 81], dtype=torch.int32, device="cuda")+ log_probs = torch.randn(81, 6, device="cuda").log_softmax(-1).requires_grad_()+ targets = torch.randint(0, 5, (4, 32), device="cuda")+ target_lengths = torch.tensor([4, 3, 2, 0], device="cuda")++ def reference(targets, target_lengths):+ leaf = log_probs.detach().clone().requires_grad_()+ loss = packed_ctc_loss(leaf, cu_seqlens, lengths, targets, target_lengths, blank=5)+ loss.sum().backward()+ return loss.detach(), leaf.grad++ # the warm-up compiles the specialization of the capture, where the width stands in for the longest target+ stream = torch.cuda.Stream()+ stream.wait_stream(torch.cuda.current_stream())+ with torch.cuda.stream(stream):+ full_width = torch.tensor([32, 0, 0, 0], device="cuda")+ packed_ctc_loss(log_probs, cu_seqlens, lengths, targets, full_width, blank=5).sum().backward()+ torch.cuda.current_stream().wait_stream(stream)+ log_probs.grad = None+ graph = torch.cuda.CUDAGraph()+ with torch.cuda.graph(graph):+ static_loss = packed_ctc_loss(log_probs, cu_seqlens, lengths, targets, target_lengths, blank=5)+ static_loss.sum().backward()+ static_grad = log_probs.grad++ batches = [+ (targets.clone(), target_lengths.clone()),+ (torch.randint(0, 5, (4, 32), device="cuda"), torch.tensor([32, 2, 1, 0], device="cuda")),+ (torch.randint(0, 5, (4, 32), device="cuda"), torch.tensor([0, 1, 2, 0], device="cuda")),+ ]+ for new_targets, new_target_lengths in batches:+ targets.copy_(new_targets)+ target_lengths.copy_(new_target_lengths)+ graph.replay()+ expected_loss, expected_grad = reference(new_targets, new_target_lengths)+ torch.testing.assert_close(static_loss, expected_loss, atol=1e-5, rtol=1e-5)+ torch.testing.assert_close(static_grad, expected_grad, atol=1e-5, rtol=1e-5)+ assert static_loss[3].item() == 0.0++ @pytest.mark.unit+ def test_rejects_target_lengths_beyond_width(self):+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ packed = pack_encoder_output(logits.log_softmax(-1), lengths)+ with pytest.raises(ValueError, match="targets width"):+ CTCLoss(num_classes=4).forward_packed(packed, targets, target_lengths + 1)++ @pytest.mark.unit+ @pytest.mark.parametrize("zero_infinity", [False, True])+ def test_rejects_negative_target_lengths(self, zero_infinity):+ # a negative length is malformed input, not an impossible alignment that zero_infinity may zero+ logits, lengths, targets, target_lengths = self._small_case("cuda")+ packed = pack_encoder_output(logits.log_softmax(-1), lengths)+ target_lengths[0] = -1+ with pytest.raises(ValueError, match="negative"):+ CTCLoss(num_classes=4, zero_infinity=zero_infinity).forward_packed(packed, targets, target_lengths)