少年何妨夢摘星
張靖義
我很喜歡十六七歲的夏天。
彼時的少年手持長弓，意氣風發地眺望遠方，端直了燕尾，搭上虎筋弦，弓弦一振，箭去不知所蹤。
十余年后，我走在路上，背后隱約傳來簌簌風聲，一回頭，一支箭正中紅心。
我，一名標準的“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（神經網絡處理器）上執行效率低。問題原因基本定位清楚，接下來就是需要思考如何優化解決這個“形狀親和匹配的問題”。
經過對數據的反復思考和分析后，我的腦海中迸發出一個新的想法——試試一種自適應的“形狀分桶算法”，通過動態調整輸入數據使其更加親和。這個方案就好像孩童在搭積木時，會先把形狀相同的積木放一起，再想辦法將異形拼接成形狀相同的方塊，然后就可以進行高效的搭建。當然，我們操作的“積木分拼”都是自動化實現的，這樣會使得矩陣乘運算更加高效。
此外，我還主導重新設計了模型主體部分的結構，在保證計算量、參數量和之前模型一致的前提下，盡量提升主體模型的親和性。提高了負載應用的親和性后，我們設計了一套針對大模型的瓶頸分析方法，采用自頂向下的思想，從集群、單機、芯片多層級多粒度地分析大模型訓練中的瓶頸，讓算子得以針對性優化，整個算子性能也得到了極大提高。
稼先社區論壇現場，與公司專家討論大模型
、海思和昇騰使能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的目標，但是算子的性能優化并沒有達到預期，我們離目標還有一段距離。
我縮在辦公區隔間的最里面，絞盡腦汁思考著各種可能的優化方案。我盯著算子性能測試結果中占比高的（視圖拷貝）操作，這是我們芯片上由于非連續轉連續引入的額外操作，它在整個計算中占據了超過10%的比例，也是最有可能優化的節點。
“能不能去掉這個操作呢？” 我自言自語。
我開始仔細研究產生操作的根因，尋找可以快速完成的優化方案。經過一番思索和實驗，我發現是客戶在循環中對一個計算結果向量做了非常規的按照奇偶索引操作，這個操作對昇騰NPU非常不親和。我們的NPU對大塊的連續內存讀寫是親和的，但是對小塊的不連續跳讀跳寫是不親和的，就好比用了一個超大的勺子，但是每次只能舀的水只有一點點，遠遠不能裝滿勺子，搬運效率非常低。
找到了根因，我稍微松了一口氣，拿出了紙筆開始寫寫畫畫起來。我將這段代碼的邏輯轉換成公式寫在了紙上，思考著如何消除這個操作。時間不知不覺地過去了，當我抬起頭時，發現整層樓的客戶都已下班，只剩下我們十多個人的開放式辦公間還亮著燈。盡管這個問題已經折磨得我筋疲力盡，但我知道不能放棄。
直接消除這個算子看起來并不容易，正向的思考陷入了死胡同，我皺起了眉頭。這時，我開始逆向地思考，尋找其他的解決方案。我習慣性地在思考時自言自語，試圖激發思維的靈感。
"降低它的執行時間？或者，減少它的調用次數？”
"調用次數！”我靈光一閃，發現了一個新的思路。
"對呀，既然消除不行，我可以通過等價的變化，把取奇偶的操作從結果轉移到計算的權重上——就好比我先把一份份離散的水排列組合起來，這樣我的大勺子就可以一次性舀走了，這樣只需要在循環開始前執行一次就可以了！”我立刻把思路分享給了昇騰產品部的大模型專家韓俊，韓老師與我一拍即合，當即拍板：“能做！馬上做！”
達成一致后，我們再無二話，立馬分頭行動起來，修改代碼，測試驗證，文檔編寫……一切高效有序，這也正是我們這個團隊幾個月來達成的完美默契。
X1項目架構訓練“三巨頭”
一轉眼到了凌晨三點，大家緊緊盯著屏幕——驗證結果出來了！那一瞬間，辦公室迸發出了壓抑許久的歡呼聲，我們終于解決了問題，攻克了攔路大魔王，達成了階段性的目標！
打破天花板，翻過那座山
“630”階段性目標的達成讓我們心情愉悅，然而此時還不能松懈，“930”達成持平競品性能的目標正是下一座山頭，橫亙在我們前進的道路上。
接下來的日子里，我們努力進行局部優化，但似乎已經觸及到了天花板，硬件峰值算力和帶寬都低于友商，物理規格的限制使我們無法超越競爭對手，我們陷入了束手無策的困境。這個問題愁得我頭發一把一把地掉，焦慮和無奈交織在心頭，不知道該如何突破這個看似無法逾越的物理限制。
一籌莫展的我和計算研究部的厙斌一起蜷縮在工位角落，思緒不由自主地飄遠了，如果達不成目標，會不會讓主管失望？想到主管，我又想起新員工答辯上，主管對我說的話：“靖義，你要培養自己的系統性思維，多從系統的角度去思考問題。”
“系統的角度……”我直起身子，“如果各個局部都不能達成目標，那確實需要各個部分配合起來……”我感覺就差臨門一腳了，抓著厙博的手問道：“該怎么配合呢，一加一總不能小于二吧？”等等，對呀，就是要1+1
我們瞬間聯想到了系統優化中經典的流水線操作。當我們優化系統性能時，流水線操作是一個常見的技術手段，流水線可以將一個復雜的任務劃分為多個子任務，并將它們按照順序連接起來，使得每個子任務可以同時進行。這樣可以提高系統的吞吐量和效率。我坐不住了，能否將那些無法掩蓋的計算和通信任務也進行流水化，讓它們相互掩蓋起來？
我們滿懷激動地將這個新想法與團隊的小伙伴們分享，并展示了它的潛力和可行性。大家聽完，展開了激烈的討論。有些人持懷疑態度，認為這個方法可能會增加計算和通信的次數并降低效率，導致訓練時間延長。但我堅信這個方法的價值，只要計算和通信能掩蓋（重疊）得比較好，就可以彌補我們通信帶寬不如友商的劣勢，而且友商的計算和通信并行時會競爭資源，導致性能劣化，這就是我們的差異化競爭力。我努力說服他們，讓他們相信這是我們邁向成功的一步。
經過反復思考和反復推敲，我們終于提出了計算通信并行的系統優化策略。在這個策略中，我們將原本龐大的計算任務拆分成了多個小的計算任務，并將通信任務也細分成了多個子任務。這樣一來，原本存在依賴關系的計算任務1完成后，和通信任務1無關的計算任務2就可以并行執行了。但是單獨看計算和通信的性能可能會有所下降，并且內存占用也會增高。進一步地，我們協同了海思的同事優化了切分后的矩陣運算的性能，中硬的同事優化了切分后通信的性能，中研和OCK（計算產品線應用使能基礎軟件平臺）的同事一起優化了顯存的占用……
經過一段時間的不懈努力和精誠合作，我們終于完成了優化方案的實施和測試。當看到系統在新的優化策略下運行時，我們欣喜若狂，性能問題得到了明顯改善，整個系統的運行速度大幅提升。這個突破讓我們感到無比振奮，我們終于達成了持平競品的目標！
最后，合作伙伴聯合創始人及總裁在發布會上宣布，“華為在圖形處理器方面的實力，目前已和N公司分庭抗禮，他們將華為NPU技術提升至可與N公司A100匹敵的水準。”
那個時刻，我由衷地笑了，但是眼前也似乎朦朧了起來，鼻頭酸酸的。是的，感謝計算產品線、海思、2012實驗室等小伙伴們的共同努力，我們做到了，這幾個月的付出沒有白費！
X1項目現場攻關留念
正如約翰赫斯金所說：“成功是用努力，而非希望造成。”我深深認同這句話。在我看來，技術創新是一個需要持續努力和深入實踐的過程。在這個過程中，積累和穩定至關重要。計算產業為我提供了廣闊的舞臺，我愿將自己不斷積累的知識、經驗和創新思維，淋漓盡致地展現在這個舞臺上。
首獲計算產業“青楊獎（2022年度最突出貢獻青年專家）”
向前看，永不言棄。我將堅定地朝著“為世界提供最強算力，讓智能無所不及”的目標前進。
少年何妨夢摘星，敢挽桑弓射玉衡！
來源/《華為人》