1.0 编程语言入门

发表于 更新于 分类于 ， ，

AI Infra 的大厦建立在编程能力之上。本文不是通用编程教程，而是专门回答一个问题：想做 AI Infra，编程底子需要打到什么程度？从 Python、C/C++、Linux 到数学基础，为每个方向划定”够用”的边界，配上面向实际场景的代码示例。

📑 目录

1. Python：AI 生态的通用语言

1.1 为什么 Python 是 AI Infra 的主力语言

你可以把编程语言想象成各行各业的工作语言。做外贸得会英语，做 AI Infra 就得会 Python——不是因为它跑得最快，而是因为整个 AI 生态圈都用它交流。

具体来说，Python 在 AI Infra 中扮演着”胶水语言”和”控制层语言”的双重角色：

  • 胶水语言：PyTorch、TensorFlow、DeepSpeed、vLLM、Triton 等几乎所有 AI 框架的用户接口都是 Python。你用 Python 定义模型结构、编排训练流程、启动推理服务，底层的高性能计算由 C++/CUDA 内核完成，但你调度这些内核用的是 Python。
  • 控制层语言：分布式训练的进程管理、数据预处理流水线、性能分析脚本、自动化测试——这些”围绕模型”的工程工作全部用 Python 完成。一个 AI Infra 工程师日常写的 Python 代码，可能比写 CUDA 代码多得多。
  • 生态优势不可替代：NumPy、pandas 做数据处理，matplotlib 画图，Jupyter Notebook 做实验，pip/conda 管理依赖。这套工具链成熟到没有任何语言能在短期内替代。

一句话总结：Python 不负责”算得快”（那是 CUDA 的事），但负责”把一切组织起来”。

1.2 需要掌握到什么程度

AI Infra 对 Python 的要求比写算法题要高，比做 Web 后端要专。你需要的不是”会写 for 循环”，而是能写出结构清晰、可调试、可复用的工程代码。以下是几个必须过关的知识点：

面向对象编程（OOP）

PyTorch 的核心抽象全建立在 OOP 之上：nn.Module 是所有模型的基类，你需要通过继承和组合来搭建网络结构。不理解类的继承、方法重写、__init__ 和 forward 的配合，就没法读懂任何一个模型实现。

装饰器（Decorator）

装饰器的本质是”不改原函数代码，给它加一层包装”。打个比方，你有一个函数负责跑矩阵乘法，你想知道它跑了多久，但又不想在函数里面加计时代码——装饰器就是帮你在函数”外面”套一个计时器。在技术上，装饰器是一个接收函数作为参数、返回新函数的高阶函数。AI Infra 中常见的用途包括：性能计时、日志记录、重试机制、@torch.no_grad() 关闭梯度计算等。

生成器（Generator）

生成器可以理解为”按需生产的流水线”：数据不是一次全部加载到内存，而是用到一条才生产一条。对于动辄几十 GB 的训练数据集，如果一股脑全读进内存，内存直接爆掉。生成器通过 yield 关键字实现惰性求值（lazy evaluation），PyTorch 的 DataLoader 底层就依赖这个机制。

多进程与多线程

  • 多线程适合 I/O 密集型任务（比如同时下载多个文件），因为 Python 的 GIL（全局解释器锁）限制了多线程在 CPU 计算上的并行能力。
  • 多进程绕开了 GIL，每个进程有独立的 Python 解释器和内存空间，适合 CPU 密集型任务（比如数据预处理、Tokenization）。分布式训练中，torch.distributed 就是基于多进程模型的——每个 GPU 对应一个独立进程。

性能 Profiling

光让代码跑起来不够，还得知道它慢在哪。Python 内置的 cProfile 可以统计每个函数的调用次数和耗时；line_profiler 可以定位到具体哪一行代码在拖后腿；torch.profiler 则能分析 PyTorch 模型训练中 CPU 和 GPU 的耗时分布。养成”写完代码先 profile 一遍”的习惯，比盲目优化高效得多。

1.3 实用代码示例

示例 1：用装饰器给函数计时

