[TTS][EasyMagpietts] Added CFG distillation

Author: artem-gorodetskii

.github/workflows/cicd-main-speech.yml

@@ -402,6 +402,9 @@ jobs:
           - runner: ${{ inputs.runner }}
             script: L2_TTS_Fast_dev_runs_EasyMagpietts_OnlinePO
             timeout: 20
+          - runner: ${{ inputs.runner }}
+            script: L2_TTS_Fast_dev_runs_EasyMagpietts_OnlineCFGDistillation
+            timeout: 20
           - runner: ${{ inputs.runner }}
             script: L2_TTS_InferEvaluate_EasyMagpietts
             timeout: 20

examples/tts/easy_magpietts.py

@@ -16,11 +16,17 @@
 import torch.multiprocessing as mp
 from omegaconf import OmegaConf, open_dict
 
-from nemo.collections.tts.models import EasyMagpieTTSModel, EasyMagpieTTSModelOnlinePO
+from nemo.collections.tts.models import EasyMagpieCFGDistillation, EasyMagpieTTSModel, EasyMagpieTTSModelOnlinePO
 from nemo.core.config import hydra_runner
 from nemo.utils import logging
 from nemo.utils.exp_manager import exp_manager
 
+_TRAIN_MODES: list[str] = [
+    "train",
+    "online_cfg_distillation_train",
+    "onlinepo_train",
+]
+
 
 @hydra_runner(config_path="conf/magpietts", config_name="easy_magpietts")
 def main(cfg):
@@ -43,8 +49,12 @@ def main(cfg):
     exp_manager(trainer, cfg.get("exp_manager", None))
 
     mode = cfg.get('mode', 'train')
+    train_modes_msg = ", ".join(_TRAIN_MODES)
+
     if mode == 'train':
         model = EasyMagpieTTSModel(cfg=cfg.model, trainer=trainer)
+    elif mode == "online_cfg_distillation_train":
+        model = EasyMagpieCFGDistillation(cfg=cfg.model, trainer=trainer)
     elif mode == 'onlinepo_train':
         model_cfg = cfg.model
         with open_dict(model_cfg):
@@ -53,14 +63,14 @@ def main(cfg):
     elif mode == 'test':
         model = EasyMagpieTTSModel(cfg=cfg.model, trainer=trainer)
     else:
-        raise NotImplementedError(f"Only train, onlinepo_train and test modes are supported. Got {mode}")
+        raise NotImplementedError(f"Only {train_modes_msg} and test modes are supported. Got {mode}.")
 
     if cfg.get("pretrained_model", None):
         model.restore_from_pretrained_checkpoint(cfg.pretrained_model)
 
     model.maybe_init_from_pretrained_checkpoint(cfg=cfg)
 
-    if mode in ['train', 'onlinepo_train']:
+    if mode in _TRAIN_MODES:
         trainer.fit(model)
     elif mode == 'test':
         trainer.test(model)

nemo/collections/tts/losses/magpietts_cfg_distillation.py

@@ -15,7 +15,7 @@
 Losses used in CFG distillation of the MagpieTTS model.
 """
 
-from typing import Generator, Optional
+from typing import Callable, Generator, Optional
 
 import torch
 from torch import Tensor, nn
@@ -29,8 +29,13 @@
     "NRMSELogitsLoss",
 ]
 
+_CODEBOOK_ORDERING_MODES: set[str] = {
+    "frame-major",
+    "codebook-major",
+}
 
-def _iter_slices(
+
+def _iter_slices_frame_major(
     num_codebooks: int,
     num_tokens_per_codebook: int,
     frame_stacking_factor: int,
@@ -48,6 +53,33 @@ def _iter_slices(
             yield fs_index, codebook, start, end, slice_mask, slice_len
 
 
+def _iter_slices_codebook_major(
+    num_codebooks: int,
+    num_tokens_per_codebook: int,
+    frame_stacking_factor: int,
+    mask: Tensor,
+) -> Generator[tuple[int, int, int, int, Tensor, Tensor], None, None]:
+    for codebook in range(num_codebooks):
+        for fs_index in range(frame_stacking_factor):
+            slice_mask = mask[:, fs_index::frame_stacking_factor].float()
+            slice_len = slice_mask.sum(dim=-1).clamp_min(1)
+
+            channel = codebook * frame_stacking_factor + fs_index
+            start = channel * num_tokens_per_codebook
+            end = start + num_tokens_per_codebook
+
+            yield fs_index, codebook, start, end, slice_mask, slice_len
+
+
+def _get_slice_iterator(mode: str) -> Callable:
+    if mode not in _CODEBOOK_ORDERING_MODES:
+        raise ValueError(
+            f"Unsupported codebook ordering {mode!r}; expected one of {sorted(_CODEBOOK_ORDERING_MODES)}."
+        )
+
+    return _iter_slices_frame_major if mode == "frame-major" else _iter_slices_codebook_major
+
+
 class KLDivergenceLoss(Loss):
     """The Kullback-Leibler divergence loss."""
 
@@ -79,11 +111,13 @@ def __init__(
         num_codebooks: int,
         num_tokens_per_codebook: int,
         frame_stacking_factor: int,
+        codebook_ordering: str = "frame-major",
     ) -> None:
         super().__init__()
         self.num_codebooks = num_codebooks
         self.num_tokens_per_codebook = num_tokens_per_codebook
         self.frame_stacking_factor = frame_stacking_factor
+        self.iter_slices: Callable = _get_slice_iterator(mode=codebook_ordering)
         self.criterion = nn.KLDivLoss(reduction="none", log_target=False)
 
     @typecheck()
@@ -111,17 +145,19 @@ def forward(
             Tensor: Scalar tensor representing the averaged masked KL divergence loss.
         """
         loss = 0.0
