chore: import upstream snapshot with attribution

swift/ray/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.

swift/ray/megatron/__init__.py

@@ -0,0 +1,47 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .driver_utils import (RayConfig, build_dataset_from_dict, compute_iter_params, estimate_dp_size,
+                               extract_iteration, merge_group_dict, parse_args_from_dict, parse_ray_yaml)
+    from .grpo_trainer import GRPOTrainer
+    from .loss import GRPOLoss, Loss
+    from .megatron_worker import MegatronWorker
+    from .pipeline import MegatronRayPipeline, register_ray_trainer
+    from .resource_pool import ResourcePool, ResourcePoolManager
+    from .rollout import RolloutAdapter, RolloutMode, RolloutReplica, VllmEngineConfig, VllmServer
+    from .worker_group import CollectMode, DispatchMode, WorkerGroup, dispatch_collect
+
+
+def __getattr__(name):
+    _imports = {
+        'RayConfig': '.driver_utils',
+        'build_dataset_from_dict': '.driver_utils',
+        'compute_iter_params': '.driver_utils',
+        'estimate_dp_size': '.driver_utils',
+        'extract_iteration': '.driver_utils',
+        'merge_group_dict': '.driver_utils',
+        'parse_args_from_dict': '.driver_utils',
+        'parse_ray_yaml': '.driver_utils',
+        'GRPOTrainer': '.grpo_trainer',
+        'MegatronRayPipeline': '.pipeline',
+        'register_ray_trainer': '.pipeline',
+        'Loss': '.loss',
+        'GRPOLoss': '.loss',
+        'ResourcePool': '.resource_pool',
+        'ResourcePoolManager': '.resource_pool',
+        'RolloutMode': '.rollout',
+        'RolloutReplica': '.rollout',
+        'VllmEngineConfig': '.rollout',
+        'VllmServer': '.rollout',
+        'RolloutAdapter': '.rollout',
+        'MegatronWorker': '.megatron_worker',
+        'CollectMode': '.worker_group',
+        'DispatchMode': '.worker_group',
+        'WorkerGroup': '.worker_group',
+        'dispatch_collect': '.worker_group',
+    }
+    if name in _imports:
+        import importlib
+        return getattr(importlib.import_module(_imports[name], __name__), name)
+    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')

swift/ray/megatron/base_trainer.py

@@ -0,0 +1,296 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Base class for Ray-based Megatron trainers (driver-side)."""
+from __future__ import annotations
+
+import os
+import ray
+import torch
+from contextlib import contextmanager
+from typing import Any, Dict, List, Tuple
+
+from swift.rl_core.data import GRPOBatch, OnPolicySample
+from swift.rl_core.resample import resample_encode_failed_inputs
+from swift.rlhf_trainers.utils import create_cyclic_iterator
+from swift.template.base import Template
+from swift.utils import JsonlWriter, get_logger
+from .driver_utils import compute_iter_params
+from .worker_group import DPDispatchedDict
+
+logger = get_logger()
+
+
+class BaseRayTrainer:
+    """Shared driver-side logic for Ray Megatron trainers.
+
+    Subclasses implement ``_prepare_state`` and ``_train_loop``.
+    """
+
+    def __init__(
+        self,
+        worker_groups: Dict[str, Any],
+        rollout_replicas: List[Any],
+        weight_sync_mode: str = 'nccl',
+        sleep_level: int = 1,
+        teacher_replicas: List[Any] = None,
+    ):
+        self.worker_groups = worker_groups
+        self.rollout_replicas = rollout_replicas
+        self._weight_sync_mode = weight_sync_mode
+        self._sleep_level = sleep_level
+        self.teacher_replicas = teacher_replicas or []
+
+    def set_data_info(self, data_info: Dict[str, Any]) -> None:
+        self._data_info = data_info
+
+    @property
+    def train_group(self):
+        return self.worker_groups['train']
+
+    @property
+    def is_colocated_rollout(self) -> bool:
+        from .rollout.replica import RolloutMode
+        if not self.rollout_replicas:
+            return False
+        return self.rollout_replicas[0].mode == RolloutMode.HYBRID
+
+    @property
+    def ckpt_manager(self):
+        if not hasattr(self, '_ckpt_manager'):
+            from .checkpoint_engine import CheckpointEngineManager
+            tg = self.train_group
+            self._ckpt_manager = CheckpointEngineManager(
+                train_actors=tg.workers,
+                rollout_replicas=self.rollout_replicas,
+                weight_sync_mode=self._weight_sync_mode,
+                is_colocated=self.is_colocated_rollout,
+                sleep_level=self._sleep_level,
+                train_group=tg,
+            )
+        return self._ckpt_manager
+
+    def _distribute_to_replicas(self, batch, params):
+        n = len(self.rollout_replicas)
+        chunk_size = (len(batch) + n - 1) // n
+        refs = []
+        for i, replica in enumerate(self.rollout_replicas):
+            shard = batch[i * chunk_size:(i + 1) * chunk_size]
+            if not shard:
+                continue
+            refs.append(replica.generate(shard, params))
+        parts = ray.get(refs)
+        result = []
+        for p in parts:
+            result.extend(p)
+        return result
+
+    @contextmanager
+    def _generation_context(self, tg, ckpt):
+        offload_model = getattr(self.args, 'offload_model', False)
+        offload_optimizer = getattr(self.args, 'offload_optimizer', False)
+        enable_offload = offload_model or offload_optimizer or self.is_colocated_rollout
+
+        if enable_offload:
+            tg.offload_to_cpu()
+        if self.is_colocated_rollout:
+            ckpt.wake_up_rollout(tags=['kv_cache'])
+
+        try:
+            yield
+        finally:
+            tg.finalize_generation()
+            if self.is_colocated_rollout:
+                ckpt.sleep_rollout()
+            if enable_offload:
+                tg.reload_to_gpu()
+
+    def _build_dataloader(self):
+        info = self._data_info
+        dataset = info['train_dataset']
+        num_gen = int(info.get('num_generations', 1) or 1)
+        spg = self._steps_per_generation
+        prompts_per_generation = max(info['global_batch_size'] * spg // max(num_gen, 1), 1)
+        self._dataloader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=prompts_per_generation,
+            shuffle=True,
+            collate_fn=info['data_collator'],
+            drop_last=True,
+        )
+        self._data_iter = create_cyclic_iterator(self._dataloader)
+        logger.info('%s driver dataloader: dataset=%d, prompts_per_generation=%d, num_gen=%d, spg=%d',
+                    type(self).__name__, len(dataset), prompts_per_generation, num_gen, spg)
+
+    def _build_resample_iterator(self) -> None:
+        """Independent cyclic prompt iterator (different shuffle order) used to replace
+        encode-failed prompts (truncation_strategy='delete') and to refill DAPO
+        dynamic_sample std=0 groups (driver-side)."""
+        info = self._data_info
+        dataset = info['train_dataset']
+        num_gen = int(info.get('num_generations', 1) or 1)
+        spg = self._steps_per_generation
+        prompts_per_generation = max(info['global_batch_size'] * spg // max(num_gen, 1), 1)
+        loader = torch.utils.data.DataLoader(
+            dataset,
+            batch_size=prompts_per_generation,
+            shuffle=True,
+            collate_fn=info['data_collator'],
+            drop_last=True,
+        )
+        self._resample_iter = create_cyclic_iterator(loader)
+
+    def _resample_failed_prompts(self, prompts: List[dict], strip_response: bool = True) -> List[dict]:
+        """Replace prompts whose encode fails with fresh ones from the resample iterator.
+        Shares the backend-agnostic loop with HF / Megatron (see ``rl_core.resample``)."""
+        return resample_encode_failed_inputs(
+            self.template,
+            self._resample_iter,
+            prompts,
+            max_resample_rounds=getattr(self, '_max_resample_rounds', 10),
+            strip_response=strip_response,
+        )
+
+    def _collate_for_workers(self, tg, samples: List[OnPolicySample],
+                             **collate_kwargs) -> Tuple['DPDispatchedDict', List['GRPOBatch']]:
+        """Driver-side collate: ``List[OnPolicySample]`` -> ``({dp_rank: [model_inputs]}, flat_grpo_batches)``.
+
+        The driver owns the whole global batch, so it does the (pure-CPU)
+        ``template.data_collator`` itself — mirroring the non-Ray Megatron path
+        where each rank encodes then collates its own micro-batches. The worker
+        receives the collated micro-batches directly (dispatch='dp') and only runs
+        the rank-local ``prepare_batch`` (PP/CP slice) + forward.
+
+        Layout: split the global batch into ``dp_size`` contiguous shards, each into
+        ``micro_batch_size`` micro-batches, collate each via the shared
+        ``collate_to_grpo_micro_batch``. The second return value is the per-micro-batch
+        ``GRPOBatch`` list IN SAMPLE ORDER (dp_rank major, then micro-batch): the same
+        objects referenced inside the dispatch dict, so the caller fills old/ref logps +
+        advantages on them and they reach ``train_step`` via the dispatch dict.
+        """
+        from swift.rlhf_trainers.utils import collate_to_grpo_micro_batch
+
+        dp_size = tg.dp_size
+        mbs = int(self.args.micro_batch_size)
+        n = len(samples)
+        if n % dp_size != 0:
+            raise ValueError(f'_collate_for_workers: batch size {n} not divisible by dp_size {dp_size}.')
+        shard_size = n // dp_size
+
+        dispatch = DPDispatchedDict()
+        flat_grpo_batches: List['GRPOBatch'] = []
+        for dp_rank in range(dp_size):
+            shard = samples[dp_rank * shard_size:(dp_rank + 1) * shard_size]
+            micro_batches = []
+            for i in range(0, len(shard), mbs):
+                chunk = shard[i:i + mbs]
+                model_inputs, grpo_batch = collate_to_grpo_micro_batch(
+                    chunk,
+                    self.template,
+                    device=self.device,
+                    padding_to=self._padding_to,
+                    router_replay_mode=getattr(self.args, 'router_replay_mode', 'disabled'),
+                    **collate_kwargs,
+                )
+                model_inputs['grpo_batch'] = grpo_batch
+                micro_batches.append(model_inputs)
+                flat_grpo_batches.append(grpo_batch)
+            dispatch[dp_rank] = micro_batches
+        return dispatch, flat_grpo_batches
+
+    def _prepare_state(self) -> None:
+        """Shared ``_prepare_state`` prefix for all Ray trainers.
+
+        Resolves the fields every driver needs (args / template / device /
+        global_batch_size / temperature / beta / steps_per_generation /
+        padding_to) from ``_data_info``. Subclasses call ``super()._prepare_state()``
+        first, then set algorithm-specific state (advantage / dynamic_sample for
+        GRPO; lmbda / teacher for GKD).
+        """
+        assert hasattr(self, '_data_info'), 'call set_data_info() before train()'
+        info = self._data_info
+        args = info['_driver_args']
+        self.args = args
+        self.template: Template = info['template']
+        self.device = torch.device('cpu')
+
+        self.global_batch_size = int(args.global_batch_size)
+        self.temperature = args.temperature
+        self.beta = args.beta
+
+        # steps_per_generation>1: one generation feeds spg training steps.
+        gen_bs = getattr(args, 'generation_batch_size', None)
+        spg = getattr(args, 'steps_per_generation', None)
+        if gen_bs is not None:
+            self._steps_per_generation = max(int(gen_bs) // self.global_batch_size, 1)
+        elif spg is not None:
+            self._steps_per_generation = int(spg)
+        else:
+            self._steps_per_generation = 1
+
+        self._padding_to = info.get('_padding_to')
+
+    def _train_loop(self, tg, train_iters, iteration) -> int:
+        raise NotImplementedError
+
+    def train(self) -> Any:
+        self._prepare_state()
+        tg = self.train_group
+        self._build_dataloader()
+        if getattr(self, '_needs_resample_iterator', False):
+            self._build_resample_iterator()
+
+        args_override = compute_iter_params(self._data_info, tg.dp_size)
+        meta = tg.setup(args_override)
+        train_iters = meta['train_iters']
+        iteration = meta['iteration']
+
+        try:
+            iteration = self._train_loop(tg, train_iters, iteration)
+        finally:
+            results = tg.finalize()
+        return results
+
+    def _maybe_log_completions(self, rollout_with_outputs: List[OnPolicySample], rewards=None, gen_step=None) -> None:
+        """Driver-side ``log_completions``: dump prompt/completion (+reward) to
+        ``output_dir/completions.jsonl``. No-op unless ``args.log_completions`` is set.
+        Completions live on the driver (rollout side), so this is the right place to log them
+        (worker on_log handles scalar metrics)."""
+        args = self.args
+        if not getattr(args, 'log_completions', False) or not rollout_with_outputs:
+            return
+
+        if getattr(self, '_completions_writer', None) is None:
+            self._completions_writer = JsonlWriter(os.path.join(args.output_dir, 'completions.jsonl'))
+        table = []
+        for i, item in enumerate(rollout_with_outputs):
+            msgs = item.messages
+            has_resp = bool(msgs) and msgs[-1].get('role') == 'assistant'
+            completion = self._decode_log_content(msgs[-1].get('content')) if has_resp else ''
+            prompt_msgs = msgs[:-1] if has_resp else msgs
+            row = {'gen_step': gen_step, 'prompt': self._format_log_prompt(prompt_msgs), 'completion': completion}
+            if rewards is not None and i < len(rewards):
+                row['reward'] = float(rewards[i])
+            table.append(row)
+        self._completions_writer.append(table)
+
+    def _format_log_prompt(self, prompt_msgs) -> str:
+        """Render the prompt as the model actually sees it (chat template applied),
+        matching the non-Ray ``_apply_chat_template_to_messages_list`` so completions.jsonl
+        is consistent across backends. Falls back to a plain role/content join if encode fails
+        (e.g. multimodal placeholders the driver template can't re-encode standalone)."""
+        from swift.template import TemplateInputs
+        try:
+            template_inputs = TemplateInputs.from_dict({'messages': [dict(m) for m in prompt_msgs]})
+            res = self.template.encode(template_inputs)
+            return self.template.safe_decode(res['input_ids'])
+        except Exception:
+            return ''.join(f"{m.get('role')}: {m.get('content')}\n" for m in prompt_msgs)
+
+    def _decode_log_content(self, content) -> str:
+        """Decode an assistant message content for logging (mirrors non-Ray)."""
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            return self.template.safe_decode(content)
+        if isinstance(content, dict) and 'input_ids' in content:
+            return self.template.safe_decode(content['input_ids'])
+        return str(content) if content is not None else ''

swift/ray/megatron/checkpoint_engine/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from .base import CheckpointEngine, MasterMetadata, TensorMeta
+from .hccl import HCCLCheckpointEngine
+from .manager import CheckpointEngineManager
+from .mixin import CheckpointEngineMixin
+from .nccl import NCCLCheckpointEngine
+
+__all__ = [
+    'CheckpointEngine',
+    'MasterMetadata',
+    'TensorMeta',
+    'CheckpointEngineMixin',
+    'CheckpointEngineManager',
+    'NCCLCheckpointEngine',
+    'HCCLCheckpointEngine',
+]

swift/ray/megatron/checkpoint_engine/base.py

@@ -0,0 +1,89 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+import socket
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple, TypedDict
+
+if TYPE_CHECKING:
+    import torch
+
+
+class TensorMeta(TypedDict):
+    """Metadata for a tensor in the weight bucket."""
+    name: str
+    shape: 'torch.Size'
+    dtype: 'torch.dtype'
+    offset: int
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        s.bind(('', 0))
+        return s.getsockname()[1]
+
+
+def _is_valid_ipv6_address(addr: str) -> bool:
+    try:
+        socket.inet_pton(socket.AF_INET6, addr)
+        return True
+    except (OSError, socket.error):
+        return False
+
+
+@dataclass
+class MasterMetadata:
+    """Metadata from the master (trainer rank 0) for topology building."""
+    zmq_ip: str
+    zmq_port: int
+    nccl_store_host: str = ''
+    nccl_store_port: int = 0
+
+
+class CheckpointEngine(ABC):
+    rank: Optional[int] = None
+
+    @abstractmethod
+    def prepare(self) -> Dict[str, Any]:
+        """Prepare the checkpoint engine before weight synchronization.
+
+        Allocate weight transfer buffers, setup communication channels,
+        and return metadata needed for topology building.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def build_topology(
+        cls,
+        trainer_world_size: int,
+        rollout_world_size: int,
+        metadata: List[Dict],
+    ) -> Tuple[Dict[str, List[Any]], Dict[str, List[Any]]]:
+        """Build communication topology between trainer and rollout workers.
+
+        Returns (trainer_kwargs, rollout_kwargs) for init_process_group().
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def init_process_group(self, **kwargs):
+        """Initialize the process group for weight synchronization."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def finalize(self):
+        """Finalize: free buffers, optionally destroy the process group."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def send_weights(self, weights: Generator[Tuple[str, 'torch.Tensor'], None, None]):
+        """Send model weights to rollout workers."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def receive_weights(self) -> AsyncGenerator[Tuple[str, 'torch.Tensor'], None]:
+        """Receive model weights from trainer."""
+        raise NotImplementedError

swift/ray/megatron/checkpoint_engine/hccl.py

@@ -0,0 +1,466 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+# Adapted from twinkle/src/twinkle/checkpoint_engine/hccl_checkpoint_engine.py
+"""HCCL-based checkpoint engine for Ascend NPU.
+
+Uses HCCL for weight payload transfer and ZMQ REQ/REP for bucket
+metadata handshakes (reliable, with timeout).
+"""
+from __future__ import annotations
+
+import os
+import time
+import torch
+import zmq
+from dataclasses import dataclass
+from typing import Any, AsyncGenerator, Generator, List, Optional, Tuple
+
+from swift.utils import get_current_device, synchronize
+from swift.utils.logger import get_logger
+from .base import CheckpointEngine, TensorMeta, _find_free_port, _is_valid_ipv6_address
+
+logger = get_logger()
+
+
+def _configure_zmq_socket(socket: zmq.Socket, timeout_ms: int, linger: int = 0) -> None:
+    """Apply timeout/linger options to a ZMQ socket."""
+    socket.setsockopt(zmq.RCVTIMEO, timeout_ms)
+    socket.setsockopt(zmq.SNDTIMEO, timeout_ms)
+    socket.setsockopt(zmq.LINGER, linger)
+
+
+@dataclass
+class HCCLMasterMetadata:
+    """Metadata from the master for HCCL process group initialization."""
+    zmq_ip: str
+    zmq_port: int
+    dist_ip: str
+    dist_port: int
+
+
+def _stateless_init_hccl(
+    master_address: str,
+    master_port: int,
+    rank: int,
+    world_size: int,
+    device: int,
+):
+    """Create a stateless HCCL communicator via vLLM's StatelessProcessGroup."""
+    import socket as _socket
+    from datetime import timedelta
+    from torch.distributed import TCPStore
+    from vllm.distributed.utils import StatelessProcessGroup
+    from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator
+
+    launch_server = (rank == 0)
+    listen_socket = None
+    listen_fd = None
+
+    if launch_server:
+        if _is_valid_ipv6_address(master_address):
+            listen_socket = _socket.socket(_socket.AF_INET6, _socket.SOCK_STREAM)
+        else:
+            listen_socket = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM)
+        listen_socket.setsockopt(_socket.SOL_SOCKET, _socket.SO_REUSEADDR, 1)
+        listen_socket.bind((master_address, master_port))
+        listen_socket.listen()
+        listen_fd = listen_socket.fileno()
+
+    store = TCPStore(
+        host_name=master_address,
+        port=master_port,
+        world_size=world_size,
+        is_master=launch_server,
+        timeout=timedelta(seconds=300),
+        use_libuv=False,
+        master_listen_fd=listen_fd,
+    )
+
+    pg = StatelessProcessGroup(
+        rank=rank,
+        world_size=world_size,
+        store=store,
+        socket=listen_socket,
+        data_expiration_seconds=3600,
+    )
+
+    return PyHcclCommunicator(pg, device=device)
+
+
+class HCCLCheckpointEngine(CheckpointEngine):
+    """HCCL checkpoint engine for Ascend NPU weight synchronization."""
+
+    def __init__(
+        self,
+        bucket_size: int = 3072 << 20,
+        group_name: str = 'swift_ckpt',
+        rebuild_group: bool = True,
+        **kwargs,
+    ) -> None:
+        self.bucket_size = bucket_size
+        self.group_name = group_name
+        self.rebuild_group = rebuild_group
+        self.pyhccl = None
+
+        self.meta_timeout_s = int(os.environ.get('SWIFT_CKPT_HCCL_META_TIMEOUT_S', '300'))
+        self.meta_timeout_ms = self.meta_timeout_s * 1000
+
+        self.device = get_current_device()
+
+        self.is_master = False
+        self.rank: Optional[int] = None
+        self.world_size: Optional[int] = None
+        self.send_buf: Optional[torch.Tensor] = None
+        self.recv_buf: Optional[torch.Tensor] = None
+        self.socket: Optional[zmq.Socket] = None
+        self._zmq_ctx: Optional[zmq.Context] = None
+
+        self._prepared = False
+        self._group_initialized = False
+        self.ip: Optional[str] = None
+        self.zmq_port: Optional[int] = None
+        self.dist_port: Optional[int] = None
+
+    def _new_socket(self, socket_type: int) -> zmq.Socket:
+        assert self._zmq_ctx is not None
+        socket = self._zmq_ctx.socket(socket_type)
+        _configure_zmq_socket(socket, timeout_ms=self.meta_timeout_ms)
+        return socket
+
+    def _recv_pyobj(self, where: str) -> Any:
+        assert self.socket is not None
+        try:
+            return self.socket.recv_pyobj()
+        except zmq.error.Again as e:
+            raise RuntimeError(f'HCCL metadata timeout ({self.meta_timeout_s}s) waiting at {where}.') from e
+
+    def _send_pyobj(self, payload: Any, where: str) -> None:
+        assert self.socket is not None
+        try:
+            self.socket.send_pyobj(payload)
+        except zmq.error.Again as e:
+            raise RuntimeError(f'HCCL metadata timeout ({self.meta_timeout_s}s) sending at {where}.') from e
+
+    def _start_zmq_server(self):
+        import ray
+        self.ip = ray.util.get_node_ip_address().strip('[]')
+        self.zmq_port = _find_free_port()
+        self.dist_port = _find_free_port()
+
+        self._zmq_ctx = zmq.Context()
+        self.socket = self._new_socket(zmq.REP)
+        if _is_valid_ipv6_address(self.ip):
+            address = f'tcp://[{self.ip}]:{self.zmq_port}'
+            self.socket.setsockopt(zmq.IPV6, 1)
+        else:
+            address = f'tcp://{self.ip}:{self.zmq_port}'
+        self.socket.bind(address)
+
+    def _connect_zmq_client(self, metadata: HCCLMasterMetadata):
+        self._zmq_ctx = zmq.Context()
+        self.socket = self._new_socket(zmq.REQ)
+        if _is_valid_ipv6_address(metadata.zmq_ip):
+            address = f'tcp://[{metadata.zmq_ip}]:{metadata.zmq_port}'
+            self.socket.setsockopt(zmq.IPV6, 1)
+        else:
+            address = f'tcp://{metadata.zmq_ip}:{metadata.zmq_port}'
+        self.socket.connect(address)
+
+    # ── Core lifecycle ───────────────────────────────────────────────────
+
+    def prepare(self) -> Optional[HCCLMasterMetadata]:
+        if self._prepared:
+            if self.is_master:
+                return HCCLMasterMetadata(
+                    zmq_ip=self.ip, zmq_port=self.zmq_port, dist_ip=self.ip, dist_port=self.dist_port)
+            return None
+
+        self.send_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device='npu')
+        self.recv_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device='npu')
+
+        if self.is_master:
+            self._start_zmq_server()
+            self._prepared = True
+            return HCCLMasterMetadata(zmq_ip=self.ip, zmq_port=self.zmq_port, dist_ip=self.ip, dist_port=self.dist_port)
+
+        self._prepared = True
+        return None
+
+    def finalize(self):
+        if self.rebuild_group:
+            if self.socket is not None:
+                try:
+                    self.socket.close()
+                except Exception as e:
+                    logger.warning('Error closing ZMQ socket: %s', e)
+                self.socket = None
+
+            if self._zmq_ctx is not None:
+                try:
+                    self._zmq_ctx.term()
+                except Exception as e:
+                    logger.warning('Error terminating ZMQ context: %s', e)
+                self._zmq_ctx = None
+
+            if self.rank is not None and self.rank >= 0 and self.pyhccl is not None:
+                try:
+                    self.pyhccl.destroyComm(self.pyhccl.comm)
+                except Exception:
+                    pass
+                self.pyhccl = None
+
+            self.rank = None
+            self.world_size = None
+            self.send_buf = None
+            self.recv_buf = None
+            self._prepared = False
+            self._group_initialized = False
+
+    @classmethod
+    def build_topology(
+        cls,
+        trainer_world_size: int,
+        rollout_world_size: int,
+        metadata: List[Any],
+    ) -> Tuple[dict, dict]:
+        master_metadata = None
+        for m in metadata:
+            if m is not None:
+                master_metadata = m
+                break
+
+        trainer_kwargs = {
+            'rank': [0] + [-1] * (trainer_world_size - 1),
+            'world_size': [rollout_world_size + 1] * trainer_world_size,
+            'master_metadata': [master_metadata] * trainer_world_size,
+        }
+        rollout_kwargs = {
+            'rank': list(range(1, rollout_world_size + 1)),
+            'world_size': [rollout_world_size + 1] * rollout_world_size,
+            'master_metadata': [master_metadata] * rollout_world_size,
+        }
+        return trainer_kwargs, rollout_kwargs
+
+    def init_process_group(self, rank: int, world_size: int, master_metadata: HCCLMasterMetadata):
+        if rank < 0:
+            self.rank = rank
+            self.world_size = world_size
+            self._group_initialized = True
+            return
+
+        if self._group_initialized and not self.rebuild_group:
+            return
+
+        if self.rebuild_group or self.pyhccl is None:
+            self.pyhccl = _stateless_init_hccl(
+                master_address=master_metadata.dist_ip,
+                master_port=master_metadata.dist_port,
+                rank=rank,
+                world_size=world_size,
+                device=self.device,
+            )
+            self.rank = rank
+            self.world_size = world_size
+        else:
+            assert self.rank == rank
+            assert self.world_size == world_size
+
+        if self.rank > 0 and self.socket is None:
+            self._connect_zmq_client(master_metadata)
+
+        signal = torch.tensor([1], dtype=torch.int8, device=get_current_device())
+        self.pyhccl.all_reduce(signal)
+
+        self._group_initialized = True
+        logger.info('HCCL init_process_group: rank=%d, world_size=%d', self.rank, self.world_size)
+
+    # ── Metadata exchange ────────────────────────────────────────────────
+
+    def _serve_bucket_requests(self, bucket_id: int, metadata: dict) -> None:
+        """Master serves bucket metadata to all receivers via REQ/REP."""
+        assert self.rank == 0 and self.world_size is not None
+        if self.world_size <= 1:
+            return
+
+        pending = set(range(1, self.world_size))
+        while pending:
+            req = self._recv_pyobj(f'NEXT request for bucket={bucket_id}')
+            if not isinstance(req, dict) or req.get('type') != 'NEXT':
+                self._send_pyobj({'ok': False, 'error': f'unexpected: {req}'}, 'NEXT reply')
+                continue
+
+            req_rank = int(req.get('rank', -1))
+            req_bucket_id = int(req.get('bucket_id', -1))
+
+            if req_rank not in pending or req_bucket_id != bucket_id:
+                self._send_pyobj({'ok': False, 'error': 'rank/bucket mismatch'}, 'NEXT reply')
+                continue
+
+            self._send_pyobj({'ok': True, 'metadata': metadata}, 'NEXT reply')
+            pending.remove(req_rank)
+
+    def _request_bucket(self, bucket_id: int) -> dict:
+        """Receiver requests bucket metadata from master via REQ/REP."""
+        assert self.rank is not None and self.rank > 0
+        self._send_pyobj({'type': 'NEXT', 'rank': self.rank, 'bucket_id': bucket_id}, f'NEXT send bucket={bucket_id}')
+        resp = self._recv_pyobj(f'NEXT recv bucket={bucket_id}')
+        if not isinstance(resp, dict) or not resp.get('ok', False):
+            raise RuntimeError(f'Metadata request failed for bucket {bucket_id}: {resp}')
+        return resp['metadata']
+
+    # ── Send / Receive ───────────────────────────────────────────────────
+
+    @torch.no_grad()
+    async def send_weights(
+        self,
+        weights: Generator[Tuple[str, torch.Tensor], None, None],
+    ):
+        assert self.rank is not None and self.rank <= 0
+        if self.rank < 0:
+            for _ in weights:
+                pass
+            return
+
+        send_buf = self.send_buf
+        start_time = time.time()
+        bucket_meta: List[dict] = []
+        offset = 0
+        bucket_id = 0
+        total_params = 0
+        total_bytes = 0
+
+        def _flush(is_last: bool):
+            nonlocal bucket_meta, offset, bucket_id, total_bytes
+            if not bucket_meta and not is_last:
+                return
+            metadata = {
+                'bucket_id': bucket_id,
+                'is_last': is_last,
+                'bucket_meta': bucket_meta,
+                'payload_size': offset,
+            }
+            self._serve_bucket_requests(bucket_id, metadata)
+            self.pyhccl.broadcast(send_buf, src=0)
+            synchronize()
+            total_bytes += offset
+            bucket_id += 1
+            bucket_meta = []
+            offset = 0
+
+        for name, weight in weights:
+            total_params += 1
+            if weight.device.type == 'cpu':
+                weight = weight.to(get_current_device())
+            if not weight.is_contiguous():
+                weight = weight.contiguous()
+
+            weight_u8 = weight.view(-1).view(torch.uint8)
+            nbytes = weight_u8.numel()
+            if nbytes == 0:
+                if offset >= self.bucket_size:
+                    _flush(is_last=False)
+                bucket_meta.append({
+                    'name': name,
+                    'shape': weight.shape,
+                    'dtype': weight.dtype,
+                    'offset': offset,
+                    'nbytes': 0,
+                    'chunk_offset': 0,
+                    'total_nbytes': 0,
+                })
+                continue
+
+            chunk_offset = 0
+            while chunk_offset < nbytes:
+                if offset >= self.bucket_size:
+                    _flush(is_last=False)
+                chunk_nbytes = min(self.bucket_size - offset, nbytes - chunk_offset)
+                send_buf[offset:offset + chunk_nbytes].copy_(weight_u8[chunk_offset:chunk_offset + chunk_nbytes])
+                bucket_meta.append({
+                    'name': name,
+                    'shape': weight.shape,
+                    'dtype': weight.dtype,
+                    'offset': offset,
+                    'nbytes': chunk_nbytes,
+                    'chunk_offset': chunk_offset,
+                    'total_nbytes': nbytes,
+                })
+                offset += chunk_nbytes
+                chunk_offset += chunk_nbytes
+
+        _flush(is_last=True)
+
+        elapsed = time.time() - start_time
+        bandwidth = total_bytes / elapsed / (1024**3) if elapsed > 0 else 0.0
+        logger.debug('HCCL send_weights done: rank=%d, params=%d, time=%.2fs, bw=%.2f GB/s', self.rank, total_params,
+                     elapsed, bandwidth)
+
+    @torch.no_grad()
+    async def receive_weights(self) -> AsyncGenerator[Tuple[str, torch.Tensor], None]:
+        assert self.rank is not None and self.rank > 0
+
+        recv_buf = self.recv_buf
+        bucket_id = 0
+        total_params = 0
+        total_bytes = 0
+        start_time = time.time()
+        partial_tensors: dict = {}
+
+        while True:
+            metadata = self._request_bucket(bucket_id)
+            self.pyhccl.broadcast(recv_buf, src=0)
+            synchronize()
+
+            bucket_meta = metadata['bucket_meta']
+            entries = bucket_meta.values() if isinstance(bucket_meta, dict) else bucket_meta
+            total_bytes += int(metadata.get('payload_size', self.bucket_size))
+
+            for meta in entries:
+                name = meta['name']
+                dtype = meta['dtype']
+                shape = meta['shape']
+                if not isinstance(shape, torch.Size):
+                    shape = torch.Size(shape)
+                offset = int(meta['offset'])
+                nbytes = int(meta.get('nbytes', dtype.itemsize * shape.numel()))
+                chunk_offset = int(meta.get('chunk_offset', 0))
+                total_nbytes = int(meta.get('total_nbytes', dtype.itemsize * shape.numel()))
+
+                if nbytes == total_nbytes and chunk_offset == 0:
+                    tensor = recv_buf[offset:offset + nbytes].view(dtype=dtype).view(shape)
+                    yield name, tensor
+                    total_params += 1
+                    continue
+
+                state = partial_tensors.get(name)
+                if state is None:
+                    state = {
+                        'buffer': torch.empty(total_nbytes, dtype=torch.uint8, device=recv_buf.device),
+                        'dtype': dtype,
+                        'shape': shape,
+                        'total': total_nbytes,
+                        'received': 0,
+                    }
+                    partial_tensors[name] = state
+
+                if nbytes > 0:
+                    state['buffer'][chunk_offset:chunk_offset + nbytes].copy_(recv_buf[offset:offset + nbytes])
+                state['received'] += nbytes
+
+                if state['received'] == state['total']:
+                    full_size = dtype.itemsize * shape.numel()
+                    tensor = state['buffer'][:full_size].view(dtype=dtype).view(shape)
+                    yield name, tensor
+                    total_params += 1
+                    del partial_tensors[name]
+
+            if metadata['is_last']:
+                if partial_tensors:
+                    pending = ', '.join(sorted(partial_tensors.keys())[:8])
+                    raise RuntimeError(f'Incomplete chunked weights: {len(partial_tensors)} pending: {pending}')
+                break
+            bucket_id += 1
+
+        elapsed = time.time() - start_time
+        bandwidth = total_bytes / elapsed / (1024**3) if elapsed > 0 else 0.0
+        logger.debug('HCCL receive_weights done: rank=%d, params=%d, time=%.2fs, bw=%.2f GB/s', self.rank, total_params,
+                     elapsed, bandwidth)

