计图适配华为昇腾910系列，迈向国产AI新生态
华为昇腾910是一款专为数据中心设计的高性能AI处理器，主要用于加速机器学习和深度学习任务。华为的CANN(Compute Architecture for Neural Networks)则是昇腾硬件与AI软件生态之间的重要桥梁，通过不断的升级优化，支持越来越丰富的AI场景和框架。2024 年10月，CANN8.0的推出标志着昇腾平台实现了更高的性能和兼容性。在此背景下，计图团队积极投身技术探索与创新实践，通过精心且深入地将计图框架与 CANN8.0 接口进行深度融合与优化，成功实现了计图框架在昇腾 910系列芯片上的适配。
适配完成后，计图框架的用户无需对计图代码进行任何额外的修改，就能实现跨硬件（英伟达 → 华为）平台的无缝迁移。AI开发者能够更加专注于算法的研发和优化工作，而无需关注底层硬件的不同。
现阶段，随着ChatGPT的成功，大语言模型(Large language model，LLM)不仅在学术研究中占据重要地位，同时也广泛应用于各行各业。为了进一步完善“昇腾+计图"在前沿应用场景中的表现，计图团队对大语言模型也做了深度适配，力求为科研人员提供了便捷的国产解决方案，也为大模型在实际应用场景中的落地提供了支持。
在推理优化方面，计图框架特别支持了大模型必备的FlashAttention高效算子。FlashAttention通过减少显存冗余访问和优化计算路径，使得在处理超长文本序列时，推理时间和显存消耗都得到了减少。这不仅提升了推理效率，也有效缓解了硬件资源瓶颈问题。现阶段，计图在昇腾上已经支持计图大模型推理库JittorLLMs中的LLaMA 7B、ChatGLM和ChatRWKV模型，并与PyTorch在昇腾平台上推理速度持平。
图2 计图在昇腾平台LLaMA推理示例
同时，计图也支持昇腾平台上大语言模型的微调。在训练过程中，计图框架充分利用其创新的统一计算图特性，能够根据不同的训练任务自动优化计算图，动态调整计算策略，以适应大语言模型在训练阶段的复杂需求。避免了传统深度学习框架中的静态图限制，使得在昇腾硬件上进行大语言模型的微调变得更加高效和灵活。此举极大地提升了训练过程的适应性，开发者可以根据不同的模型和任务需求快速调整训练配置，而无需担心硬件的底层调优。同时，计图通过对显存管理和运算调度的精细控制，确保了在大规模模型微调过程中，训练过程更加稳定，进一步优化了资源利用，提升了整体性能。
图3 计图在昇腾平台LLaMA 7B训练示例
此外，计图团队还特别重视模型的精度验证工作。在大语言模型的适配过程中，通过大量单元测试和自动化工具对算子精度进行了逐一比对，确保框架的数值精度与国外主流框架保持一致，保障了模型训练和推理结果的准确性。通过对大语言模型训练和推理的支持，计图不仅为国内AI硬件的广泛应用奠定了更加坚实的基础，为大模型的应用落地提供了有力支持。
计图框架的适配不仅限于大语言模型，还涵盖了多种模型任务类型，为华为昇腾生态注入了更多活力。在模型库迁移方面，计图已全面支持多种经典任务模型。例如，在图像分类任务中，计图支持ResNet、ViT等模型；在图像分割任务中，支持FCN、PSPNet等模型；在图像生成领域，Stable Diffusion等生成模型也可以高效运行在计图框架中。
图4 Stable diffusion运行实例，prompt为"a cute cat."
图5 基于华为910开展计图AI算法大赛
除此之外，计图在华为硬件上还支持了Jtorch，可以便捷地将大部分 pytorch 模型直接通过计图底层运行到华为硬件上。这种多样化的模型支持，使得计图在科研和工业应用中的适用范围进一步扩大。
未来，计图团队计划与华为CANN团队进一步合作，推出更多针对华为硬件优化的功能模块。我们将深度挖掘硬件潜力，通过对框架与硬件之间交互机制的优化，让硬件资源能够更加灵活、高效地被调用和利用。同时，在计算速度方面进行全方位的优化，采用先进的算法和技术手段，减少计算过程中的冗余环节，提高数据处理效率，为用户带来更加流畅、高效的使用体验。
最后，要特别感谢华为昇腾团队的领导和技术专家对计图的昇腾适配提供指导、技术支持与帮助。感谢天津市人工智能计算中心和并行科技提供的昇腾910系列芯片资源的支持。
参考文献
Shi-Min Hu, Dun Liang, Guo-Ye Yang, Guo-Wei Yang, Wen-Yang Zhou, Jittor: a novel deep learning framework with meta-operators and unified graph execution, Science China Information Science，2020, Vol. 63, No. 12, 222103.
GGC往期回顾
1. 张量场引导的机器人分层路径规划与自主探索 | CVMJ Spotlight
可通过下方二维码，关注清华大学图形学实验室，了解图形学、深度学习、Jittor框架、CVMJ期刊和CVM会议的相关资讯。