在 AI Infra 的日常工作中，你经常需要知道某个函数到底跑了多久——是数据加载慢还是模型前向传播慢。下面这个装饰器可以挂在任意函数上，自动打印执行耗时：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
import time
import functools

def timer(func):
    """打印被装饰函数的执行耗时"""
    @functools.wraps(func)  # 保留原函数的名称和文档字符串
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        elapsed = time.perf_counter() - start
        print(f"[{func.__name__}] 耗时 {elapsed:.4f}s")
        return result
    return wrapper

@timer
def simulate_data_loading(num_samples: int) -> list:
    """模拟从磁盘加载训练数据"""
    data = []
    for i in range(num_samples):
        # 模拟 I/O 延迟和数据处理
        data.append([float(x) for x in range(128)])
    return data

@timer
def simulate_forward_pass(data: list) -> float:
    """模拟模型前向传播"""
    total = 0.0
    for sample in data:
        total += sum(x * 0.01 for x in sample)
    return total

if __name__ == "__main__":
    data = simulate_data_loading(50000)
    loss = simulate_forward_pass(data)
    print(f"模拟 loss: {loss:.2f}")

运行后会输出类似：

1
2
3
[simulate_data_loading] 耗时 0.8123s
[simulate_forward_pass] 耗时 0.3456s
模拟 loss: 3200.00

一眼就能看出数据加载才是瓶颈。

示例 2：多进程加速数据预处理

训练大模型之前，通常需要对原始文本做 Tokenization。如果用单进程逐条处理，几十 GB 的数据集可能要跑好几个小时。用多进程可以把数据切成若干分片，每个 CPU 核心处理一片，大幅缩短时间：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
import multiprocessing as mp
import time
import os

def tokenize_chunk(chunk: list[str]) -> list[list[int]]:
    """
    模拟 Tokenization：将文本转换为 token ID 列表。
    实际场景中这里会调用 tokenizer.encode()。
    """
    results = []
    for text in chunk:
        # 模拟 tokenize：把每个字符转成 ASCII 码
        token_ids = [ord(c) for c in text[:512]]
        results.append(token_ids)
    return results

def parallel_tokenize(texts: list[str], num_workers: int = None) -> list[list[int]]:
    """多进程并行 Tokenization"""
    if num_workers is None:
        num_workers = min(mp.cpu_count(), 8)  # 最多用 8 个核心

    # 将数据均匀切分成 num_workers 份
    chunk_size = len(texts) // num_workers
    chunks = []
    for i in range(num_workers):
        start = i * chunk_size
        end = start + chunk_size if i < num_workers - 1 else len(texts)
        chunks.append(texts[start:end])

    # 启动进程池并行处理
    with mp.Pool(processes=num_workers) as pool:
        results = pool.map(tokenize_chunk, chunks)

    # 合并所有分片的结果
    all_tokens = []
    for chunk_result in results:
        all_tokens.extend(chunk_result)
    return all_tokens

if __name__ == "__main__":
    # 模拟 10 万条训练文本
    texts = [f"This is sample text number {i} for training." for i in range(100_000)]

    # 单进程基线
    start = time.perf_counter()
    single_result = tokenize_chunk(texts)
    single_time = time.perf_counter() - start
    print(f"单进程: {single_time:.2f}s, 处理 {len(single_result)} 条")

    # 多进程加速
    start = time.perf_counter()
    multi_result = parallel_tokenize(texts, num_workers=4)
    multi_time = time.perf_counter() - start
    print(f"4 进程: {multi_time:.2f}s, 处理 {len(multi_result)} 条")
    print(f"加速比: {single_time / multi_time:.1f}x")

示例 3：用 cProfile 定位性能瓶颈

当一段代码跑得很慢但你不确定慢在哪里时，cProfile 是最直接的诊断工具。它会统计每个函数被调用了多少次、总共花了多少时间：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
import cProfile
import pstats
import io
import math

def load_data(n: int) -> list[list[float]]:
    """模拟加载数据"""
    return [[math.sin(i * j) for j in range(256)] for i in range(n)]

