附录A 常用工具速查

附录A 常用工具速查

本附录是 AI Infra 工程师的案头手册。不会解释每个工具的设计哲学，但会给出最常用的命令和配置模板。

A.1 训练框架命令速查

A.1.1 PyTorch DistributedDataParallel (DDP)

DDP 是 PyTorch 内置的数据并行方案，适合单机多卡或小规模多机训练。

# 基本启动方式（单机 8 卡）
torchrun --nproc_per_node=8 train.py \
    --batch_size=32 \
    --learning_rate=1e-4

# 多机训练（2 台机器，每台 8 卡）
# 机器 0（主节点）
torchrun --nnodes=2 --nproc_per_node=8 \
    --rdzv_id=exp01 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

# 机器 1（工作节点）
torchrun --nnodes=2 --nproc_per_node=8 \
    --rdzv_id=exp01 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=192.168.1.100:29500 \
    train.py

# DDP 标准代码模板
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def setup_ddp():
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank())

def cleanup():
    dist.destroy_process_group()

def train():
    setup_ddp()
    
    model = MyModel().cuda()
    model = DDP(model, device_ids=[dist.get_rank()])
    
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    
    for epoch in range(num_epochs):
        sampler.set_epoch(epoch)  # ⚠️ 必须调用！
        for batch in loader:
            loss = model(batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
    
    cleanup()

Warning常见坑

忘记调用 sampler.set_epoch(epoch) 会导致每个 epoch 的数据分配相同——模型学不到新东西。

A.1.2 DeepSpeed

DeepSpeed 提供了 ZeRO 优化器，是大模型训练的核心工具。

# 安装
pip install deepspeed

# 基本启动
deepspeed --num_gpus=8 train.py \
    --deepspeed ds_config.json

# 多机启动
deepspeed --num_gpus=8 --num_nodes=2 \
    --hostfile hostfile \
    train.py --deepspeed ds_config.json

// ds_config_zero2.json - ZeRO-2 配置
{
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 5e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 5e8,
    "contiguous_gradients": true
  },
  "bf16": {
    "enabled": true
  },
  "gradient_accumulation_steps": 4,
  "train_batch_size": 256,
  "train_micro_batch_size_per_gpu": 8
}

// ds_config_zero3.json - ZeRO-3 配置（最激进的内存优化）
{
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": { "device": "cpu" },
    "offload_param": { "device": "cpu" },
    "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
  },
  "bf16": { "enabled": true },
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "gradient_accumulation_steps": "auto"
}

A.1.3 Megatron-LM

Megatron-LM 是 NVIDIA 开发的大模型训练框架，支持 3D 并行。

# Megatron-LM 启动脚本（预训练 GPT 模型）
GPUS_PER_NODE=8
NNODES=2
MASTER_ADDR=192.168.1.100
MASTER_PORT=6000

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE \
                  --nnodes $NNODES \
                  --master_addr $MASTER_ADDR \
                  --master_port $MASTER_PORT"

torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
    --num-layers 32 \
    --hidden-size 4096 \
    --num-attention-heads 32 \
    --seq-length 4096 \
    --micro-batch-size 2 \
    --global-batch-size 256 \
    --tensor-model-parallel-size 4 \
    --pipeline-model-parallel-size 2 \
    --DDP-impl local \
    --bf16 \
    --use-flash-attn \
    --optimizer adamw \
    --lr 1e-4 \
    --min-lr 1e-5 \
    --weight-decay 0.1 \
    --clip-grad 1.0 \
    --train-iters 500000 \
    --lr-decay-iters 320000 \
    --data-path dataset/indices \
    --vocab-file vocab.json \
    --merge-file merges.txt \
    --save checkpoints/gpt-7b \
    --load checkpoints/gpt-7b \
    --save-interval 5000 \
    --log-interval 10

A.1.4 PyTorch FSDP

FSDP 是 PyTorch 原生的 ZeRO-3 等价方案：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import ShardingStrategy, MixedPrecision
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy

