fix(lora): restore checker-scrambled base weights once

miles/backends/megatron_utils/lora_utils.py

@@ -187,7 +187,8 @@ def is_lora_weight_name(name: str) -> bool:
 
 def _is_adapter_param_name(name: str) -> bool:
     """Check if a parameter name belongs to a LoRA adapter (Megatron internal naming)."""
-    return "lora_" in name or (".adapter." in name and ("linear_in" in name or "linear_out" in name))
+    adapter_container = ".adapter." in name or ".adapters." in name
+    return "lora_" in name or (adapter_container and ("linear_in" in name or "linear_out" in name))
 
 
 _param_grad_buffer_patched = False

miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py

@@ -352,8 +352,9 @@ def update_weights(self) -> None:
 
         Full: pause → base non-expert (TP) → base expert (EP) → resume.
         LoRA: pause → LoRA adapter (every iteration) → resume. The frozen base is
-        never pushed; the remote rollout engines already load it from
-        ``hf_checkpoint`` at init.
+        not pushed; the remote rollout engines already load it from
+        ``hf_checkpoint`` at init. The exception is the first update with
+        ``--check-weight-update-equal``, which restores the weights scrambled at startup.
         """
         self.weight_version += 1
 
@@ -366,9 +367,12 @@ def update_weights(self) -> None:
             is_lora = getattr(self, "is_lora", False)
             is_multi_lora = is_lora and is_multi_lora_enabled(self.args)
 
-            # LoRA: base weights are frozen and already loaded by the rollout engines
-            # from ``hf_checkpoint``, so only full-param runs sync the base.
-            if not is_lora:
+            # LoRA normally keeps the frozen base on the remote engines. The checker
+            # overwrites that copy only at startup, so restore it on the first update.
+            skip_base_sync = is_lora and not (
+                getattr(self.args, "check_weight_update_equal", False) and self.weight_version == 1
+            )
+            if not skip_base_sync:
                 pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_source else None
 
                 self._gather_and_update_non_expert_weights(self._update_weight_implementation, pbar)

miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py

@@ -242,14 +242,16 @@ def update_weights(self) -> None:
 
         rank = dist.get_rank()
 
-        # TODO: implement lora weight checker
         colocate_base_persistent = getattr(self.args, "colocate", False) and not getattr(
             self.args, "offload_rollout", True
         )
+        checker_needs_base_restore = (
+            self.is_lora and getattr(self.args, "check_weight_update_equal", False) and self.weight_version == 1
+        )
         skip_base_sync = (
             self.is_lora
             and (self.use_distribute or lora_base_cpu_backup_enabled(self.args) or colocate_base_persistent)
-            and not getattr(self.args, "check_weight_update_equal", False)
+            and not checker_needs_base_restore
         )
 
         if rank == 0:

tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py

@@ -9,14 +9,15 @@
 """
 
 from argparse import Namespace
