Add MoE blockwise FP8 DTensor and compile coverage

[iamzainhuda]

Summary

• add DTensor TP support for blockwise FP8 MoE grouped GEMM, covering fwd, dgrad, and wgrad placements
• add custom-op wrappers and DTensor sharding rules for the emulated and DeepGEMM grouped GEMM paths so torch.compile can hit the production kernels
• add end-to-end FSDP2/TP coverage matching the blockwise FP8 linear tests, plus shared DTensor test utilities reused by both paths

this keeps the MoE path aligned with the blockwise FP8 linear DTensor design: shared DTensor helpers handle local tensor unwrap/rewrap, replicated metadata, placement checks, and common test parity utilities. The grouped MoE code still exercises the quantized training path rather than a BF16-only fallback.

Part of #3290

Tests

pytest -q test/prototype/moe_training/test_config.py
NCCL_SOCKET_IFNAME=lo pytest -q test/prototype/moe_training/test_fp8_blockwise_grouped_mm.py
NCCL_SOCKET_IFNAME=lo pytest -q test/prototype/moe_training/test_fp8_blockwise_fsdp2.py
NCCL_SOCKET_IFNAME=lo pytest -q test/prototype/blockwise_fp8_training/test_dtensor.py
NCCL_SOCKET_IFNAME=lo pytest -q test/prototype/blockwise_fp8_training/test_fsdp2.py

For the broader related coverage:

NCCL_SOCKET_IFNAME=lo pytest -q test/prototype/moe_training/test_fp8_blockwise_deepgemm_backend.py
NCCL_SOCKET_IFNAME=lo pytest -q test/prototype/moe_training/test_fp8_blockwise_fsdp2_tp.py

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 9140d1c with merge base cb76f29 ():

NEW FAILURES - The following jobs have failed:

• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  RuntimeError: Command docker exec -t 16a63e33a0916fabdeed17196cf7b6bb31a7e5d20e95f189a1941cc5a2ab5678 /exec failed with exit code 1
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch torchvision --index-url htt... / linux-job (gh)
  test/test_low_bit_optim.py::TestFSDP2::test_uneven_shard

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

[iamzainhuda]

@claude review

[claude]

Claude finished @iamzainhuda's task in 8m 34s —— View job

---

PR Review: Add MoE blockwise FP8 DTensor TP support

• Read changed source files
• Review correctness, bugs, and edge cases
• Review tests
• Post review feedback

Overall this is a well-structured change. The DTensor plumbing is cleanly factored into dtensor_utils.py, the non-distributed path is preserved by making every _replicate_like_dtensor/_local_tensor helper a no-op on plain tensors, and the custom-op + register_sharding registrations are consistent with the documented operand layouts. Comments below are mostly minor / questions.

---

Correctness & design

1. Lost kernel-contract documentation in the wgrad custom op (deepgemm_grouped_kernels.py ~L744)
The refactor that moved the wgrad body into _deepgemm_blockwise_scaled_grouped_mm_wgrad_custom_op dropped the detailed comment explaining why accum = torch.zeros_like(out_fp32) is required (the SM90 sm90_k_grouped_fp8_gemm_1d1d kernel hard-asserts c.has_value(), reads c/writes d, so the buffers must be distinct and c zero-seeded). The behavior is preserved but the rationale is now gone — please re-add that comment, since it's the kind of non-obvious invariant that's easy to "optimize away" later.

