简介（未完成）
Ray 是一个分布式计算框架，现在流行的RL框架如VeRL和OpenRLHF都依托Ray管理RL中复杂的Roles（比如PPO需要四个模型）和分配资源。以下是一些核心的概念：作者：不关岳岳的事
- Ray Actor：有状态的远程计算任务，一般是被ray.remote装饰器装饰的Python类，运行时是一个进程（和PPO等Actor-Critic算法的Actor不要混淆了）；
- Ray Task：无状态的远程计算任务，一般是被ray.remote装饰器装饰的Python函数，创建的局部变量仅在当前可见，对于任务的提交者不可见，因此可以视作无状态；
- 资源管理：Ray可以自动管理CPU、GPU、Mem等资源的分配（通过ray.remote装饰器或者启动的options参数可以指定指定的ray actor所需的计算资源），并且还可以设计资源组（placement group），将不同的ray actor指定放置在相同或者不同的资源位置（bundle）； 通过使用ray，verl可以方便地实现各种角色、各种并行策略的资源分配，并且实现hybrid engine等colocate策略；
- 异步执行：ray的计算是异步的，一般执行一个ray的计算任务后，ray会立刻返回任务的执行句柄Object reference，用户的代码不会阻塞，可以自行使用ray.get/ray.wait进行阻塞式/轮询式的结果获取；
- PS: 在RL训练中引入异步的概念，可以方便actor/critic/generator/rm之间互相overlap掉一些处理时间（比如actor在更新上一个batch的时候，generator已经可以生成下一个batch了）。由于o1-liked rl的主要时间卡点在rollout位置，因此将rollout 更好地aynsc化（例如充分利用线上serving集群的夜晚空闲时间）是未来 rl infra优化的方向之一；
parallelism
标准训练流程包含三个核心环节：
- 模型forward，计算loss，保存中间激活值
- 模型backward，通过中间激活值计算gradient
- 模型update，把gradient传给optimizer，更新模型weight。
模型不大的情况下，在step1和step2都不需要做通信，也就是每张卡算自己loss和gradient即可，并不会有什么影响。而在step3之前，我们需要把各卡的梯度放在一起求平均，保证得到正确的完整bs的梯度，而这个操作也就是all-reduce通信。
我们简单地把模型看成是Y=XW的矩阵乘法。将模型抽象为Y=XW的矩阵运算时，参数切分存在两种基本策略：
- 输入切分（X维度）：对应Data Parallel/Sequence Parallel
- 权重切分（W维度）：对应Tensor Parallel/Pipeline Parallel/Expert Parallel
切分X要比切分W简单的多。因为我们的模型输入往往是一个或多个规整的tensor，在batch维度可以很容易地做切分。而切分W就要头疼得多了，一旦出现诸如诸如卷积这种非典型矩阵计算，或者unet这种前后复杂的依赖关系，都要经过精心设计才行。在目前这种朴素的data parallel策略下，每块卡都拥有完整的model weight/gradient/optimizer，尺寸和单卡训练无异。而deepspeed使用的zero stage即是对这部分显存占用的优化。
- zero1中，每张卡只需要保留1/n的optimizer参数，通信量保持不变
- zero2在zero1的基础上，每张卡只需要保留1/n的graident，通信量保持不变
- zero3在zero2的基础上，每张卡只需要保留1/n的model weight，通信量变为1.5倍。
其中，zero1和zero2影响的分别是optimizer和graident，对应的是后两步，并没有影响forward部分
而Zero3模式下的训练流程演进为：
- Forward阶段：all-gather获取完整参数→计算loss→释放参数→保存中间激活
- Backward阶段：all-gather获取完整参数→计算梯度→释放参数
- Update阶段：reduce-scatter获取梯度切片→优化器更新局部参数”
要注意的是，zero123本质仍然属于data parallel，不属于model parallel的范畴，尽管zero3看起来做了模型参数的切分，但实际上计算时会先做all gather得到完整的模型参数，计算时使用的也是完整的参数和切分后的输入。
- 对比tp/pp，它们从头到尾都只存模型参数的一部分，计算时使用的是切分后的参数和完整的输入。
- 对于dp，通信的是模型参数，也就是W和它对应的weight/optimizer
- 对于tp/pp，通信的是中间激活值，例如PP需要将上一个rank计算得到的中间结果传给下一个rank继续计算。
SPMD
在典型的多卡训练场景中（如使用torchrun或accelerate launch），通过nvidia-smi可观察到每块GPU对应独立进程，这种模式本质源于SPMD（Single Program Multiple Data）架构。那么问题来了，是torchrun之类的启动脚本把它们“分配”到每张卡上的吗？实际上并不是。主流并行框架（DDP/DeepSpeed/Megatron）均基于SPMD范式：所有进程执行相同代码逻辑，通过环境变量差异自主确定行为模式，无需中心调度节点。一段经典的PyTorch分布式训练初始化的代码
import torch
import os
print(os.environ['RANK'], os.environ['WORLD_SIZE'], os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
torch.distributed.init_process_group(backend="nccl")
torch.cuda.set_device(torch.distributed.get_rank())
当我们使用torchrun启动这段代码后，会启动多个进程，每个进程有着不同的环境变量，标识它们属于哪一台机器和端口，是第几个进程和进程总数。之后torch.distributed.init_process_group会根据这些环境变量构建通信组，这是一个阻塞操作，所有进程都要完成init_process_group后才会继续往下走。最后set_device将当前进程绑定到一块gpu上，对于RANK=0的进程绑定在0号卡，RANK=1的进程绑定在1号卡，以此类推，不存在一个进程去调度安排它们的行为，它们运行的是完全相同的代码，只是用不同的rank区分开他们的行为。以naive dp为例，会发现在训练过程中并不存在各个dp rank之间对齐参数的行为，这是因为只要保证各个rank初始化时的模型参数保持一致，之后每个step的gradient一致，从而optimizer对模型参数的更新是一致的，自然每个rank的模型就是一致的。这也就引出了一个问题，SPMD的编程模式心智负担较重，相信写过Megatron的朋友都有感受，当逻辑复杂以后要考虑不同rank之间的不同行为和通信，以及避免corner case造成的stuck，一写一个不吱声，都是容易掉头发的环节。总结来说，SPMD由于没有中心控制器，在运行时更为高效，完全由worker自驱。但由于在编程模式上需要各worker运行相同的程序（根据rank领取自己的weights和data，自己算自己的部分，到点了通信一下就好了，不需要调度），灵活性不如single-controller模式。
FSDP
核心思想，将模型参数(权重、优化器状态等)在所有GPU之间分片存储，计算时，仅当某个GPU需要其他GPU上的参数时才进行通信，并且还进行了计算通信重叠的优化；
工程框架
OpenRLHF OpenRLHF源码解读：1.理解PPO单机训练