使用PyTorch Profiler進行模型性能分析，改善并加速PyTorch訓練

    

     

      

       

        來源：DeepHub IMBA
       
      
     
    
    

     

      

       

        

         
本文約3100字，建議閱讀6分鐘
         
本文中介紹了使用PyTorch Profiler來查找運行瓶頸和一些簡單的提速方法。
        
       
      
     
    

如果所有機器學習工程師都想要一樣東西，那就是更快的模型訓練——也許在良好的測試指標之后。

加速機器學習模型訓練是所有機器學習工程師想要的一件事。更快的訓練等于更快的實驗，更快的產(chǎn)品迭代，還有最重要的一點需要更少的資源，也就是更省錢。

熟悉PyTorch Profiler

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

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

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

 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)

然后就可以啟動tensorboard查看分析軌跡。如果這一步有問題，請查看是否安裝了torch-tb-profiler。

Profiler有很多不同的選項，但最重要的是activities和profile_memory，一般情況下我們只需要這兩個選項，因為啟用的選項越少，開銷就越小。

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

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

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

或者增加更精細的自定義的標簽，這里的名稱將在跟蹤中可見，我們就可以更簡單的追蹤想要的東西了。

 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顯示是這樣的。

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

反向傳播由PyTorch在一個單獨的線程中處理(上圖中的線程16893)，因此很容易識別，這部分門也控制不了，因為都是Pytorch根據(jù)我們的計算來自動進行的。（當然也可以自定義反向傳播，但是這過于復(fù)雜，一般不建議自己實現(xiàn)）

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

這是因為在數(shù)據(jù)加載過程中，GPU什么也不做，這會使可用資源利用率不足。并且在Pytorch的訓練時數(shù)據(jù)處理可以與GPU計算重疊，因為它們是獨立的部分，也就是說我們加載一個批次的時間只要與一個前向和一個反向傳播的時間相近就可以了，這樣就可以最大化的利用GPU的資源。

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

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

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

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

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

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

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

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

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

 PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"

這告訴PyTorch分配器分配可以在將來擴展的塊。但是，如果大小變化太大，它仍然可能無法解決問題。

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

比如最簡單的將數(shù)據(jù)填充到相同的大小。或者可以通過運行具有最大輸入大小的模型來預(yù)熱分配器。

內(nèi)存歷史記錄

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

為了調(diào)試它，我們可以查看分配器的內(nèi)存歷史記錄。它可以通過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)

可視化會畫出這樣的東西:

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

你可能會注意到狹窄的尖峰，這些是持續(xù)時間很短的張量，并且占據(jù)了很多空間。通過點擊一個張量，可以得到這個張量被分配到哪里的信息。我們希望的就是最小化這些峰值，因為它們限制了有效的內(nèi)存使用。檢查導致這個峰值的原因，并考慮優(yōu)化或者使用其他計算方法替代。

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

第一次運行之后的一些數(shù)據(jù)沒有被清除，所以導致內(nèi)存占用過高。通過點擊塊，可以知道這些張量是從哪里來的。在圖像中，梯度在訓練步驟之后沒有被清除，因此它們在向前傳遞過程中處于無用狀態(tài)，占用了寶貴的內(nèi)存。

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

我們知道了原因，并且可以通過Profiler來找到瓶頸，那么我們可以通過什么方法來加速訓練呢？

1、FlashAttention

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

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

當然也有其他的庫可以替換，例如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ù)冗余

如果使用多個gpu來運行訓練，基本的解決方案是使用DistributedDataParallel。生成了幾個相同的進程，并且在反向傳播期間聚合梯度。

當我們生成相同的進程時，在每個GPU上都有相同的模型和優(yōu)化器狀態(tài)，這是冗余的。可以通過跨數(shù)據(jù)分片來優(yōu)化內(nèi)存使用。

當在多個gpu上進行訓練時，每個進程在使用DDP進行訓練時都有相同數(shù)據(jù)的精確副本。可以通過實現(xiàn)以下幾個增強功能來優(yōu)化它：

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

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

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

ZeRO 2：梯度分片

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

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

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

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

如何使用FSDP？

其實很簡單。我們所需要的就是用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策略來實現(xiàn)與ZeRO2類似的行為。

3、torch.compile

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

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

 import torch
 model = torch.compile(model)

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

總結(jié)

本文中介紹了使用PyTorch Profiler來查找運行瓶頸，并且介紹了一些簡單的提速方法，雖然這篇文章沒有完整的解釋，但是里面提供的方法都是值得馬上嘗試方法，希望對大家有所幫助。