Loss Turns to NaN After Several Hundred Ticks When Using Mixed-Precision Training

[aiihn]

Hello,
I've encountered an issue where the loss turns to NaN after several hundred ticks when I enable mixed-precision training using the --fp16=True flag. And I noticed this line in the code:

ect/training/ct_training_loop.py

Line 117 in 28c1223

loss_scaling = 1, # Loss scaling factor for reducing FP16 under/overflows.

I'm wondering if I should also adjust the --ls setting for loss scaling in conjunction with the --fp16=True flag. Could you advise what value the loss scaling should be set to under these conditions?

Additionally, are there other specific settings that should be configured to optimize mixed-precision training? For example, should the "learning rate" be modified together with "loss scaling"?

If possible, could you share the commands or configuration that you typically use for mixed-precision training?

Thank you very much in advance!

[Mingxiao-Li]

I came across the same issue. I also got nan loss and nan gradient during training.

[Gsunshine]

Hi @aiihn and @Mingxiao-Li ,

Thanks for your interest in ECT! I used neither fp16 training nor loss scaling (fp32 for my experiments).

If you consider mixed precision training, I'd suggest bf16 (usually better numerical stability) and turning on the loss scaling factor. I don't have a recommended value for the loss scaling since I didn't try it. You may start with--ls=100, following EDM's config.

Thanks,
Zhengyang

[aiihn]

Thanks for your quick response! @Gsunshine

I have tried using fp16 with --ls=100, and then at tick 462 the loss still turned to NaN. When I resumed from the checkpoint of tick 450 and continued using fp16 but changed to --lr=1000, the loss turned to NaN again at tick 699.

So, I tried using torch.cuda.amp.GradScaler to automatically adjust the loss scaling, and it works well. Here is the modified code near this line:

ect/training/ct_training_loop.py

Line 257 in 28c1223

loss.mul(loss_scaling).mean().backward()

dist.print0(f'#### GradScaler enabled: {gradscaler_enabled} ####')
dist.print0(f'#### Loss scaling is not used when GradScaler is enabled ####')
if gradscaler_enabled:
    # https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#adding-gradscaler
    # https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation
    scaler = torch.cuda.amp.GradScaler()
    # TODO Save and resume scaler state when resuming training from a checkpoint, or it may affect reproducibility.
    # https://discuss.pytorch.org/t/do-i-need-to-save-the-state-dict-oof-gradscaler/95718/4

while True:

    # Accumulate gradients.
    optimizer.zero_grad(set_to_none=True)
    for round_idx in range(num_accumulation_rounds):
        with misc.ddp_sync(ddp, (round_idx == num_accumulation_rounds - 1)):
            images, labels = next(dataset_iterator)
            images = images.to(device).to(torch.float32) / 127.5 - 1
            labels = labels.to(device)

            loss = loss_fn(net=ddp, images=images, labels=labels, augment_pipe=augment_pipe)
            training_stats.report('Loss/loss', loss)
            # loss.sum().mul(loss_scaling / batch_gpu_total).backward()
            if gradscaler_enabled:
                scaler.scale(loss.mean()).backward()
            else:
                loss.mul(loss_scaling).mean().backward()

    if gradscaler_enabled:
        # TODO Is torch.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad) needed when using GradScaler?
        scaler.step(optimizer)
        scaler.update()
    else:
        # Update weights.
        # for g in optimizer.param_groups:
        #     g['lr'] = optimizer_kwargs['lr'] * min(cur_nimg / max(lr_rampup_kimg * 1000, 1e-8), 1)
        for param in net.parameters():
            if param.grad is not None:
                torch.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad)
        optimizer.step()

Hope it works for you too. @Mingxiao-Li

[Gsunshine]

Hi @aiihn ,

Thanks for the prompt reply! I appreciate your solution here. You are welcome if you'd like to submit a PR to fix this issue!

If you're tuning EDM (not EDM2), weighting & dropout from ECM (based on EDM2) may also help. I am still surprised by the 0.5 dropout rate used for ECM-M/L/XL. The learning rate is associated with batch size. You can consider sqrt(2), 2, or other scaling factors according to your budget and batch size.

