add float8 support

[bdhirsh]

repro command:

CONFIG_FILE="./torchtitan/models/llama3/train_configs/debug_model.toml" ./run_train.sh --model.name llama3_auto_parallel --parallelism.tensor_parallel_degree 4 --model.converters="float8"

my integration uses an existing float8 util that just replaces any linear layers in the user model with Float8Linear layers.

The good news

we don't actually need special subclass handling in autoparallel yet! The model weights are stored in high precision and gets quantized dynamically during the forward, so we have no subclass params in the state dict. After talking to Driss/Daniel, this is true for both float8 compute, and float8 allgathers

The bad news

The first thing that breaks is that _scaled_mm raises an error because it rquires its second arg to be column-major, and autoparallel accidentally makes all intermediate tensors contiguous during FlopCounter estimation. fix here: https://github.com/pytorch-labs/autoparallel/pull/37

After that, there are 3 problems. First, I ran the above repro on Francisco's deepseek branch, with my PR patched in above: https://github.com/pytorch-labs/autoparallel/pull/29

(1) i needed to patch this:

diff --git a/autoparallel/compute_estimation.py b/autoparallel/compute_estimation.py
index f9b2de1..d4d347b 100644
--- a/autoparallel/compute_estimation.py
+++ b/autoparallel/compute_estimation.py
@@ -43,6 +43,7 @@ DEVICE_LIMITS: Tuple[DeviceLimit, ...] = (
             # but we want the full GEMM numbers
             torch.float32: 989 // 2,
             torch.float16: 1979 // 2,
+            torch.float8_e4m3fn: 3958 // 2,
             torch.bfloat16: 1979 // 2,
             torch.int8: 3958 // 2,
         },

(2) it fails with RuntimeError: Expected both dimensions of mat2 to be divisble by 16 but got torch.Size([256, 564]). This looks like the same issue that @wconstab has been looking into, where we are trying to use an invalid sharding.

(3) there are some ops without sharding rules:

aten.clamp_max.default
aten.clamp_min.default

[tianyu-l]

Is this considering FSDP at all?

Previously we tried this for SimpleFSDP, but found the loss weren't close versus not applying Float8. Could it be that there's extra code needed when integrating Float8 with FSDP?
Anyways, @mori360 and I are looking at SimpleFSDP + float8. We can chat more when you are back.