File size: 5,800 Bytes
7134ce7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | # ray的支持
SWIFT已经支持使用ray来进行多卡或多节点训练。已有功能中对ray的支持情况如下：
| 功能 | 支持ray | 例子 | 可分配角色 |
|----------|-------|--------------------------------------------------------------------------------|-----------------|
| pt/sft | ✅ | https://github.com/modelscope/ms-swift/tree/main/examples/train/multi-node/ray | default |
| dpo | ❎ | | |
| grpo | ❎ | | |
| ppo | ❎ | | |
| megatron | ❎ | | |
| 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 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使用。
|