第1章：AI Infra 工程师的编程语言基础

AI Infra 不是只会写 Python 调包就能干的工作。从 PyTorch 模型脚本到 CUDA Kernel,从分布式启动脚本到性能 profiling,这套技术栈横跨三种语言生态——Python 描述算法、C/C++ 操控硬件、Linux 串起一切。本文系统梳理 AI Infra 工程师必须掌握的编程语言基础,以及每一项的”够用门槛”。

📑 目录

• 1. 三种语言的角色分工
• 2. Python:AI 生态的通用语言
• 3. C/C++:CUDA 的宿主语言
• 4. Linux:开发与部署的舞台
• 5. 工具链:Git / Conda / VSCode
• 自我检验清单
• 参考资料

---

1. 三种语言的角色分工

把 AI Infra 工程师的工作想象成一座三层楼——楼上、楼下、和把它们串起来的电梯井。

层	语言	角色	典型代码
上层	Python	描述模型 / 调度任务	PyTorch 训练脚本、配置加载、推理服务
下层	C/C++	操控硬件 / 极致性能	CUDA Kernel、自定义算子、推理引擎
串联	Linux + Bash	部署 / 调试 / 自动化	启动脚本、容器、CI、tmux、git

🌟 一个常见误区:很多新人以为”AI Infra = Python 高手 + 一点点 CUDA”。实际情况是:真正影响产出效率的瓶颈往往不是 Python 写得多快,而是当训练在凌晨 3 点挂掉时,你能否在 5 分钟内 SSH 上服务器、定位 GPU 死锁、检查 NCCL 日志、重启任务。Linux 和 C++ 的”够用”才是隐形分水岭。

---

2. Python:AI 生态的通用语言

2.1 必须熟练的核心特性

面向对象 (OOP)

PyTorch 的 nn.Module、Dataset、Optimizer 全部基于 OOP。能流畅写出继承、多态、super() 调用是基本盘。

import torch.nn as nn

class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads):
        super().__init__()
        self.attn = nn.MultiheadAttention(dim, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(dim, 4 * dim),
            nn.GELU(),
            nn.Linear(4 * dim, dim),
        )
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x):
        # Pre-Norm 结构,残差连接
        h = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
        return h + self.ffn(self.norm2(h))

装饰器 (Decorator)

PyTorch 的 @torch.no_grad()、@torch.jit.script、HuggingFace 的 @dataclass,以及性能 profiling 中常用的 timing decorator——不会写装饰器就读不懂这些库的源码。

import time
from functools import wraps

