企业官网建设流程全解析

SWIFT Ray 分布式训练指南：Megatron RLHF 集群编排与装饰器式角色抽象

【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600+ LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300+ MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift

本文基于 SWIFT 开源仓库的 Ray 支持文档，系统讲解两条 Ray 集成路径：Megatron 后端的 Ray 集群编排（用于 GRPO/GKD 强化与蒸馏训练）与 HF Trainer 侧基于装饰器的 Ray 角色抽象（用于采样/蒸馏等流程）。读完本文，你将掌握如何用一份 YAML 声明 train/rollout/teacher 等角色的 GPU 分配、如何通过colocate_groups与offload_*配置实现训练与推理的显存共享或隔离，以及如何用@RayHelper.worker/@RayHelper.function将普通 Python 类平滑改造成跨节点分布式执行的任务。

一、Ray 支持总览

SWIFT 在两条技术路线上支持 Ray：

1. Megatron 后端（megatron rlhf）——基于 Megatron 的 GRPO、GKD 训练可通过 Ray 完成多进程编排：

2. HF Trainer 侧（SWIFT 主流程）——pt/sft、sampling、distill等任务同样支持 Ray：

从上面的角色列可以看到两条路径的分工差异：Megatron 路径把训练、推理（rollout）、教师模型（teacher）拆成独立角色；而 SWIFT 主流程侧则把采样器（sampler）、过程奖励模型（prm）、结果奖励模型（orm）拆成可独立分配的角色。二者的共同点都是「一份配置声明角色，Ray 自动完成资源调度」。

二、Megatron Ray：何时使用

非 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 的运维负担。

选择建议：

三、Megatron Ray：快速开始

先启动 Ray 集群（单节点可省略），再提交训练任务：

# 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

仓库也提供了现成的启动脚本，例如 examples/ray/grpo/run_colocate.sh 与 examples/ray/gkd/run.sh，二者均通过export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}限定可见 GPU，再执行megatron rlhf --use_ray true --config ...。需要注意：CUDA_VISIBLE_DEVICES暴露的 GPU 总数必须不小于所选部署方式所需的总卡数（colocate 需要train.gpus，separate 需要train.gpus + rollout.gpus，独立 teacher 还需再加上teacher.gpus）。

四、Megatron Ray：GPU 分配模式

Colocate（共享 GPU）

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

colocate_groups: [[train, rollout]] offload_model: true offload_optimizer: true sleep_level: 1 train: gpus: 4 rollout: gpus: 4 # 必须和 train 相同

colocate 模式下，colocate_groups中同一组的所有角色必须设置相同的gpus（共享同一组物理设备）。训练与 rollout 在时间上交替占用显存：offload_model、offload_optimizer负责把模型与优化器状态换出到 CPU，sleep_level: 1让空闲角色释放 GPU 显存给活动角色。

Separate（独立 GPU）

训练和推理各占独立 GPU，无显存竞争：

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

separate 模式不设置colocate_groups，训练与 rollout 占用互不相交的 GPU 集合，每一步训练结束后把权重推送到 rollout 推理引擎，训练与采样可重叠执行。当单卡显存不足以同时容纳训练与推理时，这是更稳妥的选择。

五、Megatron Ray：GKD Teacher 模式

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

仓库的 examples/ray/gkd 目录提供了三种现成配置，文件名即编码了「rollout 放置方式 + teacher 模式」两个独立选择：

Colocated teacher 配置示例

teacher_model: Qwen/Qwen3.5-4B offload_teacher_model: true gkd_logits_topk: 64 # 省略该项即为 full-vocab 蒸馏

colocated teacher 与 student 共享 train GPU，通过offload_teacher_model: true在 teacher forward 之间把教师模型换出到 CPU。完整示例见 rollout_colocate_teacher_colocate.yaml。

独立 teacher GPU 组配置示例

gkd_logits_topk: 64 # REQUIRED —— 副本模式仅支持 top-k # 不要设置顶层 teacher_model（否则会额外加载一个 colocated teacher） teacher: gpus: 4 model: Qwen/Qwen3.5-4B # 这些副本服务的 teacher checkpoint vllm_engine_kwargs: {"max_logprobs": 64} # 必须 >= gkd_logits_topk

独立 teacher 以一组由 Ray 管理的 vLLM 副本运行在独立 GPU 上，按位置返回 prompt top-k logprobs，driver 每步拉取。两个关键约束：

• vllm_engine_kwargs.max_logprobs必须≥gkd_logits_topk，否则 vLLM 会拒绝prompt_logprobs请求；
• 所需 GPU 总数 = colocated 的 train+rollout 集合+teacher.gpus。

top-k 与 full-vocab 的选择

• top-k（gkd_logits_topk: K）：teacher 每位置只暴露 top-K logprobs，显存占用显著更低，适用于所有 teacher 模式；
• full-vocab（省略gkd_logits_topk）：蒸馏完整词表分布，仅限 colocated teacher，且显存开销大（需要缓存按 rank 切分的 teacher logits）。若 OOM，应切换为 top-k 或调低micro_batch_size。

OPSD（On-Policy / privileged Distillation）

在 Ray 场景下还支持 OPSD：teacher 看到与 student不同的（特权）prompt，但对同一段 on-policy 响应打分——例如 teacher 看到「问题 + 参考答案」。数据集预处理器（通过external_plugins加载）为每行数据产出teacher_prompt，损失按 mask 对齐共享响应 token：

