Rename torchao.float8.Float8Tensor to torchao.float8.Float8TrainingTensor

[jerryzh168]

Stacked PRs:

• Add all fbgemm kernel Tensors into Int4WeightOnlyConfig and Float8DynamicActivationInt4WeightConfig #2474
• Add Float8Tensor #2463
• ->Rename torchao.float8.Float8Tensor to torchao.float8.Float8TrainingTensor #2479
• Add support for float8 activation for Int4GroupwisePreshuffleTensor #2437

---

Rename torchao.float8.Float8Tensor to torchao.float8.Float8TrainingTensor

Summary:
att, since we are introducing a inference version Float8Tensor

Test Plan:
regression tests for float8 training: pytest test/float8

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit b39be5d with merge base 1fd34e4 ():

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

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

• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh) (trunk failure)
  Process completed with exit code 127.

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

[vkuzo]

can we test two more things:

1. run the overall test script for float8 training

./test/float8/test_everything.sh

2. verify torchtitan training still works: https://github.com/pytorch/torchtitan/blob/main/docs/float8.md

Also, there are some callsites for Float8Tensor in fbcode, so we should make sure this diff is landed in a way where we could fix them in the fbcode diff.

[danielvegamyhre]

can we test two more things:

1. run the overall test script for float8 training

./test/float8/test_everything.sh

2. verify torchtitan training still works: https://github.com/pytorch/torchtitan/blob/main/docs/float8.md

Also, there are some callsites for Float8Tensor in fbcode, so we should make sure this diff is landed in a way where we could fix them in the fbcode diff.

+1, for torchtitan i would suggest (1) 2d parallel w/ FSDP+TP in torchtitan for all 3 recipes, (2) for tensorwise test enabling fp8 all gather enabled and enabling precompute fp8 weight scales

[jerryzh168]

Also, there are some callsites for Float8Tensor in fbcode, so we should make sure this diff is landed in a way where we could fix them in the fbcode diff.

sg, in this case I can do a forward fix in diff train

[jerryzh168]

./test/float8/test_everything.sh

Runs without errors.

for torchtitan I have tested:

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --model.converters="float8" --float8.enable_fsdp_float8_all_gather --float8.precompute_float8_dynamic_scale_for_fsdp --float8.force_recompute_fp8_weight_in_bwd --training.compile

and

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --model.converters="float8" --float8.recipe_name rowwise --training.compile

CONFIG_FILE="./torchtitan/models/llama3/train_configs/llama3_8b.toml" ./run_train.sh --model.converters="float8" --float8.recipe_name rowwise_with_gw_hp --training.compile

is this enough?