ml-training

ML 训练

概述

此技能适用于使用 PyTorch、TensorFlow/Keras、JAX/Flax、Hugging Face Transformers/Diffusers/Accelerate/PEFT、scikit-learn、XGBoost、LightGBM、CatBoost、Spark MLlib 和 Ray 进行模型训练。首先确保正确性：验证数据、避免数据泄漏的划分、可复现的配置、有意义的指标，并在复杂的分布式训练或大量使用加速器的训练之前，先建立简单的基线。

训练准备检查清单

1. 确定任务、目标、指标、基线和验收阈值。
2. 验证数据架构、标签质量、缺失情况、重复项、类别平衡，以及训练与服务阶段的特征一致性。
3. 选择划分策略：iid 分类采用随机分层划分，相关实体采用分组划分，时序数据采用基于时间的划分，模型选择相关的结论采用嵌套交叉验证。
4. 固定环境：框架版本、CUDA/cuDNN、驱动程序、数据集快照、预处理代码、模型配置和硬件。
5. 在适用时设置随机种子，并记录非确定性操作。除非已验证确定性模式，否则不要承诺跨 GPU 或分布式内核实现逐位可复现。
6. 先用小批量数据进行过拟合测试，再运行完整基线，最后进行调优或扩展训练。

训练循环要点

对于深度学习，每个训练循环都应明确执行前向传播、损失计算、反向传播、优化器更新、调度器更新、梯度清零、指标记录、验证、检查点保存和提前停止。

生产级 PyTorch 训练循环蓝图（含 AMP、梯度累积与裁剪）

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.cuda.amp import autocast, GradScaler

def train_one_epoch(
    model: nn.Module, 
    dataloader: DataLoader, 
    optimizer: torch.optim.Optimizer, 
    criterion: nn.Module, 
    device: torch.device,
    scaler: GradScaler,
    scheduler = None,
    grad_accum_steps: int = 1,
    max_grad_norm: float = 1.0
):
    model.train()
    optimizer.zero_grad(set_to_none=True)
    total_loss = 0.0
    
    for step, (inputs, targets) in enumerate(dataloader):
        inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
        
        # Mixed precision forward pass
        with autocast(dtype=torch.float16): # or torch.bfloat16 if hardware supports it
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            # Scale loss for gradient accumulation
            loss = loss / grad_accum_steps
            
        # Scaled backpropagation
        scaler.scale(loss).backward()
        
        if (step + 1) % grad_accum_steps == 0 or (step + 1) == len(dataloader):
            # Unscales gradients before clipping
            scaler.unscale_(optimizer)
            # Gradient clipping to prevent exploding gradients
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            
            # Step optimizer and update scale factor
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)
            
            if scheduler is not None:
                scheduler.step()
                
        total_loss += loss.item() * grad_accum_steps
        
    return total_loss / len(dataloader)

@torch.inference_mode()
def validate(model: nn.Module, dataloader: DataLoader, criterion: nn.Module, device: torch.device):
    model.eval()
    total_loss = 0.0
    correct = 0
    total = 0
    
    for inputs, targets in dataloader:
        inputs, targets = inputs.to(device, non_blocking=True), targets.to(device, non_blocking=True)
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        total_loss += loss.item()
        
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
    return total_loss / len(dataloader), correct / total

在 TensorFlow/Keras 中，当回调和分布策略足够满足需求时，优先使用内置的 fit；需要自定义控制时则使用 tf.GradientTape。在 JAX/Flax 中，明确管理状态，使用纯更新函数，并为 dropout、数据增强和采样分别使用 PRNG 密钥。

对于传统机器学习，将预处理构建在流水线对象中，以便交叉验证只在训练折上应用转换。

from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.impute import SimpleImputer
from xgboost import XGBClassifier

numeric_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='median')),
    ('scaler', StandardScaler())
])

