一种用于训练AI模型的系统、方法及电子设备

本申请实施例提供一种用于训练AI模型的系统、方法以及电子设备，系统包括：至少一个训练节点，包括：片上网络、连接片上网络的多个计算模块；片上网络有多个计算模块内存连接点，用于连接多个计算模块使用的多个计算模块内存，多个计算模块内存包含具有集合通信卸载功能的计算模块内存；训练节点至少用于执行训练AI模型的张量并行任务，由训练节点内的至少两个计算模块并行计算，各个计算模块的张量并行任务数据传输到计算模块使用的具有集合通信卸载功能的计算模块内存，在具有集合通信卸载功能的计算模块内存完成各个计算模块的张量并行任务的集合通信卸载操作后，将集合通信卸载操作结果返回给各个计算模块。本申请能够提高模型训练效率。