fix(asr): warm up Numba RNNT and TDT losses before training

@@ -739,6 +739,14 @@ The loss config is based on a resolver pattern and can be used as follows:     warprnnt_numba_kwargs:       fastemit_lambda: 0.0 +Numba Transducer Warmup+^^^^^^^^^^^^^^^^^^^^^^^++RNNT models automatically warm up Numba loss kernels on CUDA before training, including supported FP16 RNNT kernels and both FP32 branches for TDT models.+This reduces memory retained during compilation without changing random states.++Custom training loops can call ``RNNTLoss.warmup(device)`` before allocating batch activations.+ FastEmit Regularization ^^^^^^^^^^^^^^^^^^^^^^^ 

@@ -31,7 +31,7 @@ import inspect import operator from dataclasses import dataclass-from typing import Any, Callable, Dict, List, Optional, Set+from typing import Any, Callable, Dict, List, Optional, Set, Union  import torch from omegaconf import DictConfig, OmegaConf@@ -419,6 +419,22 @@ def __init__(self, num_classes, reduction: str = 'mean_batch', loss_name: str =         self._force_float32 = RNNT_LOSS_RESOLVER[loss_name].force_float32         self._fp16_compat_checked = False +    def warmup(self, device: Union[str, torch.device]) -> bool:+        """Warm supported Numba RNNT/TDT kernels before allocating training activations.++        Returns whether warmup ran; False for CPU devices or unsupported backends.+        """+        if not NUMBA_RNNT_AVAILABLE:+            return False+        if isinstance(self._loss, RNNTLossNumba):+            dtypes = [torch.float32]+            if not self._force_float32 and numba_utils.is_numba_cuda_fp16_supported():+                dtypes.append(torch.float16)+            return self._loss.warmup(device, dtypes=dtypes)+        if isinstance(self._loss, TDTLossNumba):+            return self._loss.warmup(device)+        return False+     def reduce(self, losses, target_lengths):          if isinstance(losses, List):

@@ -190,6 +190,11 @@ def setup_optim_normalization(self):         # Setup normalized joint norm for model         self._optim_normalize_joint_norm = self.cfg.get('normalize_joint_norm', False) +    def on_train_start(self) -> None:+        """Warm loss kernels before the first training batch."""+        super().on_train_start()+        self.loss.warmup(self.device)+     def extract_rnnt_loss_cfg(self, cfg: Optional[DictConfig]):         """         Helper method to extract the rnnt loss name, and potentially its kwargs

@@ -28,6 +28,11 @@ # limitations under the License.  +import gc+import random+from contextlib import contextmanager+from typing import Iterator, Sequence, Union+ import torch from torch.autograd import Function from torch.nn import Module@@ -411,6 +416,32 @@ def __init__(self, blank=0, reduction='mean', fastemit_lambda: float = 0.0, clam         self.reduction = reduction         self.loss = _RNNTNumba.apply +    def warmup(self, device: Union[str, torch.device], dtypes: Sequence[torch.dtype] = (torch.float32,)) -> bool:+        """Warm contiguous CUDA RNNT kernels for each dtype before allocating training activations.++        Preserves random states and the configured loss; returns False on CPU.+        """+        device = torch.device(device)+        if device.type != 'cuda':+            return False++        def warmup(dtype):+            # Non-singleton dimensions preserve the contiguous production array layout.+            acts = torch.zeros((2, 8, 4, max(2, self.blank + 1)), device=device, dtype=dtype, requires_grad=True)+            labels = torch.full((2, 3), 1 if self.blank == 0 else 0, device=device, dtype=torch.int64)+            input_lengths = torch.full((2,), 8, device=device, dtype=torch.int64)+            label_lengths = torch.full((2,), 3, device=device, dtype=torch.int64)+            value = self(acts, labels, input_lengths, label_lengths)+            value.sum().backward()+            if not torch.isfinite(value).all() or not torch.isfinite(acts.grad).all():+                raise RuntimeError('RNNT loss warmup produced non-finite loss or gradients')++        # Return from the local function before collecting compiler cycles and temporary tensors.+        with _numba_loss_warmup(device):+            for dtype in dtypes:+                warmup(dtype)+        return True+     def forward(self, acts, labels, act_lens, label_lens):         """         log_probs: Tensor of (batch x seqLength x labelLength x outputDim) containing output from network@@ -552,6 +583,51 @@ def __init__(         self.sigma = sigma         self.omega = omega +    def warmup(self, device: Union[str, torch.device]) -> bool:+        """Warm both TDT and RNNT CUDA branches with float32 contiguous inputs before allocating training activations.++        Preserves random states and the configured loss; returns False on CPU.+        """+        device = torch.device(device)+        if device.type != 'cuda':+            return False++        def warmup():+            # Non-singleton dimensions preserve the contiguous production array layout.+            labels_count = 3+            frames = max(8, (labels_count + 1) * max(self.durations))+            vocabulary_size = max(2, self.blank + 1)+            label = 1 if self.blank == 0 else 0+            for omega in (-1.0, 2.0):+                # Force each branch even when random.uniform returns an endpoint.+                loss = TDTLossNumba(+                    blank=self.blank,+                    durations=list(self.durations),+                    reduction='mean',+                    fastemit_lambda=self.fastemit_lambda,+                    clamp=self.clamp,+                    sigma=self.sigma,+                    omega=omega,+                )+                acts = torch.zeros(+                    (2, frames, labels_count + 1, vocabulary_size + len(self.durations)),+                    device=device,+                    dtype=torch.float32,+                    requires_grad=True,+                )+                labels = torch.full((2, labels_count), label, device=device, dtype=torch.int64)+                input_lengths = torch.full((2,), frames, device=device, dtype=torch.int64)+                label_lengths = torch.full((2,), labels_count, device=device, dtype=torch.int64)+                value = loss(acts, labels, input_lengths, label_lengths)+                value.backward()+                if not torch.isfinite(value).all() or not torch.isfinite(acts.grad).all():+                    raise RuntimeError('TDT loss warmup produced non-finite loss or gradients')++        # Return from the local function before collecting compiler cycles and temporary tensors.+        with _numba_loss_warmup(device):+            warmup()+        return True+     def forward(self, acts, labels, act_lens, label_lens):         """         log_probs: Tensor of (batch x seqLength x labelLength x outputDim) containing output from network@@ -633,3 +709,21 @@ def certify_inputs(log_probs, labels, lengths, label_lengths):         raise ValueError(f"Input length mismatch! Given T: {T}, Expected max T from input lengths: {max_T}")     if U != max_U + 1:         raise ValueError(f"Output length mismatch! Given U: {U}, Expected max U from target lengths: {max_U} + 1")+++@contextmanager+def _numba_loss_warmup(device: torch.device) -> Iterator[None]:+    device_index = device.index if device.index is not None else torch.cuda.current_device()+    python_rng = random.getstate()+    try:+        with (+            torch.random.fork_rng(devices=[device_index]),+            torch.inference_mode(False),+            torch.enable_grad(),+            torch.autocast('cuda', enabled=False),+        ):+            yield+    finally:+        random.setstate(python_rng)+    gc.collect()+    torch.cuda.synchronize(device)

@@ -13,13 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. +import gc import random+import weakref  import numpy as np import pytest import torch -from nemo.collections.asr.losses.rnnt import MultiblankRNNTLossPytorch, RNNTLossPytorch, TDTLossPytorch+from nemo.collections.asr.losses.rnnt import MultiblankRNNTLossPytorch, RNNTLoss, RNNTLossPytorch, TDTLossPytorch from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_numpy import RNNTLoss as RNNTLoss_Numpy from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import (     MultiblankRNNTLossNumba,@@ -76,7 +78,110 @@ def wrap_and_call(fn, acts, labels, device):     return costs.data.cpu().numpy(), grad  +@pytest.fixture+def mock_cuda(monkeypatch):+    original_fork_rng = torch.random.fork_rng+    monkeypatch.setattr(torch.random, 'fork_rng', lambda devices: original_fork_rng(devices=[]))+    monkeypatch.setattr(torch.cuda, 'synchronize', lambda device: None)+    original_zeros = torch.zeros+    original_full = torch.full++    def zeros(*args, **kwargs):+        kwargs['device'] = 'cpu'+        return original_zeros(*args, **kwargs)++    def full(*args, **kwargs):+        kwargs['device'] = 'cpu'+        return original_full(*args, **kwargs)++    monkeypatch.setattr(torch, 'zeros', zeros)+    monkeypatch.setattr(torch, 'full', full)++ class TestRNNTLossPytorch:++    @pytest.mark.unit+    def test_warmup_cpu_is_noop(self):+        assert RNNTLossNumba(blank=3).warmup('cpu') is False++    @pytest.mark.unit+    @pytest.mark.parametrize('blank', [0, 3])+    @pytest.mark.parametrize('reduction', ['none', 'sum', 'mean'])+    @pytest.mark.parametrize('dtypes', [(torch.float32,), (torch.float32, torch.float16)])+    @pytest.mark.usefixtures('mock_cuda')+    def test_warmup_forward_backward(self, monkeypatch, blank, reduction, dtypes):+        source = RNNTLossNumba(blank=blank, reduction=reduction, fastemit_lambda=0.1, clamp=1.0)+        backwards = []+        forwards = []++        def forward(loss, acts, labels, input_lengths, label_lengths):+            assert loss is source+            assert loss.blank == source.blank+            assert loss.fastemit_lambda == source.fastemit_lambda+            assert loss.clamp == source.clamp+            assert loss.reduction == reduction+            assert torch.all(labels != blank)+            assert labels.max() < acts.shape[-1]+            assert acts.is_contiguous()+            forwards.append(acts.dtype)+            assert acts.shape[1] == input_lengths.max()+            assert acts.shape[2] == label_lengths.max() + 1+            acts.register_hook(lambda grad: backwards.append(torch.isfinite(grad).all().item()))+            values = acts.sum(dim=(1, 2, 3))+            if reduction == 'sum':+                return values.sum()+            if reduction == 'mean':+                return values.mean()+            return values++        monkeypatch.setattr(RNNTLossNumba, 'forward', forward)+        assert source.warmup('cuda:0', dtypes=dtypes)+        assert forwards == list(dtypes)+        assert backwards == [True] * len(dtypes)+        assert source.reduction == reduction++    @pytest.mark.unit+    @pytest.mark.parametrize('blank', [0, 3])+    @pytest.mark.parametrize('dtypes', [(torch.float32,), (torch.float32, torch.float16)])+    def test_warmup_cuda(self, blank, dtypes):+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)+        if not torch.cuda.is_available():+            pytest.skip('CUDA is required')+        if torch.float16 in dtypes and not numba_utils.is_numba_cuda_fp16_supported():+            pytest.skip('Numba FP16 support is required')+        device = torch.device('cuda', torch.cuda.current_device())+        source = RNNTLossNumba(blank=blank, reduction='sum', fastemit_lambda=0.1, clamp=1.0)+        python_rng = random.getstate()+        cpu_rng = torch.get_rng_state().clone()+        cuda_rng = torch.cuda.get_rng_state(device).clone()+        assert source.warmup(device, dtypes=dtypes)+        assert random.getstate() == python_rng+        assert torch.equal(torch.get_rng_state(), cpu_rng)+        assert torch.equal(torch.cuda.get_rng_state(device), cuda_rng)+        assert source.reduction == 'sum'+        assert source.fastemit_lambda == 0.1+        assert source.clamp == 1.0++    @pytest.mark.unit+    @pytest.mark.parametrize('force_float32', [False, True])+    @pytest.mark.parametrize('fp16_supported', [False, True])+    def test_warmup_public_dtypes(self, monkeypatch, force_float32, fp16_supported):+        source = RNNTLoss(num_classes=3, loss_name='warprnnt_numba')+        source._force_float32 = force_float32+        monkeypatch.setattr(numba_utils, 'is_numba_cuda_fp16_supported', lambda: fp16_supported)+        calls = []++        def warmup(loss, device, dtypes=(torch.float32,)):+            calls.append((device, list(dtypes)))+            return True++        monkeypatch.setattr(RNNTLossNumba, 'warmup', warmup)+        expected = [torch.float32]+        if not force_float32 and fp16_supported:+            expected.append(torch.float16)+        assert source.warmup('cuda:0')+        assert calls == [('cuda:0', expected)]+     @pytest.mark.unit     @pytest.mark.parametrize('device', DEVICES)     @pytest.mark.parametrize('dtype', DTYPES)@@ -543,6 +648,57 @@ def test_case_randomized_act_label(self, device):   class TestTDTLoss:++    @pytest.mark.unit+    def test_warmup_cpu_is_noop(self):+        loss = TDTLossNumba(blank=3, durations=[0, 1, 2])+        assert loss.warmup('cpu') is False++    @pytest.mark.unit+    @pytest.mark.parametrize('blank', [0, 3])+    @pytest.mark.usefixtures('mock_cuda')+    def test_warmup_exercises_both_branches(self, monkeypatch, blank):+        source = TDTLossNumba(blank=blank, durations=[0, 1, 2], reduction='sum', omega=0.1, sigma=0.05)+        branches = []+        backwards = []++        def forward(loss, acts, labels, input_lengths, label_lengths):+            branches.append(loss.omega)+            assert loss is not source+            assert loss.blank == source.blank+            assert loss.durations == source.durations+            assert loss.sigma == source.sigma+            assert torch.all(labels != blank)+            assert labels.max() < acts.shape[-1] - len(source.durations)+            assert acts.is_contiguous() and acts.dtype == torch.float32+            acts.register_hook(lambda grad: backwards.append(torch.isfinite(grad).all().item()))+            return acts.sum()++        monkeypatch.setattr(TDTLossNumba, 'forward', forward)+        assert source.warmup('cuda:0')+        assert branches == [-1.0, 2.0]+        assert backwards == [True, True]+        assert source.omega == 0.1+        assert source.reduction == 'sum'++    @pytest.mark.unit+    @pytest.mark.parametrize('blank', [0, 3])+    def test_warmup_cuda(self, blank):+        numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__)+        if not torch.cuda.is_available():+            pytest.skip('CUDA is required')+        device = torch.device('cuda', torch.cuda.current_device())+        source = TDTLossNumba(blank=blank, durations=[0, 1, 2], reduction='sum', omega=0.1)+        python_rng = random.getstate()+        cpu_rng = torch.get_rng_state().clone()+        cuda_rng = torch.cuda.get_rng_state(device).clone()+        assert source.warmup(device)+        assert random.getstate() == python_rng+        assert torch.equal(torch.get_rng_state(), cpu_rng)+        assert torch.equal(torch.cuda.get_rng_state(device), cuda_rng)+        assert source.omega == 0.1+        assert source.reduction == 'sum'+     @pytest.mark.unit     @pytest.mark.parametrize('device', CUDA_ONLY_DEVICE)     def test_case_randomized_act_label(self, device):@@ -626,5 +782,75 @@ def test_case_fixed_case_act_label(self, device):         assert np.allclose(pt_grads, expected_grads, rtol=1e-2), "td gradient mismatch."  +@pytest.mark.usefixtures('mock_cuda')+class TestNumbaLossWarmup:++    @pytest.mark.unit+    @pytest.mark.parametrize('backend', ['tdt', 'rnnt'])+    def test_warmup_restores_rng_and_collects_cycles(self, monkeypatch, backend):+        references = []+        synchronized = []++        def forward(loss, acts, labels, input_lengths, label_lengths):+            assert torch.is_grad_enabled()+            assert not torch.is_inference_mode_enabled()+            random.random()+            torch.rand(2)+            references.append(weakref.ref(acts))+            # Model a caught compilation error retaining its frame and the caller's tensors.+            failures = []+            try:+                raise RuntimeError('candidate signature failed')+            except RuntimeError as exc:+                failures.append(exc)+            return acts.sum()++        def synchronize(device):+            assert references+            assert all(reference() is None for reference in references)+            synchronized.append(device)++        monkeypatch.setattr(torch.cuda, 'synchronize', synchronize)+        loss_type = TDTLossNumba if backend == 'tdt' else RNNTLossNumba+        monkeypatch.setattr(loss_type, 'forward', forward)+        source = TDTLossNumba(blank=3, durations=[0, 1, 2]) if backend == 'tdt' else RNNTLossNumba(blank=3)+        python_rng = random.getstate()+        torch_rng = torch.get_rng_state().clone()+        gc_enabled = gc.isenabled()+        gc.disable()+        try:+            with torch.inference_mode():+                if backend == 'rnnt':+                    assert source.warmup('cuda:0', dtypes=(torch.float32, torch.float16))+                else:+                    assert source.warmup('cuda:0')+                assert torch.is_inference_mode_enabled()+            assert len(references) == 2+            assert synchronized == [torch.device('cuda:0')]+            assert random.getstate() == python_rng+            assert torch.equal(torch.get_rng_state(), torch_rng)+        finally:+            if gc_enabled:+                gc.enable()++    @pytest.mark.unit+    @pytest.mark.parametrize('backend', ['tdt', 'rnnt'])+    def test_warmup_failure_restores_rng(self, monkeypatch, backend):+        def fail(loss, acts, labels, input_lengths, label_lengths):+            random.random()+            torch.rand(2)+            raise RuntimeError('compilation failed')++        loss_type = TDTLossNumba if backend == 'tdt' else RNNTLossNumba+        monkeypatch.setattr(loss_type, 'forward', fail)+        source = TDTLossNumba(blank=3, durations=[0, 1, 2]) if backend == 'tdt' else RNNTLossNumba(blank=3)+        python_rng = random.getstate()+        torch_rng = torch.get_rng_state().clone()+        with pytest.raises(RuntimeError, match='compilation failed'):+            source.warmup('cuda:0')+        assert random.getstate() == python_rng+        assert torch.equal(torch.get_rng_state(), torch_rng)++ if __name__ == "__main__":     pytest.main([__file__])

@@ -14,6 +14,7 @@ # limitations under the License. import copy from typing import Any, Dict, List, Optional, Tuple+from unittest.mock import Mock  import pytest import torch@@ -253,6 +254,15 @@ def asr_model():   class TestEncDecRNNTModel:++    @pytest.mark.unit+    def test_loss_startup_warmup(self, asr_model, monkeypatch):+        warmup = Mock()+        monkeypatch.setattr(asr_model.loss, 'warmup', warmup)+        asr_model.on_train_start()+        assert hasattr(asr_model, '_freeze_cfg')+        warmup.assert_called_once_with(asr_model.device)+     @pytest.mark.skipif(         not NUMBA_RNNT_LOSS_AVAILABLE,         reason='RNNTLoss has not been compiled with appropriate numba version.',