[v1] fix non-persistent buffer sync for init_on_rank0 (#10820)

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