Hi @ptrblck, thanks for your answer. You are right about the typo, I corrected it for clarity.
A bit more context:
The wrapper wraps two nn.Modules, the constraint and the model. They do not have any shared parameters. They do both interact with (parts of) the total loss. The loss composition looks something like:
loss_1, loss_2 = model(input)
loss_2 = constraint(loss_2)
total_loss = loss_1 + loss_2

The model optimises with gradients from the total_loss, while the constraint optimises only part of the loss (loss_2). I hope this makes it clearer.
On your suggested solutions:
1. Using two scalers.
Would the following be valid?
model_scaled_loss = model_scaler.scale(loss)
constraint_scaled_loss = constraint_scaler.scale(loss)
loss.backward()

if (global_step + 1) % accumulate_n_batches_grad == 0:
        # OPTIMISER STEP (first unscaled)
        model_scaler.step(model_optimiser)

        # AMP SCALER UPDATE
        model_scaler.update()

        # Zero the gradients only here
        model_optimiser.zero_grad()

constraint_scaler.step(constraint_optimiser)
constraint_optimiser.zero_grad()

2. use scaler.step() on the constraint_optimizer , if the inner block wasn’t executed.
I am not sure what you mean by this. Are you saying only to execute scaler.step(constraint_optimiser) if scaler.step(model_optimiser) is not called (inner block execution)? I would think that both optimisers need explicit unscaling of the gradients, or don’t they? If both optimisers do not share parameters, would scaler.unscale_(model_optimiser) affect scaler.unscale_(constraint_optimiser) in any way then?
Hopefully this is clear enough, otherwise I can provide more information.
Cheers,
Claartje