| # 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使用。 | |