hifigan_finetune on nemo:1.8.1 crashes when trying to train with `trainer.precision=bf16`
State: closed | Pull request: no
Labels: name: bug

Running the following script
```bash
# train
python examples/tts/hifigan_finetune.py \
  --config-path=conf/hifigan \
  --config-name=hifigan \
  exp_manager.name=hifigan \
  +exp_manager.exp_dir=/experiments \
  +exp_manager.version=20220423-1956 \
  exp_manager.resume_if_exists=true \
  exp_manager.resume_ignore_no_checkpoint=true \
  +exp_manager.checkpoint_callback_params.always_save_nemo=true \
  train_dataset=/data/manifests/manifest_v4_train.json \
  validation_datasets=/data/manifests/manifest_v4_valid.json \
  trainer.devices=-1 \
  trainer.num_nodes=1 \
  trainer.strategy=ddp \
  trainer.accelerator=gpu \
  trainer.precision=bf16
```
on one A100 GPU results in the following stack trace
```python3
[NeMo W 2022-04-23 19:56:25 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pydub/utils.py:170: RuntimeWarning: Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work
      warn("Couldn't find ffmpeg or avconv - defaulting to ffmpeg, but may not work", RuntimeWarning)
    
[NeMo W 2022-04-23 19:56:26 experimental:27] Module <class 'nemo.collections.nlp.data.language_modeling.megatron.megatron_batch_samplers.MegatronPretrainingRandomBatchSampler'> is experimental, not ready for production and is not fully supported. Use at your own risk.
[NeMo W 2022-04-23 19:56:27 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/hydra/_internal/defaults_list.py:251: UserWarning: In 'hifigan': Defaults list is missing `_self_`. See https://hydra.cc/docs/upgrades/1.0_to_1.1/default_composition_order for more information
      warnings.warn(msg, UserWarning)
    
Multiprocessing is handled by SLURM.
Using bfloat16 Automatic Mixed Precision (AMP)
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
[NeMo W 2022-04-23 19:56:27 exp_manager:409] There was no checkpoint folder at checkpoint_dir :/experiments/hifigan/20220423-1956/checkpoints. Training from scratch.
[NeMo I 2022-04-23 19:56:27 exp_manager:281] Experiments will be logged at /experiments/hifigan/20220423-1956
[NeMo I 2022-04-23 19:56:27 exp_manager:647] TensorboardLogger has been set up
[NeMo W 2022-04-23 19:56:27 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:2302: LightningDeprecationWarning: `Trainer.weights_save_path` has been deprecated in v1.6 and will be removed in v1.8.
      rank_zero_deprecation("`Trainer.weights_save_path` has been deprecated in v1.6 and will be removed in v1.8.")
    
[NeMo W 2022-04-23 19:56:27 exp_manager:881] The checkpoint callback was told to monitor a validation value and trainer's max_steps was set to 25000000. Please ensure that max_steps will run for at least 10 epochs to ensure that checkpointing will not error out.
Created a temporary directory at /tmp/tmpw14tzmbq
Writing /tmp/tmpw14tzmbq/_remote_module_non_sriptable.py
[NeMo I 2022-04-23 19:56:27 data:775] Loading dataset from /data/manifests/manifest_v4_train.json.
[NeMo I 2022-04-23 19:56:27 data:799] Loaded dataset with 10322 files.
[NeMo I 2022-04-23 19:56:27 data:801] Dataset contains 16.66 hours.
[NeMo I 2022-04-23 19:56:27 data:297] Pruned 0 files. Final dataset contains 10322 files
[NeMo I 2022-04-23 19:56:27 data:299] Pruned 0.00 hours. Final dataset contains 16.66 hours.
[NeMo I 2022-04-23 19:56:27 data:775] Loading dataset from /data/manifests/manifest_v4_valid.json.

0it [00:00, ?it/s]
10322it [00:00, 181425.66it/s]
[NeMo I 2022-04-23 19:56:27 data:799] Loaded dataset with 543 files.
[NeMo I 2022-04-23 19:56:27 data:801] Dataset contains 0.84 hours.
[NeMo I 2022-04-23 19:56:27 data:297] Pruned 87 files. Final dataset contains 456 files
[NeMo I 2022-04-23 19:56:27 data:299] Pruned 0.05 hours. Final dataset contains 0.79 hours.
[NeMo I 2022-04-23 19:56:27 features:259] PADDING: 0
[NeMo I 2022-04-23 19:56:27 features:276] STFT using torch
[NeMo I 2022-04-23 19:56:27 features:278] STFT using exact pad
[NeMo I 2022-04-23 19:56:28 features:259] PADDING: 0
[NeMo I 2022-04-23 19:56:28 features:276] STFT using torch
[NeMo I 2022-04-23 19:56:28 features:278] STFT using exact pad

0it [00:00, ?it/s]
543it [00:00, 170791.68it/s]
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/1
Added key: store_based_barrier_key:1 to store for rank: 0
Rank 0: Completed store-based barrier for key:store_based_barrier_key:1 with 1 nodes.
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 1 processes
----------------------------------------------------------------------------------------------------

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
[NeMo W 2022-04-23 19:56:32 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py:327: RuntimeWarning: The lr scheduler dict contains the key(s) ['interval'], but the keys will be ignored. You need to call `lr_scheduler.step()` manually in manual optimization.
      rank_zero_warn(
    

  | Name                       | Type                     | Params
------------------------------------------------------------------------
0 | audio_to_melspec_precessor | FilterbankFeatures       | 0     
1 | trg_melspec_fn             | FilterbankFeatures       | 0     
2 | generator                  | Generator                | 13.9 M
3 | mpd                        | MultiPeriodDiscriminator | 41.1 M
4 | msd                        | MultiScaleDiscriminator  | 29.6 M
5 | feature_loss               | FeatureMatchingLoss      | 0     
6 | discriminator_loss         | DiscriminatorLoss        | 0     
7 | generator_loss             | GeneratorLoss            | 0     
------------------------------------------------------------------------
84.7 M    Trainable params
0         Non-trainable params
84.7 M    Total params
338.643   Total estimated model params size (MB)
SLURM auto-requeueing enabled. Setting signal handlers.
[NeMo W 2022-04-23 19:56:33 nemo_logging:349] /opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:240: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 256 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
      rank_zero_warn(
    
Error executing job with overrides: ['exp_manager.name=hifigan', '+exp_manager.exp_dir=/experiments', '+exp_manager.version=20220423-1956', 'exp_manager.resume_if_exists=true', 'exp_manager.resume_ignore_no_checkpoint=true', '+exp_manager.checkpoint_callback_params.always_save_nemo=true', 'train_dataset=/data/manifests/manifest_v4_train.json', 'validation_datasets=/data/manifests/manifest_v4_valid.json', 'trainer.devices=-1', 'trainer.num_nodes=1', 'trainer.strategy=ddp', 'trainer.accelerator=gpu', 'trainer.precision=bf16']
Traceback (most recent call last):
  File "examples/tts/hifigan_finetune.py", line 32, in <module>
    main()  # noqa pylint: disable=no-value-for-parameter
  File "/workspace/nemo/nemo/core/config/hydra_runner.py", line 104, in wrapper
    _run_hydra(
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 377, in _run_hydra
    run_and_report(
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 214, in run_and_report
    raise ex
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 211, in run_and_report
    return func()
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/utils.py", line 378, in <lambda>
    lambda: hydra.run(
  File "/opt/conda/lib/python3.8/site-packages/hydra/_internal/hydra.py", line 111, in run
    _ = ret.return_value
  File "/opt/conda/lib/python3.8/site-packages/hydra/core/utils.py", line 233, in return_value
    raise self._return_value
  File "/opt/conda/lib/python3.8/site-packages/hydra/core/utils.py", line 160, in run_job
    ret.return_value = task_function(task_cfg)
  File "examples/tts/hifigan_finetune.py", line 28, in main
    trainer.fit(model)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 768, in fit
    self._call_and_handle_interrupt(
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 721, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 809, in _fit_impl
    results = self._run(model, ckpt_path=self.ckpt_path)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1234, in _run
    results = self._run_stage()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1321, in _run_stage
    return self._run_train()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1343, in _run_train
    self._run_sanity_check()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1411, in _run_sanity_check
    val_loop.run()
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 153, in advance
    dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/base.py", line 204, in run
    self.advance(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 127, in advance
    output = self._evaluation_step(**kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 222, in _evaluation_step
    output = self.trainer._call_strategy_hook("validation_step", *kwargs.values())
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 1763, in _call_strategy_hook
    output = fn(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/strategies/ddp.py", line 347, in validation_step
    return self.model(*args, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 968, in forward
    output = self.module(*inputs[0], **kwargs[0])
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1111, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/pytorch_lightning/overrides/base.py", line 93, in forward
    return self.module.validation_step(*inputs, **kwargs)
  File "/workspace/nemo/nemo/collections/tts/models/hifigan.py", line 234, in validation_step
    pred_denoised = self._bias_denoise(audio_pred, audio_mel).squeeze(1)
  File "/workspace/nemo/nemo/collections/tts/models/hifigan.py", line 306, in _bias_denoise
    self.stft_bias, _ = stft(audio_bias)
  File "/workspace/nemo/nemo/collections/tts/models/hifigan.py", line 292, in stft
    comp = torch.stft(x.squeeze(1), n_fft=1024, hop_length=256, win_length=1024)
  File "/opt/conda/lib/python3.8/site-packages/torch/functional.py", line 693, in stft
    input = F.pad(input.view(extended_shape), [pad, pad], pad_mode)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/functional.py", line 4371, in _pad
    return torch._C._nn.reflection_pad1d(input, pad)
RuntimeError: "reflection_pad1d_out_template" not implemented for 'BFloat16'

Sanity Checking: 0it [00:00, ?it/s]
Sanity Checking DataLoader 0:   0%|          | 0/2 [00:00<?, ?it/s]
```