diff --git a/src/transformers/integrations/hub_kernels.py b/src/transformers/integrations/hub_kernels.py index cd035c57685d..c504fe73524b 100644 --- a/src/transformers/integrations/hub_kernels.py +++ b/src/transformers/integrations/hub_kernels.py @@ -709,6 +709,13 @@ def _build_kernel_mapping() -> dict: ), }, }, + "tdt_loss": { + "cuda": { + Mode.TRAINING | Mode.TORCH_COMPILE: LayerRepository( + repo_id="kernels-community/tdt-loss", layer_name="TDTLoss", version=1 + ), + }, + }, } return _KERNEL_MAPPING diff --git a/src/transformers/loss/loss_tdt.py b/src/transformers/loss/loss_tdt.py index a686afbb374a..b22a2aaef6a8 100644 --- a/src/transformers/loss/loss_tdt.py +++ b/src/transformers/loss/loss_tdt.py @@ -14,12 +14,14 @@ import torch +from ..integrations import use_kernel_forward_from_hub from ..utils import logging logger = logging.get_logger(__name__) +@use_kernel_forward_from_hub("tdt_loss") def tdt_loss( token_logits: torch.Tensor, duration_logits: torch.Tensor, diff --git a/src/transformers/models/parakeet/modeling_parakeet.py b/src/transformers/models/parakeet/modeling_parakeet.py index 190f09a7a49f..1bab1f2d0419 100644 --- a/src/transformers/models/parakeet/modeling_parakeet.py +++ b/src/transformers/models/parakeet/modeling_parakeet.py @@ -29,6 +29,7 @@ from ...activations import ACT2FN from ...generation import CompileConfig, GenerationMixin, GenerationMode from ...integrations import use_kernel_forward_from_hub, use_kernelized_func +from ...loss.loss_tdt import tdt_loss from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, CausalLMOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -1051,6 +1052,7 @@ def __init__(self, config: ParakeetTDTConfig): Parakeet Encoder with a TDT (Token Duration Transducer) head. """ ) +@use_kernelized_func(tdt_loss) class ParakeetForTDT(ParakeetTDTGenerationMixin, ParakeetForRNNT): config: ParakeetTDTConfig diff --git a/src/transformers/models/parakeet/modular_parakeet.py b/src/transformers/models/parakeet/modular_parakeet.py index 8d92c24a504c..7f046e25a060 100644 --- a/src/transformers/models/parakeet/modular_parakeet.py +++ b/src/transformers/models/parakeet/modular_parakeet.py @@ -23,6 +23,8 @@ from ... import initialization as init from ...activations import ACT2FN from ...generation import CompileConfig, GenerationMixin, GenerationMode +from ...integrations import use_kernelized_func +from ...loss.loss_tdt import tdt_loss from ...modeling_layers import GradientCheckpointingLayer from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, CausalLMOutput from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel @@ -900,6 +902,7 @@ def __init__(self, config: ParakeetTDTConfig): Parakeet Encoder with a TDT (Token Duration Transducer) head. """ ) +@use_kernelized_func(tdt_loss) class ParakeetForTDT(ParakeetTDTGenerationMixin, ParakeetForRNNT): config: ParakeetTDTConfig diff --git a/tests/models/parakeet/test_modeling_parakeet.py b/tests/models/parakeet/test_modeling_parakeet.py index b8696ecec7be..eab0cdffc618 100644 --- a/tests/models/parakeet/test_modeling_parakeet.py +++ b/tests/models/parakeet/test_modeling_parakeet.py @@ -19,7 +19,16 @@ from pathlib import Path from transformers import is_datasets_available, is_torch_available -from transformers.testing_utils import cleanup, require_torch, require_torchaudio, slow, torch_device +from transformers.testing_utils import ( + cleanup, + preserve_module_forwards, + require_kernels, + require_torch, + require_torch_gpu, + require_torchaudio, + slow, + torch_device, +) from ...test_configuration_common import ConfigTester from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask @@ -608,6 +617,41 @@ def create_and_check_model(self, config, inputs_dict): (self.batch_size, self.output_seq_length, self.encoder_model_tester.hidden_size), ) + def prepare_config_and_inputs_for_loss(self): + config, input_features, attention_mask = self.prepare_config_and_inputs() + labels = ids_tensor([self.batch_size, 6], self.blank_token_id) # no blank in the labels + labels[labels == self.pad_token_id] = self.pad_token_id + 1 # nor padding + blank = torch.full((self.batch_size, 1), self.blank_token_id, device=labels.device) + inputs_dict = { + "input_features": input_features, + "attention_mask": attention_mask, + "decoder_input_ids": torch.cat([blank, labels], dim=1), + "labels": labels, + } + return config, inputs_dict + + def create_and_check_use_kernels(self, config, inputs_dict): + """`use_kernels=True` swaps `tdt_loss` for the `kernels-community/tdt-loss` kernel, with the same loss.""" + model = ParakeetForTDT(config).to(torch_device).eval() + with torch.no_grad(): + expected = model(**inputs_dict).loss + + with preserve_module_forwards(model): + model.set_use_kernels(True) + self.parent.assertIn("forward", vars(tdt_loss), "`tdt_loss` was not kernelized") + with torch.no_grad(): + loss = model(**inputs_dict).loss + torch.testing.assert_close(loss, expected, rtol=1e-4, atol=1e-4) + self.parent.assertNotIn("forward", vars(tdt_loss), "`tdt_loss` was not restored") + + # Training mode also uses the kernel, which has a backward + with preserve_module_forwards(model): + model.train() + model.set_use_kernels(True) + self.parent.assertIn("forward", vars(tdt_loss), "`tdt_loss` was not kernelized in training mode") + model(**inputs_dict).loss.backward() + self.parent.assertTrue(any(p.grad is not None for p in model.parameters()), "No gradients after backward") + def prepare_config_and_inputs_for_common(self): config, input_features, attention_mask = self.prepare_config_and_inputs() decoder_input_ids = ids_tensor([self.batch_size, 1], self.vocab_size) @@ -650,6 +694,12 @@ def test_model(self): config_and_inputs = self.model_tester.prepare_config_and_inputs_for_common() self.model_tester.create_and_check_model(*config_and_inputs) + @require_torch_gpu + @require_kernels + def test_use_kernels(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs_for_loss() + self.model_tester.create_and_check_use_kernels(*config_and_inputs) + @unittest.skip(reason="ParakeetForTDT does not use inputs_embeds") def test_model_get_set_embeddings(self): pass