external_plugins: examples/train/rlhf/opsd/opsd_plugin.py # 注册 teacher_prompt teacher_model: Qwen/Qwen3.5-4B gkd_logits_topk: 64

OPSD 在 Ray 下以 top-k（gkd_logits_topk）方式同时支持 colocated teacher 与独立 teacher 副本（teacher.gpus > 0），且无需额外开关：当数据行携带非空teacher_prompt时自动激活，否则回退为普通 GKD。

六、Megatron Ray：常见参数与避坑

结合 examples/ray/README.md 的说明，以下参数需要特别注意：

• 序列长度预算：编码器预算为max_length + max_completion_length（prompt 上限max_length，on-policy 补全最多max_completion_length），请据此设置vllm_max_model_len。
• padding_free: true会把一个 micro-batch 打包进单条序列，配合tensor_model_parallel_size > 1时需同时开启sequence_parallel: true（参考 rollout_colocate_teacher_colocate.yaml 的写法）。
• 并行度与 DP 计算：数据并行大小 =gpus / (TP * PP * CP)。例如 4 卡 +tensor_model_parallel_size: 2→ DP2。
• colocate 显存释放：offload_model、offload_optimizer、offload_teacher_model与sleep_level: 1是 colocate 角色能够共存的关键。
• GRPO 专项：奖励通过reward_funcs+external_plugins接入，采样通过num_generations/steps_per_generation控制，不涉及任何teacher_*配置。

以仓库提供的 ray_grpo_colocate.yaml 为例，一份完整的 GRPO colocate 配置包含：顶层共享训练参数（模型、数据集、reward_funcs、micro_batch_size、lr、loss_type: grpo、advantage_estimator: grpo等）、colocate_groups/offload_*/sleep_level放置策略，以及train/rollout两个角色各自的 GPU 数与并行参数；separate 版本（ray_grpo_separate.yaml）则去掉colocate_groups，并将vllm_gpu_memory_utilization从 0.4 提升到 0.8（因为不再需要与训练共享显存）。

七、Swift Ray：装饰器式角色抽象的技术细节

与 Megatron 路径「一份 YAML 编排现成角色」不同，SWIFT 主流程（HF Trainer 侧）对 Ray 的支持采用了装饰器为主的技术方案。原因在于：SWIFT 内部大量复用了 transformers 和 trl 的既有实现，像 veRL 或 ROLL 那样把流程完整拆解为不同 ray 角色并不现实——强行拆解会以 ray 为中心，反而伤害非 ray 场景的支持。因此 SWIFT 以函数级别定义不同角色，角色如何被使用完全由参数决定。

核心机制在 swift/ray_utils/base.py 中实现：RayHelper提供worker与function两个静态装饰器，worker指定 Ray 集群中的角色（注册到RayHelper.worker_cls），function指定方法如何被分发执行。看下面的示例：

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 集群中的远程调用（worker 侧通过RAY_SWIFT_GROUP环境变量判断自身所属分组，driver 侧通过execute_all_sync→ray.get(execute_all_async(...))同步获取远程结果）。也可以以类为中心进行划分：

@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 注解的额外参数

function注解在 base.py 中的完整签名为：

@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 负载均衡执行（源码中对 list 入参按divmod均匀切片，非 list 则每个 worker 得到完整副本）
  • all：各个 worker 入参完全相同
  • 自定义切分方式，格式为：

    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 的拉平（源码中对 tuple 结果递归收集，list 结果逐层 flatten）
  • Callable：自定义 collect 方式，格式为：

    def my_custom_collect(result): # result是各个worker返回的列表 # 输入你想要的格式

八、Swift Ray：参数设置

理解技术细节后即可配置参数。开发者可以根据不同流程中的角色列表，设置不同的硬件搭配方式。以采样功能为例，共有 sampler、prm、orm 三个角色，可以这样配置（完整示例见 examples/sampler/sample/sampling.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为 JSON 字符串——该参数在 swift/ray_utils/arguments.py 的RayArguments中定义，__post_init__会自动对 JSON 字符串执行json.loads解析。为了配置的简便，强烈推荐使用 YAML 方式搭配 Ray 使用。

Ray 实验名称可通过ray_exp_name指定（会写入RAY_SWIFT_EXP_NAME环境变量，作为集群与 worker 命名前缀），多实验并存时可用于区分命名空间；对应的use_ray参数（RayArguments.use_ray，默认False）是所有 Ray 功能的统一开关。

九、相关文档

• GRPO 训练：Megatron GRPO 文档
• GKD 训练：GKD 文档
• Megatron 训练参数：命令行参数文档
• Megatron 快速开始：Quick Start

详细配置说明和示例见 examples/ray 目录下的 GRPO、GKD 与 README；Ray 基础设施的完整源码实现位于 swift/ray_utils（RayHelper、ResourceManager与RayArguments三个模块），深入阅读可以了解 placement group 调度、环境变量透传（WORLD_SIZE、RANK、MASTER_ADDR/PORT、CUDA_VISIBLE_DEVICES/ASCEND_VISIBLE_DEVICES映射）等底层编排细节。

【免费下载链接】swiftUse PEFT or Full-parameter to CPT/SFT/DPO/GRPO 600+ LLMs (Qwen3.6, DeepSeek-V4, GLM-5.1, InternLM3, Llama4, ...) and 300+ MLLMs (Qwen3-VL, Qwen3-Omni, InternVL3.5, Ovis2.5, GLM4.5v, Gemma4, Llava, Phi4, ...) (AAAI 2025).项目地址: https://gitcode.com/GitHub_Trending/swift1/swift