How to customize amp in pytorchightning

[MooreManor]

To use amp autocast in pytorch, the pipeline is listed below.

from torch.cuda.amp import autocast as autocast

model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        scaler.scale(loss).backward()

        scaler.step(optimizer)

        scaler.update()

I want to reimplement the same procedure in pytorch-ligtning, but I don't know where to rewrite the call of scaler.scale(loss).backward(), scaler.step(optimizer), andscaler.update() because they are encapsulated in pl. Now I want to customize these parts, but I don't know what are the corresponding apis in pl.

[akihironitta]

I want to reimplement the same procedure in pytorch-ligtning, but I don't know where to rewrite the call of scaler.scale(loss).backward(), scaler.step(optimizer), andscaler.update() because they are encapsulated in pl.

The implementation sits in pytorch_lightning.plugins.precision.native_amp.NativeMixedPrecisionPlugin:
https://github.com/Lightning-AI/lightning/blob/4c47619f38331e4a2c13c0562e5c46500cfd3d70/src/pytorch_lightning/plugins/precision/native_amp.py

[MooreManor]

@akihironitta

Maybe I didn't make it very clear in the original question. For example, I override training_dataloader to customize my own dataset. In PL, I want to replace the part optimizer.step() which is hidden in PL with scaler.scale(loss).backward(), scaler.step(optimizer), andscaler.update() . I think PL has a corresponding API.

[aniketmaurya]

Hey @MooreManor, if you're looking for how to disable the automatic optimizer update and customize that part then here is the resource you're looking for 😄

[MooreManor]

@aniketmaurya Thanks! A manual Optimization is what I want 👍

[akihironitta]

In PL, I want to replace the part optimizer.step() which is hidden in PL with scaler.scale(loss).backward(), scaler.step(optimizer), andscaler.update() . I think PL has a corresponding API.

@MooreManor Just curious why you want to replace the optimizer step call with scaler.scale(loss).backward(), scaler.step(optimizer) and scaler.update(). If you call them manually, your code will no longer be hardware-agnostic. FYI, PL has the built-in support for PyTorch's amp so that your code stays the same across different precision settings:

model = YourModule()
trainer = Trainer(precision=16)
trainer.fit(model, ...)

[MooreManor]

@akihironitta Thanks for your tips. Because I am not familiar with the mechanism how PL does in its amp optimization, I am not sure this automatic method will have the same effect as the pseudocode mentioned above. Are there any docs about PL that show the pseudocode implementation of PL AMP?

[akihironitta]

Are there any docs about PL that show the pseudocode implementation of PL AMP?

@MooreManor No, there aren't AFAIK. However, it should have the same effect as your pseudocode. Do you see any issues around amp with PL?

[kapsner]

@MooreManor did you find a solution for that? I am curious how you finally solved this, as I am facing the exact same obstacle now.