mx: expose scaling calculation methods in training UX

[vkuzo]

Summary:

To prepare MX for graduating out of prototype, exposes the scaling mode at the top level training config and recipe. The two well supported scaling modes are FLOOR and RCEIL. CEIL has not been well tested, and EVEN does not work with torch.compile yet. We may further adjust these options in future PRs.

Note that for RCEIL, the dim0 casts are not yet using hardware accelerated instructions, so overall performance is currently slightly below FLOOR. We can improve this in a future PR.

Test Plan:

unit tests:

pytest test/prototype/mx_formats/ -s -x

performance on llama 3 8b training:

// requires https://github.com/pytorch/torchtitan/pull/1512

// requires hardcoding dim1 kernel choice to CUDA
with-proxy CONFIG_FILE="torchtitan/models/llama3/train_configs/llama3_8b.toml " ./run_train.sh --model.print_after_conversion --training.compile --training.steps 50 --model.converters mx --mx.recipe_name "mxfp8_cublas"
...tps ~10.5k

with-proxy CONFIG_FILE="torchtitan/models/llama3/train_configs/llama3_8b.toml " ./run_train.sh --model.print_after_conversion --training.compile --training.steps 50 --model.converters mx --mx.recipe_name "mxfp8_cublas_rceil"
...tps ~10.3k

performance on individual cast

(pytorch_nightly) [[email protected] ~/local/ao
(20250728_mx_expose_scale)]$ python benchmarks/mx_formats/cast_bench.py
--mode dim0_mx_floor
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.9.0.dev20250724+cu128
triton version: 3.4.0
mode: dim0_mx_floor
time_us 184.38400328159332
mem_bw_gbps 4413.045391781173
(pytorch_nightly) [[email protected] ~/local/ao
(20250728_mx_expose_scale)]$ python benchmarks/mx_formats/cast_bench.py
--mode dim0_mx_rceil
M 16384 K 16384 BLOCK_SIZE 32
GPU: NVIDIA B200
torch version: 2.9.0.dev20250724+cu128
triton version: 3.4.0
mode: dim0_mx_rceil
time_us 143.39199662208557
mem_bw_gbps 5674.619191924083

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• mx: make CUDA kernel for dim1 cast in mxfp8_cublas recipe #2661
• -> mx: expose scaling calculation methods in training UX #2620

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 29354b7 with merge base d05e54f ():

NEW FAILURE - The following job has failed:

• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  test_replace_q_dq_patterns_with_quantized_linear_ops_pass

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