swift/ray/megatron/checkpoint_engine/manager.py

@@ -0,0 +1,171 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+import ray
+from typing import Any, List, Optional
+
+from swift.utils.logger import get_logger
+from .nccl import NCCLCheckpointEngine
+
+logger = get_logger()
+
+
+class CheckpointEngineManager:
+
+    def __init__(
+        self,
+        train_actors: List[Any],
+        rollout_replicas: List[Any],
+        *,
+        weight_sync_mode: str = 'nccl',
+        is_colocated: bool = False,
+        sleep_level: int = 1,
+        train_group: Any,
+    ):
+        self.train_actors = train_actors
+        self._rollout_replicas = rollout_replicas
+        self.rollout_actors = [r.primary for r in rollout_replicas]
+        self._weight_sync_mode = weight_sync_mode
+        self.is_colocated = is_colocated
+        self._train_group = train_group
+
+        if is_colocated and sleep_level >= 2:
+            logger.warning(
+                'sleep_level=%d capped to 1 in colocate mode '
+                '(out-of-process vLLM cannot safely discard all GPU memory).', sleep_level)
+            sleep_level = 1
+        self.sleep_level = sleep_level
+
+        self.base_sync_done: bool = False
+        self._model_keys: Optional[List[str]] = None
+        self._sleeping_tags: set = set()
+
+    def sync_weights(self, merge_and_sync: bool = True) -> None:
+        """Synchronize weights from training model to rollout replicas."""
+        if self.is_colocated:
+            self.sleep_rollout()
+            self.wake_up_rollout(tags=['weights'])
+        if self._weight_sync_mode == 'naive':
+            self._sync_weights_naive(merge_and_sync)
+        else:
+            self._sync_weights_nccl(merge_and_sync)
+
+    def sleep_rollout(self) -> None:
+        if self._sleeping_tags:
+            return
+        for replica in self._rollout_replicas:
+            replica.sleep(level=self.sleep_level)
+        self._sleeping_tags = {'weights', 'kv_cache'}
+        logger.debug('CheckpointEngineManager: rollout replicas sleeping (level=%d)', self.sleep_level)
+
+    def wake_up_rollout(self, tags: Optional[List[str]] = None) -> None:
+        if not self._sleeping_tags:
+            return
+        for replica in self._rollout_replicas:
+            replica.wake_up(tags=tags)
+        if tags is None:
+            self._sleeping_tags.clear()
+        else:
+            self._sleeping_tags -= set(tags)
+        logger.debug('CheckpointEngineManager: rollout wake_up tags=%s, still_sleeping=%s', tags, self._sleeping_tags)
+
+    def _sync_weights_naive(self, merge_and_sync: bool) -> None:
+        tg = self._train_group
+        adapter_only = self.base_sync_done and not merge_and_sync
+        need_merge = not adapter_only and merge_and_sync
+        if need_merge:
+            tg.merge_lora()
+        if self.is_colocated:
+            tg.offload_to_cpu()
+        try:
+            tg.update_weights(adapter_only=adapter_only)
+        finally:
+            if self.is_colocated:
+                tg.reload_to_gpu()
+            if need_merge:
+                tg.unmerge_lora()
+        if not self.base_sync_done:
+            self.base_sync_done = True
+            logger.debug('CheckpointEngineManager[naive]: initial weight sync done')
+
+    def _sync_weights_nccl(self, merge_and_sync: bool) -> None:
+        """NCCL broadcast weight sync path.
+
+        Lifecycle:
+        1. prepare_checkpoint_engine on all actors
+        2. build_topology
+        3. init_process_group on all actors (concurrent — required for TCPStore)
+        4. send_weights (train) + receive_weights (rollout) concurrently
+        5. finalize_checkpoint_engine on all actors
+        """
+        n_train = len(self.train_actors)
+        n_rollout = len(self.rollout_actors)
+
+        # 1. Prepare — train side: rank 0 is master, others are not
+        is_master_flags = [True] + [False] * (n_train - 1)
+        prepare_refs = [
+            actor.prepare_checkpoint_engine.remote(flag) for actor, flag in zip(self.train_actors, is_master_flags)
+        ]
+        prepare_results = ray.get(prepare_refs)
+        model_metadata = prepare_results[0]
+
+        # 1b. Prepare — rollout side: all non-master
+        rollout_prepare_refs = [actor.prepare_checkpoint_engine.remote(False) for actor in self.rollout_actors]
+        ray.get(rollout_prepare_refs)
+
+        # 2. Build topology
+        model_kwargs, rollout_kwargs = NCCLCheckpointEngine.build_topology(n_train, n_rollout, [model_metadata])
+
+        # 3. Init process groups (MUST be concurrent — TCPStore server
+        #    blocks until all clients connect)
+        train_init_refs = [
+            actor.init_checkpoint_process_group.remote(
+                rank=model_kwargs['rank'][i],
+                world_size=model_kwargs['world_size'][i],
+                master_metadata=model_kwargs['master_metadata'][i],
+            ) for i, actor in enumerate(self.train_actors)
+        ]
+        rollout_init_refs = [
+            actor.init_checkpoint_process_group.remote(
+                rank=rollout_kwargs['rank'][i],
+                world_size=rollout_kwargs['world_size'][i],
+                master_metadata=rollout_kwargs['master_metadata'][i],
+            ) for i, actor in enumerate(self.rollout_actors)
+        ]
+        ray.get(train_init_refs + rollout_init_refs)
+
+        # 4. Send/receive weights (concurrent)
+        adapter_only = self.base_sync_done and not merge_and_sync
+        need_merge = not adapter_only and merge_and_sync
+        peft_config = None
+        if adapter_only:
+            peft_config = ray.get(self.train_actors[0].get_peft_config_dict.remote())
+
+        if need_merge:
+            merge_refs = [actor.merge_lora.remote() for actor in self.train_actors]
+            ray.get(merge_refs)
+
+        train_send_refs = [
+            actor.send_checkpoint_weights.remote(adapter_only=adapter_only) for actor in self.train_actors
+        ]
+        rollout_recv_refs = [
+            actor.receive_checkpoint_weights.remote(
+                base_sync_done=self.base_sync_done,
+                peft_config=peft_config,
+            ) for actor in self.rollout_actors
+        ]
+        ray.get(train_send_refs + rollout_recv_refs)
+
+        if need_merge:
+            unmerge_refs = [actor.unmerge_lora.remote() for actor in self.train_actors]
+            ray.get(unmerge_refs)
+
+        # 5. Finalize
+        train_fin_refs = [actor.finalize_checkpoint_engine.remote() for actor in self.train_actors]
+        rollout_fin_refs = [actor.finalize_checkpoint_engine.remote() for actor in self.rollout_actors]
+        ray.get(train_fin_refs + rollout_fin_refs)
+
+        if not self.base_sync_done:
+            self.base_sync_done = True
+            logger.info('CheckpointEngineManager[nccl]: initial weight sync to %d replica(s) '
+                        '(lora_only=%s)', n_rollout, not merge_and_sync)

swift/ray/megatron/checkpoint_engine/mixin.py

@@ -0,0 +1,57 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+from typing import Optional
+
+from .base import CheckpointEngine, MasterMetadata
+
+
+class CheckpointEngineMixin:
+    """Mixin providing checkpoint engine lifecycle methods for Ray actors.
+
+    Add this to Ray actors (MegatronWorker, VllmServer) to give them
+    checkpoint engine capabilities.  The Manager calls these methods
+    via ``ray.remote()`` to coordinate NCCL weight sync.
+
+    Applied to ``MegatronWorker`` (send side) and ``VllmServer``
+    (receive side).
+    """
+
+    _checkpoint_engine: Optional[CheckpointEngine] = None
+    _bucket_size: int = 3072 << 20
+
+    def _get_or_create_checkpoint_engine(self) -> CheckpointEngine:
+        """Get or create the checkpoint engine instance (lazy singleton)."""
+        if self._checkpoint_engine is None:
+            from transformers.utils import is_torch_npu_available
+            if is_torch_npu_available():
+                from .hccl import HCCLCheckpointEngine
+                self._checkpoint_engine = HCCLCheckpointEngine(self._bucket_size, rebuild_group=False)
+            else:
+                from .nccl import NCCLCheckpointEngine
+                self._checkpoint_engine = NCCLCheckpointEngine(self._bucket_size)
+        return self._checkpoint_engine
+
+    def prepare_checkpoint_engine(self, is_master: bool) -> Optional[MasterMetadata]:
+        """Prepare checkpoint engine for weight sync."""
+        engine = self._get_or_create_checkpoint_engine()
+        engine.is_master = is_master
+        return engine.prepare()
+
+    def init_checkpoint_process_group(
+        self,
+        rank: int,
+        world_size: int,
+        master_metadata: MasterMetadata,
+    ) -> None:
+        engine = self._get_or_create_checkpoint_engine()
+        engine.init_process_group(
+            rank=rank,
+            world_size=world_size,
+            master_metadata=master_metadata,
+        )
+
+    def finalize_checkpoint_engine(self) -> None:
+        """Finalize checkpoint engine: release buffers, optionally destroy group."""
+        if self._checkpoint_engine is not None:
+            self._checkpoint_engine.finalize()

swift/ray/megatron/checkpoint_engine/nccl.py

@@ -0,0 +1,378 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+import asyncio
+import ray
+import time
+import torch
+import torch.distributed as dist
+from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple
+
+from swift.utils import get_current_device, synchronize
+from swift.utils.logger import get_logger
+from .base import CheckpointEngine, MasterMetadata, TensorMeta, _find_free_port, _is_valid_ipv6_address
+
+logger = get_logger()
+
+
+def _pg_broadcast(pg, tensor, src: int = 0):
+    """Broadcast tensor via raw ProcessGroupNCCL (unregistered PG)."""
+    opts = dist.BroadcastOptions()
+    opts.rootRank = src
+    work = pg.broadcast([tensor], opts)
+    work.wait()
+
+
+class BroadcastOperation:
+    """Async NCCL broadcast in a thread pool executor."""
+
+    def __init__(self, rank, pg, bucket, metadata, zmq_socket, topic):
+        self.rank = rank
+        self.pg = pg
+        self.bucket = bucket
+        self.metadata = metadata
+        self.socket = zmq_socket
+        self.topic = topic
+
+        loop = asyncio.get_running_loop()
+        self._task = loop.run_in_executor(None, self._run)
+
+    def _run(self):
+        import zmq
+        if self.rank == 0:
+            self.socket.send_string(self.topic, flags=zmq.SNDMORE)
+            self.socket.send_pyobj(self.metadata)
+        else:
+            self.socket.recv_string()
+            self.metadata = self.socket.recv_pyobj()
+        _pg_broadcast(self.pg, self.bucket, src=0)
+
+    async def wait_for_complete(self) -> dict:
+        await self._task
+        return self.metadata
+
+
+class NCCLCheckpointEngine(CheckpointEngine):
+
+    def __init__(
+        self,
+        bucket_size: int = 3072 << 20,
+        group_name: str = 'swift_ckpt',
+        rebuild_group: bool = False,
+        **kwargs,
+    ) -> None:
+        self.bucket_size = bucket_size
+        self.group_name = group_name
+        self.rebuild_group = rebuild_group
+
+        self.is_master = False
+        self.topic = 'bucket_metadata'
+
+        self.rank = None
+        self.world_size = None
+        self.send_buf = None
+        self.recv_buf = None
+        self.socket = None
+
+        self._pg = None
+        self._store = None
+        self._prepared = False
+        self._group_initialized = False
+
+    def _start_zmq_server(self):
+        import zmq
+        self.ip = ray.util.get_node_ip_address().strip('[]')
+        self.listen_port = _find_free_port()
+
+        context = zmq.Context()
+        self.socket = context.socket(zmq.PUB)
+        if _is_valid_ipv6_address(self.ip):
+            address = f'tcp://[{self.ip}]:{self.listen_port}'
+            self.socket.setsockopt(zmq.IPV6, 1)
+        else:
+            address = f'tcp://{self.ip}:{self.listen_port}'
+        self.socket.bind(address)
+
+    def _connect_zmq_client(self, metadata: MasterMetadata):
+        import zmq
+        context = zmq.Context()
+        self.socket = context.socket(zmq.SUB)
+        if _is_valid_ipv6_address(metadata.zmq_ip):
+            address = f'tcp://[{metadata.zmq_ip}]:{metadata.zmq_port}'
+            self.socket.setsockopt(zmq.IPV6, 1)
+        else:
+            address = f'tcp://{metadata.zmq_ip}:{metadata.zmq_port}'
+        self.socket.connect(address)
+        self.socket.setsockopt_string(zmq.SUBSCRIBE, self.topic)
+
+    def prepare(self) -> Optional[MasterMetadata]:
+        """Allocate buffers and start ZMQ server (master only). Idempotent."""
+        if self._prepared:
+            if self.is_master:
+                return MasterMetadata(
+                    zmq_ip=self.ip,
+                    zmq_port=self.listen_port,
+                    nccl_store_host=self._nccl_store_host,
+                    nccl_store_port=self._nccl_store_port,
+                )
+            return None
+
+        device = get_current_device()
+        self.send_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=device)
+        self.recv_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=device)
+
+        if self.is_master:
+            self._start_zmq_server()
+            self._nccl_store_host = self.ip
+            self._nccl_store_port = _find_free_port()
+            self._prepared = True
+            return MasterMetadata(
+                zmq_ip=self.ip,
+                zmq_port=self.listen_port,
+                nccl_store_host=self._nccl_store_host,
+                nccl_store_port=self._nccl_store_port,
+            )
+        else:
+            self._prepared = True
+            return None
+
+    def finalize(self):
+        """Clean up resources. Full teardown only when rebuild_group=True."""
+        if self.rebuild_group:
+            if self.socket is not None:
+                try:
+                    self.socket.close()
+                except Exception as e:
+                    logger.warning('Error closing ZMQ socket: %s', e)
+                self.socket = None
+
+            if self._pg is not None:
+                self._pg = None
+                self._store = None
+
+            self.rank = None
+            self.world_size = None
+            self.send_buf = None
+            self.recv_buf = None
+            self._prepared = False
+            self._group_initialized = False
+
+    @classmethod
+    def build_topology(
+        cls,
+        trainer_world_size: int,
+        rollout_world_size: int,
+        metadata: List[Dict],
+    ) -> Tuple[Dict[str, List[Any]], Dict[str, List[Any]]]:
+        """Build NCCL broadcast topology: trainer rank0 as source, rollout as receivers."""
+        master_metadata = metadata[0]
+
+        trainer_kwargs = {
+            'rank': [0] + [-1] * (trainer_world_size - 1),
+            'world_size': [rollout_world_size + 1] * trainer_world_size,
+            'master_metadata': [master_metadata] * trainer_world_size,
+        }
+        rollout_kwargs = {
+            'rank': list(range(1, rollout_world_size + 1)),
+            'world_size': [rollout_world_size + 1] * rollout_world_size,
+            'master_metadata': [master_metadata] * rollout_world_size,
+        }
+        return trainer_kwargs, rollout_kwargs
+
+    def init_process_group(self, rank: int, world_size: int, master_metadata: MasterMetadata):
+        """Initialize a dedicated NCCL process group for weight synchronization.
+
+        Creates a ``ProcessGroupNCCL`` directly (without registering it in
+        the default ``_World``), using a ``TCPStore`` hosted by the master
+        for rendezvous.
+
+        Idempotent when ``rebuild_group=False``.
+        """
+        import os
+
+        if rank < 0:
+            self.rank = rank
+            self.world_size = world_size
+            self._group_initialized = True
+            return
+
+        if self._group_initialized and not self.rebuild_group:
+            return
+
+        if self._pg is None:
+            self.rank = rank
+            self.world_size = world_size
+
+            os.environ['NCCL_CUMEM_ENABLE'] = '0'
+
+            is_store_master = (rank == 0)
+            self._store = dist.TCPStore(
+                host_name=master_metadata.nccl_store_host,
+                port=master_metadata.nccl_store_port,
+                world_size=world_size,
+                is_master=is_store_master,
+                wait_for_workers=True,
+            )
+            self._pg = dist.ProcessGroupNCCL(
+                self._store,
+                rank,
+                world_size,
+            )
+        else:
+            assert self.rank == rank, f'rank {rank} != self.rank {self.rank}'
+            assert self.world_size == world_size
+
+        if self.rank > 0 and self.socket is None:
+            self._connect_zmq_client(master_metadata)
+
+        barrier_tensor = torch.zeros(1, dtype=torch.int32, device=get_current_device())
+        _pg_broadcast(self._pg, barrier_tensor, src=0)
+        synchronize()
+
+        # ZMQ PUB/SUB "slow joiner" mitigation: after NCCL barrier confirms
+        # all participants are connected, give SUB sockets time to fully
+        # establish the subscription before PUB sends metadata.
+        if self.rank == 0 and self.socket is not None:
+            time.sleep(0.1)
+
+        self._group_initialized = True
+
+    # ── Send / Receive ───────────────────────────────────────────────────
+
+    @torch.no_grad()
+    async def send_weights(
+        self,
+        weights: Generator[Tuple[str, 'torch.Tensor'], None, None],
+    ):
+        """Send model weights to rollout workers via NCCL broadcast.
+
+        Uses double buffering: fill send_buf while the previous bucket
+        is being broadcast, then swap buffers.
+        """
+        assert self.rank is not None and self.rank <= 0
+
+        if self.rank < 0:
+            for name, weight in weights:
+                pass
+            return
+
+        send_buf, recv_buf = self.send_buf, self.recv_buf
+        broadcast_op = None
+        start_time = time.time()
+        bucket_meta: Dict[str, TensorMeta] = {}
+        offset = 0
+
+        for name, weight in weights:
+            if offset + weight.nbytes > self.bucket_size:
+                synchronize()
+                if broadcast_op is not None:
+                    await broadcast_op.wait_for_complete()
+
+                broadcast_op = BroadcastOperation(
+                    rank=self.rank,
+                    pg=self._pg,
+                    bucket=send_buf,
+                    metadata={
+                        'bucket_meta': bucket_meta,
+                        'is_last': False
+                    },
+                    zmq_socket=self.socket,
+                    topic=self.topic,
+                )
+                send_buf, recv_buf = recv_buf, send_buf
+                bucket_meta = {}
+                offset = 0
+
+            assert offset + weight.nbytes <= self.bucket_size, (
+                f'Weight {name}({weight.shape}, {weight.dtype}) is too large '
+                f'for bucket ({self.bucket_size / 1e6:.1f} MB). Increase bucket_size.')
+
+            bucket_meta[name] = {
+                'name': name,
+                'shape': weight.shape,
+                'dtype': weight.dtype,
+                'offset': offset,
+            }
+            send_buf[offset:offset + weight.nbytes].copy_(weight.view(-1).view(torch.uint8), non_blocking=True)
+            offset += weight.nbytes
+
+        synchronize()
+        if broadcast_op is not None:
+            await broadcast_op.wait_for_complete()
+
+        broadcast_op = BroadcastOperation(
+            rank=self.rank,
+            pg=self._pg,
+            bucket=send_buf,
+            metadata={
+                'bucket_meta': bucket_meta,
+                'is_last': True
+            },
+            zmq_socket=self.socket,
+            topic=self.topic,
+        )
+        await broadcast_op.wait_for_complete()
+
+        logger.debug('Rank %d send weights done, time cost: %.2fs', self.rank, time.time() - start_time)
+
+    @torch.no_grad()
+    async def receive_weights(self) -> AsyncGenerator[Tuple[str, 'torch.Tensor'], None]:
+        """Receive model weights from trainer via NCCL broadcast.
+
+        Uses double buffering: receive into recv_buf while processing
+        send_buf, then swap.
+
+        Yields:
+            Tuples of (name, tensor).  The tensor is a *view* into the
+            receive buffer — callers that need to keep it should clone it.
+        """
+        assert self.rank is not None and self.rank > 0
+
+        send_buf, recv_buf = self.send_buf, self.recv_buf
+        total_bytes, total_params = 0, 0
+
+        start_time = time.time()
+        broadcast_op = BroadcastOperation(
+            rank=self.rank,
+            pg=self._pg,
+            bucket=recv_buf,
+            metadata=None,
+            zmq_socket=self.socket,
+            topic=self.topic,
+        )
+        metadata = await broadcast_op.wait_for_complete()
+        total_bytes += self.bucket_size
+        total_params += len(metadata['bucket_meta'])
+        send_buf, recv_buf = recv_buf, send_buf
+
+        while not metadata['is_last']:
+            broadcast_op = BroadcastOperation(
+                rank=self.rank,
+                pg=self._pg,
+                bucket=recv_buf,
+                metadata=None,
+                zmq_socket=self.socket,
+                topic=self.topic,
+            )
+            for name, meta in metadata['bucket_meta'].items():
+                dtype, shape = meta['dtype'], meta['shape']
+                size = dtype.itemsize * shape.numel()
+                tensor = send_buf[meta['offset']:meta['offset'] + size].view(dtype=dtype).view(shape)
+                yield name, tensor
+
+            metadata = await broadcast_op.wait_for_complete()
+            total_bytes += self.bucket_size
+            total_params += len(metadata['bucket_meta'])
+            synchronize()
+            send_buf, recv_buf = recv_buf, send_buf
+
+        for name, meta in metadata['bucket_meta'].items():
+            dtype, shape = meta['dtype'], meta['shape']
+            size = dtype.itemsize * shape.numel()
+            tensor = send_buf[meta['offset']:meta['offset'] + size].view(dtype=dtype).view(shape)
+            yield name, tensor
+
+        elapsed = time.time() - start_time
+        bandwidth = total_bytes / elapsed / (1024 * 1024 * 1024) if elapsed > 0 else 0
+        logger.debug('receive_weights done: rank=%d, params=%d, time=%.2fs, bandwidth=%.2f GB/s', self.rank,
+                     total_params, elapsed, bandwidth)

swift/ray/megatron/driver_utils.py

