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..997a23020e5 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,29 @@ 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) + # 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. + 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 + 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)))