LR scheduler does not work properly due to AcceleratedOptimizer

Ever since accelerate has been introduced, the optimizer inside the Trainer is doubly wrapped - first DPOptimizer wraps the original Optimizer, and then AcceleratedOptimizer wraps the DPOptimizer. Opacus does not do the wrapping very carefully, simply passing state, defaults and param_groups by reference - this reference could easily be replaced, and as a matter of fact accelerate seems to be doing that. In contrast, the wrapping approach used by the AcceleratedOptimizer guarantees that the reference can't be easily replaced.

A consequence of this is that LR schedulers do not seem to work right now. To verify this, you can add the following statements inside the on_substep_end callback:

print('LR at AcceleratedOptimizer: {}'.format([g['lr'] for g in optimizer.param_groups]))
print('LR at DPOptimizer: {}'.format([g['lr'] for g in dp_optimizer.param_groups]))
print('LR at original Optimizer: {}'.format([g['lr'] for g in dp_optimizer.original_optimizer.param_groups]))

The first 2 follow the scheduler, but the 3rd one doesn't, and the LR actually applied once self.original_optimizer.step() is called inside of Opacus is precisely that of the original Optimizer.

I only noticed this because I was using a scheduler that started with a LR of 0, and the model parameters weren't changing at all, even if at the 2nd step already LR > 0. It should be harder to notice it in general.

Introducing a new DPOptimizer that inherits from Opacus' but sets param_groups using properties, like the AcceleratedOptimizer does, seems to fix the issue:

class DPOptimizer(opacus.optimizers.DPOptimizer):
    @property
    def param_groups(self):
        return self.original_optimizer.param_groups

    @param_groups.setter
    def param_groups(self, param_groups):
        self.original_optimizer.param_groups = param_groups

Start at the Trainer's on_substep_end callback and inspect how the optimizer, DPOptimizer, and AcceleratedOptimizer expose param_groups. Compare opacus/optimizers/optimizer.py with accelerate/src/accelerate/optimizer.py, then reproduce the reported learning-rate mismatch. Done means the scheduler's learning rate reaches the original optimizer before its step is applied.