深度强化学习在芯片设计和大语言模型中的应用

这个研究课题来自ANNA GOLDIE，我进行了总结，目前最前沿的深度深度强化学习在芯片设计和大语言模型中的应用技术。（2024年）

ANNA GOLDIE, a Senior Staff Research Scientist at Google DeepMind, where she works on Gemini and Bard.

我将这部分内容整理成书籍，目前先放出核心的总结，以供大家学习。并给出了参考文献，可以进行阅读，我正在编写的书籍内容也来自这些文献。

深度强化学习在芯片设计和大语言模型中的应用

第一章 绪论

1.1 机器学习在系统和硬件设计中的重要性

在过去的十年中,系统和硬件的发展极大地推动了机器学习的进步。现在,是时候让机器学习反哺系统和硬件设计了。Thompson等人在2020年的研究表明,自2012年以来,人工智能训练中使用的计算量每3.4个月就会增加一倍。然而,摩尔定律已经不再适用,计算资源的供给远跟不上需求的增长。因此,我们亟需找到新的方法来弥合这一鸿沟。

Kaplan等人在2020年的研究表明,计算能力是推动人工智能进步的关键。许多系统优化问题本质上都是图上的组合优化问题。例如:

• 编译器优化:输入是程序的中间表示(IR)或高级操作(HLO)图,目标是调度或融合操作。
• 芯片布局:输入是芯片组件的网表(一种图),目标是将它们放置在二维或三维网格上。
• 数据中心资源分配:输入是作业负载图,目标是将它们放置在数据中心的机架和服务器中。

与传统的启发式算法或解析方法相比,基于机器学习的方法有以下优势:

1. 机器学习模型可以学习优化目标与输入特征之间的内在关系,从而更好地权衡不同的优化目标。
2. 机器学习模型可以从求解大量问题实例的经验中不断学习和进步,其专业知识可能远超人类专家。
3. 机器学习模型可以很容易地在分布式平台上扩展,参数规模可以任意增大。

1.2 本书的主要内容

本书主要介绍了两个将深度强化学习应用于系统优化的案例研究:

1. 用于模型并行的设备放置优化
2. 用于芯片布局设计的平面规划优化

此外,本书还介绍了深度强化学习在大语言模型中的重要应用,包括:

1. 从人类反馈中学习
2. 从人工智能反馈中学习

通过这些内容,读者可以深入理解深度强化学习在实际系统优化和人工智能应用中的巨大潜力。同时,本书也分享了将研究成果落地应用的宝贵经验。

第二章 利用深度强化学习优化设备放置

2.1 设备放置问题简介

随着机器学习模型规模的不断增长,当前最先进的模型已经拥有数十亿甚至上万亿的参数。然而,硬件内存的增长速度远低于模型的增长速度。为了训练这些大型模型,我们需要将模型切分到多个设备上。而模型切分的方式对训练性能有很大的影响。

2.2 将设备放置问题建模为强化学习问题

我们可以将设备放置问题建模为一个强化学习问题:

• 状态:神经网络操作图,可用设备的列表(如CPU核心,GPU卡等)。
• 动作:决定将每个操作放置到哪个设备上。
• 奖励:模型在硬件上实际运行的耗时。

策略网络接受状态作为输入,输出一个概率分布,表示将每个操作放置到每个设备上的概率。我们可以通过在真实硬件上运行,用实际运行时间作为奖励信号来训练策略网络。

2.3 分层放置模型

在早期的工作中,我们提出了一个端到端的分层放置模型:

1. 对操作进行聚类,形成一个高层的操作组图。
2. 使用一个强化学习模型来对操作组逐个放置。
3. 使用另一个强化学习模型在组内对单个操作逐个放置。

通过分层处理,该模型可以处理非常大的神经网络计算图。实验表明,该模型产生的放置方案在硬件利用率和通信开销等方面都优于人工设计的启发式方法。

2.4 泛化性能和图神经网络

端到端分层放置模型需要为每个新的计算图从头训练,这限制了它的实用性。为了提高泛化性能,我们在后续工作中引入了图神经网络来学习设备放置策略:

1. 使用图神经网络来学习计算图的节点embedding。
2. 使用transformer编码器来对embedding序列建模。
3. 使用自回归解码器逐个产生放置决策。

实验表明,基于图神经网络的模型可以在新的计算图上产生高质量的放置方案,泛化性能大大优于之前的方法。这为将设备放置技术集成到机器学习编译器和框架中铺平了道路。

第三章 利用深度强化学习优化芯片布局设计

3.1 芯片布局设计的重要性和挑战