categorical_transformer = Pipeline(steps=[
    ('imputer', SimpleImputer(strategy='most_frequent')),
    ('onehot', OneHotEncoder(handle_unknown='ignore'))
])

preprocessor = ColumnTransformer(transformers=[
    ('num', numeric_transformer, numeric_cols),
    ('cat', categorical_transformer, categorical_cols)
])

clf_pipeline = Pipeline(steps=[
    ('preprocessor', preprocessor),
    ('classifier', XGBClassifier(use_label_encoder=False, eval_metric='logloss'))
])

优化器、学习率调度和正则化

当需要解耦权重衰减时，将 AdamW 作为可靠的深度学习默认选择；对于某些视觉任务，可使用带动量的 SGD；仅在内存压力足以构成理由时才使用 Adafactor 或 8 位优化器。对于 Transformer 类模型，除非架构文档另有说明，否则不要对偏置和归一化参数应用权重衰减。

学习率策略通常比对优化器进行细微调参更重要。常见默认方案：

• Transformer 和大批量训练采用预热加余弦衰减。
• 探索性训练采用 One-cycle 或余弦重启。
• 收敛过程波动的较小规模监督任务采用 Reduce-on-plateau。
• 增大全局批量时，按有效批量大小调整线性预热。

应通过训练集和验证集曲线检测过拟合，而不能只看最终指标。可采用数据增强、dropout、标签平滑、权重衰减、提前停止、mixup/cutmix、更严格的验证集划分或更小的模型。对于欠拟合，在增加复杂度之前，先检查标签、特征信号、损失缩放、优化器学习率、模型容量和预处理错误。

批量大小、梯度累积和精度

有效批量大小等于每设备批量大小乘以设备数，再乘以梯度累积步数。只有在验证效果保持稳定时，才通过增大批量来提高吞吐量。如果内存有限，可使用梯度累积、激活检查点、更短的序列、更小的图像、分桶、参数高效微调、优化器状态分片或卸载。

在现代加速器上使用混合精度。对于原生支持 bf16 的硬件，优先使用 bf16，因为它能减少损失缩放问题；否则使用带动态损失缩放的 fp16 AMP。必要时，将某些归约、softmax/logits 处理和指标累积等对数值敏感的操作保持在 fp32。注意 NaN、溢出、下溢和归一化不稳定问题。

分布式训练模式

选择能满足瓶颈需求的最简单分布模式：

模式	适用场景	注意事项
数据并行 / DDP	模型可在单个设备上运行；需要提高吞吐量	需要正确设置全局批次、分片采样器并同步指标
FSDP / ZeRO	模型只有在对参数或优化器状态进行分片后才能运行	检查点管理更复杂；需要调整封装和卸载策略
张量并行	单个层无法放入单个设备，或需要扩展到大型 Transformer	需要框架或运行时支持；通信开销大
流水线并行	超深模型需要按层分区	存在气泡开销；微批次调度很重要
专家并行 / 模型并行	混合效应模型或超大型模型	路由和负载均衡复杂

PyTorch DDP / FSDP 初始化蓝图

import os
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

def setup_distributed():
    # Environment variables set by torchrun
    world_size = int(os.environ["WORLD_SIZE"])
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    
    dist.init_process_group(
        backend="nccl", 
        init_method="env://", 
        world_size=world_size, 
        rank=rank
    )
    torch.cuda.set_device(local_rank)
    return rank, local_rank, world_size

def cleanup():
    dist.destroy_process_group()

# Usage DDP:
# rank, local_rank, _ = setup_distributed()
# model = MyModel().to(local_rank)
# model = DDP(model, device_ids=[local_rank])

# Usage FSDP:
# rank, local_rank, _ = setup_distributed()
# model = MyModel().to(local_rank)
# model = FSDP(model)

DeepSpeed JSON 配置蓝图（deepspeed_config.json）

