Why is the fp16 conformer not learning?
State: closed | Pull request: no
Labels: name: bug

Can you please tell me how to properly train a model with mixed precision?
I set to train with a model with such a config, but on validation I get NaN all the time.
Config:
```
name: "Conformer-CTC-BPE"

model:
  sample_rate: &sample_rate 8000

  log_prediction: false # enables logging sample predictions in the output during training
  ctc_reduction: 'mean_batch'

  train_ds:
    manifest_filepath: "../data/processed/manifests/train.json"

    sample_rate: *sample_rate
    batch_size: 2 # you may increase batch_size if your memory allows
    shuffle: true
    num_workers: &num_workers 8
    # pin_memory: false
    trim_silence: true
    max_duration: 32.9 # it is set for LibriSpeech, you may need to update it for your dataset
    min_duration: 0.1
    # tarred datasets
    is_tarred: false
    tarred_audio_filepaths: null
    shuffle_n: 2048
    # bucketing params
    bucketing_strategy: "synced_randomized"
    bucketing_batch_size: null

  validation_ds:
    manifest_filepath: "../data/processed/manifests/val.json"
    sample_rate: *sample_rate
    batch_size: 1
    shuffle: false
    num_workers: 2
    pin_memory: false

  test_ds:
    manifest_filepath: "../data/processed/manifests/test.json"
    sample_rate: *sample_rate
    batch_size: 1 # you may increase batch_size if your memory allows
    shuffle: false
    num_workers: 1
    pin_memory: false
  tokenizer:
    dir: "../models/tokenizers/tokenizer_spe_bpe_v128_max_2"
    type: bpe
  preprocessor:
    _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor
    sample_rate: ${model.sample_rate}
    normalize: "per_feature"
    window_size: 0.025
    window_stride: 0.01
    window: "hann"
    features: 80
    n_fft: 512
    log: true
    frame_splicing: 1
    dither: 0.00001
    pad_to: 0
    pad_value: 0.0

  spec_augment:
    _target_: nemo.collections.asr.modules.SpectrogramAugmentation
    freq_masks: 2 # set to zero to disable it
    # you may use lower time_masks for smaller models to have a faster convergence
    time_masks: 10 # set to zero to disable it
    freq_width: 27
    time_width: 0.05

  encoder:
    _target_: nemo.collections.asr.modules.ConformerEncoder
    feat_in: ${model.preprocessor.features}
    feat_out: -1
    n_layers: 18
    d_model: 256

    # Sub-sampling params
    subsampling: striding # vggnet or striding, vggnet may give better results but needs more memory
    subsampling_factor: 4 # must be power of 2
    subsampling_conv_channels: -1 # 176 set to -1 to make it equal to the d_model

    # Feed forward module's params
    ff_expansion_factor: 4

    # Multi-headed Attention Module's params
    self_attention_model: rel_pos # rel_pos or abs_pos
    n_heads: 4 # may need to be lower for smaller d_models
    # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention
    att_context_size: [-1, -1] # -1 means unlimited context
    xscaling: true # scales up the input embeddings by sqrt(d_model)
    untie_biases: true # unties the biases of the TransformerXL layers
    pos_emb_max_len: 5000

    # Convolution module's params
    conv_kernel_size: 31
    conv_norm_type: 'batch_norm' # batch_norm or layer_norm

    ### regularization
    dropout: 0.1 # The dropout used in most of the Conformer Modules
    dropout_emb: 0.0 # The dropout used for embeddings
    dropout_att: 0.1 # The dropout for multi-headed attention modules

  decoder:
    _target_: nemo.collections.asr.modules.ConvASRDecoder
    feat_in: null
    num_classes: -1
    vocabulary: null

  optim:
    name: adamw
    lr: 5.0
    # optimizer arguments
    betas: [0.9, 0.98]
    # less necessity for weight_decay as we already have large augmentations with SpecAug
    # you may need weight_decay for large models, stable AMP training, small datasets, or when lower augmentations are used
    # weight decay of 0.0 with lr of 2.0 also works fine
    weight_decay: 0.0

    # scheduler setup NoamAnnealing
    sched:
      name: NoamAnnealing
      d_model: ${model.encoder.d_model}
      # scheduler config override
      warmup_steps: 10000
      warmup_ratio: null
      min_lr: 1e-6

trainer:
  devices: -1 # number of GPUs, -1 would use all available GPUs
  num_nodes: 1
  max_epochs: 100
  max_steps: -1 # computed at runtime if not set
  val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations
  accelerator: gpu
  strategy: dp
  accumulate_grad_batches: 16
  gradient_clip_val: 1.0
  precision: 16 # Should be set to 16 for O1 and O2 to enable the AMP.
  amp_level: O1
  amp_backend: apex
  log_every_n_steps: 50  # Interval of logging.
  progress_bar_refresh_rate: 10
  resume_from_checkpoint: null
  num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it
  check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs
  sync_batchnorm: true
  enable_checkpointing: true  # Provided by exp_manager
  # logger: false  # Provided by exp_manager
  benchmark: false # needs to be false for models with variable-length speech input as it slows down training


exp_manager:
  exp_dir: null
  name: ${name}
  create_tensorboard_logger: true
  create_checkpoint_callback: true
  checkpoint_callback_params:
    # in case of multiple validation sets, first one is used
    monitor: "val_wer"
    mode: "min"
    save_top_k: 5
    always_save_nemo: True # saves the checkpoints as nemo files instead of PTL checkpoints

  # you need to set these two to True to continue the training
  resume_if_exists: false
  resume_ignore_no_checkpoint: false

  # You may use this section to create a W&B logger
  create_wandb_logger: false
  wandb_logger_kwargs:
    name: null
    project: null

```

