Fix tile-wise fp8 FSDP parallel size 1 all-gather

tests/engine/test_moe_train_engine_float8.py

@@ -32,15 +32,22 @@
 class TestMoEEngineFloat8(DeterministicDDPTestCase):
 
     @parametrize.parametrize(
-        "device,ep_size,hsdp_sharding_size",
+        "device,ep_size,hsdp_sharding_size,sim_tol,rtol",
         [
-            ("cuda", 1, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8"))),  # todo: test ep8 and hsdp, OOM in 8 gpus
+            ("cuda", 1, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.01),
+            # ep8 is a smoke/trend coverage for the FSDP shard-mesh-size-1 FP8 path.
+            # It shares the ep1 reference below, but is not expected to align step-by-step
+            # because EP changes routing/collective order and accumulates FP8 numeric drift.
+            # Observed 10-step loss:
+            # [2.4714, 2.4714, 1.8044, 1.5210, 0.9570, 0.6952, 0.4370, 0.3123, 0.1714, 0.1100]
+            ("cuda", 8, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.15),
         ],
     )
-    def test_tile_wise_fp8(self, device, ep_size, hsdp_sharding_size):
+    def test_tile_wise_fp8(self, device, ep_size, hsdp_sharding_size, sim_tol, rtol):
         pg = self.create_pg(device)
 
         moe_cfg = Qwen3MoE30BA3Config(
+            ep_size=ep_size,
             balancing_loss_cfg=BalancingLossConfig(),
             float8_cfg=Float8Config(
                 scaling_granularity_gemm=ScalingGranularity.TILEWISE,
@@ -102,7 +109,7 @@ def warmup_fn(x):
         losses = torch.tensor(losses)
         losses_ref = torch.tensor([2.4234, 2.4234, 1.5270, 1.1483, 0.8904, 0.6388, 0.3963, 0.2589, 0.1519, 0.1101])
 
-        self._check_loss_curve(losses, losses_ref, sim_tol=0.01, rtol=0.01)
+        self._check_loss_curve(losses, losses_ref, sim_tol=sim_tol, rtol=rtol)
         torch.cuda.empty_cache()
         try:
             dist.destroy_process_group(pg)

xtuner/v1/float8/fsdp_utils.py

@@ -378,6 +378,11 @@ def __repr__(self):
 
     def fsdp_pre_all_gather(self, mesh):
         if self._precomputed_w is not None:
+            if mesh.size() == 1:
+                # FSDP2 bypasses the collective when shard mesh size is 1 and only
+                # copies all_gather_inputs[0] into all_gather_outputs. Keep scale in
+                # metadata so post_all_gather can recover the full Float8Tensor.
+                return (self._precomputed_w,), (self._precomputed_scale,)
             return (self._precomputed_w, self._precomputed_scale), None
         assert self._precomputed_scale is not None
         if self._tensor.shape[0] >= 128 and self._tensor.shape[0] % 128 == 64:
@@ -396,6 +401,11 @@ def fsdp_pre_all_gather(self, mesh):
                 float8_dtype=self._dtype,
             )
 
+        # FSDP2 bypasses the collective when shard mesh size is 1 and only
+        # copies all_gather_inputs[0] into all_gather_outputs. Keep scale in
+        # metadata so post_all_gather can recover the full Float8Tensor.
+        if mesh.size() == 1:
+            return (w_fp8_data,), (self._precomputed_scale,)
         return (w_fp8_data, self._precomputed_scale), None
 
     def fsdp_post_all_gather(
@@ -406,7 +416,12 @@ def fsdp_post_all_gather(
         *,
         out: Optional[torch.Tensor] = None,
     ):
-        data, scale = all_gather_outputs
+        scale_from_metadata = len(all_gather_outputs) == 1
+        if scale_from_metadata:
+            (data,) = all_gather_outputs
+            (scale,) = metadata
+        else:
+            data, scale = all_gather_outputs
         dim0, dim1 = data.shape
         assert dim1 == self._tensor.shape[1], (
             f"Expected data.shape[1] == self._tensor[1], got data.shape = {data.shape}, self._tensor.shape = {self._tensor.shape}"
@@ -451,13 +466,21 @@ def fsdp_post_all_gather(
             else:
                 raise RuntimeError(f"out must be a Float8Tensor or DTensor(_local_tensor=Float8Tensor), but got {out}")
             return
+        # FSDP frees tensors returned as unsharded inner_tensors after reshard.
+        # If scale came from metadata, it aliases persistent _precomputed_scale
+        # that the next pre_all_gather will reuse. The Python Tensor reference
+        # still exists, but FSDP frees by resizing the backing storage to 0.
+        # Returning (data, scale) here would let FSDP free that persistent
+        # storage after the first reshard, and the next quant cast may read
+        # freed CUDA memory.
+        inner_tensors = (data,) if scale_from_metadata else (data, scale)
         return Float8Tensor(
             data,
             scale,
             param_dtype,
             scaling_granularity=ScalingGranularity.BLOCKWISE,  # tilewise fp8: weight is blockwise quantized
             group_size=128,
-        ), (data, scale)  # afterwards will be freed
+        ), inner_tensors  # afterwards will be freed
 
 
 def tensor_to_per_tensor_fp8_scales(