+from contextlib import nullcontext
 from dataclasses import dataclass
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
-from miles.backends.megatron_utils.lora_utils import is_lora_weight_name
+from miles.backends.megatron_utils.lora_utils import _is_adapter_param_name, is_lora_weight_name
 from miles.backends.megatron_utils.update_weight.common import _check_weight_sync_results
 from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast import (
     UpdateWeightFromDistributed,
@@ -513,3 +514,128 @@ def test_raises_when_engine_reports_failure(self):
         fake_self = self._make_self(engines=engines)
         with pytest.raises(RuntimeError, match="LoRA weight sync failed"):
             self._run(fake_self, SAMPLE_LORA_WEIGHTS)
+
+
+class TestDistBaseSyncGating:
+    """The weight checker forces a frozen-base sync that LoRA normally skips."""
+
+    @staticmethod
+    def _make_self(*, is_lora, check_equal):
+        return SimpleNamespace(
+            weight_version=0,
+            is_lora=is_lora,
+            args=Namespace(check_weight_update_equal=check_equal),
+            _is_source=True,
+            _group_name="g",
+            _update_weight_implementation=MagicMock(),
+            _pause_and_prepare_engines=MagicMock(),
+            _finalize_and_resume_engines=MagicMock(),
+            _gather_and_update_non_expert_weights=MagicMock(),
+            _gather_and_update_expert_weights=MagicMock(),
+            _update_lora_weights=MagicMock(),
+            _update_multi_lora_weights=MagicMock(),
+        )
+
+    @pytest.mark.parametrize(
+        ("is_lora", "check_equal", "expected_base_sync_counts"),
+        [
+            (True, False, (0, 0)),
+            (True, True, (1, 1)),
+            (False, False, (1, 2)),
+            (False, True, (1, 2)),
+        ],
+    )
+    def test_base_sync_gating(self, is_lora, check_equal, expected_base_sync_counts):
+        fake_self = self._make_self(is_lora=is_lora, check_equal=check_equal)
+        non_expert_sync_counts = []
+        expert_sync_counts = []
+        with (
+            patch(f"{_MIXIN_MODULE}.dist"),
+            patch(f"{_MIXIN_MODULE}.get_gloo_group", return_value=MagicMock()),
+            patch(f"{_MIXIN_MODULE}.timer", lambda *a, **k: nullcontext()),
+            patch(f"{_MIXIN_MODULE}.tqdm", MagicMock()),
+            patch("miles.utils.multi_lora.is_multi_lora_enabled", return_value=False),
+        ):
+            for _ in range(2):
+                DistBucketedWeightUpdateMixin.update_weights(fake_self)
+                non_expert_sync_counts.append(fake_self._gather_and_update_non_expert_weights.call_count)
+                expert_sync_counts.append(fake_self._gather_and_update_expert_weights.call_count)
+
+        assert tuple(non_expert_sync_counts) == expected_base_sync_counts
+        assert tuple(expert_sync_counts) == expected_base_sync_counts
+        assert fake_self._update_lora_weights.call_count == (2 if is_lora else 0)
+
+
+@pytest.mark.parametrize(
+    "name",
+    [
+        "module.decoder.layers.0.mlp.linear_fc1.adapter.linear_in.weight",
+        "module.decoder.layers.0.mlp.linear_fc1.adapters.0.linear_out.weight",
+        "module.decoder.layers.0.mlp.linear_fc1.lora_A.weight",
+    ],
+)
+def test_base_sync_filter_recognizes_single_and_multi_lora_params(name):
+    assert _is_adapter_param_name(name)
+
+
+class TestTensorBaseSyncGating:
+    """The tensor updater also restores checker-scrambled LoRA base weights only once."""
+
+    @staticmethod
+    def _make_self(*, is_lora, check_equal):
+        iterator = MagicMock()
+        iterator.get_hf_weight_chunks.side_effect = lambda *_args, weight_type: iter(
+            [SAMPLE_LORA_WEIGHTS if weight_type == "lora" else SAMPLE_BASE_ONLY_WEIGHTS]
+        )
+        return SimpleNamespace(
+            weight_version=0,
+            is_lora=is_lora,
+            args=Namespace(
+                check_weight_update_equal=check_equal,
+                colocate=True,
+                offload_rollout=False,
+                pause_generation_mode="retract",
+            ),
+            use_distribute=False,
+            rollout_engines=[MagicMock()],
+            weights_getter=MagicMock(return_value={}),
+            _hf_weight_iterator=iterator,
+            _send_base_params=MagicMock(return_value=([], None)),
+            _send_lora_params=MagicMock(return_value=([], None)),
+            _mm_tower_named_tensors=MagicMock(return_value=None),
+            _lora_base_synced=False,
+        )
+
+    @pytest.mark.parametrize(
+        ("is_lora", "check_equal", "expected_base_sync_counts"),
+        [
+            (True, False, (0, 0)),
+            (True, True, (1, 1)),
+            (False, False, (1, 2)),
+            (False, True, (1, 2)),
+        ],
+    )
+    def test_base_sync_gating(self, is_lora, check_equal, expected_base_sync_counts):
+        fake_self = self._make_self(is_lora=is_lora, check_equal=check_equal)
+        base_sync_counts = []
+        with (
+            patch(f"{_UW_MODULE}.dist") as dist_mock,
+            patch(f"{_UW_MODULE}.ray") as ray_mock,
+            patch(f"{_UW_MODULE}.get_gloo_group", return_value=MagicMock()),
+            patch(f"{_UW_MODULE}.lora_base_cpu_backup_enabled", return_value=False),
+            patch(f"{_UW_MODULE}.begin_weight_update") as begin_mock,
+            patch(f"{_UW_MODULE}.end_weight_update") as end_mock,
+            patch(f"{_UW_MODULE}._pp_assemble_full_adapter", side_effect=lambda tensors: tensors),
+            patch(f"{_UW_MODULE}.torch.cuda.ipc_collect"),
+            patch(f"{_UW_MODULE}.torch.cuda.empty_cache"),
+        ):
+            dist_mock.get_rank.return_value = 0
+            ray_mock.get.side_effect = lambda refs: refs
+            for _ in range(2):
+                UpdateWeightFromTensor.update_weights(fake_self)
+                base_sync_counts.append(fake_self._send_base_params.call_count)
+
+        assert tuple(base_sync_counts) == expected_base_sync_counts
+        assert begin_mock.call_count == expected_base_sync_counts[-1]
+        assert end_mock.call_count == expected_base_sync_counts[-1]
+        assert fake_self._send_lora_params.call_count == (2 if is_lora else 0)