@@ -0,0 +1,289 @@
+"""Driver-side helpers shared by the Ray Megatron pipeline.
+
+This module owns the "plain Python" helpers that do not belong to any
+particular trainer class:
+
+* YAML → dict config parsing and merging
+* structured config parsing from YAML group dicts
+* driver-side dataset building (via dict, no argv round-trip)
+* train/eval iteration bookkeeping
+* extracting the canonical iteration from worker results
+"""
+import json
+import os
+import torch
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+from swift.arguments import BaseArguments
+from swift.utils import get_logger, parse_args, seed_everything, to_abspath
+
+logger = get_logger()
+
+_RAY_ONLY_KEYS = frozenset({'gpus', 'colocate_groups', 'nnodes'})
+
+_PARALLEL_DEFAULTS: Dict[str, Any] = {
+    'tensor_model_parallel_size': 1,
+    'pipeline_model_parallel_size': 1,
+    'context_parallel_size': 1,
+}
+
+
+def parse_args_from_dict(class_type, cfg: Dict[str, Any]):
+    """Construct a dataclass from a config dict via HfArgumentParser."""
+    argv = _dict_to_argv(cfg)
+    args, remaining_args = parse_args(class_type, argv)
+    if remaining_args:
+        logger.warning('parse_args_from_dict: unrecognised args: %s', remaining_args)
+    return args
+
+
+def _dict_to_argv(cfg: Dict[str, Any]) -> List[str]:
+    argv: List[str] = []
+    for k, v in cfg.items():
+        if k in _RAY_ONLY_KEYS or v is None:
+            continue
+        flag = f'--{k}'
+        if isinstance(v, bool):
+            argv += [flag, str(v).lower()]
+        elif isinstance(v, (list, tuple)):
+            argv.append(flag)
+            argv += [str(item) for item in v]
+        elif isinstance(v, dict):
+            argv += [flag, json.dumps(v)]
+        else:
+            argv += [flag, str(v)]
+    return argv
+
+
+@dataclass
+class RayConfig:
+    rlhf_type: str = 'grpo'
+    colocate_groups: List[List[str]] = field(default_factory=list)
+    train_gpus: int = 0
+    rollout_gpus: int = 0
+    teacher_gpus: int = 0
+    sleep_level: int = 1
+    nnodes: int = 1
+
+    @property
+    def group_gpus(self) -> Dict[str, int]:
+        return {
+            'train': self.train_gpus,
+            'rollout': self.rollout_gpus,
+            'teacher': self.teacher_gpus,
+        }
+
+    def gpus_as_process_on_nodes(self, total_gpus: int) -> List[int]:
+        """Split ``total_gpus`` evenly across ``nnodes`` for ResourcePool."""
+        if self.nnodes <= 1:
+            return [total_gpus]
+        per_node, remainder = divmod(total_gpus, self.nnodes)
+        if remainder != 0:
+            raise ValueError(f'total_gpus={total_gpus} is not evenly divisible by nnodes={self.nnodes}')
+        return [per_node] * self.nnodes
+
+
+def parse_ray_yaml(config_path: str) -> 'tuple[RayConfig, Dict[str, Dict[str, Any]], Dict[str, Any]]':
+    """Parse a Ray YAML config into (ray_config, group_dicts, shared_dict)."""
+    import yaml
+    with open(config_path) as f:
+        raw = yaml.safe_load(f)
+
+    rlhf_type = raw.get('rlhf_type')
+    colocate_groups = raw.pop('colocate_groups', [])
+    sleep_level = int(raw.pop('sleep_level', 1))
+    nnodes = int(raw.pop('nnodes', 1))
+
+    group_configs: Dict[str, dict] = {}
+    for g in KNOWN_GROUPS:
+        group_configs[g] = raw.pop(g, {}) or {}
+
+    gpu_counts = {g: int(cfg.pop('gpus', 0)) for g, cfg in group_configs.items()}
+    shared_config = dict(raw)
+    for key, default in _PARALLEL_DEFAULTS.items():
+        shared_config.setdefault(key, default)
+
+    ray_config = RayConfig(
+        rlhf_type=rlhf_type,
+        colocate_groups=colocate_groups,
+        train_gpus=gpu_counts.get('train', 0),
+        rollout_gpus=gpu_counts.get('rollout', 0),
+        teacher_gpus=gpu_counts.get('teacher', 0),
+        sleep_level=sleep_level,
+        nnodes=nnodes,
+    )
+
+    _validate_colocate_groups(colocate_groups, gpu_counts)
+
+    return ray_config, group_configs, shared_config
+
+
+KNOWN_GROUPS = frozenset(('train', 'rollout', 'teacher'))
+
+
+def _validate_colocate_groups(
+    colocate_groups: List[List[str]],
+    gpu_counts: Dict[str, int],
+) -> None:
+    """Validate colocate_groups: ≥2 roles, known, non-overlapping, each with gpus > 0."""
+    if not colocate_groups:
+        return
+    seen: set = set()
+    for idx, group in enumerate(colocate_groups):
+        if not isinstance(group, list) or len(group) < 2:
+            raise ValueError(f'colocate_groups[{idx}] must be a list of ≥2 roles, '
+                             f'got {group!r}')
+        group_gpu_counts = set()
+        for role in group:
+            if role not in KNOWN_GROUPS:
+                raise ValueError(f'colocate_groups[{idx}] contains unknown role {role!r}; '
+                                 f'valid roles: {sorted(KNOWN_GROUPS)}')
+            if role in seen:
+                raise ValueError(f'Role {role!r} appears in multiple colocate groups')
+            seen.add(role)
+            n = gpu_counts.get(role, 0)
+            if n <= 0:
+                raise ValueError(f'Role {role!r} in colocate_groups[{idx}] has 0 GPUs; '
+                                 f'colocated roles must each have gpus > 0')
+            group_gpu_counts.add(n)
+        if len(group_gpu_counts) > 1:
+            raise ValueError(f'colocate_groups[{idx}] roles have different GPU counts '
+                             f'{dict(zip(group, [gpu_counts[r] for r in group]))}; '
+                             f'colocated roles must share the same GPU set')
+
+
+def merge_group_dict(shared: Dict[str, Any], group: Dict[str, Any]) -> Dict[str, Any]:
+    """Merge shared + group config, stripping Ray-only keys and None values."""
+    merged = {**shared, **group}
+    for k in _RAY_ONLY_KEYS:
+        merged.pop(k, None)
+    return {k: v for k, v in merged.items() if v is not None}
+
+
+def estimate_dp_size(cfg: Dict[str, Any], gpus: int) -> int:
+    """Estimate DP size from a merged group config dict."""
+    tp = cfg.get('tensor_model_parallel_size', 1)
+    pp = cfg.get('pipeline_model_parallel_size', 1)
+    cp = cfg.get('context_parallel_size', 1)
+    assert gpus % (tp * pp * cp) == 0
+    return gpus // (tp * pp * cp)
+
+
+def build_dataset_from_dict(cfg: Dict[str, Any]):
+    """Build dataset on the driver without instantiating a Megatron pipeline.
+    """
+    from swift.megatron.arguments import MegatronRLHFArguments
+    from swift.rlhf_trainers.utils import identity_data_collator
+    cfg = dict(cfg)
+    cfg['skip_megatron_init'] = True
+    args = parse_args_from_dict(MegatronRLHFArguments, cfg)
+
+    if hasattr(args, 'seed'):
+        seed_everything(args.seed)
+
+    rlhf_type = args.rlhf_type
+    if rlhf_type in ('grpo', 'gkd'):
+        args.remove_unused_columns = False
+    if args.output_dir is None:
+        args.output_dir = f'megatron_output/{args.model_suffix}'
+    args.output_dir = to_abspath(args.output_dir)
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    with torch.device('meta'):
+        _, processor = args.get_model_processor(load_model=False, download_model=args.mcore_model is None)
+
+    template = _prepare_template(args, processor)
+    train_dataset, val_dataset = _prepare_dataset(args)
+
+    data_collator = identity_data_collator if rlhf_type in ('grpo', 'gkd') else template.data_collator
+
+    # TODO: integrate val_dataset / eval_iters into the training loop
+    return {
+        'train_dataset': train_dataset,
+        'val_dataset': val_dataset,
+        'data_collator': data_collator,
+        'micro_batch_size': args.micro_batch_size,
+        'global_batch_size': args.global_batch_size,
+        'padding_free': args.padding_free,
+        'num_train_epochs': args.num_train_epochs,
+        'train_iters': args.train_iters,
+        'save_strategy': args.save_strategy,
+        'eval_iters': args.eval_iters,
+        'num_generations': args.num_generations,
+        'template': template,
+        '_driver_args': args,
+    }
+
+
+def _prepare_template(args, processor):
+    """Create template from args and processor — no pipeline object needed."""
+    template = args.get_template(processor)
+    mode_mapping = {'grpo': 'train', 'gkd': 'train', 'kto': 'kto'}
+    template.set_mode(mode_mapping.get(args.rlhf_type, 'rlhf'))
+    template.use_megatron = True
+    return template
+
+
+def _prepare_dataset(args: BaseArguments):
+    """Load and optionally encode dataset — no pipeline object needed."""
+    # Ray pipeline has no validation/eval loop yet
+    if args.split_dataset_ratio and args.split_dataset_ratio > 0:
+        logger.warning(
+            'Ray pipeline has no validation loop yet; overriding split_dataset_ratio '
+            '%s -> 0.0 (no validation split).', args.split_dataset_ratio)
+        args.split_dataset_ratio = 0.0
+    if args.val_dataset:
+        logger.warning('Ray pipeline has no validation loop yet; ignoring val_dataset=%s.', args.val_dataset)
+        args.val_dataset = []
+
+    assert args.rlhf_type in ('grpo', 'gkd')
+    return args.load_dataset()
+
+
+def compute_iter_params(data_info: Dict[str, Any], dp_size: int) -> Dict[str, Any]:
+    """Compute train_iters / eval_iters / save_steps on the driver."""
+    mbs = data_info['micro_batch_size']
+    gbs = data_info['global_batch_size']
+    step_batch_size = mbs * dp_size
+    num_gen = data_info.get('num_generations', 1)
+    train_ds = data_info.get('train_dataset')
+    val_ds = data_info.get('val_dataset')
+    train_len = len(train_ds) if train_ds is not None and hasattr(train_ds, '__len__') else 0
+    val_len = len(val_ds) if val_ds is not None and hasattr(val_ds, '__len__') else 0
+
+    result: Dict[str, Any] = {}
+
+    if data_info.get('save_strategy') == 'epoch' and train_len > 0:
+        ds_sample = train_len // step_batch_size * step_batch_size * num_gen
+        result['save_steps'] = ds_sample // gbs
+        result['eval_steps'] = result['save_steps']
+
+    train_iters = data_info.get('train_iters')
+    if data_info.get('num_train_epochs') is not None and train_len > 0:
+        ds_sample = train_len // step_batch_size * step_batch_size * num_gen
+        train_iters = ds_sample * data_info['num_train_epochs'] // gbs
+    result['train_iters'] = train_iters
+
+    eval_iters = data_info.get('eval_iters', -1)
+    if eval_iters is not None and eval_iters < 0:
+        if val_len == 0:
+            eval_iters = 0
+        else:
+            ds_sample = val_len // step_batch_size * step_batch_size * num_gen
+            eval_iters = max(ds_sample // gbs, 1)
+    if val_len > 0 and val_len < step_batch_size:
+        eval_iters = 0
+    result['eval_iters'] = eval_iters or 0
+
+    return result
+
+
+def extract_iteration(step_results) -> int:
+    """Read the canonical iteration off ``WorkerGroup.execute`` results."""
+    if not step_results:
+        return 0
+    for r in step_results:
+        if isinstance(r, dict) and 'iteration' in r:
+            return int(r['iteration'])
+    return 0

swift/ray/megatron/gkd_trainer.py

@@ -0,0 +1,380 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Driver-side GKD trainer for Ray-based Megatron training."""
+from __future__ import annotations
+
+import copy
+import random
+import ray
+import torch
+from contextlib import contextmanager
+from typing import List
+
+from swift.infer_engine.protocol import RequestConfig, RolloutOutput
+from swift.rl_core.data import GKDSample
+from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput
+from swift.rlhf_trainers.utils import parse_prompt_logprobs
+from swift.rollout import MultiTurnScheduler, invoke_async_hook, multi_turns, run_multi_turn
+from swift.utils import get_logger, remove_response
+from .base_trainer import BaseRayTrainer
+from .driver_utils import extract_iteration
+from .worker_group import DPDispatchedDict
+
+logger = get_logger()
+
+
+class GKDTrainer(BaseRayTrainer):
+
+    def _prepare_state(self) -> None:
+        super()._prepare_state()
+        args = self.args
+
+        self.sft_alpha = args.sft_alpha
+        self.gkd_logits_topk = args.gkd_logits_topk
+        # GKD on-policy schedule: each step is on-policy (student generates) with
+        # probability ``lmbda``; otherwise off-policy (distill on dataset responses).
+        self.lmbda = args.lmbda
+        self._data_source_rng = random.Random(getattr(args, 'seed', 42))
+
+        # GKD generates exactly one completion per prompt (on-policy student generation),
+        # so num_generations is always 1 here regardless of the (GRPO-oriented) default.
+        self._data_info['num_generations'] = 1
+        self._teacher_model_dir = getattr(args, 'teacher_model_dir', None) or args.teacher_model
+        self._teacher_model_server = args.teacher_model_server
+        self._teacher_use_disable_adapter = args._teacher_use_disable_adapter
+        if self._teacher_use_disable_adapter:
+            self._teacher_model_dir = None
+
+        if self._teacher_model_server and not self.teacher_replicas:
+            raise NotImplementedError('teacher_model_server is not yet supported in the Ray pipeline. '
+                                      'Use teacher_model (colocated) or teacher replicas (teacher.gpus > 0) instead.')
+
+        vp_size = getattr(args, 'virtual_pipeline_model_parallel_size', None)
+        assert vp_size is None or vp_size == 1, \
+            'Ray GKD does not support VPP (virtual_pipeline_model_parallel_size > 1).'
+
+        # truncation_strategy='delete': resample prompts whose encode fails (over max_length).
+        self.truncation_strategy = args.truncation_strategy
+        self.max_completion_length = args.max_completion_length
+        self._max_resample_rounds = getattr(args, 'max_resample_times', 3)
+        self._needs_resample_iterator = self.truncation_strategy == 'delete'
+
+        self._prepare_multi_turn()
+
+    def _prepare_multi_turn(self) -> None:
+        args = self.args
+        self._multi_turn_scheduler: MultiTurnScheduler | None = None
+        self._max_turns: int | None = getattr(args, 'max_turns', None)
+        self._enable_server_multi_turn = False
+
+        scheduler_cfg = getattr(args, 'multi_turn_scheduler', None)
+        if not scheduler_cfg:
+            return
+        if isinstance(scheduler_cfg, str):
+            if scheduler_cfg not in multi_turns:
+                raise ValueError(f'Unknown multi_turn_scheduler: {scheduler_cfg!r}; '
+                                 f'available: {list(multi_turns)}')
+            scheduler_kwargs = {'max_turns': self._max_turns}
+            tokenizer = getattr(getattr(self, 'template', None), 'tokenizer', None)
+            if tokenizer is not None:
+                scheduler_kwargs['tokenizer'] = tokenizer
+            gym_env = getattr(args, 'gym_env', None)
+            if gym_env is not None:
+                scheduler_kwargs['gym_env'] = gym_env
+            self._multi_turn_scheduler = multi_turns[scheduler_cfg](**scheduler_kwargs)
+        else:
+            assert isinstance(scheduler_cfg, MultiTurnScheduler)
+            self._multi_turn_scheduler = scheduler_cfg
+
+    def _train_loop(self, tg, train_iters, iteration):
+        ckpt = self.ckpt_manager
+        spg = self._steps_per_generation
+
+        # Initialize colocated teacher if configured (skip for self-distillation via disable_adapter)
+        if self._teacher_model_dir and not self._teacher_model_server and not self._teacher_use_disable_adapter:
+            tg.execute('init_teacher_model', self._teacher_model_dir)
+            logger.info('Colocated teacher model initialized from %s', self._teacher_model_dir)
+
+        while iteration < train_iters:
+            # One generation (a single data_source) feeds ``spg`` training steps.
+            prompt_batch = next(self._data_iter)
+            if self.truncation_strategy == 'delete':
+                prompt_batch = self._resample_failed_prompts(prompt_batch, strip_response=False)
+            data_source = self._determine_data_source()
+
+            if data_source == DataSource.STUDENT:
+                # On-policy: sync the latest weights to the rollout engine and generate.
+                ckpt.sync_weights(merge_and_sync=True)
+                with self._generation_context(tg, ckpt):
+                    rollout_batch = self._expand_for_generation(prompt_batch)
+                    completions = self._generate(rollout_batch)
+                    gkd_samples = self._postprocess_rollout(rollout_batch, completions)
+                source_items = gkd_samples
+            else:
+                # Off-policy (lmbda<1): distill on the dataset's ground-truth responses,
+                # no generation and no weight sync to the rollout engine.
+                gkd_samples = [GKDSample.from_row(item) for item in prompt_batch]
+                source_items = gkd_samples
+
+            self._maybe_log_completions(gkd_samples if data_source == DataSource.STUDENT else None, gen_step=iteration)
+
+            # Split one generation into ``spg`` chunks; each chunk is one training step
+            # (same data_source). spg=1 degenerates to a single chunk == the whole batch.
+            # n == global_batch_size * spg: the driver dataloader uses drop_last=True + a
+            # cyclic iterator (see _setup_dataloader) and GKD uses num_generations=1, so n is
+            # always an exact multiple of spg and the spg chunks tile source_items with no
+            # remainder. (max(., 1) only guards the impossible spg > n case.)
+            n = len(source_items)
+            chunk_size = max(n // spg, 1)
+            for step_idx in range(spg):
+                if iteration >= train_iters:
+                    break
+                chunk = source_items[step_idx * chunk_size:(step_idx + 1) * chunk_size]
+                if not chunk:
+                    break
+                samples = self._encode_rollout_batch(chunk)
+
+                use_colocated_teacher = self._teacher_use_disable_adapter or (self._teacher_model_dir
+                                                                              and not self._teacher_model_server)
+                if self.teacher_replicas:
+                    if data_source != DataSource.STUDENT:
+                        raise NotImplementedError('Teacher replicas currently require on-policy generation (lmbda=1). '
+                                                  'Use a colocated teacher_model for lmbda<1 (off-policy) training.')
+                    self._fetch_teacher_from_replicas(chunk, samples)
+
+                # Driver collates the student (and, for the colocated path, the teacher view)
+                # micro-batches; the worker only runs prepare_batch (PP/CP slice) + forward.
+                dispatch = self._collate_for_workers_gkd(tg, samples, data_source, with_teacher=use_colocated_teacher)
+                if use_colocated_teacher:
+                    # Teacher forwards on the worker (CP slicing keeps each rank's shard
+                    # aligned) and caches per-micro-batch; train_step attaches the cache.
+                    tg.compute_teacher_logits(dispatch)
+                results = tg.train_step(dispatch)
+                iteration = extract_iteration(results)
+
+        return iteration
+
+    def _determine_data_source(self):
+        """Pick the data source for this step (GKD on/off-policy schedule).
+
+        With probability ``lmbda`` the step is on-policy (the student generates the
+        response); otherwise it is off-policy and we distill on the dataset's
+        ground-truth response.
+        """
+        if self._data_source_rng.random() < self.lmbda:
+            return DataSource.STUDENT
+        return DataSource.DATASET
+
+    def _expand_for_generation(self, prompt_batch):
+        """Convert prompt dicts to GKDSample and strip response for generation."""
+        samples = [GKDSample.from_row(item) for item in prompt_batch]
+        for s in samples:
+            remove_response(s.messages)
+        return samples
+
+    def _generate(self, batch: List[GKDSample]) -> List[RolloutOutput]:
+        args = self.args
+        request_config = RequestConfig(
+            n=1,
+            max_tokens=args.max_completion_length,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            top_k=args.top_k,
+            stop=args.stop_words or None,
+            return_details=True,
+        )
+        # Convert samples to RolloutInferRequest at the engine boundary
+        # (same pattern as GRPO Ray trainer).
+        requests = [s.to_infer_request() for s in batch]
+
+        if self._multi_turn_scheduler is not None and not self._enable_server_multi_turn:
+            # Mode A: driver-side trainer loop. run_multi_turn mutates `messages`
+            # in place on RolloutInferRequest objects.
+            invoke_async_hook(self._multi_turn_scheduler.on_trajectory_start(requests))
+            first_turn = [
+                RolloutOutput(response=resp) for resp in self._distribute_to_replicas(requests, request_config)
+            ]
+            return run_multi_turn(
+                requests=requests,
+                first_turn_outputs=first_turn,
+                scheduler=self._multi_turn_scheduler,
+                rollout_fn=lambda reqs, cfg:
+                [RolloutOutput(response=resp) for resp in self._distribute_to_replicas(reqs, cfg)],
+                request_config=request_config,
+                max_turns=self._max_turns,
+            )
+
+        completions = self._distribute_to_replicas(requests, request_config)
+        return [RolloutOutput(response=resp) for resp in completions]
+
+    def _postprocess_rollout(self, samples, outputs):
+        """Merge rollout outputs back onto GKDSample (deepcopy to match HF path)."""
+        results = []
+        for sample, output in zip(samples, outputs):
+            if output is not None:
+                sample = copy.deepcopy(sample)
+                sample.apply_rollout_output(rollout_output=output)
+            results.append(sample)
+        return results
+
+    @contextmanager
+    def _extended_max_length(self):
+        """Temporarily extend template.max_length by max_completion_length so the
+        prompt+response (token-in-token-out) encodes without truncation.
+        """
+        template = self.template
+        original = template.max_length
+        template.max_length = original + self.args.max_completion_length
+        try:
+            yield
+        finally:
+            template.max_length = original
+
+    def _encode_rollout_batch(self, samples: List[GKDSample]):
+        """Encode samples into per-sample worker payloads.
+
+        Uses the shared ``encode_gkd_samples`` helper (same as HF / Megatron GKD)
+        so OPSD logic is fully encapsulated. Returns per-sample payloads with
+        ``encoded`` (student) and optionally ``teacher_encoded`` (OPSD teacher view).
+        """
+        from swift.rlhf_trainers.gkd_helpers import encode_gkd_samples
+        template = self.template
+        with self._extended_max_length():
+            student_list, teacher_list, has_opsd = encode_gkd_samples(samples, template)
+        result = []
+        for i in range(len(samples)):
+            payload = {'encoded': student_list[i]}
+            if has_opsd:
+                payload['teacher_encoded'] = teacher_list[i]
+            result.append(payload)
+        return result
+
+    def _collate_for_workers_gkd(self, tg, samples: List[dict], data_source, *, with_teacher: bool):
+        """Driver-side GKD collate: ``List[payload-dict]`` -> ``{dp_rank: [model_inputs]}``.
+
+        Mirrors the non-Ray GKD ``_encode_samples`` (data_collator on the rank, teacher
+        forward later via prepare_batch). Each micro-batch ``model_inputs`` carries:
+        - student forward tensors (``template.data_collator`` of ``encoded``),
+        - ``data_source`` (SFT gating in loss_func),
+        - ``teacher_model_inputs`` (colocated path): the collated teacher VIEW for the
+          worker teacher forward — OPSD uses ``teacher_encoded``, else ``encoded``,
+        - ``teacher_output`` (teacher-replicas path): the per-sample TeacherOutputs
+          (already on each sample) collated into one batched TeacherOutput.
+        """
+        from swift.megatron.utils import get_padding_to
+        from .megatron_worker import MegatronWorker
+
+        template = self.template
+        padding_to = self._padding_to if self._padding_to is not None else get_padding_to(self.args)
+        dp_size = tg.dp_size
+        mbs = int(self.args.micro_batch_size)
+        n = len(samples)
+        if n % dp_size != 0:
+            raise ValueError(f'_collate_for_workers_gkd: batch size {n} not divisible by dp_size {dp_size}.')
+        shard_size = n // dp_size
+
+        dispatch = DPDispatchedDict()
+        for dp_rank in range(dp_size):
+            shard = samples[dp_rank * shard_size:(dp_rank + 1) * shard_size]
+            micro_batches = []
+            for i in range(0, len(shard), mbs):
+                chunk = shard[i:i + mbs]
+                model_inputs = template.data_collator([s['encoded'] for s in chunk], padding_to=padding_to)
+                model_inputs['data_source'] = data_source
+                if with_teacher:
+                    has_opsd = chunk[0].get('teacher_encoded') is not None
+                    key = 'teacher_encoded' if has_opsd else 'encoded'
+                    model_inputs['teacher_model_inputs'] = template.data_collator(
+                        batch=[s[key] for s in chunk], padding_to=padding_to)
+                elif chunk[0].get('teacher_output') is not None:
+                    # Teacher-replicas path: per-sample TeacherOutputs collated on the driver
+                    # (pure tensor ops). The teacher seq length differs from the student under
+                    # OPSD, so align by mask (is_opsd) rather than padding to the student length.
+                    if getattr(self.args, 'context_parallel_size', 1) > 1:
+                        raise ValueError('Standalone teacher replicas (teacher.gpus > 0) do not support '
+                                         'context_parallel_size > 1: per-sample teacher token-logprobs are built '
+                                         'from raw sequence lengths and cannot be CP-sharded to align with the '
+                                         'student. Use a colocated teacher_model for CP>1.')
+                    has_opsd = any(s.get('teacher_encoded') is not None for s in chunk)
+                    model_inputs['teacher_output'] = MegatronWorker._collate_teacher_outputs(
+                        [s['teacher_output'] for s in chunk],
+                        self.device,
+                        padding_free=template.padding_free,
+                        target_seq_len=model_inputs['labels'].shape[-1],
+                        is_opsd=has_opsd)
+                micro_batches.append(model_inputs)
+            dispatch[dp_rank] = micro_batches
+        return dispatch
+
+    def _fetch_teacher_from_replicas(self, gkd_samples: List[GKDSample], samples):
+        """Fetch teacher logprobs from Ray teacher replicas.
+
+        Uses to_infer_request() + teacher_messages replacement (OPSD) to build
+        unified RolloutInferRequest objects, matching HF GKD's _build_teacher_requests.
+        """
+        topk = self.gkd_logits_topk
+        assert topk is not None, 'gkd_logits_topk must be set when using teacher replicas'
+
+        requests = []
+        teacher_encodeds = []  # teacher-side encoded (OPSD) or None (non-OPSD)
+        for s, sample in zip(gkd_samples, samples):
+            req = s.to_infer_request()
+            teacher_encoded = sample.get('teacher_encoded')
+            if s.teacher_messages:
+                req.messages = s.teacher_messages
+                teacher_encodeds.append(teacher_encoded)
+            else:
+                teacher_encodeds.append(None)
+            requests.append(req)
+
+        request_config = RequestConfig(prompt_logprobs=topk, max_tokens=1, temperature=0.0)
+
+        replicas = self.teacher_replicas
+        n = len(replicas)
+        chunk_size = (len(requests) + n - 1) // n
+        refs = []
+        for i, replica in enumerate(replicas):
+            shard = requests[i * chunk_size:(i + 1) * chunk_size]
+            if not shard:
+                continue
+            refs.append(replica.generate(shard, request_config))
+        parts = ray.get(refs)
+        responses = []
+        for p in parts:
+            responses.extend(p)
+
+        for sample, response, t_encoded in zip(samples, responses, teacher_encodeds):
+            parsed = parse_prompt_logprobs(response, topk=topk)
+            encoded = t_encoded if t_encoded is not None else sample['encoded']
+            teacher_labels = t_encoded.get('labels') if t_encoded is not None else None
+            sample['teacher_output'] = self._build_per_sample_teacher_output(parsed, encoded, topk, teacher_labels)
+
+    @staticmethod
+    def _build_per_sample_teacher_output(parsed, encoded, topk, labels=None):
+        """Build a per-sample TeacherOutput from parsed prompt logprobs.
+
+        For OPSD, ``encoded`` is the teacher-side encoding and ``labels`` are
+        its labels; together they let ``extract_active`` mask-align the shared response.
+        For non-OPSD, teacher and student share the same encoding, so we fall
+        back to ``encoded['labels']`` when ``labels`` is not provided.
+        """
+        if labels is None:
+            labels = encoded.get('labels')
+        lps, ixs = parsed
+        input_ids = encoded['input_ids']
+        seq_len = len(input_ids) if isinstance(input_ids, list) else input_ids.shape[-1]
+
+        parsed_len = len(lps)
+        topk_logprobs = torch.full((seq_len, topk), float('-inf'), dtype=torch.float32)
+        topk_indices = torch.zeros(seq_len, topk, dtype=torch.long)
+        length = min(parsed_len, seq_len)
+        if length > 0:
+            topk_logprobs[:length] = torch.tensor(lps[:length], dtype=torch.float32)
+            topk_indices[:length] = torch.tensor(ixs[:length], dtype=torch.long)
+
+        kwargs = dict(topk_logprobs=topk_logprobs.unsqueeze(0), topk_indices=topk_indices.unsqueeze(0))
+        if labels is not None:
+            t_labels = labels
+            if not isinstance(t_labels, torch.Tensor):
+                t_labels = torch.tensor(t_labels, dtype=torch.long)
+            kwargs['labels'] = t_labels.unsqueeze(0) if t_labels.dim() == 1 else t_labels
+        return TeacherOutput(**kwargs)

swift/ray/megatron/grpo_trainer.py

@@ -0,0 +1,593 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+import concurrent.futures
+import copy
+import os
+import torch
+import uuid
+from typing import Any, Dict, List, Optional, Tuple
+
+from swift.dataset import RowPreprocessor
+from swift.infer_engine import RequestConfig
+from swift.infer_engine.protocol import RolloutOutput
+from swift.rl_core.advantage import (compute_advantages, compute_reward_metrics, compute_teacher_kl_per_token,
+                                     expand_advantage_to_per_token)
+from swift.rl_core.data import GRPOBatch, GRPOSample
+from swift.rl_core.grpo_algorithm import compute_std_for_dynamic_sampling, score_completions
+from swift.rlhf_trainers.gkd_helpers import (TeacherServerConfig, assemble_teacher_completion_logprobs,
+                                             build_opsd_samples, build_teacher_requests, encode_teacher_view,
+                                             fetch_teacher_parsed_by_routing, parse_teacher_model_server,
+                                             remap_teacher_logps_to_student_frame,
+                                             resolve_dynamic_opd_self_distillation)
+from swift.rlhf_trainers.utils import encode_sample, make_reward_weights, resolve_reward_funcs
+from swift.rollout import MultiTurnScheduler, invoke_async_hook, multi_turns, run_multi_turn
+from swift.utils import get_logger, remove_response
+from .base_trainer import BaseRayTrainer
+from .driver_utils import extract_iteration
+
+logger = get_logger()
+
+
+class GRPOTrainer(BaseRayTrainer):
+    """Driver-side GRPO trainer."""
+
+    def _prepare_state(self) -> None:
+        super()._prepare_state()
+        args = self.args
+
+        self.num_generations = args.num_generations
+        self.advantage_estimator = args.advantage_estimator
+        self.scale_rewards = args.scale_rewards
+        self.kl_in_reward = args.kl_in_reward
+
+        self._teacher_model_dir = getattr(args, 'teacher_model_dir', None) or args.teacher_model
+        self._teacher_model_server = getattr(args, 'teacher_model_server', None)
+        self._teacher_use_disable_adapter = getattr(args, '_teacher_use_disable_adapter', False)
+        if self._teacher_use_disable_adapter:
+            self._teacher_model_dir = None
+        teacher_explicit = bool(self._teacher_model_dir or self._teacher_model_server
+                                or self._teacher_use_disable_adapter)
+        self._has_teacher_explicit = teacher_explicit
+        self._is_dynamic_self_distillation = resolve_dynamic_opd_self_distillation(
+            has_teacher_explicit=teacher_explicit,
+            is_self_distillation=not teacher_explicit,
+        )
+        self._has_teacher = teacher_explicit or self._is_dynamic_self_distillation
+        self.teacher_kl_coef = getattr(args, 'teacher_kl_coef', 1.0)
+        # Parse teacher_model_server: supports single URL and multi-teacher JSON.
+        self.use_teacher_api = self._teacher_model_server is not None
+        self.teacher_configs: List[TeacherServerConfig] = []
+        self.teacher_clients = []
+        if self.use_teacher_api:
+            self.teacher_configs = parse_teacher_model_server(self._teacher_model_server)
+            from swift.rlhf_trainers.vllm_client import VLLMInferClient
+            self.teacher_clients = [VLLMInferClient(base_urls=[cfg.url]) for cfg in self.teacher_configs]
+
+        self._prepare_rewards()
+        self._prepare_multi_turn()
+
+        # Ray supports router replay only in R3 (rollout records routed_experts, the driver
+        # collates them into the train micro-batch). R2 records during the policy logps
+        # forward, which — with driver-side collation — would not flow back into the train
+        # batch; reject it explicitly (mirrors pipeline.py's R3-only rollout wiring).
+        router_mode = getattr(args, 'router_replay_mode', 'disabled')
+        if router_mode not in ('disabled', 'R3'):
+            raise ValueError(f"Ray Megatron GRPO supports router_replay_mode in {{'disabled', 'R3'}}, "
+                             f'got {router_mode!r}. Use R3 (rollout-recorded routing) for the Ray pipeline.')
+
+        # DAPO dynamic_sample + truncation_strategy='delete' resampling (driver-side).
+        self.dynamic_sample = getattr(args, 'dynamic_sample', False)
+        self.max_resample_times = getattr(args, 'max_resample_times', 3)
+        self.truncation_strategy = args.truncation_strategy
+        self._max_resample_rounds = getattr(args, 'max_resample_times', 10)
+        self._needs_resample_iterator = self.dynamic_sample or self.truncation_strategy == 'delete'
+
+    def _prepare_multi_turn(self) -> None:
+        """Configure driver-side multi-turn scheduler (Mode A only).
+
+        Mode B (server-side scheduler) is intentionally not enabled here because
+        :class:`VllmServer.launch_server` does not yet wrap the engine via
+        ``get_rollout_engine_type`` — the server-side scheduler plumbing is a
+        separate cross-process change.  When that lands, set
+        ``self._enable_server_multi_turn`` from a new
+        ``RolloutReplica.get_engine_type()`` passthrough.
+        """
+        args = self.args
+        self._multi_turn_scheduler: Optional[MultiTurnScheduler] = None
+        self._max_turns: Optional[int] = getattr(args, 'max_turns', None)
+        self._enable_server_multi_turn = False
+
+        scheduler_cfg = getattr(args, 'multi_turn_scheduler', None)
+        if not scheduler_cfg:
+            return
+        if isinstance(scheduler_cfg, str):
+            if scheduler_cfg not in multi_turns:
+                raise ValueError(f'Unknown multi_turn_scheduler: {scheduler_cfg!r}; '
+                                 f'available: {list(multi_turns)}')
+            scheduler_kwargs = {'max_turns': self._max_turns}
+            gym_env = getattr(args, 'gym_env', None)
+            if gym_env is not None:
+                scheduler_kwargs['gym_env'] = gym_env
+            self._multi_turn_scheduler = multi_turns[scheduler_cfg](**scheduler_kwargs)
+        else:
+            assert isinstance(scheduler_cfg, MultiTurnScheduler)
+            self._multi_turn_scheduler = scheduler_cfg
+
+    def _prepare_rewards(self):
+        args = self.args
+        reward_funcs_cfg = (args.reward_funcs or []).copy()
+        if not isinstance(reward_funcs_cfg, list):
+            reward_funcs_cfg = [reward_funcs_cfg]
+
+        self.reward_funcs, self.reward_func_names = resolve_reward_funcs(reward_funcs_cfg, args=args)
+
+        # use_gym_env: gym total_reward is appended as an extra reward column so it can
+        # blend with reward_funcs via reward_weights.  When reward_funcs is empty, it becomes
+        # the single reward source.
+        self.use_gym_env = bool(getattr(args, 'use_gym_env', False))
+        if self.use_gym_env:
+            self.reward_func_names.append('gym_reward')
+
+        self.reward_weights = make_reward_weights(args.reward_weights, len(self.reward_func_names), self.device)
+        self.reward_model_plugins = [None] * len(self.reward_funcs)
+
+        if not self.reward_funcs and not self.use_gym_env and not getattr(self, '_has_teacher', False):
+            raise ValueError('GRPOTrainer: no reward functions configured '
+                             '(or pass use_gym_env: true / a teacher for OPD-RL)')
+
+    def _get_request_config(self):
+        """Build a RequestConfig for rollout generation."""
+        from swift.infer_engine.protocol import RequestConfig
+        args = self.args
+        return RequestConfig(
+            n=1,
+            max_tokens=args.max_completion_length,
+            temperature=args.temperature,
+            top_p=args.top_p,
+            top_k=args.top_k,
+            repetition_penalty=args.repetition_penalty,
+            stop=args.stop_words or None,
+            return_details=True,
+            logprobs=True,
+        )
+
+    def _train_loop(self, tg, train_iters, iteration):
+        ckpt = self.ckpt_manager
+        merge_and_sync = not self.args.vllm_enable_lora
+        spg = self._steps_per_generation
+
+        # OPD-RL: load the colocated teacher once (disable_adapter self-distillation needs none).
+        if self._teacher_model_dir and not self._teacher_use_disable_adapter:
+            tg.execute('init_teacher_model', self._teacher_model_dir)
+            logger.info('OPD-RL colocated teacher model initialized from %s', self._teacher_model_dir)
+
+        while iteration < train_iters:
+            ckpt.sync_weights(merge_and_sync=merge_and_sync)
+
+            with self._generation_context(tg, ckpt):
+                prompt_batch = next(self._data_iter)
+                if self.truncation_strategy == 'delete':
+                    prompt_batch = self._resample_failed_prompts(prompt_batch)
+                rollout_batch = self.expand_for_generation(prompt_batch)
+                completions = self._generate(rollout_batch)
+                rollout_with_outputs = self._postprocess_rollout(rollout_batch, completions)
+                rewards_per_func = self.score_completions(rollout_with_outputs)
+                # DAPO dynamic sampling: drop zero-variance (std==0) prompt groups and
+                # resample fresh prompts (the rollout engine is still awake in this context).
+                if self.dynamic_sample:
+                    rollout_with_outputs, rewards_per_func = self._dynamic_sampling(rollout_with_outputs,
+                                                                                    rewards_per_func)
+
+            self._maybe_log_completions(
+                rollout_with_outputs, rewards=rewards_per_func.sum(dim=1).tolist(), gen_step=iteration)
+
+            n_samples = len(rollout_with_outputs)
+            chunk_size = n_samples // spg
+
+            all_chunks = []  # per spg step: (dispatch, flat_grpo_batches)
+            for step_idx in range(spg):
+                chunk_start = step_idx * chunk_size
+                chunk_end = chunk_start + chunk_size
+                chunk_rollout = rollout_with_outputs[chunk_start:chunk_end]
+                chunk_samples = self.encode_rollout_batch(chunk_rollout)
+
+                dispatch, grpo_batches = self._collate_for_workers(tg, chunk_samples)
+                logps_rows = tg.compute_logps(dispatch)
+                self._scatter_logps(grpo_batches, logps_rows, 'old_per_token_logps')
+                if self.beta != 0.0:
+                    ref_rows = tg.compute_ref_logps(dispatch)
+                    self._scatter_logps(grpo_batches, ref_rows, 'ref_per_token_logps')
+                # OPD-RL: teacher logp on the sampled tokens (same frame as old/ref logps).
+                # TODO(perf): When use_teacher_api, start API requests after old_logps scatter
+                # and overlap with ref_logps computation (beta != 0) to hide API latency.
+                if self._has_teacher:
+                    if self.use_teacher_api:
+                        self._compute_teacher_api_logps(chunk_samples, grpo_batches)
+                    else:
+                        self._compute_teacher_logps(tg, chunk_samples, dispatch, grpo_batches)
+                all_chunks.append((dispatch, grpo_batches))
+
+            for step_idx in range(spg):
+                if iteration >= train_iters:
+                    break
+                dispatch, grpo_batches = all_chunks[step_idx]
+                chunk_start = step_idx * chunk_size
+                chunk_end = chunk_start + chunk_size
+                chunk_rewards_pf = rewards_per_func[chunk_start:chunk_end]
+
+                kl_values = self._compute_kl_from_batches(grpo_batches) if self.beta != 0.0 else None
+
+                chunk_advantages, rewards = self.compute_advantages(chunk_rewards_pf, kl_values=kl_values)
+                self._scatter_advantages(grpo_batches, chunk_advantages)
+
+                results = tg.train_step(
+                    dispatch, extra_metrics=self._build_grpo_log_metrics(rewards, chunk_advantages, chunk_rewards_pf))
+                iteration = extract_iteration(results)
+
+        return iteration
+
+    def _compute_teacher_logps(self, tg, chunk_samples: List[GRPOSample], dispatch,
+                               grpo_batches: List[GRPOBatch]) -> None:
+        """OPD-RL: fill each micro-batch's ``teacher_per_token_logps`` (student frame).
+
+        Non-OPSD: the teacher forwards the SAME student-collated dispatch, so the rows
+        already align to the student ``completion_mask`` frame and are scattered directly.
+        OPSD: the teacher forwards its own (teacher_prompt + same response) encoding via a
+        separate dispatch, then the teacher-frame logps are remapped onto the student frame.
+        """
+        has_opsd_batch = build_opsd_samples(chunk_samples)
+        if not has_opsd_batch:
+            if not self._has_teacher_explicit:
+                return
+            teacher_rows = tg.compute_teacher_logps(dispatch)
+            self._scatter_logps(grpo_batches, teacher_rows, 'teacher_per_token_logps')
+            return
+
+        # OPSD: encode the teacher view (teacher_prompt + shared response) and dispatch separately.
+        for s in chunk_samples:
+            s.encoded = encode_teacher_view(s, self.template)
+        teacher_dispatch, teacher_grpo_batches = self._collate_for_workers(tg, chunk_samples)
+        # Restore the student encoding so the dispatched student micro-batches stay the student frame.
+        self.encode_rollout_batch(chunk_samples)
+        teacher_rows = tg.compute_teacher_logps(teacher_dispatch)
+        self._scatter_logps(teacher_grpo_batches, teacher_rows, 'teacher_per_token_logps')
+        for student_gb, teacher_gb in zip(grpo_batches, teacher_grpo_batches):
+            student_gb.teacher_per_token_logps = remap_teacher_logps_to_student_frame(
+                teacher_gb.teacher_per_token_logps.to(student_gb.completion_mask.device),
+                teacher_gb.completion_mask.to(student_gb.completion_mask.device), student_gb.completion_mask)
+
+    def _compute_teacher_api_logps(self, chunk_samples: List[GRPOSample], grpo_batches: List[GRPOBatch]) -> None:
+        """Driver-side: fetch teacher logps from API servers, scatter into GRPOBatch.
+
+        Each sample routes to exactly one teacher by tag (single teacher = all samples). Runs on
+        the driver, so there is no distributed gather: the per-teacher requests go straight to
+        ``client.infer``. Multiple teachers infer concurrently (distinct HTTP servers).
+        """
+        from swift.rlhf_trainers.utils import parse_prompt_logprobs
+
+        build_opsd_samples(chunk_samples)
+        request_config = RequestConfig(prompt_logprobs=0, max_tokens=1, temperature=0.0)
+
+        def infer(reqs, client):
+            if not reqs:  # no sample routed to this teacher: skip the empty HTTP call
+                return []
+            responses = client.infer(reqs, request_config=request_config, use_tqdm=False)
+            return [parse_prompt_logprobs(r, topk=0) for r in responses]
+
+        requests = build_teacher_requests(chunk_samples, self.template)
+        all_rti = [s.response_token_ids for s in chunk_samples]
+        parsed = fetch_teacher_parsed_by_routing(
+            chunk_samples,
+            requests,
+            self.teacher_configs,
+            self.teacher_clients,
+            gather_fn=lambda reqs: reqs,  # driver-side: no distributed gather
+            infer_fn=infer,
+            scatter_fn=lambda reqs, parsed_global: parsed_global,  # already local
+            is_main_process=True,
+            tag_key=self.args.teacher_tag_key)
+
+        offset = 0
+        for gb in grpo_batches:
+            device = gb.completion_mask.device
+            n = gb.completion_mask.shape[0]
+            teacher_out = assemble_teacher_completion_logprobs(
+                parsed[offset:offset + n], gb.completion_mask, device, response_token_ids=all_rti[offset:offset + n])
+            gb.teacher_per_token_logps = teacher_out.topk_logprobs[..., 0]
+            offset += n
+
+    @staticmethod
+    def _scatter_logps(grpo_batches: List[GRPOBatch], rows: List[Dict[str, torch.Tensor]], key: str) -> None:
+        """Stack the flat per-sample logps rows (dp_flat, sample order) back onto each
+        micro-batch's GRPOBatch as ``[B, T]`` — the same carrier non-Ray Megatron uses.
+
+        ``completion_mask`` is NOT touched here: it was built by the driver collate and the
+        worker only forwards logps, so the existing ``gb.completion_mask`` is already correct.
+        """
+        # The worker keys ``old_per_token_logps`` rows as ``per_token_logps``; ref / teacher
+        # rows carry their destination key verbatim.
+        src_key = 'per_token_logps' if key == 'old_per_token_logps' else key
+        pos = 0
+        for gb in grpo_batches:
+            b = gb.completion_mask.shape[0]
+            chunk = rows[pos:pos + b]
+            pos += b
+            setattr(gb, key, torch.stack([r[src_key] for r in chunk], dim=0))
+        assert pos == len(rows), f'_scatter_logps: consumed {pos} rows but got {len(rows)}'
+
+    @staticmethod
+    def _align_width(x: torch.Tensor, width: int) -> torch.Tensor:
+        """Truncate or right-pad the last (token) dim of ``x`` to ``width``.
+
+        A small width drift is expected (padding alignment across micro-batches), but a large
+        gap means the teacher logps are mis-shaped (e.g. a different tokenizer) and silently
+        slicing them would corrupt the per-token KL -- guard against that.
+        """
+        cur = x.shape[-1]
+        if cur == width:
+            return x
+        assert abs(cur - width) <= 8, (f'teacher logp width {cur} differs from mask width {width} by more than the '
+                                       'padding slack; teacher/student token alignment is likely broken.')
+        if cur > width:
+            return x[..., :width]
+        return torch.nn.functional.pad(x, (0, width - cur))
+
+    def _scatter_advantages(self, grpo_batches: List[GRPOBatch], advantages: torch.Tensor) -> None:
+        """Write the advantage onto each micro-batch's GRPOBatch, expanding the per-sequence base
+        advantage to per-token ``[B, T]`` so the OPD-RL signed teacher log-ratio is added per token
+        (``adv_t = base + coef * (teacher_logp - student_logp)``). ``advantages`` is ``[N]`` in sample order."""
+        pos = 0
+        kl_sum, tok_sum = 0.0, 0.0
+        for gb in grpo_batches:
+            b = gb.completion_mask.shape[0]
+            base = advantages[pos:pos + b].to(gb.completion_mask.device)
+            pos += b
+            teacher_lp = policy_lp = None
+            if self._has_teacher and gb.teacher_per_token_logps is not None and gb.old_per_token_logps is not None:
+                # teacher / old logps share the completion_mask frame (worker forwards them on the
+                # driver-collated batch); align widths defensively to the mask before computing k3.
+                T = gb.completion_mask.shape[-1]
+                teacher_lp = self._align_width(gb.teacher_per_token_logps, T).to(gb.completion_mask.device)
+                policy_lp = self._align_width(gb.old_per_token_logps, T).to(gb.completion_mask.device)
+                k3 = compute_teacher_kl_per_token(teacher_lp, policy_lp, gb.completion_mask.to(teacher_lp.dtype))
+                kl_sum += k3.sum().item()
+                tok_sum += gb.completion_mask.sum().item()
+            gb.advantages = expand_advantage_to_per_token(
+                base,
+                gb.completion_mask,
+                teacher_per_token_logps=teacher_lp,
+                policy_per_token_logps=policy_lp,
+                teacher_kl_coef=self.teacher_kl_coef if teacher_lp is not None else 0.0,
+            )
+        assert pos == advantages.shape[0], f'_scatter_advantages: wrote {pos} but got {advantages.shape[0]}'
+        # Per-token teacher KL averaged over response tokens (monitoring only; the signal is applied
+        # per-token above). Surfaced via _build_grpo_log_metrics -> worker on_log.
+        self._last_teacher_kl = (kl_sum / tok_sum) if tok_sum > 0 else None
+
+    def _build_grpo_log_metrics(self, rewards, advantages, rewards_per_func) -> Dict[str, float]:
+        """Driver-computed GRPO metrics (reward / reward_std / adv_nonzero / per-func),
+        injected into the worker megatron on_log so all logging is unified there."""
+        reward_metrics = compute_reward_metrics(
+            rewards=rewards,
+            rewards_per_func=rewards_per_func,
+            reward_func_names=self.reward_func_names,
+            num_generations=self.num_generations,
+            scale_rewards=self.scale_rewards,
+        )
+        metrics = {
+            'reward': reward_metrics.reward_mean,
+            'reward_std': reward_metrics.reward_std,
+            'frac_reward_zero_std': reward_metrics.frac_reward_zero_std,
+            'adv_nonzero': (advantages.abs() > 1e-8).float().mean().item(),
+        }
+        if getattr(self, '_last_teacher_kl', None) is not None:
+            metrics['teacher_kl'] = self._last_teacher_kl
+        # Flatten per-function metrics into scalar values the worker can inject.
+        for name in self.reward_func_names:
+            metrics[name] = reward_metrics.per_func_mean[name]
+            metrics[f'rewards/{name}/std'] = reward_metrics.per_func_std[name]
+        return metrics
+
+    def _generate(self, samples: List[GRPOSample]) -> List[RolloutOutput]:
+        """Run a prompt batch through rollout replicas.
+
+        Returns ``List[RolloutOutput]`` (one per request). For Mode A
+        (driver-side multi-turn) the per-turn ``response_token_ids`` and
+        ``response_loss_mask`` are accumulated inside each ``RolloutOutput``.
+        """
+        request_config = self._get_request_config()
+        # Convert samples to RolloutInferRequest at the engine boundary.
+        requests = [s.to_infer_request() for s in samples]
+
+        if self._multi_turn_scheduler is not None and not self._enable_server_multi_turn:
+            # Mode A: driver-side trainer loop. run_multi_turn mutates `messages`
+            # in place on RolloutInferRequest objects.
+            invoke_async_hook(self._multi_turn_scheduler.on_trajectory_start(requests))
+            first_turn = [
+                RolloutOutput(response=resp) for resp in self._distribute_to_replicas(requests, request_config)
+            ]
+            return run_multi_turn(
+                requests=requests,
+                first_turn_outputs=first_turn,
+                scheduler=self._multi_turn_scheduler,
+                rollout_fn=lambda reqs, cfg:
+                [RolloutOutput(response=resp) for resp in self._distribute_to_replicas(reqs, cfg)],
+                request_config=request_config,
+                max_turns=self._max_turns,
+            )
+
+        # Mode B (server-side multi-turn, currently disabled) + single-turn share this path.
+        completions = self._distribute_to_replicas(requests, request_config)
+        assert len(completions) == len(requests)
+        return [RolloutOutput(response=resp) for resp in completions]
+
+    def _postprocess_rollout(self, samples: List[GRPOSample], outputs: List[RolloutOutput]) -> List[GRPOSample]:
+        if not outputs:
+            return list(samples)
+        if len(outputs) != len(samples):
+            raise RuntimeError(f'GRPOTrainer: rollout produced {len(outputs)} completions '
+                               f'for {len(samples)} samples; shapes mismatch.')
+        results = []
+        for sample, output in zip(samples, outputs):
+            if output is None:
+                results.append(sample)
+                continue
+            sample = copy.deepcopy(sample)
+            sample.apply_rollout_output(rollout_output=output)
+            results.append(sample)
+        return results
+
+    def expand_for_generation(
+        self,
+        prompt_batch: List[Dict[str, Any]],
+    ) -> List[GRPOSample]:
+        num_gen = self.num_generations
+        samples: List[GRPOSample] = []
+        for item in prompt_batch:
+            base = GRPOSample.from_row(item)
+            if base.messages:
+                remove_response(base.messages)
+            base.request_id = uuid.uuid4().hex
+            samples.append(base)
+            for _ in range(num_gen - 1):
+                dup = copy.deepcopy(base)
+                dup.request_id = uuid.uuid4().hex
+                samples.append(dup)
+        return samples
+
+    def score_completions(
+        self,
+        samples: List[GRPOSample],
+    ) -> torch.Tensor:
+        """Score completions using the backend-agnostic shared helper.
+
+        The driver-side Ray trainer already sees the global prompt/completion
+        batch, so no distributed gather is performed here.
+        """
+        return score_completions(
+            samples,
+            reward_funcs=self.reward_funcs,
+            reward_model_plugins=self.reward_model_plugins,
+            use_gym_env=self.use_gym_env,
+            device=self.device,
+        )
+
+    def compute_advantages(
+        self,
+        rewards_per_func: torch.Tensor,
+        kl_values: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return the per-sequence base advantage and rewards, both shaped [N] (N = B * num_gen).
+
+        The driver already holds every completion of each group, so no gather is needed before
+        calling the pure advantage function. The OPD-RL teacher signal is applied per-token later in
+        ``_scatter_advantages`` (see ``expand_advantage_to_per_token``).
+        """
+        return compute_advantages(
+            rewards_per_func=rewards_per_func,
+            reward_weights=self.reward_weights,
+            num_generations=self.num_generations,
+            advantage_estimator=self.advantage_estimator,
+            scale_rewards=self.scale_rewards,
+            kl_in_reward=self.kl_in_reward,
+            beta=self.beta,
+            kl_values=kl_values,
+        )
+
+    def _dynamic_sampling(
+        self,
+        samples: List[GRPOSample],
+        rewards_per_func: torch.Tensor,
+    ) -> Tuple[List[GRPOSample], torch.Tensor]:
+        num_gen = self.num_generations
+        target = len(samples)
+        valid_samples: List[GRPOSample] = []
+        valid_rewards: List[torch.Tensor] = []
+        cur_samples, cur_rewards = samples, rewards_per_func
+
+        for resample_count in range(self.max_resample_times + 1):
+            grouped_std = compute_std_for_dynamic_sampling(
+                cur_rewards,
+                self.reward_weights,
+                num_gen,
+            )
+            keep_mask = grouped_std > 0
+            for i in range(len(cur_samples)):
+                if keep_mask[i]:
+                    valid_samples.append(cur_samples[i])
+                    valid_rewards.append(cur_rewards[i])
+            logger.info('dynamic_sample round %d: kept %d/%d (std>0), accumulated %d/%d', resample_count,
+                        int(keep_mask.sum().item()), len(cur_samples), len(valid_samples), target)
+            if len(valid_samples) >= target or resample_count >= self.max_resample_times:
+                break
+            prompt_batch = next(self._resample_iter)
+            if self.truncation_strategy == 'delete':
+                prompt_batch = self._resample_failed_prompts(prompt_batch)
+            cur_samples = self.expand_for_generation(prompt_batch)
+            comp = self._generate(cur_samples)
+            cur_samples = self._postprocess_rollout(cur_samples, comp)
+            cur_rewards = self.score_completions(cur_samples)
+
+        if len(valid_samples) >= target:
+            return valid_samples[:target], torch.stack(valid_rewards[:target])
+        logger.warning('dynamic_sample: only %d/%d std>0 samples after %d retries; using original batch.',
+                       len(valid_samples), target, self.max_resample_times)
+        return samples, rewards_per_func
+
+    def _batch_encode_parallel(self, infer_requests: List[Dict[str, Any]], strict: bool):
+        max_workers = max(min(32, os.cpu_count() or 1, len(infer_requests)), 1)
+        encoded: List[Dict[str, Any]] = []
+        errors: List[Tuple[int, Exception]] = []
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as ex:
+            futures = [ex.submit(self.template.encode, req, return_length=True) for req in infer_requests]
+            concurrent.futures.wait(futures)
+            for i, fut in enumerate(futures):
+                try:
+                    encoded.append(fut.result())
+                except Exception as e:  # pragma: no cover
+                    if strict:
+                        raise
+                    errors.append((i, e))
+        return encoded, errors
+
+    def encode_rollout_batch(
+        self,
+        samples: List[GRPOSample],
+    ) -> List[GRPOSample]:
+        """Encode each sample in place and return the samples.
+
+        This is the driver → worker boundary: the same ``GRPOSample`` objects
+        cross the RPC (``tg.compute_logps`` / ``tg.train_step``) — the worker
+        feeds them straight to ``collate_to_grpo_micro_batch`` (the shared collate
+        used by HF / Megatron). Uses the shared ``encode_sample`` helper so bug
+        fixes to loss_mask / non_thinking_prefix propagate across all backends.
+        """
+        for sample in samples:
+            encoded = encode_sample(sample, self.template)
+            encoded.pop('_extra_kwargs', None)
+            sample.encoded = encoded
+        return samples
+
+    def _compute_kl_from_batches(self, grpo_batches: List[GRPOBatch]) -> Optional[torch.Tensor]:
+        """Per-sample KL = sum_t (old_lp - ref_lp) * completion_mask, in sample order.
+
+        Reads the [B, T] logps/mask off each micro-batch GRPOBatch (the unified logps
+        carrier), so the driver-side DAPO ``kl_in_reward`` penalty matches non-Ray.
+        """
+        if not (self.kl_in_reward and self.beta != 0.0):
+            return None
+        kl_values = []
+        for gb in grpo_batches:
+            old_lp, ref_lp, mask = gb.old_per_token_logps, gb.ref_per_token_logps, gb.completion_mask
+            if old_lp is None or ref_lp is None or mask is None:
+                return None
+            old_lp = old_lp.to(self.device)
+            ref_lp = ref_lp.to(self.device)
+            mask = mask.to(self.device)
+            width = min(old_lp.shape[-1], ref_lp.shape[-1], mask.shape[-1])
+            per_token_kl = (old_lp[..., :width] - ref_lp[..., :width]) * mask[..., :width].to(old_lp.dtype)
+            kl_values.append(per_token_kl.sum(dim=-1))  # [B] per-sample
+        if not kl_values:
+            return None
+        return torch.cat(kl_values, dim=0)

swift/ray/megatron/loss/__init__.py

@@ -0,0 +1,2 @@
+from .base import Loss
+from .grpo import GRPOLoss

swift/ray/megatron/loss/base.py

@@ -0,0 +1,55 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Abstract base for worker-side loss computation.
+
+A ``Loss`` subclass defines ``forward_step`` and ``loss_func`` -- the two
+methods that Megatron's pipeline-parallel scheduler calls during
+training.  Users who want to customise loss computation only need to
+subclass ``Loss`` and override these two methods; no understanding of
+the internal Megatron trainer is required.
+
+Example::
+
+    class MyLoss(Loss):
+        def __init__(self, args):
+            self.label_smoothing = args.label_smoothing
+
+        def forward_step(self, data_iterator, model):
+            batch = next(data_iterator)
+            output = model(batch['input_ids'], ...)
+            return output, partial(self.loss_func, labels=batch['labels'])
+
+        def loss_func(self, output_tensor, *, labels):
+            loss = F.cross_entropy(output_tensor, labels)
+            return loss, {'loss': loss.item()}
+
+Then register it::
+
+    register_ray_trainer('my_algo', trainer='...MyDriver', loss='...MyLoss')
+"""
+from abc import ABC, abstractmethod
+
+
+class Loss(ABC):
+    """Abstract base for worker-side loss / forward computation.
+
+    Mirrors the two methods that Megatron's PP scheduler calls:
+    ``forward_step(data_iterator, model) -> (output, loss_fn)``
+    and ``loss_func(output_tensor, **ctx) -> (loss, metrics)``.
+
+    Subclasses may wrap an existing trainer via composition for code
+    reuse (see ``GRPOLoss``) or implement these from scratch.
+    """
+
+    @abstractmethod
+    def forward_step(self, data_iterator, model):
+        """Run a single forward micro-batch through *model*.
+
+        Returns ``(output_tensor, partial(self.loss_func, ...))``.
+        """
+
+    @abstractmethod
+    def loss_func(self, output_tensor, **kwargs):
+        """Compute scalar loss + metrics from ``output_tensor``.
+
+        Returns ``(loss, metric_dict)``.
+        """

swift/ray/megatron/loss/gkd.py

@@ -0,0 +1,141 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""GKD loss for Ray-based Megatron training."""
+from __future__ import annotations
+
+import torch
+from functools import partial
+from megatron.core import mpu
+from typing import Any, Dict, List, Optional
+
+from swift.megatron.trainers.gkd_utils import cp_reduce, tp_gather_topk, vocab_parallel_topk
+from swift.megatron.trainers.utils import prepare_batch
+from swift.megatron.trainers.vocab_parallel_utils import vocab_parallel_kl_div, vocab_parallel_log_softmax
+from swift.megatron.utils import forward_step_helper
+from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss
+from swift.utils import get_current_device, to_device
+from .base import Loss
+
+
+class GKDLoss(Loss):
+    """GKD loss: JSD between student and teacher + optional SFT loss."""
+
+    def __init__(self, args):
+        self.args = args
+        self.beta = getattr(args, 'beta', 0.5)
+        self.temperature = getattr(args, 'temperature', 1.0)
+        self.sft_alpha = getattr(args, 'sft_alpha', 0.0)
+
+    def forward_step(self, data_iterator, model):
+        data = next(data_iterator)
+        teacher_output = data.pop('teacher_output', TeacherOutput())
+        data_source = data.pop('data_source', None)
+        data.pop('grpo_batch', None)  # RL signals packed in GRPOBatch (not used by GKD loss)
+        data = prepare_batch(self.args, data)
+
+        data.pop('loss_scale', None)
+        data.pop('routed_experts', None)  # MoE routing replay data (not a model forward kwarg)
+        labels = data.pop('labels', None)
+
+        # data is now clean model forward kwargs (template.encode guarantees this)
+        student_output = model(**data)
+        return student_output, partial(
+            self.loss_func,
+            labels=labels,
+            teacher_output=teacher_output,
+            data_source=data_source,
+            model=model,
+        )
+
+    # ------------------------------------------------------------------
+    # Teacher logit computation (called from MegatronWorker thin wrapper)
+    # ------------------------------------------------------------------
+
+    def compute_teacher_logits(
+        self,
+        teacher_model: torch.nn.Module,
+        teacher_micro_batches: List[Dict[str, Any]],
+        args,
+    ) -> List[TeacherOutput]:
+        gkd_logits_topk = getattr(args, 'gkd_logits_topk', None)
+        device = get_current_device()
+        outputs: List[TeacherOutput] = []
+        with torch.no_grad():
+            for teacher_model_inputs in teacher_micro_batches:
+                collated = to_device(dict(teacher_model_inputs), device)
+                teacher_data = prepare_batch(args, collated)
+                teacher_data.pop('loss_scale', None)
+                # Labels are the teacher's (OPSD) or == student's (non-OPSD); prepare_batch
+                # shifts them so extract_active mask-aligns the shared response.
+                labels = teacher_data.pop('labels', None)
+                teacher_logits = forward_step_helper(teacher_model, teacher_data)
+                if teacher_logits is None:
+                    # PP non-last stage: no logits; placeholder keeps micro-batch alignment.
+                    outputs.append(TeacherOutput())
+                    continue
+                teacher_logits = teacher_logits.detach()
+                if gkd_logits_topk is not None:
+                    topk_logits, topk_indices = vocab_parallel_topk(teacher_logits, k=gkd_logits_topk)
+                    outputs.append(TeacherOutput(topk_logprobs=topk_logits, topk_indices=topk_indices, labels=labels))
+                else:
+                    outputs.append(TeacherOutput(full_logits=teacher_logits, labels=labels))
+                del collated
+        return outputs
+
+    # ------------------------------------------------------------------
+    # Loss computation
+    # ------------------------------------------------------------------
+
+    def loss_func(self, output_tensor, *, labels, teacher_output, data_source=None, model=None):
+        args = self.args
+        student_logits = output_tensor
+
+        jsd_total, jsd_num_valid = gkd_loss(
+            student_logits,
+            teacher_output,
+            labels,
+            self.beta,
+            self.temperature,
+            gather_fn=tp_gather_topk,
+            log_softmax_fn=vocab_parallel_log_softmax,
+            kl_div_fn=vocab_parallel_kl_div)
+        jsd_loss_val = cp_reduce(jsd_total, jsd_num_valid, cp_size=args.context_parallel_size)
+
+        loss = jsd_loss_val
+        sft_loss = None
+        # SFT loss only applies to ground-truth (dataset) responses; skip it on
+        # student-generated (on-policy) responses, matching the non-ray GKD trainer.
+        if self.sft_alpha > 0 and data_source != DataSource.STUDENT:
+            # Vocab-parallel-aware SFT loss: route through ``model.compute_language_model_loss``
+            # (mirrors the non-ray GKD trainer). Naive ``torch.nn.functional.cross_entropy``
+            # would index out of bounds on the TP-sharded local vocab when TP>1.
+            assert model is not None, 'sft_alpha>0 requires the model handle from forward_step'
+            unwrapped = model
+            while hasattr(unwrapped, 'module'):
+                unwrapped = unwrapped.module
+            if hasattr(unwrapped, 'language_model'):
+                unwrapped = unwrapped.language_model
+            logits_sbv = student_logits.transpose(0, 1).contiguous()
+            per_token_loss = unwrapped.compute_language_model_loss(labels, logits_sbv)
+            loss_mask = labels != -100
+            sft_loss_sum = (per_token_loss * loss_mask).sum()
+            sft_loss_count = loss_mask.sum().float()
+            if args.context_parallel_size > 1:
+                sft_stats = torch.stack([sft_loss_sum, sft_loss_count])
+                torch.distributed.all_reduce(
+                    sft_stats, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group())
+                sft_loss_sum, sft_loss_count = sft_stats[0], sft_stats[1]
+            sft_loss = sft_loss_sum / sft_loss_count if sft_loss_count > 0 else sft_loss_sum * 0
+            loss = loss + self.sft_alpha * sft_loss
+
+        metric = {'loss': loss.detach().clone()}
+        if sft_loss is not None:
+            metric['jsd_loss'] = jsd_loss_val.detach().clone()
+            metric['sft_loss'] = sft_loss.detach().clone()
+
+        dp_group = mpu.get_data_parallel_group()
+        reporting = torch.stack(list(metric.values()), dim=0)
+        torch.distributed.all_reduce(reporting, torch.distributed.ReduceOp.AVG, group=dp_group)
+        metric = {k: reporting[i] for i, k in enumerate(metric.keys())}
+
+        loss = loss / mpu.get_context_parallel_world_size()
+        return loss, metric

swift/ray/megatron/loss/grpo.py

@@ -0,0 +1,75 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""GRPO loss for Ray-based Megatron training.
+
+``GRPOLoss`` reuses ``MegatronGRPOTrainer.forward_step`` and
+``loss_func`` via a dummy trainer instance (created with ``__new__``
+to skip heavy ``__init__`` side-effects like vLLM / reward setup).
+This is an internal implementation detail for code reuse --
+users writing custom losses do NOT need to understand or replicate
+this pattern; they simply subclass ``Loss`` and implement
+``forward_step`` / ``loss_func``.
+"""
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from .base import Loss
+
+
+class GRPOLoss(Loss):
+    """GRPO loss registered in the pipeline registry.
+
+    Builds a minimal ``MegatronGRPOTrainer`` stub that only holds
+    algorithm parameters (beta, epsilon, …) and reuses its
+    ``forward_step`` / ``loss_func`` without duplicating code.
+
+    To define a custom loss, subclass ``Loss``, override
+    ``forward_step`` / ``loss_func``, and pass the dotted path to
+    ``register_ray_trainer(..., loss='your.module.YourLoss')``.
+    """
+
+    def __init__(self, args):
+        self._dummy = self._create_dummy_trainer(args)
+
+    @staticmethod
+    def _create_dummy_trainer(args):
+        """Create a minimal MegatronGRPOTrainer for loss computation only.
+
+        Skips the heavy __init__ side-effects (vLLM, reward, model init)
+        by using __new__ and manually initialising only the fields that
+        ``forward_step`` / ``loss_func`` actually read.
+        """
+        import torch
+
+        from swift.megatron.trainers.grpo_trainer import MegatronGRPOTrainer
+        from swift.utils import is_last_rank
+        cls = MegatronGRPOTrainer
+        dummy = cls.__new__(cls)
+        dummy.args = args
+        dummy._setup_teacher()
+        dummy._init_grpo_params()
+        dummy._prepare_metrics()
+        dummy.log_rollout_offpolicy_metrics = args.log_rollout_offpolicy_metrics
+        dummy.disable_rollout_importance_sampling = False
+        dummy.enable_routing_replay = args.router_replay_mode != 'disabled'
+        dummy.micro_batch_size = args.micro_batch_size
+        dummy.temperature = args.temperature
+        dummy.is_main_process = is_last_rank()
+        dummy.process_index = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+        dummy.world_size = torch.distributed.get_world_size() if torch.distributed.is_initialized() else 1
+        dummy._step = 0
+        dummy.max_completion_length = args.max_completion_length
+
+        class _AlwaysTraining:
+            training = True
+
+        dummy.unwrapped_models = [_AlwaysTraining()]
+        return dummy
+
+    def forward_step(self, data_iterator, model):
+        from swift.megatron.trainers.grpo_trainer import MegatronGRPOTrainer
+        cls = MegatronGRPOTrainer
+        return cls.forward_step(self._dummy, data_iterator, model)
+
+    def loss_func(self, output_tensor, *, data: Dict[str, Any]):
+        return self._dummy.loss_func(output_tensor, data=data)

swift/ray/megatron/megatron_worker.py

@@ -0,0 +1,640 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+import os
+import torch
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+
+from swift.rl_core.data import GRPOBatch
+from swift.utils import gc_collect, get_current_device, get_logger
+from .checkpoint_engine import CheckpointEngineMixin
+from .worker_group import dispatch_collect
+
+if TYPE_CHECKING:
+    from swift.rlhf_trainers.gkd_loss import TeacherOutput
+
+logger = get_logger()
+
+# --- worker-side data-flow types ------------------------------------------
+# Since collation moved to the driver, the worker consumes already-collated
+# micro-batches and returns per-sample logps:
+#
+# ModelInputs: a collated micro-batch (driver-side ``collate_to_grpo_micro_batch``)
+#   fed to ``model(**model_inputs)``. Carries the model-forward tensors plus
+#   ``grpo_batch`` (GRPOBatch) and optional ``teacher_output`` / ``teacher_model_inputs``
+#   / ``data_source``.
+# LogpsRow: one per-sample logps result returned worker→driver (collected via
+#   ``collect='dp_flat'``), e.g. ``{'per_token_logps': ..., 'completion_mask': ...}``.
+ModelInputs = Dict[str, Any]
+LogpsRow = Dict[str, torch.Tensor]
+
+
+def _import_class(dotted_path: str):
+    """Import a class from a dotted module path like ``'a.b.ClassName'``."""
+    import importlib
+    mod_path, cls_name = dotted_path.rsplit('.', 1)
+    return getattr(importlib.import_module(mod_path), cls_name)
+
+
+def _make_lifecycle_trainer(args, template):
+    from swift.megatron.trainers.rlhf_mixin import MegatronRLHFTrainer
+
+    class _LifecycleTrainer(MegatronRLHFTrainer):
+
+        def forward_step(self, data_iterator, model):
+            return None
+
+    return _LifecycleTrainer(args, template)
+
+
+class MegatronWorker(CheckpointEngineMixin):
+
+    def __init__(self):
+        self._megatron = None
+        self._loss_fn = None
+        self._args = None
+        self._pipeline = None
+        self.rollout = None
+        self._checkpoint_engine = None
+        self._bucket_size: int = 3072 << 20
+        self.actor = None  # TrainableModelWorker
+        self.ref = None  # MegatronModelWorker (explicit ref for full fine-tune)
+        self.teacher = None  # MegatronModelWorker (colocated teacher)
+
+    def init_actor(
+        self,
+        cfg: Dict[str, Any],
+        loss_cls_path: Optional[str] = None,
+        rollout_config: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Initialise the training (actor) model, optimizer, and optionally the rollout adapter.
+
+        This only sets up the actor model for training. Ref and teacher models
+        are initialized separately (ref in _init_trainable, teacher via
+        init_teacher_model).
+
+        Args:
+            cfg: Merged config dict (shared + group overrides).
+            rollout_config: When provided, creates an internal RolloutAdapter.
+        """
+        from swift.megatron.arguments import MegatronRLHFArguments
+        from swift.megatron.pipelines.train.rlhf import MegatronRLHF
+        from .driver_utils import parse_args_from_dict
+
+        args = parse_args_from_dict(MegatronRLHFArguments, cfg)
+        self._pipeline = MegatronRLHF(args)
+        self._loss_cls_path = loss_cls_path
+        self._args = self._pipeline.args
+
+        self._init_trainable()
+
+        if rollout_config:
+            self._init_rollout_adapter(rollout_config)
+
+    def _init_trainable(self):
+        from .model_worker import TrainableModelWorker
+
+        pipeline = self._pipeline
+        args = pipeline.args
+        self._megatron = _make_lifecycle_trainer(args, pipeline.template)
+
+        if self._loss_cls_path:
+            loss_cls = _import_class(self._loss_cls_path)
+            self._loss_fn = loss_cls(args)
+            self._megatron.forward_step = self._loss_fn.forward_step
+
+        self.actor = TrainableModelWorker(args, self._megatron)
+        if args.tuner_type == 'full' and self._megatron.ref_models:
+            from .model_worker import MegatronModelWorker
+            self.ref = MegatronModelWorker(args, self._megatron.ref_models)
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def init_teacher_model(self, model_dir: str):
+        """Load a colocated teacher model (same parallelism as student)."""
+        from .model_worker import MegatronModelWorker
+
+        # Prefer the worker's own resolved teacher_model_dir; bridge.load_weights needs a
+        # real local path to locate safetensors (a raw model id yields an empty state dict).
+        model_dir = getattr(self._args, 'teacher_model_dir', None) or model_dir
+        self.teacher = MegatronModelWorker.from_pretrained(self._args, model_dir)
+        logger.info('Colocated teacher model loaded from %s', model_dir)
+        if getattr(self._args, 'offload_teacher_model', False):
+            self.teacher.offload_to_cpu()
+
+    @dispatch_collect(dispatch='dp', collect='first')
+    def compute_teacher_logits(self, micro_batches: List[ModelInputs]) -> None:
+        """Forward the teacher on each driver-collated micro-batch's ``teacher_model_inputs``
+        and cache one batched ``TeacherOutput`` per micro-batch (worker-local).
+
+        Cached (never returned to the driver): for context parallel each rank forwards its
+        own sequence shard, so the cache is already the correct local slice. ``train_step``
+        attaches it to the matching micro-batch (same dispatch dict ⇒ same order).
+        """
+        teacher_inputs = [mi['teacher_model_inputs'] for mi in micro_batches]
+        if getattr(self._args, '_teacher_use_disable_adapter', False):
+            # Self-distillation (LoRA): teacher = student base model with the LoRA adapter
+            # disabled — no separate teacher loaded.
+            from contextlib import ExitStack
+            megatron = self._megatron
+            with ExitStack() as stack:
+                for m in megatron.peft_models:
+                    stack.enter_context(m.disable_adapter())
+                self._cached_teacher_logits = self._loss_fn.compute_teacher_logits(megatron.unwrapped_models[0],
+                                                                                   teacher_inputs, self._args)
+        else:
+            assert self.teacher is not None, 'Teacher model not initialized. Call init_teacher_model first.'
+            with self.teacher.loaded_context():
+                self._cached_teacher_logits = self._loss_fn.compute_teacher_logits(self.teacher.models[0],
+                                                                                   teacher_inputs, self._args)
+        gc_collect()
+
+    def get_parallel_info(self) -> Dict[str, Any]:
+        from megatron.core import mpu
+        info = {
+            'dp_rank':
+            mpu.get_data_parallel_rank(),
+            'dp_size':
+            mpu.get_data_parallel_world_size(),
+            'is_collector':
+            (mpu.get_tensor_model_parallel_rank() == 0
+             and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1
+             and mpu.get_context_parallel_rank() == 0),
+        }
+        return info
+
+    def get_padding_to(self) -> Optional[int]:
+        """Delegates to ``swift.megatron.utils.get_padding_to`` (handles SP, CP, fp8)."""
+        from swift.megatron.utils import get_padding_to
+        return get_padding_to(self._args)
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def setup(self, args_override: Dict[str, Any]) -> Dict[str, Any]:
+        """Apply pre-computed args from driver, then set up model training.
+
+        The driver is responsible for computing train_iters, eval_iters,
+        save_steps, etc.  The worker just applies the overrides.
+        """
+        megatron = self._megatron
+        for k, v in args_override.items():
+            if v is not None:
+                setattr(megatron.args, k, v)
+        megatron.setup_model_training()
+        return {
+            'train_iters': megatron.args.train_iters,
+            'iteration': megatron.state.iteration,
+        }
+
+    @dispatch_collect(dispatch='dp', collect='all')
+    def train_step(self,
+                   micro_batches: List[ModelInputs],
+                   extra_metrics: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        from swift.utils import to_device
+        megatron = self._megatron
+        args = megatron.args
+        assert isinstance(micro_batches, list), \
+            f'train_step expects List[ModelInputs], got {type(micro_batches).__name__}'
+        self._inject_extra_metrics(extra_metrics)
+        # GKD colocated teacher: attach the per-micro-batch TeacherOutput cached by
+        # compute_teacher_logits (worker-local, already CP-correct), aligned by order.
+        cached_teacher = getattr(self, '_cached_teacher_logits', None)
+        if cached_teacher is not None:
+            for mi, t_out in zip(micro_batches, cached_teacher):
+                if t_out is not None:
+                    mi['teacher_output'] = t_out
+            self._cached_teacher_logits = None
+        # Driver-side collate produces CPU tensors; move each micro-batch to the GPU.
+        device = get_current_device()
+        moved: List[ModelInputs] = []
+        for mi in micro_batches:
+            mi.pop('teacher_model_inputs', None)  # consumed by compute_teacher_logits
+            grpo_batch = mi.pop('grpo_batch', None)
+            teacher_output = mi.pop('teacher_output', None)  # GPU (cache) or CPU (replicas)
+            mi = to_device(mi, device)
+            if grpo_batch is not None:
+                mi['grpo_batch'] = grpo_batch.to_device(device)
+            if teacher_output is not None:
+                mi['teacher_output'] = teacher_output.to_device(device)
+            moved.append(mi)
+        micro_batches = moved
+        data_iterator = iter(micro_batches)
+        assert len(micro_batches) == args.num_microbatches, (
+            f'Worker got {len(micro_batches)} micro-batches but args.num_microbatches='
+            f'{args.num_microbatches}; check per_device_generation_batch_size / micro_batch_size config.')
+        router_replay_mode = getattr(args, 'router_replay_mode', 'disabled')
+        need_routing_replay = router_replay_mode != 'disabled'
+        RouterReplay = None
+        if need_routing_replay:
+            try:
+                from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+                RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+            except ImportError:
+                need_routing_replay = False
+
+        try:
+            megatron.run_train_step(data_iterator, None)
+        finally:
+            if need_routing_replay and RouterReplay is not None:
+                RouterReplay.clear_global_indices()
+                RouterReplay.clear_global_router_replay_action()
+        del data_iterator
+        gc_collect()
+        return self._extract_step_metrics(megatron)
+
+    def _inject_extra_metrics(self, extra_metrics) -> None:
+        """Inject driver-computed metrics (reward, MathAccuracy, data_source, ...) into
+        the megatron trainer's ``_train_metrics`` so they flow through the standard
+        ``on_log`` path (console PrintCallback + tensorboard + swanlab), unifying ALL
+        logging in the worker's megatron callbacks (the driver no longer prints metrics).
+
+        Values are stored as ``[sum, count]`` pairs to match ``_aggregated_metrics`` /
+        ``_log_callback`` (which divides sum/count), so a per-step scalar logs as itself.
+        """
+        if not extra_metrics:
+            return
+        megatron = self._megatron
+        tm = getattr(megatron, '_train_metrics', None)
+        if tm is None:
+            tm = megatron._train_metrics = {}
+        device = get_current_device()
+        for k, v in extra_metrics.items():
+            if v is None:
+                continue
+            add = torch.tensor([float(v), 1.0], dtype=torch.float32, device=device)
+            tm[k] = tm[k] + add if k in tm else add
+
+    @staticmethod
+    def _extract_step_metrics(megatron) -> Dict[str, Any]:
+        """Extract training metrics from the last logged step.
+
+        After ``run_train_step``, Megatron's ``on_log`` stores metrics
+        in ``_last_logged_metrics``.  We extract numeric values and
+        normalize key names (e.g. ``learning_rate`` → ``lr``) so the
+        driver receives a clean dict for aggregation and logging.
+        """
+        result: Dict[str, Any] = {'iteration': megatron.state.iteration}
+        logged = getattr(megatron, '_last_logged_metrics', None) or {}
+        for k, v in logged.items():
+            if isinstance(v, (int, float)):
+                result[k] = v
+            else:
+                try:
+                    result[k] = float(v)
+                except (TypeError, ValueError):
+                    continue
+        if 'learning_rate' in result:
+            result['lr'] = result.pop('learning_rate')
+        return result
+
+    @dispatch_collect(dispatch='dp', collect='dp_flat')
+    def compute_logps(self, micro_batches: List[ModelInputs]) -> List[LogpsRow]:
+        """Compute per-token logps under the current policy model.
+
+        Receives this dp_rank's collated micro-batches (driver-side collate, CPU)
+        and runs the rank-local forward; returns one row per sample.
+        """
+        model = self._megatron.unwrapped_models[0]
+        return self._compute_logps_micro_batches(micro_batches, model, 'per_token_logps')
+
+    @dispatch_collect(dispatch='dp', collect='dp_flat')
+    def compute_ref_logps(self, micro_batches: List[ModelInputs]) -> List[LogpsRow]:
+        """Compute per-token logps under the frozen reference model."""
+        if self.ref is not None:
+            return self._compute_logps_micro_batches(micro_batches, self.ref.models[0], 'ref_per_token_logps')
+        with self.actor.null_ref_context() as ref_models:
+            return self._compute_logps_micro_batches(micro_batches, ref_models[0], 'ref_per_token_logps')
+
+    @dispatch_collect(dispatch='dp', collect='dp_flat')
+    def compute_teacher_logps(self, micro_batches: List[ModelInputs]) -> List[LogpsRow]:
+        """OPD-RL: per-token teacher logp on the sampled tokens (token-in-token-out).
+
+        Same forward/frame as ``compute_logps`` so the teacher logp aligns with the policy's;
+        same-model LoRA self-distillation disables the student's adapter, otherwise the colocated
+        teacher (loaded via ``init_teacher_model``) is used.
+        """
+        if getattr(self._args, '_teacher_use_disable_adapter', False):
+            from contextlib import ExitStack
+            with ExitStack() as stack:
+                for m in self._megatron.peft_models:
+                    stack.enter_context(m.disable_adapter())
+                return self._compute_logps_micro_batches(micro_batches, self._megatron.unwrapped_models[0],
+                                                         'teacher_per_token_logps')
+        # Dynamic self-distillation (teacher is None): teacher = student (same weights
+        # including LoRA). No offload/load needed.
+        model = self.teacher.models[0] if self.teacher else self._megatron.unwrapped_models[0]
+        with (self.teacher.loaded_context() if self.teacher else nullcontext()):
+            return self._compute_logps_micro_batches(micro_batches, model, 'teacher_per_token_logps')
+
+    def _compute_logps_micro_batches(
+        self,
+        micro_batches: List[ModelInputs],
+        model,
+        output_key: str,
+    ) -> List[LogpsRow]:
+        from swift.utils import to_device
+        device = get_current_device()
+        rows: List[LogpsRow] = []
+        for model_inputs in micro_batches:
+            grpo_batch = model_inputs.pop('grpo_batch').to_device(device)
+            model_inputs = to_device(model_inputs, device)
+            model_inputs['grpo_batch'] = grpo_batch  # _compute_logps pops it again
+            out = self._compute_logps(model_inputs, model, output_key)
+            if output_key not in out:
+                continue
+            rows.extend(
+                self._split_logps_rows(
+                    out[output_key],
+                    output_key,
+                    routed_experts=out.get('routed_experts'),
+                    seq_lengths=grpo_batch.seq_lengths))
+        return rows
+
+    @staticmethod
+    def _split_logps_rows(
+        logps: Optional[torch.Tensor],
+        key: str,
+        *,
+        routed_experts: Optional[torch.Tensor] = None,
+        seq_lengths: Optional[torch.Tensor] = None,
+    ) -> List[LogpsRow]:
+        if logps is None:
+            return []
+        routed_rows: List[Optional[torch.Tensor]] = [None] * int(logps.shape[0])
+        if routed_experts is not None:
+            routed = routed_experts.detach().cpu() if isinstance(routed_experts, torch.Tensor) \
+                else torch.as_tensor(routed_experts)
+            if routed.dim() > 0 and routed.shape[0] == logps.shape[0]:
+                routed_rows = [routed[i] for i in range(logps.shape[0])]
+            elif routed.dim() > 1 and routed.shape[0] == 1 and seq_lengths is not None:
+                seq_cpu = seq_lengths.detach().cpu().tolist()
+                start = 0
+                for i in range(logps.shape[0]):
+                    seq_len = int(seq_cpu[i])
+                    end = start + seq_len
+                    routed_rows[i] = routed[0, start:end]
+                    start = end
+        rows: List[LogpsRow] = []
+        for i in range(logps.shape[0]):
+            item: LogpsRow = {key: logps[i].detach().cpu()}
+            if routed_rows[i] is not None:
+                item['routed_experts'] = routed_rows[i].detach().cpu()
+            rows.append(item)
+        return rows
+
+    def _compute_logps(self, model_inputs: ModelInputs, model, output_key: str) -> Dict[str, torch.Tensor]:
+        """Compute per-token logps for a single collated micro-batch."""
+        from swift.rlhf_trainers.utils import pad_logps_back_to_batch
+        megatron = self._megatron
+        args = self._args
+        temperature = getattr(args, 'temperature', 1.0)
+        # grpo_batch carries the per-batch masks/seq_lengths; pop it so what
+        # remains is the pure ``model(**model_inputs)`` forward kwargs.
+        grpo_batch: GRPOBatch = model_inputs.pop('grpo_batch')
+        seq_lengths = grpo_batch.seq_lengths
+        batch_size = grpo_batch.completion_mask.shape[0]
+        max_seq_len = grpo_batch.completion_mask.shape[1]
+
+        enable_routing_replay = bool(getattr(megatron, 'enable_routing_replay', False))
+        router_mode = getattr(args, 'router_replay_mode', 'disabled')
+
+        RouterReplay = None
+        RouterReplayAction = None
+        if enable_routing_replay:
+            try:
+                from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction
+            except ImportError:
+                enable_routing_replay = False
+
+        if enable_routing_replay and RouterReplay is not None:
+            if router_mode == 'R2':
+                RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)
+            elif router_mode == 'R3':
+                RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)
+
+        routing_topk_idx = None
+        try:
+            logps_packed, routing_topk_idx = megatron.compute_per_token_logps(
+                model, iter([model_inputs]), temperature=temperature)
+        finally:
+            if enable_routing_replay and RouterReplay is not None:
+                RouterReplay.clear_global_indices()
+                RouterReplay.clear_global_router_replay_action()
+
+        out: Dict[str, torch.Tensor] = {}
+        if logps_packed is not None:
+            if args.padding_free:
+                logps, _ = pad_logps_back_to_batch(
+                    logps_rmpad=logps_packed,
+                    logits_to_keep=max_seq_len,
+                    batch_size=batch_size,
+                    seq_lengths=seq_lengths)
+            else:
+                logps = logps_packed
+            out[output_key] = logps.detach().cpu()
+        if routing_topk_idx is not None:
+            out['routed_experts'] = routing_topk_idx.detach().cpu()
+        return out
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def finalize(self) -> Dict[str, Any]:
+        from swift.utils import is_last_rank
+        megatron = self._megatron
+        megatron.finalize_training()
+        self._pipeline._handle_trainer_state(megatron, is_last_rank())
+        state = megatron.state
+        return {
+            'last_model_checkpoint': state.last_model_checkpoint,
+            'best_model_checkpoint': state.best_model_checkpoint,
+            'best_metric': state.best_metric,
+        }
+
+    def _init_rollout_adapter(self, rollout_config: Dict[str, Any]) -> None:
+        """Create the internal RolloutAdapter.
+
+        The adapter lazily resolves the VllmServer handle via named actor,
+        so it can be created before the server is fully started.
+        """
+        from .rollout.adapter import RolloutAdapter
+
+        tp = rollout_config['rollout_tp_size']
+        dp = rollout_config['rollout_dp_size']
+        world_per_replica = tp * dp
+        rank = int(os.environ.get('RANK', '0'))
+        replica_rank = rank // world_per_replica
+        rollout_rank = rank % world_per_replica
+        bucket_mb = rollout_config.get('bucket_size_mb', 2048)
+
+        self.rollout = RolloutAdapter(
+            replica_rank=replica_rank,
+            rollout_rank=rollout_rank,
+            bucket_size_mb=bucket_mb,
+        )
+        logger.info('MegatronWorker[rank=%s]: rollout adapter created (replica=%d, rollout_rank=%d)', rank,
+                    replica_rank, rollout_rank)
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def merge_lora(self):
+        """Merge LoRA adapters into base weights (must be called before offload)."""
+        megatron = self._megatron
+        if megatron.args.tuner_type in ('lora', 'lora_llm'):
+            megatron.merge_lora_adapters()
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def unmerge_lora(self):
+        """Unmerge LoRA adapters to restore training state (call after reload)."""
+        megatron = self._megatron
+        if megatron.args.tuner_type in ('lora', 'lora_llm'):
+            megatron.unmerge_lora_adapters()
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def update_weights(self, adapter_only: bool = False):
+        """Push training weights to rollout via IPC (streaming).
+
+        All TP ranks must call export_weights (contains TP collectives).
+        Only the primary rank sends; others drain the iterator.
+
+        For full-weight sync with LoRA, the caller must ensure merge_lora()
+        was called beforehand and unmerge_lora() is called after reload.
+
+        Args:
+            adapter_only: When True, export only LoRA adapter weights
+                (peft_format=True) and pass peft_config to vLLM for
+                TensorLoRARequest loading. When False, export full
+                merged weights.
+        """
+        megatron = self._megatron
+        target_device = 'cpu' if megatron.args.offload_bridge else None
+
+        if adapter_only:
+            weight_iter = megatron.bridge.export_weights(
+                megatron.unwrapped_models, target_device=target_device, peft_format=True)
+            peft_config = self.get_peft_config_dict()
+            lora_names = None
+        else:
+            weight_iter = megatron.bridge.export_weights(megatron.unwrapped_models, target_device=target_device)
+            peft_config = None
+            lora_names = self._resolve_lora_param_names()
+
+        if self.rollout is None:
+            # No rollout adapter attached just drain the
+            # iterator so all TP ranks finish the collective export.
+            for _ in weight_iter:
+                pass
+            return
+
+        self.rollout.update_weights(
+            weight_iter,
+            vllm_lora_param_names=lora_names,
+            peft_config=peft_config,
+            base_sync_done=adapter_only,
+        )
+        self.rollout.reset_prefix_cache()
+
+    def _resolve_lora_param_names(self) -> Optional[set]:
+        """Get vLLM param names for LoRA mapping, if applicable."""
+        megatron = self._megatron
+        if not (megatron.args.tuner_type == 'lora' and megatron.args.vllm_enable_lora):
+            return None
+        raw_names = self.rollout.get_model_param_names()
+        if not raw_names:
+            return None
+        from swift.rlhf_trainers.utils import expand_vllm_param_name_aliases
+        expanded = expand_vllm_param_name_aliases(set(raw_names))
+        stripped = set()
+        for n in expanded:
+            stripped.add(n)
+            if n.startswith('model.'):
+                stripped.add(n[len('model.'):])
+        return stripped
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def finalize_generation(self):
+        if self.rollout is not None:
+            self.rollout.reset_prefix_cache()
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def offload_to_cpu(self):
+        self.actor.offload_to_cpu()
+        if self.ref is not None:
+            self.ref.offload_to_cpu()
+        if self.teacher is not None:
+            self.teacher.offload_to_cpu()
+
+    @dispatch_collect(dispatch='broadcast', collect='first')
+    def reload_to_gpu(self):
+        self.actor.reload_to_gpu()
+        if self.ref is not None:
+            self.ref.reload_to_gpu(load_grad=False)
+        # When offload_teacher_model is set, the teacher is managed by compute_teacher_logits
+        # (loaded only for the teacher forward), so keep it on CPU here.
+        if self.teacher is not None and not getattr(self._args, 'offload_teacher_model', False):
+            self.teacher.reload_to_gpu(load_grad=False)
+
+    @staticmethod
+    def _align_seq_len(t, target_len, pad_val=0):
+        """Pad or truncate a tensor along dim=1 to target_len. Works for 2D [B,S] and 3D [B,S,*]."""
+        cur = t.shape[1]
+        if cur == target_len:
+            return t
+        if cur < target_len:
+            pad = (0, target_len - cur) if t.dim() == 2 else (0, 0, 0, target_len - cur)
+            return torch.nn.functional.pad(t, pad, value=pad_val)
+        return t[:, :target_len]
+
+    @staticmethod
+    def _collate_teacher_outputs(
+        teacher_outputs: List['TeacherOutput'],
+        device: torch.device,
+        padding_free: bool = False,
+        target_seq_len: Optional[int] = None,
+        is_opsd: bool = False,
+    ) -> 'TeacherOutput':
+        """Collate per-sample TeacherOutputs into a batched one (driver-side).
+
+        For non-OPSD: each tensor is aligned to target_seq_len (pad or truncate).
+        For OPSD: teacher keeps its own length (target_seq_len ignored).
+        """
+        from swift.rlhf_trainers.gkd_loss import TeacherOutput
+        effective_target = None if is_opsd else target_seq_len
+        pad_vals = {'topk_logprobs': float('-inf'), 'labels': -100}
+        fields = ('full_logits', 'topk_logprobs', 'topk_indices', 'labels')
+        kwargs = {}
+        for field in fields:
+            tensors = [getattr(t, field) for t in teacher_outputs]
+            tensors = [t for t in tensors if t is not None]
+            if not tensors:
+                continue
+            pad_val = pad_vals.get(field, 0)
+            if effective_target is not None:
+                tensors = [MegatronWorker._align_seq_len(t, effective_target, pad_val) for t in tensors]
+            if padding_free:
+                non_empty = [t for t in tensors if t.shape[0] > 0]
+                kwargs[field] = torch.cat(non_empty, dim=1).to(device)
+            else:
+                kwargs[field] = torch.cat(tensors, dim=0).to(device)
+        return TeacherOutput(**kwargs)
+
+    def send_checkpoint_weights(self, adapter_only: bool = False) -> None:
+        """Export and send model weights via NCCL checkpoint engine."""
+        import asyncio
+        megatron = self._megatron
+        engine = self._get_or_create_checkpoint_engine()
+        target_device = 'cpu' if megatron.args.offload_bridge else None
+        weight_iter = megatron.bridge.export_weights(
+            megatron.unwrapped_models, target_device=target_device, peft_format=adapter_only)
+        asyncio.run(engine.send_weights(weight_iter))
+
+    def get_peft_config_dict(self) -> dict:
+        """Return the PEFT config for LoRA-only sync."""
+        from dataclasses import asdict
+        peft_config = self._megatron.unwrapped_models[0].peft_config['default']
+        return asdict(peft_config)
+
+    def shutdown(self):
+        self.rollout = None
+        self._megatron = None
+        self._loss_fn = None
+        self._checkpoint_engine = None
+        self.actor = None
+        self.ref = None
+        self.teacher = None

