diff --git a/examples/tts/conf/magpietts/easy_magpietts_lhotse.yaml b/examples/tts/conf/magpietts/easy_magpietts_lhotse.yaml index b06482868d08..d20723a6e751 100644 --- a/examples/tts/conf/magpietts/easy_magpietts_lhotse.yaml +++ b/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 \ No newline at end of file + resume_ignore_no_checkpoint: true diff --git a/nemo/collections/tts/models/easy_magpietts.py b/nemo/collections/tts/models/easy_magpietts.py index 0f4fd215cb8f..fcf62c8a4c43 100644 --- a/nemo/collections/tts/models/easy_magpietts.py +++ b/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 diff --git a/nemo/collections/tts/models/easy_magpietts_preference_optimization.py b/nemo/collections/tts/models/easy_magpietts_preference_optimization.py index 34e632030e6f..1032e654a72b 100644 --- a/nemo/collections/tts/models/easy_magpietts_preference_optimization.py +++ b/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``. diff --git a/tests/collections/tts/models/test_easy_magpietts_validation_offload.py b/tests/collections/tts/models/test_easy_magpietts_validation_offload.py new file mode 100644 index 000000000000..a577b9d3b675 --- /dev/null +++ b/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())