[1/N] Add optional NCCL m2n weight sync backend

tests/unit/backends/megatron_utils/update_weight/test_nccl_xfer_layout.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vime.backends.megatron_utils.update_weight.nccl_xfer_layout import analyze_nccl_xfer_layout
+
+
+@pytest.mark.unit
+def test_column_parallel_weight_shards_dim0():
+    decision = analyze_nccl_xfer_layout(
+        "model.layers.0.self_attn.q_proj.weight",
+        (4096, 4096),
+        torch.bfloat16,
+    )
+
+    assert decision.supported
+    assert decision.shard_tensor_dim == 0
+    assert not decision.replicated
+
+
+@pytest.mark.unit
+def test_row_parallel_weight_shards_dim1():
+    decision = analyze_nccl_xfer_layout(
+        "model.layers.0.self_attn.o_proj.weight",
+        (4096, 4096),
+        torch.float16,
+    )
+
+    assert decision.supported
+    assert decision.shard_tensor_dim == 1
+
+
+@pytest.mark.unit
+def test_replicated_1d_weight():
+    decision = analyze_nccl_xfer_layout(
+        "model.layers.0.input_layernorm.weight",
+        (4096,),
+        torch.float32,
+    )
+
+    assert decision.supported
+    assert decision.replicated
+    assert decision.shard_tensor_dim is None
+
+
+@pytest.mark.unit
+def test_grouped_moe_expert_weight_shards_expert_dim():
+    decision = analyze_nccl_xfer_layout(
+        "model.layers.0.mlp.experts.gate_proj.weight",
+        (16, 4096, 4096),
+        torch.bfloat16,
+    )
+
+    assert decision.supported
+    assert decision.shard_tensor_dim == 0
+
+
+@pytest.mark.unit
+def test_ungrouped_moe_expert_weight_is_unsupported():
+    decision = analyze_nccl_xfer_layout(
+        "model.layers.0.mlp.experts.0.gate_proj.weight",
+        (4096, 4096),
+        torch.bfloat16,
+    )
+
+    assert not decision.supported
+    assert "rank-3" in decision.reason
+
+
+@pytest.mark.unit
+def test_tensor_rank_greater_than_three_is_unsupported():
+    decision = analyze_nccl_xfer_layout(
+        "model.layers.0.weird_packed.weight",
+        (2, 3, 4, 5),
+        torch.float16,
+    )
+
+    assert not decision.supported
+    assert "exceeds NCCL Xfer limit" in decision.reason
+
+
+@pytest.mark.unit
+def test_compressed_tensors_quantization_is_unsupported():
+    decision = analyze_nccl_xfer_layout(
+        "model.layers.0.self_attn.q_proj.weight",
+        (4096, 4096),
+        torch.float16,
+        quantization_config={"quant_method": "compressed-tensors"},
+    )
+
+    assert not decision.supported
+    assert "compressed-tensors" in decision.reason

tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_nccl_xfer.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+from argparse import Namespace
+from dataclasses import dataclass
+
+import pytest
+
+
+@dataclass
+class DummyAvailability:
+    available: bool
+    reason: str | None = None
+
+
+class DummyFallback:
+    instances = []
+
+    def __init__(self, *args, **kwargs):
+        self.update_calls = 0
+        self.connect_calls = []
+        self.disconnect_calls = 0
+        DummyFallback.instances.append(self)
+
+    def connect_rollout_engines(self, *args, **kwargs):
+        self.connect_calls.append((args, kwargs))
+
+    def disconnect_rollout_engines(self):
+        self.disconnect_calls += 1
+
+    def update_weights(self):
+        self.update_calls += 1
+
+
+@pytest.fixture
+def upw(monkeypatch):
+    from vime.backends.megatron_utils.update_weight import update_weight_from_nccl_xfer as mod
+
+    DummyFallback.instances.clear()
+    monkeypatch.setattr(mod, "UpdateWeightFromDistributed", DummyFallback)
+    return mod
+
+
+@pytest.mark.unit
+def test_falls_back_when_native_bridge_unavailable(upw, monkeypatch, caplog):
+    monkeypatch.setattr(
+        upw,
+        "get_nccl_xfer_availability",
+        lambda: DummyAvailability(False, "missing pybind bridge"),
+    )
+    updater = upw.UpdateWeightFromNcclXfer(
+        Namespace(),
+        model=[],
+        weights_getter=lambda: {},
+        model_name="qwen",
+        quantization_config=None,
+    )
+
+    updater.update_weights()
+
+    assert DummyFallback.instances[-1].update_calls == 1
+    assert "missing pybind bridge" in caplog.text
+
+
+@pytest.mark.unit
+def test_connect_and_disconnect_delegate_to_broadcast_fallback(upw):
+    updater = upw.UpdateWeightFromNcclXfer(
+        Namespace(),
+        model=[],
+        weights_getter=lambda: {},
+        model_name="qwen",
+        quantization_config=None,
+    )
+    fallback = DummyFallback.instances[-1]
+
+    updater.connect_rollout_engines(["engine"], "lock", engine_gpu_counts=[1], engine_gpu_offsets=[0])
+    updater.disconnect_rollout_engines()
+
+    assert fallback.connect_calls == [((["engine"], "lock"), {"engine_gpu_counts": [1], "engine_gpu_offsets": [0]})]
+    assert fallback.disconnect_calls == 1

vime/backends/megatron_utils/actor.py

@@ -34,6 +34,7 @@
 from .model import forward_only, initialize_model_and_optimizer, save, train
 from .update_weight.common import named_params_and_buffers
 from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed
