torch.inference_mode switches`aten.linear.default, this is not supported`

[michaelfeil]

First, congrats for the repo - looks great

I discovered that switching between torch.no_grad and torch.inference_mode leads to a switch to aten.linear.default. Feel free to use this feedback - its likely expected.

# dependencies
'2.4.0.dev20240316+cu121'
float8_experimental commit 88e9e507c56e59c5f17edf513ecbf621b46fc67d

from transformers import AutoModel, AutoTokenizer, PreTrainedModel
import torch

from float8_experimental.float8_linear_utils import (
    swap_linear_with_float8_linear,
)
from float8_experimental.float8_dynamic_linear import Float8DynamicLinear

"""
pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 --upgrade
"""

model_name = "michaelfeil/bge-small-en-v1.5"
model: PreTrainedModel = AutoModel.from_pretrained(model_name)
model_orig: PreTrainedModel = AutoModel.from_pretrained(model_name).cuda()

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokens = tokenizer.batch_encode_plus(["hi","there is no"], return_tensors="pt", padding=True, truncation=True)
tokens = {k: v.cuda() for k,v in tokens.items()}

model = model.to("cpu")
swap_linear_with_float8_linear(model, Float8DynamicLinear)
model = model.to("cuda")

with torch.no_grad():
    out1 = model_orig.forward(**tokens)["last_hidden_state"]
    out2 = model.forward(**tokens)["last_hidden_state"]

print(out1 - out2, out1, out2, (out1-out2).mean())

with torch.inference_mode():
   # breaks with NotImplementedError: attempting to run aten.linear.default, this is not supported
    out2 = model.forward(**tokens)["last_hidden_state"]

[drisspg]

Hey, thanks for opening this issue! Let me take a look at this

[tgaddair]

Hey @drisspg thanks for looking into this. Any updates you're able to share on this? Fixing this would be a enabler for a lot of serving use cases.

[drisspg]

Hey, thanks for repinging 😊! I just added a one liner that unblocked the repro, still need to tidy some things up.

I do think that this will allow you to experiment further with fp8 for inference, however I do not think that this is the ideal usage. FP8 will benefit from what is usual done in other int8 quantization techniques.

Ex. Float8DynamicLinear's weight param will always recompute its scale since the weight is updated during training. However, for inference this is frozen and we should make this static. Dynamic quantization for the activations is still appropriate but one could want to statically calibrate this activation on example data.

TorchAO is one such project that aims to provide a number of quantization techniques.

We do have plans to add a better training to inference solution for FP8! So stay tuned but, in the meantime I will work to land the linked PR

[vkuzo]

pytorch/ao#564