# FSDP 配置
fsdp_config = {
    "sharding_strategy": ShardingStrategy.FULL_SHARD,  # ZeRO-3 等价
    "mixed_precision": MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
        buffer_dtype=torch.bfloat16,
    ),
    "auto_wrap_policy": size_based_auto_wrap_policy(min_num_params=1e6),
    "cpu_offload": None,  # 可配置 CPU offload
    "limit_all_gathers": True,  # 限制同时进行的 all-gather
}

model = MyModel()
model = FSDP(model, **fsdp_config)

# ⚠️ 保存 checkpoint 时需要特殊处理
# 全量保存：
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
    state_dict = model.state_dict()
    if dist.get_rank() == 0:
        torch.save(state_dict, "checkpoint.pt")

A.2 推理引擎命令速查

A.2.1 vLLM

# 安装
pip install vllm

# 基本启动（OpenAI 兼容 API）
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-chat-hf \
    --tensor-parallel-size 2 \
    --port 8000

# 使用 QLoRA
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-chat-hf \
    --quantization awq \
    --dtype half \
    --max-model-len 4096 \
    --gpu-memory-utilization 0.9 \
    --tensor-parallel-size 1

# 性能调优参数
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-13b-chat-hf \
    --tensor-parallel-size 2 \
    --enable-chunked-prefill \     # 分块 prefill
    --max-num-batched-tokens 8192 \ # 批处理 token 上限
    --swap-space 4 \               # KV Cache swap（GB）
    --gpu-memory-utilization 0.92

# vLLM Python API
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    tensor_parallel_size=2,
    max_model_len=4096,
    gpu_memory_utilization=0.9,
    enforce_eager=False,  # 使用 CUDA Graph
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.95,
    max_tokens=512,
)

outputs = llm.generate(["你好，请介绍一下自己。"], sampling_params)

A.2.2 NVIDIA Triton Inference Server

# config.pbtxt - Triton 模型配置
name: "llama-7b"
platform: "vllm"
max_batch_size: 0  # vLLM 内部处理 batching

input [
  {
    name: "prompt"
    data_type: TYPE_STRING
    dims: [1]
  },
  {
    name: "max_tokens"
    data_type: TYPE_INT32
    dims: [1]
    optional: true
  }
]

output [
  {
    name: "text"
    data_type: TYPE_STRING
    dims: [1]
  }
]

dynamic_batching {
  preferred_batch_size: [4, 8, 16]
  max_queue_delay_microseconds: 100000
}

# Triton 启动
docker run --gpus all -p 8000:8000 -p 8001:8001 -p 8002:8002 \
    -v $(pwd)/model_repo:/models \
    nvcr.io/nvidia/tritonserver:24.01-py3 \
    tritonserver --model-repository=/models

# 性能测试
perf_analyzer -m llama-7b -u localhost:8000 \
    --input-data prompt.json \
    --concurrency-range 1:32:2

A.2.3 TensorRT-LLM

# TensorRT-LLM 构建引擎（模型转换）
python build.py \
    --model_dir ./llama-7b-hf \
    --dtype float16 \
    --use_gpt_attention_plugin \
    --use_gemm_plugin \
    --output_dir ./trt_engines/llama-7b \
    --world_size 1 \
    --tp_size 1 \
    --pp_size 1 \
    --max_batch_size 8 \
    --max_input_len 2048 \
    --max_output_len 512 \
    --max_beam_width 1

# 启动 Triton + TensorRT-LLM 后端
python3 tensorrt_llm/serve.py \
    --engine_dir ./trt_engines/llama-7b \
    --tokenizer_dir ./llama-7b-hf \
    --port 8000

A.2.4 SGLang

# 安装
pip install sglang

# 启动服务
python -m sglang.launch_server \
    --model-path meta-llama/Llama-2-7b-chat-hf \
    --tp 2 \
    --port 30000

