add FSDP and TP tests for Float8BlockwiseLinear by iamzainhuda · Pull Request #4295 · pytorch/ao

iamzainhuda · 2026-04-17T20:13:21Z

Summary

This adds distributed coverage for Float8BlockwiseLinear on top of the DTensor sharding rules we already added for the blockwise FP8 kernels. The goal is to validate the full module-level training path before wiring the kernels into a real torchtitan run.

Added a torchao::blockwise_scaled_mm wrapper around the fallback _scaled_mm path and registered DTensor sharding for it. This was necessary because TP on Float8BlockwiseLinear(use_triton=False) routes through aten._scaled_mm_v2, and that op did not have the DTensor sharding propagation we needed. Without this, the Triton path worked under TP, but the _scaled_mm fallback path failed.

Added the following tests to the following files:

test/prototype/blockwise_fp8_training/test_dtensor.py
- test_linear_tp_parity
  - Adds a 2-GPU tensor-parallel parity test for a quantized ToyModel.
  - Uses the standard FFN TP plan:
    - w1: colwise
    - w2: colwise
    - out_proj: rowwise
  - Runs forward, backward, and optimizer step parity against an unsharded reference for both:
    - _scaled_mm
    - Triton GEMM
test/prototype/blockwise_fp8_training/test_fsdp2.py
- test_fsdp2_parity
  - Add a 2-GPU FSDP2 parity test for a quantized ToyModel.
  - fully_shard() is applied on the FFN submodules plus the parent module.
test/prototype/blockwise_fp8_training/test_fsdp2_tp.py
- Add a manual 4-GPU FSDP2 + TP parity test on a 2D mesh with (dp, tp) = (2, 2).
- Apply TP on the FFN first, then FSDP2 on the DP mesh.
- This is the closest test to how it will be integrated into torchtitan.

All of the three tests use the same testing metholdogy:

build identical reference and distributed ToyModel instances
quantize both with Float8BlockwiseLinearConfig()
run both execution paths by toggling use_triton
feed the same per-replica inputs to reference and distributed models
compare:
- forward outputs
- losses
- gradients after the expected distributed reduction
- parameter values after optimizer step

Testing

2-GPU DTensor + TP coverage:

NCCL_SOCKET_IFNAME=lo pytest test/prototype/blockwise_fp8_training/test_dtensor.py -v

2-GPU FSDP2 coverage:

NCCL_SOCKET_IFNAME=lo pytest test/prototype/blockwise_fp8_training/test_fsdp2.py -v

4-GPU FSDP2 + TP coverage:

NCCL_SOCKET_IFNAME=lo NCCL_DEBUG=WARN torchrun --standalone --nproc_per_node=4 test/prototype/blockwise_fp8_training/test_fsdp2_tp.py

pytorch-bot · 2026-04-17T20:13:24Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 1 New Failure, 8 Pending

As of commit 96989e1 with merge base 15f2e97 ():

NEW FAILURE - The following job has failed:

Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio mslk --index-url https://download.... / linux-job (gh)
RuntimeError: Command docker exec -t 742e7917a1e565efd4b772329dbbe0df8ad3dcba3fffd93e9c8a21d48c1b9ecd /exec failed with exit code 1

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

iamzainhuda · 2026-04-20T22:23:46Z