2. Output shape now derived from scale shapes (deepgemm_grouped_kernels.py L744-748)
The allocation changed from (len(group_sizes), lhs.dim, rhs.dim) to (ks_tensor.numel(), a_s.shape[0], b_s.shape[0]). This relies on a_s.shape[0] == lhs.dim and b_s.shape[0] == rhs.dim after _prepare_k_grouped_*_operand (the RHS scale is even .transpose'd earlier). It looks correct given the prep functions, but an explicit assert tying a_s.shape[0]/b_s.shape[0] back to the operand .dim would make this self-documenting and catch layout regressions early.

3. ks_tensor.tolist() D2H sync inside the custom op (deepgemm_grouped_kernels.py L750)
group_sizes = ks_tensor.tolist() forces a device→host sync. This matches prior behavior (the old offset_plan.group_sizes property also did .tolist()), and being inside a custom_op it's opaque to torch.compile, so it's fine — just flagging since the config docstring elsewhere calls out avoiding D2H syncs for compile. No action needed.

4. Duplicate Shard(1) output strategies in the emulated sharding rule (grouped_kernels.py _emulated_grouped_mm_shardings)
This op serves both b.ndim==3 (fwd/dgrad, out (M,N)) and b.ndim==2 (wgrad, out (E,N,K)), and the strategy list contains two entries that both produce output [Shard(1)] with different input placements. DTensor's propagation disambiguates on the input placements, so this should be fine, but a single static strategy list spanning two different output ranks is fragile. Worth a comment noting the rule is selected by input-placement match, not output shape.

5. Sharding maps look internally consistent. I cross-checked the swap between the DeepGEMM and non-DeepGEMM quant ops: non-DeepGEMM forward (_grouped_transposed_rhs) outputs (E,K,N) → preserve shardings, while DeepGEMM forward outputs (E,N,K) → transpose_kn shardings (and vice-versa for dgrad). That matches the docstrings and the K/N transpose in the DeepGEMM kernels. 👍

---

Minor / nits

• dtensor_utils.py: imports mix torch.distributed._tensor (private) and torch.distributed.tensor (public). DTensor is importable from the public torch.distributed.tensor as well — prefer the public path consistently to avoid the _tensor private import. Same mix appears in grouped_mm.py.
• _replicate_like_dtensor(A, B_t) then _replicate_like_dtensor(B_t, A) (grouped_mm.py L238-239): correct, but a one-line comment that the second call is a no-op once A is already lifted would help readers (the ordering dependency isn't obvious).
• config.py: Float8TrainingRecipe.FP8_BLOCKWISE forces pad_token_groups_for_grouped_mm=True and kernel_preference=AUTO. Given _to_fp8_blockwise_then_scaled_grouped_mm asserts kernel_preference in (AUTO, EMULATED), a user passing any other KernelPreference via the config will hit an assert rather than a clear error — consider validating in __post_init__ / from_recipe.

---

Tests

• test_fp8_blockwise_grouped_mm.py (+600) gives solid single-device coverage; good.
• test_fp8_blockwise_fsdp2_tp.py is explicitly a manual 4-GPU torchrun test that does not run in CI (WORLD_SIZE != 4 skip). That's a reasonable call for a multi-GPU TP path, but it means the DTensor sharding rules — the bulk of this PR — are not exercised by CI. Consider whether a 2-GPU gated job (or a single-process fake-mesh test that just validates register_sharding propagation) is feasible, so these strategy tables don't silently rot.
• _fp8_blockwise_distributed_test_utils.py: the reference path uses torch._grouped_mm while the dist path uses the quantized config — comparing fp8 against bf16 reference via SQNR thresholds (20-30) is the right approach. 👍

No blocking issues from my read; (1) and (5) are the most important to address (restore the contract comment; the sharding logic looks right). The lack of CI coverage for the TP path is the main risk.
• branch fp8-blockwise-moe-dtensor-tp

[iamzainhuda]

@claude review this PR again

[claude]

Claude finished @iamzainhuda's task in 2m 13s —— View job

---

PR Re-Review: Add MoE blockwise FP8 DTensor TP support

• Read changed source files
• Check whether prior review feedback was addressed
• Review correctness, bugs, and edge cases
• Post review feedback

The branch is still at a single commit (0db5f7f), so this re-review revisits the same code as before. Below is the current status of each earlier point plus a few fresh observations.

---

Status of prior feedback

1. (Still open) Missing accum/c.has_value() contract comment — deepgemm_grouped_kernels.py:744-749
The _deepgemm_blockwise_scaled_grouped_mm_wgrad_custom_op still allocates out_fp32 and accum = torch.zeros_like(out_fp32) with no comment explaining why a separate zero-seeded c buffer is mandatory (the SM90 sm90_k_grouped_fp8_gemm_1d1d kernel hard-asserts c.has_value(), reads c and writes d, so the two buffers must be distinct and c zero-initialized). This is the most important item — it's a non-obvious invariant that's easy to "optimize away" later. Please restore it. Fix this →

2. (Partially addressed) Output shape derived from scale shapes — deepgemm_grouped_kernels.py:744-748
The launcher deepgemm_blockwise_scaled_grouped_mm_wgrad now asserts a.numel() == total_tokens * lhs.dim / b.numel() == total_tokens * rhs.dim (lines 704-709), which is good. But the custom_op body itself still derives the output as (ks_tensor.numel(), a_s.shape[0], b_s.shape[0]) with no assert tying a_s.shape[0]/b_s.shape[0] to the operand dims. Under torch.compile the op is the dispatch boundary, so a one-line assert here would make the contract self-documenting and catch a layout regression at the kernel edge.

3. (Resolved, no action) ks_tensor.tolist() D2H sync — matches prior behavior and is opaque to compile inside the custom op. Fine.

4. (Addressed) Emulated sharding rule comments — grouped_kernels.py:65-148
_emulated_grouped_mm_shardings now documents the b.ndim==3 (fwd/dgrad) vs b.ndim==2 (wgrad) cases and labels each strategy. Two entries still emit output [Shard(1)] with different input placements (forward output-feature at L89 vs wgrad output-dim at L119); DTensor disambiguates on input placements so this is correct, and the comments now make that selection intent clear. 👍

5. (Confirmed correct) DeepGEMM vs non-DeepGEMM sharding swap — re-verified: preserve vs transpose_kn shardings line up with the K/N transpose in the forward/dgrad quant ops. 👍

---

Still-open minor items

• Private torch.distributed._tensor import still present in dtensor_utils.py:8 and grouped_mm.py:10 (mixed with the public torch.distributed.tensor). DTensor is importable from the public path — worth unifying. (Note _fp8_blockwise_distributed_test_utils.py:13 imports it the private way too.)
• config.py kernel_preference validation — Float8TrainingOpConfig.from_recipe sets kernel_preference=AUTO, but a user can still construct the config with any other KernelPreference. Since _to_fp8_blockwise_then_scaled_grouped_mm only accepts AUTO/EMULATED, a different value surfaces as a bare assert deep in the op rather than a clear config-time error. A check in __post_init__ when fp8_grouped_mm_recipe == "blockwise" would fail fast with a useful message. Fix this →

---

Fresh observations

• grouped_mm.py:238-239 — the lift comment ("TP commonly shards only the expert weights, so lift replicated activations/metadata onto the same mesh before dispatch") was added and reads well. The ordering dependency _replicate_like_dtensor(A, B_t) then _replicate_like_dtensor(B_t, A) is now understandable from that comment. Good.
• CI coverage gap (unchanged risk) — test_fp8_blockwise_fsdp2_tp.py remains gated to a manual 4-GPU torchrun run (WORLD_SIZE != 4 → skip). The DTensor sharding rules and custom-op register_sharding tables — the bulk of this PR — are therefore not exercised by CI. The single-device coverage in test_fp8_blockwise_grouped_mm.py (+600) is solid for the kernels, but the sharding strategy tables can silently rot. A single-process fake-mesh test that just validates register_sharding propagation (no real collectives, runs on 1 GPU) would catch placement-table regressions cheaply — worth considering even if the full 2D-mesh test stays manual.
• CI failures — the two reported failures (test_low_bit_optim.py::TestFSDP2::test_uneven_shard and a docker exec failure) appear unrelated to this PR's code paths and look like infra/flaky failures. Worth a rebase + re-run to confirm they're not introduced here.

---

No blocking correctness issues from my read; the sharding logic still checks out. The top action item remains (1) — restoring the wgrad kernel-contract comment — followed by the config-validation and CI-coverage suggestions.
• branch fp8-blockwise-moe-dtensor-tp

[vkuzo]

https://github.com/pytorch/ao/tree/main/torchao/testing is the right top level dir for this, test utilities should not be mixed with product code. Same for other files in this PR

[iamzainhuda]

makes sense, moved. kept the suite-specific builders local under test/prototype

[vkuzo]

hmm, this is unexpected. Do we know why this is needed?

[iamzainhuda]

at the blockwise_scaled_mm custom-op boundary: DTensor local dispatch / redistribution can materialize the local RHS shard in a regular contiguous layout even when the logical operand is the transposed FP8 RHS produced by the quantizer, and _scaled_mm_v2 requires a physically column-major RHS

[iamzainhuda]

i actually saw this issue in a job i ran internally, i can make this fix in a seperate PR if thats cleaner

[vkuzo]

do we need TP for experts? what's the motivation for that?