refactor(update-weight): move the transfer protocols into training_utils/weight_update

[yueming-yuan]

Part of #1360. Stacked on #2753 (review only the last four commits). The final PR of the step-5 stack: the whole weight-update engine now lives in training_utils/weight_update/.

What

• Commit 1 — decouple the colocated protocol's backend imports: lora_base_cpu_backup_enabled (pure args inspection) moves from megatron_utils/lora_utils.py to miles/utils/lora.py next to is_lora_enabled; both callers follow. The from ..sglang import ... compat shim is inlined as direct sglang imports (megatron_utils/sglang.py stays behind as the megatron actor's dep manager).
• Commit 2 — the move (git mv, imports/path-constants updated, zero logic diff):

from megatron_utils/update_weight/	to training_utils/weight_update/
update_weight_from_distributed/broadcast.py	broadcast.py
update_weight_from_distributed/delta.py	delta.py
update_weight_from_distributed/p2p.py	p2p.py
update_weight_from_distributed/p2p_transfer_utils.py	p2p_transfer_utils.py
update_weight_from_rdt.py	rdt.py
update_weight_from_tensor.py	colocate.py

Three test files whose subjects all moved follow into tests/fast/backends/training_utils/weight_update/ (the lock test is renamed test_broadcast_lock.py to match broadcast.py).

• Two trailing commits: pre-commit (isort) normalization and one stale docstring path.

What stays in megatron_utils/update_weight/: the megatron iterator implementations (hf_weight_iterator*.py) and common.py — backend code behind the iterator interface.

Mechanical Move

The transform script reproduces this PR's entire tree from its base commit (verify_mechanical_refactor scaffold: applies the transform to a fresh worktree at base, runs pre-commit, diffs against the PR head — verified PASS locally). Script inline below; run from the repo root.

transform_move_weight_update_protocols.py

#!/usr/bin/env python3
"""Reproducible transform: move the weight-transfer protocols into training_utils/weight_update/.

Commit A re-homes the colocated protocol's two backend imports
(lora_base_cpu_backup_enabled -> miles/utils/lora.py; sglang compat imports inlined).
Commit B is the pure directory move plus import-path fixes.

Run from the repo root:  python3 /tmp/transform_move_weight_update_protocols.py
"""
import sys
from pathlib import Path

sys.path.append(".claude/skills/mechanical-refactor-verify")
from mechanical_refactor_verify_utils import (
    exec_command,
    git_add_and_commit,
    verify_mechanical_refactor,
)

BASE_COMMIT = "1be433260"
TARGET_COMMIT = "de8453d82"

MEG_UW = "miles/backends/megatron_utils/update_weight"
NEUTRAL = "miles/backends/training_utils/weight_update"
MEG_TESTS = "tests/fast/backends/megatron_utils"
NEUTRAL_TESTS = "tests/fast/backends/training_utils/weight_update"


def _replace(path: Path, old: str, new: str) -> None:
    content = path.read_text()
    assert old in content, f"{path}: anchor not found: {old!r}"
    path.write_text(content.replace(old, new))


