KTransformers项目中KExpertsCPU模块的FP16精度问题分析与解决方案

问题背景

在KTransformers项目实现DeepSeek-V2-Lite-Chat模型的推理过程中，发现了一个重要的精度问题：当使用A100-80G 8卡环境运行模型时，在BF16精度下模型能够正常输出结果，但在FP16或FP32精度下却会出现输出为空的情况。

问题定位

经过技术团队深入分析，发现问题出在KExpertsCPU模块的实现上。具体表现为：

1. 当模型第一层不使用非共享专家时，各精度下的计算结果基本一致
2. 但在使用非共享专家的情况下，FP16与BF16的输入虽然相近，输出却存在显著差异
3. 即使在FP16模式下使用BF16的输入数据，输出结果仍然与BF16模式下的结果差异很大

根本原因

技术团队发现问题的核心在于KExpertsCPU模块中数据类型处理的硬编码问题：

1. 当前实现中，GPU向CPU传输数据时强制使用了BF16数据类型
2. 在ktransformers/operators/experts.py文件的KExpertsCPU.load函数中，MOEConfig的最后一个参数被硬编码为30(对应BF16)
3. KExpertsCPU.output_cpu的数据类型设置也存在不一致问题

解决方案

针对这一问题，技术团队提出了明确的修复方案：

1. 修改MOEConfig的最后一个参数，从30(bf16)调整为1(fp16)
2. 同步调整KExpertsCPU.output_cpu的数据类型设置
3. 确保GPU到CPU的数据传输能够正确识别和保持原始精度设置

技术影响

这一问题的修复将带来以下改进：

1. 使模型能够正确支持FP16精度模式下的推理
2. 提高模型在不同精度设置下的结果一致性
3. 增强框架对不同硬件配置的适应性

总结

KTransformers项目中的这一精度问题展示了混合精度计算中数据类型一致性的重要性。通过精确控制各计算环节的数据类型转换，可以确保模型在不同精度设置下都能获得稳定可靠的推理结果。技术团队已经将该修复纳入下一个版本更新计划，将为用户提供更稳定、更灵活的多精度支持。