[SimpleFSDP] SimpleFSDP silently ignores `enable_fsdp_float8_all_gather=True`

[aditvenk]

Bug description

When used with torchao's Float8LinearConfig(enable_fsdp_float8_all_gather=True), SimpleFSDP all-gathers weights in bf16 instead of fp8 — 2× the expected comm bandwidth, no warning. FSDP2 honors the flag correctly.

SimpleFSDP's ReplicateComputation uses DTensor.redistribute(Replicate()), which dispatches to _c10d_functional.all_gather_into_tensor. This op isn't in torchao's _ops_to_preserve_subclass, so WeightWithDynamicFloat8CastTensor is stripped before the collective. FSDP2 avoids this by going through fsdp_pre_all_gather / fsdp_post_all_gather hooks.

Repro: https://gist.github.com/aditvenk/9108e099a75150be64184aba40c984ea

torchrun --nproc_per_node=2 repro.py

SimpleFSDP's ReplicateComputation needs to account for quantized weights analogous to FSDP2.

Repro script is for Float8, but same issue should apply to other types like MXFP8 also

Versions

Reproducible on latest PyTorch nightly.

[aditvenk]

cc @SherlockNoMad , @yiming0416

[aditvenk]

=== SimpleFSDP ===
  Trace: /tmp/fp8_ag_traces/SimpleFSDP_rank0.json
  _c10d_functional::all_gather_into_tensor
    dtypes=['c10::BFloat16', 'Scalar', '']  dims=[[128, 256], [], []]
  c10d::allgather_into_tensor_coalesced_
    dtypes=['TensorList', 'TensorList', '', 'Scalar']  dims=[[[256, 256]], [[128, 256]], [], []]
  nccl:all_gather_into_tensor_coalesced
    dtypes=['c10::BFloat16']  dims=[[128, 256]]
  Status: BF16 all-gather (fp8 flag ignored)

>>> FSDP2 with enable_fsdp_float8_all_gather=True
USDT:2026-04-28 00:24:37 2769:2769 SyncActivityProfilerHandler.cpp:52] profiler_start
USDT:2026-04-28 00:24:37 2770:2770 SyncActivityProfilerHandler.cpp:52] profiler_start
USDT:2026-04-28 00:24:37 2769:2769 SyncActivityProfilerHandler.cpp:59] profiler_stop
USDT:2026-04-28 00:24:37 2770:2770 SyncActivityProfilerHandler.cpp:59] profiler_stop
gpu-dev-722cab1f:2770:2770 [1] NCCL INFO TUNER/Plugin: Closing tuner: 'nccl_ofi_tuner'

=== FSDP2 ===
  Trace: /tmp/fp8_ag_traces/FSDP2_rank0.json
  fsdp::all_gather_copy_in
    dtypes=['TensorList', 'c10::Float8_e4m3fn', 'ScalarList', 'Scalar', 'Scalar']  dims=[[[32768]], [65536], [], [], []]
  c10d::_allgather_base_
    dtypes=['c10::Float8_e4m3fn', 'c10::Float8_e4m3fn', '', 'Scalar', 'Scalar']  dims=[[65536], [32768], [], [], []]
  FSDP::all_gather (fc)
    dtypes=['c10::Float8_e4m3fn']  dims=[[32768]]
  Status: FP8 all-gather active

============================================================
SUMMARY
============================================================
FSDP2      -> FP8 AG
SimpleFSDP -> BF16 AG

BUG CONFIRMED: SimpleFSDP silently ignores enable_fsdp_float8_all_gather=True

[tianyu-l]

We are deprecating this field due to not seeing adoption of tensorwise fp8. cc @acisseJZhong

[acisseJZhong]

landing #3032 soon and this field is removed!

[aditvenk]

@tianyu-l , @acisseJZhong -- isn't this a field in torchao API?
Regardless of whether we are doing tensorwise or channelwise or blockwise, if the gather has to occur on quantized dtype, don't we still need some kind of handling for this in SimpleFSDP's parameterization?

[tianyu-l]

other than tensorwise, other recipes don't always support quantize -> all-gather, so we'll have to do all-gather -> quantize, hence no need of this field.

[weifengpy]

just learned we deprecated tensor-wise float8 all-gather. Are we prioritizing row-wise and block-wise now?

[tianyu-l]

IIUC torchao team has been prioritizing mxfp8, and will work on nvfp4, but is simultaneously building block-wise fp8. Rowwise was built earlier. cc @danielvegamyhre