Add module-swap UX for INT8 mixed-precision training

[gau-nernst]

Background

The current INT8 mixed-precision training recipe conceptually is an op-modifier - torch.matmul is replaced with dynamic_int8_mm. It is implemented with tensor subclass, though it doesn't use any tensor subclass-specific features, such as quantized storage and quantized FSDP all-gather. Having an alternative module-swap UX would have the following benefits:

1. State dict remains plain tensor. This is beneficial for model checkpointing, as well as some complex use cases such as shard-then-load pre-trained weights for FSDP fine-tuning (cc Integrate INT8 mixed-precision from torchao 0.7 torchtune#1552)
2. It can be easier to hack/compose with other techniques. e.g. use NF4 weight for storage and INT8 matmul for compute -> QLoRA integration

Usage

from torchao import quantize_
from torchao.prototype.quantized_training import int8_mixed_precision_training

model = ...
## nn.Linear -> Int8MixedPrecisionTrainingLinear
quantize_(model, int8_mixed_precision_training(module_swap=True))

Benchmarks

Pre-train Llama2-1B on 4070Ti SUPER. torch==2.6.0.dev20241029. No regression. Module swap has the same perf as tensor subclass

python benchmarks/quantized_training/pretrain_llama2.py --seed 2024 --bf16_model --compile --quantize int8_mixed_precision_module_swap --model 1B --activation_checkpointing

Pre-train Llama3-8B with torchtitan, 2x A100, torch==2.6.0.dev20241104+cu124. No regression. Module swap has the same perf as tensor subclass

Fine-tune Llama3-1B QLoRA with torchtune (using pytorch/torchtune@main...gau-nernst:qlora)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 ca8c85a with merge base f99b667 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/whl/nightl... / 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.

[vkuzo]

It is implemented with tensor subclass, though it doesn't use any tensor subclass-specific features, such as quantized storage and quantized FSDP all-gather. Having an alternative module-swap UX would have the following benefits:

You can still use tensor subclass inside of your module swap to interact with FSDP and TP/SP and get low precision all_gather. See Float8Linear for reference. I think this is a good way to go for training in general.

[gau-nernst]

You can still use tensor subclass inside of your module swap

The way I see it is that, whatever can be done with module swap, can be done with tensor subclass. (maybe it's better to hold persistent states with modules? like for delayed scaling. But at least for my use cases, I don't need persistent states). So using both module swap + tensor subclass feels redundant to me.

[vkuzo]

lgtm