Enable offloading of validation models to CPU during training runs for EMTTS

examples/tts/conf/magpietts/easy_magpietts_lhotse.yaml

@@ -8,6 +8,8 @@ quadratic_duration: 20
 
 model:
   use_lhotse: true
+  # Keep validation-only scorers and waveform codec modules on CPU between validation epochs.
+  offload_validation_models: true
 
   # Decoder backend selection
   # Options: "huggingface" (default), "nemotron_h"
@@ -224,4 +226,4 @@ exp_manager:
     always_save_nemo: true
     filename: '${name}--{${exp_manager.checkpoint_callback_params.monitor}:.4f}-{step}-{epoch}'
   resume_if_exists: true
-  resume_ignore_no_checkpoint: true
+  resume_ignore_no_checkpoint: true

nemo/collections/tts/models/easy_magpietts.py

@@ -129,6 +129,7 @@ def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
 
         # Validation inference with metrics (optional)
         self.run_val_inference = cfg.get('run_val_inference', False)
+        self.offload_validation_models = cfg.get('offload_validation_models', True)
         self.use_multilingual_asr = cfg.get('use_multilingual_asr', False)
         if self.run_val_inference:
             logging.info("Loading eval models for validation inference (ASR and speaker verification)...")
@@ -162,6 +163,138 @@ def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
             self._utmos_calculator = UTMOSv2Calculator(device='cpu')
             logging.info("UTMOSv2 calculator initialized for validation naturalness scoring")
 
