如何提交 DeepSpeed 训练任务¶

根据 DeepSpeed 官方文档，我们推荐使用修改代码的方式实现。

即使用 deepspeed.init_distributed() 代替 torch.distributed.init_process_group(...)。然后运行命令使用 torchrun，提交为 Pytorch 分布式任务，既可运行 DeepSpeed 任务。

是的，你可以使用 torchrun 运行你的 DeepSpeed 训练脚本。 torchrun 是 PyTorch 提供的一个实用工具，用于分布式训练。你可以结合 torchrun 和 DeepSpeed API 来启动你的训练任务。

以下是一个使用 torchrun 运行 DeepSpeed 训练脚本的示例：

编写训练脚本：

train.py

import torch
import deepspeed
from torch.utils.data import DataLoader

# 模型和数据加载
model = YourModel()
train_dataset = YourDataset()
train_dataloader = DataLoader(train_dataset, batch_size=32)

# 配置文件路径
deepspeed_config = "deepspeed_config.json"

# 创建 DeepSpeed 训练引擎
model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config_params=deepspeed_config
)

# 训练循环
for batch in train_dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

创建 DeepSpeed 配置文件：

deepspeed_config.json

{
  "train_batch_size": 32,
  "gradient_accumulation_steps": 1,
  "fp16": {
    "enabled": true,
    "loss_scale": 0
  },
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.00015,
      "betas": [0.9, 0.999],
      "eps": 1e-08,
      "weight_decay": 0
    }
  }
}

使用 torchrun 或者 baizectl 运行训练脚本：
```
torchrun train.py
```
通过这种方式，你可以结合 PyTorch 的分布式训练功能和 DeepSpeed 的优化技术，从而实现更高效的训练。您可以在 Notebook 中，使用 baizectl 提交命令：
```
baizectl job submit --pytorch --workers 2 -- torchrun train.py
```