[dist] fix: preserve persistent buffers during meta random initialization

scripts/ci/accelerator/suites/npu.json

@@ -399,6 +399,16 @@
       "tests/utils/test_save_safetensor_utils.py"
     ]
   },
+  {
+    "name": "fsdp-cpu-offload-native-unit",
+    "suite": "unit",
+    "args": [
+      "-v",
+      "-s",
+      "-x",
+      "tests/utils/test_rank0_load_and_broadcast_weights.py::test_fsdp_cpu_offload_training_keeps_buffers_on_compute_device"
+    ]
+  },
   {
     "name": "npu-e2e-40",
     "suite": "e2e",

tests/checkpoints/test_dcp_checkpointer.py

@@ -505,23 +505,44 @@ def test_build_parallelize_model_forwards_should_skip_hf_weight_load(self, monke
         assert result is parallelized_model
         assert parallelize_fsdp2.call_args.kwargs["should_skip_hf_weight_load"] is True
 
-    def test_dcp_resume_preserves_nonpersistent_buffers_and_forward(self, monkeypatch, tmp_path):
+    def test_cpu_offload_dcp_resume_restores_persistent_and_nonpersistent_buffers(self, monkeypatch, tmp_path):
+        """CPU regression for the accelerator-only FSDP2 offload subprocess tests.
+
+        The persistent registration is checkpoint-authoritative while the
+        non-persistent-only alias keeps its config-derived value.
+        """
+
         class ModelWithDerivedBuffer(nn.Module):
             _no_split_modules = []
 
             def __init__(self):
                 super().__init__()
                 self.weight = nn.Parameter(torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
-                self.register_buffer("scale", torch.tensor([0.25, 2.0]), persistent=False)
+                shared_scale = torch.tensor([0.25, 2.0])
+                self.register_buffer("scale", shared_scale, persistent=False)
+                self.register_buffer("checkpointed_scale", shared_scale)
+                offset = torch.tensor([1.5, -0.5])
+                self.register_buffer("offset", offset)
+                self.register_buffer("offset_alias", offset)
+                runtime_cache = torch.tensor([7.0])
+                self.register_buffer("runtime_cache", runtime_cache, persistent=False)
+                self.register_buffer("runtime_cache_alias", runtime_cache, persistent=False)
 
             def forward(self, x):
-                return (x @ self.weight) * self.scale
+                return (x @ self.weight) * self.scale + self.offset
 
             def init_weights(self):
                 raise AssertionError("DCP resume must not initialize model parameters")
 
         original = ModelWithDerivedBuffer()
         assert "scale" not in original.state_dict()
+        assert "checkpointed_scale" in original.state_dict()
+        assert "offset" in original.state_dict()
+        assert "runtime_cache" not in original.state_dict()
+        assert original.scale is original.checkpointed_scale
+        assert original.offset is original.offset_alias
+        assert original.runtime_cache is original.runtime_cache_alias
+        original.scale.copy_(torch.tensor([0.5, 3.0]))
         inputs = torch.tensor([[2.0, -1.0]])
         expected_output = original(inputs)
         checkpoint_dir = tmp_path / "dcp"
@@ -531,6 +552,9 @@ def init_weights(self):
             resumed = ModelWithDerivedBuffer()
         assert resumed.weight.is_meta
         assert not resumed.scale.is_meta
+        assert resumed.scale is resumed.checkpointed_scale
+        assert resumed.offset is resumed.offset_alias
+        assert resumed.runtime_cache is resumed.runtime_cache_alias
         parallel_state = SimpleNamespace(any_extra_parallel_enabled=False, extra_parallel_names=[], fsdp_mesh=None)
         monkeypatch.setattr(torch_parallelize, "get_parallel_state", lambda: parallel_state)
         monkeypatch.setattr(torch_parallelize, "fully_shard", lambda *args, **kwargs: None)
@@ -541,11 +565,19 @@ def init_weights(self):
             weights_path="unused-hf-path",
             mixed_precision=SimpleNamespace(enable=False),
             should_skip_hf_weight_load=True,
+            enable_fsdp_offload=True,
+            fsdp_offload_pin_memory=False,
             init_device="meta",
         )
         dcp.load({"model": resumed}, checkpoint_id=checkpoint_dir)
 
+        assert resumed.scale is resumed.checkpointed_scale
+        assert resumed.offset is resumed.offset_alias
+        assert resumed.runtime_cache is resumed.runtime_cache_alias
         torch.testing.assert_close(resumed.scale, original.scale, rtol=0, atol=0)
+        torch.testing.assert_close(resumed.checkpointed_scale, original.checkpointed_scale, rtol=0, atol=0)
+        torch.testing.assert_close(resumed.offset, original.offset, rtol=0, atol=0)
+        torch.testing.assert_close(resumed.runtime_cache, original.runtime_cache, rtol=0, atol=0)
         torch.testing.assert_close(resumed(inputs), expected_output, rtol=0, atol=0)