Update Float8Tensor for GRPO training in unsloth

[andrewor14]

Summary: Support a few extra ops called during GRPO loop in unsloth/vllm for Float8Tensor.

Test Plan:

python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_fp8_matmul_lora_variants
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_to_dtype_layout
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_has_compatible_shallow_copy_type
python test/quantization/quantize_/workflows/float8/test_float8_tensor.py -k test_transpose

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 1 Unrelated Failure

As of commit 82012af with merge base f856d36 ():

NEW FAILURES - The following jobs have failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  test/integration/test_integration.py::TestExport::test_export_float8
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  test/integration/test_integration.py::TestExport::test_export_float8

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/sparsity/test_sparse_api.py::TestQuantSemiSparse::test_sparse_marlin_compile_True

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

[vkuzo]

plz clean up _float8_mm_impl

[vkuzo]

instead of this flag, just transpose at the callsite to match the meaning of matmul

[andrewor14]

The reason behind this flag is to prevent unnecessary double transpose when we call linear, since the fbgemm op expects the weight to be in the linear format (already transposed). So if we don't have this flag:

1. linear calls _float8_mm_impl(input, weight.t())
2. _float8_mm_impl calls weight.t() before calling torch.ops.fbgemm.f8f8bf16

If we just transpose the weight for linear, we may end up slowing linear down, is that OK?

[vkuzo]

how about

1. define _float8_mm_impl as the float8 version of torch.mm, as the lowest level shared code of maybe quantizing the input and then choosing a gemm
2. all other functions (matmul, linear, etc) call _float8_mm_impl

it's a bit confusing to have two different paths for linear and matmul

[andrewor14]

This was done to avoid double transpose in the linear path (which doesn't happen today, see this comment). I agree that ideally everything should go through _float8_mm_impl, but doing so may add overhead for the linear path, should I go ahead and merge the implementations anyway?

[andrewor14]

Ok, just did some benchmark on the double transpose thing, seems like it didn't introduce much overhead. I refactored the code to the way you suggested, please have another look