Paddle原始碼之記憶體管理技術

Aurelius84發表於2020-12-08
原始碼記憶體

前言

在深度學習模型訓練中，每次迭代過程中都涉及到Tensor的建立和銷燬，伴隨著的是記憶體的頻繁 malloc和free操作，可能對模型訓練帶來不必要的 overhead。

在主流的深度學習框架中，會藉助 chunk 機制的記憶體池管理技術來避免這一點。通過實事先統一申請不同 chunk size 的記憶體，並記錄到記憶體池中。建立一個Tensor時，若記憶體池中存在滿足需求的可用記憶體，則直接分配。銷燬一個Tensor時，並不馬上free掉還給系統，而是標記為可用狀態，放在記憶體池供下個Tensor使用。

通過記憶體池管理技術，可以有效減少頻繁的malloc和free操作，避免不必要的overhead。

技術實現

chunk

每個chunk代表一段連續的儲存空間。不同的chunk按照地址升序組成雙向連結串列。每個chunk只有兩種狀態：空閒、已佔用。不存在部分使用的中間態。

image-20201208200543809

在Paddle中，記憶體池統一通過 BuddyAllocator類來管理，下面逐一剖析相關實現。成員變數包括：

private:
    /*
     * 預設的記憶體分配器，支援CPUAllocator、GPUAllocator、CUDAPinnedAllocator。
     */
    std::unique_ptr<SystemAllocator> system_allocator_;

    // 用於表示一個記憶體段的資訊
    using IndexSizeAddress = std::tuple<size_t, size_t, void*>;
    // 藉助有序的 set 存放可用的記憶體段
    using PoolSet = std::set<IndexSizeAddress>;
    
    PoolSet pool_; // 記憶體池，存放可用的不同 chunk size的記憶體資訊
    PoolSet chunks_; // 記憶體池。存放從系統重新申請的記憶體塊

從BuddyAllocator的成員變數可以看出，不同BuddyAllocator物件可以管理不同型別的記憶體池，比如 CPU記憶體池、GPU記憶體池、CUDAPinned記憶體池。

建構函式顯式需要一個SystemAllocator來初始化：

public:
    BuddyAllocator(std::unqiue_ptr<SystemAllocator> system_allocator, size_t min_chunk_size, size_t max_chunk_size);

記憶體申請

BuddyAllocator如何避免記憶體頻繁的malloc和free操作呢？

申請記憶體時：

void* BuddyAllocator::Alloc(size_t unaligned_size){
  // step 1: 做記憶體對齊，保證申請的記憶體大小都是 min_chunk_size的整數倍
  size_t size = align(unaligned_size+sizeof(MemoryBlock::Desc), min_chunk_size_);
  
  // 加鎖
  std::lock_guard<std::mutex> lock(mutex_);
  
  // step 2: 如果申請記憶體超過 max_chunk_size_, 則交由system_allocator完成
  if(size > max_chunk_size_){
    return SystemAlloc(size);
  }
  
  // step 3: 否則，去記憶體池查詢是否有滿足大小的可用記憶體塊
  auto it = FindExistChunk(size);
  
  // step 4: 若找不到，則向系統申請新記憶體塊，並記錄到記憶體池中
  if(it == pool_.end()){
    it = RefillPool(size);
    if(it == pool_.end()){
      return nullptr;
    }
  }else{
    VLOG(10)<<;
  }
  
  // step 5: 更新記憶體池 size 相關資訊
  total_used_ += size;
  total_free_ -= size;
  
  // step 6: 若申請的size小於記憶體塊實際大小，則把多餘的部分切分掉，新建一個記憶體塊放到記憶體池中
  return reinterpret_cast<MemoryBlock*>(SplitToAlloc(it, size))->Data();
}

記憶體釋放

此處並非真正的將記憶體歸還給系統，而是將記憶體塊從佔用狀態標記為可用狀態，並放到記憶體池中開放出去。

