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

tests/checkpoints/test_random_init_buffers.py

@@ -0,0 +1,92 @@
+# 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.
+
+"""Config-derived buffers must survive random initialization from meta."""
+
+import pytest
+import torch
+from torch import nn
+
+from scripts.ci.accelerator.launch_plan import LaunchSpec
+from veomni.distributed.torch_parallelize import _materialize_and_load_weights
+from veomni.models.checkpoint.weights import init_empty_weights
+from veomni.utils.device import get_device_id, get_device_type, get_torch_device
+
+
+class BufferModel(nn.Module):
+    def __init__(self, overwrite=False):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(2))
+        self.register_buffer("persistent_value", torch.tensor([2.0, 5.0]))
+        self.register_buffer("runtime_value", torch.tensor([3.0, 7.0]), persistent=False)
+        self.register_buffer("optional_value", None)
+        self.child = nn.Module()
+        self.child.register_buffer("persistent_value", torch.tensor([11.0, 13.0]))
+        self.overwrite = overwrite
+        self.init_calls = 0
+
+    def init_weights(self):
+        self.init_calls += 1
+        with torch.no_grad():
+            self.weight.fill_(19)
+            if self.overwrite:
+                self.child.persistent_value.fill_(23)
+
+
+def _check_materialization(device, *, overwrite, skip_hf):
+    with init_empty_weights():
+        model = BufferModel(overwrite=overwrite)
+    _materialize_and_load_weights(
+        model,
+        None,
+        str(device),
+        should_skip_hf_weight_load=skip_hf,
+        is_peft_model=False,
+        adapter_path=None,
+        broadcast_from_rank0=False,
+    )
+    assert model.weight.device.type == device.type
+    assert model.persistent_value.device.type == device.type
+    assert model.runtime_value.device.type == device.type
+    assert model.child.persistent_value.device.type == device.type
+    assert model.optional_value is None
+    torch.testing.assert_close(model.runtime_value.cpu(), torch.tensor([3.0, 7.0]), rtol=0, atol=0)
+    if skip_hf:
+        # A later DCP load is authoritative for persistent buffers and params.
+        assert model.init_calls == 0
+    else:
+        assert model.init_calls == 1
+        torch.testing.assert_close(model.weight.cpu(), torch.full((2,), 19.0), rtol=0, atol=0)
+        torch.testing.assert_close(model.persistent_value.cpu(), torch.tensor([2.0, 5.0]), rtol=0, atol=0)
+        child = [23.0, 23.0] if overwrite else [11.0, 13.0]
+        torch.testing.assert_close(model.child.persistent_value.cpu(), torch.tensor(child), rtol=0, atol=0)
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
+@pytest.mark.parametrize("overwrite", [False, True])
+@pytest.mark.parametrize("skip_hf", [False, True])
+def test_random_init_and_resume_buffer_materialization_cpu(overwrite, skip_hf):
+    _check_materialization(torch.device("cpu"), overwrite=overwrite, skip_hf=skip_hf)
+
+
+@pytest.mark.launch_plan(provider=LaunchSpec(device="accelerator"))
+def test_random_init_and_resume_buffer_materialization_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 overwrite in (False, True):
+        for skip_hf in (False, True):
+            _check_materialization(device, overwrite=overwrite, skip_hf=skip_hf)

veomni/distributed/torch_parallelize.py

@@ -53,7 +53,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 without discarding config-derived buffers.
 
     ``init_empty_weights()`` patches ``register_parameter`` only, so a
@@ -65,11 +67,14 @@ def _to_empty_preserving_nonpersistent_buffers(model: nn.Module, device: str) ->
     exposes ``original_inv_freq``. Buffers outside that shape -- Gemma3's
     per-layer-type ``{type}_inv_freq`` and its ``embed_scale``, the Omni audio
     tower's sinusoidal ``positional_embedding`` -- have nothing else to restore
-    them. Persistent buffers are left to whichever loader runs next.
+    them. Persistent buffers are preserved for random initialization and left
+    to the loader during checkpoint recovery. Shared registration aliases follow
+    ``to_empty`` semantics and are not reconstructed by this helper.
     """
     buffers = []
     for module in model.modules():
-        for name in module._non_persistent_buffers_set:
+        names = module._buffers.keys() if preserve_persistent_buffers else module._non_persistent_buffers_set
+        for name in names:
             buffer = module._buffers.get(name)
             if buffer is None:
                 continue
@@ -79,7 +84,7 @@ 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
@@ -203,7 +208,9 @@ def _materialize_and_load_weights(
         # 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)
+        _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: