Add support for quantize_() with Float8Linear module

[jainapurva]

Added support for quantize_() API to work with models trained with float8, using Float8Linear.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1344

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit f66e7ed with merge base ed76e9c ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  test/prototype/test_parametrization.py::TestFakeSparsity::test_jit_trace

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[drisspg]

I think makes sense, can you add some tests

Some alternatives would be to maintain the Float8Linear structure and then swap the weight class dtype (feels weird).

I think this is good argument for subclasses since you can maintain structure and then assert all low_p subclasses have a dequant method and call that which will convert to fp32. I know Christian wants tensor.to(torch.float32) to do this but I think its tooo magical

cc @vkuzo

[vkuzo]

can you just move your code snippet from the other file here:

if isinstance(model, Float8Linear):
    with torch.device("meta"):
        new_module = nn.Linear(model.in_features, model.out_features)
    new_module.weight = model.weight
    new_module.bias = model.bias        
    model = new_module

and not need any changes to torchao/float8?

[jerryzh168]

@vkuzo what do you think about having dequantizing a model as a separate API? it feels a bit weird to have this logic in _replace_with_custom_fn_if_matches_filter which is supposed to be a simple module replacement function I feel.

[vkuzo]

my 5c -

1. we can make it work now without adding any public APIs, with minimal increase in complexity
2. if it's important to have a public API for "remove low precision training from a model", we can have that conversation in parallel

wdyt

[jerryzh168]

the motivation for adding a new API is making the dequantizing step more explicit for user, instead of hide it in a module replacement function.

but agree this can happen in parallel. also it's probably not worth spending time to discuss as of now, and wait until there are more use cases might be better

[vkuzo]

would be good to hear some motivation on why this needs a public API versus doing the same thing without a new API

[jerryzh168]

LGTM