Unable resume training from last checkpoint TTS Vit PTL 2.0 [docker nvcr.io/nvidia/nemo:23.08]
State: closed | Pull request: no
Labels: name: bug

name: TTS

**Describe the bug**

I'm trying to resume training from last trained checkpoint (vit tts) but i get following error

```
[NeMo W 2024-02-05 13:09:16 experimental:26] `<class 'nemo.collections.tts.models.vits.VitsModel'>` is experimental and not rea
dy for production yet. Use at your own risk.
Creating ClassifyFst grammars. This might take some time...
[NeMo W 2024-02-05 13:09:42 experimental:26] `<class 'nemo.collections.tts.g2p.models.i18n_ipa.IpaG2p'>` is experimental and no
t ready for production yet. Use at your own risk.
[NeMo W 2024-02-05 13:09:45 i18n_ipa:124] apply_to_oov_word=None, This means that some of words will remain unchanged if they a
re not handled by any of the rules in self.parse_one_word(). This may be intended if phonemes and chars are both valid inputs,
otherwise, you may see unexpected deletions in your input.
[NeMo W 2024-02-05 13:09:45 experimental:26] `<class 'nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPAToken
izer'>` is experimental and not ready for production yet. Use at your own risk.
[NeMo I 2024-02-05 13:09:45 dataset:228] Loading dataset from /path/to/manifests/train_manifest_text_normed.json.
25087it [00:05, 4743.57it/s]
[NeMo I 2024-02-05 13:09:50 dataset:266] Loaded dataset with 25087 files.
[NeMo I 2024-02-05 13:09:50 dataset:268] Dataset contains 35.62 hours.
[NeMo I 2024-02-05 13:09:50 dataset:376] Pruned 0 files. Final dataset contains 25087 files
[NeMo I 2024-02-05 13:09:50 dataset:378] Pruned 0.00 hours. Final dataset contains 35.62 hours.
[NeMo I 2024-02-05 13:09:51 dataset:228] Loading dataset from /path/to/manifests/val_manifest_text_normed.json.
100it [00:00, 4799.47it/s]
[NeMo I 2024-02-05 13:09:51 dataset:266] Loaded dataset with 100 files.
[NeMo I 2024-02-05 13:09:51 dataset:268] Dataset contains 0.15 hours.
[NeMo I 2024-02-05 13:09:51 dataset:376] Pruned 0 files. Final dataset contains 100 files
[NeMo I 2024-02-05 13:09:51 dataset:378] Pruned 0.00 hours. Final dataset contains 0.15 hours.
[NeMo I 2024-02-05 13:09:51 features:289] PADDING: 1
[NeMo I 2024-02-05 13:09:51 features:297] STFT using exact pad
[NeMo W 2024-02-05 13:09:51 nemo_logging:349] /usr/local/lib/python3.10/dist-packages/torch/nn/utils/weight_norm.py:30: UserWar
ning: torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.                              warnings.warn("torch.nn.utils.weight_norm is deprecated in favor of torch.nn.utils.parametrizations.weight_norm.")

Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/2                                                                          Creating ClassifyFst grammars. This might take some time...
25087it [00:05, 4910.99it/s]
100it [00:00, 4593.43it/s]                                                                                                     Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/2                                                                          ----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 2 processes
----------------------------------------------------------------------------------------------------

[NeMo W 2024-02-05 13:10:59 nemo_logging:349] /usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpo
int.py:617: UserWarning: Checkpoint directory /path/to/vits/resultVitsGermanTTS
/VITS/2024-02-04_10-10-06/checkpoints exists and is not empty.
      rank_zero_warn(f"Checkpoint directory {dirpath} exists and is not empty.")