+from .update_weight.update_weight_from_nccl_xfer import UpdateWeightFromNcclXfer
 from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor
 
 logging.getLogger("megatron").setLevel(logging.WARNING)
@@ -136,6 +137,8 @@ def init(
 
         if self.args.colocate:
             update_weight_cls = UpdateWeightFromTensor
+        elif getattr(self.args, "non_colocate_weight_sync_backend", "broadcast") == "nccl-xfer":
+            update_weight_cls = UpdateWeightFromNcclXfer
         else:
             update_weight_cls = UpdateWeightFromDistributed
         self.weight_updater = update_weight_cls(

vime/backends/megatron_utils/update_weight/nccl_xfer_bindings/__init__.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class NcclXferAvailability:
+    available: bool
+    reason: str | None = None
+
+
+def get_nccl_xfer_availability() -> NcclXferAvailability:
+    """Return whether a native NCCL Xfer bridge is available to Python.
+
+    The vendored ``nccl_xfer`` tree currently exposes a C/CUDA API only. This
+    shim prevents the opt-in backend from silently pretending to transfer
+    payloads until a pybind/torch-extension bridge is added.
+    """
+
+    return NcclXferAvailability(
+        available=False,
+        reason="native NCCL Xfer Python bridge is not implemented",
+    )
+
+
+def reshard_with_window(*args, **kwargs):
+    del args, kwargs
+    raise NotImplementedError("native NCCL Xfer Python bridge is not implemented")

vime/backends/megatron_utils/update_weight/nccl_xfer_layout.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+NCCL_XFER_MAX_TENSOR_DIMS = 3
+
+_COLUMN_PARALLEL_SUFFIXES = (
+    "q_proj.weight",
+    "k_proj.weight",
+    "v_proj.weight",
+    "gate_proj.weight",
+    "up_proj.weight",
+    "q_a_proj.weight",
+    "q_b_proj.weight",
+    "kv_a_proj_with_mqa.weight",
+    "kv_b_proj.weight",
+)
+_ROW_PARALLEL_SUFFIXES = ("o_proj.weight", "down_proj.weight")
+_VOCAB_PARALLEL_NAMES = ("embed_tokens.weight", "lm_head.weight")
+
+_SUPPORTED_DTYPES = {
+    dtype
+    for dtype in (
+        torch.int8,
+        torch.uint8,
+        getattr(torch, "float8_e4m3fn", None),
+        getattr(torch, "float8_e5m2", None),
+        torch.float16,
+        torch.bfloat16,
+        torch.int32,
+        torch.uint32,
+        torch.float32,
+        torch.int64,
+        torch.uint64,
+        torch.float64,
+    )
+    if dtype is not None
+}
+
+
+@dataclass(frozen=True)
+class NcclXferLayoutDecision:
+    supported: bool
+    reason: str | None = None
+    shard_tensor_dim: int | None = None
+    replicated: bool = False
+
+
+def analyze_nccl_xfer_layout(
+    name: str,
+    shape: Sequence[int],
+    dtype: torch.dtype,
+    *,
+    quantization_config: Mapping[str, Any] | object | None = None,
+) -> NcclXferLayoutDecision:
+    """Classify whether a model weight can use the MVP NCCL Xfer layout mapping.
+
+    This mirrors the RFC's first-pass placement rules. It is intentionally
+    conservative: unsupported cases should fall back to the existing broadcast
+    path instead of entering a partially implemented native transfer.
+    """
+
+    if _get_quant_method(quantization_config) == "compressed-tensors":
+        return NcclXferLayoutDecision(False, "compressed-tensors quantized weights require broadcast fallback")
+
+    ndim = len(shape)
+    if ndim == 0:
+        return NcclXferLayoutDecision(False, f"{name}: scalar tensors are not supported by NCCL Xfer reshard")
+    if ndim > NCCL_XFER_MAX_TENSOR_DIMS:
+        return NcclXferLayoutDecision(False, f"{name}: tensor rank {ndim} exceeds NCCL Xfer limit of 3")
+    if dtype not in _SUPPORTED_DTYPES:
+        return NcclXferLayoutDecision(False, f"{name}: dtype {dtype} is not supported by NCCL Xfer")
+
+    if ".experts." in name:
+        if ndim != 3:
+            return NcclXferLayoutDecision(
+                False,
+                f"{name}: MoE expert tensors must be grouped as rank-3 [num_experts, out, in]",
+            )
+        return NcclXferLayoutDecision(True, shard_tensor_dim=0)
+
+    if ndim < 2:
+        return NcclXferLayoutDecision(True, shard_tensor_dim=None, replicated=True)
+
+    if _endswith_any(name, _COLUMN_PARALLEL_SUFFIXES) or _contains_any(name, _VOCAB_PARALLEL_NAMES):
+        return NcclXferLayoutDecision(True, shard_tensor_dim=0)
+
+    if _endswith_any(name, _ROW_PARALLEL_SUFFIXES):
+        return NcclXferLayoutDecision(True, shard_tensor_dim=1)
+
+    return NcclXferLayoutDecision(True, shard_tensor_dim=None, replicated=True)
+
+
+def _endswith_any(name: str, suffixes: Sequence[str]) -> bool:
+    return any(name.endswith(suffix) for suffix in suffixes)
+
+
+def _contains_any(name: str, needles: Sequence[str]) -> bool:
+    return any(needle in name for needle in needles)
+
+
+def _get_quant_method(quantization_config: Mapping[str, Any] | object | None) -> str | None:
+    if quantization_config is None:
+        return None
+    if isinstance(quantization_config, Mapping):
+        return quantization_config.get("quant_method")
+    return getattr(quantization_config, "quant_method", None)