From 974eca1c12d028c97508bd5430ae27d2eafa8e79 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Sat, 10 Oct 2026 18:01:47 +0000 Subject: [PATCH 1/2] Fix the FSDP2 export deadlock at world > 1 A forward leaves the root FSDP module's parameters unsharded, because FSDP2 skips the root's post-forward reshard, and FSDP2's state_dict() pre-hook reshards a module's whole parameter group. In collect_export_tensors only the unit's owner calls state_dict(), so the owner of the non-layer unit resharded the root on its rank alone. The next gather window then all-gathered parameters on that rank that the other ranks still held as plain tensors, a collective they never joined, and the export hung (rank 0 in collect_export_tensors, the rest at the barrier). Both the streaming (owner="share") and rank-0 (owner="rank0") FSDP2 exports were affected. Reshard every FSDP module on every rank before the loop. It is local, and it makes the owner's pre-hook a no-op. The two GPU tests that time out in the nightly 2-GPU run, test_fsdp2_streaming_export_matches_reference[nvfp4|fp8], pass again. Add a two-rank CPU/gloo unit test, which runs in PR CI, so the multi-rank path no longer depends on the nightly to catch a regression. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- .../export/unified_export_hf_streaming.py | 9 +++++ .../export/test_fsdp2_parallel_export.py | 38 ++++++++++++++++++- 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 9188bb5780e..26896da92f5 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -550,6 +550,15 @@ def collect_export_tensors( from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _assert_fsdp2_owns_every_mesh_dim(model) + # A forward leaves the root's parameters unsharded (FSDP2 skips the root's post-forward + # reshard), and FSDP2's state_dict() pre-hook reshards a module's whole parameter group. Only + # the owner calls state_dict() below, so it would reshard the root on its rank alone, and the + # next gather window would then gather parameters on that rank that the others see as plain + # tensors -- a collective they never join. Resharding locally on every rank first makes the + # pre-hook a no-op. + for module in model.modules(): + if isinstance(module, FSDPModule): + module.reshard() my_rank, world = _dist.rank(), _dist.size() names = module_name_maps(model) ctx = ExportContext( diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index 18846c9844d..503fbbd05a6 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -18,16 +18,22 @@ The pipeline is world-agnostic: with ``torch.distributed`` uninitialized, ``rank/size`` are ``0/1``, the unshard window is a no-op, and rank 0 owns every unit. So these cover unit enumeration, per-tensor export, the shard writer, index naming and tied-weight dropping without a -GPU. The multi-rank unshard collective needs real FSDP2 and is covered in ``tests/gpu``. +GPU. ``test_streaming_export_two_cpu_ranks`` adds one two-rank gloo run; the multi-rank paths +on GPU are covered in ``tests/gpu``. """ import copy +import faulthandler import json +from functools import partial from pathlib import Path import pytest import torch +from _test_utils.torch.distributed.utils import spawn_multiprocess_job from safetensors.torch import load_file +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard import modelopt.torch.opt as mto import modelopt.torch.quantization as mtq @@ -38,7 +44,7 @@ from transformers import AutoModelForCausalLM, LlamaConfig -def _tiny_quantized_llama(quant_cfg=None, tie=True): +def _tiny_quantized_llama(quant_cfg=None, tie=True, mesh=None): mto.enable_huggingface_checkpointing() torch.manual_seed(0) cfg = LlamaConfig( @@ -52,6 +58,10 @@ def _tiny_quantized_llama(quant_cfg=None, tie=True): architectures=["LlamaForCausalLM"], ) model = AutoModelForCausalLM.from_config(cfg).eval() + if mesh is not None: + for layer in model.model.layers: + fully_shard(layer, mesh=mesh) + fully_shard(model, mesh=mesh) calib = [torch.randint(0, 128, (1, 8)) for _ in range(4)] return mtq.quantize(model, quant_cfg or mtq.FP8_DEFAULT_CFG, lambda m: [m(x) for x in calib]) @@ -143,3 +153,27 @@ def test_streaming_export_extra_state_dict_mtp(tmp_path): loaded = _load_all(d) assert "model.layers.99.mtp.weight" in loaded assert torch.equal(loaded["model.layers.99.mtp.weight"], extra["model.layers.99.mtp.weight"]) + + +def _two_rank_streaming_export(rank, size, export_dir): + # A deadlocked rank exits with its stack instead of hanging the run on the join. + faulthandler.dump_traceback_later(45, exit=True) + # Calibrating after sharding leaves the root's parameters unsharded: FSDP2 skips the root's + # post-forward reshard. That is the state in which the owner's state_dict() used to reshard + # them on its rank alone, so the next gather window ran a collective the other rank skipped. + model = _tiny_quantized_llama(tie=False, mesh=init_device_mesh("cpu", (size,))) + _export_fsdp2_checkpoint_streaming(model, torch.bfloat16, export_dir=export_dir) + torch.distributed.barrier() + faulthandler.cancel_dump_traceback_later() + if rank != 0: + return + expected, _ = _export_transformers_checkpoint(_tiny_quantized_llama(tie=False), torch.bfloat16) + exported = _load_all(Path(export_dir)) + assert set(exported) == set(expected) + for key, tensor in expected.items(): + assert torch.equal(exported[key].float(), tensor.float()), key + + +def test_streaming_export_two_cpu_ranks(tmp_path, skip_on_windows): + """Two gloo ranks export what a single process does, without a rank waiting on the other.""" + spawn_multiprocess_job(2, partial(_two_rank_streaming_export, export_dir=str(tmp_path))) From 143a71c165a4bc28068c9a04cb5728e305bb8287 Mon Sep 17 00:00:00 2001 From: Chenjie Luo Date: Mon, 12 Oct 2026 21:31:21 +0000 Subject: [PATCH 2/2] Build the two-rank test's reference on every rank The reference checkpoint was built and exported on rank 0 alone, after rank 1 had returned, while both still belonged to the two-rank process group. Nothing on that path issues a collective today, but one would wait on a partner that has already left. Build and export it on every rank before the sharded export, so any collective it issues has a partner and the deadlock guard covers it. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Chenjie Luo --- tests/unit/torch/export/test_fsdp2_parallel_export.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/export/test_fsdp2_parallel_export.py b/tests/unit/torch/export/test_fsdp2_parallel_export.py index 503fbbd05a6..997a23020e5 100644 --- a/tests/unit/torch/export/test_fsdp2_parallel_export.py +++ b/tests/unit/torch/export/test_fsdp2_parallel_export.py @@ -158,6 +158,9 @@ def test_streaming_export_extra_state_dict_mtp(tmp_path): def _two_rank_streaming_export(rank, size, export_dir): # A deadlocked rank exits with its stack instead of hanging the run on the join. faulthandler.dump_traceback_later(45, exit=True) + # The unsharded reference is built and exported on every rank, so any collective its + # calibration or export issues has a partner. + expected, _ = _export_transformers_checkpoint(_tiny_quantized_llama(tie=False), torch.bfloat16) # Calibrating after sharding leaves the root's parameters unsharded: FSDP2 skips the root's # post-forward reshard. That is the state in which the owner's state_dict() used to reshard # them on its rank alone, so the next gather window ran a collective the other rank skipped. @@ -167,7 +170,6 @@ def _two_rank_streaming_export(rank, size, export_dir): faulthandler.cancel_dump_traceback_later() if rank != 0: return - expected, _ = _export_transformers_checkpoint(_tiny_quantized_llama(tie=False), torch.bfloat16) exported = _load_all(Path(export_dir)) assert set(exported) == set(expected) for key, tensor in expected.items():