芯片是人工智能系统的核心,而芯片的布局设计在很大程度上决定了芯片的性能、功耗和成本。当前的芯片设计流程通常需要1.5到2年的时间,这远跟不上人工智能算法的迭代速度。如果我们能将芯片设计周期缩短到几周,就可以极大地提高人工智能系统的性能。此外,如果能够同时优化神经网络架构和芯片布局,我们就有可能发现全新的高效架构。

3.2 芯片布局问题简介

芯片布局可以建模为一个图上的组合优化问题:

• 输入:芯片组件网表,包括宏块、标准单元、它们之间的连线关系等。
• 输出:将每个组件放置到一个二维网格上的位置。
• 优化目标:最小化计算延迟、功耗、面积等,同时满足布线拥塞度等约束条件。

芯片布局问题的解空间非常巨大,即使是简化版的问题,其状态空间也可达到10^9000。传统的布局方法主要有三类:

1. 基于划分的方法,如Fiduccia-Mattheyses算法。
2. 基于随机搜索或爬山法的方法,如模拟退火算法。
3. 基于解析优化的方法,如RePlAce方法。

这些方法在特定的问题实例上可以产生不错的结果,但很难扩展到新的问题实例,而且求解速度也难以满足当前芯片设计的需求。

3.3 基于强化学习的芯片布局方法

我们提出了第四类芯片布局方法:基于学习的布局方法。具体来说,我们将芯片布局建模为一个马尔可夫决策过程:

• 状态:当前的布局状态,包括已经放置的组件位置,和当前待放置的组件。
• 动作:选择一个网格位置,将当前组件放置到这个位置上。
• 奖励:布局完成后,根据线长、拥塞度等指标计算的布局质量。

策略网络接受状态特征,输出在每个网格位置上放置当前组件的概率。我们使用近端策略优化(PPO)算法来训练策略网络,通过试错搜索来最大化累积奖励。在搜索过程中,我们使用基于物理模拟的布局结果作为奖励信号来指导训练。我们的方法在TPU芯片的布局设计中取得了超越人类专家的结果,并被应用于最新一代的TPU芯片中。

3.4 提高芯片布局策略的泛化性能

与设备放置问题类似,如果我们为每个芯片网表训练一个全新的策略网络,其实用性会大大降低。理想情况下,布局系统应该能够在新的芯片网表上快速产生高质量的布局,而无需从头开始训练。为了提高策略网络的泛化性能,我们采取了以下方法:

1. 使用图神经网络来学习芯片网表的节点embedding。
2. 在一个芯片布局数据集上预训练奖励预测模型,使其能够准确评估布局质量。
3. 将预训练的奖励模型集成到策略网络和价值网络中,提高其泛化能力。

实验表明,在一个芯片布局数据集上预训练的策略网络,可以在新的芯片网表上快速产生接近最优的布局,并且随着数据集规模的增加,zero-shot性能还会进一步提升。这为芯片设计流程的自动化和智能化铺平了道路。

第四章 基于人类反馈的强化学习及其在大语言模型中的应用

4.1 从人类偏好中学习

在很多应用领域中,直接给出回报函数是很困难的。Christiano等人在2017年提出了一种解决方案:让人类来对模型的行为进行打分排序,然后训练一个奖励模型来拟合人类的偏好,最后用这个奖励模型作为强化学习的目标函数。通过这种方法,即使环境交互的样本很少,我们也可以训练出符合人类偏好的策略。

4.2 人类反馈在对话系统中的应用

基于人类反馈的强化学习(RLHF)已经成为训练大规模语言模型的重要手段,被广泛应用于GPT-3、ChatGPT、Claude等系统中。RLHF的一般流程如下:

1. 收集人类对语言模型输出的偏好数据。
2. 训练一个奖励模型来拟合人类偏好。
3. 用奖励模型作为目标函数,对预训练的语言模型进行RL微调。

实验表明,经过RLHF训练的语言模型在指令遵循、上下文学习等方面有显著的提升,同时生成的内容也更加安全、有帮助。此外,在代码生成任务上,RLHF也可以提高语言模型的编程能力。在MMLU、Lambada等常见的语言理解基准上,RLHF训练的模型也取得了更好的表现。

第五章 基于人工智能反馈的强化学习及其在大语言模型中的应用

5.1 从人工智能反馈中学习的动机

随着语言模型接近甚至超越人类水平,用人类反馈来指导训练变得越来越困难。为了进一步提升语言模型的性能,Bai等人在2022年提出了一种新颖的方法:基于人工智能反馈的强化学习(RLAIF)。