def normalize(data: list[list[float]]) -> list[list[float]]:
    """逐行归一化（故意用低效实现来演示 profiling）"""
    result = []
    for row in data:
        row_min = min(row)
        row_max = max(row)
        span = row_max - row_min if row_max != row_min else 1.0
        normalized = [(x - row_min) / span for x in row]
        result.append(normalized)
    return result

def compute_mean(data: list[list[float]]) -> list[float]:
    """计算每行均值"""
    return [sum(row) / len(row) for row in data]

def pipeline():
    """完整的数据处理流水线"""
    data = load_data(2000)
    normed = normalize(data)
    means = compute_mean(normed)
    return means

if __name__ == "__main__":
    # 用 cProfile 分析 pipeline 的性能
    profiler = cProfile.Profile()
    profiler.enable()

    result = pipeline()

    profiler.disable()

    # 按累计耗时排序，只看前 15 行
    stream = io.StringIO()
    stats = pstats.Stats(profiler, stream=stream).sort_stats("cumulative")
    stats.print_stats(15)
    print(stream.getvalue())

输出会清楚地告诉你 normalize 函数占了绝大部分时间，因为它对每一行都做了三次遍历（min、max、列表推导）。知道瓶颈在哪，才能有的放矢地优化——比如用 NumPy 向量化替换逐元素循环。

2. C/C++：CUDA 编程的宿主语言

2.1 AI Infra 对 C++ 的需求层级

先说结论：AI Infra 不要求你成为 C++ 大师，但要求你能”读得懂、改得动”。

这里有一个常见的误解——很多人看到 CUDA、PyTorch 底层、推理引擎都是 C++ 写的，就觉得必须把 C++ 学到精通模板元编程的程度才能入行。实际上，AI Infra 对 C++ 的需求可以分成三个层级：

层级 能力要求 对应场景
基础（必须） 能读懂 C/C++ 代码逻辑，理解指针、内存分配、编译链接 阅读 CUDA host 端代码、理解 PyTorch C++ 扩展
进阶（推荐） 能写简单的 C++ 函数并编译运行，理解类和结构体 编写 CUDA kernel 的 host 端包装、修改现有算子代码
深入（锦上添花） 模板、STL 容器、RAII、移动语义 深度参与推理引擎开发（vLLM C++ backend、TensorRT-LLM）

大多数 AI Infra 工程师的日常工作停留在前两个层级。你需要做的是：拿到一段 CUDA 代码，能看懂它在做什么、数据怎么从 CPU 搬到 GPU、kernel 怎么启动的。这比”从零写一个 C++ 项目”简单得多。

2.2 必须理解的核心概念

指针与内存管理

指针是 C/C++ 的灵魂，也是很多人的噩梦。但你可以先建立一个简单的直觉：变量是一个”盒子”，里面装着数据；指针是一张”纸条”，上面写着盒子的地址。通过这张纸条，你可以找到盒子、读取或修改里面的内容。

在技术上，指针就是一个存储内存地址的变量。int *p = &x 表示 p 存储了变量 x 的内存地址，*p 则是通过这个地址访问 x 的值（解引用）。

为什么这对 AI Infra 很重要？因为 CUDA 编程的核心操作之一就是在 CPU 和 GPU 之间搬运数据，而搬运数据的接口全靠指针：

  • cudaMalloc(&d_ptr, size) —— 在 GPU 上分配一块内存，把地址存进 d_ptr
  • cudaMemcpy(d_ptr, h_ptr, size, cudaMemcpyHostToDevice) —— 把 CPU 指针 h_ptr 指向的数据搬到 GPU 指针 d_ptr 指向的位置
  • cudaFree(d_ptr) —— 释放 GPU 内存

如果你不理解指针，上面这三行代码就完全看不懂。

编译链接过程

Python 是解释执行的——写完代码直接 python xxx.py 就能跑。C/C++ 则需要先编译（把源代码翻译成机器码）再链接（把多个编译好的模块拼接成一个可执行文件），最后才能运行。

这个过程可以类比翻译出版一本书：编译就像把每一章从中文翻译成英文（源文件 -> 目标文件），链接就像把翻译好的各章合订成一本完整的书（目标文件 -> 可执行文件），中间还要确保章节之间的引用对得上（符号解析）。

