diff --git a/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py b/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py index e6da80445..ee539ac01 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py @@ -373,6 +373,8 @@ class FSDP2Engine: init_mode = getattr(model, "_init_mode", "init_on_default") if init_mode == "init_on_rank0": + non_persistent_buffers = self._save_non_persistent_buffers(model) if self.rank == 0 else {} + if getattr(model.config, "tie_word_embeddings", False): model.tie_weights() @@ -391,6 +393,13 @@ class FSDP2Engine: # Broadcast the full state dict from the global rank-0 process to all ranks in this group. options = StateDictOptions(full_state_dict=True, cpu_offload=True, broadcast_from_rank0=True) set_model_state_dict(model, full_sd, options=options) + self._restore_non_persistent_buffers(model, non_persistent_buffers) + if self.world_size > 1: + for module in model.modules(): + for buffer_name in sorted(module._non_persistent_buffers_set): + buffer = getattr(module, buffer_name, None) + if buffer is not None: + torch.distributed.broadcast(buffer, src=0) if self.rank == 0: logger.info("init_on_rank0 sync complete.") diff --git a/tests_v1/plugins/trainer_plugins/distributed/test_fsdp2.py b/tests_v1/plugins/trainer_plugins/distributed/test_fsdp2.py index d1b0d7bf8..cf4e6718f 100644 --- a/tests_v1/plugins/trainer_plugins/distributed/test_fsdp2.py +++ b/tests_v1/plugins/trainer_plugins/distributed/test_fsdp2.py @@ -18,12 +18,15 @@ Validates that the FSDP2 meta loading path behaves correctly for tied weights and non-persistent buffers by comparing it with the standard non-meta path. """ +from types import SimpleNamespace + import torch from transformers import AutoConfig from llamafactory.v1.accelerator.interface import DistributedInterface from llamafactory.v1.config.arg_parser import get_args from llamafactory.v1.core.model_engine import ModelEngine +from llamafactory.v1.plugins.trainer_plugins.distributed import fsdp2 as fsdp2_module from llamafactory.v1.plugins.trainer_plugins.distributed.fsdp2 import FSDP2Engine @@ -42,6 +45,50 @@ def collect_non_persistent_buffers(model): return result +def test_fsdp2_rank0_syncs_non_persistent_buffers(monkeypatch): + expected = torch.tensor([1.0, 0.5, 0.25]) + + class Rank0Model(torch.nn.Module): + def __init__(self, rank): + super().__init__() + self.config = SimpleNamespace(tie_word_embeddings=False) + self._init_mode = "init_on_rank0" + self.weight = torch.nn.Parameter(torch.ones(1)) + buffer = expected.clone() if rank == 0 else torch.empty_like(expected, device="meta") + self.register_buffer("inv_freq", buffer, persistent=False) + + def to_empty(self, *, device, recurse=True): + self.inv_freq = torch.zeros_like(expected, device=device) + return self + + def set_model_state_dict(model, state_dict, **kwargs): + assert torch.count_nonzero(model.inv_freq) == 0 + + monkeypatch.setattr(fsdp2_module, "get_current_accelerator", lambda: torch.device("cpu")) + monkeypatch.setattr(fsdp2_module, "set_model_state_dict", set_model_state_dict) + + for rank in (0, 1): + engine = object.__new__(FSDP2Engine) + engine.rank = rank + engine.world_size = 2 + monkeypatch.setattr(engine, "prepare_model", lambda model: model) + monkeypatch.setattr(engine, "_warmup_grad_norm", lambda model: None) + + def broadcast(buffer, src): + assert src == 0 + if rank == 0: + assert torch.equal(buffer, expected) + else: + assert torch.count_nonzero(buffer) == 0 + buffer.copy_(expected) + + monkeypatch.setattr(torch.distributed, "broadcast", broadcast) + model = engine.shard_model(Rank0Model(rank)) + + assert "inv_freq" not in model.state_dict() + assert torch.equal(model.inv_freq, expected) + + def test_fsdp2_meta_loading_buffers_and_tied_weights(): """Verify non-persistent buffers and tied weights consistency after meta load.""" # 1. Initialize DistributedInterface for single process