Fix tile-wise fp8 FSDP parallel size 1 all-gather#1948
Conversation
|
@claude review |
jayhenry
changed the title
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| @@ -102,7 +104,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]) | |||
There was a problem hiding this comment.
Claude: Nit: The ep_size=8 case compares against the same losses_ref as ep_size=1 but with 12x higher relative tolerance (rtol=0.12 vs 0.01). At 12% tolerance the assertion becomes quite loose — it mostly checks that loss decreases, not that it matches a known-good trajectory.
Consider capturing a dedicated losses_ref from a validated ep_size=8 run and tightening rtol back down. That would catch regressions that a 12% window would miss (e.g., subtle numerical drift in expert routing under FP8).
| 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) |
There was a problem hiding this comment.
Claude: Good catch — this comment accurately describes the hazard and the fix is correct.
One subtle detail worth noting: the .contiguous() call at line 438 does not guarantee a copy — when the tensor is already contiguous it returns self. So in the common case (shape[0] >= 128 and shape[0] % 128 == 0), scale after .contiguous() still aliases _precomputed_scale, making this inner_tensors guard load-bearing (not just defensive). Well done.
|
Claude: ## Summary Targeted fix for FSDP2's all-gather bypass when shard mesh size is 1 (ep_size=8 configuration). When IssuesNit
VerdictAPPROVE — The fix is correct, well-scoped, and the reasoning in the comments is accurate. The inner_tensors guard is load-bearing (not just defensive) because |
2 participants
Summary
test_tile_wise_fp8to run withep_size=8._precomputed_scaleas an FSDP-managed inner tensor, since FSDP frees inner tensor storage after reshard.Test
python -m py_compile tests/engine/test_moe_train_engine_float8.py xtuner/v1/float8/fsdp_utils.pypytest -v -s tests/engine/test_moe_train_engine_float8.py::TestMoEEngineFloat8::test_tile_wise_fp81 passed, 31 warnings in 85.44s