CUDA 代码的编译稍有特殊：.cu 文件通过 nvcc（NVIDIA 的编译器）处理，它会把 GPU 代码和 CPU 代码分别编译，最后合在一起。你在 AI Infra 项目中见到的 CMakeLists.txt 或 setup.py 中的编译配置，本质上就是在告诉编译器：哪些文件需要 nvcc 编译，哪些用 g++，怎么链接在一起。

基本的类和结构体

C++ 的 class 和 struct 用于把相关的数据和操作打包在一起。在 CUDA 相关代码中，你经常会看到用结构体来描述 kernel 的参数：

1
2
3
4
5
6
7
struct AttentionParams {
    float *q, *k, *v, *output;  // 指向 Q、K、V 和输出矩阵的指针
    int batch_size;
    int seq_len;
    int head_dim;
    float scale;                 // 1/sqrt(head_dim)
};

不需要深入理解 C++ 的构造函数、运算符重载、虚函数这些高级特性。能看懂结构体在打包什么数据、类的方法在做什么操作，就够用了。

2.3 典型代码示例

下面是一个简化的 CUDA host 端代码结构，展示了 CPU 和 GPU 之间最基本的交互流程——分配显存、上传数据、启动 kernel、取回结果、释放显存。这就是你在阅读几乎所有 CUDA 项目时都会碰到的骨架：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#include <cstdio>
#include <cuda_runtime.h>

// GPU 端的计算函数（kernel）：每个线程负责一个元素的向量加法
// __global__ 关键字表示这个函数在 GPU 上执行，由 CPU 端发起调用
__global__ void vector_add(const float *a, const float *b, float *c, int n) {
    // 计算当前线程的全局索引
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        c[idx] = a[idx] + b[idx];
    }
}

int main() {
    const int N = 1024;
    const size_t bytes = N * sizeof(float);

    // ---- 第 1 步：在 CPU（Host）上准备数据 ----
    float *h_a = new float[N];
    float *h_b = new float[N];
    float *h_c = new float[N];
    for (int i = 0; i < N; i++) {
        h_a[i] = static_cast<float>(i);
        h_b[i] = static_cast<float>(i * 2);
    }

    // ---- 第 2 步：在 GPU（Device）上分配显存 ----
    float *d_a, *d_b, *d_c;
    cudaMalloc(&d_a, bytes);
    cudaMalloc(&d_b, bytes);
    cudaMalloc(&d_c, bytes);

    // ---- 第 3 步：把数据从 CPU 搬到 GPU ----
    cudaMemcpy(d_a, h_a, bytes, cudaMemcpyHostToDevice);
    cudaMemcpy(d_b, h_b, bytes, cudaMemcpyHostToDevice);

    // ---- 第 4 步：启动 GPU kernel ----
    // <<<grid_size, block_size>>> 是 CUDA 特有的语法
    // 256 个线程为一组（block），共需 ceil(N/256) 个 block
    int threads_per_block = 256;
    int num_blocks = (N + threads_per_block - 1) / threads_per_block;
    vector_add<<<num_blocks, threads_per_block>>>(d_a, d_b, d_c, N);

    // ---- 第 5 步：把结果从 GPU 搬回 CPU ----
    cudaMemcpy(h_c, d_c, bytes, cudaMemcpyDeviceToHost);

    // ---- 第 6 步：验证结果 ----
    printf("h_c[0] = %.1f, h_c[1023] = %.1f\n", h_c[0], h_c[1023]);
    // 预期输出: h_c[0] = 0.0, h_c[1023] = 3069.0

    // ---- 第 7 步：释放显存和内存 ----
    cudaFree(d_a);
    cudaFree(d_b);
    cudaFree(d_c);
    delete[] h_a;
    delete[] h_b;
    delete[] h_c;

    return 0;
}

编译和运行：

1
2
nvcc -o vector_add vector_add.cu
./vector_add

