B站文生視頻模型工程實踐
一、前言
近年來,AI 內容生成(AIGC)領域的快速發展令人雀躍,OpenAI 在 2023 年初推出大型語言模型(LLM)GPT-4 受到了學術界和工業界的極大關注。OpenAI 隨后在 2024 年初推出文生視頻(T2V)模型Sora,能夠根據文本指令制作出具有現實風格和富有想象力的場景視頻,更是展示了令人驚喜的“世界模擬器”能力。
B站作為UGC內容豐富的視頻網站,在視頻生成模型領域有著天然數據優勢和廣泛應用場景。在此之前我們已經有了一段時間的LLM模型訓練經驗,文生視頻模型結構、語料以及訓練過程有一定的差異性,本文重點介紹B站TTV團隊在文生視頻模型上積極探索后的經驗及感悟。
二、TTV model
在OpenAI提供的公開信息中,Sora模型實際上是一個Diffusion Model+Transformer架構。本文基于B站在生成式TTV在自研道路上的探索、結合行業進展和工程實踐,先后嘗試了幾種TTV(Text to video,后面將簡稱為TTV)模型。文中重點介紹由colossal-ai發布的類Sora模型Open-Sora,以及由智譜AI發布的CogVideoX模型。
2.1 OpenSora
OpenSora的核心是Stdit(Spatial Temporal Diffusion Transformer)結構。
DiT(Diffusion Transformer)模型是一種結合了去噪擴散概率模型(DDPM)和Transformer架構的擴散模型,通過模擬數據的逐步去噪過程來生成新的文本,在此基礎上發展出了STDiT,STDiT模型就是一種使用Cross-Attention的DiT變體。
STDiT模型 的特點包括:
- 融合空間 - 時間注意力機制:結構巧妙地串聯起二維空間注意力模塊和一維時間注意力模塊,能夠在捕捉視頻幀內的空間特征的同時,精準模擬視頻幀與幀間的時序關聯。
- 交叉注意力集中模塊:緊隨在時間注意力模塊后,確保文本語義與生成視頻的深度對齊。
- 計算資源需求降低:相較于全注意力機制顯著減少了計算資源需求。
- 利用預訓練權重:能夠更好地利用預訓練好的圖像 DiT 權重遷移學習至視頻場景。
圖2-1 opensora模型架構
如圖2-1所示,OpenSora的Transformer Block的部分包括了Spatial Attention(空間維度上的注意力計算)、Temporal Attention(時間維度上的注意力計算),以及通過Cross Attention將兩個維度信息進行交叉計算,此外整個架構還包括一個預訓練好的視頻編碼器(VAE),一個文本編碼器(T5)。整個模型的訓練階段如圖2-2所示,首先采用預訓練好的 Variational Autoencoder (VAE) 的編碼器將視頻、圖片數據進行壓縮,將對應的文本描述通過Text Encoder進行embedding化,然后在壓縮之后的隱層空間中與文本嵌入(text embedding)一起訓練 STDiT擴散模型。
圖2-2 opensora訓練流程示意圖
Open-Sora 模型的訓練采用多階段訓練方法,包括大規模圖像預訓練、大規模視頻預訓練和高質量視頻數據微調等。在數據收集和預處理方面,涉及到對視頻美學、動態性、運鏡、字幕等信息的處理,構建高質量的訓練數據集。
其功能特點有:支持視頻數據預處理、加速訓練、推理等全套流程;提供視頻切割和字幕工具,支持剪輯和 T5 文本調節;實現可變長寬比、可變分辨率和可變時長等功能;能夠生成各種風格的視頻內容,支持多種生成任務等。
2.2 CogVideoX
CogVideoX【8】使用了由智譜提出的3D Causal VAE以及專家Transformer模塊,訓練流程與STDIT基本一致,都是通過VAE與Text Encoder將視頻數據與文本描述進行embedding化后,輸入Transformer Block進行訓練。與STDiT不同的是,視頻embedding與文本embedding會直接拼接起來做 Full Attention 計算,以便更好地對齊視覺和語義信息。但是這兩種模態的特征空間差異顯著,它們的嵌入甚至可能具有不同的數值范圍。為了在同一序列中更好地處理它們,Transformer模塊采用了專家自適應層歸一化(Expert Adaptive Layernorm)來獨立處理每種模態,如下圖所示。處理后的模態信息會繼續拼接到一起進行3D Full Attention進行計算,在Attention計算前使用了 3D-RoPE將相對位置信息添加到視覺模態上。
圖2-3 cogvideo模型結構
三、工程實踐
3.1 數據存儲與加載
TTV模型訓練相較于LLM模型,訓練數據集有較大差異性。處理后的訓練數據包括視頻切片、圖片和對應的文字描述。相比于LLM語料庫,TTV單個數據文件較小、但整體數據量大,給訓練造成了一些不便。
此前的LLM模型訓練數據均以5G一個分片整塊存儲在boss上,每個訓練epoch開始前會拉取一個5G的分片到訓練機器上,IO效率較高。但視頻切片的數據平均大小在1M左右,如果在一個epoch的開始拉取所需的訓練數據,boss拉取小文件的效率無法滿足我們的訓練需求,會造成GPU資源閑置。
結合上述特點,我們和B站基礎架構存儲團隊討論過多個方案:
方案一:HDFS+文件打包
相較于boss,HDFS能夠提供更大的數據訪問的帶寬上限。HDFS 的負載通常較為平穩,但也存在抖動場景,且DN和NN受任務影響明顯,當前HDFS數據99%為2副本存儲,當有2臺DN受到任務影響時,數據讀寫性能會有明顯下降。此外因為小文件的特點,拉取效率較難滿足訓練要求,基于這個特點,我們考慮將多個小文件打包在一起,形成一個大chunk file,加速HDFS拉取效率,此外在訓練的腳本中,加入一個dataset reader,解析相應的chunk file,重新還原成小視頻、圖片文件。顯然這個辦法可以解決文件拉取效率問題,并且還節約了一定的存儲空間和文件句柄,對存儲媒介友好。但是存在對代碼的侵入,以及chunk內容調整不靈活的問題,我們短期沒有采用。
方案二:HDFS->backend & Alluxio->frontend
Alluxio 是一個開源的分布式內存文件系統,能大幅提升數據訪問性能,加速大數據分析和處理任務。在機器學習訓練中經常使用alluxio,快速提供數據給模型。并且他能靈活適配,可與多種存儲系統(如 HDFS、S3、Azure Blob Storage 等)和計算框架(如 Hadoop、Spark、Flink 等)集成,適應不同的大數據架構。經基架同學一同努力,最終我們采用了Alluxion Fused+HDFS backend方案,加載數據的方式近似于訪問本地文件系統,訓練框架無感。alluxio定期同步HDFS上的數據到本地,能較好的滿足我們海量小文件的讀寫以及快速上線的訴求。
3.2 數據預處理優化
TTV數據的預處理過程更為繁瑣,除了對視頻的文字描述進行tokenizer外,還需要對視頻數據進行抽幀、歸一化和vae編碼,尤其是vae編碼這一步在單epoch訓練中會占到總時長的18%,并且極為消耗顯存。為了優化這部分的速度,我們做了兩部優化。
- 數據并行
參考訓練中的數據并行策略(Data Parallel),在訓練任務啟動的初始階段,初始化一個包含全部顯卡的全局通信組,將一個epoch訓練中最耗算力的編碼部分拆分到所有顯卡上,待每張卡處理完成后再將所有的結果gather到0號卡上進行聚合和保存,后續的訓練階段只需讀取處理好的embedding數據即可。
- VAE離線化
考慮到一份訓練數據可能會在后續下訓練任務中被反復使用,并且數據處理非常消耗算力和顯存,因此在第一步優化的基礎上,我們會把聚合好的訓練結果持久化,后續直接從HDFS\Alluxio中讀取,整體顯存和訓練時間都有較大優化。
上述過程如圖3-1所示:
圖3-1 數據讀取、預處理示意
3.3 模型并行優化
雖然TTV模型通常參數量較小,但視頻訓練數據更為復雜,序列長度較長。因此TTV模型訓練過程中的激活內存成本極高,導致訓練速度明顯更慢,且限制了訓練數據的清晰度與抽幀率的提高。
在Transformer Block中,激活顯存占用的大頭為Attention的計算,在不使用Flash-attention的情況下,attention激活占用與序列長度的平方成正比,而使用了時空Cross-Attention設計的STDIT模型,訓練過程中更是包含了3次attention的計算,使得激活占用更是大大增加。針對Transformer訓練過程中,序列長度過長導致的激活占用高的問題,業界已有一些序列并行方案提被出,序列并行是一種專門為跨多個設備分發長序列和激活而設計的技術,以下四種是主要的序列并行技術, Ring-Attention、Megatron-SP、DeepSpeed-Ulysses和Dyncamic SP,首先將介紹下各方案的設計。
業界方案
- Ring Attention
Ring Attention 的核心思想是將輸入序列分割成多個塊,并將這些塊分布在多個計算設備上進行并行處理,通過使用 Online Softmax 機制,在不保留完整序列長度的情況下計算注意力分數。每個設備首先對自己的數據塊進行局部的自注意力計算,然后將關鍵信息(key-value pair)傳遞給下一個設備,通過一個環形的數據傳輸策略,實現在不增加單個設備內存負擔的情況下處理超長序列。然而,環形注意力對 P2P 通信的依賴在高延遲環境中可能效率較低。
- Megatron-SP
Megatron框架中的的序列并行,是在張量并行(Tensor Parallelism)的基礎上,將Transformer Block中的LayerNorm以及Dropout層的輸入按Sequence Length維度進行了切分,使得各個設備上面只需要做一部分的Dropout和LayerNorm。雖然減少了每張卡上的激活占用,但在通信過程中引入了額外的all-gather和reduce-scater操作。因在設計上依托切分注意力頭來實現并行,使用上也會受到注意力頭數量的限制。
圖3-2 megatron-sp
- DeepSpeed-Ulysses
DeepSpeed-Ulysses 【7】引入了一種創新的方法,通過利用全對全(all-to-all)集體通信來訓練長序列。在處理長序列時,它將查詢(query)、鍵(key)和值(value)矩陣在注意力頭之間進行切分,但保留原始注意力計算結構。這個過程通過兩組全對全通信來實現,這兩組通信在序列分割和注意力頭分割之間交替進行。這樣的設計使得在處理長序列時,能夠在保持計算結構的同時,有效地在多個 GPU 之間分配數據,減少通信開銷。
圖3-3 deepspeed-ulysses
- Dyncamic SP
與以上三種針對單序列維度內的并行性方案不同,DSP(Dynamic Sequence Parallel)【6】方案是針對多維序列的并行問題所設計。在多維Transformer模塊中,每個序列維度的計算其實是獨立進行的,因此可以分別在不同緯度切分和計算,僅在維度計算交換的時間點使用高效的全對全操作(all-to-all)來為中間序列切換并行維度并重新進行動態切分。這種方法使 DSP 能夠獨立于模塊內的計算邏輯外,消除了模塊內許多不必要的通信。
圖3-4 dsp
相關實踐
- 基于OpenSora SP實現
OpenSora模型使用了時空交叉注意力機制,會分別計算視頻空間維度的Attention、時間維度的Attention,并最后和文本embedding進行交叉Attention計算,其設計更適合使用DSP方案。具體實現上,我們會在空間注意力計算前,先在空間維度進行切分,SP并行組內的各rank分別計算不同的seq片段。隨后在時間維度注意力計算前,進行一次all-to-all通信進行同步,并交換切分緯度到時間維度上。因為與文本embedding的交叉注意力計算,只與空間信息有關,因此可在交叉Attention計算后再進行一次all-to-all通信來同步計算結果。
- 基于CogvideoX的SP實現
考慮到CogVideoX模型雖然使用了視頻信息與文字信息,但兩種embedding是在單一維度進行拼接,并進行全局Attention計算的,因此本身屬于單維Transformer。對于單維Transformer,我們選擇實現了DS-UIysses方案,在實現上:
a. 先在Transformer Block之前,沿sequence維度進行切分,但只對視頻的hidden state的seq維進行切分。文字embedding的部分不切分,并與切分后的視頻embedding拼在一起。
b. 相對位置編碼是需要對全序列長度進行計算的,因此在每個3D-Rope計算前,進行一次all-to-all通信,回收sequence維度的信息。
c. 因為每個seq段都拼接了文字embedding,因此首先需要通過remove_extra_encoder方法,移除每段seq冗余的text embedding。在計算完Attention后,進行一次all-to-all通信,在attention head維度回收信息,并回到sequence切分的狀態。在MLP Block計算前,通過 add_text_encoder補上每段seq都有的文字embedding部分。
d. 在所有的Transformer Block后,進行gather sequence操作,合并SP組各組上的計算結果。
通過使用SP并行策略,CogVideoX由單機16卡只能訓練45幀1080p的數據,提升至可訓練221幀1080p的訓練數據。
圖3-5 基于CogVideoX的SP實現示意圖
四、文生視頻模型在NPU架構上的工程實踐
目前我們的訓練算力構成為GPU + NPU,但GPU與NPU在芯片設計上差異較大,軟件棧和生態也存在較大區別,因此需要做相應的適配和優化才能充分利用NPU資源進行訓練。
4.1 基礎適配
- 模型適配
目的是讓模型能夠在npu環境下開箱啟用(訓練),保證pipeline可運行。其主要工作為檢查依賴cuda硬件及精度在npu下受限的算子,查詢入參匹配的npu-wrapped算子,進行替換與輕量的代碼適配,例如替換T5模塊中的LayerNorm為NPURmsNorm、在nn.Conv3d算子中顯式指定使用torch.float16精度等。
- 框架移植
我們的訓練框架中部分模塊開采用megatron-core作為加速手段,需要移植到對應NPU的版本,主要參考華為Megatron-NPU倉庫的范例和實現,進行移植。
- 精度驗證
可訓練后,需要進行精度的驗證。精度驗證需注意保證GPU版本與NPU版本的第三方依賴保持一致,并固定代碼中隨機的部分,例如隨機數設置、數據的抽樣、Vision encoder中的加躁信息等。可以借助華為的精度驗證工具,dump主要算子的輸入輸出,并逐步進行API級別、模塊級別與整網級別的精度驗證。
不過值得注意的是,因為底層實現有所區別,配置不同,以及fusion的使用與否等等,loss是無法完全對齊的。下面列出幾個常見的,會引起loss差異的參數以供參考。
- transformer_impl: local -> transformer_engine
- attention_softmax_in_fp32: False -> True
- apply_rope_fusion: False -> True
- rotary_fusion: False -> True
- swiglu_fusion: False -> True
案例
基于內部基座版本:global batch size 160,Run 1曲線代表NPU,Run 2曲線代表GPU。
2940-steps:
Loss diff max = 0.46540455520153046, diff mean = 0.006955942064507858
圖片
1500-steps:
Loss diff max = 0.2761184275150299, diff mean = 0.008173752022286256
圖片
對比趨勢和diff,基本一致,可以確認模型參數基本可用,更細粒度則需要做到分層精度對比以及生成視頻benchmark對比。
基于CogVideo-5B基座版本:8卡,25幀,480P
圖片
4.2 優化
- profiler性能分析調優
A.后向耗時問題定位
下圖profiler中可以發現,在后向過程中執行了一次前向,占用接近20%,但其實我們并未顯式地配置重計算策略。
后續分析代碼發現,訓練中會默認使用torch.utils.checkpoint對全層數的transformer進行重計算,但基于如下幾點原因可以對其優化:
a.顯存仍未用盡,有tradeoff的余地
b.需要重計算的層數應該可控
c.重計算作用在一個融合算子上,其中涉及MLP,GN,QKV計算,attention計算等,但其實只有attention計算使用該策略的價值最高
經過一系列的代碼調整后,最終訓練速度提升約12%。
B.融合算子替換
根據profiling分析,紅框中gelu算?實際執?時是以多個?算?拼接的形式下發和執?的。
圖片
可以使?融合算?F.gelu 進?替換,優化下發和執?。跟據profiling中的API調?棧可以定位到該算?主要出現在T5模塊和megatron中。
C.使用連續張量減少切分操作
圖片
圖片
如上圖profiling所示,耗時占比最大的3個op中,第二第三都是計算密集型的,可以從策略或算法的角度優化。第一位StridedSlice是訪存密集型,改變使用StridedSlice 的上層算子的輸入tensor的存儲方式(使用torch.contiguous)即可連續分配,優化原理說明如下:
torch.contiguous操作將非連續張量轉換為連續張量,后續的切片和訪問操作大幅簡化,甚至避免StridedSlice的調用。
主要體現在三種場景:
其一,連續張量的切片不需要依賴復雜的stride信息。硬件可以更高效地預取數據,提高計算速度。
其二,非連續張量的 StridedSlice 需要動態計算目標地址,頻繁調用可能帶來性能瓶頸。連續張量的切片是直接基于線性偏移量完成的,減少了計算需求。
其三,某些操作(如 view)要求張量是連續的。如果張量已經是連續的,相關操作無需隱式調用 StridedSlice 或創建新的張量拷貝,直接提升性能。
在我們的訓練中得到加速的主要是第三種場景。上右profiling是優化后的結果。
- FlashAttention(NPU)
在torch2.0之前的時代,業界都會使用樸素的standard_attention進行注意力機制的計算,但當其attention_mask為精度f32時會引起巨大的顯存占用及 NaN-bug。后續torch推出了一個改良算子scaled_dot_product_attention(sdpa),進行算子融合,優化內存使用,適合長序列,內置支持 Dropout 和 Causal Mask。現今,主流的大模型結構中都會采用FlashAttention,是一種sdpa的實現,相較于原版scaled_dot_product_attention,其對QKT進行了分塊處理,避 免存儲完整的注意力矩陣,可處理更長的序列。
圖片
在遷移至NPU架構后,我們逐步將原有的注意力機制代碼,改進為基于 NPU底層實現的FlashAttention,其中涉及的抽象、封裝,引入必要的設計模式等。以下僅對flashAttention作簡單介紹。
FlashAttention本質是算子和數據的融合:
- 將多個算子合并為一個,簡化計算過程,減少計算量,提高計算效率。
- 將多個中間結果合并為一個,減少內存占用,提高內存利用率;減少不同算子之間的傳輸,提高數據處理效率。
- 簡化代碼實現,減少代碼量,提高代碼可讀性和可維護性。
FlashAttention實現原理【1】:三基石
- Tiling切片:利用高速SRAM代替內存,但SRAM內存小,無法一次性完成所有數據的注意力計算,需要進行分塊計算,對應上文中的QKT分塊。
- 重計算:放棄中間結果寫回,需要使用時重新計算,用計算換訪存。
- Kernel Fusion:將多個操作融合為一個操作,基于Tiling利用一個kernel完成整個計算,對應上文中的算子融合。
以下是前向、反向過程的公式描述,具體細節不在此討論。
前向Forward: FlashAttentionScore[2]
后向Backward: FlashAttentionScoreGrad[3]
- 虛擬內存特性:expandable_segments【4】
一般情況下,由PyTorch自己管理虛擬地址與物理地址映射,降低內存碎片。其原理大致如下:對于大于 2MB 的分配,分配器會調用 aclrtMalloc,以獲取與用戶請求大小完全相同的內存分配。后續計算中,如果這些分配中的某些部分空閑,它們可以被重新用于其他請求。這種方式在程序多次請求相同大小或者是該大小整數倍的內存時效果很好。許多深度學習模型的行為符合這一特點。
然而,有一種常見的例外情況是,批次大小在每次迭代中會略有變化,例如在批量推理中。當程序最初以批次大小N運行時,會為該大小進行合適的內存分配。如果后續運行的批次大小變為N?1,現有的內存分配仍然足夠使用。然而,如果批次大小變為N+1,則需要進行新的內存分配,這些后續分配并非所有張量的大小都相同。一些張量的大小可能是(N+1)×A,而另一些可能是(N+1)×A×B,其中A和B是模型中與批次無關的維度。由于分配器會在現有分配足夠大時重用它們,因此某些(N+1)×A 的分配可能會勉強適應已經存在的N×B×A 段,盡管不完全匹配。當模型運行時,這些段會被部分填充,導致在段末尾留下不可用的空閑內存切片。最終,分配器可能需要調用 aclrtMalloc 為新的(N+1)×A×B 段分配內存。如果沒有足夠的內存,就會拋出異常,結束程序。對于有 50 層以上的模型,這種模式可能會重復 50 多次,從而產生許多小的內存碎片。
通過分析訓練時采集的profiling文件,發現我們研發中常見的顯存瓶頸符合上述場景描述:
a.動態shape場景,比如VAE中的升降采樣,shape隨step增加而增大,從而導致顯存塊不能復用,碎片上升
b.transformer一般由多層組成,其帶來的大量激活數據與激活數據size不統一,會給顯存復用帶來困難,不僅會影響性能,嚴重時導致OOM
采用NPU自研的“內存池擴展段”功能,有助于優化上述場景。
“內存池擴展段”在最初創建一個內存段后,可按需擴展其大小。與每次分配都創建一個新的內存段不同,它嘗試為每個流(stream)創建一個內存段,并根據需求動態增長。當程序運行到批次大小N+1的情況時,這些分配會很好地排列到單個大型內存段中,直到段被填滿。然后,分配器會請求更多內存,并將其追加到該段的末尾。這種方式不會產生太多無法使用的內存碎片,因此更有可能成功找到所需的內存。
五、后續工作方向
1.引入流水線并行(Pipeline Parallelism)
如上文所述,現在已采用的優化方案多集中在tensor計算的某些維度,比如sp/cp針對的是tensor的序列維度,flash_attention針對的tensor的特征維度。在業界主流優化方案中,以layer或模塊作為優化對象的Pipeline Parallelism,PP也是值得探索的方向。
優點:
- 將模型分成多個階段,每個階段只在一個NPU上運行,每個NPU只需要存儲它負責的部分模型,而不是整個模型。這大大減少了顯存消耗。類似FSDP與ZERO機制。
- 模型并行MP以及張量并行TP會導致頻繁的跨設備通信,而PP通過流水線操作僅傳遞每個階段最終結果,減少了通信負擔。
- 擴展性強,比如transformer這類的堆疊結構
挑戰:
- 流水線填充與同步延遲:前向傳播和反向傳播之間存在同步延遲(Pipeline Bubble)。增加batch size是常見的優化策略,但可能導致內存壓力增大。
- 負載均衡:如果模型切分不均勻,會導致某些NPU過載,而其他NPU閑置。
- 實現復雜性:需要在代碼層面將模型和數據流切分為多個階段,并設計高效的通信方案。
2.突破torch.nn.GroupNorm的限制
在處理VAE encoding時,雖然我們在序列維度進行了切分,但在處理groupnorm時,常規做法需要統計全序列的數值才可計算。但全序列groupnorm,計算時的激活tensor過大,單份顯存~26.9G = 128 * 9 * 720 * 1088 * 16 * 2 / 1024 / 1024 / 1024,如果后續增加視頻幀數或是視頻分辨率,則將對NPU的64G顯存上限是一個不小的挑戰。
一種方案是對group維度進行切分,有2個關注點:
- 序列維度合并后才能進行group維度的切分,在此過程中是否會發生用于存儲中間結果的tensor無法申請到顯存的情況
- 對預訓練的groupnorm權重進行提取,并新建對應的batchnorm算子,計算后,還需還原序列維度的切分狀態
另一種方案,全序列groupnorm計算需傳輸高維張量。但從原理上,只需要獲得序列切分后的統計值(低維)即可,單卡獲取全量統計值后,可以處理自身的序列切分數據,其中有3個難點:
- 本質上是使用理論公式對groupnorm進行函數級別的分拆實現,涉及eps,有偏無偏參數,統計值精度等細節對齊,以保證與torch.nn.GroupNorm的誤差可接受
- 因為是一個FusedOp-reverse的過程,會造成速度的損失,需考量可接受度
- 與第一種方案類似,需對預訓練的groupnorm權重進行提取;另外在某些場景下,還需要實現backward代碼,實現難度較高
3.分層ZERO3
Deepspeed 的 zero-3:是一種用于深度學習優化的技術。在分布式訓練框架下,zero-3 將訓練狀態(包括權重、梯度和優化器狀態)分布到不同的顯卡上,以優化顯存利用。與 zero-2 相比,它還對模型參數進行了分區,顯存減少幅度與使用的 GPU 數量成正比。在TTV這種顯存為瓶頸的場景,可以支持更大參數量的模型,提高batch的大小,優化訓練效率。
參考文獻
- https://arxiv.org/pdf/2205.14135
- https://gitee.com/ascend/cann-ops-adv/blob/master/docs/FlashAttentionScore.md
- https://gitee.com/ascend/cann-ops-adv/blob/master/docs/FlashAttentionScoreGrad.md#https://gitee.com/link?target=https%3A%2F%2Fcreativecommons.org%2Flicenses%2Fby%2F4.0%2Flegalcode
- https://gitee.com/ascend/pytorch/blob/master/torch_npu/csrc/core/npu/NPUCachingAllocator.cpp#L240
- https://arxiv.org/abs/2205.05198
- https://arxiv.org/abs/2309.14509
- https://arxiv.org/abs/2403.10266
- https://www.alphaxiv.org/abs/2408.06072v1