0it [00:00, ?it/s][NeMo I 2024-02-05 13:10:59 dataset:228] Loading dataset from /path/to/manifests/train_manifest_text_normed.json.
25087it [00:01, 23103.62it/s]
25087it [00:01, 23063.96it/s]
[NeMo I 2024-02-05 13:11:00 dataset:266] Loaded dataset with 25087 files.
[NeMo I 2024-02-05 13:11:00 dataset:268] Dataset contains 35.62 hours.
[NeMo I 2024-02-05 13:11:00 dataset:376] Pruned 0 files. Final dataset contains 25087 files
[NeMo I 2024-02-05 13:11:00 dataset:378] Pruned 0.00 hours. Final dataset contains 35.62 hours.
Restoring states from the checkpoint path at resultVitsGermanTTS/VITS/2024-02-04_10-10-06/checkpoints/VITS--loss_gen_all=28.129
7-epoch=37-last.ckpt
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1]
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name                       | Type                     | Params
------------------------------------------------------------------------
0 | audio_to_melspec_processor | FilterbankFeatures       | 0
1 | feat_matching_loss         | FeatureMatchingLoss      | 0
2 | disc_loss                  | DiscriminatorLoss        | 0
3 | gen_loss                   | GeneratorLoss            | 0
4 | kl_loss                    | KlLoss                   | 0
5 | net_g                      | SynthesizerTrn           | 39.7 M
6 | net_d                      | MultiPeriodDiscriminator | 46.7 M
------------------------------------------------------------------------
86.4 M    Trainable params
0         Non-trainable params
86.4 M    Total params
345.615   Total estimated model params size (MB)
Restored all states from the checkpoint at resultVitsGermanTTS/VITS/2024-02-04_10-10-06/checkpoints/VITS--loss_gen_all=28.1297-
epoch=37-last.ckpt
0it [00:00, ?it/s][NeMo I 2024-02-05 13:11:14 dataset:228] Loading dataset from /path/to/manifests/train_manifest_text_normed.json.
25087it [00:01, 22758.76it/s]
25087it [00:01, 21809.31it/s]
[NeMo I 2024-02-05 13:11:15 dataset:266] Loaded dataset with 25087 files.
[NeMo I 2024-02-05 13:11:15 dataset:268] Dataset contains 35.62 hours.
[NeMo I 2024-02-05 13:11:15 dataset:376] Pruned 0 files. Final dataset contains 25087 files
[NeMo I 2024-02-05 13:11:15 dataset:378] Pruned 0.00 hours. Final dataset contains 35.62 hours.
Epoch 37:   0%|                                                                                        | 0/375 [00:00<?, ?it/s]
[NeMo W 2024-02-05 13:11:21 nemo_logging:349] /usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/training_epoch_lo
op.py:151: UserWarning: You're resuming from a checkpoint that ended before the epoch ended. This can cause unreliable results
if further training is done. Consider using an end-of-epoch checkpoint
      rank_zero_warn(

Epoch 37:   0%|                                                                            | 0/375 [00:17<?, ?it/s, v_num=xpkcE
rror executing job with overrides: ['model.sample_rate=22050', 'train_dataset=/path/to/manifests/train_manifest_text_normed.json', 'validation_datasets=/path/to/manifests/val_manifest_text_normed.json', 'phoneme_dict_path=tts_dataset_files/ipa_cmudict-0.7b_nv23
.01.txt', 'heteronyms_path=tts_dataset_files/heteronyms-052722', 'trainer.max_epochs=5000', 'trainer.accelerator=gpu', 'trainer
.check_val_every_n_epoch=2', 'trainer.strategy=ddp_find_unused_parameters_true', 'exp_manager.exp_dir=resultVitsGermanTTS', 'mo
del.text_normalizer.lang=de', '+exp_manager.create_wandb_logger=true', '+exp_manager.wandb_logger_kwargs.name=VITS_DS',
'+exp_manager.wandb_logger_kwargs.project=VITS_GermanTTSCustomDS', '++exp_manager.resume_if_exists=true', '+exp_manager.checkpo
int_callback_params.dirpath=resultVitsGermanTTS/VITS/2024-02-04_10-10-06/checkpoints', 'trainer.devices=2']
Traceback (most recent call last):
  File "/workspace/nemo/examples/tts/vits.py", line 29, in main
    trainer.fit(model)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launc
h
    return function(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 980, in _run
    results = self._run_stage()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1023, in _run_stage
    self.fit_loop.run()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/fit_loop.py", line 202, in run
    self.advance()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/fit_loop.py", line 355, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/training_epoch_loop.py", line 134, in run
    self.on_advance_end()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/training_epoch_loop.py", line 249, in on_advance_end
    self.val_loop.run()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 122, in run
    return self.on_run_end()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 258, in on_run_end
    self._on_evaluation_end()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 303, in _on_evaluation_end
    call._call_callback_hooks(trainer, hook_name, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 194, in _call_callback_hooks
    fn(trainer, trainer.lightning_module, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 311, in on_validation_en
d
    self._save_topk_checkpoint(trainer, monitor_candidates)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 358, in _save_topk_check
point
    raise MisconfigurationException(m)
lightning_fabric.utilities.exceptions.MisconfigurationException: `ModelCheckpoint(monitor='loss_gen_all')` could not find the m
onitored key in the returned metrics: ['epoch', 'step']. HINT: Did you call `log('loss_gen_all', value)` in the `LightningModul
e`?

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
```

**Steps/Code to reproduce bug**

```
WORKING_DIR='/workspace/nemo'

wandb login api-key

cd /path/to/vits

python $WORKING_DIR/examples/tts/vits.py \
  --config-path="$WORKING_DIR/examples/tts/conf/" \
  --config-name=vits \
  model.sample_rate=22050 \ train_dataset="/path/to/manifests/train_manifest_text_normed.json" \
validation_datasets="/path/to/manifests/val_manifest_text_normed.json" \
  phoneme_dict_path="tts_dataset_files/ipa_cmudict-0.7b_nv23.01.txt" \
  heteronyms_path="tts_dataset_files/heteronyms-052722" \
  trainer.max_epochs=5000 \
  trainer.accelerator="gpu" \
  trainer.check_val_every_n_epoch=2 \
  trainer.strategy="ddp_find_unused_parameters_true" \
  exp_manager.exp_dir="resultVitsGermanTTS" \
  model.text_normalizer.lang="de" \
  +exp_manager.create_wandb_logger=true \
  +exp_manager.wandb_logger_kwargs.name="VITS_DS" \
  +exp_manager.wandb_logger_kwargs.project="VITS_GermanTTSCustomDS" \
  ++exp_manager.resume_if_exists=true \
  +exp_manager.checkpoint_callback_params.save_top_k=20 \
  +exp_manager.checkpoint_callback_params.dirpath="resultVitsGermanTTS/VITS/2024-02-04_10-10-06/checkpoints" \
  trainer.devices=2
```

**Environment overview**

 - Environment location: Docker nvcr.io/nvidia/nemo:23.08
```
docker run --runtime=nvidia -it --rm -v /home/ubuntu/DATA:/mnt/e --shm-size=16g \
-p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit \
stack=67108864 nvcr.io/nvidia/nemo:23.08
```

**Additional Note**
@stevehuang52 @KunalDhawan
Could you take a look

I guess the bug is related to the pytorch lightning (PTL) version upgrade. Since nvcr.io/nvidia/nemo:23.08, PTL is upgraded to 2.0, so there may be some back-compatibility issues for older models trained under PTL>2.0. 

Could you pls try `nvcr.io/nvidia/nemo:23.06` or lower versions of nemo container?

> I guess the bug is related to the pytorch lightning (PTL) version upgrade. Since nvcr.io/nvidia/nemo:23.08, PTL is upgraded to 2.0, so there may be some back-compatibility issues for older models trained under PTL>2.0.
> 
> Could you pls try `nvcr.io/nvidia/nemo:23.06` or lower versions of nemo container?

Thanks for your reply. I use docker 23.06 and everything works fine again