# ray的支持

## Megatron Ray

Megatron 后端支持通过 Ray 进行 GRPO 和 GKD 训练：

| 功能 | 例子 | 可分配角色 |
|------|------|-----------|
| megatron grpo | https://github.com/modelscope/ms-swift/tree/main/examples/ray/grpo | train/rollout |
| megatron gkd  | https://github.com/modelscope/ms-swift/tree/main/examples/ray/gkd  | train/rollout/teacher |

### 何时使用

非 Ray Megatron（`megatron rlhf`）和 Ray Megatron（`megatron rlhf --use_ray true`）训练功能相同

核心区别在于**部署方式**：

- **非 Ray Megatron**：通过 torchrun 启动。推理可选 colocate（同进程）或 server（手动启动 vLLM server）模式。多节点需要在每个节点手动配置 `MASTER_ADDR/PORT` 并分别启动 torchrun 和 vLLM server。
- **Ray Megatron**：通过一份 YAML 声明各角色的 GPU 数量（`train.gpus`、`rollout.gpus`、`teacher.gpus`），Ray 自动完成进程创建、GPU 分配和跨节点调度，无需手动管理多个进程。

两者都支持训练和推理的 GPU 隔离（非 Ray 通过 `vllm_mode=server`，Ray 通过 YAML 配置 separate 模式），功能上等价。Ray 的优势是将多进程的编排自动化——在多节点场景下，免去逐节点手动启动 torchrun 和 vLLM server 的运维负担。

**选择建议：**

| 场景 | 建议 |
|------|------|
| 单机训练 | **非 Ray** — 更简单 |
| 多节点集群 | **Ray** — 自动跨节点调度，一份 YAML 一键启动 |

### 快速开始

```bash
# 1. 启动 Ray 集群（单节点可省略）
ray start --head                        # 主节点
ray start --address=<head_ip>:6379      # 其他节点

# 2. 提交训练
megatron rlhf --use_ray true --config examples/ray/grpo/ray_grpo_colocate.yaml
```

### GPU 分配模式

**Colocate（共享 GPU）**— 训练和推理共享同一组 GPU，交替使用，通过 sleep/wake 释放显存：

```yaml
colocate_groups: [[train, rollout]]
offload_model: true
offload_optimizer: true
sleep_level: 1

train:
  gpus: 4
rollout:
  gpus: 4    # 必须和 train 相同
```

**Separate（独立 GPU）**— 训练和推理各占独立 GPU，无显存竞争：

```yaml
# 不设置 colocate_groups
train:
  gpus: 4
rollout:
  gpus: 4    # 独立的 4 张卡
```

### GKD Teacher 模式

| 模式 | 配置方式 | top-k | full-vocab |
|------|---------|:-----:|:----------:|
| Colocated teacher | 设置 `teacher_model` + `offload_teacher_model: true` | ✅ | ✅ |
| 独立 teacher GPU 组 | 添加 `teacher:` 组并设置 `gpus`、`model` | ✅ | ❌ |

- **Colocated teacher**：teacher 是 Megatron 模型，与 student 共享同一组 GPU 和相同的并行参数，通过 offload 交替释放显存。
- **独立 teacher GPU 组**：teacher 是独立的 vLLM 推理引擎，运行在单独 GPU 上，并行参数独立配置（`vllm_tensor_parallel_size`）。
- **top-k**：蒸馏损失仅在 teacher 概率最高的 k 个 token 上计算（通过 `gkd_logits_topk` 设置），显存占用更低，但会丢弃长尾分布信息。
- **full-vocab**：蒸馏损失在完整词表上计算，保留完整分布信息，但显存占用较高。

### 相关文档

更多文档请参考

- **GRPO 训练**：[Megatron GRPO 文档](../Megatron-SWIFT/GRPO.md)
- **GKD 训练**：[GKD 文档](../Megatron-SWIFT/GKD.md)
- **Megatron 训练参数**：[命令行参数文档](../Megatron-SWIFT/Command-line-parameters.md)
- **Megatron 快速开始**：[Quick Start](../Megatron-SWIFT/Quick-start.md)

