來(lái)源：DeepHub IMBA
       
      
     
    
    

     

      

       

        

         
本文約3100字，建議閱讀6分鐘
         
本文中介紹了使用PyTorch Profiler來(lái)查找運(yùn)行瓶頸和一些簡(jiǎn)單的提速方法。
        
       
      
     
    

如果所有機(jī)器學(xué)習(xí)工程師都想要一樣?xùn)|西，那就是更快的模型訓(xùn)練——也許在良好的測(cè)試指標(biāo)之后。

加速機(jī)器學(xué)習(xí)模型訓(xùn)練是所有機(jī)器學(xué)習(xí)工程師想要的一件事。更快的訓(xùn)練等于更快的實(shí)驗(yàn)，更快的產(chǎn)品迭代，還有最重要的一點(diǎn)需要更少的資源，也就是更省錢(qián)。

熟悉PyTorch Profiler

在進(jìn)行任何優(yōu)化之前，你必須了解代碼的某些部分運(yùn)行了多長(zhǎng)時(shí)間。Pytorch profiler是一個(gè)用于分析訓(xùn)練的一體化工具。它可以記錄：

CPU操作時(shí)間、CUDA內(nèi)核計(jì)時(shí)、內(nèi)存消耗歷史。

要記錄事件，只需要將訓(xùn)練嵌入到分析器上下文中，如下所示：

 import torch.autograd.profiler as profiler
 with profiler.profile(  activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],  on_trace_ready=torch.profiler.tensorboard_trace_handler('./logs'), ) as prof:  train(args)

然后就可以啟動(dòng)tensorboard查看分析軌跡。如果這一步有問(wèn)題，請(qǐng)查看是否安裝了torch-tb-profiler。

Profiler有很多不同的選項(xiàng)，但最重要的是activities和profile_memory，一般情況下我們只需要這兩個(gè)選項(xiàng)，因?yàn)閱⒂玫倪x項(xiàng)越少，開(kāi)銷(xiāo)就越小。

如果只想要分析CUDA內(nèi)核執(zhí)行時(shí)間，那么關(guān)閉CPU分析和所有其他功能也是可以的。因?yàn)樵谶@種模式下，我們可以理解為顯卡能力的真實(shí)評(píng)測(cè)。

為了方便分析，我們可以為每一步操作指定名稱(chēng)，例如

 with profiler.record_function("forward_pass"):  result = model(**batch)
 with profiler.record_function("train_step"):  step(**result)

或者增加更精細(xì)的自定義的標(biāo)簽，這里的名稱(chēng)將在跟蹤中可見(jiàn)，我們就可以更簡(jiǎn)單的追蹤想要的東西了。

 with profiler.record_function("transformer_layer:self_attention"):  data = self.self_attention(**data)
 ...
 with profiler.record_function("transformer_layer:encoder_attention"):  data = self.encoder_attention(**data, **encoder_data)

查看PyTorch Traces

收集完信息后，tensorboard顯示是這樣的。

訓(xùn)練的過(guò)程一般包括：數(shù)據(jù)加載、前向傳播、反向傳播。

反向傳播由PyTorch在一個(gè)單獨(dú)的線(xiàn)程中處理(上圖中的線(xiàn)程16893)，因此很容易識(shí)別，這部分門(mén)也控制不了，因?yàn)槎际荘ytorch根據(jù)我們的計(jì)算來(lái)自動(dòng)進(jìn)行的。（當(dāng)然也可以自定義反向傳播，但是這過(guò)于復(fù)雜，一般不建議自己實(shí)現(xiàn)）

首先看看數(shù)據(jù)加載：對(duì)于數(shù)據(jù)加載我們希望時(shí)間接近于零。

這是因?yàn)樵跀?shù)據(jù)加載過(guò)程中，GPU什么也不做，這會(huì)使可用資源利用率不足。并且在Pytorch的訓(xùn)練時(shí)數(shù)據(jù)處理可以與GPU計(jì)算重疊，因?yàn)樗鼈兪仟?dú)立的部分，也就是說(shuō)我們加載一個(gè)批次的時(shí)間只要與一個(gè)前向和一個(gè)反向傳播的時(shí)間相近就可以了，這樣就可以最大化的利用GPU的資源。

這里可以很容易地識(shí)別GPU空閑的區(qū)域-查看性能分析器跟蹤中的GPU Est. SM效率和GPU利用率數(shù)字。沒(méi)有活動(dòng)的區(qū)域是我們的關(guān)注點(diǎn)，因?yàn)镚PU什么都不做。

如果使用PyTorch DataLoader，則可以通過(guò)指定num_workers來(lái)多線(xiàn)程處理數(shù)據(jù)。如果您使用IterableDataset，則會(huì)更復(fù)雜，因?yàn)閿?shù)據(jù)將被復(fù)制。這個(gè)問(wèn)題可以通過(guò)使用get_worker_info()來(lái)解決，需要以某種方式調(diào)整迭代，以便每個(gè)worker接收不同的、不相交的行，所以這個(gè)比較麻煩，一般盡量避免IterableDataset。