swift/ray/megatron/model_worker.py

@@ -0,0 +1,127 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Single-model worker abstraction for Ray-based Megatron training.
+
+MegatronModelWorker wraps one Megatron model (inference-only).
+TrainableModelWorker extends it with training capabilities via _LifecycleTrainer.
+"""
+from __future__ import annotations
+
+import torch
+from contextlib import contextmanager
+from typing import Any, Dict, Optional, Sequence
+
+from swift.utils import gc_collect, get_current_device, get_logger, to_device
+
+logger = get_logger()
+
+
+class MegatronModelWorker:
+    """Wraps a single Megatron model with inference / offload interfaces.
+
+    Two creation paths:
+    - ``__init__(args, models)``: wraps models already created in the current
+      process (e.g. ref models created by MegatronRLHFTrainer.prepare_model).
+    - ``from_pretrained(args, model_dir)``: independently loads a new model
+      from disk (e.g. colocated teacher with different weights).
+    """
+
+    def __init__(self, args, models, bridge=None):
+        self.args = args
+        self.models = models
+        self.bridge = bridge
+
+    @classmethod
+    def from_pretrained(cls, args, model_dir):
+        """Load an inference-only model (ref / teacher) from disk."""
+        from transformers import AutoConfig
+
+        from swift.megatron.model import get_mcore_model
+
+        hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True)
+        models = get_mcore_model(args, hf_config)
+        for m in models:
+            if not args.use_cpu_initialization:
+                m.cuda(torch.cuda.current_device())
+            m.requires_grad_(False)
+            m.eval()
+        models[0].config.bridge.load_weights(models, model_dir)
+        return cls(args, models, bridge=models[0].config.bridge)
+
+    def compute_per_token_logps(self, data_iterator, temperature=1.0, enable_routing_replay=False):
+        from swift.megatron.trainers.utils import compute_per_token_logps_fn
+        return compute_per_token_logps_fn(
+            self.models[0],
+            self.args,
+            data_iterator,
+            temperature=temperature,
+            enable_routing_replay=enable_routing_replay)
+
+    def offload_to_cpu(self):
+        from swift.megatron.trainers.utils import offload_megatron_model_to_cpu
+        offload_megatron_model_to_cpu(self.models)
+        gc_collect()
+
+    def reload_to_gpu(self, load_grad=False):
+        from swift.megatron.trainers.utils import load_megatron_model_to_gpu
+        load_megatron_model_to_gpu(self.models, load_grad=load_grad)
+
+    @contextmanager
+    def loaded_context(self, load_grad=False):
+        """Temporarily load model to GPU, offload on exit.
+
+        No-op if offloading is not configured. Use this to bracket
+        inference on an offloaded model (e.g. teacher forward).
+        """
+        if not getattr(self.args, 'offload_teacher_model', False):
+            yield
+            return
+        self.reload_to_gpu(load_grad=load_grad)
+        try:
+            yield
+        finally:
+            self.offload_to_cpu()
+
+
+class TrainableModelWorker(MegatronModelWorker):
+    """Trainable model wrapping a _LifecycleTrainer with optimizer / training step.
+
+    Note: the lifecycle_trainer dependency on MegatronRLHFTrainer is a
+    transitional design. Future refactoring should extract the needed
+    capabilities (optimizer, model wrapping, ref model context) into
+    standalone components so the ray module no longer depends on the
+    non-ray trainer hierarchy.
+    """
+
+    def __init__(self, args, lifecycle_trainer):
+        self._trainer = lifecycle_trainer
+        super().__init__(args, lifecycle_trainer.wrapped_models, lifecycle_trainer.bridge)
+
+    @property
+    def trainer(self):
+        return self._trainer
+
+    @property
+    def unwrapped_models(self):
+        return self._trainer.unwrapped_models
+
+    def set_forward_step(self, fn):
+        self._trainer.forward_step = fn
+
+    def run_train_step(self, data_iterator):
+        self._trainer.run_train_step(data_iterator, None)
+
+    def null_ref_context(self):
+        return self._trainer.null_ref_context()
+
+    def offload_to_cpu(self):
+        from swift.megatron.trainers.utils import offload_megatron_model_to_cpu, offload_megatron_optimizer
+        offload_megatron_model_to_cpu(self._trainer.wrapped_models)
+        if getattr(self._trainer, 'optimizer', None) and self.args.offload_optimizer:
+            offload_megatron_optimizer(self._trainer.optimizer)
+        gc_collect()
+
+    def reload_to_gpu(self, load_grad=True):
+        from swift.megatron.trainers.utils import load_megatron_model_to_gpu, load_megatron_optimizer
+        load_megatron_model_to_gpu(self._trainer.wrapped_models)
+        if getattr(self._trainer, 'optimizer', None) and self.args.offload_optimizer:
+            load_megatron_optimizer(self._trainer.optimizer)

swift/ray/megatron/pipeline.py

@@ -0,0 +1,409 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import importlib
+import os
+import ray
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional
+
+from swift.utils import get_logger
+from .base_trainer import BaseRayTrainer
+from .driver_utils import (build_dataset_from_dict, compute_iter_params, estimate_dp_size, merge_group_dict,
+                           parse_ray_yaml)
+
+logger = get_logger()
+
+_TRAINER_REGISTRY: Dict[str, Dict[str, Any]] = {
+    'grpo': {
+        'trainer': 'swift.ray.megatron.grpo_trainer.GRPOTrainer',
+        'loss': 'swift.ray.megatron.loss.grpo.GRPOLoss',
+    },
+    'gkd': {
+        'trainer': 'swift.ray.megatron.gkd_trainer.GKDTrainer',
+        'loss': 'swift.ray.megatron.loss.gkd.GKDLoss',
+    },
+}
+
+
+def register_ray_trainer(
+    rlhf_type: str,
+    trainer: str,
+    loss: Optional[str] = None,
+):
+    """Register a custom algorithm for the Ray pipeline.
+
+    Args:
+        rlhf_type: Algorithm identifier (e.g. ``'grpo'``).
+        trainer: Dotted path to the driver-side trainer class.
+            The class must accept ``(worker_groups, rollout_replicas)``
+            and expose ``set_data_info()`` / ``train()`` methods.
+            Example: ``'swift.ray.megatron.grpo_trainer.GRPOTrainer'``
+        loss: Dotted path to a ``Loss`` subclass that defines
+            ``forward_step`` + ``loss_func``.
+            Pass ``None`` to use the internal trainer's forward_step.
+    """
+    _TRAINER_REGISTRY[rlhf_type] = {'trainer': trainer, 'loss': loss}
+
+
+class MegatronRayPipeline:
+
+    def __init__(self, config_path: str):
+        self.ray_config, group_configs, shared_config = parse_ray_yaml(config_path)
+        shared_config['use_ray'] = True
+
+        self.rlhf_type = self.ray_config.rlhf_type
+        if self.rlhf_type not in _TRAINER_REGISTRY:
+            raise ValueError('Unknown rlhf_type %r. Available: %s' % (self.rlhf_type, list(_TRAINER_REGISTRY)))
+
+        self.group_cfgs: Dict[str, Dict[str, Any]] = {
+            g: merge_group_dict(shared_config, gd)
+            for g, gd in group_configs.items()
+        }
+        self.shared_cfg = {k: v for k, v in shared_config.items() if v is not None}
+
+        self._entry = _TRAINER_REGISTRY[self.rlhf_type]
+        self.resource_pool_manager = None
+        self.worker_groups: Dict[str, Any] = {}
+        self.rollout_replicas: List[Any] = []
+        self.teacher_replicas: List[Any] = []
+
+    def init(self) -> None:
+        # Initialize Ray, create resource pools, spawn workers and replicas.
+        self._data_info = self._build_dataset()
+        self._compute_train_iters()
+        ray.init(ignore_reinit_error=True)
+        self._create_pools()
+        self._init_worker_groups()
+        with self._colocate_offload_ctx():
+            self._init_rollout_replicas()
+        self._init_teacher_replicas()
+        self._driver_trainer: BaseRayTrainer = self._create_trainer()
+        self._driver_trainer.set_data_info(self._data_info)
+
+    def train(self) -> Any:
+        """Run the training loop.  Requires ``init()`` to have been called."""
+        if not hasattr(self, '_driver_trainer'):
+            raise RuntimeError('MegatronRayPipeline.train(): call init() first')
+        return self._driver_trainer.train()
+
+    def run(self) -> Any:
+        """Convenience: ``init()`` + ``train()`` + ``shutdown()``."""
+        self.init()
+        try:
+            return self.train()
+        finally:
+            self._shutdown()
+
+    def _build_dataset(self) -> Dict[str, Any]:
+        # Merge train group config (tuner_type, lora_rank, etc.) into shared cfg so that
+        # _check_teacher can detect LoRA self-distillation (teacher_model == model + lora).
+        cfg = {**self.shared_cfg, **self.group_cfgs.get('train', {})}
+        data_info = build_dataset_from_dict(cfg)
+        return data_info
+
+    def _compute_train_iters(self):
+        train_cfg = self.group_cfgs.get('train')
+        gpus = self.group_gpus.get('train', 0)
+        assert train_cfg is not None and gpus > 0
+        dp_size = estimate_dp_size(train_cfg, gpus)
+        iter_params = compute_iter_params(self._data_info, dp_size)
+        train_iters = iter_params.get('train_iters')
+        assert train_iters is not None and train_iters > 0
+
+        self.group_cfgs['train']['train_iters'] = train_iters
+
+    def _create_pools(self):
+        from .resource_pool import ResourcePool, ResourcePoolManager
+
+        colocated_sets = {frozenset(g) for g in self.colocate_groups}
+        pool_mapping: Dict[str, ResourcePool] = {}
+        assigned: set = set()
+
+        for colocated in colocated_sets:
+            # Colocated roles share one GPU set, so they must request the same number
+            # of gpus. Validate explicitly (and avoid relying on frozenset order).
+            gpus_by_role = {g: self.group_gpus.get(g, 0) for g in colocated}
+            distinct = set(gpus_by_role.values())
+            if distinct == {0}:
+                continue
+            if len(distinct) > 1:
+                raise ValueError(f'Colocated roles must request the same number of gpus, but got '
+                                 f'{gpus_by_role}. Set an equal `gpus` for all roles in {sorted(colocated)}.')
+            gpus = distinct.pop()
+            pon = self.ray_config.gpus_as_process_on_nodes(gpus)
+            shared = ResourcePool(pon, max_colocate_count=len(colocated))
+            for g in colocated:
+                pool_mapping[g] = shared
+                assigned.add(g)
+
+        for name, gpus in self.group_gpus.items():
+            if name in assigned or gpus <= 0:
+                continue
+            pon = self.ray_config.gpus_as_process_on_nodes(gpus)
+            pool_mapping[name] = ResourcePool(pon)
+
+        self.resource_pool_manager = ResourcePoolManager(pool_mapping)
+        self.resource_pool_manager.create_all()
+
+    def _is_rollout_hybrid(self) -> bool:
+        """True if rollout shares its pool with train (HYBRID mode)."""
+        return any('rollout' in cg and 'train' in cg for cg in self.colocate_groups)
+
+    def _init_worker_groups(self):
+        self._validate_grpo_train_batch_params()
+
+        self._spawn_train_group('train')
+
+        train_wg = self.worker_groups['train']
+        padding_vals = train_wg.broadcast('get_padding_to')
+        self._data_info['_padding_to'] = next((v for v in padding_vals if v is not None), None)
+
+    def _validate_grpo_train_batch_params(self) -> None:
+        """Early validation of GRPO batch params before spawning workers.
+
+        Resolves generation_batch_size / steps_per_generation, then validates
+        DP alignment — all via static helpers in RLHFMegatronArgumentsMixin.
+        """
+        if self.rlhf_type != 'grpo':
+            return
+        train_cfg = self.group_cfgs.get('train')
+        train_gpus = self.group_gpus.get('train', 0)
+        if not train_cfg or train_gpus <= 0:
+            return
+
+        cfg = dict(train_cfg)
+        global_batch_size = cfg.get('global_batch_size')
+        if global_batch_size <= 0:
+            return
+
+        micro_batch_size = cfg.get('micro_batch_size', 1)
+        num_generations = cfg.get('num_generations', 8)
+
+        from swift.megatron.arguments.megatron_args import RLHFMegatronArgumentsMixin
+        generation_batch_size, _ = RLHFMegatronArgumentsMixin.resolve_generation_batch_size(
+            cfg.get('generation_batch_size'), cfg.get('steps_per_generation'), global_batch_size, num_generations)
+
+        dp_size = estimate_dp_size(cfg, train_gpus)
+        RLHFMegatronArgumentsMixin.validate_batch_dp_alignment(generation_batch_size, num_generations, dp_size,
+                                                               micro_batch_size, train_gpus)
+
+    def _spawn_train_group(self, role: str) -> None:
+        from .megatron_worker import MegatronWorker
+        from .worker_group import WorkerGroup
+
+        pool = self.resource_pool_manager.get_pool(role)
+        cfg = dict(self.group_cfgs.get(role, {}))
+        cfg.setdefault('rlhf_type', self.ray_config.rlhf_type)
+        worker_cls = ray.remote(num_gpus=0)(MegatronWorker)
+        wg = WorkerGroup.from_pool(role, pool, worker_cls=worker_cls)
+
+        loss_cls = self._entry.get('loss')
+        rollout_config = self._build_rollout_config_for_workers() if self._is_rollout_hybrid() else None
+        wg.broadcast('init_actor', cfg, loss_cls_path=loss_cls, rollout_config=rollout_config)
+        wg.build_dispatch_info(worker_cls=MegatronWorker)
+
+        self.worker_groups[role] = wg
+        logger.info('MegatronWorker group [%s] on %d GPUs', role, pool.world_size)
+
+    @contextmanager
+    def _colocate_offload_ctx(self):
+        """Offload train workers during vLLM init (colocate only)."""
+        need = self._is_rollout_hybrid() and bool(self.shared_cfg.get('offload_model', True))
+        colocated_wgs = [
+            wg for role, wg in self.worker_groups.items() if need and any(role in g for g in self.colocate_groups)
+        ]
+        for wg in colocated_wgs:
+            wg.broadcast('offload_to_cpu')
+        try:
+            yield
+        finally:
+            for wg in colocated_wgs:
+                wg.broadcast('reload_to_gpu')
+
+    def _init_rollout_replicas(self) -> None:
+        rollout_gpus = self.group_gpus.get('rollout', 0)
+        if rollout_gpus <= 0:
+            self.rollout_replicas = []
+            return
+
+        from .rollout.replica import RolloutReplica
+
+        rollout_cfg = self._with_router_replay_rollout_config(self.group_cfgs.get('rollout', {}))
+        is_hybrid = self._is_rollout_hybrid()
+        pool = self.resource_pool_manager.get_pool('train' if is_hybrid else 'rollout')
+
+        template_kwargs = self._get_template_kwargs_for_rollout()
+        self.rollout_replicas = RolloutReplica.create_replicas(
+            rollout_cfg=rollout_cfg,
+            rollout_gpus=rollout_gpus,
+            pool=pool,
+            is_hybrid=is_hybrid,
+            sleep_level=self.ray_config.sleep_level,
+            template_kwargs=template_kwargs,
+        )
+
+    def _init_teacher_replicas(self) -> None:
+        teacher_gpus = self.group_gpus.get('teacher', 0)
+        if teacher_gpus <= 0:
+            self.teacher_replicas = []
+            return
+
+        from .rollout.replica import RolloutReplica
+
+        teacher_cfg = self.group_cfgs.get('teacher', {})
+        pool = self.resource_pool_manager.get_pool('teacher')
+        template_kwargs = self._get_template_kwargs_for_rollout()
+        args = self._data_info.get('_driver_args')
+        if args is not None:
+            base_len = template_kwargs.get('max_length') or getattr(args, 'max_length')
+            template_kwargs = dict(template_kwargs)
+            template_kwargs['max_length'] = base_len + args.max_completion_length
+        self.teacher_replicas = RolloutReplica.create_replicas(
+            rollout_cfg=teacher_cfg,
+            rollout_gpus=teacher_gpus,
+            pool=pool,
+            is_hybrid=False,
+            sleep_level=0,
+            template_kwargs=template_kwargs,
+            actor_name_prefix='swift_teacher_server',
+        )
+
+    def _with_router_replay_rollout_config(self, rollout_cfg: Dict[str, Any]) -> Dict[str, Any]:
+        cfg = dict(rollout_cfg or {})
+        args = self._data_info.get('_driver_args')
+        router_mode = getattr(args, 'router_replay_mode', 'disabled') if args is not None else 'disabled'
+        if router_mode != 'R3':
+            return cfg
+
+        from swift.rlhf_trainers.utils import check_vllm_version_ge
+        if not check_vllm_version_ge('0.14.0'):
+            raise ValueError('router_replay_mode=R3 requires vLLM>=0.14.0 to return routed_experts.')
+
+        engine_kwargs = dict(cfg.get('vllm_engine_kwargs') or {})
+        engine_kwargs.setdefault('enable_return_routed_experts', True)
+
+        # https://github.com/vllm-project/vllm/pull/39917
+        import vllm
+        from packaging import version
+        vllm_version = vllm.__version__
+        if vllm_version is not None and version.parse('0.21.0rc1') <= version.parse(vllm_version) <= version.parse(
+                '0.21.0'):
+            engine_kwargs.setdefault('async_scheduling', False)
+
+        cfg['vllm_engine_kwargs'] = engine_kwargs
+        return cfg
+
+    def _get_template_kwargs_for_rollout(self) -> Dict[str, Any]:
+        """Extract template config for vLLM alignment (padding_free=False, sp=1)."""
+        args = self._data_info.get('_driver_args')
+        if args is None:
+            return {}
+        kwargs = args.get_template_kwargs()
+        kwargs['padding_free'] = False
+        kwargs['sequence_parallel_size'] = 1
+        return kwargs
+
+    def _build_rollout_config_for_workers(self) -> Optional[Dict[str, Any]]:
+        """Build rollout config dict for MegatronWorker._init_rollout_adapter.
+
+        Returns None if no rollout GPUs are configured.
+        """
+        rollout_gpus = self.group_gpus.get('rollout', 0)
+        if rollout_gpus <= 0:
+            return None
+        rollout_cfg = self.group_cfgs.get('rollout', {})
+        bucket_mb = int(os.environ.get('SWIFT_RAY_WEIGHT_BUCKET_MB', '2048'))
+        return {
+            'rollout_tp_size': rollout_cfg.get('vllm_tensor_parallel_size', 1),
+            'rollout_dp_size': rollout_cfg.get('vllm_data_parallel_size', 1),
+            'bucket_size_mb': bucket_mb,
+        }
+
+    def _create_trainer(self):
+        cls_path = self._entry['trainer']
+        mod_path, cls_name = cls_path.rsplit('.', 1)
+        mod = importlib.import_module(mod_path)
+        trainer_cls = getattr(mod, cls_name)
+        weight_sync_mode = self._get_weight_sync_mode()
+        sleep_level = self._resolve_sleep_level()
+        return trainer_cls(
+            self.worker_groups,
+            self.rollout_replicas,
+            weight_sync_mode=weight_sync_mode,
+            sleep_level=sleep_level,
+            teacher_replicas=self.teacher_replicas)
+
+    def _resolve_sleep_level(self) -> int:
+        """Colocate: honor user config; separate: force 0."""
+        user_level = self.ray_config.sleep_level
+        if self._is_rollout_hybrid():
+            return user_level
+        if user_level != 0:
+            logger.warning('sleep_level=%d ignored in separate mode (vLLM stays resident). '
+                           'Overriding to 0.', user_level)
+        return 0
+
+    def _get_weight_sync_mode(self) -> str:
+        """Colocate: naive (IPC); separate: nccl (broadcast)."""
+        if self._is_rollout_hybrid():
+            return 'naive'
+        return 'nccl'
+
+    @property
+    def group_gpus(self) -> Dict[str, int]:
+        return self.ray_config.group_gpus
+
+    @property
+    def colocate_groups(self) -> List[List[str]]:
+        return self.ray_config.colocate_groups
+
+    def _shutdown(self):
+        """Best-effort teardown — each step swallows exceptions so a
+        failure in one stage does not skip the remaining cleanup."""
+        for replica in self.rollout_replicas:
+            try:
+                replica.shutdown()
+            except Exception as e:  # noqa: BLE001
+                logger.warning('RolloutReplica shutdown failed: %s', e)
+        self.rollout_replicas = []
+
+        for replica in self.teacher_replicas:
+            try:
+                replica.shutdown()
+            except Exception as e:  # noqa: BLE001
+                logger.warning('TeacherReplica shutdown failed: %s', e)
+        self.teacher_replicas = []
+
+        seen: set = set()
+        for wg in self.worker_groups.values():
+            if id(wg) not in seen:
+                seen.add(id(wg))
+                try:
+                    wg.shutdown()
+                except Exception as e:  # noqa: BLE001
+                    logger.warning('WorkerGroup shutdown failed: %s', e)
+        self.worker_groups.clear()
+
+        if self.resource_pool_manager is not None:
+            try:
+                self.resource_pool_manager.destroy_all()
+            except Exception as e:  # noqa: BLE001
+                logger.warning('destroy_all placement groups failed: %s', e)
+        ray.shutdown()
+
+
+def main():
+    import sys
+    argv = sys.argv[1:]
+    config_path = None
+    for i, arg in enumerate(argv):
+        if arg == '--config' and i + 1 < len(argv):
+            config_path = argv[i + 1]
+            break
+    if config_path is None:
+        raise ValueError('Usage: python -m swift.ray.megatron.pipeline --config <yaml>')
+
+    return MegatronRayPipeline(config_path).run()
+
+
+if __name__ == '__main__':
+    main()

swift/ray/megatron/resource_pool.py

@@ -0,0 +1,168 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Tuple
+
+from swift.utils import get_logger
+
+logger = get_logger()
+
+
+def sort_pgs_by_node_ip(pgs: List[Any]) -> List[Any]:
+    """Sort placement groups by node IP for deterministic rank assignment.
+
+    Sorts PGs by node IP for deterministic ordering.
+    """
+    import ray
+
+    node_ip = {node['NodeID']: node['NodeManagerAddress'] for node in ray.nodes()}
+    pg_ip = {}
+    for pg in pgs:
+        specs = ray._private.state.state.placement_group_table(pg.id)
+        node_id = specs['bundles_to_node_id'][0]
+        pg_ip[pg.id] = node_ip[node_id]
+    return sorted(pgs, key=lambda pg: pg_ip[pg.id])
+
+
+@dataclass
+class ResourcePool:
+    """A pool of GPU resources backed by multiple Ray placement groups.
+
+    Args:
+        process_on_nodes: GPUs per node. ``[8]`` = 8 GPUs on 1 node,
+            ``[4, 4]`` = 8 GPUs across 2 nodes.
+        max_colocate_count: How many WorkerGroups share these GPUs.
+    """
+
+    process_on_nodes: List[int]
+    max_colocate_count: int = 1
+
+    pgs: List[Any] = field(default_factory=list, repr=False, init=False)
+    node_ips: List[str] = field(default_factory=list, repr=False, init=False)
+    bundle_infos: List[Tuple[str, str]] = field(default_factory=list, repr=False, init=False)
+
+    @property
+    def world_size(self) -> int:
+        return sum(self.process_on_nodes)
+
+    @property
+    def num_nodes(self) -> int:
+        return len(self.process_on_nodes)
+
+    @property
+    def visible_devices(self) -> List[int]:
+        """Physical GPU ordinals: flat list across all nodes."""
+        return [int(info[1]) if info[1] else i for i, info in enumerate(self.bundle_infos)]
+
+    def create(self, device_name: str = 'GPU'):
+        """Create one PG per node with STRICT_PACK strategy."""
+        import ray
+        from ray.util.placement_group import placement_group
+
+        if device_name == 'npu':
+            device_name = 'NPU'
+        elif device_name == 'cuda':
+            device_name = 'GPU'
+
+        bundle_template = {
+            device_name: 1,
+            'CPU': max(self.max_colocate_count, 1),
+        }
+
+        pgs = []
+        for n_gpus in self.process_on_nodes:
+            bundles = [bundle_template.copy() for _ in range(n_gpus)]
+            pg = placement_group(bundles, strategy='STRICT_PACK')
+            pgs.append(pg)
+
+        ray.get([pg.ready() for pg in pgs])
+
+        self.pgs = sort_pgs_by_node_ip(pgs)
+        self._discover_bundle_infos()
+
+    def _discover_bundle_infos(self):
+        """Probe each bundle's accelerator_id via lightweight actors.
+
+        Swift-specific: needed because Swift uses num_gpus=0 + explicit
+        CUDA_VISIBLE_DEVICES (torchrun style).
+        """
+        import os
+        import ray
+        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+        from transformers.utils import is_torch_npu_available
+
+        @ray.remote(num_gpus=0.01, num_cpus=0.01)
+        def _probe_bundle():
+            ctx = ray.get_runtime_context()
+            acc_ids = ctx.get_accelerator_ids()
+            gpu_id = ''
+            for key in ('GPU', 'NPU'):
+                ids = acc_ids.get(key, [])
+                if ids:
+                    gpu_id = ids[0]
+                    break
+            return ctx.get_node_id(), gpu_id
+
+        all_infos: List[Tuple[str, str]] = []
+        node_id_to_ip = {node['NodeID']: node['NodeManagerAddress'] for node in ray.nodes()}
+
+        for pg_idx, pg in enumerate(self.pgs):
+            refs = [
+                _probe_bundle.options(
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(
+                        placement_group=pg, placement_group_bundle_index=i), ).remote()
+                for i in range(self.process_on_nodes[pg_idx])
+            ]
+            results = ray.get(refs)
+            for r in results:
+                all_infos.append(r)
+
+        vis_key = 'ASCEND_RT_VISIBLE_DEVICES' if is_torch_npu_available() else 'CUDA_VISIBLE_DEVICES'
+        parent_cvd = os.environ.get(vis_key, '')
+        if parent_cvd:
+            phys_ids = [x.strip() for x in parent_cvd.split(',')]
+            all_infos = [(nid, phys_ids[int(gid)] if gid.isdigit() and int(gid) < len(phys_ids) else gid)
+                         for nid, gid in all_infos]
+
+        self.bundle_infos = all_infos
+        seen: set = set()
+        node_ips = []
+        for nid, _ in all_infos:
+            if nid not in seen:
+                seen.add(nid)
+                node_ips.append(node_id_to_ip.get(nid, ''))
+        self.node_ips = node_ips
+        logger.info('ResourcePool: %d PG(s), %d bundles, node_ips=%s', len(self.pgs), len(all_infos), self.node_ips)
+
+    def destroy(self):
+        if self.pgs:
+            import ray
+            for pg in self.pgs:
+                try:
+                    ray.util.remove_placement_group(pg)
+                except Exception:  # noqa: BLE001
+                    pass
+            self.pgs = []
+
+
+class ResourcePoolManager:
+    """Manages multiple ResourcePools, deduplicating shared pools (colocate)."""
+
+    def __init__(self, pool_mapping: Dict[str, 'ResourcePool']):
+        self._pools = pool_mapping
+
+    def get_pool(self, group_name: str) -> 'ResourcePool':
+        return self._pools[group_name]
+
+    def create_all(self):
+        seen: set = set()
+        for pool in self._pools.values():
+            if id(pool) not in seen:
+                seen.add(id(pool))
+                pool.create()
+
+    def destroy_all(self):
+        seen: set = set()
+        for pool in self._pools.values():
+            if id(pool) not in seen:
+                seen.add(id(pool))
+                pool.destroy()

swift/ray/megatron/rollout/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .adapter import RolloutAdapter
+    from .ray_vllm_engine import RayVllmEngine
+    from .replica import RolloutMode, RolloutReplica, VllmEngineConfig
+    from .vllm_server import VllmServer
+    from .weight_transfer import BucketedWeightSender
+
+
+def __getattr__(name):
+    _imports = {
+        'RolloutAdapter': '.adapter',
+        'RayVllmEngine': '.ray_vllm_engine',
+        'RolloutMode': '.replica',
+        'RolloutReplica': '.replica',
+        'VllmEngineConfig': '.replica',
+        'VllmServer': '.vllm_server',
+        'BucketedWeightSender': '.weight_transfer',
+    }
+    if name in _imports:
+        import importlib
+        return getattr(importlib.import_module(_imports[name], __name__), name)
+    raise AttributeError(f'module {__name__!r} has no attribute {name!r}')

swift/ray/megatron/rollout/adapter.py

@@ -0,0 +1,123 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+from __future__ import annotations
+
+import asyncio
+import logging
+import ray
+import time
+import torch
+from typing import Any, Generator, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+def _is_ipc_supported() -> bool:
+    """Check if CUDA IPC is supported (GPU=True, NPU=fallback to SHM)."""
+    from transformers.utils import is_torch_npu_available
+    return not is_torch_npu_available()
+
+
+class RolloutAdapter:
+
+    def __init__(
+        self,
+        *,
+        replica_rank: int = 0,
+        rollout_rank: int = 0,
+        bucket_size_mb: int = 2048,
+    ):
+        self.replica_rank = replica_rank
+        self.rollout_rank = rollout_rank
+        self.bucket_size_mb = bucket_size_mb
+        self.is_primary = (rollout_rank == 0)
+        self.use_shm = not _is_ipc_supported()
+        self.zmq_handle = (f'ipc:///tmp/swift-rollout-zmq-replica-{replica_rank}-rank-{rollout_rank}.sock')
+        self._server_handle = None
+        # Persistent CUDA IPC buffer reused across all update_weights() syncs so the IPC
+        # handle stays stable (the vLLM worker's mapping cache hits) and no IPC mapping
+        # leaks per step.
+        self._ipc_buffer: Optional[torch.Tensor] = None
+
+    @property
+    def server_handle(self) -> Any:
+        if self._server_handle is None:
+            self._server_handle = ray.get_actor(f'swift_rollout_server_{self.replica_rank}_0')
+        return self._server_handle
+
+    def update_weights(
+        self,
+        weight_iter: Generator[Tuple[str, torch.Tensor], None, None],
+        *,
+        vllm_lora_param_names: Optional[set] = None,
+        peft_config: Optional[dict] = None,
+        base_sync_done: bool = False,
+    ) -> None:
+        """Push training weights to vLLM via ZMQ IPC.
+
+        Only the primary rank (rollout_rank == 0) sends weights;
+        other ranks are no-op.
+
+        Args:
+            weight_iter: Iterator of (name, tensor) from bridge.export_weights.
+            vllm_lora_param_names: When set, remap dense param names to
+                LoRA-wrapped names (*.base_layer.weight) for full-weight sync
+                with vllm_enable_lora.
+            peft_config: When provided with base_sync_done=True, vLLM loads
+                weights as a LoRA adapter via TensorLoRARequest.
+            base_sync_done: Indicates this is an adapter-only sync after
+                the initial full weight sync has completed.
+        """
+        if vllm_lora_param_names:
+            from swift.rlhf_trainers.utils import add_base_layer_suffix_by_param_names
+            weight_iter = add_base_layer_suffix_by_param_names(weight_iter, vllm_lora_param_names)
+
+        if not self.is_primary:
+            for _ in weight_iter:
+                pass
+            return
+
+        from ..rollout.weight_transfer import BucketedWeightSender
+
+        start_time = time.time()
+
+        # Lazily (re)allocate the persistent IPC buffer; reused across syncs so the
+        # handle signature stays stable and the worker-side IPC cache hits.
+        external_buffer = None
+        if not self.use_shm:
+            from swift.utils import get_current_device
+            bucket_size = self.bucket_size_mb << 20
+            if self._ipc_buffer is None or self._ipc_buffer.numel() < bucket_size:
+                self._ipc_buffer = torch.empty(bucket_size, dtype=torch.uint8, device=get_current_device())
+            external_buffer = self._ipc_buffer
+
+        async def _do_ipc_sync():
+            sender = BucketedWeightSender(
+                zmq_handle=self.zmq_handle,
+                bucket_size_mb=self.bucket_size_mb,
+                use_shm=self.use_shm,
+                external_buffer=external_buffer,
+            )
+            try:
+                async with sender:
+                    rpc_ref = self.server_handle.update_weights_ipc.remote(self.zmq_handle, self.use_shm, 600,
+                                                                           peft_config, base_sync_done)
+                    await sender.handshake()
+                    await sender.send_weights(weight_iter)
+                    await asyncio.get_running_loop().run_in_executor(None, ray.get, rpc_ref)
+            finally:
+                sender.cleanup()
+
+        asyncio.run(_do_ipc_sync())
+        logger.debug('RolloutAdapter: update_weights done (replica=%d, adapter_only=%s, %.2fs)', self.replica_rank,
+                     base_sync_done,
+                     time.time() - start_time)
+
+    def reset_prefix_cache(self) -> None:
+        if not self.is_primary:
+            return
+        ray.get(self.server_handle.reset_prefix_cache.remote())
+
+    def get_model_param_names(self) -> List[str]:
+        if not self.is_primary:
+            return []
+        return ray.get(self.server_handle.get_model_param_names.remote()) or []

swift/ray/megatron/rollout/ray_vllm_engine.py

@@ -0,0 +1,228 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""RayVllmEngine — vLLM engine wrapper for Ray rollout actors.
+
+Reuses ``VllmEngine``'s template / processor initialisation but creates
+the ``AsyncLLM`` engine inside a **persistent event-loop thread** so that
+vLLM's ZMQ-based ``AsyncMPClient`` (which caches event-loop-bound tasks
+and sockets) stays consistent across all async operations.
+
+``VllmServer`` holds a ``RayVllmEngine`` instance and delegates all
+engine operations to it.
+"""
+import asyncio
+import os
+import threading
+import torch
+from typing import Any, Dict, List, Optional
+
+from swift.rlhf_trainers.utils import set_expandable_segments
+from swift.utils import gc_collect
+from swift.utils.logger import get_logger
+
+logger = get_logger()
+
+
+class RayVllmEngine:
+
+    def __init__(
+        self,
+        model_id: str,
+        *,
+        tensor_parallel_size: int = 1,
+        gpu_memory_utilization: float = 0.9,
+        max_model_len: Optional[int] = None,
+        max_num_seqs: int = 256,
+        enable_sleep_mode: bool = False,
+        enable_lora: bool = False,
+        max_lora_rank: int = 8,
+        enable_prefix_caching: bool = False,
+        enforce_eager: bool = False,
+        trust_remote_code: bool = True,
+        dtype: str = 'auto',
+        load_format: str = 'auto',
+        template_kwargs: Optional[Dict[str, Any]] = None,
+        **engine_kwargs,
+    ):
+        os.environ.setdefault('VLLM_USE_V1', '1')
+        os.environ.setdefault('VLLM_WORKER_MULTIPROC_METHOD', 'spawn')
+        os.environ.setdefault('VLLM_ENGINE_ITERATION_TIMEOUT_S', '86400')
+
+        self.model_id = model_id
+        self.tp_size = tensor_parallel_size
+        self.enable_sleep_mode = enable_sleep_mode
+        self.enable_lora = enable_lora
+
+        distributed_executor_backend = 'mp' if tensor_parallel_size > 1 else None
+
+        extra_engine_kwargs = dict(engine_kwargs)
+        extra_engine_kwargs['worker_extension_cls'] = ('swift.pipelines.infer.rollout.WeightSyncWorkerExtension')
+
+        from swift.model import get_processor
+        from swift.template import get_template
+        processor = get_processor(model_id_or_path=model_id, download_model=True)
+        template = get_template(processor, **(template_kwargs or {}))
+        self.template = template
+        self.tokenizer = template.tokenizer
+
+        # --- Start persistent event loop thread ---
+        self._loop = asyncio.new_event_loop()
+        self._loop_thread = threading.Thread(
+            target=self._run_event_loop,
+            daemon=True,
+            name='RayVllmEngine-EventLoop',
+        )
+        self._loop_dead = threading.Event()
+        self._loop_exception: Optional[BaseException] = None
+        self._loop_thread.start()
+
+        # --- Create VllmEngine (with engine) inside the event loop ---
+        self._vllm_engine = self._run_in_loop(
+            self._create_engine_async(
+                model_id,
+                template=template,
+                tensor_parallel_size=tensor_parallel_size,
+                gpu_memory_utilization=gpu_memory_utilization,
+                max_model_len=max_model_len,
+                max_num_seqs=max_num_seqs,
+                enable_sleep_mode=enable_sleep_mode,
+                enable_lora=enable_lora,
+                max_lora_rank=max_lora_rank,
+                enable_prefix_caching=enable_prefix_caching,
+                enforce_eager=enforce_eager,
+                load_format=load_format,
+                distributed_executor_backend=distributed_executor_backend,
+                logprobs_mode='processed_logprobs',
+                engine_kwargs=extra_engine_kwargs,
+            ))
+        self.engine = self._vllm_engine.engine
+
+        logger.info(
+            'RayVllmEngine: ready (model=%s, tp=%d, sleep=%s, load_format=%s)',
+            model_id,
+            tensor_parallel_size,
+            enable_sleep_mode,
+            load_format,
+        )
+
+    @staticmethod
+    async def _create_engine_async(model_id: str, *, template, **kwargs):
+        from swift.infer_engine import VllmEngine
+        engine = VllmEngine(
+            model_id,
+            use_async_engine=True,
+            template=template,
+            **kwargs,
+        )
+        return engine
+
+    def _run_event_loop(self):
+        asyncio.set_event_loop(self._loop)
+        try:
+            self._loop.run_forever()
+        except Exception as exc:
+            logger.exception('RayVllmEngine event loop died unexpectedly')
+            self._loop_exception = exc
+            self._loop_dead.set()
+
+    def _run_in_loop(self, coro):
+        if self._loop_dead.is_set():
+            cause = self._loop_exception
+            raise RuntimeError('RayVllmEngine event loop is no longer running. '
+                               f'Original error: {type(cause).__name__}: {cause}'
+                               if cause else 'RayVllmEngine event loop is no longer running') from cause
+        future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+        return future.result()
+
+    # ------------------------------------------------------------------
+    # Sleep / Wake-up
+    # ------------------------------------------------------------------
+
+    def sleep(self, level: int = 2):
+        if not self.enable_sleep_mode:
+            return
+        self._run_in_loop(self.engine.sleep(level=level))
+        gc_collect()
+        set_expandable_segments(True)
+        logger.debug('RayVllmEngine: sleeping at level %d', level)
+
+    def wake_up(self, tags: Optional[List[str]] = None):
+        if not self.enable_sleep_mode:
+            return
+        if tags is None or 'kv_cache' in tags:
+            gc_collect()
+            set_expandable_segments(False)
+        self._run_in_loop(self.engine.wake_up(tags=tags))
+        logger.debug('RayVllmEngine: woke up with tags %s', tags)
+
+    def generate_batch(
+        self,
+        infer_requests: List[Any],
+        request_config: Optional[Any] = None,
+    ) -> List[Any]:
+        from swift.infer_engine.protocol import RequestConfig
+        if request_config is None:
+            request_config = RequestConfig()
+
+        async def _gen():
+            tasks = [self._vllm_engine.infer_async(req, request_config) for req in infer_requests]
+            return await asyncio.gather(*tasks)
+
+        return list(self._run_in_loop(_gen()))
+
+    def get_model_param_names(self) -> List[str]:
+        """Return parameter names from vLLM model via collective_rpc."""
+
+        async def _get():
+            result = await self.engine.collective_rpc('get_state_keys')
+            if result and isinstance(result[0], list):
+                return result[0]
+            return []
+
+        return self._run_in_loop(_get())
+
+    def reset_prefix_cache(self):
+        self._run_in_loop(self.engine.reset_prefix_cache())
+
+    def update_weights_ipc(
+        self,
+        zmq_handle: str,
+        use_shm: bool = False,
+        timeout_s: int = 600,
+        peft_config: Optional[dict] = None,
+        base_sync_done: bool = False,
+    ):
+        """Trigger the vLLM worker extension's ``update_weights_from_ipc``."""
+
+        async def _rpc():
+            rpc_kwargs = {
+                'use_shm': use_shm,
+                'zmq_handle': zmq_handle,
+                'timeout_s': timeout_s,
+            }
+            if peft_config is not None and base_sync_done:
+                rpc_kwargs['peft_config'] = peft_config
+                rpc_kwargs['base_sync_done'] = base_sync_done
+            return await asyncio.wait_for(
+                self.engine.collective_rpc('update_weights_from_ipc', kwargs=rpc_kwargs),
+                timeout=timeout_s,
+            )
+
+        self._run_in_loop(_rpc())
+
+    def shutdown(self):
+        if self.engine is not None:
+            try:
+                self.engine.shutdown()
+            except Exception as e:  # noqa: BLE001
+                logger.warning('RayVllmEngine shutdown error: %s', e)
+            self.engine = None
+
+        if self._loop is not None:
+            try:
+                self._loop.call_soon_threadsafe(self._loop.stop)
+                if self._loop_thread is not None:
+                    self._loop_thread.join(timeout=5)
+            except Exception:
+                pass
+
+        gc_collect()

