From 7ecf5359cadf40cea3354ca224751b13c1b2dd92 Mon Sep 17 00:00:00 2001 From: Arist12 Date: Fri, 21 Aug 2026 18:27:11 +0000 Subject: [PATCH 1/2] fix(lora): restore checker-scrambled base weights once --- miles/backends/megatron_utils/lora_utils.py | 3 +- .../update_weight_from_distributed/mixin.py | 14 +- .../update_weight_from_tensor.py | 6 +- .../test_lora_qwen2.5_0.5B_disaggregated.py | 131 ++++++++++++++++++ .../test_lora_weight_sync_validation.py | 128 ++++++++++++++++- 5 files changed, 273 insertions(+), 9 deletions(-) create mode 100644 tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py diff --git a/miles/backends/megatron_utils/lora_utils.py b/miles/backends/megatron_utils/lora_utils.py index 7fe3804e7a7..2c4905241ff 100644 --- a/miles/backends/megatron_utils/lora_utils.py +++ b/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 diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py index bbeea97f815..9cf8f77640e 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py +++ b/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) diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 43373c45aa5..9c847ab425f 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/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: diff --git a/tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py b/tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py new file mode 100644 index 00000000000..eb30cd50b41 --- /dev/null +++ b/tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py @@ -0,0 +1,131 @@ +"""E2E test for LoRA training with Qwen2.5-0.5B on GSM8K, disaggregated. + +The colocate sibling (test_lora_qwen2.5_0.5B.py) covers the CUDA/HIP-IPC weight-sync path. +This one covers the other transport: actor and rollout engines on disjoint GPUs, adapters +pushed over the NCCL/RCCL broadcast group via UpdateWeightFromDistributed. That path has its +own base/adapter sync logic, so a colocate-only suite leaves it unguarded. + +Validates: + - LoRA model setup via Bridge + - base-weight sync under --check-weight-update-equal, which --ci-test turns on + - a rollout after the first optimizer update + - Training completes without errors + +Requires: 4 GPUs, Qwen2.5-0.5B-Instruct model, GSM8K dataset. +Triggered by label: run-ci-lora +""" + +import os + +from tests.ci.ci_register import register_cuda_ci + +import miles.utils.external_utils.command_utils as U + +register_cuda_ci(est_time=400, suite="stage-c-4-gpu-h200", labels=["lora"]) + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 4 +ACTOR_GPUS = 2 +ROLLOUT_GPUS = NUM_GPUS - ACTOR_GPUS + + +def prepare(): + U.exec_command_cpu("mkdir -p /root/models /root/datasets") + U.exec_command_cpu(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command_cpu("hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/datasets/gsm8k") + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " "--megatron-to-hf-mode bridge " + + lora_args = "--lora-rank 32 " "--lora-alpha 32 " "--lora-dropout 0.0 " '--target-modules "all-linear" ' + + rollout_args = ( + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 2 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 8 " + "--rollout-max-response-len 1024 " + "--rollout-temperature 1.0 " + "--global-batch-size 32 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 4096 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--kl-coef 0.00 " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-5 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_args = "--rollout-num-gpus-per-engine 1 " "--sglang-mem-fraction-static 0.4 " + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--calculate-per-token-loss " + "--use-miles-router " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {ACTOR_GPUS} " + f"--rollout-num-gpus {ROLLOUT_GPUS} " + "--update-weight-transfer-mode broadcast " + ) + + train_args = ( + f"{ckpt_args} " + f"{lora_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{sglang_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py b/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py index 1c59f200917..df455b5a7c6 100644 --- a/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py +++ b/tests/fast/backends/megatron_utils/test_lora_weight_sync_validation.py @@ -9,6 +9,7 @@ """ from argparse import Namespace +from contextlib import nullcontext from dataclasses import dataclass from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -16,7 +17,7 @@ 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) From 6d56ec36049a2a6b4e7c234edae36a8441370871 Mon Sep 17 00:00:00 2001 From: Arist12 Date: Sat, 22 Aug 2026 00:18:18 +0000 Subject: [PATCH 2/2] test(lora): drop the disaggregated E2E from this PR Registering a new 400s 4-GPU E2E is a CI-budget decision that belongs with the reviewer, not with a weight-sync fix. The gating logic is already covered by the fast tests in this PR, and the disaggregated transport was validated by hand on 8x MI350X. --- .../test_lora_qwen2.5_0.5B_disaggregated.py | 131 ------------------ 1 file changed, 131 deletions(-) delete mode 100644 tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py diff --git a/tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py b/tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py deleted file mode 100644 index eb30cd50b41..00000000000 --- a/tests/e2e/lora/test_lora_qwen2.5_0.5B_disaggregated.py +++ /dev/null @@ -1,131 +0,0 @@ -"""E2E test for LoRA training with Qwen2.5-0.5B on GSM8K, disaggregated. - -The colocate sibling (test_lora_qwen2.5_0.5B.py) covers the CUDA/HIP-IPC weight-sync path. -This one covers the other transport: actor and rollout engines on disjoint GPUs, adapters -pushed over the NCCL/RCCL broadcast group via UpdateWeightFromDistributed. That path has its -own base/adapter sync logic, so a colocate-only suite leaves it unguarded. - -Validates: - - LoRA model setup via Bridge - - base-weight sync under --check-weight-update-equal, which --ci-test turns on - - a rollout after the first optimizer update - - Training completes without errors - -Requires: 4 GPUs, Qwen2.5-0.5B-Instruct model, GSM8K dataset. -Triggered by label: run-ci-lora -""" - -import os - -from tests.ci.ci_register import register_cuda_ci - -import miles.utils.external_utils.command_utils as U - -register_cuda_ci(est_time=400, suite="stage-c-4-gpu-h200", labels=["lora"]) - -MODEL_NAME = "Qwen2.5-0.5B-Instruct" -MODEL_TYPE = "qwen2.5-0.5B" -NUM_GPUS = 4 -ACTOR_GPUS = 2 -ROLLOUT_GPUS = NUM_GPUS - ACTOR_GPUS - - -def prepare(): - U.exec_command_cpu("mkdir -p /root/models /root/datasets") - U.exec_command_cpu(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") - U.exec_command_cpu("hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/datasets/gsm8k") - - -def execute(): - ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " "--megatron-to-hf-mode bridge " - - lora_args = "--lora-rank 32 " "--lora-alpha 32 " "--lora-dropout 0.0 " '--target-modules "all-linear" ' - - rollout_args = ( - "--prompt-data /root/datasets/gsm8k/train.parquet " - "--input-key messages " - "--label-key label " - "--apply-chat-template " - "--rollout-shuffle " - "--rm-type math " - "--num-rollout 2 " - "--rollout-batch-size 8 " - "--n-samples-per-prompt 8 " - "--rollout-max-response-len 1024 " - "--rollout-temperature 1.0 " - "--global-batch-size 32 " - ) - - perf_args = ( - "--tensor-model-parallel-size 1 " - "--sequence-parallel " - "--pipeline-model-parallel-size 1 " - "--context-parallel-size 1 " - "--expert-model-parallel-size 1 " - "--expert-tensor-parallel-size 1 " - "--use-dynamic-batch-size " - "--max-tokens-per-gpu 4096 " - ) - - grpo_args = ( - "--advantage-estimator grpo " - "--kl-loss-coef 0.00 " - "--kl-loss-type low_var_kl " - "--kl-coef 0.00 " - "--entropy-coef 0.00 " - "--eps-clip 0.2 " - "--eps-clip-high 0.28 " - ) - - optimizer_args = ( - "--optimizer adam " - "--lr 1e-5 " - "--lr-decay-style constant " - "--weight-decay 0.1 " - "--adam-beta1 0.9 " - "--adam-beta2 0.98 " - ) - - sglang_args = "--rollout-num-gpus-per-engine 1 " "--sglang-mem-fraction-static 0.4 " - - ci_args = "--ci-test " - - misc_args = ( - "--attention-dropout 0.0 " - "--hidden-dropout 0.0 " - "--accumulate-allreduce-grads-in-fp32 " - "--attention-softmax-in-fp32 " - "--attention-backend flash " - "--calculate-per-token-loss " - "--use-miles-router " - "--actor-num-nodes 1 " - f"--actor-num-gpus-per-node {ACTOR_GPUS} " - f"--rollout-num-gpus {ROLLOUT_GPUS} " - "--update-weight-transfer-mode broadcast " - ) - - train_args = ( - f"{ckpt_args} " - f"{lora_args} " - f"{rollout_args} " - f"{optimizer_args} " - f"{grpo_args} " - f"{U.get_default_wandb_args(__file__)} " - f"{perf_args} " - f"{sglang_args} " - f"{ci_args} " - f"{misc_args} " - ) - - U.execute_train( - train_args=train_args, - num_gpus_per_node=NUM_GPUS, - megatron_model_type=MODEL_TYPE, - ) - - -if __name__ == "__main__": - prepare() - for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): - os.environ.pop(proxy_var, None) - execute()