IREE项目多GPU编译配置问题分析与解决方案

问题背景

在IREE编译器项目中，当用户尝试编译一个针对8个GPU设备(Llama 70B模型，TP8配置，2048输入长度)的MLIR文件时，遇到了设备声明缺失的错误。错误信息显示编译器无法找到已声明的设备#hal.device.promise<@__device_7>，这表明在多GPU环境下的设备配置存在问题。

错误分析

从错误日志中可以清楚地看到，编译器在执行流操作时无法解析设备引用。具体表现为：

1. 编译器在处理stream.cmd.execute操作时，发现引用了未声明的设备承诺(device promise)
2. 错误指向了特定的设备标识符@__device_7
3. 调用栈显示问题发生在处理分页注意力KV缓存聚集操作时

根本原因

经过技术分析，问题的根本原因在于编译命令中缺少对多GPU设备的正确配置。原始编译命令仅指定了目标设备为HIP(ROCm)平台，但未明确配置8个独立的GPU设备。

解决方案

正确的做法是在编译命令中显式指定每个GPU设备。修正后的编译命令应包含以下关键参数：

--iree-hal-target-device="hip[0]" \
--iree-hal-target-device="hip[1]" \
--iree-hal-target-device="hip[2]" \
--iree-hal-target-device="hip[3]" \
--iree-hal-target-device="hip[4]" \
--iree-hal-target-device="hip[5]" \
--iree-hal-target-device="hip[6]" \
--iree-hal-target-device="hip[7]" \

这些参数明确告诉IREE编译器系统中有8个可用的HIP设备，索引从0到7。这样编译器就能正确地为每个设备生成相应的代码和设备承诺声明。

完整编译命令示例

iree-compile \
    input.mlir \
    --iree-hip-target=gfx942 \
    -o=output.vmfb \
    --iree-hal-target-device="hip[0]" \
    --iree-hal-target-device="hip[1]" \
    --iree-hal-target-device="hip[2]" \
    --iree-hal-target-device="hip[3]" \
    --iree-hal-target-device="hip[4]" \
    --iree-hal-target-device="hip[5]" \
    --iree-hal-target-device="hip[6]" \
    --iree-hal-target-device="hip[7]" \
    --iree-opt-level=O3 \
    --iree-hal-indirect-command-buffers=true \
    --iree-stream-resource-memory-model=discrete \
    --iree-hal-memoization=true \
    --iree-dispatch-creation-propagate-collapse-across-expands=true \
    --iree-codegen-enable-default-tuning-specs=true \
    --iree-stream-affinity-solver-max-iterations=1024

技术要点

1. 设备声明：IREE编译器需要明确知道系统中可用的设备数量和类型，以便正确分配计算任务和资源。

2. HIP设备索引：在ROCm平台上，设备通过索引号标识，从0开始连续编号。

3. 多GPU支持：IREE通过显式设备声明支持多GPU并行计算，这对于大模型推理至关重要。

4. 设备承诺：IREE使用设备承诺机制来管理跨设备的资源分配和任务调度。

总结

在IREE项目中处理多GPU编译任务时，正确配置目标设备是至关重要的。对于像Llama 70B这样的大型模型，特别是在TP8(张量并行度为8)配置下，必须确保编译命令中包含了所有参与计算的GPU设备声明。这个案例展示了IREE编译器在多设备环境下的配置要求，为开发者处理类似问题提供了参考。