From ff45cc768d54e8650db4c35bda3e7c2dd23bf3cd Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Thu, 23 Jul 2026 21:19:27 +0800
Subject: [PATCH 01/12] [dist] fix: keep FSDP2 buffers on compute device with
 CPU offload

---
 tests/checkpoints/test_dcp_checkpointer.py    |  32 +-
 .../test_rank0_load_and_broadcast_weights.py  | 301 +++++++++++++++++-
 veomni/distributed/torch_parallelize.py       |  90 +++++-
 3 files changed, 397 insertions(+), 26 deletions(-)

diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py
index 2f4af7d785..6a51aac6be 100644
--- a/tests/checkpoints/test_dcp_checkpointer.py
+++ b/tests/checkpoints/test_dcp_checkpointer.py
@@ -174,23 +174,38 @@ 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):
         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"
@@ -200,6 +215,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)
@@ -210,11 +228,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)
 
     @pytest.mark.parametrize("parallelize", [build_parallelize_model, parallelize_model_fsdp2])
diff --git a/tests/utils/test_rank0_load_and_broadcast_weights.py b/tests/utils/test_rank0_load_and_broadcast_weights.py
index 05b70984ea..e980b60a08 100644
--- a/tests/utils/test_rank0_load_and_broadcast_weights.py
+++ b/tests/utils/test_rank0_load_and_broadcast_weights.py
@@ -6,15 +6,17 @@
 
 import pytest
 import torch
+import torch.distributed.checkpoint as dcp
 from safetensors.torch import save_file
 from torch import distributed as dist
 from torch import nn
+from torch.distributed.device_mesh import DeviceMesh
 
 from veomni.arguments import DataArguments, ModelArguments, TrainingArguments, parse_args
 from veomni.distributed.parallel_plan import ParallelPlan
 from veomni.distributed.parallel_state import init_parallel_state
-from veomni.distributed.torch_parallelize import build_parallelize_model
-from veomni.models.module_utils import load_model_weights, rank0_load_and_broadcast_weights
+from veomni.distributed.torch_parallelize import _move_model_buffers_to_device, build_parallelize_model
+from veomni.models.module_utils import init_empty_weights, load_model_weights, rank0_load_and_broadcast_weights
 from veomni.utils import helper
 from veomni.utils.device import get_device_type, get_dist_comm_backend, get_torch_device
 
@@ -34,8 +36,8 @@
 class BroadcastTestArguments:
     weights_path: str = ""
     device_type: str = get_device_type()
-    backend: str = get_dist_comm_backend()
-    mode: str = "broadcast"  # "broadcast" | "load_weights"
+    backend: str = ""
+    mode: str = "broadcast"  # "broadcast" | "load_weights" | "offload_training"
 
 
 @dataclass
@@ -172,6 +174,24 @@ def init_weights(self):
         self.linear3.bias.data.fill_(2.0)
 
 
