G. 显存不足¶

分数：200 分

RuntimeError: CUDA error: out of memory

背景描述¶

众所周知，LLM 模型非常大，大到显存放不下，但是为了提高模型的能力，我们还是希望能用有限的显存里跑尽可能大的模型，。我们可以选择将模型参数卸载到 CPU 内存中，按需加载到显存中进行计算。但同时我们还想提高 GPU 的利用率，毕竟看着计算单元摸鱼总是让人不爽的。

现在我们从大模型训推中最简单的部分开始，求输入序列的 next token或者说loss，这个环节在推理中称为 prefill，在训练中称为前向传播。

任务描述¶

在本题中，你需要使用一张显存为 48GB 的 AMD W7800 GPU 运行 BF16精度的Qwen3-32B模型，求一组输入序列的 loss 。Qwen3-32B 的模型参数大小约为 64GB，同时还需要额外的显存放 KVCache 和激活值等数据，因此显存远远不够，需要将模型参数卸载到 CPU 内存中，按需加载到 GPU 显存中进行计算。

我们要求使用完整的 BF16 精度的模型进行计算，不允许使用量化、剪枝等技术。

handout中有两个文件，prefill.py和interact.py。你需要修改prefill.py，我们会使用interact.py来调用你的prefill.py，通过标准输入，以 json 数组字符串形式传入输入文本，期望获取 loss 结果。

interact.py会等待prefill.py输出ready 后，输入文本并开始计时，因此加载模型和预处理的时间不计入评分。

interact.py调用prefill.py 的命令行参数如下：

python3 prefill.py --model_path <模型路径>

<模型路径>是模型文件所在的目录，目录下包含了模型的config.json、 safetensors 权重、 tokenizer 等文件。模型是从https://huggingface.co/Qwen/Qwen3-32B 下载的。

评测环境¶

评测容器镜像：cr.hpc.lcpu.dev/hpcgame/rocm-prefill:latest，基础发行版为Rocky Linux 10.1，软件列表如下。进入容器后，你可以使用source /root/prefill-handout/.venv/bin/activate 激活 Python 虚拟环境。注意，构建镜像用的 handout.tar.gz 是旧版，请使用平台附件的新版 handout 。

flash_attn-2.8.3：基于上游 flash_attn编译，通过设置环境变量FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE 启用。
uv （以下依赖都是通过 uv 安装的）
AMD ROCm 7.1
Python 3.13
Pytorch 2.10.0
Transformers 4.57.6

如何做题¶

容器里已经有了基本的 uv 环境，很多包不需要重复安装
在容器中，只保留 /root/prefill-handout中的.venv 目录，删除其他文件，然后将平台附件中的 handout 解压到文件夹下
在 /root/prefill-handout目录下运行uv sync，同步环境
修改 prefill.py，实现按需加载模型参数到 GPU 进行计算
权重和数据是以 PVC 的形式提供。使用 hpcgame cli时，可以通过-v problem-prefill-data挂载。默认会被挂载到/mnt/problem-prefill-data目录下，模型和数据分别在该目录下的models/Qwen3-32B和datasets/NuminaMath-CoT 目录中。

输入输出格式¶

输入和交互格式¶

程序可以用一定的时间进行预处理，该部分用时不影响评分。程序准备好后需要输出一行 ready 到标准输出，表示可以接受输入并开始计时。

从标准输入读取一行包含了所有输入的文本的 json 数组字符串。

输出格式¶

按照输入文本数组顺序，输出对应的 loss 值，每个 loss 占一行。

评分方法¶

输出的 loss 结果与参考实现的误差需要通过 BFloat16 精度的数值比较，要求相对误差不超过 5%，或者在数值接近 0 时，允许绝对误差不超过 1e-2，伪代码：

def is_close(a: bfloat16, b: bfloat16) -> bool:
    if abs(b) > 1e-3:
        return abs(a - b) / abs(b) <= 0.05
    else:
        return abs(a - b) <= 1e-2

记程序的运行时间（从输出 ready 到输出 loss 结果）为 T 。满分时间 75s，超时时间 130 秒。

评分公式如下：

$S(T) = \begin{cases} 100 & T \leq 75 \\[6pt] 100 \cdot \exp\!\left(-\left(\dfrac{T - 75}{25}\right)^{\!2}\right) & T > 75 \end{cases}$

其中：

T：程序运行时间（秒），从输出 ready 到输出 loss 结果
S(T)：最终得分（满分 100）
75s：满分时间，T ≤ 75s 时得满分
τ = 25：衰减系数，控制超时后分数下降的速度

得分参考¶

T（秒）	得分
≤ 75	100
80	96.1
85	85.2
90	67.0
100	36.0
110	14.1
120	2.0
130	≈ 0

提示¶

Pytorch 提供了 cuda stream 和 event 相关的 API，建议阅读以下 Pytorch 文档：

附件¶

handout-prefill