diff --git a/tests/utils/test_nccl_xfer_layout.py b/tests/utils/test_nccl_xfer_layout.py new file mode 100644 index 000000000..8ec5fdf3f --- /dev/null +++ b/tests/utils/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 diff --git a/tests/utils/test_update_weight_from_nccl_xfer.py b/tests/utils/test_update_weight_from_nccl_xfer.py new file mode 100644 index 000000000..7407c8905 --- /dev/null +++ b/tests/utils/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 diff --git a/vime/backends/megatron_utils/actor.py b/vime/backends/megatron_utils/actor.py index 8a568f55a..0801dd05b 100644 --- a/vime/backends/megatron_utils/actor.py +++ b/vime/backends/megatron_utils/actor.py @@ -36,6 +36,7 @@ from .update_weight.common import named_params_and_buffers from .update_weight.update_weight_from_disk import UpdateWeightFromDisk 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) @@ -141,6 +142,8 @@ def init( self.args.update_weight_mode == "full" ), "--update-weight-mode=delta is not supported with --colocate" update_weight_cls = UpdateWeightFromTensor + elif getattr(self.args, "non_colocate_weight_sync_backend", "broadcast") == "nccl-xfer": + update_weight_cls = UpdateWeightFromNcclXfer elif self.args.update_weight_mode == "delta": # Lazy import: the delta module pulls DeltaEncoding/DeltaParam/DeltaSpec from # vllm, which only exist on newer images. Importing eagerly would break old diff --git a/vime/backends/megatron_utils/update_weight/nccl_xfer_bindings/__init__.py b/vime/backends/megatron_utils/update_weight/nccl_xfer_bindings/__init__.py new file mode 100644 index 000000000..d3ebc5e71 --- /dev/null +++ b/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") diff --git a/vime/backends/megatron_utils/update_weight/nccl_xfer_layout.py b/vime/backends/megatron_utils/update_weight/nccl_xfer_layout.py new file mode 100644 index 000000000..a9e53d4c3 --- /dev/null +++ b/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) diff --git a/vime/backends/megatron_utils/update_weight/update_weight_from_nccl_xfer.py b/vime/backends/megatron_utils/update_weight/update_weight_from_nccl_xfer.py new file mode 100644 index 000000000..3ac45b5d8 --- /dev/null +++ b/vime/backends/megatron_utils/update_weight/update_weight_from_nccl_xfer.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import logging +from argparse import Namespace +from collections.abc import Callable, Mapping, Sequence + +import torch +from ray.actor import ActorHandle + +from .common import named_params_and_buffers +from .nccl_xfer_bindings import get_nccl_xfer_availability +from .nccl_xfer_layout import analyze_nccl_xfer_layout +from .update_weight_from_distributed import UpdateWeightFromDistributed + +logger = logging.getLogger(__name__) + + +class UpdateWeightFromNcclXfer: + """Opt-in NCCL Xfer updater with explicit broadcast fallback. + + The first implementation wires the control plane and layout checks but does + not fake native payload transfer. Until a Python bridge for + ``ncclXferReshardWithWindow`` exists, updates delegate to the existing + non-colocated broadcast backend with a clear log message. + """ + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + self.args = args + self.model = model + self.quantization_config = quantization_config + self._fallback = UpdateWeightFromDistributed( + args, + model, + weights_getter, + model_name=model_name, + quantization_config=quantization_config, + ) + self._fallback_logged = False + self._fallback_reason_cache = self._determine_fallback_reason() + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + self._fallback.connect_rollout_engines( + rollout_engines, + rollout_engine_lock, + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + ) + + def disconnect_rollout_engines(self) -> None: + self._fallback.disconnect_rollout_engines() + + @torch.no_grad() + def update_weights(self) -> None: + reason = self._fallback_reason() + if reason is not None: + if not self._fallback_logged: + logger.warning("NCCL Xfer weight sync unavailable; falling back to broadcast: %s", reason) + self._fallback_logged = True + self._fallback.update_weights() + return + + raise NotImplementedError("native NCCL Xfer weight transfer is not implemented") + + def _fallback_reason(self) -> str | None: + return self._fallback_reason_cache + + def _determine_fallback_reason(self) -> str | None: + availability = get_nccl_xfer_availability() + if not availability.available: + return availability.reason or "native NCCL Xfer bridge is unavailable" + + for name, tensor in named_params_and_buffers(self.args, self.model): + decision = analyze_nccl_xfer_layout( + name, + tuple(tensor.shape), + tensor.dtype, + quantization_config=self.quantization_config, + ) + if not decision.supported: + return decision.reason + return None diff --git a/vime/utils/arguments.py b/vime/utils/arguments.py index 704aeebe5..d645f724c 100644 --- a/vime/utils/arguments.py +++ b/vime/utils/arguments.py @@ -78,6 +78,16 @@ def add_cluster_arguments(parser): "Turning this on will also set --offload to true." ), ) + parser.add_argument( + "--non-colocate-weight-sync-backend", + choices=["broadcast", "nccl-xfer"], + default="broadcast", + help=( + "Weight sync backend for non-colocated rollout engines. " + "'broadcast' uses the existing vLLM NCCL gather-broadcast path; " + "'nccl-xfer' opts into the experimental NCCL Xfer reshard path with broadcast fallback." + ), + ) parser.add_argument( "--offload", action="store_true", @@ -1960,6 +1970,9 @@ def vime_validate_args(args): if args.offload_rollout is None: args.offload_rollout = False + if args.colocate and getattr(args, "non_colocate_weight_sync_backend", "broadcast") != "broadcast": + raise ValueError("--non-colocate-weight-sync-backend=nccl-xfer is only valid when --colocate is disabled.") + if args.use_critic: args.offload_train = True