janestreet/gtc2026

# 五个瓶颈，五个修复：我们如何避免浪费 GPU 性能 本仓库包含了我们在 [GTC 演讲](https://www.nvidia.com/gtc/session-catalog/sessions/gtc26-s82065/) 中用于获取计时和分析数据的所有代码。 ## 训练优化 你可以直接查看 train*.py 文件来查看原始代码，但查看 train*.diff 文件可以更直观地了解我们具体需要修改哪些内容才能启用特定的优化。 训练中应用的完整优化序列 - train_00_bad.py：无需任何特殊操作即可获得的训练效果 - train_01_baseline.py - 启用 torch.compile - 启用 tensor cores - train_02_vectorized.py：启用响应者指标向量化 - train_03_cuda.py - 对响应者指标进行手动 kernel fusion - 详情请参阅 sum_many.* - train_04_stream.py：使用 CUDA streams 重叠 dataloading 与计算 - train_05_cutile.py - 使用感知 padding 的 GEMM kernel 来减少 matmul 计算 - 具体实现请参阅 padded_matmul.py ## 推理优化 同样地，我们提供了原始推理文件和 .diff 文件，以便查看 CUDA graph 优化。 - inference_00_basic.py：基础解码推理 - inference_01_cuda_graph.py：应用 CUDA graphs ## 复现步骤 你需要一台满足以下条件的系统： 1. 一块 NVIDIA GPU 2. 在你的 `PATH` 中有一个正常工作的 [CUDA toolkit](https://developer.nvidia.com/cuda-downloads) 安装 3. 一个能够编译 torch 扩展的 C 工具链（例如，你需要 `python3-dev` 包） 在运行脚本之前，你还需要安装我们的 Python 依赖项。 ``` # 创建并激活虚拟环境 python -m venv .venv source .venv/bin/activate # 安装支持 CUDA 的 PyTorch（调整 cu130 以匹配您的 CUDA 版本）。 # 有关选项，请参见 https://pytorch.org/get-started/locally/。 pip install torch --index-url https://download.pytorch.org/whl/cu130 # 安装剩余的依赖项 pip install -r requirements.txt ``` 安装完成后，你可以直接运行训练脚本，例如： ``` python train_01_baseline.py ``` 为了在 RAM 较小的 GPU 上运行，你可能需要使用更小的模型或更小的 batch size。考虑传入更小的 `--num-blocks` 或更小的 `--max-num-tokens` 来降低 RAM 需求。例如，以下命令应该可以在 A100 或 L40S 上运行： ``` python train_01_baseline.py --max-num-tokens 1024 ```

标签：CUDA, PyTorch, Vectored Exception Handling, 人工智能, 凭据扫描, 性能优化, 检测绕过, 模型训练, 用户模式Hook绕过, 逆向工具