def timeit(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        t0 = time.perf_counter()
        result = fn(*args, **kwargs)
        elapsed = (time.perf_counter() - t0) * 1000
        print(f"[{fn.__name__}] elapsed: {elapsed:.2f} ms")
        return result
    return wrapper

@timeit
def heavy_compute(x):
    return x @ x.T  # 矩阵乘法

生成器与迭代器

DataLoader 内部全靠生成器避免一次性把整个数据集加载到内存。理解 yield 是看懂 streaming dataloader 和大文件流式处理的前提。

def stream_jsonl(path):
    with open(path) as f:
        for line in f:
            yield json.loads(line)

# 内存只占一行,即使文件 100GB 也不会 OOM
for sample in stream_jsonl("massive_dataset.jsonl"):
    process(sample)

多进程与多线程

由于 GIL 的存在,Python 的多线程对 CPU 密集任务无效,但对 I/O 密集(数据加载、网络请求)依然有用。多进程才是 CPU 密集型并行的正确姿势——torch.utils.data.DataLoader 的 num_workers 参数底层就是 multiprocessing。

场景	推荐方案
数据预处理(图像 decode / tokenize)	multiprocessing 或 DataLoader(num_workers>0)
网络请求并发(下载、API 调用)	asyncio 或 concurrent.futures.ThreadPoolExecutor
纯 CPU 数值计算	NumPy / PyTorch (内部 C++ 释放 GIL)
跨进程共享 GPU 张量	torch.multiprocessing (注意 share_memory_())

性能 profiling 三件套

# 1. cProfile:函数级耗时
import cProfile
cProfile.run('train_one_epoch()', 'profile.out')

# 2. line_profiler:行级耗时(需 pip install line_profiler)
@profile
def slow_function():
    ...

# 3. memory_profiler:内存追踪
from memory_profiler import profile

2.2 常见反模式

❌ 隐式数据类型转换导致性能崩塌

# 反例:循环中频繁在 list 和 tensor 间转换
losses = []
for batch in loader:
    loss = model(batch)
    losses.append(loss.item())  # GPU→CPU 同步,数千次累计很慢

# 正例:在 GPU 上累加,最后一次性同步
total_loss = torch.zeros(1, device='cuda')
for batch in loader:
    total_loss += model(batch)
final = total_loss.item()

❌ 在主进程做重 CPU 任务阻塞 GPU

数据预处理放在 __getitem__ 里 + num_workers > 0,让 GPU 永远不等数据。

2.3 推荐资料

类型	资料	说明
官方	Python 官方教程	从语法到标准库
视频	Real Python 系列	高质量进阶教程
书籍	《Fluent Python》(流畅的 Python)	进阶必读,深入装饰器/迭代器/元类
博客	Yongbo Wang:CPython GIL 详解	理解多线程为什么没用

---

3. C/C++:CUDA 的宿主语言

CUDA Kernel 本身写在 .cu 文件里(语法接近 C++),但调用 Kernel、管理显存、和 Python 互通,都需要 C/C++ 基本功。

3.1 你必须看得懂的特性

指针与内存管理

float* h_data = (float*)malloc(n * sizeof(float));     // 主机内存
float* d_data = nullptr;
cudaMalloc(&d_data, n * sizeof(float));                 // 设备内存
cudaMemcpy(d_data, h_data, n * sizeof(float),
           cudaMemcpyHostToDevice);                     // 数据搬运
// ... kernel launch ...
cudaMemcpy(h_data, d_data, n * sizeof(float),
           cudaMemcpyDeviceToHost);
cudaFree(d_data);
free(h_data);

记住一个核心区分:h_ 前缀 = host (CPU) 内存,d_ 前缀 = device (GPU) 内存。两者地址空间隔离,不能直接互访。

模板基础(读懂即可,不用精通)

PyTorch C++ 扩展和 CUDA Kernel 大量使用模板做类型分发:

// AT_DISPATCH_FLOATING_TYPES 是模板宏,自动生成 float / double 两份 kernel
AT_DISPATCH_FLOATING_TYPES(input.scalar_type(), "my_kernel", ([&] {
    my_kernel<scalar_t><<<grid, block>>>(
        input.data_ptr<scalar_t>(),
        output.data_ptr<scalar_t>(),
        n);
}));

RAII 与智能指针

C++11 之后强烈推荐 std::unique_ptr / std::shared_ptr 替代裸 new/delete,避免内存泄漏。

std::unique_ptr<float[]> buffer(new float[n]);
// 离开作用域自动释放,无需手动 delete

3.2 编译链接最小知识

# CUDA 编译
nvcc -O3 -arch=sm_80 my_kernel.cu -o my_kernel

# C++ 与 CUDA 混编
nvcc -O3 -arch=sm_80 main.cpp my_kernel.cu -o app

# 链接 cuBLAS / cuDNN
nvcc main.cu -lcublas -lcudnn -o app

# 查看寄存器和共享内存使用(性能调优时必看)
nvcc -Xptxas -v my_kernel.cu

工具	作用
nvcc	NVIDIA 的 CUDA 编译器(底层调用 g++/clang)
cmake	跨平台构建工具,大型项目首选
make	经典构建工具,简单项目够用
ldd	查看二进制依赖的动态库
nm	查看 .so / .a 中的符号

3.3 PyTorch C++ 扩展最小示例

很多自定义算子的工程模式是:Python 写训练逻辑,C++/CUDA 写性能关键算子,通过 pybind11 互通。

// my_op.cpp
#include <torch/extension.h>

torch::Tensor my_add(torch::Tensor a, torch::Tensor b) {
    return a + b;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("my_add", &my_add, "Add two tensors");
}

# setup.py
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

setup(name='my_op',
      ext_modules=[CppExtension('my_op', ['my_op.cpp'])],
      cmdclass={'build_ext': BuildExtension})

3.4 推荐资料

类型	资料	说明
教程	LearnCpp.com	从 0 到 C++ 现代特性,免费
书籍	《Effective Modern C++》	C++11/14 最佳实践
官方	PyTorch C++ Extension Tutorial	写自定义算子的入门
视频	CppCon 历年 talk	最前沿的 C++ 工程实践

---

4. Linux:开发与部署的舞台

GPU 服务器几乎清一色跑 Linux。能在 Linux 上”游刃有余”是 AI Infra 工程师的隐形门槛。

4.1 必备命令清单

文件与目录

ls -la /home/user/projects     # 详细列表
find . -name "*.cu" -type f    # 按名字找文件
du -sh */                      # 查看每个子目录的占用
df -h                          # 查看磁盘剩余
tree -L 2                      # 树状结构(需安装)

进程与资源

top / htop                     # 实时进程监控
ps aux | grep python           # 查找 Python 进程
kill -9 <pid>                  # 强制杀进程
nvidia-smi                     # GPU 状态
nvidia-smi -l 1                # 每秒刷新一次
nvidia-smi pmon -i 0           # 查看 GPU 0 上每个进程的占用

网络与服务

ss -tlnp                       # 查看监听端口(替代 netstat)
ssh -i ~/.ssh/id_rsa -p 22 user@host
scp -r local/dir user@host:/remote/path
rsync -avz --exclude='.git' src/ user@host:dst/

文本处理(神器三连)

grep -rn "FlashAttention" .              # 递归搜索
sed -i 's/old/new/g' file.py             # 行内替换
awk -F',' '{print $2}' data.csv          # 按列处理

4.2 Shell 脚本必会模式

循环批量提交任务

#!/bin/bash
for lr in 1e-4 3e-4 1e-3; do
    for bs in 32 64 128; do
        echo "Launching lr=$lr bs=$bs"
        python train.py --lr $lr --batch_size $bs \
                        --output_dir runs/lr${lr}_bs${bs} \
                        > logs/lr${lr}_bs${bs}.log 2>&1 &
    done
done
wait  # 等所有后台任务完成

变量与条件判断

NODE=${NODE:-0}                # 默认值
MASTER_ADDR=$(hostname -I | awk '{print $1}')

if [[ -z "$CUDA_VISIBLE_DEVICES" ]]; then
    export CUDA_VISIBLE_DEVICES=0,1,2,3
fi

if nvidia-smi > /dev/null 2>&1; then
    echo "GPU OK"
else
    echo "No GPU detected"; exit 1
fi

4.3 会话管理:tmux

tmux 是远程开发的命脉——SSH 断开后任务不死,关键命令:

操作	快捷键 / 命令
新建会话	tmux new -s train
列出会话	tmux ls
重连	tmux a -t train
分离(后台保留)	Ctrl+b 然后按 d
创建窗口	Ctrl+b c
切换窗口	Ctrl+b 0/1/2...
上下分屏	Ctrl+b "
左右分屏	Ctrl+b %

4.4 系统监控与排查

# CPU 与内存
free -h                        # 内存使用
vmstat 1                       # 内存/进程统计实时刷新

# 磁盘 I/O
iostat -x 1                    # 磁盘 I/O 统计
iotop                          # 看哪个进程在读写

# 网络流量
iftop -i eth0                  # 实时网络带宽
ss -s                          # socket 统计

# 系统日志
dmesg | tail -50               # 内核日志,GPU 掉卡或 OOM 都在这
journalctl -u nvidia-persistenced  # systemd 服务日志

🍎 实战:训练突然变慢怎么排查

1. nvidia-smi:GPU 利用率低?显存满了?温度过高降频?
2. htop:CPU 是否打满(数据预处理瓶颈)?内存是否快爆?
3. iostat:磁盘 I/O 是否成瓶颈(读 dataset 太慢)?
4. iftop:多机训练时 IB 带宽是否打满?
5. dmesg | tail:有没有 OOM Kill 或 ECC 错误?

---

5. 工具链:Git / Conda / VSCode

5.1 Git:版本管理

场景	命令
拉代码	git clone <url>
创建分支	git checkout -b feat/new-kernel
提交	git add -A && git commit -m "feat: ..."
推送	git push -u origin feat/new-kernel
同步主干	git fetch && git rebase origin/main
二分查 bug	git bisect start && git bisect bad && git bisect good <commit>
紧急保存	git stash / git stash pop

5.2 Conda / venv:Python 环境管理

# Conda(适合多项目隔离)
conda create -n vllm python=3.10 -y
conda activate vllm
pip install vllm

# uv(更快的 pip 替代)
pip install uv
uv venv .venv && source .venv/bin/activate
uv pip install torch transformers

环境管理的金科玉律:每个项目独立环境,绝不在 base 里装东西。否则不同项目的 PyTorch 版本会互相打架。

5.3 VSCode Remote-SSH:本地编辑 + 远程运行

AI Infra 的标准工作流:本地 Mac/笔记本写代码,远程 GPU 服务器跑训练。

1. VSCode 装 Remote - SSH 扩展
2. ~/.ssh/config 配置好别名:

Host gpu1
    HostName 192.168.1.100
    Port 22
    User user
    IdentityFile ~/.ssh/id_rsa

3. 命令面板 Remote-SSH: Connect to Host → 选 gpu1 → 直接打开远程目录,所有保存自动同步,断点调试也能用。

---

✅ 自我检验清单

• Python OOP 实战:不看资料能写出一个继承 nn.Module 的 Transformer Block,正确实现 __init__ 和 forward
• 装饰器手写:能现场写一个 @timeit 装饰器,统计任意函数耗时
• 生成器辨析:解释为什么 DataLoader 用生成器而不是 list,并能改写一个 streaming dataloader
• 多进程数据加载:能解释 num_workers > 0 时为什么训练加速,以及什么情况下会死锁
• C++ 读写:能读懂一个 CUDA host 端代码(malloc / memcpy / kernel launch / free)并解释每一步
• nvcc 命令:不查资料能编译一个简单 .cu 文件,并知道 -arch=sm_80 是什么意思
• Linux 排查:训练突然变慢,能在 5 分钟内用 nvidia-smi/htop/iostat/dmesg 缩小到瓶颈
• tmux 熟练度:能在 tmux 中开 3 个窗口分别跑训练 / 监控 GPU / 看日志,并知道如何在 SSH 断开后重连
• Git 高阶:能用 git bisect 在 100 个 commit 中定位引入回归的那个,并能熟练 rebase 解冲突
• Shell 脚本:能写一个 bash 脚本批量跑 9 组超参组合,自动生成日志目录,任务结束后聚合结果

---

📚 参考资料

Python

• 官方教程:https://docs.python.org/3/tutorial/
• Real Python:https://realpython.com/
• 《Fluent Python》第二版:Luciano Ramalho

C++ / CUDA

• LearnCpp:https://www.learncpp.com/
• PyTorch C++ Extension:https://pytorch.org/tutorials/advanced/cpp_extension.html
• CUDA C++ Programming Guide:https://docs.nvidia.com/cuda/cuda-c-programming-guide/

Linux

• The Missing Semester(MIT):https://missing.csail.mit.edu/ —— 教科书没教但每个工程师都该会的工具
• The Linux Command Line(William Shotts):免费 PDF
• tmux 官方 Wiki:https://github.com/tmux/tmux/wiki

工具链

• Pro Git:https://git-scm.com/book
• Conda 文档:https://docs.conda.io/
• VSCode Remote-SSH:https://code.visualstudio.com/docs/remote/ssh