详细配置说明和示例见 [examples](https://github.com/modelscope/ms-swift/tree/main/examples/ray)。

## Swift Ray

SWIFT 的 HF Trainer 侧也支持使用 ray 来进行多卡或多节点训练：

| 功能       | 支持ray | 例子                                                                             | 可分配角色           |
|----------|-------|--------------------------------------------------------------------------------|-----------------|
| pt/sft   | ✅     | https://github.com/modelscope/ms-swift/tree/main/examples/train/multi-node/ray | default         |
| dpo      | ❎     |                                                                                |                 |
| grpo     | ❎     |                                                                                |                 |
| ppo      | ❎     |                                                                                |                 |
| sampling | ✅     | https://github.com/modelscope/ms-swift/tree/main/examples/sampler/distill      | sampler/prm/orm |
| distill  | ✅     | https://github.com/modelscope/ms-swift/tree/main/examples/sampler/sample       | sampler/prm/orm |

### 技术细节

在叙述参数设置之前，我们有必要先行讲一下技术细节。由于SWIFT的内部当前使用了大量transformers和trl的已有实现，像veRL或ROLL一样拆解为不同的ray角色是不现实的，而且拆解后会以ray为中心，对非ray的场景的支持会不良。
因此SWIFT采取了装饰器为主的技术方案，以函数级别定义了不同角色，这些角色可以在参数中被定义如何使用。看下面的例子：

```python
from swift.ray_utils import RayHelper

@RayHelper.worker(group=['model1', 'model2'])
class MyTrainer:

    def __init__(self, args):
        self._prepare_model1()
        self._prepare_model2()
        self._prepare_datasets()

    @RayHelper.function(group='model1')
    def _prepare_model1(self):
        ...

    @RayHelper.function(group='model2')
    def _prepare_model2(self):
        ...

    @RayHelper.function(group='model1')
    def rollout(self, inputs):
        return self.model1.generate(inputs)

    @RayHelper.function(group='model2')
    def forward_model2(self, inputs):
        loss = self.model2.forward(inputs)
        loss.backward()

    def _prepare_datasets(self):
        self.dataset = ...

    def train(self):
        for batch in DataLoader(self.dataset):
            generated = self.rollout(batch)
            self.forward_model2(generated)
            ...


if __name__ == '__main__':
    ...
    MyTrainer(args).train()
```

RayHelper会将被装饰的方法分配到不同的硬件集群中，本地调用会被平滑转换到ray集群中进行远程调用。也可以以类为中心进行划分：

```python

@RayHelper.worker(group=['model1'])
class Model1:
    ...

    @RayHelper.function(group='model1')
    def rollout(self):
        ...

@RayHelper.worker(group=['model2'])
class Model2:
    ...

    @RayHelper.function(group='model2')
    def forward_and_optimize(self):
        ...


class Trainer:
    ...
```

SWIFT对ray的支持本质上是使用@worker和@function两个注解的组合使用，worker指定ray集群的角色，function指定如何分配数据。

function注解有额外的几个参数：
```python
    @staticmethod
    def function(group: str,
                 dispatch: Union[Literal['slice', 'all'], Callable] = 'all',
                 execute: Literal['first', 'all'] = 'all',
                 collect: Union[Literal['none', 'flatten'], Callable] = 'none'):
```

- dispatch: 如何分配调用入参
  - slice：对入参切分，也就是worker负载均衡执行
  - all：各个worker入参完全相同
  - 自定义切分方式，格式为：
    ```python
        def my_custom_slice(n, i, data):
            # n是worker数量，i是当前worker索引，data是原始入参
            # 返回第i个的入参
    ```
- execute: 如何执行
  - first: rank0执行，此时slice和Callable方式切分无效
  - all: 全部执行

- collect: 如何收集返回数据
  - none：原样返回，格式为各个worker返回值的列表
  - flatten: 将worker返回的结果进行拉平，支持tuple的拉平
  - Callable: 自定义collect方式，格式为：
    ```python
        def my_custom_collect(result):
            # result是各个worker返回的列表
            # 输入你想要的格式
    ```

### 参数设置

讲完技术细节后，可以将参数配置了。开发者可以根据不同的流程中的角色列表，设置不同的硬件搭配方式，例如采样功能中，共有三个角色，sampler、prm、orm，可以这样配置：

```yaml
device_groups:
  nproc_per_node: 4
  sample_group:
    device: GPU
    ranks: list(range(0, 2))
    workers:
      - sampler
  rm_group:
    device: GPU
    ranks: list(range(2, 4))
    workers:
      - prm
      - orm
```

- nproc_per_node: ray集群中需要的每个node的最小卡数。
xxx_group: 每个ray组的名称，可以随意指定
  - device: 设备类型，当前支持GPU/CPU等。
  - ranks: 当前组分配到哪些ranks上。如果是CPU，ranks只能为整数，代表共需要多少进程，如果是GPU，可以为`[0,1,2,3]`, `4`, `list(range(0, 4))`等格式。
  - workers: 哪些角色分配到当前组中。

所有可用的角色可以见本文最上面的表。

如果使用命令行，device_groups也可以以`--device_groups xxx`方式传入，xxx为jsonstring。为了配置的简便，我们强烈推荐使用yaml方式搭配ray使用。