Output:
```
[NeMo W 2022-09-24 19:03:30 nemo_logging:349] /home/user/projects/conformer/venv/lib/python3.8/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:611: UserWarning: Checkpoint directory /home/maxikon/projects/conformer/models/nemo/conformer/trained exists and is not empty.
      rank_zero_warn(f"Checkpoint directory {dirpath} exists and is not empty.")
    
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

[NeMo I 2022-09-24 19:03:32 modelPT:597] Optimizer config = AdamW (
    Parameter Group 0
        amsgrad: False
        betas: [0.8, 0.98]
        capturable: False
        eps: 1e-08
        foreach: None
        lr: 5.0
        maximize: False
        weight_decay: 0.0
    )
[NeMo I 2022-09-24 19:03:32 lr_scheduler:910] Scheduler "<nemo.core.optim.lr_scheduler.NoamAnnealing object at 0x7f98981dc3a0>" 
    will be used during training (effective maximum steps = 38900) - 
    Parameters : 
    (d_model: 256
    warmup_steps: 10000
    warmup_ratio: null
    min_lr: 1.0e-06
    max_steps: 38900
    )


  | Name              | Type                              | Params
------------------------------------------------------------------------
0 | preprocessor      | AudioToMelSpectrogramPreprocessor | 0     
1 | encoder           | ConformerEncoder                  | 30.5 M
2 | decoder           | ConvASRDecoder                    | 33.2 K
3 | loss              | CTCLoss                           | 0     
4 | spec_augmentation | SpectrogramAugmentation           | 0     
5 | _wer              | WERBPE                            | 0     
------------------------------------------------------------------------
30.5 M    Trainable params
0         Non-trainable params
30.5 M    Total params
122.154   Total estimated model params size (MB)
[NeMo W 2022-09-24 19:03:32 nemo_logging:349] /home/user/projects/conformer/venv/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 12 which is the number of cpus on this machine) in the DataLoader init to improve performance.
      rank_zero_warn(
    

Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic

Epoch 0: 19%
1470/7790 [04:32<19:32, 5.39it/s, loss=565, v_num=29]

[NeMo W 2022-09-24 19:03:38 nemo_logging:349] /home/user/projects/conformer/venv/lib/python3.8/site-packages/torch/optim/lr_scheduler.py:124: UserWarning: Seems like optimizer.step() has been overridden after learning rate scheduler initialization. Please, make sure to call optimizer.step() before lr_scheduler.step(). See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
      warnings.warn("Seems like optimizer.step() has been overridden after learning rate scheduler "
    

Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 1024.0
Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 512.0
Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 256.0
```

![image](https://user-images.githubusercontent.com/47638600/192102647-77e193e8-a8e0-428d-91da-5193d1609ddc.png)
![image](https://user-images.githubusercontent.com/47638600/192103008-62b3b96c-16d2-475f-ae2c-7f807d8712d5.png)
![image](https://user-images.githubusercontent.com/47638600/192103007-9dd084f5-7023-4f66-8670-ccd76e5c66f9.png)

**Environment details**
- Linux Ubuntu 20.04
- PyTorch 11.6
- Python 3.8
- NeMo 1.11.0 (or 1.8.2)