I tried this branch earlier today to include different loss scales for each loss in my progressive growing GAN, but I had issues getting it working correctly. What I ended up with was something similar to this:

d_optimizer.zero_grad()
with amp.scale_loss(wasserstein_distance, self.d_optimizer, loss_id=0) as scaled_loss:
                    scaled_loss.backward(retain_graph=True)
with amp.scale_loss(epsilon_penalty, self.d_optimizer, loss_id=1) as scaled_loss:
                    scaled_loss.backward(retain_graph=True)
with amp.scale_loss(gradient_penalty, self.d_optimizer, loss_id=2) as scaled_loss:
                    scaled_loss.backward()
d_optimzier.step()

This seems to be working with a single GPU, but when I use apex DistributedDataParallel (2 GPU's) it is very unstable, and I have no idea why :/

I also took a look at the advanced AMP usage for multiple backward passes for a single optimizer, where it says that you should use delay_unscale=True. I tried this (by having delay_unscale=True for the first two backward passes), but it seems like it was using the last loss's loss_scaler to unscale all the grads. So the model was diverging immediately.

Bold of you to test it already...That usage looks correct. I'm glad to hear it's working on single-GPU and that the new syntax enables your use case. I'm still trying it with multi-GPU.

I haven't updated documentation for the new branch, but guidance will be that using delay_unscale is no longer necessary and should only be used as a minor performance optimization if you're really sure what you're doing.

@hukkelas I notice you're using gradient penalty. I don't think Apex DDP supports gradient penalty properly, so your instability may not be the fault of the new branch. I still haven't finished testing the new branch with distributed, but it seems to work ok so far. Can you try torch.nn.parallel.DistributedDataParallel instead? Replace

model, optimizer = amp.initialize(model, optimizer, opt_level=XX, num_losses=YY)
model = apex.parallel.DistributedDataParallel(model)

with

model, optimizer = amp.initialize(model, optimizer, opt_level=XX, num_losses=YY)
model = torch.nn.parallel.DistributedDataParallel(model,
    device_ids=[args.local_rank],
    output_device=args.local_rank)

Hmm, didn't know that, thanks!

I was using Apex DDP, switched to torch.nn.parallel.DistributedDataParallel now. From some quick tests it looks better.

I just recently started using multi-gpu training for my GAN. I think most of my previous instability, with a single GPU, has come from the gradient penalty calculation. I made a hacky solution to introduce independent loss scaling for it previously, but this branch would remove a lot of ugly code :)

Glad it appears to be working now. My local distributed tests/tests on local models have also succeeded so far so I'll probably merge this branch today.

When you say you are scaling the gradient penalty calculation, is this the line you are referring to:

with amp.scale_loss(gradient_penalty, self.d_optimizer, loss_id=2) as scaled_loss:
                    scaled_loss.backward()

? Scaling the loss for this call^ is fine.

Or are you referring to an earlier call to torch.autograd.grad(...., create_graph=True, only_inputs=True) to create out-of-place gradients used to construct the gradient_penalty scalar? Scaling the loss for that particular call would not be supported by the new branch either (Amp has no visibility of the newly created out-of-place gradients, so it has no way to check them for infs).

Before this branch, I independently scaled my gradient penalty like this:

logits = discriminator(x_hat, condition, landmarks)
logits = logits.sum() * loss_scaler.loss_scale()
grad = torch.autograd.grad(
    outputs=logits,
    inputs=x_hat,
    grad_outputs=torch.ones(logits.shape).to(fake_data.dtype).to(fake_data.device),
    create_graph=True
)[0] 
grad = grad.view(x_hat.shape[0], -1)
if check_overflow(grad):
    print("Overflow in gradient penalty calculation.")
    loss_scaler._loss_scale /= 2
    print("Scaling down loss to:", loss_scaler._loss_scale)
    return None
grad = grad / loss_scaler.loss_scale()

grad_penalty = ((grad.norm(p=2, dim=1) - 1)**2)

check_overfow is doing the same checks as scale_check_overflow_python in scaler.py. If there is an overflow, I just skip the current batch for all GPUs.

From my understanding, since AMP has no visibility of the torch.autograd.grad call, I should keep this independent scaling of the gradient penalty, and also include the independent loss scaling as I showed earlier?