+class ToyBufferModel(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(4, 2)
+        self.register_buffer("offset", torch.tensor([0.25, -0.5]))
+        self.register_buffer("scale", torch.tensor(0.75), persistent=False)
+        self.config = SimpleNamespace(tie_word_embeddings=False)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        return self.linear(inputs) * self.scale + self.offset
+
+    def init_weights(self) -> None:
+        nn.init.xavier_uniform_(self.linear.weight)
+        self.linear.bias.data.fill_(0.125)
+        self.offset.copy_(torch.tensor([0.25, -0.5], device=self.offset.device))
+        self.scale.fill_(0.75)
+
+
 def _write_checkpoint(checkpoint_dir: Path, is_parallel: bool) -> str:
     torch.manual_seed(0)
     if is_parallel:
@@ -188,6 +208,15 @@ def _write_checkpoint(checkpoint_dir: Path, is_parallel: bool) -> str:
     return str(checkpoint_dir)
 
 
+def _write_buffer_checkpoint(checkpoint_dir: Path) -> str:
+    torch.manual_seed(2026)
+    model = ToyBufferModel()
+    model.init_weights()
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    save_file(model.state_dict(), str(checkpoint_dir / "model.safetensors"))
+    return str(checkpoint_dir)
+
+
 def _clone_to_cpu(tensor: torch.Tensor) -> torch.Tensor:
     if DTensor is not None and isinstance(tensor, DTensor):
         assert Replicate is not None
@@ -210,13 +239,226 @@ def _assert_model_parameters_on_device(model: nn.Module, expected_device_type: s
         )
 
 
+def _assert_model_buffers_on_device(model: nn.Module, expected_device_type: str) -> None:
+    for name, buffer in model.named_buffers():
+        local_buffer = buffer.to_local() if DTensor is not None and isinstance(buffer, DTensor) else buffer
+        assert local_buffer.device.type == expected_device_type, (
+            f"Expected buffer {name} to be on {expected_device_type}, got {local_buffer.device.type}."
+        )
+
+
+def _assert_model_gradients_on_device(model: nn.Module, expected_device_type: str) -> None:
+    gradient_count = 0
+    for name, parameter in model.named_parameters():
+        if parameter.grad is None:
+            continue
+        gradient_count += 1
+        gradient = parameter.grad
+        local_gradient = gradient.to_local() if DTensor is not None and isinstance(gradient, DTensor) else gradient
+        assert local_gradient.device.type == expected_device_type, (
+            f"Expected gradient {name} to be on {expected_device_type}, got {local_gradient.device.type}."
+        )
+    assert gradient_count > 0, "Expected the training step to produce gradients."
+
+
+def _assert_optimizer_moments_on_device(optimizer: torch.optim.Optimizer, expected_device_type: str) -> None:
+    assert optimizer.state, "Expected the optimizer step to initialize optimizer state."
+    moment_count = 0
+    for state in optimizer.state.values():
+        for name in ("exp_avg", "exp_avg_sq"):
+            assert name in state, f"Expected optimizer state to contain {name}."
+            moment_count += 1
+            moment = state[name]
+            local_moment = moment.to_local() if DTensor is not None and isinstance(moment, DTensor) else moment
+            assert local_moment.device.type == expected_device_type, (
+                f"Expected optimizer state {name} to be on {expected_device_type}, got {local_moment.device.type}."
+            )
+    assert moment_count > 0, "Expected at least one optimizer moment tensor."
+
+
+def test_move_model_buffers_to_device_preserves_parameters_and_persistence() -> None:
+    model = nn.Sequential(nn.Linear(2, 2))
+    model.register_buffer("persistent_buffer", torch.ones(2))
+    runtime_buffer = torch.ones(3)
+    model[0].register_buffer("runtime_buffer", runtime_buffer, persistent=False)
+    model.register_buffer("runtime_buffer_alias", runtime_buffer, persistent=False)
+    model.register_buffer("optional_buffer", None)
+
+    parameter_devices = {name: parameter.device for name, parameter in model.named_parameters()}
+    moved_count, moved_bytes = _move_model_buffers_to_device(model, torch.device("meta"))
+
+    assert moved_count == 2
+    assert moved_bytes == 5 * torch.tensor([], dtype=torch.float32).element_size()
+    assert {name: parameter.device for name, parameter in model.named_parameters()} == parameter_devices
+    assert model.persistent_buffer.device.type == "meta"
+    assert model[0].runtime_buffer.device.type == "meta"
+    assert model.runtime_buffer_alias is model[0].runtime_buffer
+    assert model.optional_buffer is None
+    assert "persistent_buffer" in model.state_dict()
+    assert "0.runtime_buffer" not in model.state_dict()
+    assert "runtime_buffer_alias" not in model.state_dict()
+
+
+@pytest.mark.skipif(DTensor is None, reason="DTensor required")
+def test_move_model_buffers_to_device_preserves_dtensor_aliases(tmp_path: Path) -> None:
+    if dist.is_initialized():
+        pytest.skip("Test requires ownership of the default process group.")
+
+    dist.init_process_group(
+        backend="gloo",
+        init_method=f"file://{tmp_path / 'process_group'}",
+        rank=0,
+        world_size=1,
+    )
+    try:
+        mesh = DeviceMesh("cpu", [0])
+        buffer = DTensor.from_local(torch.ones(3), mesh, [Replicate()])
+        model = nn.Module()
+        model.register_buffer("distributed_buffer", buffer)
+        model.register_buffer("distributed_buffer_alias", buffer, persistent=False)
+
+        moved_count, moved_bytes = _move_model_buffers_to_device(model, torch.device("meta"))
+
+        assert moved_count == 1
+        assert moved_bytes == 3 * torch.tensor([], dtype=torch.float32).element_size()
+        assert isinstance(model.distributed_buffer, DTensor)
+        assert model.distributed_buffer.to_local().device.type == "meta"
+        assert model.distributed_buffer_alias is model.distributed_buffer
+    finally:
+        dist.destroy_process_group()
+
+
+def run_fsdp_cpu_offload_training_test(args: Arguments) -> None:
+    weights_path = Path(args.test.weights_path)
+    if not weights_path.exists():
+        raise ValueError("`--test.weights_path` must point to an existing directory.")
+
+    get_torch_device().set_device(args.train.local_rank)
+    dist.init_process_group(backend=args.test.backend or get_dist_comm_backend())
+
+    parallel_state = init_parallel_state(
+        dp_size=args.train.accelerator.dp_size,
+        dp_replicate_size=args.train.accelerator.dp_replicate_size,
+        dp_shard_size=args.train.accelerator.dp_shard_size,
+        tp_size=args.train.accelerator.tp_size,
+        pp_size=args.train.accelerator.pp_size,
+        cp_size=args.train.accelerator.cp_size,
+        ulysses_size=args.train.accelerator.ulysses_size,
+        dp_mode=args.train.accelerator.fsdp_config.fsdp_mode,
+    )
+
+    try:
+        reference_model = ToyBufferModel().to(get_device_type())
+        load_model_weights(reference_model, str(weights_path), init_device=get_device_type())
+
+        with init_empty_weights():
+            fsdp_model = ToyBufferModel()
+        assert all(parameter.is_meta for parameter in fsdp_model.parameters())
+        fsdp_model = build_parallelize_model(
+            fsdp_model,
+            weights_path=str(weights_path),
+            init_device=args.train.init_device,
+            mixed_precision=args.train.accelerator.fsdp_config.mixed_precision,
+            enable_gradient_checkpointing=False,
+            enable_fsdp_offload=True,
+            fsdp_offload_pin_memory=False,
+            basic_modules=[],
+        )
+
+        _assert_model_parameters_on_device(fsdp_model, "cpu")
+        _assert_model_buffers_on_device(fsdp_model, get_device_type())
+
+        fsdp_optimizer = torch.optim.AdamW(fsdp_model.parameters(), lr=1e-2, weight_decay=0.0)
+        reference_optimizer = torch.optim.AdamW(reference_model.parameters(), lr=1e-2, weight_decay=0.0)
+        inputs = torch.tensor(
+            [[0.5, -1.0, 1.5, 2.0], [-0.25, 0.75, 1.25, -1.5]],
+            device=get_device_type(),
+        )
+
+        fsdp_output = fsdp_model(inputs)
+        reference_output = reference_model(inputs)
+        torch.testing.assert_close(fsdp_output, reference_output, atol=1e-6, rtol=1e-6)
+
+        fsdp_output.square().mean().backward()
+        reference_output.square().mean().backward()
+        _assert_model_gradients_on_device(fsdp_model, "cpu")
+        fsdp_optimizer.step()
+        reference_optimizer.step()
+
+        _assert_model_parameters_on_device(fsdp_model, "cpu")
+        _assert_model_buffers_on_device(fsdp_model, get_device_type())
+        _assert_optimizer_moments_on_device(fsdp_optimizer, "cpu")
+
+        with torch.no_grad():
+            resumed_offset = torch.tensor([1.5, -2.0], device=get_device_type())
+            fsdp_model.offset.copy_(resumed_offset)
+            reference_model.offset.copy_(resumed_offset)
+
+        fsdp_updated_output = fsdp_model(inputs)
+        reference_updated_output = reference_model(inputs)
+        torch.testing.assert_close(fsdp_updated_output, reference_updated_output, atol=1e-5, rtol=1e-5)
+
+        checkpoint_dir = weights_path.parent / "dcp_cold_resume"
+        from veomni.checkpoint.dcp_checkpointer import ModelState
+
+        dcp.save(
+            {"model": ModelState(fsdp_model, parallel_state=parallel_state)},
+            checkpoint_id=checkpoint_dir,
+        )
+
+        with init_empty_weights():
+            resumed_model = ToyBufferModel()
+        assert all(parameter.is_meta for parameter in resumed_model.parameters())
+        resumed_model = build_parallelize_model(
+            resumed_model,
+            weights_path="unused-hf-path",
+            init_device=args.train.init_device,
+            mixed_precision=args.train.accelerator.fsdp_config.mixed_precision,
+            enable_gradient_checkpointing=False,
+            enable_fsdp_offload=True,
+            fsdp_offload_pin_memory=False,
+            should_skip_hf_weight_load=True,
+            basic_modules=[],
+        )
+        _assert_model_parameters_on_device(resumed_model, "cpu")
+        _assert_model_buffers_on_device(resumed_model, get_device_type())
+        torch.testing.assert_close(
+            resumed_model.scale,
+            torch.tensor(0.75, device=get_device_type()),
+            atol=0.0,
+            rtol=0.0,
+        )
+
+        dcp.load(
+            {"model": ModelState(resumed_model, parallel_state=parallel_state)},
+            checkpoint_id=checkpoint_dir,
+        )
+        _assert_model_parameters_on_device(resumed_model, "cpu")
+        _assert_model_buffers_on_device(resumed_model, get_device_type())
+        torch.testing.assert_close(resumed_model.offset, resumed_offset, atol=0.0, rtol=0.0)
+        torch.testing.assert_close(
+            resumed_model.scale,
+            torch.tensor(0.75, device=get_device_type()),
+            atol=0.0,
+            rtol=0.0,
+        )
+        torch.testing.assert_close(resumed_model(inputs), reference_updated_output, atol=1e-5, rtol=1e-5)
+
+        dist.barrier()
+    finally:
+        dist.destroy_process_group()
+        from veomni.distributed.parallel_state import clear_parallel_state
+
+        clear_parallel_state()
+
+
 def run_rank0_broadcast_test(args: Arguments) -> None:
     weights_path = Path(args.test.weights_path)
     if not weights_path.exists():
         raise ValueError("`--test.weights_path` must point to an existing directory.")
 
     get_torch_device().set_device(args.train.local_rank)
-    dist.init_process_group(backend=args.test.backend)
+    dist.init_process_group(backend=args.test.backend or get_dist_comm_backend())
 
     init_parallel_state(
         dp_size=args.train.accelerator.dp_size,
@@ -399,7 +641,10 @@ def run_rank0_broadcast_test(args: Arguments) -> None:
         clear_parallel_state()
 
 
-@pytest.mark.skipif(not dist.is_available(), reason="torch.distributed required")
+@pytest.mark.skipif(
+    not dist.is_available() or get_device_type() == "cpu",
+    reason="Distributed FSDP2 weight loading test requires an accelerator",
+)
 @pytest.mark.parametrize("cpu_offload", [False, True], ids=["no_offload", "cpu_offload"])
 def test_load_dist_model_weights_matches_standard(tmp_path: Path, cpu_offload: bool) -> None:
     checkpoint_dir = tmp_path / "ckpt"
@@ -451,7 +696,7 @@ def run_load_weights_test(args: Arguments) -> None:
         raise ValueError("`--test.weights_path` must point to an existing directory.")
 
     get_torch_device().set_device(args.train.local_rank)
-    dist.init_process_group(backend=args.test.backend)
+    dist.init_process_group(backend=args.test.backend or get_dist_comm_backend())
 
     init_parallel_state(
         dp_size=args.train.accelerator.dp_size,
@@ -566,7 +811,10 @@ def run_load_weights_test(args: Arguments) -> None:
         clear_parallel_state()
 
 
-@pytest.mark.skipif(not dist.is_available(), reason="torch.distributed required")
+@pytest.mark.skipif(
+    not dist.is_available() or get_device_type() == "cpu",
+    reason="Distributed FSDP2 weight loading test requires an accelerator",
+)
 @pytest.mark.parametrize("cpu_offload", [False, True], ids=["no_offload", "cpu_offload"])
 def test_load_weights_no_scatter(tmp_path: Path, cpu_offload: bool) -> None:
     """
@@ -607,9 +855,44 @@ def test_load_weights_no_scatter(tmp_path: Path, cpu_offload: bool) -> None:
     assert result.returncode == 0
 
 
+@pytest.mark.skipif(
+    not dist.is_available() or get_device_type() == "cpu",
+    reason="FSDP2 CPU-offload device residency requires an accelerator",
+)
+def test_fsdp_cpu_offload_training_keeps_buffers_on_compute_device(tmp_path: Path) -> None:
+    if get_torch_device().device_count() < 2:
+        pytest.skip("FSDP2 CPU-offload residency gate requires at least two accelerator devices.")
+
+    weights_path = _write_buffer_checkpoint(tmp_path / "ckpt")
+    world_size = 2
+    port = 12345 + random.randint(0, 100)
+    command = [
+        "torchrun",
+        f"--nproc_per_node={world_size}",
+        f"--master_port={port}",
+        "tests/utils/test_rank0_load_and_broadcast_weights.py",
+        "--model.config_path=test",
+        "--data.train_path=tests",
+        "--train.checkpoint.output_dir=.tests/cache",
+        "--train.accelerator.fsdp_config.fsdp_mode=fsdp2",
+        "--train.init_device=meta",
+        "--train.accelerator.fsdp_config.mixed_precision.enable=False",
+        "--train.gradient_checkpointing.enable=False",
+        f"--test.weights_path={weights_path}",
+        f"--test.device_type={get_device_type()}",
+        f"--test.backend={get_dist_comm_backend()}",
+        "--test.mode=offload_training",
+    ]
+
+    result = subprocess.run(command, check=True)
+    assert result.returncode == 0
+
+
 if __name__ == "__main__":
     args = parse_args(Arguments)
-    if getattr(args.test, "mode", "broadcast") == "load_weights":
+    if getattr(args.test, "mode", "broadcast") == "offload_training":
+        run_fsdp_cpu_offload_training_test(args)
+    elif getattr(args.test, "mode", "broadcast") == "load_weights":
         run_load_weights_test(args)
     else:
         run_rank0_broadcast_test(args)
diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 63dcebdab6..652a658d63 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -29,7 +29,7 @@
 from ..arguments import MixedPrecisionConfig
 from ..models import load_model_weights, load_model_weights_ep_sharded, rank0_load_and_broadcast_weights
 from ..utils import logging
-from ..utils.device import IS_NPU_AVAILABLE, get_device_type
+from ..utils.device import IS_NPU_AVAILABLE, get_device_id, get_device_type
 from .checkpoint import CheckpointFunction
 from .chunk_mbs import apply_chunk_mbs
 from .parallel_plan import get_runtime_parallel_plan
@@ -49,25 +49,76 @@ def _reset_hf_initialized_flag(module: nn.Module) -> None:
 
 
 def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) -> None:
-    """Materialize parameters without discarding config-derived buffers.
+    """Materialize parameters while preserving buffer alias groups.
 
-    Distributed checkpoints restore ``state_dict()``, which excludes buffers
-    registered with ``persistent=False``. Snapshot those buffers before
-    ``to_empty()`` so values such as rotary ``inv_freq`` survive the DCP resume
-    materialization path without duplicating persistent buffers that DCP loads.
+    ``Module.to_empty()`` replaces each registered buffer independently, which
+    breaks aliases. Record every registration before materialization and bind
+    each original alias group to one materialized tensor afterwards.
+
+    Distributed checkpoints do not restore groups containing only
+    nonpersistent buffers, so preserve their config-derived values. If a group
+    has any persistent registration, the checkpoint is authoritative and will
+    populate the shared materialized tensor during the subsequent DCP load.
     """
-    buffers = []
+
+    buffer_groups: dict[int, list[tuple[nn.Module, str]]] = {}
+    preserved_nonpersistent_buffers: dict[int, torch.Tensor] = {}
     for module in model.modules():
-        for name in module._non_persistent_buffers_set:
-            buffer = module._buffers.get(name)
-            if buffer is not None:
-                buffers.append((module, name, buffer.detach().clone()))
+        for name, buffer in module._buffers.items():
+            if buffer is None:
+                continue
+            buffer_id = id(buffer)
+            buffer_groups.setdefault(buffer_id, []).append((module, name))
+
+    for buffer_id, registrations in buffer_groups.items():
+        if all(name in module._non_persistent_buffers_set for module, name in registrations):
+            module, name = registrations[0]
+            preserved_nonpersistent_buffers[buffer_id] = module._buffers[name].detach().clone()
 
     model.to_empty(device=device)
 
-    for module, name, buffer in buffers:
-        materialized_buffer = module._buffers[name]
-        materialized_buffer.copy_(buffer.to(device=materialized_buffer.device, dtype=materialized_buffer.dtype))
+    for buffer_id, registrations in buffer_groups.items():
+        if buffer_id in preserved_nonpersistent_buffers:
+            canonical_buffer = preserved_nonpersistent_buffers[buffer_id].to(device=device)
+        else:
+            module, name = registrations[0]
+            canonical_buffer = module._buffers[name]
+        for module, name in registrations:
+            module._buffers[name] = canonical_buffer
+
+
+@torch.no_grad()
+def _move_model_buffers_to_device(model: nn.Module, device: torch.device) -> tuple[int, int]:
+    """Move registered buffers without moving CPU-offloaded parameters.
+
+    PyTorch's ``Module._apply`` replaces buffer tensor objects during device
+    conversion, so external references are not guaranteed to retain object
+    identity. This helper follows that contract while memoizing by object
+    identity so multiple registrations of the same buffer remain aliased.
+    Tensor subclasses such as DTensor keep their type through ``Tensor.to``.
+    """
+
+    moved_count = 0
+    moved_bytes = 0
+    moved_by_id: dict[int, torch.Tensor] = {}
+    for module in model.modules():
+        for name, buffer in module._buffers.items():
+            if buffer is None:
+                continue
+            buffer_id = id(buffer)
+            if buffer_id in moved_by_id:
+                module._buffers[name] = moved_by_id[buffer_id]
+                continue
+            if buffer.device == device:
+                moved_by_id[buffer_id] = buffer
+                continue
+
+            moved_count += 1
+            moved_bytes += buffer.numel() * buffer.element_size()
+            moved_buffer = buffer.to(device=device)
+            moved_by_id[buffer_id] = moved_buffer
+            module._buffers[name] = moved_buffer
+    return moved_count, moved_bytes
 
 
 def _check_extra_parallel_dim0_divisibility(model: "nn.Module", para_name: str, ep_fsdp_size: int) -> bool:
@@ -535,6 +586,17 @@ def parallelize_model_fsdp2(
                     fqn_to_index_mapping=fqn_to_index_mapping,
                 )
 
+    if enable_fsdp_cpu_offload:
+        device_type = get_device_type()
+        compute_device = (
+            torch.device(device_type, get_device_id()) if device_type != "cpu" else torch.device(device_type)
+        )
+        moved_buffers, moved_buffer_bytes = _move_model_buffers_to_device(model, compute_device)
+        logger.info_rank0(
+            "FSDP2 CPU offload kept registered buffers on the compute device: "
+            f"moved_buffers={moved_buffers}, moved_bytes={moved_buffer_bytes}."
+        )
+
     # Register grad norm clipping method for FSDP2
     from .fsdp2 import clip_grad_norm as clip_grad_norm_fn
 

From a1348602ac75c3c97756e83c106411b18b86b304 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 10 Aug 2026 18:31:13 +0800
Subject: [PATCH 02/12] fix: keep FSDP2 offload buffers on compute device

---
 .github/workflows/gpu_unit_tests.yml            |  3 +++
 .github/workflows/npu_unit_tests.yml            |  3 +++
 tests/checkpoints/test_dcp_checkpointer.py      |  5 +++++
 .../test_rank0_load_and_broadcast_weights.py    | 17 ++++++++++++++++-
 veomni/distributed/torch_parallelize.py         | 10 ++++++++--
 5 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml
index ddd7f99c96..a66f76154b 100644
--- a/.github/workflows/gpu_unit_tests.yml
+++ b/.github/workflows/gpu_unit_tests.yml
@@ -160,6 +160,9 @@ jobs:
       - name: Run e2e dcp save and load test
         run: |
           uv run --frozen pytest -s -x tests/checkpoints/test_trainer_saveload.py
+      - name: Run DCP checkpointer CPU-offload regression tests
+        run: |
+          uv run --frozen pytest -s -x tests/checkpoints/test_dcp_checkpointer.py
       - name: Run distributed tests (dummy_forward)
         run: |
           uv run --frozen pytest -s -x tests/distributed/test_dummy_forward.py
diff --git a/.github/workflows/npu_unit_tests.yml b/.github/workflows/npu_unit_tests.yml
index 60091a78a7..5f0bc41eaa 100644
--- a/.github/workflows/npu_unit_tests.yml
+++ b/.github/workflows/npu_unit_tests.yml
@@ -176,6 +176,9 @@ jobs:
       - name: Run e2e dcp save and load test
         run: |
           uv run --frozen pytest -v -s -x tests/checkpoints/test_trainer_saveload.py
+      - name: Run DCP checkpointer CPU-offload regression tests
+        run: |
+          uv run --frozen pytest -v -s -x tests/checkpoints/test_dcp_checkpointer.py
       - name: Run distributed tests (dummy_forward)
         run: |
           uv run --frozen pytest -v -s -x tests/distributed/test_dummy_forward.py
diff --git a/tests/checkpoints/test_dcp_checkpointer.py b/tests/checkpoints/test_dcp_checkpointer.py
index 6a51aac6be..b03fe6b018 100644
--- a/tests/checkpoints/test_dcp_checkpointer.py
+++ b/tests/checkpoints/test_dcp_checkpointer.py
@@ -175,6 +175,11 @@ def test_build_parallelize_model_forwards_should_skip_hf_weight_load(self, monke
         assert parallelize_fsdp2.call_args.kwargs["should_skip_hf_weight_load"] is True
 
     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 = []
 
diff --git a/tests/utils/test_rank0_load_and_broadcast_weights.py b/tests/utils/test_rank0_load_and_broadcast_weights.py
index e980b60a08..ec567309b9 100644
--- a/tests/utils/test_rank0_load_and_broadcast_weights.py
+++ b/tests/utils/test_rank0_load_and_broadcast_weights.py
@@ -312,7 +312,13 @@ def test_move_model_buffers_to_device_preserves_dtensor_aliases(tmp_path: Path)
     )
     try:
         mesh = DeviceMesh("cpu", [0])
-        buffer = DTensor.from_local(torch.ones(3), mesh, [Replicate()])
+        buffer = DTensor.from_local(
+            torch.ones(3),
+            mesh,
+            [Shard(0)],
+            shape=torch.Size([6]),
+            stride=(1,),
+        )
         model = nn.Module()
         model.register_buffer("distributed_buffer", buffer)
         model.register_buffer("distributed_buffer_alias", buffer, persistent=False)
@@ -320,6 +326,8 @@ def test_move_model_buffers_to_device_preserves_dtensor_aliases(tmp_path: Path)
         moved_count, moved_bytes = _move_model_buffers_to_device(model, torch.device("meta"))
 
         assert moved_count == 1
+        # ``buffer.numel()`` is the global logical size (6), while only the
+        # local shard (3 elements) is transferred by ``Tensor.to``.
         assert moved_bytes == 3 * torch.tensor([], dtype=torch.float32).element_size()
         assert isinstance(model.distributed_buffer, DTensor)
         assert model.distributed_buffer.to_local().device.type == "meta"
@@ -647,6 +655,9 @@ def run_rank0_broadcast_test(args: Arguments) -> None:
 )
 @pytest.mark.parametrize("cpu_offload", [False, True], ids=["no_offload", "cpu_offload"])
 def test_load_dist_model_weights_matches_standard(tmp_path: Path, cpu_offload: bool) -> None:
+    """Run the accelerator-only subprocess path; CPU buffer semantics are covered
+    by ``test_cpu_offload_dcp_resume_restores_persistent_and_nonpersistent_buffers``.
+    """
     checkpoint_dir = tmp_path / "ckpt"
     weights_path = _write_checkpoint(checkpoint_dir, is_parallel=True)
 
@@ -820,6 +831,10 @@ def test_load_weights_no_scatter(tmp_path: Path, cpu_offload: bool) -> None:
     """
     Regression test for https://github.com/ByteDance-Seed/VeOmni/issues/637.
 
+    This launches FSDP2 with an accelerator because CPU-only FSDP2 does not
+    exercise the device-residency behavior. The CPU checkpoint path has a
+    separate regression test in ``tests/checkpoints/test_dcp_checkpointer.py``.
+
     Ensures load_model_weights with src_data_rank=None (all-ranks-read path)
     loads bit-for-bit correct parameters into an FSDP2 model.
     The rank0_load_and_broadcast_weights path is tested separately in
diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 652a658d63..6279d20a12 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -94,7 +94,9 @@ def _move_model_buffers_to_device(model: nn.Module, device: torch.device) -> tup
     PyTorch's ``Module._apply`` replaces buffer tensor objects during device
     conversion, so external references are not guaranteed to retain object
     identity. This helper follows that contract while memoizing by object
-    identity so multiple registrations of the same buffer remain aliased.
+    identity so multiple registrations of the same buffer remain aliased. The
+    byte count is based on each buffer's local storage, which is the actual
+    memory moved for a DTensor rather than its global logical shape.
     Tensor subclasses such as DTensor keep their type through ``Tensor.to``.
     """
 
@@ -114,7 +116,11 @@ def _move_model_buffers_to_device(model: nn.Module, device: torch.device) -> tup
                 continue
 
             moved_count += 1
-            moved_bytes += buffer.numel() * buffer.element_size()
+            local_buffer = buffer
+            to_local = getattr(buffer, "to_local", None)
+            if callable(to_local):
+                local_buffer = to_local()
+            moved_bytes += local_buffer.numel() * local_buffer.element_size()
             moved_buffer = buffer.to(device=device)
             moved_by_id[buffer_id] = moved_buffer
             module._buffers[name] = moved_buffer

From 8d1738cf132b51aef4093a8718b2f31878e66be5 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 10 Aug 2026 20:08:35 +0800
Subject: [PATCH 03/12] fix: resolve conflicts with upstream main

---
 .github/workflows/gpu_unit_tests.yml    | 10 ++++++----
 veomni/distributed/torch_parallelize.py |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml
index a66f76154b..a3db7800e7 100644
--- a/.github/workflows/gpu_unit_tests.yml
+++ b/.github/workflows/gpu_unit_tests.yml
@@ -134,6 +134,7 @@ jobs:
           uv run --frozen pytest -s -x tests/models/test_model_registry.py
           uv run --frozen pytest -s -x tests/models/test_models_patch.py
           uv run --frozen pytest -s -x tests/models/test_vlm_trainer.py
+          uv run --frozen pytest -s -x tests/lora/test_vlm_lora.py
           uv run --frozen pytest -s -x tests/models/test_padded_packed_loss.py
           uv run --frozen pytest -s -x tests/models/test_models_logits_equal_v5.py
           uv run --frozen pytest -s -x tests/models/test_model_forward_no_implicit_sync.py
@@ -160,12 +161,10 @@ jobs:
       - name: Run e2e dcp save and load test
         run: |
           uv run --frozen pytest -s -x tests/checkpoints/test_trainer_saveload.py
-      - name: Run DCP checkpointer CPU-offload regression tests
-        run: |
-          uv run --frozen pytest -s -x tests/checkpoints/test_dcp_checkpointer.py
-      - name: Run distributed tests (dummy_forward)
+      - name: Run distributed tests
         run: |
           uv run --frozen pytest -s -x tests/distributed/test_dummy_forward.py
+          uv run --frozen pytest -s -x tests/distributed/test_torch_compile.py
       - name: Run data tests
         run: |
           uv run --frozen pytest -s -x tests/data
@@ -207,3 +206,6 @@ jobs:
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_trainer.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep2.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep_sharded_stream_load.py
+      - name: Run DCP checkpointer CPU-offload regression tests
+        run: |
+          uv run --frozen pytest -s -x tests/checkpoints/test_dcp_checkpointer.py
diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 6279d20a12..7a94aca21f 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -592,7 +592,7 @@ def parallelize_model_fsdp2(
                     fqn_to_index_mapping=fqn_to_index_mapping,
                 )
 
-    if enable_fsdp_cpu_offload:
+    if materialize_device == "cpu":
         device_type = get_device_type()
         compute_device = (
             torch.device(device_type, get_device_id()) if device_type != "cpu" else torch.device(device_type)

From d12960ece856936c072cbf0919959052958c8c81 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 10 Aug 2026 20:13:14 +0800
Subject: [PATCH 04/12] fix: resolve FSDP2 merge conflict with upstream main

---
 veomni/distributed/torch_parallelize.py | 49 +++++++++++++++++++------
 1 file changed, 38 insertions(+), 11 deletions(-)

diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 7a94aca21f..f9a53d7cb1 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -34,7 +34,7 @@
 from .chunk_mbs import apply_chunk_mbs
 from .parallel_plan import get_runtime_parallel_plan
 from .parallel_state import get_parallel_state
-from .torch_compile import CompileConfig, compile_decoder_blocks, validate_compile_config_for_fsdp2
+from .torch_compile import CompileConfig, compile_decoder_blocks, validate_compile_runtime
 from .utils import sort_fqn_by_submodule_first
 
 
@@ -127,6 +127,24 @@ def _move_model_buffers_to_device(model: nn.Module, device: torch.device) -> tup
     return moved_count, moved_bytes
 
 
+def _veomni_shard_placement_fn(param: "nn.Parameter") -> Optional[Shard]:
+    """``fully_shard(..., shard_placement_fn=...)`` hook: shard on a model-chosen
+    dimension instead of the FSDP2 default ``Shard(0)``.
+
+    A model opts a specific parameter in by tagging it with a
+    ``_veomni_fsdp_shard_dim`` attribute in its own ``__init__`` (e.g. DeepSeek-V4's
+    compressor/indexer ``position_bias`` -- shape ``(compress_rate, head_dim*k)``
+    where ``compress_rate`` can be as small as 4, but the trailing dim is a large,
+    reliably-divisible power of 2). Sharding such a param on dim-0 across a large
+    FSDP world leaves most ranks with a genuinely empty local shard; redirecting it
+    to a dimension that divides evenly avoids that at no memory cost. Returning
+    ``None`` for untagged params keeps the FSDP2 default (``Shard(0)``).
+    """
+    dim = getattr(param, "_veomni_fsdp_shard_dim", None)
+    return Shard(dim) if dim is not None else None
+
+
+
 def _check_extra_parallel_dim0_divisibility(model: "nn.Module", para_name: str, ep_fsdp_size: int) -> bool:
     """Return whether EP-local dim-0 can be evenly sharded by ``ep_fsdp_size``."""
     parallel_plan = getattr(model, "get_parallel_plan", None)
@@ -297,23 +315,32 @@ def parallelize_model_fsdp2(
     logger.info_rank0(f"extra_parallel layer pairs: {layer_pairs}")
 
     if compile_config.enable:
-        if get_device_type() != "cuda":
-            raise RuntimeError("train.torch_compile.enable is CUDA-only for now.")
-        if parallel_state.any_extra_parallel_enabled:
-            raise RuntimeError(
-                "train.torch_compile.enable currently does not support ExtraParallel models because EP all-to-all "
-                "communication may be captured inside compiled blocks."
-            )
-        validate_compile_config_for_fsdp2(compile_config, enable_reshard_after_forward)
+        validate_compile_runtime(
+            compile_config,
+            device_type=get_device_type(),
+            fsdp_enabled=parallel_state.fsdp_enabled,
+            fsdp_mode=parallel_state.dp_mode,
+            any_extra_parallel_enabled=parallel_state.any_extra_parallel_enabled,
+            enable_reshard_after_forward=enable_reshard_after_forward,
+        )
 
-        compiled_count = compile_decoder_blocks(model, compile_config)
+        compiled_count = compile_decoder_blocks(
+            model,
+            compile_config,
+            sequence_parallel_enabled=parallel_state.sp_enabled,
+            async_enabled=parallel_state.async_enabled,
+        )
         if compiled_count == 0:
             raise RuntimeError("train.torch_compile.enable found no decoder blocks to compile.")
         model._veomni_compile_enabled = True
         model._veomni_compile_uses_cuda_graphs = compile_config.uses_cuda_graphs()
 
     # Step 2: Update fsdp2 kwargs
-    fsdp_kwargs = {"mesh": parallel_state.fsdp_mesh, "reshard_after_forward": enable_reshard_after_forward}
+    fsdp_kwargs = {
+        "mesh": parallel_state.fsdp_mesh,
+        "reshard_after_forward": enable_reshard_after_forward,
+        "shard_placement_fn": _veomni_shard_placement_fn,
+    }
     # prepare mp_policy kwargs
     if mixed_precision.enable:
         mp_policy = MixedPrecisionPolicy(

From 1059c6f8f2f2337ef0f44873183cac9075746350 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 10 Aug 2026 20:14:43 +0800
Subject: [PATCH 05/12] chore: align FSDP2 file with upstream main before
 branch update

---
 veomni/distributed/torch_parallelize.py | 109 +++++++-----------------
 1 file changed, 29 insertions(+), 80 deletions(-)

diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index f9a53d7cb1..4897d041aa 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -29,7 +29,7 @@
 from ..arguments import MixedPrecisionConfig
 from ..models import load_model_weights, load_model_weights_ep_sharded, rank0_load_and_broadcast_weights
 from ..utils import logging
-from ..utils.device import IS_NPU_AVAILABLE, get_device_id, get_device_type
+from ..utils.device import IS_NPU_AVAILABLE, get_device_type
 from .checkpoint import CheckpointFunction
 from .chunk_mbs import apply_chunk_mbs
 from .parallel_plan import get_runtime_parallel_plan
@@ -49,82 +49,25 @@ def _reset_hf_initialized_flag(module: nn.Module) -> None:
 
 
 def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) -> None:
-    """Materialize parameters while preserving buffer alias groups.
+    """Materialize parameters without discarding config-derived buffers.
 
-    ``Module.to_empty()`` replaces each registered buffer independently, which
-    breaks aliases. Record every registration before materialization and bind
-    each original alias group to one materialized tensor afterwards.
-
-    Distributed checkpoints do not restore groups containing only
-    nonpersistent buffers, so preserve their config-derived values. If a group
-    has any persistent registration, the checkpoint is authoritative and will
-    populate the shared materialized tensor during the subsequent DCP load.
+    Distributed checkpoints restore ``state_dict()``, which excludes buffers
+    registered with ``persistent=False``. Snapshot those buffers before
+    ``to_empty()`` so values such as rotary ``inv_freq`` survive the DCP resume
+    materialization path without duplicating persistent buffers that DCP loads.
     """
-
-    buffer_groups: dict[int, list[tuple[nn.Module, str]]] = {}
-    preserved_nonpersistent_buffers: dict[int, torch.Tensor] = {}
+    buffers = []
     for module in model.modules():
-        for name, buffer in module._buffers.items():
-            if buffer is None:
-                continue
-            buffer_id = id(buffer)
-            buffer_groups.setdefault(buffer_id, []).append((module, name))
-
-    for buffer_id, registrations in buffer_groups.items():
-        if all(name in module._non_persistent_buffers_set for module, name in registrations):
-            module, name = registrations[0]
-            preserved_nonpersistent_buffers[buffer_id] = module._buffers[name].detach().clone()
+        for name in module._non_persistent_buffers_set:
+            buffer = module._buffers.get(name)
+            if buffer is not None:
+                buffers.append((module, name, buffer.detach().clone()))
 
     model.to_empty(device=device)
 
-    for buffer_id, registrations in buffer_groups.items():
-        if buffer_id in preserved_nonpersistent_buffers:
-            canonical_buffer = preserved_nonpersistent_buffers[buffer_id].to(device=device)
-        else:
-            module, name = registrations[0]
-            canonical_buffer = module._buffers[name]
-        for module, name in registrations:
-            module._buffers[name] = canonical_buffer
-
-
-@torch.no_grad()
-def _move_model_buffers_to_device(model: nn.Module, device: torch.device) -> tuple[int, int]:
-    """Move registered buffers without moving CPU-offloaded parameters.
-
-    PyTorch's ``Module._apply`` replaces buffer tensor objects during device
-    conversion, so external references are not guaranteed to retain object
-    identity. This helper follows that contract while memoizing by object
-    identity so multiple registrations of the same buffer remain aliased. The
-    byte count is based on each buffer's local storage, which is the actual
-    memory moved for a DTensor rather than its global logical shape.
-    Tensor subclasses such as DTensor keep their type through ``Tensor.to``.
-    """
-
-    moved_count = 0
-    moved_bytes = 0
-    moved_by_id: dict[int, torch.Tensor] = {}
-    for module in model.modules():
-        for name, buffer in module._buffers.items():
-            if buffer is None:
-                continue
-            buffer_id = id(buffer)
-            if buffer_id in moved_by_id:
-                module._buffers[name] = moved_by_id[buffer_id]
-                continue
-            if buffer.device == device:
-                moved_by_id[buffer_id] = buffer
-                continue
-
-            moved_count += 1
-            local_buffer = buffer
-            to_local = getattr(buffer, "to_local", None)
-            if callable(to_local):
-                local_buffer = to_local()
-            moved_bytes += local_buffer.numel() * local_buffer.element_size()
-            moved_buffer = buffer.to(device=device)
-            moved_by_id[buffer_id] = moved_buffer
-            module._buffers[name] = moved_buffer
-    return moved_count, moved_bytes
+    for module, name, buffer in buffers:
+        materialized_buffer = module._buffers[name]
+        materialized_buffer.copy_(buffer.to(device=materialized_buffer.device, dtype=materialized_buffer.dtype))
 
 
 def _veomni_shard_placement_fn(param: "nn.Parameter") -> Optional[Shard]:
@@ -144,6 +87,20 @@ def _veomni_shard_placement_fn(param: "nn.Parameter") -> Optional[Shard]:
     return Shard(dim) if dim is not None else None
 
 
+def _move_buffers_to_device(model: nn.Module, device: str) -> None:
+    """Place every buffer on ``device``.
+
+    ``CPUOffloadPolicy`` only offloads parameters, gradients and optimizer states:
+    ``fully_shard`` puts buffers on the FSDP device once at wrap time and never stages
+    them again. Meta init defers that move, so a CPU materialization leaves buffers such
+    as the DeepSeek-V4 hash-router ``tid2eid`` table or rotary ``inv_freq`` on CPU while
+    the activations are on the accelerator.
+    """
+    for module in model.modules():
+        for name, buffer in module._buffers.items():
+            if buffer is not None and buffer.device.type != device:
+                module._buffers[name] = buffer.to(device)
+
 
 def _check_extra_parallel_dim0_divisibility(model: "nn.Module", para_name: str, ep_fsdp_size: int) -> bool:
     """Return whether EP-local dim-0 can be evenly sharded by ``ep_fsdp_size``."""
@@ -620,15 +577,7 @@ def parallelize_model_fsdp2(
                 )
 
     if materialize_device == "cpu":
-        device_type = get_device_type()
-        compute_device = (
-            torch.device(device_type, get_device_id()) if device_type != "cpu" else torch.device(device_type)
-        )
-        moved_buffers, moved_buffer_bytes = _move_model_buffers_to_device(model, compute_device)
-        logger.info_rank0(
-            "FSDP2 CPU offload kept registered buffers on the compute device: "
-            f"moved_buffers={moved_buffers}, moved_bytes={moved_buffer_bytes}."
-        )
+        _move_buffers_to_device(model, get_device_type())
 
     # Register grad norm clipping method for FSDP2
     from .fsdp2 import clip_grad_norm as clip_grad_norm_fn

From e8bc2d6f697c7d83e3d410e141b1a51962afdff6 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 10 Aug 2026 20:19:08 +0800
Subject: [PATCH 06/12] chore: align CI workflows with upstream main before
 branch update

---
 .github/workflows/gpu_unit_tests.yml | 3 ---
 .github/workflows/npu_unit_tests.yml | 4 +---
 2 files changed, 1 insertion(+), 6 deletions(-)

diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml
index a3db7800e7..0f8ae52d7d 100644
--- a/.github/workflows/gpu_unit_tests.yml
+++ b/.github/workflows/gpu_unit_tests.yml
@@ -206,6 +206,3 @@ jobs:
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_trainer.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep2.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep_sharded_stream_load.py
-      - name: Run DCP checkpointer CPU-offload regression tests
-        run: |
-          uv run --frozen pytest -s -x tests/checkpoints/test_dcp_checkpointer.py
diff --git a/.github/workflows/npu_unit_tests.yml b/.github/workflows/npu_unit_tests.yml
index 5f0bc41eaa..eadcfd44fe 100644
--- a/.github/workflows/npu_unit_tests.yml
+++ b/.github/workflows/npu_unit_tests.yml
@@ -160,6 +160,7 @@ jobs:
           uv run --frozen pytest -v -s -x tests/models/test_models_patch.py
           uv run --frozen pytest -v -s -x tests/models/test_padded_packed_loss.py
           uv run --frozen pytest -v -s -x tests/models/test_vlm_trainer.py
+          uv run --frozen pytest -v -s -x tests/lora/test_vlm_lora.py
           uv run --frozen pytest -v -s -x tests/models/test_checkpoint_tensor_converter.py
       - name: Run fsdp2/extra_parallel+fsdp2 (e.g. ep+fsdp2, mb+fsdp2, ep+emb+fsdp2) clip grad norm test
         run: |
@@ -176,9 +177,6 @@ jobs:
       - name: Run e2e dcp save and load test
         run: |
           uv run --frozen pytest -v -s -x tests/checkpoints/test_trainer_saveload.py
-      - name: Run DCP checkpointer CPU-offload regression tests
-        run: |
-          uv run --frozen pytest -v -s -x tests/checkpoints/test_dcp_checkpointer.py
       - name: Run distributed tests (dummy_forward)
         run: |
           uv run --frozen pytest -v -s -x tests/distributed/test_dummy_forward.py

From 5fc05b9bf81a6a095c2cf0498d1c447d9dc2a6e6 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 10 Aug 2026 20:28:29 +0800
Subject: [PATCH 07/12] fix: reapply FSDP2 CPU-offload fix after syncing
 upstream main

---
 .github/workflows/gpu_unit_tests.yml    |   3 +
 .github/workflows/npu_unit_tests.yml    |   3 +
 veomni/distributed/torch_parallelize.py | 109 +++++++++++++++++-------
 3 files changed, 86 insertions(+), 29 deletions(-)

diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml
index 0f8ae52d7d..a3db7800e7 100644
--- a/.github/workflows/gpu_unit_tests.yml
+++ b/.github/workflows/gpu_unit_tests.yml
@@ -206,3 +206,6 @@ jobs:
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_trainer.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep2.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep_sharded_stream_load.py
+      - name: Run DCP checkpointer CPU-offload regression tests
+        run: |
+          uv run --frozen pytest -s -x tests/checkpoints/test_dcp_checkpointer.py
diff --git a/.github/workflows/npu_unit_tests.yml b/.github/workflows/npu_unit_tests.yml
index eadcfd44fe..bad64641ac 100644
--- a/.github/workflows/npu_unit_tests.yml
+++ b/.github/workflows/npu_unit_tests.yml
@@ -177,6 +177,9 @@ jobs:
       - name: Run e2e dcp save and load test
         run: |
           uv run --frozen pytest -v -s -x tests/checkpoints/test_trainer_saveload.py
+      - name: Run DCP checkpointer CPU-offload regression tests
+        run: |
+          uv run --frozen pytest -v -s -x tests/checkpoints/test_dcp_checkpointer.py
       - name: Run distributed tests (dummy_forward)
         run: |
           uv run --frozen pytest -v -s -x tests/distributed/test_dummy_forward.py
diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 4897d041aa..f9a53d7cb1 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -29,7 +29,7 @@
 from ..arguments import MixedPrecisionConfig
 from ..models import load_model_weights, load_model_weights_ep_sharded, rank0_load_and_broadcast_weights
 from ..utils import logging
-from ..utils.device import IS_NPU_AVAILABLE, get_device_type
+from ..utils.device import IS_NPU_AVAILABLE, get_device_id, get_device_type
 from .checkpoint import CheckpointFunction
 from .chunk_mbs import apply_chunk_mbs
 from .parallel_plan import get_runtime_parallel_plan
@@ -49,25 +49,82 @@ def _reset_hf_initialized_flag(module: nn.Module) -> None:
 
 
 def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) -> None:
-    """Materialize parameters without discarding config-derived buffers.
+    """Materialize parameters while preserving buffer alias groups.
 
-    Distributed checkpoints restore ``state_dict()``, which excludes buffers
-    registered with ``persistent=False``. Snapshot those buffers before
-    ``to_empty()`` so values such as rotary ``inv_freq`` survive the DCP resume
-    materialization path without duplicating persistent buffers that DCP loads.
+    ``Module.to_empty()`` replaces each registered buffer independently, which
+    breaks aliases. Record every registration before materialization and bind
+    each original alias group to one materialized tensor afterwards.
+
+    Distributed checkpoints do not restore groups containing only
+    nonpersistent buffers, so preserve their config-derived values. If a group
+    has any persistent registration, the checkpoint is authoritative and will
+    populate the shared materialized tensor during the subsequent DCP load.
     """
-    buffers = []
+
+    buffer_groups: dict[int, list[tuple[nn.Module, str]]] = {}
+    preserved_nonpersistent_buffers: dict[int, torch.Tensor] = {}
     for module in model.modules():
-        for name in module._non_persistent_buffers_set:
-            buffer = module._buffers.get(name)
-            if buffer is not None:
-                buffers.append((module, name, buffer.detach().clone()))
+        for name, buffer in module._buffers.items():
+            if buffer is None:
+                continue
+            buffer_id = id(buffer)
+            buffer_groups.setdefault(buffer_id, []).append((module, name))
+
+    for buffer_id, registrations in buffer_groups.items():
+        if all(name in module._non_persistent_buffers_set for module, name in registrations):
+            module, name = registrations[0]
+            preserved_nonpersistent_buffers[buffer_id] = module._buffers[name].detach().clone()
 
     model.to_empty(device=device)
 
-    for module, name, buffer in buffers:
-        materialized_buffer = module._buffers[name]
-        materialized_buffer.copy_(buffer.to(device=materialized_buffer.device, dtype=materialized_buffer.dtype))
+    for buffer_id, registrations in buffer_groups.items():
+        if buffer_id in preserved_nonpersistent_buffers:
+            canonical_buffer = preserved_nonpersistent_buffers[buffer_id].to(device=device)
+        else:
+            module, name = registrations[0]
+            canonical_buffer = module._buffers[name]
+        for module, name in registrations:
+            module._buffers[name] = canonical_buffer
+
+
+@torch.no_grad()
+def _move_model_buffers_to_device(model: nn.Module, device: torch.device) -> tuple[int, int]:
+    """Move registered buffers without moving CPU-offloaded parameters.
+
+    PyTorch's ``Module._apply`` replaces buffer tensor objects during device
+    conversion, so external references are not guaranteed to retain object
+    identity. This helper follows that contract while memoizing by object
+    identity so multiple registrations of the same buffer remain aliased. The
+    byte count is based on each buffer's local storage, which is the actual
+    memory moved for a DTensor rather than its global logical shape.
+    Tensor subclasses such as DTensor keep their type through ``Tensor.to``.
+    """
+
+    moved_count = 0
+    moved_bytes = 0
+    moved_by_id: dict[int, torch.Tensor] = {}
+    for module in model.modules():
+        for name, buffer in module._buffers.items():
+            if buffer is None:
+                continue
+            buffer_id = id(buffer)
+            if buffer_id in moved_by_id:
+                module._buffers[name] = moved_by_id[buffer_id]
+                continue
+            if buffer.device == device:
+                moved_by_id[buffer_id] = buffer
+                continue
+
+            moved_count += 1
+            local_buffer = buffer
+            to_local = getattr(buffer, "to_local", None)
+            if callable(to_local):
+                local_buffer = to_local()
+            moved_bytes += local_buffer.numel() * local_buffer.element_size()
+            moved_buffer = buffer.to(device=device)
+            moved_by_id[buffer_id] = moved_buffer
+            module._buffers[name] = moved_buffer
+    return moved_count, moved_bytes
 
 
 def _veomni_shard_placement_fn(param: "nn.Parameter") -> Optional[Shard]:
@@ -87,20 +144,6 @@ def _veomni_shard_placement_fn(param: "nn.Parameter") -> Optional[Shard]:
     return Shard(dim) if dim is not None else None
 
 
-def _move_buffers_to_device(model: nn.Module, device: str) -> None:
-    """Place every buffer on ``device``.
-
-    ``CPUOffloadPolicy`` only offloads parameters, gradients and optimizer states:
-    ``fully_shard`` puts buffers on the FSDP device once at wrap time and never stages
-    them again. Meta init defers that move, so a CPU materialization leaves buffers such
-    as the DeepSeek-V4 hash-router ``tid2eid`` table or rotary ``inv_freq`` on CPU while
-    the activations are on the accelerator.
-    """
-    for module in model.modules():
-        for name, buffer in module._buffers.items():
-            if buffer is not None and buffer.device.type != device:
-                module._buffers[name] = buffer.to(device)
-
 
 def _check_extra_parallel_dim0_divisibility(model: "nn.Module", para_name: str, ep_fsdp_size: int) -> bool:
     """Return whether EP-local dim-0 can be evenly sharded by ``ep_fsdp_size``."""
@@ -577,7 +620,15 @@ def parallelize_model_fsdp2(
                 )
 
     if materialize_device == "cpu":
-        _move_buffers_to_device(model, get_device_type())
+        device_type = get_device_type()
+        compute_device = (
+            torch.device(device_type, get_device_id()) if device_type != "cpu" else torch.device(device_type)
+        )
+        moved_buffers, moved_buffer_bytes = _move_model_buffers_to_device(model, compute_device)
+        logger.info_rank0(
+            "FSDP2 CPU offload kept registered buffers on the compute device: "
+            f"moved_buffers={moved_buffers}, moved_bytes={moved_buffer_bytes}."
+        )
 
     # Register grad norm clipping method for FSDP2
     from .fsdp2 import clip_grad_norm as clip_grad_norm_fn

From 601ec99409a5bb4d0289b88a8960ab26a8e3f13a Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Tue, 25 Aug 2026 12:14:10 +0800
Subject: [PATCH 08/12] [ci] fix: run DCP checkpointer tests once per workflow

---
 .github/workflows/gpu_unit_tests.yml | 3 ---
 .github/workflows/npu_unit_tests.yml | 3 ---
 2 files changed, 6 deletions(-)

diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml
index 07e604d3c6..01566f6291 100644
--- a/.github/workflows/gpu_unit_tests.yml
+++ b/.github/workflows/gpu_unit_tests.yml
@@ -216,6 +216,3 @@ jobs:
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_trainer.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep2.py
           uv run --frozen pytest -s -x tests/lora/test_moe_lora_ep_sharded_stream_load.py
-      - name: Run DCP checkpointer CPU-offload regression tests
-        run: |
-          uv run --frozen pytest -s -x tests/checkpoints/test_dcp_checkpointer.py
diff --git a/.github/workflows/npu_unit_tests.yml b/.github/workflows/npu_unit_tests.yml
index 721b1b02c6..f531cd694b 100644
--- a/.github/workflows/npu_unit_tests.yml
+++ b/.github/workflows/npu_unit_tests.yml
@@ -180,9 +180,6 @@ jobs:
       - name: Run e2e dcp save and load test
         run: |
           uv run --frozen pytest -v -s -x tests/checkpoints/test_trainer_saveload.py
-      - name: Run DCP checkpointer CPU-offload regression tests
-        run: |
-          uv run --frozen pytest -v -s -x tests/checkpoints/test_dcp_checkpointer.py
       - name: Run distributed tests (dummy_forward)
         run: |
           uv run --frozen pytest -v -s -x tests/distributed/test_dummy_forward.py

From 804d743ef09c1ff8cd33e57ff99478fcddb472c9 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Fri, 23 Oct 2026 16:28:31 +0800
Subject: [PATCH 09/12] [ci] test: schedule native CPU-offload buffer
 regression on NPU

---
 scripts/ci/accelerator/suites/npu.json | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/scripts/ci/accelerator/suites/npu.json b/scripts/ci/accelerator/suites/npu.json
index 5202287834..854f08c138 100644
--- a/scripts/ci/accelerator/suites/npu.json
+++ b/scripts/ci/accelerator/suites/npu.json
@@ -396,6 +396,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-01",
     "suite": "e2e",

From 23a86a02730f80e58b1a7af6490ffc990efbf447 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 26 Oct 2026 19:17:34 +0800
Subject: [PATCH 10/12] [checkpoint] fix: preserve shared buffers during HF
 loading

---
 .../models/base/test_shared_buffer_loading.py | 232 ++++++++++++++++++
 veomni/distributed/torch_parallelize.py       |   3 +
 veomni/models/checkpoint/weights.py           |  80 +++++-
 3 files changed, 302 insertions(+), 13 deletions(-)
 create mode 100644 tests/models/base/test_shared_buffer_loading.py

diff --git a/tests/models/base/test_shared_buffer_loading.py b/tests/models/base/test_shared_buffer_loading.py
new file mode 100644
index 0000000000..cdb0bf8f2f
--- /dev/null
+++ b/tests/models/base/test_shared_buffer_loading.py
@@ -0,0 +1,232 @@
+# Copyright 2026 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""HF materialization must preserve shared-buffer values and registrations."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+from safetensors.torch import save_file
+from torch import nn
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.tensor import Replicate, distribute_tensor
+
+from scripts.ci.accelerator.launch_plan import LaunchSpec
+from tests.tools.launch_utils import spawn_with_rendezvous
+from veomni.distributed.parallel_plan import ParallelPlan
+from veomni.distributed.parallel_state import _init_parallel_state, clear_parallel_state
+from veomni.distributed.torch_parallelize import _materialize_and_load_weights, _move_model_buffers_to_device
+from veomni.models.checkpoint import weights
+from veomni.models.checkpoint.convert import ConvertedCheckpointTensor
+
+
+pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+_PERSISTENCE = [(True, False), (False, True), (False, False), (True, True)]
+
+
+class _SharedBuffers(nn.Module):
+    def __init__(self, root_persistent, child_persistent, layout="direct"):
+        super().__init__()
+        self.config = SimpleNamespace(tie_word_embeddings=False)
+        self.layout = layout
+        self.weight = nn.Parameter(torch.zeros(2))
+        self.child = nn.Module()
+        shared = torch.tensor([7.0, 11.0])
+        self.register_buffer("root_buffer", shared, persistent=root_persistent)
+        self.child.register_buffer("child_buffer", shared, persistent=child_persistent)
+        self.register_buffer("optional", None)
+
+    def get_parallel_plan(self):
+        # Buffers are read whole even by the EP streaming loader.
+        return ParallelPlan(extra_parallel_plan={"ep": {}})
+
+    @staticmethod
+    def _create_checkpoint_tensor_converter(model):
+        return None if model.layout == "direct" else _BufferAliasConverter(model.layout)
+
+
+class _BufferAliasConverter:
+    def __init__(self, layout):
+        self.layout = layout
+        self.pending = []
+
+    def can_handle(self, name):
+        return name.startswith("legacy.")
+
+    def convert(self, name, tensor):
+        converted = ConvertedCheckpointTensor(name=name.removeprefix("legacy."), tensor=tensor)
+        if self.layout == "finalize":
+            self.pending.append(converted)
+            return None
+        return converted
+
+    def finalize(self):
+        pending, self.pending = self.pending, []
+        return pending
+
+    def fused_expert_target(self, name):
+        return None
+
+    def for_expert_range(self, start, num_local):
+        return _BufferAliasConverter(self.layout)
+
+
+def _write_checkpoint(directory, root_persistent, child_persistent, layout="direct"):
+    directory.mkdir(exist_ok=True)
+    payload = {"weight": torch.tensor([1.0, 2.0])}
+    for name, persistent in (("root_buffer", root_persistent), ("child.child_buffer", child_persistent)):
+        if persistent:
+            prefix = "" if layout == "direct" else "legacy."
+            payload[prefix + name] = torch.tensor([21.0, 34.0])
+    save_file(payload, str(directory / "model.safetensors"))
+
+
+def _assert_shared_values(model, root_persistent, child_persistent):
+    root, child = model.root_buffer, model.child.child_buffer
+    expected = torch.tensor([21.0, 34.0] if root_persistent or child_persistent else [7.0, 11.0])
+    local = root.to_local() if hasattr(root, "to_local") else root
+    assert torch.equal(local, expected)
+    assert root.dtype == expected.dtype
+    assert root is child
+    assert model.optional is None
+    expected_keys = {"weight"}
+    if root_persistent:
+        expected_keys.add("root_buffer")
+    if child_persistent:
+        expected_keys.add("child.child_buffer")
+    assert set(model.state_dict()) == expected_keys
+    assert torch.equal(model.weight, torch.tensor([1.0, 2.0]))
+    _move_model_buffers_to_device(model, torch.device("cpu"))
+    assert model.root_buffer is model.child.child_buffer
+    model.root_buffer.fill_(123)
+    child = model.child.child_buffer
+    local = child.to_local() if hasattr(child, "to_local") else child
+    assert torch.equal(local, torch.full((2,), 123.0))
+
+
+@pytest.fixture
+def cpu_group(tmp_path, monkeypatch):
+    assert not dist.is_initialized()
+    clear_parallel_state()
+    dist.init_process_group("gloo", init_method=f"file://{tmp_path / 'rendezvous'}", rank=0, world_size=1)
+    _init_parallel_state(device_type="cpu")
+    monkeypatch.setattr(weights, "_get_communication_device", lambda _: torch.device("cpu"))
+    try:
+        yield
+    finally:
+        clear_parallel_state()
+        dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
+@pytest.mark.parametrize("loader", ["plain", "rank0", "ep_sharded"])
+@pytest.mark.parametrize("layout", ["direct", "rename", "finalize"])
+def test_hf_shared_buffer_values_and_identity(tmp_path, cpu_group, root_persistent, child_persistent, loader, layout):
+    _write_checkpoint(tmp_path, root_persistent, child_persistent, layout)
+    with weights.init_empty_weights():
+        model = _SharedBuffers(root_persistent, child_persistent, layout)
+    loaders = {
+        "plain": weights.load_model_weights,
+        "rank0": weights.rank0_load_and_broadcast_weights,
+        "ep_sharded": weights.load_model_weights_ep_sharded,
+    }
+    loaders[loader](model, str(tmp_path), init_device="cpu")
+    _assert_shared_values(model, root_persistent, child_persistent)
+
+
+@pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
+def test_hf_shared_dtensor_buffers_rebind_after_dispatch(tmp_path, cpu_group, root_persistent, child_persistent):
+    _write_checkpoint(tmp_path, root_persistent, child_persistent)
+    with weights.init_empty_weights():
+        model = _SharedBuffers(root_persistent, child_persistent)
+    shared = distribute_tensor(model.root_buffer, init_device_mesh("cpu", (1,)), [Replicate()])
+    model.root_buffer = model.child.child_buffer = shared
+    weights.load_model_weights(model, str(tmp_path), init_device="cpu", dtensor_factory=distribute_tensor)
+    _assert_shared_values(model, root_persistent, child_persistent)
+
+
+@pytest.mark.parametrize("first_persistent,second_persistent", _PERSISTENCE)
+@pytest.mark.parametrize("reverse", [False, True])
+def test_mapping_load_preserves_cross_child_buffers(tmp_path, first_persistent, second_persistent, reverse):
+    parent = nn.Module()
+    shared = torch.tensor([7.0, 11.0])
+    parent.register_buffer("root_alias", shared, persistent=False)
+    paths = {}
+    for name, persistent in (("first", first_persistent), ("second", second_persistent)):
+        with weights.init_empty_weights():
+            child = nn.Module()
+            child.weight = nn.Parameter(torch.zeros(2))
+        child.config = SimpleNamespace(tie_word_embeddings=False)
+        child.register_buffer("buffer", shared, persistent=persistent)
+        parent.add_module(name, child)
+        directory = tmp_path / name
+        directory.mkdir()
+        payload = {"weight": torch.tensor([1.0, 2.0])}
+        if persistent:
+            payload["buffer"] = torch.tensor([21.0, 34.0])
+        save_file(payload, str(directory / "model.safetensors"))
+        paths[name] = str(directory)
+    if reverse:
+        paths = dict(reversed(paths.items()))
+    _materialize_and_load_weights(
+        parent,
+        paths,
+        "cpu",
+        should_skip_hf_weight_load=False,
+        is_peft_model=False,
+        adapter_path=None,
+        broadcast_from_rank0=False,
+    )
+    expected = torch.tensor([21.0, 34.0] if first_persistent or second_persistent else [7.0, 11.0])
+    assert torch.equal(parent.root_alias, expected)
+    assert parent.root_alias is parent.first.buffer is parent.second.buffer
+    expected_keys = {"first.weight", "second.weight"}
+    if first_persistent:
+        expected_keys.add("first.buffer")
+    if second_persistent:
+        expected_keys.add("second.buffer")
+    assert set(parent.state_dict()) == expected_keys
+    for name in ("first", "second"):
+        assert torch.equal(getattr(parent, name).weight, torch.tensor([1.0, 2.0]))
+    parent.first.buffer.fill_(123)
+    assert torch.equal(parent.second.buffer, torch.full((2,), 123.0))
+
+
+def _rank0_alias_worker(rank, directory, endpoint):
+    dist.init_process_group("gloo", store=endpoint.connect(), rank=rank, world_size=2)
+    _init_parallel_state(dp_size=2, device_type="cpu")
+    # This is a CPU/Gloo test on both CPU and accelerator CI runners.
+    weights._get_communication_device = lambda _: torch.device("cpu")
+    try:
+        for layout in ("direct", "rename", "finalize"):
+            for index, (root_persistent, child_persistent) in enumerate(_PERSISTENCE):
+                with weights.init_empty_weights():
+                    model = _SharedBuffers(root_persistent, child_persistent, layout)
+                weights.rank0_load_and_broadcast_weights(
+                    model, str(directory / f"{layout}-{index}"), init_device="cpu"
+                )
+                _assert_shared_values(model, root_persistent, child_persistent)
+    finally:
+        clear_parallel_state()
+        dist.destroy_process_group()
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(nproc=2, device="cpu"))
+def test_rank0_shared_buffers_on_receiving_rank(tmp_path):
+    for layout in ("direct", "rename", "finalize"):
+        for index, persistence in enumerate(_PERSISTENCE):
+            _write_checkpoint(tmp_path / f"{layout}-{index}", *persistence, layout)
+    spawn_with_rendezvous(_rank0_alias_worker, (tmp_path,), nprocs=2)
diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 119d8e6ad4..98ce063473 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -30,6 +30,7 @@
 from ..arguments import MixedPrecisionConfig
 from ..arguments.arguments_types import validate_low_precision_reduce_scatter_comm
 from ..models import load_model_weights, load_model_weights_ep_sharded, rank0_load_and_broadcast_weights
+from ..models.checkpoint.weights import _get_buffer_aliases, _restore_buffer_aliases
 from ..utils import logging, recompute_utils
 from ..utils.device import IS_NPU_AVAILABLE, get_device_id, get_device_type
 from .checkpoint import CheckpointFunction
@@ -220,6 +221,7 @@ def _materialize_and_load_weights(
         skip_map = dict(module_skip_hf_weight_load or {})
         peft_map = dict(module_is_peft_model or {})
         adapter_map = dict(module_adapter_path or {})
+        buffer_aliases = _get_buffer_aliases(model)
         for name, path in weights_path.items():
             _materialize_and_load_weights(
                 children[name],
@@ -234,6 +236,7 @@ def _materialize_and_load_weights(
                 fqn_to_index_mapping=fqn_to_index_mapping,
                 ep_sharded_stream_load=ep_sharded_stream_load,
             )
+            _restore_buffer_aliases(model, buffer_aliases, source_prefix=f"{name}.")
         leftover_meta = [
             name
             for name, child in children.items()
diff --git a/veomni/models/checkpoint/weights.py b/veomni/models/checkpoint/weights.py
index f20a9ada6a..a78fe43d7b 100644
--- a/veomni/models/checkpoint/weights.py
+++ b/veomni/models/checkpoint/weights.py
@@ -312,6 +312,34 @@ def _dispatch_parameter(
         module._parameters[local_name].data.copy_(tensor)
 
 
+def _get_buffer_aliases(model: "nn.Module") -> Dict[str, str]:
+    """Record tensor-identity aliases before ``to_empty`` splits registrations."""
+    groups: Dict[int, List[str]] = {}
+    for name, buffer in model.named_buffers(remove_duplicate=False):
+        groups.setdefault(id(buffer), []).append(name)
+    return {name: names[0] for names in groups.values() if len(names) > 1 for name in names}
+
+
+def _restore_buffer_aliases(model: "nn.Module", aliases: Dict[str, str], source_prefix: Optional[str] = None) -> None:
+    """Rebind aliases after dispatch, including dispatch that replaces DTensors.
+
+    Composed-model loaders can select the just-loaded child as the source, so
+    the next child's snapshot sees the loaded value rather than stale config.
+    Registration persistence flags are deliberately left unchanged.
+    """
+    sources = {}
+    for name, canonical in aliases.items():
+        if source_prefix is None:
+            sources.setdefault(canonical, canonical)
+        elif name.startswith(source_prefix):
+            sources.setdefault(canonical, name)
+    buffers = {canonical: model.get_buffer(name) for canonical, name in sources.items()}
+    for name, canonical in aliases.items():
+        if canonical in buffers:
+            module, local_name = _find_submodule(model, name)
+            module._buffers[local_name] = buffers[canonical]
+
+
 def _dispatch_buffer(
     module: "nn.Module",
     name: str,
@@ -417,6 +445,7 @@ def load_model_weights(
     """
     Loads pre-trained model states in transformers' format.
     """
+    buffer_aliases = _get_buffer_aliases(model)
     buffer_dict = {name: buffer.clone() for name, buffer in model.named_buffers()}
     parameter_names_to_load = {name for name, _ in model.named_parameters()}
     model.to_empty(device=init_device)
@@ -448,8 +477,9 @@ def load_model_weights(
     state_dict_iterators = _load_state_dict(weights_path, skip_key=_skip_without_loading_fn(model, converter))
 
     def _dispatch_kv(name: str, tensor: "torch.Tensor") -> None:
-        if name in buffer_dict.keys():  # persistent buffers
-            buffer_dict[name] = tensor.clone()
+        buffer_name = buffer_aliases.get(name, name)
+        if buffer_name in buffer_dict:  # persistent buffers, including noncanonical aliases
+            buffer_dict[buffer_name] = tensor.clone()
         elif name in parameter_names_to_load:
             parameter_names_to_load.remove(name)
             _dispatch_parameter(model, name, tensor, dtensor_factory, parallel_plan, dtensor_to_cpu)
@@ -492,7 +522,12 @@ def _dispatch_kv(name: str, tensor: "torch.Tensor") -> None:
         )
 
     post_process_after_weight_loading(
-        model, buffer_dict, parameter_names_to_load, dtensor_factory, dtensor_to_cpu=dtensor_to_cpu
+        model,
+        buffer_dict,
+        parameter_names_to_load,
+        dtensor_factory,
+        dtensor_to_cpu=dtensor_to_cpu,
+        buffer_aliases=buffer_aliases,
     )
 
     fqn_to_index_mapping = kwargs.get("fqn_to_index_mapping")
@@ -768,6 +803,7 @@ def load_model_weights_ep_sharded(
                     f"'{_pname}' pattern '{_fqn_pattern}' uses Shard({_dim})."
                 )
 
+    buffer_aliases = _get_buffer_aliases(model)
     buffer_dict = {name: buffer.clone() for name, buffer in model.named_buffers()}
     params_by_name = dict(model.named_parameters())
     param_shapes = {name: tuple(p.shape) for name, p in params_by_name.items()}
@@ -895,8 +931,9 @@ def _dispatch_whole_converted(converted, source: str) -> None:
                 f"ep_sharded_stream_load: converter emitted ExtraParallel key '{dest}' from non-expert "
                 f"source '{source}'. Per-expert keys must go through fused_expert_target / for_expert_range."
             )
-        if dest in buffer_dict:
-            buffer_dict[dest] = converted.tensor.clone()
+        buffer_name = buffer_aliases.get(dest, dest)
+        if buffer_name in buffer_dict:
+            buffer_dict[buffer_name] = converted.tensor.clone()
             n_buf += 1
         elif dest in parameter_names_to_load:
             _dispatch_parameter(model, dest, converted.tensor, dtensor_factory, parallel_plan, dtensor_to_cpu)
@@ -963,8 +1000,9 @@ def _dispatch_whole_converted(converted, source: str) -> None:
                     if converted is not None:
                         _dispatch_whole_converted(converted, bare_name)
                     continue
-                if name in buffer_dict:  # persistent buffers: read whole
-                    buffer_dict[name] = f.get_tensor(raw_name).clone()
+                buffer_name = buffer_aliases.get(name, name)
+                if buffer_name in buffer_dict:  # persistent buffers: read whole, including aliases
+                    buffer_dict[buffer_name] = f.get_tensor(raw_name).clone()
                     n_buf += 1
                     continue
                 if name not in parameter_names_to_load:
@@ -1062,7 +1100,7 @@ def _dispatch_whole_converted(converted, source: str) -> None:
         # whole-tensor loader does.
         for result in expert_converter.finalize():
             dest = _apply_peft_override(result.name)
-            if dest in parameter_names_to_load or dest in buffer_dict:
+            if dest in parameter_names_to_load or buffer_aliases.get(dest, dest) in buffer_dict:
                 raise NotImplementedError(
                     f"ep_sharded_stream_load: expert-range converter emitted '{dest}' at finalize time; "
                     f"rank-local finalize outputs cannot be dispatched."
@@ -1094,7 +1132,12 @@ def _dispatch_whole_converted(converted, source: str) -> None:
         )
 
     post_process_after_weight_loading(
-        model, buffer_dict, parameter_names_to_load, dtensor_factory, dtensor_to_cpu=dtensor_to_cpu
+        model,
+        buffer_dict,
+        parameter_names_to_load,
+        dtensor_factory,
+        dtensor_to_cpu=dtensor_to_cpu,
+        buffer_aliases=buffer_aliases,
     )
 
     fqn_to_index_mapping = kwargs.get("fqn_to_index_mapping")
@@ -1223,6 +1266,7 @@ def rank0_load_and_broadcast_weights(
         logger.warning_once("Distributed environment not initialized, falling back to load_model_weights.")
         return load_model_weights(model, weights_path, init_device, dtensor_factory, **kwargs)
 
+    buffer_aliases = _get_buffer_aliases(model)
     buffer_dict = {name: buffer.clone() for name, buffer in model.named_buffers()}
     parameter_names_to_load = {name for name, _ in model.named_parameters()}
     model.to_empty(device=init_device)
@@ -1275,8 +1319,9 @@ def _broadcast_and_dispatch(name, shape, dtype, tensor):
         if broadcast_as_bytes:
             tensor = _view_from_bytes(broadcast_tensor, dtype, shape)
 
-        if name in buffer_dict:
-            buffer_dict[name] = tensor.detach().clone()
+        buffer_name = buffer_aliases.get(name, name)
+        if buffer_name in buffer_dict:
+            buffer_dict[buffer_name] = tensor.detach().clone()
         elif name in parameter_names_to_load:
             parameter_names_to_load.discard(name)
             _dispatch_parameter(model, name, tensor, dtensor_factory, parallel_plan, dtensor_to_cpu)
@@ -1333,7 +1378,7 @@ def _chunk_and_broadcast_and_dispatch(name, shape, dtype, tensor):
         """Broadcast a single (name, tensor) from rank0 and dispatch it."""
         logger.info_rank0(f"rank0_load_and_broadcast_weights: chunking and broadcasting {name=}")
 
-        assert name not in buffer_dict, f"Buffer {name} should not be chunked."
+        assert buffer_aliases.get(name, name) not in buffer_dict, f"Buffer {name} should not be chunked."
         assert name in parameter_names_to_load, f"Unexpected key in state dict: {name}."
 
         if global_rank == 0:
@@ -1640,7 +1685,12 @@ def _chunk_and_broadcast_and_dispatch(name, shape, dtype, tensor):
         )
 
     post_process_after_weight_loading(
-        model, buffer_dict, parameter_names_to_load, dtensor_factory, dtensor_to_cpu=dtensor_to_cpu
+        model,
+        buffer_dict,
+        parameter_names_to_load,
+        dtensor_factory,
+        dtensor_to_cpu=dtensor_to_cpu,
+        buffer_aliases=buffer_aliases,
     )
 
     fqn_to_index_mapping = kwargs.get("fqn_to_index_mapping")
@@ -1656,6 +1706,7 @@ def post_process_after_weight_loading(
     parameter_names_left: Optional[set[str]] = None,
     dtensor_factory: Optional[Callable[["torch.Tensor", Any, Any], "torch.Tensor"]] = None,
     dtensor_to_cpu: bool = False,
+    buffer_aliases: Optional[Dict[str, str]] = None,
 ):
     """
     shared logic after weight loading that handles buffer, missing weight keys and tied embedding weights.
@@ -1691,6 +1742,9 @@ def post_process_after_weight_loading(
             logger.info_rank0(f"Failed to tie embeddings: {e}")
             raise RuntimeError("Failed to tie input/output embeddings") from e
 
+    if buffer_aliases:
+        _restore_buffer_aliases(model, buffer_aliases)
+
 
 def _normalize_save_dtype(save_dtype: Optional[Union[str, "torch.dtype"]]) -> Optional["torch.dtype"]:
     if isinstance(save_dtype, str):

From 75ec4732db91a720ce6011dfff2733462fa9b090 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 26 Oct 2026 19:33:16 +0800
Subject: [PATCH 11/12] [checkpoint] fix: retain constructor buffers during
 random initialization

---
 .../models/base/test_shared_buffer_loading.py | 76 +++++++++++++++++++
 veomni/distributed/torch_parallelize.py       | 29 ++++---
 2 files changed, 95 insertions(+), 10 deletions(-)

diff --git a/tests/models/base/test_shared_buffer_loading.py b/tests/models/base/test_shared_buffer_loading.py
index cdb0bf8f2f..1e9cd344e2 100644
--- a/tests/models/base/test_shared_buffer_loading.py
+++ b/tests/models/base/test_shared_buffer_loading.py
@@ -84,6 +84,82 @@ def for_expert_range(self, start, num_local):
         return _BufferAliasConverter(self.layout)
 
 
+class _RandomInitSharedBuffers(_SharedBuffers):
+    def __init__(self, root_persistent, child_persistent, overwrite_buffer=False):
+        super().__init__(root_persistent, child_persistent)
+        self.register_buffer("standalone", torch.tensor([5.0, 8.0]))
+        self.overwrite_buffer = overwrite_buffer
+        self.init_calls = 0
+
+    def init_weights(self):
+        self.init_calls += 1
+        with torch.no_grad():
+            self.weight.copy_(torch.tensor([1.0, 2.0]))
+            if self.overwrite_buffer:
+                self.child.child_buffer.copy_(torch.tensor([31.0, 53.0]))
+
+
+@pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
+@pytest.mark.parametrize("overwrite_buffer", [False, True])
+def test_random_init_preserves_real_buffers(root_persistent, child_persistent, overwrite_buffer):
+    with weights.init_empty_weights():
+        model = _RandomInitSharedBuffers(root_persistent, child_persistent, overwrite_buffer)
+    _materialize_and_load_weights(
+        model,
+        None,
+        "cpu",
+        should_skip_hf_weight_load=False,
+        is_peft_model=False,
+        adapter_path=None,
+        broadcast_from_rank0=False,
+    )
+    expected = torch.tensor([31.0, 53.0] if overwrite_buffer else [7.0, 11.0])
+    assert model.init_calls == 1
+    assert model.root_buffer is model.child.child_buffer
+    assert torch.equal(model.root_buffer, expected)
+    assert torch.equal(model.standalone, torch.tensor([5.0, 8.0]))
+    assert torch.equal(model.weight, torch.tensor([1.0, 2.0]))
+
+
+@pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
+def test_resume_keeps_checkpoint_authority(monkeypatch, root_persistent, child_persistent):
+    with weights.init_empty_weights():
+        model = _RandomInitSharedBuffers(root_persistent, child_persistent)
+    to_empty = model.to_empty
+
+    def sentinel_to_empty(*, device):
+        result = to_empty(device=device)
+        for buffer in model.buffers():
+            buffer.fill_(-1234)
+        return result
+
+    monkeypatch.setattr(model, "to_empty", sentinel_to_empty)
+    _materialize_and_load_weights(
+        model,
+        "must-not-read-hf",
+        "cpu",
+        should_skip_hf_weight_load=True,
+        is_peft_model=False,
+        adapter_path=None,
+        broadcast_from_rank0=False,
+    )
+    expected = torch.full((2,), -1234.0) if root_persistent or child_persistent else torch.tensor([7.0, 11.0])
+    assert model.init_calls == 0
+    assert model.root_buffer is model.child.child_buffer
+    assert torch.equal(model.root_buffer, expected)
+    assert torch.equal(model.standalone, torch.full((2,), -1234.0))
+    checkpoint = {"weight": torch.tensor([1.0, 2.0]), "standalone": torch.tensor([42.0, 43.0])}
+    if root_persistent:
+        checkpoint["root_buffer"] = torch.tensor([21.0, 34.0])
+    if child_persistent:
+        checkpoint["child.child_buffer"] = torch.tensor([21.0, 34.0])
+    model.load_state_dict(checkpoint, strict=True)
+    expected = torch.tensor([21.0, 34.0] if root_persistent or child_persistent else [7.0, 11.0])
+    assert model.root_buffer is model.child.child_buffer
+    assert torch.equal(model.root_buffer, expected)
+    assert torch.equal(model.standalone, torch.tensor([42.0, 43.0]))
+
+
 def _write_checkpoint(directory, root_persistent, child_persistent, layout="direct"):
     directory.mkdir(exist_ok=True)
     payload = {"weight": torch.tensor([1.0, 2.0])}
diff --git a/veomni/distributed/torch_parallelize.py b/veomni/distributed/torch_parallelize.py
index 98ce063473..80abae8741 100644
--- a/veomni/distributed/torch_parallelize.py
+++ b/veomni/distributed/torch_parallelize.py
@@ -54,7 +54,9 @@ def _reset_hf_initialized_flag(module: nn.Module) -> None:
         _reset_hf_initialized_flag(child)
 
 
-def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) -> None:
+def _to_empty_preserving_nonpersistent_buffers(
+    model: nn.Module, device: str, *, preserve_persistent_buffers: bool = False
+) -> None:
     """Materialize parameters while preserving buffer alias groups.
 
     ``Module.to_empty()`` replaces each registered buffer independently, which
@@ -69,6 +71,10 @@ def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) ->
     registration, the checkpoint is authoritative and will populate the shared
     materialized tensor during the subsequent DCP load.
 
+    For random initialization, ``preserve_persistent_buffers=True`` preserves
+    all real buffer values instead: no checkpoint will restore them. The
+    subsequent ``init_weights()`` call can still overwrite these initial values.
+
     HF's ``_init_weights`` recomputes a rope table only for a module that
     exposes ``original_inv_freq``. Buffers outside that shape -- Gemma3's
     per-layer-type ``{type}_inv_freq`` and its ``embed_scale`` -- have nothing
@@ -76,7 +82,7 @@ def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) ->
     """
 
     buffer_groups: dict[int, list[tuple[nn.Module, str]]] = {}
-    preserved_nonpersistent_buffers: dict[int, torch.Tensor] = {}
+    preserved_buffers: dict[int, torch.Tensor] = {}
     for module in model.modules():
         for name, buffer in module._buffers.items():
             if buffer is None:
@@ -87,7 +93,9 @@ def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) ->
     for buffer_id, registrations in buffer_groups.items():
         module, name = registrations[0]
         buffer = module._buffers[name]
-        if not all(name in module._non_persistent_buffers_set for module, name in registrations):
+        if not preserve_persistent_buffers and not all(
+            name in module._non_persistent_buffers_set for module, name in registrations
+        ):
             continue
         if buffer.is_meta:
             # A buffer derived from a parameter (``self.weight.detach()``) is
@@ -95,18 +103,18 @@ def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) ->
             # No model does this today; say so rather than let ``to_empty()``
             # leave uninitialized memory behind unannounced.
             logger.warning_rank0(
-                f"Non-persistent buffer {name!r} on {type(module).__name__} is on meta and cannot be "
+                f"Buffer {name!r} on {type(module).__name__} is on meta and cannot be "
                 "preserved across materialization; it will hold uninitialized memory unless init_weights() sets it."
             )
             continue
-        preserved_nonpersistent_buffers[buffer_id] = buffer.detach().clone()
+        preserved_buffers[buffer_id] = buffer.detach().clone()
 
     model.to_empty(device=device)
 
     for buffer_id, registrations in buffer_groups.items():
         module, name = registrations[0]
         canonical_buffer = module._buffers[name]
-        snapshot = preserved_nonpersistent_buffers.get(buffer_id)
+        snapshot = preserved_buffers.get(buffer_id)
         if snapshot is not None:
             canonical_buffer.copy_(snapshot.to(device=canonical_buffer.device, dtype=canonical_buffer.dtype))
         for module, name in registrations:
@@ -263,10 +271,11 @@ def _materialize_and_load_weights(
                 "Skipping pretrained weight load for checkpoint resume; "
                 "parameters will be restored from the distributed checkpoint."
             )
-        # Preserve non-persistent buffers on the random-init path too: HF's
-        # ``_init_weights`` recomputes only rope tables shaped the way it expects,
-        # so the rest would train on whatever ``to_empty()`` left behind.
-        _to_empty_preserving_nonpersistent_buffers(model, materialize_device)
+        # Random init has no checkpoint to restore even persistent buffers;
+        # preserve their constructor values before calling model.init_weights().
+        _to_empty_preserving_nonpersistent_buffers(
+            model, materialize_device, preserve_persistent_buffers=not should_skip_hf_weight_load
+        )
         _reset_hf_initialized_flag(model)
         # Random init is unnecessary when the checkpoint will overwrite every parameter.
         if not should_skip_hf_weight_load:

From b53ccf5f8b9a272b202f44e5e3df5be6cbb4b096 Mon Sep 17 00:00:00 2001
From: YLlllllllllll <114622851+YLlllllllllll@users.noreply.github.com>
Date: Mon, 26 Oct 2026 20:18:50 +0800
Subject: [PATCH 12/12] [checkpoint] test: cover random buffer materialization
 on accelerators

---
 .../models/base/test_shared_buffer_loading.py | 40 +++++++++++++++----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/tests/models/base/test_shared_buffer_loading.py b/tests/models/base/test_shared_buffer_loading.py
index 1e9cd344e2..c20ff43678 100644
--- a/tests/models/base/test_shared_buffer_loading.py
+++ b/tests/models/base/test_shared_buffer_loading.py
@@ -31,9 +31,10 @@
 from veomni.distributed.torch_parallelize import _materialize_and_load_weights, _move_model_buffers_to_device
 from veomni.models.checkpoint import weights
 from veomni.models.checkpoint.convert import ConvertedCheckpointTensor
+from veomni.utils.device import get_device_id, get_device_type, get_torch_device
 
 
-pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+_CPU_PLAN = LaunchSpec(device="cpu")
 _PERSISTENCE = [(True, False), (False, True), (False, False), (True, True)]
 
 
@@ -99,15 +100,13 @@ def init_weights(self):
                 self.child.child_buffer.copy_(torch.tensor([31.0, 53.0]))
 
 
-@pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
-@pytest.mark.parametrize("overwrite_buffer", [False, True])
-def test_random_init_preserves_real_buffers(root_persistent, child_persistent, overwrite_buffer):
+def _check_random_init(root_persistent, child_persistent, overwrite_buffer, device):
     with weights.init_empty_weights():
         model = _RandomInitSharedBuffers(root_persistent, child_persistent, overwrite_buffer)
     _materialize_and_load_weights(
         model,
         None,
-        "cpu",
+        str(device),
         should_skip_hf_weight_load=False,
         is_peft_model=False,
         adapter_path=None,
@@ -116,11 +115,33 @@ def test_random_init_preserves_real_buffers(root_persistent, child_persistent, o
     expected = torch.tensor([31.0, 53.0] if overwrite_buffer else [7.0, 11.0])
     assert model.init_calls == 1
     assert model.root_buffer is model.child.child_buffer
-    assert torch.equal(model.root_buffer, expected)
-    assert torch.equal(model.standalone, torch.tensor([5.0, 8.0]))
-    assert torch.equal(model.weight, torch.tensor([1.0, 2.0]))
+    assert torch.equal(model.root_buffer.cpu(), expected)
+    assert torch.equal(model.standalone.cpu(), torch.tensor([5.0, 8.0]))
+    assert torch.equal(model.weight.cpu(), torch.tensor([1.0, 2.0]))
+    assert model.optional is None
+    assert all(tensor.device == device for tensor in (*model.parameters(), *model.buffers()))
+
+
+@pytest.mark.launch_plan(provider=_CPU_PLAN)
+@pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
+@pytest.mark.parametrize("overwrite_buffer", [False, True])
+def test_random_init_preserves_real_buffers(root_persistent, child_persistent, overwrite_buffer):
+    _check_random_init(root_persistent, child_persistent, overwrite_buffer, torch.device("cpu"))
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec())
+def test_random_init_materializes_buffers_on_accelerator():
+    device_type = get_device_type()
+    if device_type == "cpu":
+        pytest.skip("Requires an accelerator runtime")
+    assert get_torch_device().is_available(), "Selected accelerator runtime is unavailable"
+    device = torch.device(device_type, get_device_id())
+    for root_persistent, child_persistent in _PERSISTENCE:
+        for overwrite_buffer in (False, True):
+            _check_random_init(root_persistent, child_persistent, overwrite_buffer, device)
 
 
+@pytest.mark.launch_plan(provider=_CPU_PLAN)
 @pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
 def test_resume_keeps_checkpoint_authority(monkeypatch, root_persistent, child_persistent):
     with weights.init_empty_weights():
@@ -207,6 +228,7 @@ def cpu_group(tmp_path, monkeypatch):
         dist.destroy_process_group()
 
 
+@pytest.mark.launch_plan(provider=_CPU_PLAN)
 @pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
 @pytest.mark.parametrize("loader", ["plain", "rank0", "ep_sharded"])
 @pytest.mark.parametrize("layout", ["direct", "rename", "finalize"])
@@ -223,6 +245,7 @@ def test_hf_shared_buffer_values_and_identity(tmp_path, cpu_group, root_persiste
     _assert_shared_values(model, root_persistent, child_persistent)
 
 
+@pytest.mark.launch_plan(provider=_CPU_PLAN)
 @pytest.mark.parametrize("root_persistent,child_persistent", _PERSISTENCE)
 def test_hf_shared_dtensor_buffers_rebind_after_dispatch(tmp_path, cpu_group, root_persistent, child_persistent):
     _write_checkpoint(tmp_path, root_persistent, child_persistent)
@@ -234,6 +257,7 @@ def test_hf_shared_dtensor_buffers_rebind_after_dispatch(tmp_path, cpu_group, ro
     _assert_shared_values(model, root_persistent, child_persistent)
 
 
+@pytest.mark.launch_plan(provider=_CPU_PLAN)
 @pytest.mark.parametrize("first_persistent,second_persistent", _PERSISTENCE)
 @pytest.mark.parametrize("reverse", [False, True])
 def test_mapping_load_preserves_cross_child_buffers(tmp_path, first_persistent, second_persistent, reverse):