5.2 RLAIF的一般流程

RLAIF的核心思想是,用一个"宪法"来描述语言模型应该遵循的原则,然后基于这个宪法用模型自身来评估和修正模型的输出,形成一个自我指导的训练过程:

1. 从原始语言模型中采样对话。
2. 用一个评判模型来评估对话的质量,选出更好的回复。
3. 训练一个奖励模型来拟合评判模型的偏好。
4. 用奖励模型微调原始模型,重复上述过程。

相比RLHF,RLAIF不需要人类参与数据标注和偏好评判的过程。在实际任务中,RLAIF训练的模型可以媲美甚至超越RLHF训练的模型,展现了巨大的潜力。

5.3 监督学习和强化学习的互补作用

实验表明,将监督学习和强化学习结合起来,可以集两者的优势于一身:

• 监督学习使用人工标注的高质量数据,可以有效引导模型的初始行为。
• 强化学习通过试错搜索,可以在更大的对话空间中进一步提炼模型的策略。

在实践中,通常采用两阶段的训练流程:先用监督学习初始化模型,然后再用强化学习进行微调。联合训练可以大大提高模型生成安全、合理、有帮助内容的能力。

第六章 思考与展望

通过上述案例研究,我们看到了机器学习,特别是深度强化学习在系统优化和人工智能应用中的巨大潜力。随着机器学习模型和优化算法的不断进步,我们有望看到更多令人兴奋的研究成果和应用案例。以下是一些值得关注和思考的方向:

1. 进一步提高学习算法的泛化能力和样本效率,使其能够更快地适应新的问题实例。
2. 开发更加通用和模块化的problem formulation,使得强化当然,我会继续补充和丰富这本书的内容。以下是一些额外的章节和细节:

第六章 思考与展望

3. 开发更加通用和模块化的问题表示方法,使得强化学习算法可以更容易地应用于不同的组合优化问题。图神经网络在这方面展现了巨大的潜力。

4. 探索机器学习和其他优化方法(如数学规划、元启发式算法等)的结合,利用不同方法的互补优势,提高求解的效率和质量。

5. 加强机器学习系统的可解释性和可控性,使其能够满足实际应用中的各种约束和需求。对偶学习和逻辑推理等技术可能会有所帮助。

6. 在更多的实际系统和应用中验证机器学习优化的价值,积累经验教训,并建立相应的基准测试和评估体系。

7. 探索机器学习在硬件设计自动化、编译器优化、操作系统资源管理等传统系统问题上的应用,推动系统软硬件与人工智能的进一步融合。

6.1 机器学习驱动的软硬件协同设计

传统的硬件和软件设计通常是解耦的:硬件设计者根据一些抽象的计算模型来设计通用的硬件架构,而软件开发者则根据具体的硬件接口和性能模型来优化软件实现。这种解耦的模式在一定程度上限制了软硬件的效率。

借助机器学习,特别是端到端的联合优化技术,我们有望打破软硬件之间的抽象界限,实现软硬件的协同设计和优化。这里给出一个联合优化神经网络加速器和计算图的示例:

class HardwareEncoder(nn.Module):
    """硬件架构的图神经网络编码器"""
    
    def __init__(self, node_dim, edge_dim, hidden_dim):
        super().__init__()
        self.node_encoder = nn.Linear(node_dim, hidden_dim)
        self.edge_encoder = nn.Linear(edge_dim, hidden_dim)
        self.conv1 = HeteroConv({"node": self.node_encoder, 
                                 "edge": self.edge_encoder},
                                aggr="mean")
        self.conv2 = HeteroConv({"node": nn.Linear(hidden_dim, hidden_dim), 
                                 "edge": nn.Linear(hidden_dim, hidden_dim)},
                                aggr="mean")
                                
    def forward(self, hardware_graph):
        hardware_graph = self.conv1(hardware_graph)
        hardware_graph = self.conv2(hardware_graph)
        return hardware_graph
    
class SoftwareEncoder(nn.Module):
    """计算图的图神经网络编码器"""
    
    def __init__(self, node_dim, edge_dim, hidden_dim):
        super().__init__()
        self.node_encoder = nn.Linear(node_dim, hidden_dim)
        self.edge_encoder = nn.Linear(edge_dim, hidden_dim)
        self.conv1 = HeteroConv({"node": self.node_encoder, 
                                 "edge": self.edge_encoder},
                                aggr="mean")
        self.conv2 = HeteroConv({"node": nn.Linear(hidden_dim, hidden_dim), 
                                 "edge": nn.Linear(hidden_dim, hidden_dim)},
                                aggr="mean")
                                
    def forward(self, software_graph):
        software_graph = self.conv1(software_graph)
        software_graph = self.conv2(software_graph)
        return software_graph
    
