Nano: quantization API with onnx renew design

[TheaperDeng]

Current status (without onnx quantization)

For pytorch-lightning user:

pl_model = Trainer.compile(model, loss, optim)  # skip if you have a pl model
trainer.fit(pl_model, dataloader)
pl_model_quantized = trainer.quantize(pl_model, dataloader)
pl_model_quantized(x)  # quantized inference

For pytorch user (potentially, we do not support it yet but it can be valid simply):

# >>>>>>>> start of pytorch training loop >>>>>>>>>>>
# ...
# <<<<<<<< end  of pytorch training loop <<<<<<<<<<<<
model_quantized = trainer.quantize(model, dataloader)
model_quantized(x)  # quantized inference

Issue

1. pl_model_quantized and model_quantized are torch.fx.graph_module.GraphModule. An unfamilar type to users. Users can not:

• use onnx.export to trace the quantized model.

• use normal way to save or load the model.

• continually train on this quantized model

2. It is extremly hard to integrate onnx quantization with INC to this API

New revised API usage (with onnx quantization)

For pytorch-lightning user:

pl_model = Trainer.compile(model, loss, optim, onnx=True/False)  # skip if you have a pl model
trainer.fit(pl_model, dataloader)
pl_model = trainer.quantize(pl_model, dataloader, onnx=True/False)

For pytorch user:

model = Trainer.compile(model, onnx=True/False)
# >>>>>>>> start of pytorch training loop >>>>>>>>>>>
# ...
# <<<<<<<< end  of pytorch training loop <<<<<<<<<<<<
model = trainer.quantize(model, dataloader, onnx=True/False)

pl_model and model are still pytorch-lightning model, then the prediction can be

# predict with pytorch fp32
pl_model.eval()
with torch.no_grad():
    pl_model(x)
# or
pl_model.inference(x, backend=None)

# predict with pytorch int8
pl_model.eval(quantize=True)
with torch.no_grad():
    pl_model(x)
# or
pl_model.inference(x, backend=None, quantize=True)

# predict with onnx fp32
pl_model.eval_onnx()
with torch.no_grad():
    pl_model(x)
# or
pl_model.inference(x, backend="onnx")

# predict with onnx int8
pl_model.eval_onnx(quantize=True)
with torch.no_grad():
    pl_model(x)
# or
pl_model.inference(x, backend="onnx", quantize=True)

We should also provide an option to return the "raw" result to users, where onnx quantization will return a onnx model( onnx.onnx_ml_pb2.ModelProto ) and pytorch fx quantization will return a fx model( torch.fx.GraphModule ).

model = trainer.quantize(..., raw_return=True)  # defaultly False

So basically three PRs will be raised separately for this issue:

• change the Trainer.quantize(..., raw_return=False) to return a pl model and support the easy inference api (.eval(quantized=True)) Nano quantize inference API for pytorch #3866
• add support to onnx's quantization
• add support of save and load

[jason-dai]

q_model = trainer.quantize(model, dataloader, onnx=True)

1. Why not have trainer.quantize(model, dataloader) always return a pytorch-lightning model instead of torch.fx.graph_module.GraphModule?

2. After calling quantize, the reasonable behavior should be using int8 for prediction by default?

q_model.eval()
with torch.no_grad():
    q_model(x)  # predict with pytorch int8

q_model.eval(quantize=false)
with torch.no_grad():
    q_model(x)  # predict with pytorch fp32

3. What's the expect behavior when calling trainer.fit(q_model, ...) in this case?

---

trainer.save_checkpoint("example.ckpt") # save all int8/fp32 models

This only works after calling trainer.fit, yes? I think we need to deal with three different cases here:

