【大模型】- 扩展和分布式训练

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def train_ddp(rank, world_size):
    setup(rank, world_size)
    
    model = GPT(config).to(rank)
    model = DDP(model, device_ids=[rank])
    
    # 训练循环
    for iter in range(max_iters):
        xb, yb = get_batch("train")
        xb, yb = xb.to(rank), yb.to(rank)
        
        logits, loss = model(xb, yb)
        loss.backward()
        
        # DDP 自动平均梯度
        optimizer.step()
        optimizer.zero_grad()
    
    cleanup()

class ModelParallel(nn.Module):
    def __init__(self, model, device_map):
        super().__init__()
        self.devices = list(device_map.keys())
        
        # 分割模型
        self.layers = nn.ModuleDict()
        for name, layer in model.named_children():
            device = device_map[name]
            self.layers[name] = layer.to(device)
    
    def forward(self, x):
        for name, layer in self.layers.items():
            x = x.to(next(layer.parameters()).device)
            x = layer(x)
        return x

class TensorParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        self.world_size = world_size
        
        # 分割权重
        self.weight = nn.Parameter(
            torch.randn(out_features // world_size, in_features)
        )
    
    def forward(self, x):
        # 每个设备计算部分输出
        output = x @ self.weight.t()
        
        # 所有设备收集完整输出
        dist.all_reduce(output)
        return output

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for iter in range(max_iters):
    xb, yb = get_batch("train")
    
    with autocast():
        logits, loss = model(xb, yb)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

accumulation_steps = 4  # 等效批次大小 = 32 * 4 = 128

for iter in range(max_iters // accumulation_steps):
    optimizer.zero_grad()
    
    for micro_step in range(accumulation_steps):
        xb, yb = get_batch("train")
        
        with autocast():
            logits, loss = model(xb, yb)
            loss = loss / accumulation_steps
        
        scaler.scale(loss).backward()
    
    scaler.step(optimizer)
    scaler.update()

from torch.utils.checkpoint import checkpoint_sequential

class GPTWithCheckpointing(nn.Module):
    def forward(self, idx):
        tok_emb = self.token_embedding(idx)
        pos_emb = self.position_embedding(torch.arange(idx.size(1)))
        x = tok_emb + pos_emb
        
        # 使用梯度检查点
        x = checkpoint_sequential(self.blocks, segments=4, input=x)
        
        x = self.ln_f(x)
        return self.head(x)

import deepspeed

ds_config = {
    "train_batch_size": 128,
    "gradient_accumulation_steps": 4,
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 3,  # 分割参数、梯度和优化器状态
        "offload_optimizer": {"device": "cpu"},
        "offload_param": {"device": "cpu"},
    }
}

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

def save_checkpoint(model, optimizer, epoch, loss):
    torch.save({
        "epoch": epoch,
        "model_state_dict": model.module.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "loss": loss,
    }, f"checkpoint-{epoch}.pt")

def load_checkpoint(model, optimizer, path):
    checkpoint = torch.load(path)
    model.module.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
    return checkpoint["epoch"]

# 1. 梯度减少精度
torch.set_float32_matmul_precision("high")

# 2. 使用 flash attention
from torch.nn.functional import scaled_dot_product_attention

# 3. 分页注意力（vLLM）
# 使用 vLLM 库处理长序列

# 4. 卸载到 CPU
model = model.to("cpu")
optimizer = DeepSpeedCPUAdam(model.parameters())

# 使用 wandb 或 tensorboard
import wandb

wandb.init(project="llm-training", config=config)

for iter in range(max_iters):
    # ... 训练代码 ...
    
    if iter % log_interval == 0:
        wandb.log({
            "train/loss": loss.item(),
            "train/learning_rate": optimizer.param_groups[0]["lr"],
            "train/gradient_norm": gradient_norm,
        })