+
+
+@torch.library.custom_op("torchao::blockwise_scaled_mm", mutates_args=())
+def blockwise_scaled_mm(


moved here because i wanted to register sharding rules for the scaled_mm. felt it made sense to move these definitions in kernel.py

danielvegamyhre

LGTM with some minor comments

danielvegamyhre · 2026-04-22T22:23:12Z


 EPS = 1e-12

+_SCALING_TYPE = getattr(F, "ScalingType", getattr(torch._C, "_ScalingType", None))


why is this necessary? torch changes between versions?

yeah, this is to guard against different torch versions. older pytorch has torch._C._ScalingType and newer has torch.nn.functional.ScalingType

danielvegamyhre · 2026-04-22T22:23:39Z


+_SCALING_TYPE = getattr(F, "ScalingType", getattr(torch._C, "_ScalingType", None))
+BLOCKWISE_1X128_SCALING_TYPE = (
+    _SCALING_TYPE.BlockWise1x128 if _SCALING_TYPE is not None else 4


what are the magic variables 4/5? can we define them as a constant or at least add a comment to make it more easily interpretable

danielvegamyhre · 2026-04-22T22:28:06Z

+
+
+def _pad_blockwise_128x128_scale_k_major(scale: torch.Tensor) -> torch.Tensor:
+    # cuBLASLt requires BLK128x128 scales to be K-major with the K stride padded to a multiple of 4.


Does the multiple of 4 requirement stem from TMA where the tensor's global stride must be a multiple of 16 bytes? i.e., 4 fp32 scales = 16 bytes

If possible, we should explain why this requirement exists in the comment so we have a full understanding

i'm unsure if this is because of TMA but i've updated it to add more detail directly from NVIDIA's cuBLAS docs:

CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F scaling mode, the scaling factors are -major and the stride between the consecutive columns must be a multiple of 4. Let , where the denotes rounding up to the nearest multiple of 4.

reference:
https://docs.nvidia.com/cuda/cublas/index.html#scaling-factors-layouts

danielvegamyhre · 2026-04-22T22:37:00Z

+    if not _is_row_major(a):
+        a = a.contiguous()
+    if not _is_column_major(b):
+        b = b.t().contiguous().t()


can we just assert the data is in the expected layout, instead of silently doing these expensive copies to transform it (which hurts perf)?

i.e., if someone imports and uses blockwise_scaled_mm and calls it with inputs in the wrong layouts, it should fail loudly with clear error messages describing the necessary layouts, rather than "succeed" with bad perf, leaving the user potentially confused why fp8 is not producing speedup vs bf16 when rooflines for the shapes suggest it should.

yup agreed and changed to assertions. oversight on my part where it's unclear to user/hurts perf

danielvegamyhre · 2026-04-22T22:51:41Z

    scale_recipe_a,
    scale_b: torch.Tensor,
    scale_recipe_b,
+    triton_scale_b: torch.Tensor | None = None,


is the new triton_scale_b arg necessary because the Triton gem requires that the B tensor scale be in a different layout than what torch _scaled_mm requires?

if so, in a follow up, we should try to address that - having 2 separate scale_b params for different kernels is pretty gnarly

alternatively, if the fixes in core _scaled_mm are going to land soon (that will enable to stop using the triton kernel at all), then we can leave this triton_scale_b arg for now and just delete all the triton gemm stuff once the fix in core lands.

yeah, agreed its not great. the triton and scaled_mm paths expect the RHS scale tensor in different physical layouts for the grad_weight case. the triton 1x128_128x1 kernel consumes b_s in row-major layout whereas scaled_mm expects the transposed scale layout, so with the arg we let each path receive the layout it expects without adding a copy/transpose on the hot path

i believe the changes in core will be landed soon so we can get rid of the triton gemm work when it does. pytorch/pytorch#180668

ok sounds good. please make sure this is explained in the docstring somewhere.

yup added a comment explaining it

iamzainhuda added 2 commits April 16, 2026 22:16

add fsdp, tp, and fsdp+tp tests

5bf0826

change tests everything is good now

c357e1f

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 17, 2026

iamzainhuda added the module: training quantize_ api training flow label Apr 20, 2026

working clean ver

a7420f6

iamzainhuda marked this pull request as ready for review April 20, 2026 22:22

iamzainhuda requested review from jerryzh168 and vkuzo as code owners April 20, 2026 22:22

iamzainhuda requested a review from danielvegamyhre April 20, 2026 22:22

iamzainhuda commented Apr 20, 2026

View reviewed changes

iamzainhuda added 2 commits April 20, 2026 22:26

ruff formatting

c5ae2e1

fix fsdp+tp test inclusion in CI

da172d5

danielvegamyhre approved these changes Apr 22, 2026

View reviewed changes

danielvegamyhre added this to the FP8 Blockwise Training milestone Apr 24, 2026

clarity

96989e1

iamzainhuda merged commit 9058b58 into main Apr 24, 2026
22 of 23 checks passed

iamzainhuda deleted the fsdp-tp-tests branch April 24, 2026 17:38



		@torch.library.custom_op("torchao::blockwise_scaled_mm", mutates_args=())
		def blockwise_scaled_mm(


		EPS = 1e-12

		_SCALING_TYPE = getattr(F, "ScalingType", getattr(torch._C, "_ScalingType", None))



		def _pad_blockwise_128x128_scale_k_major(scale: torch.Tensor) -> torch.Tensor:
		# cuBLASLt requires BLK128x128 scales to be K-major with the K stride padded to a multiple of 4.

Uh oh!

Conversation

iamzainhuda commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing

Uh oh!

pytorch-bot Bot commented Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❗ 1 Active SEVs

❌ 1 New Failure, 8 Pending

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danielvegamyhre left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danielvegamyhre Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

iamzainhuda commented Apr 17, 2026 •

edited

Loading

pytorch-bot Bot commented Apr 17, 2026 •

edited

Loading

danielvegamyhre Apr 22, 2026 •

edited

Loading