def transform(dir_root: Path) -> None:
    # --- Commit A: decouple the colocated protocol's backend imports ---

    # Move lora_base_cpu_backup_enabled (megatron lora_utils -> miles/utils/lora.py):
    # extract the function block by anchors, delete it from the source.
    lora_utils = dir_root / "miles/backends/megatron_utils/lora_utils.py"
    content = lora_utils.read_text()
    start = content.index("def lora_base_cpu_backup_enabled")
    end = content.index("def sglang_lora_target_all_sentinel")
    func_block = content[start:end].rstrip() + "\n"
    lora_utils.write_text(content[:start] + content[end:])

    neutral_lora = dir_root / "miles/utils/lora.py"
    neutral_lora.write_text(neutral_lora.read_text().rstrip() + "\n\n\n" + func_block)

    # Callers follow the function to its new home.
    _replace(
        dir_root / "miles/backends/sglang_utils/sglang_engine.py",
        "from miles.backends.megatron_utils.lora_utils import (\n"
        "    convert_target_modules_to_hf,\n"
        "    lora_base_cpu_backup_enabled,\n"
        "    sglang_lora_target_all_sentinel,\n"
        ")",
        "from miles.backends.megatron_utils.lora_utils import (\n"
        "    convert_target_modules_to_hf,\n"
        "    sglang_lora_target_all_sentinel,\n"
        ")",
    )
    _replace(
        dir_root / "miles/backends/sglang_utils/sglang_engine.py",
        "from miles.utils.lora import LORA_ADAPTER_NAME, lora_rollout_enabled",
        "from miles.utils.lora import LORA_ADAPTER_NAME, lora_base_cpu_backup_enabled, lora_rollout_enabled",
    )

    colocate_src = dir_root / f"{MEG_UW}/update_weight_from_tensor.py"
    _replace(
        colocate_src,
        "from miles.backends.megatron_utils.lora_utils import lora_base_cpu_backup_enabled",
        "from miles.utils.lora import lora_base_cpu_backup_enabled",
    )
    # Inline the sglang compat imports (same shim as megatron_utils/sglang.py, which
    # stays behind as the megatron actor's dep manager).
    _replace(
        colocate_src,
        "from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer",
        "from sglang.srt.utils import MultiprocessingSerializer\n"
        "\n"
        "try:\n"
        "    from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket  # type: ignore[import]\n"
        "except ImportError:\n"
        "    from sglang.srt.model_executor.model_runner import FlattenedTensorBucket  # type: ignore[import]",
    )

    git_add_and_commit(
        "refactor(update-weight): re-home the colocated protocol's backend imports",
        cwd=str(dir_root),
    )

    # --- Commit B: mechanical move into training_utils/weight_update/ ---

    moves = [
        (f"{MEG_UW}/update_weight_from_distributed/broadcast.py", f"{NEUTRAL}/broadcast.py"),
        (f"{MEG_UW}/update_weight_from_distributed/delta.py", f"{NEUTRAL}/delta.py"),
        (f"{MEG_UW}/update_weight_from_distributed/p2p.py", f"{NEUTRAL}/p2p.py"),
        (f"{MEG_UW}/update_weight_from_distributed/p2p_transfer_utils.py", f"{NEUTRAL}/p2p_transfer_utils.py"),
        (f"{MEG_UW}/update_weight_from_rdt.py", f"{NEUTRAL}/rdt.py"),
        (f"{MEG_UW}/update_weight_from_tensor.py", f"{NEUTRAL}/colocate.py"),
        (f"{MEG_TESTS}/test_update_weight_from_distributed_lock.py", f"{NEUTRAL_TESTS}/test_broadcast_lock.py"),
        (f"{MEG_TESTS}/test_lora_weight_sync_validation.py", f"{NEUTRAL_TESTS}/test_lora_weight_sync_validation.py"),
        (f"{MEG_TESTS}/test_lora_update_weight.py", f"{NEUTRAL_TESTS}/test_lora_update_weight.py"),
    ]
    for src, dst in moves:
        exec_command(f"git mv {src} {dst}", cwd=str(dir_root))
    exec_command(f"git rm -q {MEG_UW}/update_weight_from_distributed/__init__.py", cwd=str(dir_root))

    # Intra-package imports: the moved files are now siblings.
    _replace(
        dir_root / f"{NEUTRAL}/rdt.py",
        "from .update_weight_from_distributed.p2p_transfer_utils import",
        "from .p2p_transfer_utils import",
    )
    _replace(
        dir_root / f"{NEUTRAL}/colocate.py",
        "from .update_weight_from_distributed.broadcast import",
        "from .broadcast import",
    )

    # Factory import paths.
    protocol = dir_root / f"{NEUTRAL}/protocol.py"
    for old, new in [
        (
            "from miles.backends.megatron_utils.update_weight.update_weight_from_tensor import UpdateWeightFromTensor",
            "from miles.backends.training_utils.weight_update.colocate import UpdateWeightFromTensor",
        ),
        (
            "from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast import (\n"
            "            UpdateWeightFromDistributed,\n"
            "        )",
            "from miles.backends.training_utils.weight_update.broadcast import UpdateWeightFromDistributed",
        ),
        (
            "from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.delta import (\n"
            "            UpdateWeightFromDiskDelta,\n"
            "        )",
            "from miles.backends.training_utils.weight_update.delta import UpdateWeightFromDiskDelta",
        ),
        (
            "from miles.backends.megatron_utils.update_weight.update_weight_from_rdt import UpdateWeightFromRDT",
            "from miles.backends.training_utils.weight_update.rdt import UpdateWeightFromRDT",
        ),
        (
            "from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.p2p import UpdateWeightP2P",
            "from miles.backends.training_utils.weight_update.p2p import UpdateWeightP2P",
        ),
    ]:
        _replace(protocol, old, new)

    # Test module-path constants and imports.
    lock_test = dir_root / f"{NEUTRAL_TESTS}/test_broadcast_lock.py"
    _replace(lock_test, "update_weight_from_distributed/broadcast.py", "training_utils/weight_update/broadcast.py")
    _replace(
        lock_test,
        "from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast import (",
        "from miles.backends.training_utils.weight_update.broadcast import (",
    )
    _replace(
        lock_test,
        '_MODULE = "miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast"',
        '_MODULE = "miles.backends.training_utils.weight_update.broadcast"',
    )

    validation_test = dir_root / f"{NEUTRAL_TESTS}/test_lora_weight_sync_validation.py"
    _replace(
        validation_test,
        "from miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast import (\n"
        "    UpdateWeightFromDistributed,\n"
        ")",
        "from miles.backends.training_utils.weight_update.broadcast import UpdateWeightFromDistributed",
    )
    _replace(
        validation_test,
        "from miles.backends.megatron_utils.update_weight.update_weight_from_tensor import UpdateWeightFromTensor",
        "from miles.backends.training_utils.weight_update.colocate import UpdateWeightFromTensor",
    )
    _replace(
        validation_test,
        '_UW_MODULE = "miles.backends.megatron_utils.update_weight.update_weight_from_tensor"',
        '_UW_MODULE = "miles.backends.training_utils.weight_update.colocate"',
    )
    _replace(
        validation_test,
        '_BROADCAST_MODULE = "miles.backends.megatron_utils.update_weight.update_weight_from_distributed.broadcast"',
        '_BROADCAST_MODULE = "miles.backends.training_utils.weight_update.broadcast"',
    )
    _replace(
        validation_test,
        "``miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py``",
        "``miles/backends/training_utils/weight_update/colocate.py``",
    )

    lifecycle_test = dir_root / f"{MEG_TESTS}/test_shared_ppo_lifecycle.py"
    _replace(
        lifecycle_test,
        'p2p_module_name = "miles.backends.megatron_utils.update_weight.update_weight_from_distributed.p2p"',
        'p2p_module_name = "miles.backends.training_utils.weight_update.p2p"',
    )
    _replace(
        lifecycle_test,
        'importlib.import_module("miles.backends.megatron_utils.update_weight.update_weight_from_distributed")',
        'importlib.import_module("miles.backends.training_utils.weight_update")',
    )

    git_add_and_commit(
        "refactor(update-weight): move the transfer protocols into training_utils/weight_update",
        cwd=str(dir_root),
    )


if __name__ == "__main__":
    verify_mechanical_refactor(
        base_commit=BASE_COMMIT,
        target_commit=TARGET_COMMIT,
        transform=transform,
    )

[claude]

Claude Code Review

This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.

_{Tip: disable this comment in your organization's Code Review settings.}

[guapisolo]

codex cmt

[P3] Update documentation paths moved by this PR

docs/advanced/fault-tolerance.md still points to miles/backends/megatron_utils/update_weight/update_weight_from_distributed/p2p.py, and docs/advanced/lora.md still points to miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py and miles/backends/megatron_utils/update_weight/update_weight_from_distributed/. Those paths are deleted by this PR, so readers following the implementation guidance land on files that no longer exist. Please update them to miles/backends/training_utils/weight_update/protocols/p2p.py, miles/backends/training_utils/weight_update/protocols/cuda_ipc.py, and miles/backends/training_utils/weight_update/protocols/.

[guapisolo]

This comment maybe deprecated after our mgt upgrade. non blocking. I will help fix this later.