class JointModel(nn.Module):
    """联合优化硬件和计算图的端到端模型"""
    
    def __init__(self, hardware_encoder, software_encoder, hidden_dim):
        super().__init__()
        self.hardware_encoder = hardware_encoder
        self.software_encoder = software_encoder
        self.predictor = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
        
    def forward(self, hardware_graph, software_graph):
        hardware_embed = self.hardware_encoder(hardware_graph)
        software_embed = self.software_encoder(software_graph)
        joint_embed = torch.cat([hardware_embed.mean(dim=0), 
                                 software_embed.mean(dim=0)], dim=-1)
        out = self.predictor(joint_embed)
        return out

在这个示例中,我们使用图神经网络分别对硬件架构图和计算图进行编码,然后将二者的embedding拼接起来,用一个预测器来预测整体的性能指标。通过端到端的梯度下降,我们可以同时优化硬件架构的超参数(如存储空间、带宽等)和计算图的结构(如算子融合、数据布局等),从而得到一个软硬协同优化的设计。

以上只是一个简单的示例,实践中的软硬协同优化可能涉及更加复杂的设计空间和优化目标,比如能效比、可靠性、可编程性等。此外,我们还可以考虑制造工艺、封装方式等底层因素,以及编译器、指令集等上层抽象,将优化问题进一步扩展为全栈综合优化。

6.2 机器学习驱动的系统安全

传统的系统安全主要依赖于人工设计的防御规则和机制,如访问控制、异常检测、冗余容错等。这些方法在应对已知的攻击模式时往往能够提供较好的保护,但对于未知的攻击方式和变种,其有效性则大打折扣。近年来,对抗机器学习逐渐成为系统安全领域的新焦点。一方面,我们可以利用机器学习来自动化攻防对抗的流程,提高系统的安全性和鲁棒性;另一方面,我们也需要考虑机器学习模型本身的脆弱性和可能带来的新风险。

在这里,我用一个对抗样本检测的例子来说明机器学习在系统安全中的应用:

class FeatureSqueezing(nn.Module):
    """特征压缩处理"""
    
    def __init__(self, bit_depth):
        super().__init__()
        self.bit_depth = bit_depth
        
    def forward(self, x):
        x_int = torch.round(x * (2**self.bit_depth - 1))
        x_float = x_int / (2**self.bit_depth - 1)
        return x_float
    
class AdversarialDetector(nn.Module):
    """对抗样本检测器"""
    
    def __init__(self, base_model, squeeze_model):
        super().__init__()
        self.base_model = base_model
        self.squeeze_model = squeeze_model
        
    def forward(self, x):
        y = self.base_model(x)
        y_squeeze = self.base_model(self.squeeze_model(x))
        diff = torch.abs(y - y_squeeze)
        return diff.mean()
    
detector = AdversarialDetector(resnet50(), FeatureSqueezing(4))

这个例子展示了一种简单的对抗样本检测方法:将输入图像经过特征压缩(如降低颜色位深)后,比较处理前后预测结果的差异。直觉上,自然图像的分类结果对这种变换应该是稳定的,而对抗样本则可能发生较大改变。我们可以设置一个阈值,当差异超过阈值时,就判定样本为对抗样本,进而采取相应的防御措施。

类似地,我们还可以设计一些其他的输入变换,如平移、旋转、噪声等,构建一个"对抗样本检测器集成"。通过联合多个检测器的判断,我们可以进一步提高对抗样本检测的精度和召回率。此外,我们还可以考虑自适应攻防的场景,即攻击者会针对现有的防御手段产生新的对抗样本。为了应对这种情况,我们需要引入在线学习和强化学习,通过不断地生成对抗样本并调整防御策略,来实现攻防双方的博弈和进化。

第七章 结语

纵观全书,我们探讨了机器学习,特别是深度强化学习在芯片设计、资源调度、大语言模型等领域的创新应用。从这些案例中,我们看到了机器学习给传统系统和软件带来的新的可能性。同时,我们也认识到,要将这些可能性转化为实际的价值,还需要克服不少困难和挑战。

展望未来,我认为机器学习与系统、软件的结合将会不断深化,催生出更多令人兴奋的研究成果和应用案例。一些值得关注的方向包括:

1. 机器学习驱动的软硬件协同设计:打破传统的软硬件分层抽象,实现端到端的性能和效率优化。

