被显存墙卡住的深夜：MI300X 上的 LLaMA-7B 训练实录

做过大模型训练的朋友都懂，最搞心态的不是模型不收敛，而是代码跑着跑着突然崩出一句 CUDA out of memory（在 AMD 环境下则是 HIP out of memory）。上周我在调试 LLaMA-7B 的全量微调任务时，就结结实实撞上了这堵墙。手里这块 AMD Instinct MI300X 明明有着 192GB 的恐怖显存，理论上能轻松吞下更大的模型，但在实际训练流程中，显存占用曲线却像坐过山车一样直冲红线，导致 OOM 报错。

这次踩坑经历让我意识到，硬件的大显存只是基础，如何配合 ROCm 软件栈进行精细化的内存管理，才是释放 MI300X 性能的关键。今天就把这次从“爆显存”到“丝滑运行”的排查和优化过程复盘一下，希望能给正在折腾 ROCm 的算法工程师们一点参考。

定位元凶：ROCm 工具链下的显存泄漏追踪

起初我以为是模型层数设多了，但检查配置文件后发现参数完全在理论范围内。这时候盲目改代码无异于大海捞针，必须得靠数据说话。在 NVIDIA 生态里大家习惯用 nvidia-smi 或者 torch.cuda.memory_summary()，而在 ROCm 环境下，rocprof 和 rocm-smi 组合拳才是神器。

我先在训练脚本启动时挂上了 rocprof 进行内核级分析：

rocprof --stats -o trace_output.csv python train_llama.py

同时，在另一个终端实时监控显存状态：

watch -n 1 rocm-smi --showmemuse --showpower

监控日志里有个细节非常可疑：随着 iteration 推进，显存占用呈阶梯式上升，但在 optimizer step 之后并没有完全回落。通过查看 trace_output.csv 中的内存分配事件，我发现大量的临时缓冲区（Temporary Buffers）在每次反向传播时被分配，却未被及时释放。进一步检查代码，发现是自定义的数据加载器中，某些中间张量意外地保留了 grad 属性，导致计算图无法完整释放。

修复这个逻辑漏洞后，基线显存占用下降了约 15%，但距离稳定运行还有差距。峰值显存依然会在处理长序列样本时瞬间击穿 190GB 阈值。

极限压测：梯度检查点与 Batch Size 的博弈

解决了泄漏问题，剩下的就是硬碰硬的显存峰值优化。LLaMA-7B 这种体量的模型，全量微调时的激活值（Activations）占用极其惊人。MI300X 虽然显存大，但也经不住无节制的消耗。这里我主要用了两招：梯度检查点（Gradient Checkpointing）和动态 Batch Size 调整。

开启梯度检查点

这是以时间换空间的经典策略。在 PyTorch (ROCm 版) 中，开启它非常简单，只需在模型包装时加上 gradient_checkpointing_enable()。这会迫使模型在反向传播时重新计算部分前向过程的激活值，从而大幅减少显存驻留。

from transformers import AutoModelForCausalLM, TrainingArguments

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
# 关键一步：开启梯度检查点
model.gradient_checkpointing_enable()

training_args = TrainingArguments(
    output_dir="./llama-7b-ft",
    per_device_train_batch_size=4,  # 初始设定
    gradient_accumulation_steps=8,
    fp16=False,  # MI300X 建议尝试 bf16
    bf16=True,
    # ...其他参数
)

开启后，显存峰值立刻出现了断崖式下跌。原本需要存储整个序列的激活值，现在只需要存储部分检查点，显存占用减少了近 40%。

寻找最佳 Batch Size

显存省下来了，接下来就是填满它。MI300X 的优势在于大显存允许我们使用更大的 Micro-Batch Size，从而提高 GPU 利用率。我写了一个简单的脚本，逐步增加 per_device_train_batch_size，观察 rocm-smi 的反馈。

最终我将单卡 Batch Size 锁定在 12，配合 8 步梯度累积，等效 Batch Size 达到 96。这个配置下，显存利用率维持在 95% 左右，既没有浪费 MI300X 的宝贵资源，又留出了少量余量应对突发波动。

优化成果与推荐配置

经过这一轮“手术”，训练任务终于能稳定跑起来了。对比优化前后的显存曲线，变化非常明显：优化前曲线呈锯齿状且不断上移直至 OOM；优化后曲线平稳，每个 step 结束都能干净利落地回落到基线。

(注：此处为示意描述，实际博客可插入 rocprof 生成的图表)

以下是我在 MI300X 上运行 LLaMA-7B 全量微调的最终推荐配置表，大家可以按需取用：

这次实战让我深刻体会到，AMD Instinct MI300X 确实是大模型训练的利器，尤其是其 192GB 显存给了我们要命的“容错空间”。但要想真正跑满性能，离不开对 ROCm 工具链的熟练掌握和对内存机制的精细调优。如果你也在从 CUDA 迁移到 ROCm，别被初期的环境配置劝退，一旦跑通第一个 stable run，那种掌控巨大算力的感觉是非常爽的。