[torchao float8tensor]

[crcrpar]

What does this PR do?

Improve the tensor subclass support of #1394 for TorchAo float8.

note: pytorch/ao#1339 is needed

my environment

• torch: 2.6.0a0+git62eea62
• nvfuser: 0.2.23+gitbb05859
• torchao: 0.7.0+gitb2e42ff6
• CUDA device: RTX 6000 Ada Generation
• Driver Version: 560.35.03
• CUDA Version: 12.6

[t-vi]

@crcrpar if you merge main, the pt nightly distributed ci tests should be fixed.

[crcrpar]

This change should be in #1394

[crcrpar]

#1485

[crcrpar]

should be in #1394

[crcrpar]

This change should also be in #1394

[crcrpar]

should be in #1394

[crcrpar]

This failure doesn't feel easy to fix to me. So I made this into a script:

import torch
import torch.nn as nn
from torchao.float8 import convert_to_float8_training
from thunder.dynamo import ThunderCompiler
from thunder.dynamo.splitter import SubgraphInfo
from thunder.tests.make_tensor import make_tensor


def main():
    batch_size, in_features, out_features = 16, 32, 64

    device = torch.device("cuda")
    dtype = torch.float32

    model = nn.Linear(in_features, out_features, bias=False, device=device, dtype=dtype)
    fp8_model = convert_to_float8_training(model)
    x = make_tensor((batch_size, in_features), device=device, dtype=dtype)
    expected = fp8_model(x)

    backend = ThunderCompiler()
    jitted = torch.compile(fp8_model, backend=backend)
    actual = jitted(x)

    backend.save_reproducer_to_folder("./debug_torchao_with_thunderfx", use_pytest_benchmark=True)
    print(f"{len(backend.subgraph_infos) = }")
    subgraph: SubgraphInfo
    for subgraph in backend.subgraph_infos:
        print(f"# {len(subgraph.thunder_compiled_fns) = }")

    torch.testing.assert_close(actual, expected)


if __name__ == "__main__":
    main()

note that pytorch/ao#1339 is needed at the moment.

Below, I put the console output of the script above:

% python debug_thunderfx_torchao_fp8.py
/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/thunder/dynamo/compiler.py:21: UserWarning: The ThunderCompiler is in active development and may not work as expected. Please report any issues you encounter to the Lightning Thunder team.
  warnings.warn(
len(backend.subgraph_infos) = 1
# len(subgraph.thunder_compiled_fns) = 0
Traceback (most recent call last):
  File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/debug_thunderfx_torchao_fp8.py", line 34, in <module>
    main()
  File "/home/mkozuki/ghq/github.com/Lightning-AI/lightning-thunder/debug_thunderfx_torchao_fp8.py", line 30, in main
    torch.testing.assert_close(actual, expected)
  File "/home/mkozuki/ghq/github.com/crcrpar/pytorch/torch/testing/_comparison.py", line 1530, in assert_close
    raise error_metas[0].to_error(msg)
AssertionError: Tensor-likes are not close!

Mismatched elements: 388 / 1024 (37.9%)
Greatest absolute difference: 0.18639898300170898 at index (1, 61) (up to 1e-05 allowed)
Greatest relative difference: 1.9664803743362427 at index (10, 33) (up to 1.3e-06 allowed)

So it seems that thunder.jit isn't used for this program but the numeric is diverging.

[kshitij12345]

Can you check the result to see if they stay the same between different invocations. (Maybe due to low precision, the results could be different).

expected = fp8_model(x)
actual = fp8_model(x)
torch.testing.assert_close(actual, expected)

[t-vi]

But please add a comment why expected and actual are both from calling the same model rather than one model and a reference.

[crcrpar]

For thunderfx, I updated the test to check the parity between inductor and ThunderCompiler, not eager and ThunderFX.

[crcrpar]

I feel #1539 related in that both of them somehow mistakenly push things to the fallback, not thunder.jit.