# 使用 RadixAttention 加速多轮对话
python -m sglang.launch_server \
    --model-path meta-llama/Llama-2-7b-chat-hf \
    --tp 2 \
    --enable-radix-attention \
    --cache-queue-size 10000

A.2.5 llama.cpp

# 下载并编译
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp && make -j

# 模型转换（GGUF 格式）
python convert.py models/llama-7b/ --outtype f16

# 量化
./quantize models/llama-7b/ggml-model-f16.gguf \
    models/llama-7b/ggml-model-q4_k_m.gguf q4_k_m

# 推理（CLI）
./main -m models/llama-7b/ggml-model-q4_k_m.gguf \
    -p "Once upon a time" \
    -n 512 \
    -t 8  # CPU 线程数

# 启动服务（OpenAI 兼容 API）
./server -m models/llama-7b/ggml-model-q4_k_m.gguf \
    --port 8080 \
    --ctx-size 4096 \
    --threads 8

A.2.6 Ollama

# 安装
curl -fsSL https://ollama.com/install.sh | sh

# 拉取模型
ollama pull llama3:8b
ollama pull mistral:7b
ollama pull qwen:7b

# 运行
ollama run llama3:8b

# 启动 API 服务
ollama serve  # 默认 localhost:11434

# 自定义模型（Modelfile）
cat > Modelfile <<EOF
FROM llama3:8b
PARAMETER temperature 0.7
PARAMETER num_ctx 4096
SYSTEM "你是一个专业的 AI 助手。"
EOF

ollama create my-assistant -f Modelfile
ollama run my-assistant

A.3 集群管理工具速查

A.3.1 Kubernetes（GPU 调度）

# GPU Pod 定义
apiVersion: v1
kind: Pod
metadata:
  name: llm-training
spec:
  containers:
  - name: trainer
    image: pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
    resources:
      limits:
        nvidia.com/gpu: 8
        memory: 512Gi
      requests:
        nvidia.com/gpu: 8
        memory: 512Gi
    volumeMounts:
    - name: data
      mountPath: /data
    - name: dshm  # 共享内存（DataLoader 需要）
      mountPath: /dev/shm
  volumes:
  - name: data
    persistentVolumeClaim:
      claimName: training-data
  - name: dshm
    emptyDir:
      medium: Memory
      sizeLimit: 128Gi

A.3.2 Slurm

# Slurm 提交训练任务
sbatch --job-name=llm-train \
    --partition=gpu-h100 \
    --nodes=4 \
    --gres=gpu:H100:8 \
    --ntasks-per-node=8 \
    --cpus-per-task=16 \
    --mem=512G \
    --time=72:00:00 \
    --output=logs/%j.out \
    train.slurm

# train.slurm 脚本
#!/bin/bash
#SBATCH --job-name=llm-train

module load cuda/12.1
module load nccl/2.18

source activate llm-env

# 获取分配的节点
srun hostname | sort > hostfile

torchrun \
    --nnodes=$SLURM_JOB_NUM_NODES \
    --nproc_per_node=8 \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$(head -1 hostfile):29500 \
    train.py --deepspeed ds_config.json

# 交互式分配
salloc --partition=gpu-h100 --nodes=2 --gres=gpu:H100:8 --time=4:00:00

A.3.3 Volcano（K8s 批调度）

# Volcano Job 定义
apiVersion: batch.volcano.sh/v1alpha1
kind: Job
metadata:
  name: llm-training-job
spec:
  minAvailable: 2
  schedulerName: volcano
  policies:
  - event: PodEvicted
    action: RestartJob
  tasks:
  - replicas: 2
    name: trainer
    template:
      spec:
        containers:
        - name: pytorch
          image: pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
          resources:
            limits:
              nvidia.com/gpu: 8
          command:
          - torchrun
          - --nnodes=2
          - --nproc_per_node=8
          - train.py
        restartPolicy: OnFailure

A.4 监控与调试工具速查

A.4.1 NVIDIA Nsight Systems

