少年何妨梦摘星                                     张靖义我很喜欢十六七岁的夏天。彼时的少年手持长弓，意气风发地眺望远方，端直了燕尾，搭上虎筋弦，弓弦一振，箭去不知所踪。十余年后，我走在路上，背后隐约传来簌簌风声，一回头，一支箭正中红心。我，一名标准的“95后”，成长在中国互联网和科技高速发展的十年里，说起来，算得上是第一代真正的数字原住民。当时十六七岁、正在读高中的我，耳边常常萦绕着“信息时代”“AI”这些新奇的词汇。它们激发了我的好奇与想象，我挣脱思维的缰绳，幻想着十年后的样子，我们的家里，会不会已经有了《钢铁侠》里的同款人工智能助手“贾维斯”？《007》主角们用VR视角远程操作车辆进行驾驶，是否已经成为了现实？《星球大战》的机器人在空气中投射出全息投影人像，能不能让我跨越时间空间的限制……我越想越兴奋，天马行空的丰富想象、翻涌而上的庞杂思绪，宛如矢箭般漫射开来。也许有一天，想象触及之处，科技真的能点石成金。我第一次模模糊糊地感知到科技的澎湃力量，既然要做梦，不如做大一点，少年何妨梦摘星！当年射出的矢箭，此刻“正中红心”上大学后，行动派的我果断选择了AI算法作为专业方向，一头扎进这个奇妙的世界。当初枯燥陌生的技术名词，如今成了我日日相见的熟人，当初艰深晦涩的领域，如今成了我甘之如饴的宝藏。我开始不断挑战自己，向更高的目标和更难的课题发起冲锋。通过不懈的努力，本科时我便以第一作者的身份，在ECCV（欧洲计算机视觉国际会议）发表了论文。后来随着多篇顶会论文陆续发表，我也逐渐在自己的细分领域“小有名气”。2021年6月，正当我硕士即将毕业，准备出国读博时，华为的HR找到了我，邀请我和昇腾计算产业的技术主管聊一下，结果这一聊可不得了，直接就聊出了炽热的火花，不仅如此，还顺利“聊”出了一场面试。面试官告诉我，在人工智能行业，全世界可能只有华为一家公司能够有芯片、硬件、软件、算法等AI全栈的能力。“昇腾将成为智能计算的底座，”他将目光投向我，“这不仅是我们突破封锁的机会，将来，我们还要为世界提供第二算力选择。你有兴趣吗？”那一刻，背后簌簌风声由远及近，我一怔，随即恍然。无需回头，多年前的那支箭，此刻正中红心。在迷雾中找到一丝光亮那时，大模型方兴未艾，作为国内最大AI算力集群的鹏城云脑II期才刚刚建成，华为“昇腾”这个新兴品牌，正式开始在AI全栈全场景上急速狂飙。入职华为后，我期待着即将展开的伟大事业——为昇腾构筑AI全栈竞争力，然而现实则比想象中的要骨感。要优化一个完整的AI系统，仅了解上层的算法知识是不够的，软件方面有框架、算子、操作系统等；硬件方面还有计算芯片、通信组网……这些底层的知识对于我而言，是陌生的领域和概念，而底层技术和系统优化，与我之前的专业背景更是有着天壤之别。我意识到，我的技能树几乎全部点在了AI算法上，与此相对的是，我对底层系统和优化的了解仅仅限于纸面上的浅薄知识，这让我压力倍增，一筹莫展。每天的工作仿佛陷入了一团迷雾之中，虽然我一直在努力，投入了大量的时间和精力研究各种资料和参考文献，但仍然没有取得突破。所有困难似乎都能将手无寸铁的我轻易击倒，有时是调试代码时出现的错误令我无所适从，有时是性能优化的挑战让我陷入绝望。我感到心力交瘁，满心的自我怀疑和否定不知该如何化解，为什么公司和主管没有发挥我的“AI算法”长处，而是让我去了解完全陌生的底层系统优化？我的主管春生哥和导师周老师察觉到了我的困惑，一天中午，他们喊我一起去会议室聊聊。见我进来了，周老师和春生哥热情地和我打招呼，春生哥笑着把我拉过去，鼓励道：“靖义，你可是我们当中最懂AI的，别妄自菲薄。”周老师换上了严肃的语气，语重心长地对我说：“计算进程中的转变，由应用来驱动架构的演进已成为主流，既懂算法应用又懂底层的人拥有独特的优势。所以靖义，你除了要发扬你的长板外，还要学习补齐你的短板，发挥你在上层算法应用的优势，去驱动我们整体软硬件架构的演进。业界既懂硬件又懂AI算法的全栈人才很少，公司对你抱有很高的期待，你要尽快成长起来。”这番话如醍醐灌顶，点亮了我前进的方向。我开始意识到，我原本的思维高度着实太过狭隘，公司已不再是象牙塔，要想达成我的终极目标，我需要成为一名全能的系统专家。不仅限于底层或硬件优化或是单纯的AI算法，我要通过站在更高的维度，在应用的角度来深入理解和优化整个系统，为华为开发出更高效、稳定的解决方案。小试牛刀，挑战超算界“诺贝尔奖”任总常说，“板凳坐得十年冷”，我理解并不是说在十年里等待机会，而是抓住这个时间充实提高自己，当机会来临的时候，才能把握住。我开始沉浸式“充电”，为日后厚积薄发储备能量。2021年到2022年，在负载建模项目中，我带领团队设计并实现了昇腾首个大模型负载仿真建模工具，负载建模让我对硬件和软件系统有了深刻的认识，并在这个过程中深刻地了解了昇腾作为一款DSA（领域专用加速器）在各种硬件参数上的权衡取舍，理解了各种不同的硬件形态、组网形态会如何影响上层业务，理解了大模型在计算、通信、内存和集群线性度上的系统瓶颈。有了充分积累之后，很快，机会就来了。我的第一个重点工作是作为技术专家参与并主导“鹏城神农大模型”的设计和优化。鹏城云脑2作为当时国内最大的AI超算集群，拥有1E FP16算力，急需一个重量级的大科学应用来证明自己。我们要在鹏城云脑2上挑战号称超算界的“诺贝尔奖”的戈登贝尔奖，利用全国产化算力构筑业界第一个千亿级蛋白质生成大模型，用以预测新冠病毒的变异和辅助抗新冠病毒肽类药物的设计。这个任务的挑战非常大，首先，我们在整体蛋白质模型设计上的经验几乎为零，也没有在昇腾+昇思上训练千亿MoE（混合专家模型）大模型的成功经验，这就像我们要设计一款汽车，却没有设计图纸可参考，在方案设计上属于摸着石头过河。其次就是模型的性能差，我们在打榜时的算力利用率不足9%，而当时业界sumit集群类似任务的算力利用率超过17%，这个差距不仅是数字的8%的差距，而是纯纯代际的差距，会导致我们的训练效率比别人落后近一倍！我们用着入门级汽车的发动机，去和友商轿跑的发动机比拼，性能优化的难度不言而喻。必须提升算力利用率，这是我们接下来的重点发力方向。接手这项工作后，我对大模型的瓶颈进行了全面分析。我让自己“蜗居”在办公室的安静角落里，一门心思专注地研究着芯片的性能指标和运行情况。纵览众多数据，我发现这其中，矩阵乘运算的效率尤其低下。矩阵乘运算是大模型的核心，大模型中99%的运算都是它。这个问题让我感到沮丧，因为时间如此紧迫，要在短时间内将matmul（矩阵乘）的性能提升到目标水平，实在不容易。我决定深入分析问题的方案和根本原因，开始仔细研究矩阵乘运算的性能瓶颈，我发现，其中一个主要原因是模型输入和权重的形状不亲和——这样形状的输入在NPU（神经网络处理器）上执行效率低。问题原因基本定位清楚，接下来就是需要思考如何优化解决这个“形状亲和匹配的问题”。经过对数据的反复思考和分析后，我的脑海中迸发出一个新的想法——试试一种自适应的“形状分桶算法”，通过动态调整输入数据使其更加亲和。这个方案就好像孩童在搭积木时，会先把形状相同的积木放一起，再想办法将异形拼接成形状相同的方块，然后就可以进行高效的搭建。当然，我们操作的“积木分拼”都是自动化实现的，这样会使得矩阵乘运算更加高效。此外，我还主导重新设计了模型主体部分的结构，在保证计算量、参数量和之前模型一致的前提下，尽量提升主体模型的亲和性。提高了负载应用的亲和性后，我们设计了一套针对大模型的瓶颈分析方法，采用自顶向下的思想，从集群、单机、芯片多层级多粒度地分析大模型训练中的瓶颈，让算子得以针对性优化，整个算子性能也得到了极大提高。稼先社区论坛现场，与公司专家讨论大模型MindSpore、海思和昇腾使能Lab同时投入了大量的时间和精力，一起对方案进行了改进和优化，甚至在中秋和国庆期间也将脑细胞全部贡献给了方案落地。我在鹏城云脑2用4096颗NPU训练了快一个月，我的导师和领导打趣道：“靖义，你们的实操训练可是烧掉了深圳黄金地段的一套房呢。”此时，我的目标只有一个：The sooner the better（越快越好）！我们的大模型对数据分布非常敏感，尤其是对于一些低频出现的高价值蛋白质序列，我们的模型并不能很好地学习。这个问题让我感到非常困惑和焦虑，因为这些低频序列对于我们的任务至关重要。我回想起自己在研究生阶段关于小样本学习的研究经验，“一条孤立的小样本在刚开始很难训练，那我只要检索一些和它类似的样本，就可以提高相似样本的训练数量，降低训练的难度！”我和诺亚的专家尹伊淳一起提出了一种基于检索和交叉注意力机制的方法，这个方法可以帮助预训练模型更好地学习小样本蛋白序列。经过连续两个月的艰苦攻关，我们终于取得了重要的突破，算力利用率从9%提高到了35%，并且“从0到1”地设计了神农大模型加速短肽发现全流程和新冠病毒变异预测全流程，通过基础模型预训练，下游微调与可控性质生成、序列结构预测、分子动力学模拟和生物学实验，我们发现了两条新的抗阴性菌肽和抗新冠肽，将短肽类药物的筛选发现耗时从原来的数月、数年，降低到几天。这意味着在药物研发过程中，医学工作者们可以更快地发现具有潜在治疗作用的短肽药物，加速新药的开发和上市。我们的这项成果，最终成功入围了戈登贝尔奖的“Final list（提名）”，并在华为全联接大会上发布。经历了这些磨炼，我不仅熟悉了各种性能优化工具和技术，积累了大模型性能优化的经验，同时也具备了解决复杂问题的能力，能沉稳应对各种性能瓶颈、资源限制的挑战了。强敌环绕， 这一战终于到来2023年，OpenAI发布了聊天型AI模型ChatGPT，这一消息引起了巨大的轰动。大模型的时代正式到来，各种基于ChatGPT的应用如雨后春笋般涌现。大模型的崛起引发了算力需求的爆发性增长，算力是大模型时代最宝贵的资源，“为世界提供第二算力选择”不再是一句口号，而是切实摆在眼前的战略机遇。时间就是生命，错过这个历史的节点可能就是前功尽弃！横在前面的竞争对手是已经打磨了十多年的N公司GPU（图形处理器），而我们昇腾自己的NPU却还是一个蹒跚学步的孩子。X1项目，是我们打开互联网市场的里程碑式项目，也是难度极高的项目。项目的客户一直以来都被公认是我们遇到的最“挑剔”的客户，他们对于产品与解决方案的要求非常严格。除了昇腾软件上的不成熟，我们的硬件芯片在某些指标上也落后于竞争对手，客户的诉求模型没有完全发挥出我们芯片的长处，这使得我们的开箱整体性能只有竞品的0.25倍，这是一个巨大的差距。面对这样的局面，我临危受命，被领导安排到X1项目组，负责整体的性能优化工作。我肩负着两项重要的任务：主线任务就是在NPU上把客户的大模型优化到持平竞品的性能；支线任务就是详细总结X1的性能优化经验，在未来的X2、X3场景中能做到批量复制。即便前方尽是强敌环绕，狭路相逢，我们也要敢于出鞘亮剑。这一战，终将会到来。接下来，我们仿佛进入了电影《碟中谍》的紧张氛围中，我对训练架构师和现场SA（解决方案架构师）开玩笑说，我们就是电影中“不可能任务小组”，这活儿可比电影里的更真实更刺激！未来说不定我们的每一次的优化都如同电影中的一次秘密行动，需要我们精密的计划和默契的配合。和客户沟通后，我们按照挑战目标一步一步行动，计划在2023年6月30日达成0.8倍竞品性能，9月30日达成1倍（持平）竞品的性能。我们从6月1日开始进行系统优化，在这期间，我们对系统进行了全面的性能分析和评估，确定了存在的瓶颈和优化的方向，也采取了一系列的措施来提高系统的性能和效率，但是事情的进展并不如想象中顺利，我们的算子耗时比竞品高了太多，非常影响阶段性目标的达成。6月28日晚上10点，办公室里弥漫着紧张的氛围。我们原本计划在这个时间点庆祝达成630的目标，但是算子的性能优化并没有达到预期，我们离目标还有一段距离。我缩在办公区隔间的最里面，绞尽脑汁思考着各种可能的优化方案。我盯着算子性能测试结果中占比高的viewcopy（视图拷贝）操作，这是我们芯片上由于非连续转连续引入的额外操作，它在整个计算中占据了超过10%的比例，也是最有可能优化的节点。“能不能去掉这个操作呢？” 我自言自语。我开始仔细研究产生viewcopy操作的根因，寻找可以快速完成的优化方案。经过一番思索和实验，我发现是客户在循环中对一个计算结果向量做了非常规的按照奇偶索引操作，这个操作对昇腾NPU非常不亲和。我们的NPU对大块的连续内存读写是亲和的，但是对小块的不连续跳读跳写是不亲和的，就好比用了一个超大的勺子，但是每次只能舀的水只有一点点，远远不能装满勺子，搬运效率非常低。找到了根因，我稍微松了一口气，拿出了纸笔开始写写画画起来。我将这段代码的逻辑转换成公式写在了纸上，思考着如何消除这个操作。时间不知不觉地过去了，当我抬起头时，发现整层楼的客户都已下班，只剩下我们十多个人的开放式办公间还亮着灯。尽管这个问题已经折磨得我筋疲力尽，但我知道不能放弃。直接消除这个算子看起来并不容易，正向的思考陷入了死胡同，我皱起了眉头。这时，我开始逆向地思考，寻找其他的解决方案。我习惯性地在思考时自言自语，试图激发思维的灵感。"降低它的执行时间？或者，减少它的调用次数？”"调用次数！”我灵光一闪，发现了一个新的思路。"对呀，既然消除不行，我可以通过等价的变化，把取奇偶的操作从结果转移到计算的权重上——就好比我先把一份份离散的水排列组合起来，这样我的大勺子就可以一次性舀走了，这样只需要在循环开始前执行一次viewcopy就可以了！”我立刻把思路分享给了昇腾产品部的大模型专家韩俊，韩老师与我一拍即合，当即拍板：“能做！马上做！”达成一致后，我们再无二话，立马分头行动起来，修改代码，测试验证，文档编写……一切高效有序，这也正是我们这个团队几个月来达成的完美默契。X1项目架构训练“三巨头”一转眼到了凌晨三点，大家紧紧盯着屏幕——验证结果出来了！那一瞬间，办公室迸发出了压抑许久的欢呼声，我们终于解决了问题，攻克了拦路大魔王，达成了阶段性的目标！打破天花板，翻过那座山“630”阶段性目标的达成让我们心情愉悦，然而此时还不能松懈，“930”达成持平竞品性能的目标正是下一座山头，横亘在我们前进的道路上。接下来的日子里，我们努力进行局部优化，但似乎已经触及到了天花板，硬件峰值算力和带宽都低于友商，物理规格的限制使我们无法超越竞争对手，我们陷入了束手无策的困境。这个问题愁得我头发一把一把地掉，焦虑和无奈交织在心头，不知道该如何突破这个看似无法逾越的物理限制。一筹莫展的我和计算研究部的厍斌一起蜷缩在工位角落，思绪不由自主地飘远了，如果达不成目标，会不会让主管失望？想到主管，我又想起新员工答辩上，主管对我说的话：“靖义，你要培养自己的系统性思维，多从系统的角度去思考问题。”“系统的角度……”我直起身子，“如果各个局部都不能达成目标，那确实需要各个部分配合起来……”我感觉就差临门一脚了，抓着厍博的手问道：“该怎么配合呢，一加一总不能小于二吧？”等等，对呀，就是要1+1<2！我们瞬间联想到了系统优化中经典的流水线操作。当我们优化系统性能时，流水线操作是一个常见的技术手段，流水线可以将一个复杂的任务划分为多个子任务，并将它们按照顺序连接起来，使得每个子任务可以同时进行。这样可以提高系统的吞吐量和效率。我坐不住了，能否将那些无法掩盖的计算和通信任务也进行流水化，让它们相互掩盖起来？我们满怀激动地将这个新想法与团队的小伙伴们分享，并展示了它的潜力和可行性。大家听完，展开了激烈的讨论。有些人持怀疑态度，认为这个方法可能会增加计算和通信的次数并降低效率，导致训练时间延长。但我坚信这个方法的价值，只要计算和通信能掩盖（重叠）得比较好，就可以弥补我们通信带宽不如友商的劣势，而且友商的计算和通信并行时会竞争资源，导致性能劣化，这就是我们的差异化竞争力。我努力说服他们，让他们相信这是我们迈向成功的一步。经过反复思考和反复推敲，我们终于提出了计算通信并行的系统优化策略。在这个策略中，我们将原本庞大的计算任务拆分成了多个小的计算任务，并将通信任务也细分成了多个子任务。这样一来，原本存在依赖关系的计算任务1完成后，和通信任务1无关的计算任务2就可以并行执行了。但是单独看计算和通信的性能可能会有所下降，并且内存占用也会增高。进一步地，我们协同了海思的同事优化了切分后的矩阵运算的性能，中硬的同事优化了切分后通信的性能，中研和OCK（计算产品线应用使能基础软件平台）的同事一起优化了显存的占用……经过一段时间的不懈努力和精诚合作，我们终于完成了优化方案的实施和测试。当看到系统在新的优化策略下运行时，我们欣喜若狂，性能问题得到了明显改善，整个系统的运行速度大幅提升。这个突破让我们感到无比振奋，我们终于达成了持平竞品的目标！最后，合作伙伴联合创始人及总裁在发布会上宣布，“华为在图形处理器方面的实力，目前已和N公司分庭抗礼，他们将华为NPU技术提升至可与N公司A100匹敌的水准。”那个时刻，我由衷地笑了，但是眼前也似乎朦胧了起来，鼻头酸酸的。是的，感谢计算产品线、海思、2012实验室等小伙伴们的共同努力，我们做到了，这几个月的付出没有白费！X1项目现场攻关留念正如约翰赫斯金所说：“成功是用努力，而非希望造成。”我深深认同这句话。在我看来，技术创新是一个需要持续努力和深入实践的过程。在这个过程中，积累和稳定至关重要。计算产业为我提供了广阔的舞台，我愿将自己不断积累的知识、经验和创新思维，淋漓尽致地展现在这个舞台上。首获计算产业“青杨奖（2022年度最突出贡献青年专家）”向前看，永不言弃。我将坚定地朝着“为世界提供最强算力，让智能无所不及”的目标前进。少年何妨梦摘星，敢挽桑弓射玉衡！来源/《华为人》-心声社区公众号转载请注明作者和出处