Cheers,
Zhengyang

[Mingxiao-Li]

Thanks for your quick response! @Gsunshine

I have tried using fp16 with --ls=100, and then at tick 462 the loss still turned to NaN. When I resumed from the checkpoint of tick 450 and continued using fp16 but changed to --lr=1000, the loss turned to NaN again at tick 699.

So, I tried using torch.cuda.amp.GradScaler to automatically adjust the loss scaling, and it works well. Here is the modified code near this line:

ect/training/ct_training_loop.py

Line 257 in 28c1223

loss.mul(loss_scaling).mean().backward()

dist.print0(f'#### GradScaler enabled: {gradscaler_enabled} ####')
dist.print0(f'#### Loss scaling is not used when GradScaler is enabled ####')
if gradscaler_enabled:
    # https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html#adding-gradscaler
    # https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-accumulation
    scaler = torch.cuda.amp.GradScaler()
    # TODO Save and resume scaler state when resuming training from a checkpoint, or it may affect reproducibility.
    # https://discuss.pytorch.org/t/do-i-need-to-save-the-state-dict-oof-gradscaler/95718/4

while True:

    # Accumulate gradients.
    optimizer.zero_grad(set_to_none=True)
    for round_idx in range(num_accumulation_rounds):
        with misc.ddp_sync(ddp, (round_idx == num_accumulation_rounds - 1)):
            images, labels = next(dataset_iterator)
            images = images.to(device).to(torch.float32) / 127.5 - 1
            labels = labels.to(device)

            loss = loss_fn(net=ddp, images=images, labels=labels, augment_pipe=augment_pipe)
            training_stats.report('Loss/loss', loss)
            # loss.sum().mul(loss_scaling / batch_gpu_total).backward()
            if gradscaler_enabled:
                scaler.scale(loss.mean()).backward()
            else:
                loss.mul(loss_scaling).mean().backward()

    if gradscaler_enabled:
        # TODO Is torch.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad) needed when using GradScaler?
        scaler.step(optimizer)
        scaler.update()
    else:
        # Update weights.
        # for g in optimizer.param_groups:
        #     g['lr'] = optimizer_kwargs['lr'] * min(cur_nimg / max(lr_rampup_kimg * 1000, 1e-8), 1)
        for param in net.parameters():
            if param.grad is not None:
                torch.nan_to_num(param.grad, nan=0, posinf=1e5, neginf=-1e5, out=param.grad)
        optimizer.step()

Hope it works for you too. @Mingxiao-Li

Thanks for your solutions. Using GradScaler and bf16 together during training solves the nan loss problem in my case, but the loss increases after several hundreds of iterations. Do you come across this issue ?

[aiihn]

@Mingxiao-Li I haven't tried using GradScaler with bf16, only with fp16. Although bf16 has a wider range compared to fp16, it has lower precision for fractional values. Since GradScaler already addresses the overflow and underflow problems caused by fp16, I suspect using the lower precision of bf16 might not be necessary. I've submitted a PR; maybe you could try it out and help test if it works correctly?

[Mingxiao-Li]

@Mingxiao-Li I haven't tried using GradScaler with bf16, only with fp16. Although bf16 has a wider range compared to fp16, it has lower precision for fractional values. Since GradScaler already addresses the overflow and underflow problems caused by fp16, I suspect using the lower precision of bf16 might not be necessary. I've submitted a PR; maybe you could try it out and help test if it works correctly?

Thanks. I will check that. I am trying to apply ECM on pretrained dit model, which uses iddpm parameterisation. Not sure if my issues come from the parameterisation or the precision.

[1999kevin]

@Mingxiao-Li Hi, Mingxiao. When applying ECM on DiT, do you encounter an issue that the loss will suddenly increase to a much higher value in the later training stage (for example when self.stage in ECMLoss class is higher than 5)? I guess the current ECM for DiT may be broken due to the different timestamp embedding method with the U-Net for EDM.