读这段代码时，抓住核心流程就行：cudaMalloc -> cudaMemcpy(H2D) -> kernel<<<...>>>() -> cudaMemcpy(D2H) -> cudaFree。这五步是所有 CUDA 程序的骨架，无论多复杂的推理引擎，底层都在重复这个模式。

3. Linux 基础：AI Infra 的工作台

3.1 为什么必须熟悉 Linux

一句话：AI Infra 的开发和部署几乎 100% 发生在 Linux 上。

原因很实际：

  • GPU 驱动和 CUDA 工具链：NVIDIA 的 GPU 驱动、CUDA Toolkit、cuDNN 等核心依赖主要面向 Linux 开发和优化。虽然 Windows 也能跑 CUDA，但几乎没有生产环境这么做。
  • 服务器环境：训练集群、推理服务器全是 Linux（通常是 Ubuntu 或 CentOS）。你不可能在这些机器上装个 Windows 桌面。
  • 容器化部署：Docker 容器是 Linux 原生技术，AI 模型的部署大量依赖容器。
  • 开源工具链：PyTorch 编译、DeepSpeed 安装、vLLM 部署——这些工具的文档和 CI 都默认 Linux 环境，遇到问题搜到的解答也几乎都是 Linux 的。

你可以把 Linux 命令行想象成 AI Infra 工程师的”工作台”——就像木匠离不开锯子和刨子，你离不开终端和 Shell。不需要成为 Linux 内核开发者，但必须能在命令行下自如地完成日常工作。

3.2 日常高频操作清单

以下是 AI Infra 工程师每天都在用的操作，每一项你都应该熟练到不用查文档：

SSH 远程连接

训练服务器通常不在你面前，你需要通过 SSH 远程登录。建议配置好 ~/.ssh/config，这样就不用每次输入完整的用户名和 IP：

1
2
3
4
5
6
# ~/.ssh/config 示例
Host gpu-server
    HostName 10.0.1.100
    User zhangsan
    Port 22
    IdentityFile ~/.ssh/id_rsa

配置好之后 ssh gpu-server 一条命令就能连上。

tmux 会话管理

训练一个大模型可能要跑好几天。如果你的 SSH 连接断了，没有 tmux 保护的进程会直接被杀死。tmux 的作用是在服务器上创建一个”持久会话”，断开连接后进程继续跑，重新连上就能恢复现场。

1
2
3
4
tmux new -s train       # 创建名为 train 的会话
tmux attach -t train    # 重新连接到 train 会话
tmux ls                 # 列出所有会话
# 在 tmux 内按 Ctrl+B 然后按 D，脱离会话但不终止进程

conda/pip 环境管理

不同项目可能依赖不同版本的 PyTorch、CUDA。用 conda 创建隔离的虚拟环境是基本素养：

1
2
3
conda create -n llm python=3.10
conda activate llm
pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu121

nvidia-smi 查看 GPU 状态

这是你用得最多的命令之一。一眼看出哪些卡在用、显存占了多少、GPU 利用率如何：

1
2
3
nvidia-smi                        # 快照式查看
watch -n 1 nvidia-smi             # 每秒刷新，持续监控
nvidia-smi topo -m                # 查看 GPU 之间的互联拓扑

git 版本管理

代码改了什么、什么时候改的、出了问题怎么回滚——这些全靠 git：

1
2
3
4
git clone <repo>
git checkout -b feature/optimize-kernel
git add . && git commit -m "optimize softmax kernel"
git push origin feature/optimize-kernel

bash 脚本批量提交任务

在集群上跑实验，经常需要用脚本批量启动不同配置的训练任务：

1
2
3
4
5
6
7
8
9
10
#!/bin/bash
for lr in 1e-4 3e-4 1e-3; do
    for bs in 16 32 64; do
        echo "Starting: lr=${lr}, batch_size=${bs}"
        python train.py --lr $lr --batch_size $bs \
            --output_dir results/lr${lr}_bs${bs} &
    done
done
wait
echo "All experiments done."

3.3 实用命令速查表