-        student_log_probs = student_logits.log_softmax(dim=-1)
-        teacher_probs = teacher_logits.softmax(dim=-1)
 
-        for _, _, start, end, slice_mask, slice_len in _iter_slices(
+        for _, _, start, end, slice_mask, slice_len in self.iter_slices(
             self.num_codebooks,
             self.num_tokens_per_codebook,
             self.frame_stacking_factor,
             mask,
         ):
-            teacher_probs_slice = teacher_probs[:, :, start:end]
-            student_log_probs_slice = student_log_probs[:, :, start:end]
+            # Normalize within this head only. Normalizing over the full
+            # concatenated dimension would incorrectly make
+            # independent codebook heads compete for probability mass.
+            student_log_probs_slice = student_logits[:, :, start:end].log_softmax(dim=-1)
+            teacher_probs_slice = teacher_logits[:, :, start:end].softmax(dim=-1)
+
             slice_loss = self.criterion(input=student_log_probs_slice, target=teacher_probs_slice)
             slice_loss = slice_loss.sum(dim=-1)
             slice_loss = (slice_loss * slice_mask).sum(dim=-1) / slice_len
@@ -166,11 +202,13 @@ def __init__(
         num_codebooks: int,
         num_tokens_per_codebook: int,
         frame_stacking_factor: int,
+        codebook_ordering: str = "frame-major",
     ) -> None:
         super().__init__()
         self.num_codebooks = num_codebooks
         self.num_tokens_per_codebook = num_tokens_per_codebook
         self.frame_stacking_factor = frame_stacking_factor
+        self.iter_slices: Callable = _get_slice_iterator(mode=codebook_ordering)
         self.criterion = nn.CrossEntropyLoss(reduction="none")
 
     @typecheck()
@@ -199,7 +237,7 @@ def forward(
         """
         loss = 0.0
 
-        for fs_index, codebook, start, end, slice_mask, slice_len in _iter_slices(
+        for fs_index, codebook, start, end, slice_mask, slice_len in self.iter_slices(
             self.num_codebooks,
             self.num_tokens_per_codebook,
             self.frame_stacking_factor,
@@ -250,11 +288,13 @@ def __init__(
         num_codebooks: int,
         num_tokens_per_codebook: int,
         frame_stacking_factor: int,
+        codebook_ordering: str = "frame-major",
     ) -> None:
         super().__init__()
         self.num_codebooks = num_codebooks
         self.num_tokens_per_codebook = num_tokens_per_codebook
         self.frame_stacking_factor = frame_stacking_factor
+        self.iter_slices: Callable = _get_slice_iterator(mode=codebook_ordering)
         self.eps = 1e-8
         self.criterion = nn.MSELoss(reduction="none")
 
@@ -286,7 +326,7 @@ def forward(
         student_logits = student_logits.masked_fill(inf_mask, 0.0)
         loss = 0.0
 
-        for _, _, start, end, slice_mask, slice_len in _iter_slices(
+        for _, _, start, end, slice_mask, slice_len in self.iter_slices(
             self.num_codebooks,
             self.num_tokens_per_codebook,
             self.frame_stacking_factor,

nemo/collections/tts/models/__init__.py

@@ -15,6 +15,7 @@
 from nemo.collections.tts.models.aligner import AlignerModel
 from nemo.collections.tts.models.audio_codec import AudioCodecModel
 from nemo.collections.tts.models.easy_magpietts import EasyMagpieTTSModel
+from nemo.collections.tts.models.easy_magpietts_cfg_distillation import EasyMagpieCFGDistillation
 from nemo.collections.tts.models.easy_magpietts_inference import EasyMagpieTTSInferenceModel
 from nemo.collections.tts.models.easy_magpietts_preference_optimization import EasyMagpieTTSModelOnlinePO
 from nemo.collections.tts.models.fastpitch import FastPitchModel
@@ -45,4 +46,5 @@
     "MagpieTTSModelOfflinePODataGen",
     "MagpieTTSModelOfflinePO",
     "MagpieTTSModelOnlinePO",
+    "EasyMagpieCFGDistillation",
 ]