void BuddyAllocator::Free(void* p){
  // step 1: 將指標轉換為記憶體塊指標
  auto block = static_cast<MemoryBlock*>(p)->MetaData();
  
  std::lock_guard<std::mutex> lock(mutex_);
  // step 2: 獲取記憶體塊的詳細元資訊，釋放記憶體需要
  auto* desc = cache_.LoadDesc(block);
  if(desc->get_type() == MemoryBlock::HUGE_CHUNK){
    // 在前面申請大記憶體時，也是交由system_allocator完成的，解鈴還須繫鈴人
    system_allocator_->Free(block, desc->get_totoal_size(), desc->get_index());
    // 刪除記憶體塊對應的元資訊
    cache_.Invalidate(block);
    return;
  }
  
  // step 3: 若待釋放記憶體塊大小在[min_chunk_size_, max_chunk_size_]之間
  block->MarkAsFree(&cache_);  // 修改元資訊，標記為 可用 狀態
  
  // step 4: 更新總記憶體資訊
  total_used_ -= desc->get_total_size();
  total_free += desc->get_total_size();
  
  // step 5: 看是否可以將此記憶體塊與左右空閒的記憶體塊合併，避免記憶體碎片
  MemoryBlock* right_buddy = block->GetRightBuddy(&cache_);
  if(right_buddy){
    auto rb_desc = cache_.LoadDesc(right_buddy);
    if(rb_desc->get_type() == MemoryBlock::FREE_CHUNK){
      pool_.erase(IndexSizedAddress(rb_desc->get_index(), rb_desc->get_total_size(), right_buddy));
      block->Merge(&cache_, right_buddy);
    }
  }
   
  MemoryBlock* left_buddy = block->GetLeftBuddy(&cache_);
  // .... (省略對前序記憶體塊的合併操作)
  
  // step 6: 將合併後的記憶體塊放入到可用記憶體池中
  pool_.insert(IndexSizeAddress(desc->get_index(), desc->get_total_size(), block));
}

記憶體歸還

此階段才是真正的將記憶體歸還給作業系統，此過程分為兩個步驟：

  • 把後來的、通過system_allocator_申請的記憶體 free掉（呼叫Release函式）
  • 析構BuddyAllocator物件時，對記憶體池剩餘的記憶體 free掉（呼叫解構函式）

我們先看第一階段 Release邏輯：

uint64_t BuddyAllocator::Release(){
  // 先加鎖
  std::lock_guard<std::mutex> lock(mutex_);
  int num = 0; // 標記後來新增申請的記憶體塊
  uint64_t bytes = 0; // 統計總共可釋放的記憶體
  bool del_flag = false;
  // step 1: 有序遍歷可用記憶體池中的每個記憶體塊
  for(auto iter = pool_.begin(); iter != pool_.end()){
    auto remain_size = std::get<1>(*iter);
    auto remain_ptr = std::get<2>(*iter);
    
    for(auto& chunk : chunks_){
      auto init_size = std::get<1>(chunk);
      auto init_ptr = std::get<2>(chunk);
      // step 2: 若在之前的chunks_記錄中找到地址一樣，空間一樣的chunk
      if(init_size = remain_size && init_ptr == remain_ptr){
        ++num;
        bytes += init_size;
        total_free_ -= init_size;
        auto block = static_cast<MemoryBlock*>(init_ptr);
        // step 3: 則歸還記憶體給系統，標記為此記憶體塊為可回收狀態
        system_allocator_->Free(init_ptr, init_size, std::get<0>(chunk));
        cache_.Invalidate(block);
        del_flag = true;
        break;
      }
    }
    // step 4: 對於標記為可回收狀態的記憶體塊，從記憶體池中移除
    if(del_flag){
      iter = pool_.erase(iter);
    }else{
      iter++;
    }
  }
  return bytes;
}

Release支援被顯式呼叫，以歸還未用到的記憶體給作業系統。

當BuddyAllocator物件在模型訓練結束後，會被析構掉。析構時需要保證之前申請的記憶體必須正確的歸還給作業系統，否則會導致記憶體洩露。

BuddyAllocator::~BuddyAllocator(){
  while(!pool.empty()){
    // step 1: 遍歷記憶體池中所有的記憶體塊
    auto block = static_cast<MemoryBlock*>(std::get<2>(pool_.begin()));
    auto desc = cache_.LoadDesc(block);
    // step 2: Free掉，歸還給系統
    system_allocator_->Free(block, desc->get_total_size(), desc->get_index());
    // step 3: 刪除元資訊
    cache_.Invalidata(block);
    pool_.erase(pool_.begin());
  }
}

參考資料

  1. Paddle 框架原始碼

相關文章