DeepSpeed 分布式训练框架

import deepspeed

model_engine, optimizer, train_dataloader, _ = deepspeed.initialize(
    model=model,                           # PyTorch 模型
    optimizer=optimizer,                   # 可选，DeepSpeed 也可从 config 创建
    args=None,                             # 命令行参数
    config_params=ds_config,               # DeepSpeed 配置（dict 或文件路径）
    model_parameters=model.parameters(),   # 模型参数（用于创建优化器）
    dist_init_backend="nccl",              # 分布式后端
)

# 前向传播
outputs = model_engine(inputs)

# 反向传播
model_engine.backward(loss)

# 优化器步进（包含梯度累积、梯度裁剪、ZeRO 同步等）
model_engine.step()

# 保存检查点
model_engine.save_checkpoint(save_dir, tag=None)

# 加载检查点
model_engine.load_checkpoint(load_dir, tag=None)

# 获取当前损失缩放
scale = model_engine.get_global_grad_norm()

# 销毁引擎（清理资源）
model_engine.destroy()

model_engine.backward(loss)

model_engine.step()

import deepspeed
import torch

# 初始化
model_engine, optimizer, train_dataloader, _ = deepspeed.initialize(
    model=model,
    config_params=ds_config,
)

# 训练循环
for epoch in range(num_epochs):
    for batch in train_dataloader:
        # 将数据移到 GPU
        inputs = batch["input_ids"].to(model_engine.device)
        labels = batch["labels"].to(model_engine.device)

        # 前向传播
        outputs = model_engine(inputs, labels=labels)
        loss = outputs.loss

        # 反向传播（DeepSpeed 处理损失缩放和梯度分片）
        model_engine.backward(loss)

        # 优化器步进（DeepSpeed 处理梯度累积和裁剪）
        model_engine.step()

import deepspeed

# 在模型中启用激活检查点
model = deepspeed.checkpointing.checkpoint(
    model,
    num_checkpoints=10,    # 检查点数量
)

import deepspeed

# 初始化推理引擎
model = deepspeed.init_inference(
    model=model,
    mp_size=1,                              # 模型并行大小
    dtype=torch.float16,                    # 推理精度
    replace_method="auto",                  # 自动替换内核
    replace_with_kernel_inject=True,        # 注入优化内核
)

# 推理
outputs = model(input_ids)

from transformers import TrainingArguments, Trainer

# 方法一：通过 TrainingArguments 指定配置文件
training_args = TrainingArguments(
    output_dir="./output",
    deepspeed="ds_config.json",    # 指定 DeepSpeed 配置文件路径
    per_device_train_batch_size=2,
    num_train_epochs=3,
)

# 方法二：将 deepspeed 配置嵌入 TrainingArguments
training_args = TrainingArguments(
    output_dir="./output",
    deepspeed={
        "zero_optimization": {"stage": 2},
        "fp16": {"enabled": True},
        "train_batch_size": 16,
        "gradient_accumulation_steps": 4,
    },
)

# 使用 Trainer 训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.train()

import deepspeed
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# ========== 1. 创建 ds_config ==========
ds_config = {
    "train_batch_size": 16,
    "train_micro_batch_size_per_gpu": 2,
    "gradient_accumulation_steps": 4,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 2e-5,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 0.01,
        }
    },
    "bf16": {
        "enabled": True,     # 使用 bf16 混合精度
    },
    "zero_optimization": {
        "stage": 2,                          # ZeRO-2
        "offload_optimizer": {
            "device": "cpu",                 # 优化器状态卸载到 CPU
            "pin_memory": True,
        },
        "allgather_partitions": True,
        "overlap_comm": True,
        "reduce_scatter": True,
        "contiguous_gradients": True,
    },
    "gradient_clipping": 1.0,
    "steps_per_print": 10,
}

# ========== 2. 加载模型 ==========
model_name = "meta-llama/Llama-2-7b-hf"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# ========== 3. 初始化 DeepSpeed ==========
model_engine, optimizer, train_dataloader, _ = deepspeed.initialize(
    model=model,
    config_params=ds_config,
    model_parameters=model.parameters(),
)

# ========== 4. 训练循环 ==========
model_engine.train()
for epoch in range(3):
    for batch in train_dataloader:
        input_ids = batch["input_ids"].to(model_engine.device)
        labels = batch["labels"].to(model_engine.device)

        outputs = model_engine(input_ids=input_ids, labels=labels)
        loss = outputs.loss

        model_engine.backward(loss)
        model_engine.step()

    # 保存检查点
    model_engine.save_checkpoint(f"./checkpoints/epoch_{epoch}")

import deepspeed
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

ds_config = {
    "train_batch_size": 64,
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 16,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 1e-5,
            "betas": [0.9, 0.95],
            "eps": 1e-8,
            "weight_decay": 0.1,
        }
    },
    "bf16": {
        "enabled": True,
    },
    "zero_optimization": {
        "stage": 3,                          # ZeRO-3: 参数分片
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True,
        },
        "offload_param": {
            "device": "cpu",                 # 参数也卸载到 CPU
            "pin_memory": True,
        },
        "overlap_comm": True,
        "contiguous_gradients": True,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": True,
    },
    "gradient_clipping": 1.0,
    "activation_checkpointing": {
        "partition_activations": True,
        "contiguous_memory_optimization": True,
        "number_checkpoints": 10,
    },
}

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-hf")

model_engine, optimizer, train_dataloader, _ = deepspeed.initialize(
    model=model,
    config_params=ds_config,
)

# 训练循环同上

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset

# ds_config.json 已准备
training_args = TrainingArguments(
    output_dir="./deepspeed-output",
    deepspeed="ds_config.json",       # DeepSpeed 配置文件
    per_device_train_batch_size=2,
    num_train_epochs=3,
    logging_steps=10,
    save_strategy="epoch",
    bf16=True,                         # 与 ds_config 中的 bf16 对应
)

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer.pad_token = tokenizer.eos_token

dataset = load_dataset("tatsu-lab/alpaca", split="train[:5000]")

def tokenize_fn(examples):
    return tokenizer(examples["instruction"], truncation=True, max_length=512, padding="max_length")

tokenized_dataset = dataset.map(tokenize_fn, batched=True)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),
)

trainer.train()

import torch
import deepspeed
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型
model_name = "meta-llama/Llama-2-7b-hf"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 初始化推理引擎
ds_engine = deepspeed.init_inference(
    model=model,
    mp_size=1,                          # 模型并行大小
    dtype=torch.float16,
    replace_with_kernel_inject=True,    # 注入优化内核
)

# 推理
prompt = "Explain the concept of recursion:"
inputs = tokenizer(prompt, return_tensors="pt").to(ds_engine.module.device)

with torch.no_grad():
    outputs = ds_engine.module.generate(
        **inputs,
        max_new_tokens=200,
        temperature=0.7,
        do_sample=True,
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))