內(nèi)存分配器 memory allocator

當(dāng)你在CUDA設(shè)備上使用PyTorch分配張量時(shí)，PyTorch將使用緩存分配器。這里是CUDA的執(zhí)行機(jī)制：cudaMalloc和cudaFree的操作比較昂貴，我們要盡量避免。所以PyTorch會(huì)嘗試重用以前通過(guò)cudaMalloc塊分配的，如果PyTorch的分配器有一個(gè)合適的塊可用，它會(huì)直接給出它，而不調(diào)用cudaMalloc。這樣cudaMalloc只在開(kāi)始時(shí)被調(diào)用。

但是如果你處理的是可變長(zhǎng)度的數(shù)據(jù)（比如文本數(shù)據(jù)），不同的正向傳播將需要不同大小的中間張量。因此，PyTorch的分配器可能沒(méi)有適當(dāng)?shù)目捎脭?shù)據(jù)塊。在這種情況下，分配器會(huì)調(diào)用cudaFree釋放以前分配的塊，為新的分配釋放空間。

然后分配器再次開(kāi)始構(gòu)建它的緩存，進(jìn)行大量的cudaMalloc，這是一個(gè)昂貴的操作，但是可以通過(guò)tensorboard分析器查看器的內(nèi)存分析器部分來(lái)發(fā)現(xiàn)這個(gè)問(wèn)題。

可以看到與分配器的保留內(nèi)存相對(duì)應(yīng)的紅線(xiàn)不斷變化。這意味著PyTorch分配器不能有效地處理分配請(qǐng)求。而當(dāng)分配程序在沒(méi)有頻繁調(diào)用的情況下處理分配時(shí)，紅線(xiàn)是完全筆直的，如下圖所示：

我們?nèi)绾谓鉀Q呢？

第一件值得嘗試的事情是設(shè)置PyTorch相對(duì)較新的分配器模式:

 PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

這告訴PyTorch分配器分配可以在將來(lái)擴(kuò)展的塊。但是，如果大小變化太大，它仍然可能無(wú)法解決問(wèn)題。

所以我們智能手動(dòng)來(lái)進(jìn)行優(yōu)化，那就是是使數(shù)據(jù)形狀一致。這樣分配器就更容易找到合適的數(shù)據(jù)塊進(jìn)行重用。

比如最簡(jiǎn)單的將數(shù)據(jù)填充到相同的大小?；蛘呖梢酝ㄟ^(guò)運(yùn)行具有最大輸入大小的模型來(lái)預(yù)熱分配器。

內(nèi)存歷史記錄

我們想要最大化的使用所有可用的GPU內(nèi)存——這讓我們能夠運(yùn)行大量數(shù)據(jù)，并更快地處理數(shù)據(jù)。但是在某些時(shí)候，當(dāng)增加批處理太大時(shí)，將遇到CUDA內(nèi)存不足錯(cuò)誤。是什么導(dǎo)致了這個(gè)錯(cuò)誤?

為了調(diào)試它，我們可以查看分配器的內(nèi)存歷史記錄。它可以通過(guò)PyTorch記錄，然后在https://pytorch.org/memory_viz上可視化。

• Start: torch.cuda.memory._record_memory_history(max_entries=100000)
• Save:
  torch.cuda.memory._dump_snapshot(file_name)
• Stop:
  torch.cuda.memory._record_memory_history(enabled=None)

可視化會(huì)畫(huà)出這樣的東西:

x軸表示時(shí)間，y軸表示已使用的總內(nèi)存，彩色塊表示張量。它顯示了張量何時(shí)被分配，何時(shí)被釋放。

你可能會(huì)注意到狹窄的尖峰，這些是持續(xù)時(shí)間很短的張量，并且占據(jù)了很多空間。通過(guò)點(diǎn)擊一個(gè)張量，可以得到這個(gè)張量被分配到哪里的信息。我們希望的就是最小化這些峰值，因?yàn)樗鼈兿拗屏擞行У膬?nèi)存使用。檢查導(dǎo)致這個(gè)峰值的原因，并考慮優(yōu)化或者使用其他計(jì)算方法替代。

除了峰值之外，很容易檢測(cè)到內(nèi)存泄漏：

第一次運(yùn)行之后的一些數(shù)據(jù)沒(méi)有被清除，所以導(dǎo)致內(nèi)存占用過(guò)高。通過(guò)點(diǎn)擊塊，可以知道這些張量是從哪里來(lái)的。在圖像中，梯度在訓(xùn)練步驟之后沒(méi)有被清除，因此它們?cè)谙蚯皞鬟f過(guò)程中處于無(wú)用狀態(tài)，占用了寶貴的內(nèi)存。

提高模型速度，減少內(nèi)存使用

我們知道了原因，并且可以通過(guò)Profiler來(lái)找到瓶頸，那么我們可以通過(guò)什么方法來(lái)加速訓(xùn)練呢？

1、FlashAttention