swift/ray/megatron/rollout/replica.py

@@ -0,0 +1,284 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import ray
+from collections import defaultdict
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+from swift.rlhf_trainers.args_mixin import VllmArguments
+from swift.utils.logger import get_logger
+
+if TYPE_CHECKING:
+    from ..resource_pool import ResourcePool
+
+logger = get_logger()
+
+
+class RolloutMode(str, Enum):
+    HYBRID = 'hybrid'
+    STANDALONE = 'standalone'
+
+
+@dataclass
+class VllmEngineConfig(VllmArguments):
+    model: str = ''
+    sleep_level: int = 0
+    vllm_enable_lora: bool = False
+    trust_remote_code: bool = True
+    dtype: str = 'auto'
+    load_format: str = 'auto'
+
+    # override
+    vllm_enable_prefix_caching: bool = True
+
+    def __post_init__(self):
+        VllmArguments.__post_init__(self)
+
+    @classmethod
+    def from_rollout_cfg(cls, rollout_cfg: Dict[str, Any], *, sleep_level: int = 0) -> 'VllmEngineConfig':
+        """Build from merged rollout config dict."""
+        cfg = rollout_cfg or {}
+        known_fields = {f.name for f in cls.__dataclass_fields__.values()}
+
+        kwargs: Dict[str, Any] = {}
+        for key, val in cfg.items():
+            if key in known_fields and val is not None:
+                kwargs[key] = val
+
+        if sleep_level > 0:
+            kwargs['sleep_level'] = sleep_level
+
+        if cfg.get('tuner_type', 'full') == 'lora':
+            kwargs.setdefault('vllm_enable_lora', True)
+
+        return cls(**kwargs)
+
+    def to_launch_kwargs(self, rollout_mode: str, template_kwargs: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        kw: Dict[str, Any] = {
+            'model_id': self.model,
+            'dtype': self.dtype,
+            'rollout_mode': rollout_mode,
+            'tensor_parallel_size': self.vllm_tensor_parallel_size,
+            'gpu_memory_utilization': self.vllm_gpu_memory_utilization,
+            'max_num_seqs': self.vllm_max_num_seqs,
+            'enforce_eager': self.vllm_enforce_eager,
+            'trust_remote_code': self.trust_remote_code,
+            'load_format': self.load_format,
+            'enable_sleep_mode': (self.sleep_level > 0),
+            'enable_lora': self.vllm_enable_lora,
+            'max_lora_rank': self.vllm_max_lora_rank,
+            'data_parallel_size': self.vllm_data_parallel_size,
+            'enable_prefix_caching': self.vllm_enable_prefix_caching,
+        }
+        if self.vllm_max_model_len is not None:
+            kw['max_model_len'] = self.vllm_max_model_len
+        if template_kwargs:
+            kw['template_kwargs'] = template_kwargs
+        extra = self.vllm_engine_kwargs
+        if isinstance(extra, dict):
+            kw.update(extra)
+        return kw
+
+
+class RolloutReplica:
+    """One vLLM rollout replica on Ray (single or multi-node)."""
+
+    def __init__(
+        self,
+        config: VllmEngineConfig,
+        mode: RolloutMode = RolloutMode.HYBRID,
+        replica_rank: int = 0,
+        template_kwargs: Optional[Dict[str, Any]] = None,
+        actor_name_prefix: str = 'swift_rollout_server',
+    ) -> None:
+        self.config = config
+        self.mode = mode
+        self.replica_rank = replica_rank
+        self.template_kwargs = template_kwargs
+        self.actor_name_prefix = actor_name_prefix
+        self._servers: List[Any] = []
+
+    @classmethod
+    def create_replicas(
+        cls,
+        rollout_cfg: Dict[str, Any],
+        rollout_gpus: int,
+        pool: 'ResourcePool',
+        is_hybrid: bool,
+        sleep_level: int = 0,
+        template_kwargs: Optional[Dict[str, Any]] = None,
+        actor_name_prefix: str = 'swift_rollout_server',
+    ) -> List['RolloutReplica']:
+        """Factory: create all rollout replicas from pipeline config.
+
+        Uses two-phase initialization for parallelism:
+        Phase 1 — spawn all Ray actors (fast, non-blocking)
+        Phase 2 — launch_server on all replicas concurrently
+        """
+        config = VllmEngineConfig.from_rollout_cfg(rollout_cfg, sleep_level=sleep_level)
+        world_size_per_replica = config.vllm_tensor_parallel_size * config.vllm_data_parallel_size
+        if world_size_per_replica > rollout_gpus:
+            raise ValueError(f'tp*dp ({world_size_per_replica}) exceeds rollout GPUs ({rollout_gpus})')
+        if rollout_gpus % world_size_per_replica != 0:
+            raise ValueError(f'rollout GPUs ({rollout_gpus}) must be divisible by '
+                             f'tp*dp ({world_size_per_replica})')
+        n_replicas = rollout_gpus // world_size_per_replica
+        mode = RolloutMode.HYBRID if is_hybrid else RolloutMode.STANDALONE
+
+        replicas: List['RolloutReplica'] = []
+        bundle_infos = pool.bundle_infos
+        for i in range(n_replicas):
+            offset = i * world_size_per_replica
+            replica_infos = bundle_infos[offset:offset + world_size_per_replica]
+            nodes = {info[0] for info in replica_infos}
+            gpus_per_node = (world_size_per_replica if len(nodes) == 1 else world_size_per_replica // len(nodes))
+            replica = cls(
+                config, mode=mode, replica_rank=i, template_kwargs=template_kwargs, actor_name_prefix=actor_name_prefix)
+            replica._spawn_actors(replica_infos, gpus_per_node)
+            replicas.append(replica)
+
+        cls._parallel_launch_all(replicas)
+
+        logger.info('Rollout: %d replica(s) in %s mode (tp=%d, dp=%d, total_gpus=%d)', n_replicas, mode.value.upper(),
+                    config.vllm_tensor_parallel_size, config.vllm_data_parallel_size, rollout_gpus)
+        return replicas
+
+    @classmethod
+    def _parallel_launch_all(cls, replicas: List['RolloutReplica']) -> None:
+        """Phase 2: launch vLLM engines on all replicas in parallel."""
+        all_refs = []
+        for replica in replicas:
+            refs = replica._launch_engines_async()
+            all_refs.extend(refs)
+        if all_refs:
+            ray.get(all_refs)
+            for replica in replicas:
+                logger.info('RolloutReplica[replica=%d, mode=%s]: launched %d server(s) (tp=%d, model=%s)',
+                            replica.replica_rank, replica.mode.value, len(replica._servers),
+                            replica.config.vllm_tensor_parallel_size, replica.config.model)
+
+    def _spawn_actors(
+        self,
+        worker_infos: List[Tuple[str, str]],
+        gpus_per_node: int,
+    ) -> None:
+        """Phase 1: create Ray actors without starting engines.
+
+        ``num_gpus=0`` + ``NOSET_CVD`` + explicit visible-device env so
+        the actor sees exactly the GPUs from pool bundles; NodeAffinity
+        pins each actor to the correct node.
+        """
+        from ray.runtime_env import RuntimeEnv
+        from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+        from transformers.utils import is_torch_npu_available
+
+        from .vllm_server import VllmServer
+
+        visible_key = 'ASCEND_RT_VISIBLE_DEVICES' if is_torch_npu_available() else 'CUDA_VISIBLE_DEVICES'
+
+        node_groups = self._group_by_node(worker_infos)
+        self._nnodes = len(node_groups)
+
+        actor_cls = ray.remote(num_gpus=0, num_cpus=1)(VllmServer)
+
+        for node_rank, (node_id, gpu_ids) in enumerate(node_groups):
+            cvd = ','.join(gpu_ids)
+            env_vars: Dict[str, str] = {
+                'VLLM_USE_V1': '1',
+                'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES': '1',
+                'RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES': '1',
+                'RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES': '1',
+                'RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES': '1',
+                'NCCL_CUMEM_ENABLE': '0',
+                visible_key: cvd,
+            }
+            handle = actor_cls.options(
+                scheduling_strategy=NodeAffinitySchedulingStrategy(node_id=node_id, soft=False),
+                runtime_env=RuntimeEnv(env_vars=env_vars),
+                name=f'{self.actor_name_prefix}_{self.replica_rank}_{node_rank}',
+                max_concurrency=10,
+            ).remote(
+                node_rank=node_rank,
+                nnodes=self._nnodes,
+                gpus_per_node=gpus_per_node,
+                cuda_visible_devices=cvd,
+            )
+            self._servers.append(handle)
+
+    def _launch_engines_async(self) -> List[ray.ObjectRef]:
+        """Phase 2: issue launch_server calls, return ObjectRefs (non-blocking)."""
+        launch_kw = self.config.to_launch_kwargs(self.mode.value, template_kwargs=self.template_kwargs)
+        nnodes = getattr(self, '_nnodes', 1)
+
+        if nnodes > 1:
+            master_address, master_port, dp_rpc_port = ray.get(self._servers[0].get_master_address.remote())
+            refs = [
+                server.launch_server.remote(
+                    master_address=master_address, master_port=master_port, dp_rpc_port=dp_rpc_port, **launch_kw)
+                for server in self._servers
+            ]
+        else:
+            refs = [self._servers[0].launch_server.remote(**launch_kw)]
+        return refs
+
+    @staticmethod
+    def _group_by_node(worker_infos: List[Tuple[str, str]]) -> List[Tuple[str, List[str]]]:
+        """Group worker infos by node, preserving order.
+
+        Returns list of ``(node_id, [accelerator_id, ...])`` in node
+        encounter order.
+        """
+        ordered_nodes: List[str] = []
+        node_gpus: Dict[str, List[str]] = defaultdict(list)
+        for node_id, acc_id in worker_infos:
+            if node_id not in node_gpus:
+                ordered_nodes.append(node_id)
+            node_gpus[node_id].append(acc_id)
+        return [(nid, node_gpus[nid]) for nid in ordered_nodes]
+
+    @property
+    def primary(self) -> Any:
+        """The node_rank=0 ``VllmServer`` actor handle.
+
+        Callers that need ``sleep`` / ``wake_up`` / ``reset_prefix_cache``
+        / ``update_weights_ipc`` / ``update_weights_direct`` talk to
+        this handle directly.
+        """
+        if not self._servers:
+            raise RuntimeError('RolloutReplica: not launched yet')
+        return self._servers[0]
+
+    @property
+    def servers(self) -> List[Any]:
+        """All server actor handles (one per node)."""
+        return list(self._servers)
+
+    def sleep(self, level: int = 1) -> None:
+        """Put all servers in this replica to sleep."""
+        refs = [server.sleep.remote(level) for server in self._servers]
+        ray.get(refs)
+
+    def wake_up(self, tags=None) -> None:
+        """Wake all servers in this replica."""
+        refs = [server.wake_up.remote(tags=tags) for server in self._servers]
+        ray.get(refs)
+
+    def generate(
+        self,
+        infer_requests: List[Any],
+        request_config: Any = None,
+    ) -> ray.ObjectRef:
+        """Submit generation to the primary server, returns an ObjectRef."""
+        return self.primary.generate.remote(infer_requests, request_config)
+
+    def shutdown(self) -> None:
+        for server in self._servers:
+            try:
+                ray.get(server.shutdown.remote(), timeout=30)
+            except Exception as e:  # noqa: BLE001
+                logger.warning('RolloutReplica shutdown error: %s', e)
+            try:
+                ray.kill(server, no_restart=True)
+            except Exception:  # noqa: BLE001
+                pass
+        self._servers = []

swift/ray/megatron/rollout/vllm_server.py

@@ -0,0 +1,251 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""VllmServer — Ray Actor hosting a vLLM engine via :class:`RayVllmEngine`.
+
+Thin Ray-actor shell that delegates to :class:`RayVllmEngine` for all
+engine operations.  ``RolloutReplica`` wraps this class with
+``ray.remote`` when spawning real actors.
+
+Multi-node topology::
+
+    Node 0 (node_rank=0)
+    ├── MegatronWorker actors        — training processes
+    └── VllmServer actor (primary)   — runs full server + API
+
+    Node 1 (node_rank=1)
+    ├── MegatronWorker actors
+    └── VllmServer actor (headless)  — participates in TP, no API
+
+"""
+import asyncio
+import os
+import socket
+import torch
+from transformers.utils import is_torch_npu_available
+from typing import Any, Dict, List, Optional, Tuple
+
+from swift.utils import gc_collect, get_logger
+from ..checkpoint_engine import CheckpointEngineMixin
+
+logger = get_logger()
+
+
+def _parse_bool_env(name: str, default: bool) -> bool:
+    val = os.environ.get(name)
+    if val is None:
+        return default
+    return val.lower() in ('1', 'true')
+
+
+def _get_free_port(address: str = '') -> int:
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind((address, 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+class VllmServer(CheckpointEngineMixin):
+
+    def __init__(
+        self,
+        node_rank: int = 0,
+        nnodes: int = 1,
+        gpus_per_node: int = 8,
+        cuda_visible_devices: str = '',
+    ) -> None:
+        self._engine = None
+        self._node_rank = node_rank
+        self._nnodes = nnodes
+        self._gpus_per_node = gpus_per_node
+
+        if cuda_visible_devices:
+            key = 'ASCEND_RT_VISIBLE_DEVICES' if is_torch_npu_available() else 'CUDA_VISIBLE_DEVICES'
+            os.environ[key] = cuda_visible_devices
+
+        self._server_address = None
+        self._master_address: Optional[str] = None
+        self._master_port: Optional[int] = None
+        self._dp_rpc_port: Optional[int] = None
+
+        if node_rank == 0:
+            import ray as _ray
+            self._server_address = _ray.util.get_node_ip_address()
+            self._master_address = self._server_address
+            self._master_port = _get_free_port(self._server_address)
+            self._dp_rpc_port = _get_free_port(self._server_address)
+
+    def get_master_address(self) -> Tuple[str, int, int]:
+        """Return ``(master_address, master_port, dp_rpc_port)`` from node_rank=0."""
+        return self._master_address, self._master_port, self._dp_rpc_port
+
+    def launch_server(
+        self,
+        *,
+        model_id: str,
+        rollout_mode: str = 'hybrid',
+        tensor_parallel_size: int = 1,
+        gpu_memory_utilization: float = 0.9,
+        max_model_len: Optional[int] = None,
+        max_num_seqs: int = 256,
+        enable_sleep_mode: bool = False,
+        enable_lora: bool = False,
+        max_lora_rank: int = 8,
+        enable_prefix_caching: bool = False,
+        enforce_eager: bool = False,
+        trust_remote_code: bool = True,
+        dtype: str = 'auto',
+        load_format: str = 'auto',
+        master_address: Optional[str] = None,
+        master_port: Optional[int] = None,
+        dp_rpc_port: Optional[int] = None,
+        data_parallel_size: int = 1,
+        template_kwargs: Optional[Dict[str, Any]] = None,
+        **engine_kwargs,
+    ) -> Dict[str, Any]:
+        if self._node_rank != 0:
+            self._master_address = master_address
+            self._master_port = master_port
+            self._dp_rpc_port = dp_rpc_port
+            import ray as _ray
+            self._server_address = _ray.util.get_node_ip_address()
+
+        extra_engine_kwargs = dict(engine_kwargs)
+
+        if self._nnodes > 1:
+            extra_engine_kwargs['nnodes'] = self._nnodes
+            extra_engine_kwargs['node_rank'] = self._node_rank
+            extra_engine_kwargs['master_addr'] = self._master_address
+            extra_engine_kwargs['master_port'] = self._master_port
+
+        if data_parallel_size > 1:
+            assert self._gpus_per_node % tensor_parallel_size == 0, (
+                'gpus_per_node should be divisible by vllm_tensor_parallel_size')
+            dp_size_local = self._gpus_per_node // tensor_parallel_size
+            extra_engine_kwargs['data_parallel_size'] = data_parallel_size
+            extra_engine_kwargs['data_parallel_size_local'] = dp_size_local
+            extra_engine_kwargs['data_parallel_start_rank'] = self._node_rank * dp_size_local
+            extra_engine_kwargs['data_parallel_address'] = self._master_address
+            extra_engine_kwargs['data_parallel_rpc_port'] = self._dp_rpc_port
+
+        from .ray_vllm_engine import RayVllmEngine
+
+        self._engine = RayVllmEngine(
+            model_id,
+            tensor_parallel_size=tensor_parallel_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+            max_model_len=max_model_len,
+            max_num_seqs=max_num_seqs,
+            enable_sleep_mode=enable_sleep_mode,
+            enable_lora=enable_lora,
+            max_lora_rank=max_lora_rank,
+            enable_prefix_caching=enable_prefix_caching,
+            enforce_eager=enforce_eager,
+            trust_remote_code=trust_remote_code,
+            dtype=dtype,
+            load_format=load_format,
+            template_kwargs=template_kwargs,
+            **extra_engine_kwargs,
+        )
+        logger.info(
+            'VllmServer[mode=%s, node_rank=%d/%d]: engine initialized (model=%s, tp=%d)',
+            rollout_mode,
+            self._node_rank,
+            self._nnodes,
+            model_id,
+            tensor_parallel_size,
+        )
+        return {
+            'model_id': model_id,
+            'tp_size': tensor_parallel_size,
+            'node_rank': self._node_rank,
+        }
+
+    def generate(
+        self,
+        infer_requests: List[Any],
+        request_config: Any = None,
+    ) -> List[Any]:
+        return self._engine.generate_batch(infer_requests, request_config)
+
+    def sleep(self, level: int = 2) -> None:
+        self._engine.sleep(level=level)
+
+    def wake_up(self, tags: Optional[List[str]] = None) -> None:
+        self._engine.wake_up(tags=tags)
+
+    def get_model_param_names(self) -> List[str]:
+        return self._engine.get_model_param_names()
+
+    def reset_prefix_cache(self) -> None:
+        self._engine.reset_prefix_cache()
+
+    def update_weights_ipc(
+        self,
+        zmq_handle: str,
+        use_shm: bool = False,
+        timeout_s: int = 600,
+        peft_config: Optional[dict] = None,
+        base_sync_done: bool = False,
+    ) -> None:
+        self._engine.update_weights_ipc(
+            zmq_handle, use_shm=use_shm, timeout_s=timeout_s, peft_config=peft_config, base_sync_done=base_sync_done)
+
+    def receive_checkpoint_weights(
+        self,
+        base_sync_done: bool = False,
+        peft_config: Optional[dict] = None,
+    ) -> None:
+        """Receive weights via NCCL broadcast and stream into vLLM."""
+        engine = self._get_or_create_checkpoint_engine()
+        # CUDA defaults to IPC (SHM disabled) to avoid frequent SharedMemory
+        # warnings and long-run instability. NPU keeps SHM as default.
+        use_shm = _parse_bool_env('SWIFT_RAY_NCCL_RECV_USE_SHM', default=is_torch_npu_available())
+
+        async def _receive_and_load():
+            from .weight_transfer import BucketedWeightSender
+            zmq_handle = f'ipc:///tmp/swift-nccl-recv-{os.getpid()}.sock'
+            bucket_mb = int(os.environ.get('SWIFT_RAY_WEIGHT_BUCKET_MB', '2048'))
+            sender = BucketedWeightSender(
+                zmq_handle=zmq_handle,
+                bucket_size_mb=bucket_mb,
+                use_shm=use_shm,
+            )
+
+            async def _weight_stream():
+                async for name, tensor in engine.receive_weights():
+                    if use_shm and tensor.device.type != 'cpu':
+                        tensor = tensor.cpu()
+                    yield name, tensor
+
+            try:
+                async with sender:
+                    rpc_kwargs = {
+                        'use_shm': use_shm,
+                        'zmq_handle': zmq_handle,
+                        'timeout_s': 600,
+                    }
+                    if peft_config is not None and base_sync_done:
+                        rpc_kwargs['peft_config'] = peft_config
+                        rpc_kwargs['base_sync_done'] = base_sync_done
+                    rpc_task = asyncio.ensure_future(
+                        self._engine.engine.collective_rpc(
+                            'update_weights_from_ipc',
+                            kwargs=rpc_kwargs,
+                        ))
+                    # Allow collective_rpc to schedule and bind its ZMQ socket
+                    await asyncio.sleep(0)
+                    await sender.handshake()
+                    await sender.send_weights_async(_weight_stream())
+                    await rpc_task
+            finally:
+                sender.cleanup()
+
+        self._engine._run_in_loop(_receive_and_load())
+
+    def shutdown(self) -> None:
+        self._checkpoint_engine = None
+        if self._engine is not None:
+            self._engine.shutdown()
+            self._engine = None
+        gc_collect()

swift/ray/megatron/rollout/weight_transfer.py

@@ -0,0 +1,186 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+"""Weight transfer utilities for training → rollout weight synchronization.
+
+BucketedWeightSender is used by the **training** side (MegatronWorker
+via RolloutAdapter) to ship weights to vLLM's WeightSyncWorkerExtension
+through ZMQ IPC.  It was originally in vllm_server.py but belongs here
+because it is a training-side concern, not a rollout engine concern.
+"""
+from __future__ import annotations
+
+import asyncio
+import os
+import torch
+import uuid
+from typing import Any, Dict, Optional
+
+from swift.utils import get_current_device, synchronize
+from swift.utils.logger import get_logger
+
+logger = get_logger()
+
+
+class BucketedWeightSender:
+    """Streams model weights to vLLM worker via ZMQ IPC with bucketed transfer."""
+
+    def __init__(
+        self,
+        zmq_handle: str,
+        bucket_size_mb: int = 512,
+        use_shm: bool = False,
+        timeout_s: int = 600,
+        external_buffer: Optional[torch.Tensor] = None,
+    ):
+        self.zmq_handle = zmq_handle
+        self.bucket_size = int(bucket_size_mb) << 20
+        self.use_shm = use_shm
+        self.timeout_ms = int(timeout_s * 1000)
+        self.socket = None
+        self.buffer = None
+        self.shm = None
+        self._pending_handshake = None
+        # When provided, reuse this caller-owned persistent GPU buffer instead of
+        # allocating a fresh one per sync. Reusing the same storage keeps the CUDA IPC
+        # handle signature stable across sync rounds, so the vLLM worker's IPC-mapping
+        # cache hits and no new mapping is leaked each step.
+        self._external_buffer = external_buffer
+        self._owns_buffer = True
+
+    async def __aenter__(self):
+        self._init_socket_and_buffer()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self.cleanup()
+
+    def _init_socket_and_buffer(self):
+        """Bind the REQ socket and allocate the bucket buffer."""
+        import zmq
+
+        ctx = zmq.Context.instance()
+        self.socket = ctx.socket(zmq.REQ)
+        self.socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms)
+        self.socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms)
+        self.socket.setsockopt(zmq.LINGER, 0)
+        self.socket.bind(self.zmq_handle)
+
+        from torch.multiprocessing.reductions import reduce_tensor
+
+        if not self.use_shm:
+            if self._external_buffer is not None and self._external_buffer.numel() >= self.bucket_size:
+                # Reuse the persistent buffer -> stable IPC handle -> worker cache hit.
+                self.buffer = self._external_buffer
+                self._owns_buffer = False
+            else:
+                self.buffer = torch.empty(self.bucket_size, dtype=torch.uint8, device=get_current_device())
+                self._owns_buffer = True
+            self._pending_handshake = reduce_tensor(self.buffer)
+        else:
+            from multiprocessing import shared_memory
+            shm_name = f'swift_weights_{uuid.uuid4().hex}'
+            self.shm = shared_memory.SharedMemory(name=shm_name, create=True, size=self.bucket_size)
+            self.buffer = torch.frombuffer(self.shm.buf, dtype=torch.uint8)
+            self._pending_handshake = {'name': shm_name, 'size': self.bucket_size}
+
+    async def handshake(self):
+        if self._pending_handshake is None:
+            raise RuntimeError('BucketedWeightSender.handshake() called before enter or twice: '
+                               'the handshake payload is consumed on first call, a second handshake '
+                               'within the same ``async with`` block would ship stale metadata.')
+        loop = asyncio.get_running_loop()
+        payload = self._pending_handshake
+        self._pending_handshake = None
+
+        def _send_recv():
+            self.socket.send_pyobj(payload)
+            return self.socket.recv()
+
+        await loop.run_in_executor(None, _send_recv)
+
+    async def _stream_weights_inner(self, items_iter, is_async: bool = False):
+        """Shared bucketing logic for sync and async weight iterators."""
+        loop = asyncio.get_running_loop()
+        offset = 0
+        bucket_meta: Dict[str, Dict[str, Any]] = {}
+        n_weights = 0
+
+        def _zmq_send_recv(payload):
+            self.socket.send_pyobj(payload)
+            return self.socket.recv()
+
+        async def _flush(is_last: bool):
+            nonlocal offset, bucket_meta
+            if not bucket_meta and not is_last:
+                return
+            if self.buffer.device.type != 'cpu':
+                synchronize()
+            await loop.run_in_executor(None, _zmq_send_recv, {'bucket_meta': bucket_meta, 'is_last': is_last})
+            offset = 0
+            bucket_meta = {}
+
+        async def _process(name, weight):
+            nonlocal offset, n_weights
+            if self.use_shm and weight.device.type != 'cpu':
+                weight = weight.cpu()
+            if not weight.is_contiguous():
+                weight = weight.contiguous()
+            nbytes = int(weight.nbytes)
+            if nbytes > self.bucket_size:
+                raise RuntimeError(f'Weight {name} ({tuple(weight.shape)}, {weight.dtype}) '
+                                   f'is {nbytes} bytes, exceeding bucket size ({self.bucket_size}).')
+            if offset + nbytes > self.bucket_size:
+                await _flush(False)
+            bucket_meta[name] = {
+                'name': name,
+                'shape': weight.shape,
+                'dtype': weight.dtype,
+                'offset': offset,
+            }
+            self.buffer[offset:offset + nbytes].copy_(weight.view(-1).view(torch.uint8), non_blocking=True)
+            offset += nbytes
+            n_weights += 1
+
+        if is_async:
+            async for name, weight in items_iter:
+                await _process(name, weight)
+        else:
+            for name, weight in items_iter:
+                await _process(name, weight)
+
+        await _flush(True)
+        logger.debug('BucketedWeightSender: sent %d weights', n_weights)
+
+    async def send_weights(self, weights):
+        """Stream weights into buckets.  Accepts ``dict`` or iterator."""
+        items = weights.items() if isinstance(weights, dict) else weights
+        await self._stream_weights_inner(items, is_async=False)
+
+    async def send_weights_async(self, async_weights):
+        """Stream weights from an async generator into buckets."""
+        await self._stream_weights_inner(async_weights, is_async=True)
+
+    def cleanup(self):
+        if self.socket is not None:
+            self.socket.close()
+            self.socket = None
+        if self.zmq_handle.startswith('ipc://'):
+            ipc_path = self.zmq_handle[len('ipc://'):]
+            try:
+                if os.path.exists(ipc_path):
+                    os.remove(ipc_path)
+            except OSError:
+                pass
+        if self._owns_buffer:
+            del self.buffer
+        self.buffer = None
+        if self.shm is not None:
+            try:
+                self.shm.close()
+                self.shm.unlink()
+            except (FileNotFoundError, BufferError):
+                pass
+            self.shm = None
+        self._pending_handshake = None
+        from swift.utils import gc_collect, ipc_collect
+        gc_collect()
+        ipc_collect()

swift/ray/megatron/worker_group.py

@@ -0,0 +1,441 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import ray
+import socket
+import torch
+from enum import Enum
+from ray.runtime_env import RuntimeEnv
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from transformers.utils import is_torch_npu_available
+from typing import TYPE_CHECKING, Any, Dict, List, Union
+
+from swift.utils.logger import get_logger
+
+if TYPE_CHECKING:
+    from .resource_pool import ResourcePool
+
+logger = get_logger()
+
+
+class DispatchMode(str, Enum):
+    BROADCAST = 'broadcast'
+    DP = 'dp'
+    DP_SPLIT = 'dp_split'
+
+
+class CollectMode(str, Enum):
+    ALL = 'all'
+    DP = 'dp'
+    DP_FLAT = 'dp_flat'
+    FIRST = 'first'
+
+
+_DC_ATTR = '_dispatch_collect_meta'
+
+
+class DPDispatchedDict(dict):
+    """Marker dict for dp-dispatched data, keyed by dp_rank."""
+
+
+def _is_dp_dispatched(value) -> bool:
+    return isinstance(value, DPDispatchedDict)
+
+
+def dispatch_collect(
+    dispatch: Union[DispatchMode, str] = DispatchMode.BROADCAST,
+    collect: Union[CollectMode, str] = CollectMode.ALL,
+):
+    """Decorator declaring default dispatch/collect for a worker method."""
+    dispatch = DispatchMode(dispatch) if isinstance(dispatch, str) else dispatch
+    collect = CollectMode(collect) if isinstance(collect, str) else collect
+
+    def decorator(fn):
+        setattr(fn, _DC_ATTR, {'dispatch': dispatch, 'collect': collect})
+        return fn
+
+    return decorator
+
+
+def _slice_dp(value: Any, dp_size: int) -> Any:
+    """Split *value* into ``DPDispatchedDict``{dp_rank: chunk}.
+
+    Handles tensors, lists, and dicts (recursed into); scalars and
+    un-sliceable values are broadcast as-is.  Raises on empty or
+    sub-dp_size inputs rather than silently collapsing to broadcast —
+    callers must pad upstream.
+    """
+    if isinstance(value, torch.Tensor):
+        n = value.shape[0]
+        if n == 0:
+            raise ValueError(f'_slice_dp got empty tensor (shape={tuple(value.shape)})')
+        if n < dp_size:
+            raise ValueError(f'_slice_dp: tensor first dim {n} < dp_size {dp_size}.  '
+                             f'Pad the batch upstream or use dispatch="broadcast".')
+        if n % dp_size != 0:
+            raise ValueError(f'_slice_dp: tensor first dim {n} not divisible by dp_size {dp_size}. '
+                             f'Pad the batch upstream or use dispatch="broadcast".')
+        parts = value.chunk(dp_size)
+        return DPDispatchedDict({i: p for i, p in enumerate(parts)})
+    if isinstance(value, list):
+        n = len(value)
+        if n == 0:
+            raise ValueError('_slice_dp got empty list')
+        if n < dp_size:
+            raise ValueError(f'_slice_dp: list length {n} < dp_size {dp_size}. '
+                             f'Pad the batch upstream or use dispatch="broadcast".')
+        if n % dp_size != 0:
+            raise ValueError(f'_slice_dp: list length {n} not divisible by dp_size {dp_size}. '
+                             f'Pad the batch upstream or use dispatch="broadcast".')
+        chunk_size = n // dp_size
+        result = DPDispatchedDict()
+        for i in range(dp_size):
+            result[i] = value[i * chunk_size:(i + 1) * chunk_size]
+        return result
+    if isinstance(value, tuple):
+        n = len(value)
+        if n == 0:
+            raise ValueError('_slice_dp got empty tuple')
+        if n < dp_size:
+            raise ValueError(f'_slice_dp: tuple length {n} < dp_size {dp_size}. '
+                             f'Pad the batch upstream or use dispatch="broadcast".')
+        if n % dp_size != 0:
+            raise ValueError(f'_slice_dp: tuple length {n} not divisible by dp_size {dp_size}. '
+                             f'Pad the batch upstream or use dispatch="broadcast".')
+        chunk_size = n // dp_size
+        result = DPDispatchedDict()
+        for i in range(dp_size):
+            result[i] = value[i * chunk_size:(i + 1) * chunk_size]
+        return result
+    if isinstance(value, dict):
+        if _is_dp_dispatched(value):
+            return value
+        splits = {k: _slice_dp(v, dp_size) for k, v in value.items()}
+        any_split = any(_is_dp_dispatched(v) for v in splits.values())
+        if not any_split:
+            return value
+        result = DPDispatchedDict()
+        for k, v in splits.items():
+            if _is_dp_dispatched(v):
+                for dp_r, chunk in v.items():
+                    result.setdefault(dp_r, {})[k] = chunk
+            else:
+                for dp_r in range(dp_size):
+                    result.setdefault(dp_r, {})[k] = v
+        return result
+    return value
+
+
+_DISPATCH_FNS: Dict[str, Any] = {}
+_COLLECT_FNS: Dict[str, Any] = {}
+
+
+def _register_builtin_modes():
+    _DISPATCH_FNS[DispatchMode.BROADCAST] = WorkerGroup._dispatch_broadcast
+    _DISPATCH_FNS[DispatchMode.DP] = WorkerGroup._dispatch_dp
+    _DISPATCH_FNS[DispatchMode.DP_SPLIT] = WorkerGroup._dispatch_dp_split
+    _COLLECT_FNS[CollectMode.ALL] = WorkerGroup._collect_all
+    _COLLECT_FNS[CollectMode.DP] = WorkerGroup._collect_dp
+    _COLLECT_FNS[CollectMode.DP_FLAT] = WorkerGroup._collect_dp_flat
+    _COLLECT_FNS[CollectMode.FIRST] = WorkerGroup._collect_first
+
+
+class WorkerGroup:
+    """A group of Ray actors with dispatch / collect helpers.
+
+    After workers are up, call :meth:`build_dispatch_info` to query
+    ``get_parallel_info`` on every actor, cache the dp layout, and
+    bind decorated methods on the group instance.  Subsequent calls
+    like ``wg.train_step(batch)`` dispatch + collect automatically
+    using the metadata attached by :func:`dispatch_collect`.
+    """
+
+    def __init__(self, name: str, worker_handles: List[Any]):
+        self.name = name
+        self._workers = list(worker_handles)
+        self._dp_rank_map: List[int] = []
+        self._collect_mask: List[bool] = []
+        self._dp_size: int = 0
+
+    @property
+    def world_size(self) -> int:
+        return len(self._workers)
+
+    @property
+    def workers(self) -> List[Any]:
+        return list(self._workers)
+
+    @property
+    def dp_size(self) -> int:
+        if self._dp_size == 0:
+            raise RuntimeError(f'WorkerGroup[{self.name}]: build_dispatch_info() has '
+                               'not been called yet (dp_size is unknown).')
+        return self._dp_size
+
+    def __len__(self) -> int:
+        return len(self._workers)
+
+    def build_dispatch_info(self, worker_cls, rpc: str = 'get_parallel_info'):
+        """Query workers' parallel info and bind decorated methods."""
+        infos = ray.get([getattr(w, rpc).remote() for w in self._workers])
+        if len(infos) != self.world_size:
+            raise RuntimeError(f'WorkerGroup[{self.name}]: expected {self.world_size} info entries, got {len(infos)}')
+        self._dp_rank_map = [int(i['dp_rank']) for i in infos]
+        self._collect_mask = [bool(i['is_collector']) for i in infos]
+        sizes = {int(i['dp_size']) for i in infos}
+        if len(sizes) != 1:
+            raise ValueError(f'WorkerGroup[{self.name}]: inconsistent dp_size across workers: {sizes}')
+        self._dp_size = sizes.pop()
+        self._bind_decorated_methods(worker_cls)
+
+        logger.info('WorkerGroup[%s]: dp_size=%d, world_size=%d, collectors=%d', self.name, self._dp_size,
+                    self.world_size, sum(self._collect_mask))
+
+    def _bind_decorated_methods(self, worker_cls) -> List[str]:
+        """Bind methods carrying ``_dispatch_collect_meta`` onto the group.
+
+        We only bind methods that explicitly opt in via
+        :func:`dispatch_collect`; other helpers on the worker class
+        stay private to the actor so callers can't accidentally talk
+        to them through the group.  Any name already defined on the
+        group (methods like ``execute`` / ``broadcast``, properties
+        like ``dp_size``) is skipped with a warning — the caller can
+        still reach it via ``wg.execute(name, ...)``.
+        """
+        bound = []
+        for attr_name in dir(worker_cls):
+            if attr_name.startswith('_'):
+                continue
+            method = getattr(worker_cls, attr_name, None)
+            if method is None or not callable(method):
+                continue
+            meta = getattr(method, _DC_ATTR, None)
+            if meta is None:
+                continue
+            if hasattr(type(self), attr_name) or attr_name in self.__dict__:
+                logger.warning(
+                    'WorkerGroup[%s]: worker method %r collides with an existing '
+                    'attribute on WorkerGroup; call wg.execute(%r, ...) instead.', self.name, attr_name, attr_name)
+                continue
+            setattr(self, attr_name, self._make_bound(attr_name, meta['dispatch'], meta['collect']))
+            bound.append(attr_name)
+        return bound
+
+    def _make_bound(self, method_name: str, default_dispatch: str, default_collect: str):
+        wg = self
+
+        def _bound(*args, dispatch=None, collect=None, **kwargs):
+            return wg.execute(
+                method_name,
+                *args,
+                dispatch=dispatch if dispatch is not None else default_dispatch,
+                collect=collect if collect is not None else default_collect,
+                **kwargs,
+            )
+
+        _bound.__name__ = method_name
+        _bound.__qualname__ = f'{wg.name}.{method_name}'
+        _bound.__doc__ = (f'Bound remote method {method_name} '
+                          f'(dispatch={default_dispatch}, collect={default_collect})')
+        return _bound
+
+    def execute(
+        self,
+        method_name: str,
+        *args,
+        dispatch: Union[DispatchMode, str] = DispatchMode.BROADCAST,
+        collect: Union[CollectMode, str] = CollectMode.ALL,
+        **kwargs,
+    ) -> Any:
+        """Call a remote method with configurable dispatch/collect."""
+        dispatch = DispatchMode(dispatch) if isinstance(dispatch, str) else dispatch
+        collect = CollectMode(collect) if isinstance(collect, str) else collect
+        per_worker = self._dispatch(dispatch, args, kwargs)
+        futures = [getattr(w, method_name).remote(*a, **kw) for w, (a, kw) in zip(self._workers, per_worker)]
+        return self._collect(collect, ray.get(futures))
+
+    def broadcast(self, method_name: str, *args, **kwargs) -> List[Any]:
+        """Same args to every worker, block, return list."""
+        return self.execute(method_name, *args, dispatch=DispatchMode.BROADCAST, collect=CollectMode.ALL, **kwargs)
+
+    def _dispatch(self, mode: Union[DispatchMode, str], args: tuple, kwargs: dict) -> List[tuple]:
+        mode = DispatchMode(mode) if isinstance(mode, str) else mode
+        fn = _DISPATCH_FNS.get(mode)
+        if fn is None:
+            raise ValueError(f'Unknown dispatch mode: {mode!r}; registered: {list(_DISPATCH_FNS)}')
+        return fn(self, args, kwargs)
+
+    def _dispatch_broadcast(self, args: tuple, kwargs: dict) -> List[tuple]:
+        return [(args, kwargs)] * self.world_size
+
+    def _dispatch_dp(self, args: tuple, kwargs: dict) -> List[tuple]:
+        result = []
+        for dp_r in self._dp_rank_map:
+            worker_args = tuple(a[dp_r] if _is_dp_dispatched(a) else a for a in args)
+            worker_kwargs = {k: v[dp_r] if _is_dp_dispatched(v) else v for k, v in kwargs.items()}
+            result.append((worker_args, worker_kwargs))
+        return result
+
+    def _dispatch_dp_split(self, args: tuple, kwargs: dict) -> List[tuple]:
+        dp = self.dp_size
+        split_args = tuple(_slice_dp(a, dp) for a in args)
+        split_kwargs = {k: _slice_dp(v, dp) for k, v in kwargs.items()}
+        return self._dispatch_dp(split_args, split_kwargs)
+
+    def _collect(self, mode: Union[CollectMode, str], results: List[Any]) -> Any:
+        mode = CollectMode(mode) if isinstance(mode, str) else mode
+        fn = _COLLECT_FNS.get(mode)
+        if fn is None:
+            raise ValueError(f'Unknown collect mode: {mode!r}; registered: {list(_COLLECT_FNS)}')
+        return fn(self, results)
+
+    def _collect_all(self, results: List[Any]) -> List[Any]:
+        return results
+
+    def _collect_dp(self, results: List[Any]) -> Dict[int, Any]:
+        collected: Dict[int, Any] = {}
+        for r, dp_r, is_coll in zip(results, self._dp_rank_map, self._collect_mask):
+            if is_coll and r is not None:
+                collected[dp_r] = r
+        return collected
+
+    def _collect_dp_flat(self, results: List[Any]) -> List[Any]:
+        """Collect from DP collectors and flatten into a single ordered list.
+
+        Equivalent to ``_collect_dp`` followed by sorting by rank and
+        concatenating lists, which is the most common access pattern on
+        the driver side.
+        """
+        per_rank = self._collect_dp(results)
+        flat: List[Any] = []
+        for rk in sorted(per_rank.keys()):
+            part = per_rank[rk]
+            if isinstance(part, list):
+                flat.extend(part)
+            elif part is not None:
+                flat.append(part)
+        return flat
+
+    def _collect_first(self, results: List[Any]) -> Any:
+        for r, is_coll in zip(results, self._collect_mask):
+            if is_coll and r is not None:
+                return r
+        return None
+
+    def shutdown(self, timeout: float = 30.0) -> None:
+        """Best-effort shutdown of every worker and kill the Ray actors."""
+        if not self._workers:
+            return
+        pending = []
+        for w in self._workers:
+            fn = getattr(w, 'shutdown', None)
+            if fn is None:
+                continue
+            try:
+                pending.append(fn.remote())
+            except Exception as e:  # noqa: BLE001
+                logger.warning('WorkerGroup[%s] shutdown dispatch failed: %s', self.name, e)
+        if pending:
+            try:
+                ray.get(pending, timeout=timeout)
+            except Exception as e:  # noqa: BLE001
+                logger.warning('WorkerGroup[%s] shutdown timed out / raised: %s', self.name, e)
+        for w in self._workers:
+            try:
+                ray.kill(w, no_restart=True)
+            except Exception as e:  # noqa: BLE001
+                logger.warning('WorkerGroup[%s] ray.kill failed: %s', self.name, e)
+        self._workers = []
+
+    @staticmethod
+    def _get_device_env_config() -> Dict[str, str]:
+        """Return platform-specific environment variable names.
+
+        Supports CUDA (GPU) and Ascend (NPU).
+        """
+        if is_torch_npu_available():
+            return {
+                'visible_devices_key': 'ASCEND_RT_VISIBLE_DEVICES',
+                'device_max_connections_key': 'HCCL_DEVICE_MAX_CONNECTIONS',
+            }
+        return {
+            'visible_devices_key': 'CUDA_VISIBLE_DEVICES',
+            'device_max_connections_key': 'CUDA_DEVICE_MAX_CONNECTIONS',
+        }
+
+    @classmethod
+    def from_pool(
+        cls,
+        name: str,
+        resource_pool: 'ResourcePool',
+        worker_cls: str,
+    ) -> 'WorkerGroup':
+        """Spawn actors on a :class:`ResourcePool`.
+
+        Iterates PGs, discovers master on PG[0], creates workers with
+        env vars for distributed init.
+
+        Swift uses ``num_gpus=0`` + ``NOSET_CVD`` + explicit CVD
+        (torchrun style for Megatron TP/PP visibility).
+        """
+        master_addr, master_port = cls._discover_master(resource_pool)
+        dev_cfg = cls._get_device_env_config()
+        vis = resource_pool.visible_devices
+        world_size = resource_pool.world_size
+        workers = []
+
+        rank = 0
+        for pg_idx, pg in enumerate(resource_pool.pgs):
+            local_ws = resource_pool.process_on_nodes[pg_idx]
+            node_gpu_ids = [str(vis[rank + j]) for j in range(local_ws)]
+            node_cvd = ','.join(node_gpu_ids)
+
+            for local_rank in range(local_ws):
+                env_vars = {
+                    'RANK': str(rank),
+                    'LOCAL_RANK': str(local_rank),
+                    'WORLD_SIZE': str(world_size),
+                    'LOCAL_WORLD_SIZE': str(local_ws),
+                    'MASTER_ADDR': master_addr,
+                    'MASTER_PORT': str(master_port),
+                    dev_cfg['device_max_connections_key']: '1',
+                    dev_cfg['visible_devices_key']: node_cvd,
+                    'RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES': '1',
+                    'RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES': '1',
+                    'RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES': '1',
+                    'NCCL_CUMEM_ENABLE': '0',
+                    'RAY_SWIFT_GROUP': f'default,{name}',
+                }
+                w = worker_cls.options(
+                    num_gpus=0,
+                    scheduling_strategy=PlacementGroupSchedulingStrategy(
+                        placement_group=pg, placement_group_bundle_index=local_rank),
+                    runtime_env=RuntimeEnv(env_vars=env_vars),
+                ).remote()
+                workers.append(w)
+                rank += 1
+        return cls(name, workers)
+
+    @staticmethod
+    def _discover_master(resource_pool: 'ResourcePool'):
+        """Find master IP + free port on PG[0] bundle[0].
+
+        Discovers a free port on the first bundle of the first PG.
+        """
+
+        @ray.remote(num_gpus=0, num_cpus=0.01)
+        def _probe():
+            addr = ray.util.get_node_ip_address()
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind(('', 0))
+            port = sock.getsockname()[1]
+            sock.close()
+            return addr, port
+
+        pg = resource_pool.pgs[0]
+        return ray.get(
+            _probe.options(
+                scheduling_strategy=PlacementGroupSchedulingStrategy(
+                    placement_group=pg, placement_group_bundle_index=0), ).remote())
+
+
+_register_builtin_modes()