类别 命令 用途
文件操作 ls -lah 列出文件详情（含隐藏文件、人类可读大小）
du -sh dir/ 查看目录占用的磁盘空间
df -h 查看磁盘剩余空间
find . -name "*.pt" -size +1G 找出当前目录下大于 1GB 的 .pt 文件
tar -czf backup.tar.gz model/ 打包压缩模型目录
scp -r user@server:~/model ./ 从远程服务器下载模型文件
进程管理 ps aux | grep python 查看所有 Python 进程
kill -9 <PID> 强制终止进程
nohup python train.py & 后台运行（不如 tmux 好用，但偶尔需要）
top / htop 实时查看 CPU 和内存使用
GPU 状态 nvidia-smi GPU 快照信息
nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv 结构化查询 GPU 状态
nvidia-smi topo -m GPU 互联拓扑（NVLink / PCIe）
nvcc --version 查看 CUDA 编译器版本
环境管理 conda env list 列出所有虚拟环境
pip list | grep torch 查看已安装的 torch 版本
echo $CUDA_HOME 查看 CUDA 安装路径
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH 添加库搜索路径
网络与传输 wget <url> / curl -O <url> 下载文件
rsync -avz src/ user@server:dst/ 增量同步文件（比 scp 更高效）
ssh -L 8888:localhost:8888 server SSH 端口转发（访问远程 Jupyter）
文本处理 tail -f train.log 实时查看训练日志输出
grep "loss" train.log | tail -20 过滤日志中的 loss 信息
wc -l dataset.jsonl 统计数据集行数

4. 数学基础：够用就好

AI Infra 不是做算法研究，你不需要推导定理、证明收敛性。但数学直觉是理解模型行为和优化策略的底层支撑——比如看到一个矩阵维度变换要能秒反应结果形状，看到 Softmax 要知道它在把数值映射成概率。下面列出”够用”标准下最核心的数学知识。

4.1 线性代数

核心需求：对矩阵运算有维度直觉

大模型的计算核心就是矩阵乘法。Transformer 里的 Attention（QK^T 矩阵乘法）、FFN（两层线性变换）、Embedding 查表——归根结底都是矩阵运算。你需要的不是证明矩阵分解定理，而是以下这种”维度直觉”：

看到 (B, S, H) x (H, V)，能立刻知道结果形状是 (B, S, V)。

这里 B 是 batch size（批次大小），S 是序列长度，H 是隐藏维度，V 是词表大小。矩阵乘法的规则是”前面矩阵的列数必须等于后面矩阵的行数”，结果矩阵取”前面的行数 x 后面的列数”。批次维度 B 不参与乘法，直接保留。

实际例子：LLaMA-7B 模型最后一层把隐藏状态 (B, S, 4096) 乘以词表投影矩阵 (4096, 32000)，得到 (B, S, 32000) 的 logits。如果你没有这个维度直觉，就没法理解模型参数量是怎么算出来的，也没法理解张量并行是沿哪个维度切的。

建议学习内容：

  • 矩阵乘法的定义和维度规则
  • 转置的含义（行列互换，用于计算 QK^T）
  • 分块矩阵的概念（理解 GPU 上矩阵乘法为什么要做 tiling）
  • 向量的点积和外积

4.2 概率论与统计

Softmax 的概率解释

Softmax 函数的作用是把一组任意实数变成一组概率值——所有输出都在 0 到 1 之间，加起来等于 1。你可以把它想象成”投票归一化”：模型对每个候选 token 打了一个分，Softmax 把这些分数转化成”选择每个 token 的概率”。

数学定义：给定输入向量 z = (z_1, z_2, …, z_n)，Softmax 的第 i 个输出为：

$$\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}}$$

在 AI Infra 中，Softmax 几乎无处不在——Attention 的权重计算（QK^T 的结果经过 Softmax 得到注意力权重）、推理时的采样（logits 经过 Softmax 得到下一个 token 的概率分布）。而且因为它涉及指数运算和求和，如何高效稳定地计算 Softmax 是 CUDA 算子优化的经典课题。

交叉熵损失

训练大模型的目标函数几乎都是交叉熵损失（Cross-Entropy Loss）。直觉上，它衡量的是”模型预测的概率分布和真实答案之间的差距”——模型越确信正确答案，损失越小；模型把概率分散到错误选项上，损失就大。

