From ac8a6a191e42d7310ba0e84dfd15ef9b540c011a Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Wed, 8 Jul 2026 11:59:18 +0000 Subject: [PATCH] Fix tile-wise fp8 FSDP ep8 all-gather --- tests/engine/test_moe_train_engine_float8.py | 15 ++++++++--- xtuner/v1/float8/fsdp_utils.py | 27 ++++++++++++++++++-- 2 files changed, 36 insertions(+), 6 deletions(-) diff --git a/tests/engine/test_moe_train_engine_float8.py b/tests/engine/test_moe_train_engine_float8.py index 14996e5cf0..8a20c1515e 100644 --- a/tests/engine/test_moe_train_engine_float8.py +++ b/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) diff --git a/xtuner/v1/float8/fsdp_utils.py b/xtuner/v1/float8/fsdp_utils.py index 644be106ec..548bee65f1 100644 --- a/xtuner/v1/float8/fsdp_utils.py +++ b/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(