2. 机器学习驱动的系统安全:利用机器学习来自动化攻防对抗流程,提高系统鲁棒性,同时保证机器学习自身的安全性和隐私性。

3. 机器学习驱动的系统自动化:利用机器学习来自动优化系统配置、故障诊断、资源管理等任务,提高系统的自治水平。

4. 机器学习驱动的软件工程:利用机器学习来辅助需求分析、程序合成、测试生成、性能调优等软件开发任务,提高开发效率和质量。

5. 机器学习驱动的人机协同:利用机器学习来构建更加自然、高效、个性化的人机交互系统,让机器成为人的得力助手和伙伴。

当然,上面只是一些初步的想法,未来可能有更多创新的方向和机会等待我们去探索。我相信,只要我们秉持开放、好奇、务实的态度,勇于尝试、善于总结、敢于创新,一定能够推动机器学习与系统、软件的进一步融合,让计算机系统变得更加智能、高效、安全、易用,更好地服务于人类社会的进步。

让我们共同期待这个激动人心的未来!

参考文献

1. Thompson, N.C., Greenewald, K., Lee, K. and Manso, G.F., 2020. The computational limits of deep learning. arXiv preprint arXiv:2007.05558.

2. Kaplan, J., McCandlish, S., Henighan, T., Brown, T.B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J. and Amodei, D., 2020. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361.

3. Mirhoseini, A., Goldie, A., Yazgan, M., Jiang, J., Songhori, E., Wang, S., Lee, Y.J., Johnson, E., Pathak, O., Bae, S. and Laudon, J., 2020. Chip placement with deep reinforcement learning. arXiv preprint arXiv:2004.10746.

4. Mirhoseini, A., Goldie, A., Yazgan, M., Jiang, J., Songhori, E., Wang, S., Lee, Y.J., Johnson, E., Pathak, O., Nazi, A. and Pak, J., 2021. A graph placement methodology for fast chip design. Nature, 594(7862), pp.207-212.

5. Addanki, R., Bojja Venkatakrishnan, S., Gupta, S., Mao, H. and Alizadeh, M., 2019. Placeto: Learning generalizable device placement algorithms for distributed machine learning. arXiv preprint arXiv:1906.08879.

6. Paliwal, A., Gimeno, F., Nair, V., Li, Y., Lubin, M., Kohli, P. and Vinyals, O., 2019. Reinforced genetic algorithm learning for optimizing computation graphs. arXiv preprint arXiv:1905.02494.

7. Christiano, P., Leike, J., Brown, T.B., Martic, M., Legg, S. and Amodei, D., 2017. Deep reinforcement learning from human preferences. arXiv preprint arXiv:1706.03741.

8. Parth Sarthi, Salman Abdullah, Aditi Tuli, Shubh Khanna, Anna Goldie, Christopher D. Manning. RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval. ICLR 2024.

9. Ganguli et al. The Capacity for Moral Self-Correction in Large Language Models. arXiv 2023.

10. Bai et al. Constitutional AI: Harmlessness from AI Feedback. arXiv 2022.

11. Bowman et al. Measuring Progress on Scalable Oversight for Large Language Models. arXiv 2022.

12. Azalia Mirhoseini*, Anna Goldie* (co-first author), Mustafa Yazgan, Joe Jiang, Ebrahim Songhori, Shen Wang, Young-Joon Lee, Eric Johnson, Omkar Pathak, Azade Nazi, Jiwoo Pak, Andy Tong, Kavya Srinivasa, William Hang, Emre Tuncer, Quoc V. Le, James Laudon, Richard Ho, Roger Carpenter, Jeff Dean. A Graph Placement Methodology for Fast Chip Design (Open Access). Nature.

13. Anna Goldie and Azalia Mirhoseini. “Placement Optimization with Deep Reinforcement Learning”. ISPD 2020.

14. Yanqi Zhou, Sudip Roy, Amirali Abdolrashidi, Daniel Wong, Peter C Ma, Qiumin Xu, Ming Zhong, Hanxiao Liu, Anna Goldie, Azalia Mirhoseini, James Laudon. “A Generalized Approach to End-to-End Computational Graph Optimization”. NeurIPS 2020 (Oral).

15. Azalia Mirhoseini*, Anna Goldie* (co-first author), Hieu Pham, Benoit Steiner, Quoc V. Le, Jeff Dean. “A Hierarchical Model for Device Placement,” ICLR 2018.

16. Denny Britz*, Anna Goldie* (co-first author), Thang Luong, Quoc V. Le. “Massive exploration of neural machine translation architectures,” EMNLP 2017.