Getting NAN loss during ASR training on new language

[mehadi92]

Hi, I'm getting nan loss when fine-tuning conformer ctc large model for new language,

We use mixed precession lr 1.0, weight_decay 0, docker image nvcr.io/nvidia/nemo:22.01

When NaN start coming the lr was 0.000049

Here are the training configs,

  cfg.train_ds.batch_size = 64
  cfg.train_ds.num_workers = 32
  cfg.train_ds.pin_memory = True
  cfg.train_ds.use_start_end_token = True
  cfg.train_ds.trim_silence = False
  cfg.train_ds.max_duration = 18.5
  cfg.train_ds.is_tarred = False

with open_dict(model.cfg.optim):
  model.cfg.optim.lr = 1.0
  model.cfg.optim.weight_decay = 0
  model.cfg.optim.sched.warmup_steps = None 
  model.cfg.optim.sched.warmup_ratio = 0.10
  model.cfg.optim.sched.min_lr = 1e-9

with open_dict(model.cfg.spec_augment):
  model.cfg.spec_augment.freq_masks = 2
  model.cfg.spec_augment.freq_width = 25
  model.cfg.spec_augment.time_masks = 10
  model.cfg.spec_augment.time_width = 0.05
model.spec_augmentation = model.from_config_dict(model.cfg.spec_augment)

trainer = ptl.Trainer(
    accelerator=accelerator,
    max_epochs=EPOCHS,
    accumulate_grad_batches=8,
    log_every_n_steps=10,
    check_val_every_n_epoch=1,
    val_check_interval=.10,
    strategy="ddp_find_unused_parameters_false",
    logger=False,
    checkpoint_callback=False,
    gpus=-1,
    precision=16,
)

Here is the training log,

Epoch 8:  90%|████████▉ | 41819/46535 [15:41:13<1:46:08,  1.35s/it, loss=nan, v_num=ting]
Epoch 8, global step 50466: val_wer was not in top 3
Epoch 8: 100%|█████████▉| 46473/46535 [17:25:49<01:23,  1.35s/it, loss=nan, v_num=ting]
Epoch 8, global step 51040: val_wer was not in top 3
Epoch 9:  10%|▉         | 4587/46535 [1:43:55<15:50:26,  1.36s/it, loss=nan, v_num=ting]
Epoch 8, global step 51613: val_wer was not in top 3
Epoch 9:  20%|█▉        | 9241/46535 [3:28:24<14:01:06,  1.35s/it, loss=nan, v_num=ting]
Epoch 9, global step 52188: val_wer was not in top 3
Epoch 9:  30%|██▉       | 13895/46535 [5:12:51<12:14:54,  1.35s/it, loss=nan, v_num=ting]
Epoch 9, global step 52761: val_wer was not in top 3
Epoch 9:  39%|███▉      | 18354/46535 [6:59:39<10:44:21,  1.37s/it, loss=nan, v_num=ting]
Epoch 9, global step 53335: val_wer was not in top 3

Should I not use mixed precession?
decrease lr and add decay value?

Thanks

[titu1994]

Fp16 stability is dependent on many factors including global batch size, gradient norm, type of loss etc. We have some stability fixes in 1.18 and 1.19 but it is not a perfect solution and will still depend on training recipe

[mehadi92]

So is it safe to use precision=32 to bypass NAN loss?

[titu1994]

Yes it should be fine. BF16 should also be ok

[VahidooX]

Some notes:
cfg.train_ds.use_start_end_token should be False. We have trained all of the recent models without start and end tokens.
train_ds.num_workers is very high and may slow down your trianing. You may use 4 to 8.

[mehadi92]

@VahidooX Thanks for your suggestion We will try this in our next training