【技术实现步骤摘要】
本专利技术涉及计算任务映射领域，尤其是涉及一种基于强化学习的晶圆级架构任务映射方法、装置及存储介质。
技术介绍
1、随着大语言模型获得突破性进展，引领了当前人工智能进入一个新的阶段。
2、大语言模型具有非常多的参数量，达到了上千亿的级别，如此的多参数，其训练过程需要消耗数千块高性能显卡数月之久。
3、为了满足大语言模型的训练，有研究提出了晶圆级架构，其提供了高速互联的带宽、极低通信延迟、大规模的算力。然而如何构建配套的软件栈，从而高效的使用这样的架构尚未有相关工作。其中关键的一个问题是计算任务的自动映射与调度的问题。
4、例如一些现有技术公开了针对晶圆级架构的ai任务自动映射与调度方法，将该问题抽象为一个优化问题，然后利用搜索优化的方法去求解，利用一般搜索优化问题得出可行的映射方案。
5、然而晶圆级架构有着超大的搜索空间，例如tesla dojo单个晶圆级硬件就有8750个节点。面对如此巨大的搜索空间，现有方案的搜索效率不高，在可接受时间范围内得到的方案硬件执行效率并没有很好的保障。
技术实现思路
1、本专利技术的目的就是为了提供一种基于强化学习的晶圆级架构任务映射方法、装置及存储介质，将晶圆级架构任务映射问题转化为强化学习问题，从而可以利用强化学习方法进行求解，可以更加快速得到更优的映射方案。
2、本专利技术的目的可以通过以下技术方案来实现：
3、一种基于强化学习的晶圆级架构任务映射方法，包括：
5、基于计算图的每个节点的邻接矩阵和每个节点的信息，提取每个节点的特征向量得到计算图的特征表示；
6、构建强化学习网络：将映射策略作为强化学习状态，将映射策略对应的代价作为强化学习奖励，将基于映射策略生成的中间状态向量作为强化学习状态，其中，所述映射策略表征计算图的各节点所分配的硬件资源；
7、训练构建的强化学习网络；
8、基于训练好的强化学习网络，输入需要进行任务映射的计算图，得到任务映射方案。
9、进一步的，所述方法还包括：
10、设置约束条件，其中，所述约束条件对强化学习网络输出的映射策略进行过滤。
11、更进一步的，所述约束条件基于硬件限制得到。
12、再更进一步的，所述硬件限制包括内存限制、拓扑逻辑限制和计算图依赖限制的一种或多种。
13、更进一步的，对于通过过滤的映射策略，代价为与计算任务的总耗时正相关的函数值；
14、对于未通过过滤的映射策略，代价被配置为无穷大的值。
15、进一步的，对于没有基于约束条件过滤的实施方案中，所述代价为与计算任务的总耗时正相关的函数值。
16、进一步的，所述映射策略的得到过程具体为：将中间状态向量和计算图的特征表示拼接后得到映射策略。
17、进一步的，所述基于训练好的强化学习网络，输入需要进行任务映射的计算图，得到任务映射方案，具体包括：
18、输入需要进行任务映射的计算图；
19、将需要进行任务映射的计算图作为训练集，对训练好的强化学习网络进行微调训练；
20、采用经过微调训练的强化学习网络得到任务映射方案。
21、另一方面，本专利技术提供一种基于强化学习的晶圆级架构任务映射装置，包括存储器、处理器，以及存储于所述存储器中的程序，所述处理器执行所述程序时实现如上述的方法。
22、再一方面，本专利技术提供一种存储介质，其上存储有程序，所述程序被执行时实现如上述的方法。
23、与现有技术相比，本专利技术具有以下有益效果：
24、1、将晶圆级架构任务映射问题转化为强化学习问题，从而可以利用强化学习方法进行求解，可以更加快速得到更优的映射方案。
25、2、通过设计约束条件，可以尽量保证得到的映射策略是可行的，从而能够得到充分的训练反馈，提高准确性和训练速度。
26、3、代价为与计算任务的总耗时正相关的函数值，可以针对计算任务的总耗时最小进行训练，从而减小最终映射方案的总耗时。
27、4、具有较强的通用性和扩展性，可广泛应用于不同类型的计算任务和计算资源。
本文档来自技高网...【技术保护点】
1.一种基于强化学习的晶圆级架构任务映射方法，其特征在于，包括：
2.根据权利要求1所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述方法还包括：
3.根据权利要求2所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述约束条件基于硬件限制得到。
4.根据权利要求3所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述硬件限制包括内存限制、拓扑逻辑限制和计算图依赖限制的一种或多种。
5.根据权利要求2所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，
6.根据权利要求1所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述代价为与计算任务的总耗时正相关的函数值。
7.根据权利要求1所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述映射策略的得到过程具体为：将中间状态向量和计算图的特征表示拼接后得到映射策略。
8.根据权利要求1所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述基于训练好的强化学习网络，输入需要进行任务
9.一种基于强化学习的晶圆级架构任务映射装置，包括存储器、处理器，以及存储于所述存储器中的程序，其特征在于，所述处理器执行所述程序时实现如权利要求1-8中任一所述的方法。
10.一种存储介质，其上存储有程序，其特征在于，所述程序被执行时实现如权利要求1-8中任一所述的方法。
...【技术特征摘要】
1.一种基于强化学习的晶圆级架构任务映射方法，其特征在于，包括：
2.根据权利要求1所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述方法还包括：
3.根据权利要求2所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述约束条件基于硬件限制得到。
4.根据权利要求3所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述硬件限制包括内存限制、拓扑逻辑限制和计算图依赖限制的一种或多种。
5.根据权利要求2所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，
6.根据权利要求1所述的一种基于强化学习的晶圆级架构任务映射方法，其特征在于，所述代价为与计算任务的总耗时正...
【专利技术属性】
技术研发人员：代旭，姜申飞，胡杨，韩慧明，王磊，郝培霖，
申请(专利权)人：上海人工智能创新中心，
类型：发明
国别省市：
还没有人留言评论。发表了对其他浏览者有用的留言会获得科技券。