超越英伟达GPU：苹果自研AI大模型技术突破
时间：2024-08-05 阅读：94 评论：443
过去几年中爆发的AI算力需求，英伟达GPU成为了最大的赢家，无论是自动驾驶还是AI大模型，比如特斯拉、小鹏、OpenAI、字节跳动等厂商都在争相抢购英伟达GPU产品，获得足够的算力构建自己的计算中心，用于训练不同应用的ai模型。AI领域众星捧月也造就了英伟达市值曾突破三万亿美元，并短暂登顶全球第一的位置。不过苹果却反其道而行之，在7月底苹果公司揭露的官方论文显示，苹果自研大模型AFM背后，完全没有英伟达GPU的影子。在苹果近期揭露的一篇技术论文中，详细介绍了苹果在端侧和服务器侧的大模型。在大模型预训练的阶段，苹果基础模型是在AXLearn框架上进行训练的，据苹果介绍，AXLearn框架是苹果公司在2023年发布的开源项目，这个框架建立在JAX和XLA的基础上，允许模型在各种硬件和云平台上进行高效和可扩展性训练，包括TPU以及云端和本地的GPU。苹果采用了数据并行、张量并行、序列并行和完全分片数据并行（FSDP）的组合来沿多个维度扩展训练，如数据规模、模型规模和序列长度。其中，苹果的AFM服务器端大模型是苹果规模最大的语言模型，该模型在8192个TPUv4芯片上进行了预训练，这些芯片被配置成8*1024个芯片集群，通过数据中心网络（DCN）连接。预训练有三个阶段，首先使用 6.3 万亿个Tokens开始，接着使用 1 万亿个Tokens继续训练，最后使用 1000 亿个Tokens进行上下文长度的扩展。而在AFM的端侧模型上，苹果对其进行了大幅修剪，论文披露AFM端侧模型是一个拥有30亿参数的模型，该模型是从64亿参数的服务器模型中蒸馏而来，而这个服务器模型则是在完整的6.3万亿参数量上进行训练。与服务器端模型不同的是，AFM端侧模型采用的芯片是谷歌的TPUv5，论文中的信息显示，AFM端侧模型是在一个由2048个TPUv5p芯片组成的集群上进行训练的。谷歌在去年12月发布了TPUv5p，面向云端AI加速，谷歌称之为“迄今为止最强大、可扩展性最强和最灵活的人工智能加速器”。TPUv5p在bfloat16精度下可以提供459 teraFLOPS（每秒可执行459万亿次浮点运算）算力；在Int8精度下，TPU v5p可以提供918 teraOPS（每秒可执行918万亿次整数运算）；支持95GB的HBM内存，带宽高达2.76 TB/s。相比上一代的TPU v4，TPUv5p实现了每秒浮点运算次数提高一倍，内存带宽是前代的三倍，训练大模型速度提升2.8倍，并且性价比是前代的2.1倍。除了苹果之外，目前采用谷歌TPU系列芯片进行大模型训练的还有谷歌自家的Gemini、PaLM，以及openai前副总裁创立的Anthropic所推出的Claude大模型，上个月Anthropic发布的Llama 3.1 405B还被认为是最强开源大模型。苹果、谷歌、Anthropic的实例，证明了TPU在大模型训练上的能力。但相比于英伟达，TPU目前在大模型领域的应用依旧只是冰山一角，背后更多的大模型公司，包括OpenAI、特斯拉、字节跳动等巨头，主力AI数据中心依然是普遍采用英伟达GPU。一直以来，围绕CUDA打造的软件生态，是英伟达在GPU领域最大的护城河，尤其是随着目前AI领域的发展加速，市场火爆，英伟达GPU+CUDA的开发生态则更加稳固，AMD、英特尔等厂商虽然在努力追赶，但目前还未能看到有威胁英伟达地位的可能。但市场的火爆难免吸引更多的玩家入局，对英伟达发起挑战，或者说是在AI广阔的市场空间中，希望分得一杯羹。首先是英伟达在GPU领域的最大对手AMD ，今年一月有研究人员在Frontier超算集群上，利用其中8%左右的GPU，训练出一个GPT 3.5级别规模的大模型。而Frontier超算集群是完全基于AMD硬件的，由37888个MI250X GPU和9472个Epyc 7A53 CPU组成，这次研究也突破了在AMD硬件上突破了先进分布式训练模型的难点，为AMD平台训练大模型验证了可行性。同时，CUDA生态也在逐步击破，今年7月英国公司Spectral Compute推出了可以为AMD GPU原生编译CUDA源代码的方案，大幅提高了AMD GPU对CUDA的兼容效率。英特尔的Gaudi 3 也在发布时直接对标英伟达H100，并宣称在模型训练速度和推理速度上分别比英伟达H100提高了40%和50%。除了芯片巨头外，也不乏来自初创公司的冲击。比如Groq推出的LPU、Cerebras推出的Wafer Scale Engine 3、Etched推出的Sohu等等。国内方面，有走多卡集群训练路线的初创公司，比如摩尔线程在今年6月宣布与羽人科技合作成功实现了摩尔线程夸娥（KUAE）千卡智算集群与羽人系列模型解决方案的训练兼容适配，高效完成了70亿参数羽人大语言模型YuRen-7b的训练测试。摩尔线程夸娥方案基于全功能MTT S4000 GPU，该GPU采用了第三代MUSA内核，单卡支持48GB显存容量和768GB/s的显存带宽，FP16算力为100TFLOPS。值得一提的是，MTT S4000计算卡借助摩尔线程自研开发工具，可以充分兼容现有CUDA软件生态，实现CUDA代码零成本迁移到MUSA平台。天数智芯也与智源研究院、爱特云翔合作，提供天垓100加速卡、构建算力集群及全程技术支持，实现基于自主通用GPU的大模型CodeGen（高效编码）项目，通过中文描述来生成可用的C、Java、Python代码以实现高效编码。另外值得一提的是，国内还有一家走TPU路线的AI芯片公司——中昊芯英。该公司在2023年底推出了国内首款量产TPU AI训练芯片“刹那”，据称在处理大模型训练和推理任务时相比英伟达A100，性能提高近150%，能耗下降30%，单位算力成本仅为A100的42%。当然，除了芯片公司，据现有的信息，目前主流的云服务供应商，比如前面提到的谷歌，还有亚马逊、微软、Meta、阿里巴巴、字节跳动、百度、华为等都有自研芯片的布局，其中还包括用于AI大模型训练的芯片。从长远来看，自研芯片是云服务供应商有效降低算力成本的方式之一，当AI大模型训练成为了云计算的重要用途时，自研AI训练芯片自然也是云服务厂商的长远之计。苹果作为消费电子巨头已经迈出了重要的一步，即摆脱对英伟达的算力依赖，并且还有大量的挑战者正在蠢蠢欲动。星星之火，可以燎原，英伟达在AI训练领域的地位，可能没有表面看起来那么稳固。文章来源微信公众号： 电子发烧友
