Ch13 §04 并发与并行 · 讲解稿¶
章节定位: Ch13 多核时代。教"用上所有算力": 同步原语 / CUDA / 分布式 / MapReduce。 承接: §03 进程/线程 → §04 跨多核 + 跨多机 走向: §05 编程语言 (类型安全) / 大模型训练 篇幅: ~4200 字 / 阅读 12 分钟 / 讲解 25-30 分钟
一 · 一句话总览¶
并发 vs 并行: 并发 = 多任务交替 (单核可), 并行 = 同时执行 (需多核)。AI 大模型训练 = GPU 并行 (CUDA SIMT) + 多卡并行 (DDP) + 多机并行 (FSDP/Megatron)。
二 · 心智模型¶
- 并发: 一心多用 (打电话 + 查邮件)
- 并行: 多双手同时干 (工厂流水线)
AI 训练 = 大规模并行。
三 · 章节主线 (5 段)¶
第 1 段 · 同步原语 (8 分钟)¶
3.1.1 互斥锁 (Mutex)¶
确保只有一个线程进入临界区。
import threading
lock = threading.Lock()
counter = 0
def increment():
global counter
with lock: # 自动获取/释放
for _ in range(100000):
counter += 1
threads = [threading.Thread(target=increment) for _ in range(4)]
for t in threads: t.start()
for t in threads: t.join()
print(counter) # 应为 400000
3.1.2 信号量 (Semaphore)¶
允许 N 个线程同时进入 (N 元信号量)。
import threading
sem = threading.Semaphore(3) # 3 个槽
def access(i):
sem.acquire()
print(f"Thread {i} accessing")
sem.release()
3.1.3 条件变量¶
线程等待某条件成立。
3.1.4 原子操作¶
硬件级不可中断。Python queue.Queue 内部用原子 + 锁。
3.1.5 常见并发问题¶
| 问题 | 现象 | 解 |
|---|---|---|
| 竞态 | 数据不一致 | 锁 |
| 死锁 | 互相等待 | 排序 / 超时 |
| 活锁 | 不断让步 | 随机退避 |
| 饥饿 | 长期得不到 | 公平调度 |
| 优先级反转 | 高优先级等低优先级 | 优先级继承 |
第 2 段 · GPU 与 CUDA (10 分钟)¶
3.2.1 CUDA 编程模型¶
Host (CPU) Device (GPU)
┌────────┐ ┌────────────┐
│ Serial │ │ Parallel │
│ Code │ ────→ │ Kernels │
└────────┘ └────────────┘
3.2.2 Thread 层次¶
| 层级 | 大小 | 说明 |
|---|---|---|
| Thread | 1 | 单线程 |
| Warp | 32 | 调度单位 |
| Block | 32-1024 | 含多 warp |
| Grid | 多 block | 整个 kernel 启动 |
索引:
3.2.3 向量加法 kernel¶
__global__ void vecAdd(float* a, float* b, float* c, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) c[idx] = a[idx] + b[idx];
}
// Host 启动
int n = 1000000;
int threads = 256;
int blocks = (n + threads - 1) / threads;
vecAdd<<<blocks, threads>>>(a, b, c, n);
3.2.4 内存层级¶
| 类型 | 作用域 | 速度 |
|---|---|---|
| Register | Thread | 极快 |
| Shared | Block | 快 (~100× global) |
| Local | Thread | 慢 (实际在 global) |
| Global | 全 Grid | 慢 (但大) |
| Constant | 全 Grid | 缓存读 |
__global__ void matmul(float* a, float* b, float* c, int n) {
__shared__ float tile[16][16]; // 共享内存
int bx = blockIdx.x, by = blockIdx.y;
int tx = threadIdx.x, ty = threadIdx.y;
// 加载 tile 到共享内存
tile[ty][tx] = a[(by * 16 + ty) * n + bx * 16 + tx];
__syncthreads();
// 计算
float sum = 0;
for (int k = 0; k < n; k++) {
sum += tile[ty][k] * tile[k][tx];
}
c[(by * 16 + ty) * n + bx * 16 + tx] = sum;
}
3.2.5 PyTorch CUDA¶
import torch
# GPU tensor
x = torch.randn(1024, 1024, device='cuda')
# 隐式 CUDA 调用
y = x @ x.T # matmul on GPU
# 同步
torch.cuda.synchronize()
# 计时
torch.cuda.Event(enable_timing=True)
第 3 段 · 分布式并行 (10 分钟)¶
3.3.1 Amdahl 定律¶
\[S(n) = \frac{1}{(1-p) + p/n}\]
\(p\) = 可并行比例。即使 \(p = 0.95\), 100 核加速 ≤ 16×。
3.3.2 并行策略¶
| 策略 | 切什么 | 适用 |
|---|---|---|
| 数据并行 (DP) | 切 batch | 模型小, 数据多 |
| 模型并行 (MP) | 切模型 | 模型大, 单卡装不下 |
| 流水线 (PP) | 切深度 | 深层网络 |
| 张量并行 (TP) | 切单层 | 大 hidden |
3.3.3 集合通信¶
| 操作 | 含义 |
|---|---|
| AllReduce | 同步到所有 (梯度) |
| Reduce | 同步到主节点 |
| Broadcast | 主→从 |
| AllGather | 每节点拿全量 |
| Reduce-Scatter | 每节点拿分片 |
3.3.4 NCCL¶
NVIDIA Collective Communications Library, GPU 优化的集合通信。
import torch.distributed as dist
# 初始化
dist.init_process_group(backend='nccl')
# AllReduce
tensor = torch.randn(1024, device='cuda')
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
3.3.5 PyTorch DDP¶
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 启动: torchrun --nproc_per_node=4 train.py
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = Model().to(local_rank)
model = DDP(model, device_ids=[local_rank])
# 训练 loop 同普通
3.3.6 FSDP / ZeRO¶
切模型参数到多卡, 每卡只持有 1/N 参数。
70B 模型 ~140 GB (FP16) → 8 卡 ZeRO-3 每卡 ~17 GB。
第 4 段 · MapReduce 与大数据 (8 分钟)¶
3.4.1 编程模型¶
# Mapper
def map_fn(record):
for word in record.split():
yield (word, 1)
# Reducer
def reduce_fn(key, values):
return (key, sum(values))
3.4.2 关键设计¶
- 数据本地化: 计算靠近数据
- 容错: 重试失败的 task
- 水平扩展: 加机器 = 加能力
3.4.3 演进¶
3.4.4 PyTorch + Ray¶
import ray
ray.init()
@ray.remote(num_gpus=1)
def train_one_model(config):
# 训练一个模型
return model
# 并行 8 个 trial
results = ray.get([train_one_model.remote(cfg) for cfg in configs])
第 5 段 · 实践: 大模型训练 (8 分钟)¶
3.5.1 LLM 训练基础设施¶
3.5.2 启动命令¶
# 单机 8 卡
torchrun --nproc_per_node=8 train.py
# 多机 32 卡 (4 机 × 8)
torchrun \
--nproc_per_node=8 \
--nnodes=4 \
--node_rank=0 \
--master_addr=head-node \
--master_port=29500 \
train.py
3.5.3 FSDP 训练 LLaMA¶
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
model = FSDP(
model,
sharding_strategy=ShardingStrategy.FULL_SHARD, # ZeRO-3
mixed_precision=MixedPrecision(
param_dtype=torch.bfloat16,
reduce_dtype=torch.float32
)
)
3.5.4 检查点¶
# 保存
torch.save({
'model': model.state_dict(),
'opt': opt.state_dict(),
'epoch': epoch,
}, 'ckpt.pt')
# FSDP 保存 (需要特殊处理)
from torch.distributed.fsdp import StateDictType, FullStateDictConfig
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, FullStateDictConfig()):
state = model.state_dict()
torch.save(state, 'ckpt.pt')
四 · 易错点 (5 条)¶
- Python GIL: 多线程对 CPU 密集任务无效, 用
multiprocessing。 - CUDA 同步: GPU 操作默认异步, 必须
torch.cuda.synchronize()才计时准确。 - AllReduce 是同步: DDP 每次 backward 后等所有节点就绪, 网络慢 = 训练慢。
- NCCL 端口冲突: 多任务同节点需不同
MASTER_PORT。 - FSDP checkpoint: 必须用
state_dict_type, 直接保存失效。
五 · 跨章连接¶
| 章节 | 关联 |
|---|---|
| Ch03 §05 梯度 | GPU 加速 |
| Ch07 §04 Transformer | DDP 训练 GPT |
| Ch12 §05 EGNN | GPU 分子模拟 |
| Ch11 §05 VLA | 边缘推理 |
六 · 自测 5 题¶
- Amdahl 公式? \(p=0.9, n=10\) 加速比?
- 死锁 vs 活锁 vs 饥饿?
- CUDA warp/block/grid 关系?
- AllReduce vs Reduce 区别?
- ZeRO-3 vs DDP 显存节省?
七 · 一句话总结¶
并发并行 = 把算力用满。同步原语 (锁/信号量) → CUDA (GPU SIMT) → 分布式 (DDP/FSDP) → 大模型训练 (Megatron/DeepSpeed)。AI 时代 = 并行计算时代。
八 · 板书建议¶
[板书 1: 同步]
Mutex / Semaphore / CV / Atomic
死锁四条件
[板书 2: CUDA]
Thread → Warp(32) → Block → Grid
Shared / Global / Register
[板书 3: Amdahl]
S = 1 / (1-p + p/n)
[板书 4: 分布式]
DP / TP / PP / FSDP
NCCL AllReduce
九 · 参考资料¶
- NVIDIA CUDA Programming Guide
- PyTorch Distributed 文档
- ZeRO 论文 (Microsoft, 2020)
- Megatron-LM 论文 (NVIDIA, 2019)
- Ch07 §04 Transformer (本套书)