[moe training] set token group alignment size to 16 for fp8 training test

[danielvegamyhre]

In pytorch/torchtitan#1503 the default TOKEN_GROUP_ALIGNMENT_SIZE_M was changed from 16 (required for fp8) to 8 (minimum for bf16). See PR description for details.

Thus, in our fp8 training tests, we need to set it to 16. This is required so that
each logically distinct gemm in the grouped gemm grad_weight = grad_output_t @ input
has the contraction dim be divisible by 16. 16 byte alignment is required for the slowest moving dim (stride 1), so 16 bytes / 1 byte per element in fp8 = 16 elements.

Test plan

Test: pytest test/prototype/moe_training/test_training.py

Error without change:

E       torch.AcceleratorError: CUDA error: device-side assert triggered
E       Search for `cudaErrorAssert' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.
E       Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

torchao/prototype/moe_training/scaled_grouped_mm.py:259: AcceleratorError
---------------------------------------------------------------------------------- Captured stderr call ----------------------------------------------------------------------------------
/pytorch/aten/src/ATen/native/cuda/GroupMMCommon.cuh:64: prepare_grouped_gemm_data: block: [0,0,0], thread: [0,0,0] Assertion `delta % align == 0 && "expected input tensor dynamic dimension byte size to be non-negative multiple of 16\n"` failed.
/pytorch/aten/src/ATen/native/cuda/GroupMMCommon.cuh:64: prepare_grouped_gemm_data: block: [0,0,0], thread: [1,0,0] Assertion `delta % align == 0 && "expected input tensor dynamic dimension byte size to be non-negative multiple of 16\n"` failed.
/pytorch/aten/src/ATen/native/cuda/GroupMMCommon.cuh:64: prepare_grouped_gemm_data: block: [0,0,0], thread: [2,0,0] Assertion `delta % align == 0 && "expected input tensor dynamic dimension byte size to be non-negative multiple of 16\n"` failed.
/pytorch/aten/src/ATen/native/cuda/GroupMMCommon.cuh:64: prepare_grouped_gemm_data: block: [0,0,0], thread: [3,0,0] Assertion `delta % align == 0 && "expected input tensor dynamic dimension byte size to be non-negative multiple of 16\n"` failed.
______________________________________________________________

With change, tests pass.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 1 Cancelled Job

As of commit 3a87a56 with merge base 7dbc816 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh)
  test/test_low_bit_optim.py::TestFSDP2::test_uneven_shard
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/test_low_bit_optim.py::TestFSDP2::test_uneven_shard
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  Process completed with exit code 1.

CANCELLED JOB - The following job was cancelled. Please retry:

• Run TorchAO Experimental Tests / test-cpu-ops (linux.arm64.2xlarge) (gh)
  ##[error]The operation was canceled.

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