JAX+STAX模型GPU显存占用超出预期的原因及排查方案

1. GPU实际显存占用高于预期的原因、追踪方式与规避方案 #

超额显存的去向 #

观测到的显存差值来自三类JAX profiler默认不统计的开销，不属于无意义的内存泄漏：

• JIT计算图的临时缓冲区开销：@jit包装的训练步在首次运行时，XLA会为正向传播、反向梯度计算、优化器更新全流程预分配所有需要的临时内存，包括卷积层反向传播必须保留的中间激活值、梯度张量、cuDNN/cuBLAS算子运行所需的工作空间。之前统计的5.5GB只是参数、输入数据、优化器状态这类持久存储的占用，没有覆盖临时部分：以batch_size=100、输入200x200分辨率、两层64通道卷积的配置计算，仅第一层卷积的正向输出张量单批就占约1GB显存，反向传播需要保留所有层的中间激活，再加Adam优化器为每个参数存储的一阶、二阶动量张量，总峰值显存达到14.6GB完全符合计算逻辑。
• XLA运行时的显存申请机制：即使设置XLA_PYTHON_CLIENT_PREALLOCATE=false，XLA也只会关闭默认启动时占满90%显存的行为，仍会按照计算图的峰值显存需求一次性向CUDA驱动申请对应大小的显存块，申请后不会随单步运行结束主动归还给系统，因此外部工具检测到的显存值是XLA申请的峰值显存，不是当前时刻的实际使用值。
• 编译配置与运行配置不匹配：初始化参数时传入的input_shapebatch维度为100，XLA会默认按照这个batch大小做内存规划，哪怕后续跑更小的batch，首次编译时申请的显存也不会释放。

显存追踪方法 #

• 不要依赖纯Python层面的显存检测工具，这类工具只能拿到CUDA驱动上报的总申请值，无法区分显存的分配主体。
• 要拆分显存占用构成，优先用NVIDIA Nsight Systems启动训练脚本，勾选CUDA追踪、内存分配追踪、CPU采样三个选项，可以直接抓取到CUDA上下文基础占用、XLA持久数组占用、算子临时工作空间占用、中间激活占用的全部分类数据，不受JAX默认设备设置影响。
• JAX自带的profiler默认只统计JAX托管的持久设备数组，不会统计临时缓冲区和底层算子库的开销，这部分数值和实际总占用有差距是正常现象。

超额显存规避方案 #

• 计算显存预算时不要只统计参数和输入数据，必须把卷积中间激活、梯度张量、优化器动量项算入总需求：卷积层的激活显存和batch_size、输入分辨率、通道数线性正相关，batch_size从10升到100时激活显存直接翻10倍，这是batch_size=100时首个step就OOM的核心原因。可以用jax.checkpoint包裹卷积块，牺牲20%左右的计算速度，换中间激活显存占用下降50%以上。
• 不要通过把默认设备设为CPU的方式控制显存占用，跨设备搬运数据会打断XLA的内存复用逻辑，反而会抬高总显存需求。正确配置是默认设备设为GPU，搭配XLA_PYTHON_CLIENT_PREALLOCATE=false和TF_FORCE_GPU_ALLOW_GROWTH=true两个环境变量，让XLA按需逐步申请显存。
• 触发JIT编译时传入和实际训练一致的batch大小，避免XLA按照更大的batch维度预留多余显存。

2. 同时采集CPU、GPU两侧Profiling数据的方法 #

JAX profiler默认和启动时指定的默认后端绑定，默认设备设为CPU时不会加载GPU侧的追踪钩子，自然拿不到GPU数据，不需要单独维护两版代码，按以下方式配置即可同时采集两侧数据：

• 移除代码中jax.config.update('jax_platform_name', 'cpu')的设置，将默认设备改回GPU，需要把数据留在CPU时直接调用jax.device_put(数组, jax.devices('cpu')[0])手动指定存放位置即可。
• 启动训练前调用jax.profiler.start_trace(trace_log_dir, trace_cpu=True, trace_gpu=True)开启追踪，跑3-5个训练步后调用jax.profiler.stop_trace()结束采集，用TensorBoard打开生成的日志目录，就能同时看到CPU侧主机内存、GPU侧显存的时序变化，以及每个算子的内存占用、耗时数据。
• 如果需要更细粒度的底层内存数据，直接用Nsight Systems启动脚本即可，不需要修改JAX侧的任何配置，就能拿到从CPU内存到GPU显存的全链路分配记录。

内容的提问来源于stack exchange，提问作者Suraja