AO and Automated Mixed Precision

[bhack]

Can we clarify in the readme what are the best practices to use ao at inference with a pytorch AMP trainer model/checkpoint?

[vkuzo]

@bhack , would love to learn some more context - is there an issue when you use the checkpoint directly?

[bhack]

Yes exactly it is the case when we use a checkpoint trained with AMP and we are entering in the AO world to optimize it for inference.

[bhack]

Generally in the repo before AO I just see model.eval().to(torch.bfloat16).
Is this the best practice?
If yes we could add some notes in the documentation.

[vkuzo]

Sorry for the late response, I was on holiday last week. There are no special things you need to do when using amp + torchao for converting the model to inference - just use torchao according to the documentation. Would you mind clarifying your question if that does not answer it?

[bhack]

Just to make an example from the Doc it is not clear if this is the best practice or not:

       model.eval()
       quantize_(model, int8_dynamic_activation_int8_weight())
       unwrap_tensor_subclass(model)
       with torch.amp.autocast(enabled=use_amp, dtype=input_dtype, device_type="cuda"):
           ......
           torch.export.export(....)
           aoti_compile_and_package(...)

[vkuzo]

I see, looks like your question is whether using autocast for quantized model inference is supported / recommended. I'll tag @jerryzh168 on this one.

[jerryzh168]

Just to make an example from the Doc it is not clear if this is the best practice or not:

       model.eval()
       quantize_(model, int8_dynamic_activation_int8_weight())
       unwrap_tensor_subclass(model)
       with torch.amp.autocast(enabled=use_amp, dtype=input_dtype, device_type="cuda"):
           ......
           torch.export.export(....)
           aoti_compile_and_package(...)

we haven't tested this, but this seems reasonable to me, and

1. the quantized kernel should be able to handle the activation dtype from autocast
2. for inference, people typically already starts with these lower precision dtypes like bfloat16, so not sure what is the use case here

[bhack]

For 2:
Do you mean model.eval().to(torch.bfloat16) without AMP context?

[jerryzh168]

For 2: Do you mean model.eval().to(torch.bfloat16) without AMP context?

Yes

[bhack]

Ok, it think it is better to add a note somewhere as it is quite common to have models trained with AMP.

[jerryzh168]

Ok, it think it is better to add a note somewhere as it is quite common to have models trained with AMP.

Do you know the e2e flow for a model trained with AMP? do people

1. train with AMP, save the state_dict
2. load the state_dict and quantize the model etc.?