{
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto",
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true,
    "cpu_offload": false
  },
  "fp16": {
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 16,
    "hysteresis": 2,
    "min_loss_scale": 1
  },
  "optimizer": {
    "type": "AdamW",
    "params": {
      "lr": "auto",
      "betas": [0.9, 0.999],
      "eps": 1e-8,
      "weight_decay": "auto"
    }
  }
}

检查点与容错

可恢复的检查点应包含模型权重、优化器状态、调度器状态、缩放器状态、轮次/步数、可行时保存的 RNG 状态、分词器/预处理器、标签映射、配置、指标历史记录，以及代码或提交标识符。分别保存验证集表现最佳的检查点和最新检查点。在长时间训练前测试恢复流程。将产物存储在持久存储中，而非临时工作节点磁盘上。

PyTorch 稳健的检查点保存与加载

def save_checkpoint(state, is_best, checkpoint_dir="checkpoints"):
    os.makedirs(checkpoint_dir, exist_ok=True)
    filepath = os.path.join(checkpoint_dir, "last_checkpoint.pt")
    torch.save(state, filepath)
    if is_best:
        shutil.copyfile(filepath, os.path.join(checkpoint_dir, "best_model.pt"))

def load_checkpoint(filepath, model, optimizer=None, scheduler=None, scaler=None):
    checkpoint = torch.load(filepath, map_location="cpu")
    model.load_state_dict(checkpoint["state_dict"])
    if optimizer is not None and "optimizer" in checkpoint:
        optimizer.load_state_dict(checkpoint["optimizer"])
    if scheduler is not None and "scheduler" in checkpoint:
        scheduler.load_state_dict(checkpoint["scheduler"])
    if scaler is not None and "scaler" in checkpoint:
        scaler.load_state_dict(checkpoint["scaler"])
    return checkpoint.get("epoch", 0), checkpoint.get("best_metric", float("inf"))

不平衡与含噪数据

针对类别不平衡，首先选择能反映目标的指标：PR-AUC、F1 变体、平衡准确率、指定精确率下的召回率、成本加权效用或校准度。可使用类别权重、焦点损失、平衡采样、阈值调整或数据收集。避免在拆分数据前过采样。针对标签噪声，应检查高损失样本、使用鲁棒损失函数、重新标注关键样本并报告不确定性。

扩展与高级训练

务实地应用缩放定律：在扩展前，先估算质量受数据、模型还是计算限制。跟踪损失与计算量、数据规模和参数量之间的关系。对于大型模型，可考虑课程学习、数据质量筛选、去重、序列打包、Token 预算核算以及计算最优权衡。对于 RLHF 或强化学习，应在在线优化前验证奖励质量、离线评估、安全约束和分布偏移。

常见故障诊断

• 训练损失没有下降：标签错误、参数被冻结、学习率不当、损失函数与激活函数不匹配、数据类型溢出、未调用 train() 模式，或输入流水线损坏。
• 验证表现远差于训练表现：无数据泄漏的拆分可能会暴露过拟合、分布偏移、数据增强不匹配或正则化差距。
• 出现 NaN：学习率过高、fp16 溢出、输入无效、损失不稳定、除以零或梯度爆炸。
• 训练缓慢：数据加载器成为瓶颈、批次过小、CPU 转换、日志过多、同步、精度未优化或使用网络存储。
• 结果不可复现：数据/环境未固定版本、内核具有非确定性、随机数据增强、分布式采样器存在差异，或预处理状态被隐藏。

资料来源

• PyTorch 分布式和 AMP 文档：https://pytorch.org/docs/stable/
• TensorFlow 训练与分布指南：https://www.tensorflow.org/guide
• JAX 和 Flax 文档：https://jax.readthedocs.io/ 和 https://flax.readthedocs.io/
• Hugging Face Transformers、Accelerate 和 PEFT 文档：https://huggingface.co/docs
• scikit-learn 模型评估和流水线文档：https://scikit-learn.org/stable/user_guide.html