首先可以使用FlashAttention來(lái)計(jì)算點(diǎn)積注意力來(lái)提高效率。如果你沒(méi)有聽(tīng)說(shuō)過(guò)它，它是一種計(jì)算精確的點(diǎn)積注意力的方法，并且不需要明確地構(gòu)建注意力矩陣。這優(yōu)化了GPU的io操作，提高了速度，也極大地減少了內(nèi)存消耗。

但是FlashAttention僅適用于兼容硬件上的fp16和bf16精度。那就是NVIDIA Ampere, Hooper以上的GPU

當(dāng)然也有其他的庫(kù)可以替換，例如XFormers，和NV自己的Transformer Engine。

新版本的PyTorch也內(nèi)置了FlashAttention的支持，在文檔中：

torch.backends.cuda.enable_flash_sdp(): Globally enables or disables FlashAttention.

2、 FSDP 優(yōu)化多gpu數(shù)據(jù)冗余

如果使用多個(gè)gpu來(lái)運(yùn)行訓(xùn)練，基本的解決方案是使用DistributedDataParallel。生成了幾個(gè)相同的進(jìn)程，并且在反向傳播期間聚合梯度。

當(dāng)我們生成相同的進(jìn)程時(shí)，在每個(gè)GPU上都有相同的模型和優(yōu)化器狀態(tài)，這是冗余的?？梢酝ㄟ^(guò)跨數(shù)據(jù)分片來(lái)優(yōu)化內(nèi)存使用。

當(dāng)在多個(gè)gpu上進(jìn)行訓(xùn)練時(shí)，每個(gè)進(jìn)程在使用DDP進(jìn)行訓(xùn)練時(shí)都有相同數(shù)據(jù)的精確副本?？梢酝ㄟ^(guò)實(shí)現(xiàn)以下幾個(gè)增強(qiáng)功能來(lái)優(yōu)化它：

ZeRO 1 ：分片優(yōu)化器狀態(tài)

當(dāng)使用DDP進(jìn)行訓(xùn)練時(shí)，每個(gè)進(jìn)程都擁有優(yōu)化器狀態(tài)的完整副本。對(duì)于zer01，可以讓每個(gè)rank只保留優(yōu)化器狀態(tài)的一部分。在反向傳播期間，每個(gè)rank只需要收集與其參數(shù)相關(guān)的優(yōu)化器狀態(tài)來(lái)進(jìn)行優(yōu)化步驟。這種冗余的減少有助于節(jié)省內(nèi)存。

??在Adam的情況下，它保存的參數(shù)大約是模型大小的兩倍，將優(yōu)化器狀態(tài)分片為8個(gè)rank意味著每個(gè)rank只存儲(chǔ)總狀態(tài)大小的四分之一(2/8)。

ZeRO 2：梯度分片

除對(duì)優(yōu)化器狀態(tài)進(jìn)行分片外，還可以修改優(yōu)化器步驟來(lái)切分梯度。我們可以將所有與該rank持有的狀態(tài)相關(guān)的梯度集合起來(lái)，計(jì)算優(yōu)化步驟，然后將部分參數(shù)的優(yōu)化步驟發(fā)送給所有其他rank

現(xiàn)在每個(gè)rank不需要保存一個(gè)完整的梯度副本，這樣可以進(jìn)一步降低峰值內(nèi)存消耗。

ZeRO 3 ：模型參數(shù)分片

我么不需要在每個(gè)rank上存儲(chǔ)模型的完整副本，我們將在向前和向后期間及時(shí)獲取所需的參數(shù)。在大型模型的情況下，這些優(yōu)化可以顯著降低內(nèi)存消耗

如何使用FSDP？

其實(shí)很簡(jiǎn)單。我們所需要的就是用FSDP包裹模型：

 import torch import torch.nn as nn import torch.optim as optim from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

 model = FSDP(model)
 # it's critical to get parameters from the wrapped model # as only a portion of them returned (sharded part) optimizer = optim.Adam(model.parameters())
 # consuct training as usual train(model, optimizer)

可以指定FSDP的分片策略。例如可以選擇SHARD_GRAD_OP策略來(lái)實(shí)現(xiàn)與ZeRO2類(lèi)似的行為。

3、torch.compile

這是最簡(jiǎn)單也是最直接的優(yōu)化方式了，只要啟用torch compile，它就可以將代碼的速度提高幾個(gè)百分點(diǎn)。

在Torch2.0中增加了compile方法，他會(huì)跟蹤執(zhí)行圖，并嘗試將其編譯成一種有效的格式，以便幾乎無(wú)需Python調(diào)用即可執(zhí)行模型。

 import torch
 model = torch.compile(model)

也就是說(shuō)，2.0以后只要你的模型能用compile那么就用compile吧。

總結(jié)

本文中介紹了使用PyTorch Profiler來(lái)查找運(yùn)行瓶頸，并且介紹了一些簡(jiǎn)單的提速方法，雖然這篇文章沒(méi)有完整的解釋?zhuān)抢锩嫣峁┑姆椒ǘ际侵档民R上嘗試方法，希望對(duì)大家有所幫助。