deepspeed的使用demo（自用、更新）

准备：创建3个文件

分别是deepspeed的配置文件、项目文件（可以是多个文件）、脚本文件。

具体的demo代码如下：

import deepspeed
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

class Moduel(nn.Moudle):
    pass

class MyDataset(Dataset):
    pass

if __name__ == "__main__":
    # 1. deepspeed 的初始化,加载deepspeed配置文件
    """
    deepspeed.init_distributed() 主要作用是设置分布式训练所需的环境，包括初始化 GPU 设备、
    通信策略（如 NCCL、MPI 等）、进程间通信等。通过调用这个方法，你可以确保每个参与训练的设备都被正确地识别并设置，
    以便在后续训练过程中能够正确进行数据的分发、模型参数的同步等操作。
    """
    deepspeed.init_distributed()
    with open("deepspeed.json", encoding="utf-8") as f:
        deepspeed_config = f.read()
        deepspeed_config = eval(deepspeed_config)

    # 2.定义模型和，数据集和损失函数
    model = Moduel()
    train_dataset = MyDataset()
    loss_fn = nn.CrossEntropyLoss()

    # 3.将模型和数据集交给deepspeed管理
    """
    传入：模型，训练集，deepspeed配置，模型参数
    输出：新的模型，优化器， 数据加载器， 学习率策略
    """
    new_model, optimizer, train_dataloader, lr_scheduler  = deepspeed.initialize(model=model, training_data=train_dataset, config=deepspeed_config, model_parameters=model.parameters())

    # 4. 开始训练，指定训练模型为deepspeed初始化之后的新的模型
    new_model = new_model.train()

    """
    数据指定设备：将数据加载器返回的特征和标签移到指定设备中，新的模型会自动调用到指定设备
    反向传播：这里的反向传播是将loss传入新的模型中进行反向传播。torch中的反向传播是用loss.backward()
    梯度更新：这里的梯度更新是用new_model.step()来进行的。torch中的梯度更新是用优化器->optimizer.step()
    """
    for x, y in train_dataloader:
        x = x.to(new_model.device)
        y = y.to(new_model.device)

        output = new_model(x)

        loss = loss_fn(output, y)

        new_model.backward(loss)
        new_model.step()