1. Newly fitted (fp32) model - I guess the user can use standard PL or PyTorch method to save and load?
2. Model accelerated using ONNX (o_model=Trainer.compile(model, onnx=True)) - I guess there is no need to handle it (except that the user may want to save the onnx file)?
3. Model accelerated using quantization (q_model = trainer.quantize(model, dataloader, onnx=True) - How should a user save and load q_model?

[TheaperDeng]

Save and load

reference to https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-for-inference

save

pytorch user:

# torch model
torch.save(model.state_dict(), PATH)  # fp32 pytorch model
torch.save(model.quantized_state_dict(), PATH)  # int8 pytorch model

# onnx
torch.to_onnx(...)  # fp32 onnx file export
torch.to_onnx(..., quantized=True)  # int8 onnx file export

pytorch lightning user (also support all operations as pytorch users):

# additional lightning fashion
trainer.save_checkpoint("example.ckpt")  # save all int8/fp32 models for onnx/pytorch

load

since onnx file is easy to load we only consider torch model.

pytorch user:

# torch model
model = TheModelClass(*args, **kargs)
model = Trainer.compile(model)
model.load_state_dict(torch.load(PATH))  # fp32
model.load_quantized_state_dict(torch.load(PATH))  # int8

pytorch lightning user (also support all operations as pytorch users):

# pytorch-lightning user
pl_model = Trainer.compile(model, loss, optim, onnx=True)
pl_model.load_from_checkpoint(checkpoint_path="example.ckpt")  # load all int8/fp32 models for onnx/pytorch

[TheaperDeng]

q_model = trainer.quantize(model, dataloader, onnx=True)

1. Why not have trainer.quantize(model, dataloader) always return a pytorch-lightning model instead of torch.fx.graph_module.GraphModule?
2. After calling quantize, the reasonable behavior should be using int8 for prediction by default?

q_model.eval()
with torch.no_grad():
    q_model(x)  # predict with pytorch int8

q_model.eval(quantize=false)
with torch.no_grad():
    q_model(x)  # predict with pytorch fp32

3. What's the expect behavior when calling trainer.fit(q_model, ...) in this case?

trainer.save_checkpoint("example.ckpt") # save all int8/fp32 models

This only works after calling trainer.fit, yes? I think we need to deal with three different cases here:

1. Newly fitted (fp32) model - I guess the user can use standard PL or PyTorch method to save and load?
2. Model accelerated using ONNX (o_model=Trainer.compile(model, onnx=True)) - I guess there is no need to handle it (except that the user may want to save the onnx file)?
3. Model accelerated using quantization (q_model = trainer.quantize(model, dataloader, onnx=True) - How should a user save and load q_model?

1. Yes, we will always return a pytorch-lightning model unless users state raw_return=True.
2. I think it might be clearer for our users if we fix quantized=False. But I have not come to a conclusion.
3. stated in Nano: quantization API with onnx renew design #3806 (comment)

[jason-dai]

Save and load

reference to pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-for-inference

save

pytorch user:

# torch model
torch.save(model.state_dict(), PATH)  # fp32 pytorch model
torch.save(model.quantized_state_dict(), PATH)  # int8 pytorch model

# onnx
torch.to_onnx(...)  # fp32 onnx file export
torch.to_onnx(..., quantized=True)  # int8 onnx file export

pytorch lightning user (also support all operations as pytorch users):

# additional lightning fashion
trainer.save_checkpoint("example.ckpt")  # save all int8/fp32 models for onnx/pytorch

This only works if the user has called trainer.fit first, yes?

load

since onnx file is easy to load we only consider torch model.

pytorch user:

# torch model
model = TheModelClass(*args, **kargs)
model = Trainer.compile(model)
model.load_state_dict(torch.load(PATH))  # fp32
model.load_quantized_state_dict(torch.load(PATH))  # int8

What if the user passes an unquantized state to load_quantized_state_dict? Can we detect that?

pytorch lightning user (also support all operations as pytorch users):

# pytorch-lightning user
pl_model = Trainer.compile(model, loss, optim, onnx=True)
pl_model.load_from_checkpoint(checkpoint_path="example.ckpt")  # load all int8/fp32 models for onnx/pytorch

[TheaperDeng]

Save and load

reference to pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-for-inference

save

pytorch user:

# torch model
torch.save(model.state_dict(), PATH)  # fp32 pytorch model
torch.save(model.quantized_state_dict(), PATH)  # int8 pytorch model

# onnx
torch.to_onnx(...)  # fp32 onnx file export
torch.to_onnx(..., quantized=True)  # int8 onnx file export

pytorch lightning user (also support all operations as pytorch users):

# additional lightning fashion
trainer.save_checkpoint("example.ckpt")  # save all int8/fp32 models for onnx/pytorch

This only works if the user has called trainer.fit first, yes?

We may do something to bind the model when users call trainer.quantize. But yes, users at least should do some thing ( quantize/fit) on this model by the trainer to be called save_checkpoint. And that is why it should only work for a pl user.

load

since onnx file is easy to load we only consider torch model.
pytorch user:

# torch model
model = TheModelClass(*args, **kargs)
model = Trainer.compile(model)
model.load_state_dict(torch.load(PATH))  # fp32
model.load_quantized_state_dict(torch.load(PATH))  # int8

What if the user passes an unquantized state to load_quantized_state_dict? Can we detect that?

I'm not 100% sure, but I think we can(and we should).

pytorch lightning user (also support all operations as pytorch users):

# pytorch-lightning user
pl_model = Trainer.compile(model, loss, optim, onnx=True)
pl_model.load_from_checkpoint(checkpoint_path="example.ckpt")  # load all int8/fp32 models for onnx/pytorch

[TheaperDeng]

So basically three PRs will be raised separately for this issue:

• change the Trainer.compile(..., raw_return=False) to return a pl model and support the easy inference api (.eval(quantized=True))
• add support to onnx's quantization
• add support of save and load