$$L = -\sum_{i=1}^{n} y_i \log(p_i)$$

其中 y 是真实标签（one-hot 向量），p 是模型预测的概率分布。在语言模型中，这退化成 L = -log(p_correct)：只看模型给正确 token 分配了多少概率。

理解交叉熵的关键在于理解 -log(p) 的行为：当 p 接近 1 时（模型很确信），-log(p) 接近 0（损失很小）；当 p 接近 0 时（模型猜错了），-log(p) 趋向无穷大（损失爆炸）。

4.3 微积分

链式法则与反向传播

深度学习的训练依赖梯度下降：计算损失函数对每个参数的梯度，然后沿梯度反方向更新参数。而计算梯度的核心工具就是链式法则。

链式法则可以用”多米诺骨牌”来理解：模型是一连串函数的嵌套 f(g(h(x)))，你要知道最终结果对最初输入的变化率，只需要把每一步的变化率乘起来：

$$\frac{df}{dx} = \frac{df}{dg} \cdot \frac{dg}{dh} \cdot \frac{dh}{dx}$$

反向传播（Backpropagation）就是从输出层开始，逐层应用链式法则，把梯度”反向传回”每一层。PyTorch 的 loss.backward() 自动帮你完成这个过程。

梯度的含义

梯度在直觉上就是”如果这个参数往正方向动一点点，损失会变大还是变小、变多少”。梯度为正表示参数增大会让损失增大（所以要减小参数），梯度为负则相反。

为什么 AI Infra 要关心梯度？因为：

  • 混合精度训练中的梯度溢出：FP16 的表示范围很小，如果梯度值极小（例如 1e-8），FP16 直接下溢变成 0，模型就学不动了。这就是为什么混合精度训练需要 Loss Scaling——先把损失放大，让梯度值回到 FP16 能表示的范围内，更新参数时再缩放回来。
  • 分布式训练中的梯度同步：DDP 的核心操作就是把各卡的梯度做 AllReduce（求平均），理解梯度是什么才能理解这个同步操作在做什么。

📝 总结

AI Infra 的编程基础可以用一句话概括：Python 主攻、C++ 辅助、Linux 是战场、数学是直觉。

四个方向的优先级和投入比例：

方向 优先级 投入比例 一句话要求
Python 最高 40% 写得出工程级代码，能 profile，会多进程
Linux 高 25% 在命令行下自如工作，不依赖图形界面
C/C++ 中 20% 读得懂 CUDA host 代码，理解指针和内存
数学 中 15% 有维度直觉，理解 Softmax、梯度、矩阵乘法

不要追求”全部学精通再开始”。AI Infra 的正确学习姿势是带着目标学基础——先把上面的内容过一遍达到”够用”，然后尽快进入 CUDA 编程、分布式训练等核心领域，在实践中遇到不会的再回来补。

🎯 自我检验清单

用以下清单检验自己的基础是否达标。每一条都应该能在不查资料的情况下完成：

  • 能独立写一个带装饰器、多进程的 Python 数据预处理脚本，并用 cProfile 定位性能瓶颈
  • 能读懂一段 CUDA host 端代码，说清 cudaMalloc、cudaMemcpy、kernel<<<...>>>()、cudaFree 分别在做什么
  • 能在 Linux 服务器上通过 SSH 登录、用 tmux 管理长时间运行的训练任务、用 conda 创建隔离环境
  • 能看到 (B, S, H) x (H, V) 立刻说出结果形状 (B, S, V)，并解释为什么张量并行是沿 H 或 V 维度切分的
  • 能用白话解释 Softmax 在做什么（把分数变成概率），以及交叉熵损失为什么在正确答案概率低时特别大
  • 能解释反向传播的链式法则，以及为什么 FP16 训练需要 Loss Scaling（防止梯度下溢）
  • 能写一个简单的 bash 脚本批量启动多组超参数实验，并用 nvidia-smi 监控 GPU 资源
  • 能用 git 完成基本的版本管理操作：创建分支、提交代码、处理合并冲突

📚 参考资料