+        self._mark_validation_models_ddp_ignored()
+
+    def _mark_validation_models_ddp_ignored(self) -> None:
+        """Exclude exactly the tensors belonging to modules that will move to CPU."""
+
+        ignored_names = self._validation_model_ddp_ignored_names()
+        if ignored_names:
+            torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model(self, ignored_names)
+
+    def _validation_modules_for_offload(self) -> list[tuple[str, nn.Module]]:
+        """Return the unique module objects moved between CPU and GPU for validation."""
+
+        modules = []
+        codec_model = getattr(self, '_codec_model', None)
+        if isinstance(codec_model, nn.Module):
+            modules.extend(
+                (f'_codec_model.{name}', module)
+                for name, module in codec_model.named_children()
+                if name != 'vector_quantizer'
+            )
+        modules.extend(
+            (name, module)
+            for name in ('_eval_asr_model', '_eval_speaker_verification_model', 'whisper_model')
+            if isinstance(module := getattr(self, name, None), nn.Module)
+        )
+
+        unique_modules = []
+        seen = set()
+        for name, module in modules:
+            if id(module) not in seen:
+                unique_modules.append((name, module))
+                seen.add(id(module))
+        return unique_modules
+
+    def _validation_model_ddp_ignored_names(self) -> list[str]:
+        """Find every parameter/buffer alias that resolves to an offloaded module tensor."""
+
+        if not self._should_offload_validation_models():
+            return []
+        modules = self._validation_modules_for_offload()
+        parameter_ids = {id(parameter) for _, module in modules for parameter in module.parameters()}
+        buffer_ids = {id(buffer) for _, module in modules for buffer in module.buffers()}
+        ignored_names = {
+            name for name, parameter in self.named_parameters(remove_duplicate=False) if id(parameter) in parameter_ids
+        }
+        ignored_names.update(
+            name for name, buffer in self.named_buffers(remove_duplicate=False) if id(buffer) in buffer_ids
+        )
+        return sorted(ignored_names)
+
+    def _uses_validation_models_during_training(self) -> bool:
+        """Whether validation/generation models are part of the training objective."""
+
+        return False
+
+    def _should_offload_validation_models(self) -> bool:
+        return self.offload_validation_models and not self._uses_validation_models_during_training()
+
+    def _move_codec_waveform_modules(self, device: torch.device | str) -> list[str]:
+        """Move codec components not needed for cached-code conversion.
+
+        The original vector quantizer must remain with the training model because
+        the codec converter uses it on every cached-code batch. The waveform
+        encoder/decoder and the codec's training-only auxiliaries can move independently.
+        """
+
+        moved = []
+        for name, module in self._validation_modules_for_offload():
+            if not name.startswith('_codec_model.'):
+                continue
+            module.to(device)
+            moved.append(name)
+        return moved
+
+    def _move_validation_models(self, device: torch.device | str) -> None:
+        """Move validation scorers and waveform codec components to the requested device."""
+
+        moved = self._move_codec_waveform_modules(device)
+        for name, module in self._validation_modules_for_offload():
+            if name.startswith('_codec_model.'):
+                continue
+            module.to(device)
+            moved.append(name)
+
+        # UTMOSv2 is deliberately constructed on CPU and scores saved waveforms
+        # there, so moving it to the validation GPU would only increase GPU usage.
+        if moved:
+            logging.info("Moved validation-only modules to %s: %s", device, ", ".join(moved))
+
+    def _offload_validation_models(self) -> None:
+        """Move validation modules to CPU after refreshing all DDP buffer exclusions."""
+
+        if not self._should_offload_validation_models():
+            return
+        self._refresh_validation_model_ddp_ignores()
+        self._move_validation_models(torch.device('cpu'))
+        self._refresh_validation_model_ddp_ignores()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def _refresh_validation_model_ddp_ignores(self) -> None:
+        """Apply offloaded-tensor exclusions to the constructed DDP wrapper and verify them."""
+
+        if not self._should_offload_validation_models() or self.trainer is None:
+            return
+        ddp_model = getattr(getattr(self.trainer, 'strategy', None), 'model', None)
+        if not isinstance(ddp_model, torch.nn.parallel.DistributedDataParallel):
+            return
+
+        ignored_names = set(self._validation_model_ddp_ignored_names())
+        ddp_model.parameters_to_ignore.update(ignored_names)
+        ddp_model._assign_modules_buffers()
+        offloaded_buffer_ids = {
+            id(buffer) for _, module in self._validation_modules_for_offload() for buffer in module.buffers()
+        }
+        synchronized_offloaded_buffers = [
+            name for name, buffer in ddp_model.named_module_buffers.items() if id(buffer) in offloaded_buffer_ids
+        ]
+        if synchronized_offloaded_buffers:
+            raise RuntimeError(
+                "DDP still includes validation-only buffers that will be offloaded: "
+                + ", ".join(synchronized_offloaded_buffers)
+            )
+        logging.info(
+            "Excluded %d validation-only parameter/buffer names from DDP synchronization",
+            len(ignored_names),
+        )
+
+    def on_train_start(self):
+        super().on_train_start()
+        self._offload_validation_models()
+
     def _get_state_dict_keys_to_exclude(self):
         return super()._get_state_dict_keys_to_exclude() + [
             '_speaker_verification_model',
@@ -1278,6 +1411,11 @@ def process_batch(
         )
 
     def training_step(self, batch, batch_idx):
+        uses_uncached_audio = 'context_audio_codes' not in batch or 'audio_codes' not in batch
+        if uses_uncached_audio and self._should_offload_validation_models():
+            audio = batch.get('context_audio', batch.get('audio'))
+            self._move_codec_waveform_modules(audio.device)
+
         if 'context_audio_codes' in batch:
             context_audio_codes = batch['context_audio_codes']
             context_audio_codes_lens = batch['context_audio_codes_lens']
@@ -1296,6 +1434,8 @@ def training_step(self, batch, batch_idx):
             audio_lens = batch['audio_lens']
             audio_codes, audio_codes_lens = self._codec_helper.audio_to_codes(audio, audio_lens)
 
+        if uses_uncached_audio and self._should_offload_validation_models():
+            self._move_codec_waveform_modules(torch.device('cpu'))
         if (
             self.cfg.get("use_multiturn_dataset", False)
             and batch["user_audio_turn_splitted"] is not None
@@ -1859,6 +1999,8 @@ def on_fit_start(self):
             self._generate_codec_silence_buffer()
 
     def on_validation_epoch_start(self) -> None:
+        if self._should_offload_validation_models():
+            self._move_validation_models(self.device)
         if torch.distributed.is_initialized():
             self.trainer.strategy.model.require_backward_grad_sync = False
 
@@ -1917,6 +2059,7 @@ def collect_if_exists(key):
                     )
 
         self.validation_step_outputs.clear()  # free memory
+        self._offload_validation_models()
 
         if torch.distributed.is_initialized():
             self.trainer.strategy.model.require_backward_grad_sync = True

nemo/collections/tts/models/easy_magpietts_preference_optimization.py

@@ -65,6 +65,9 @@ class EasyMagpieTTSModelOnlinePO(EasyMagpieTTSModel):
     5. Add auxiliary phoneme loss from the same forward pass with GT phoneme tokens.
     """
 
+    def _uses_validation_models_during_training(self) -> bool:
+        return True
+
     def __init__(self, cfg: DictConfig, trainer: 'Trainer' = None):
         """Initialize the online PO model, including the frozen reference model, reward ASR/speaker
         verification models, optional UTMOSv2 scorer, and all PO hyper-parameters from ``cfg``.

tests/collections/tts/models/test_easy_magpietts_validation_offload.py

@@ -0,0 +1,122 @@
+# 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.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from nemo.collections.tts.models.easy_magpietts import EasyMagpieTTSModel
+from nemo.collections.tts.models.easy_magpietts_preference_optimization import EasyMagpieTTSModelOnlinePO
+
+
+class TrackingModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.moves = []
+
+    def to(self, device, *args, **kwargs):
+        self.moves.append(torch.device(device))
+        return super().to(device, *args, **kwargs)
+
+
+def make_model() -> EasyMagpieTTSModel:
+    model = object.__new__(EasyMagpieTTSModel)
+    torch.nn.Module.__init__(model)
+    model.offload_validation_models = True
+    model._codec_model = torch.nn.Module()
+    model._codec_model.audio_encoder = TrackingModule()
+    model._codec_model.vector_quantizer = TrackingModule()
+    model._codec_model.audio_decoder = TrackingModule()
+    model._codec_converter = TrackingModule()
+    model._eval_asr_model = TrackingModule()
+    model._eval_speaker_verification_model = TrackingModule()
+    model.whisper_model = TrackingModule()
+    model.register_parameter("training_parameter", torch.nn.Parameter(torch.ones(1)))
+    return model
+
+
+@pytest.mark.unit
+def test_validation_model_offload_keeps_training_quantizers_resident():
+    model = make_model()
+
+    model._move_validation_models(torch.device('cpu'))
+
+    assert model._codec_model.audio_encoder.moves == [torch.device('cpu')]
+    assert model._codec_model.audio_decoder.moves == [torch.device('cpu')]
+    assert model._eval_asr_model.moves == [torch.device('cpu')]
+    assert model._eval_speaker_verification_model.moves == [torch.device('cpu')]
+    assert model.whisper_model.moves == [torch.device('cpu')]
+    assert model._codec_model.vector_quantizer.moves == []
+    assert model._codec_converter.moves == []
+
+
+@pytest.mark.unit
+def test_validation_epoch_start_restores_offloaded_modules(monkeypatch):
+    model = make_model()
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
+    monkeypatch.setattr(EasyMagpieTTSModel, "device", property(lambda _: torch.device("cpu")), raising=False)
+
+    model.on_validation_epoch_start()
+
+    assert model._codec_model.audio_encoder.moves == [torch.device("cpu")]
+    assert model._eval_asr_model.moves == [torch.device("cpu")]
+
+
+@pytest.mark.unit
+def test_online_po_keeps_reward_models_available_during_training():
+    assert EasyMagpieTTSModel._uses_validation_models_during_training(None) is False
+    assert EasyMagpieTTSModelOnlinePO._uses_validation_models_during_training(None) is True
+
+
+@pytest.mark.unit
+def test_offloaded_modules_are_ignored_by_ddp():
+    model = make_model()
+    model._codec_model.audio_decoder = torch.nn.Sequential(torch.nn.Linear(2, 2))
+    aliased_scorer = torch.nn.Module()
+    aliased_scorer.register_buffer('running_state', torch.ones(2))
+    model.scorer_alias = aliased_scorer
+    model._eval_asr_model = aliased_scorer
+
+    model._mark_validation_models_ddp_ignored()
+
+    ignored = set(model._ddp_params_and_buffers_to_ignore)
+    assert '_codec_model.audio_decoder.0.weight' in ignored
+    assert 'scorer_alias.running_state' in ignored
+    assert '_eval_asr_model.running_state' in ignored
+    assert '_codec_model.vector_quantizer' not in ignored
+    assert 'training_parameter' not in ignored
+
+
+@pytest.mark.unit
+def test_validation_offload_refreshes_ddp_ignores_for_lazy_buffers():
+    model = make_model()
+    model._eval_asr_model = torch.nn.Module()
+    model._eval_asr_model.register_buffer('initial_state', torch.ones(2))
+    model._mark_validation_models_ddp_ignored()
+
+    ddp_model = object.__new__(torch.nn.parallel.DistributedDataParallel)
+    torch.nn.Module.__init__(ddp_model)
+    ddp_model.broadcast_buffers = True
+    ddp_model.module = model
+    ddp_model.parameters_to_ignore = set(model._ddp_params_and_buffers_to_ignore)
+    ddp_model._assign_modules_buffers()
+    model._trainer = SimpleNamespace(strategy=SimpleNamespace(model=ddp_model))
+
+    model._eval_asr_model.register_buffer('lazy_state', torch.ones(2))
+    model._offload_validation_models()
+    ddp_model._assign_modules_buffers()
+
+    assert '_eval_asr_model.lazy_state' not in ddp_model.named_module_buffers
+    assert 'training_parameter' in dict(model.named_parameters())