# 系统级性能分析
nsys profile \
    --output=train_profile \
    --trace=cuda,nvtx,osrt,cudnn,cublas \
    --stats=true \
    python train.py

# 生成报告
nsys stats train_profile.nsys-rep

# GUI 查看
nsys-ui train_profile.nsys-rep

A.4.2 NVIDIA Nsight Compute

# 单 kernel 分析
ncu --set full \
    --target-processes all \
    --output kernel_profile \
    python train.py

# 只分析最耗时的 kernel
ncu --kernel-name regex:".*attention.*" \
    --launch-skip 10 \
    --launch-count 5 \
    python train.py

A.4.3 PyTorch Profiler

from torch.profiler import profile, record_function, ProfilerActivity

# 基础用法
with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, batch in enumerate(loader):
        if step >= 5:  # 采样前几步即可
            break
        loss = model(batch)
        loss.backward()
        optimizer.step()
        prof.step()  # ⚠️ 必须调用

# 打印结果
print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=20,
))

# 导出 Chrome Trace（可视化）
prof.export_chrome_trace("trace.json")

# 检查通信开销
print(prof.key_averages(group_by_input_shape=True).table(
    sort_by="self_cpu_memory_usage",
    row_limit=10,
))

A.4.4 NCCL Tests

# AllReduce 带宽测试
mpirun -np 8 \
    --bind-to none \
    ./build/all_reduce_perf -b 8 -e 128M -f 2 -g 8

# AllGather 带宽测试
mpirun -np 8 \
    ./build/all_gather_perf -b 1K -e 64M -f 2 -g 8

# Point-to-Point 带宽测试（检查 NVLink / IB）
./build/sendrecv_perf -b 4 -e 128M -f 2 -g 1

# 多机 AllReduce 测试
mpirun -np 16 \
    -H node1:8,node2:8 \
    -bind-to none \
    -x NCCL_SOCKET_IFNAME=eth0 \
    -x NCCL_IB_HCA=mlx5 \
    ./build/all_reduce_perf -b 1M -e 128M -f 2 -g 8

# NCCL 调试环境变量
export NCCL_DEBUG=INFO          # 打印详细日志
export NCCL_DEBUG_SUBSYS=ALL    # 所有子系统日志
export NCCL_NET_GDR_LEVEL=PHB   # GPUDirect RDMA 级别
export NCCL_IB_TIMEOUT=22       # IB 超时
export NCCL_IB_RETRY_CNT=7      # IB 重试次数
export NCCL_SOCKET_NTHREADS=4   # Socket 线程数
export NCCL_NSOCKS_PERTHREAD=4  # 每线程 Socket 数

Tip调试流程建议

遇到多机训练问题时的标准排查顺序：(1) 检查网络连通性 ping + nc -zv; (2) 检查 GPU 可见性 nvidia-smi; (3) 运行 NCCL test 确认带宽; (4) 检查防火墙/IB 配置; (5) 最后才是看代码。

小结

本附录列出了 AI Infra 工程师最常用的工具命令。几点使用建议：

1. 不要死记命令——理解原理，善用 --help 和文档。
2. 建立自己的工具脚本库——把常用组合封装成脚本。
3. 关注工具版本——AI Infra 工具更新极快，升级前务必看 changelog。
4. 性能基准要自己跑——不要完全依赖官方 benchmark，你的硬件和模型组合是唯一的。

延伸阅读

• PyTorch 分布式训练文档: https://pytorch.org/tutorials/beginner/dist_overview.html
• DeepSpeed 文档: https://www.deepspeed.ai/
• Megatron-LM GitHub: https://github.com/NVIDIA/Megatron-LM
• vLLM 文档: https://docs.vllm.ai/
• NVIDIA Nsight 文档: https://docs.nvidia.com/nsight-systems/
• NCCL 文档: https://docs.nvidia.com/deeplearning/nccl/
• Volcano 项目: https://volcano.sh/en/docs/