作者｜鄭建華

(資料圖片)

1?

不同Tensor類型的存儲(chǔ)管理方式

Lazy Tensor 的存儲(chǔ)是由 Runtime 和 Actor 等對象管理的。靜態(tài)圖完成編譯后，需要多少個(gè)對象、多少存儲(chǔ)空間都是確定的，Runtime 等在初始化時(shí)會(huì)分配存儲(chǔ)，在退出時(shí)回收資源。

Eager 模式下，Global Tensor 可以視為對 Local Tensor 的分布式封裝，EagerGlobalTensorImpl?在本地的數(shù)據(jù)是一個(gè)

EagerLocalTensorImpl?對象?？梢酝ㄟ^考察 EagerLocalTensorImpl 來理解 eager 模式下 tensor 的存儲(chǔ)管理。

參考的示例代碼如下：

import numpy as np import oneflow as flow a = np.random.randn(1, 4) flow.tensor(a, device=flow.device("cpu"), dtype=flow.float)

2?

Tensor 存儲(chǔ)相關(guān)類的關(guān)系

EagerLocalTensorImpl 的存儲(chǔ)相關(guān)的類關(guān)系如下。

后續(xù)會(huì)順著示例代碼的執(zhí)行過程，看看圖中的對象都是在何時(shí)、如何構(gòu)造的，存儲(chǔ)被誰持有、如何分配并釋放。

3

通過虛擬機(jī)指令為 Tensor 分配存儲(chǔ)

tensor 的構(gòu)造函數(shù)通過 Python C API 注冊為 PyTensorObject_init，由 functional::_legacy_tensor_ctor?

根據(jù)簽名進(jìn)行轉(zhuǎn)發(fā)。

示例代碼對應(yīng)的是 TensorWithDataFunctor

，調(diào)用 MakeLocalTensorFromData?構(gòu)造 tensor，在這個(gè)函數(shù)中通過調(diào)用 functional::Empty以及 EmptyFunctor分配存儲(chǔ)。在 EmptyFunctor 中把相關(guān)屬性都存到 attrs，然后調(diào)用 OpInterpUtil::Dispatch在 vm 指令的執(zhí)行準(zhǔn)備過程中分配存儲(chǔ)。

EmptyFunctor 返回的 tensor 是一個(gè)只有存儲(chǔ)空間、不含數(shù)據(jù)的對象。數(shù)據(jù)拷貝在后面由 CopyLocalTensorFromUntypedArray

完成。

3.1 存儲(chǔ)相關(guān)對象的構(gòu)造

因?yàn)槭?eager 模式下的 local tensor，OpInterpUtil::Dispatch 會(huì)被轉(zhuǎn)發(fā)到 NaiveInterpret執(zhí)行。對于示例代碼，這個(gè)函數(shù)的輸入?yún)?shù)如下：

inputs 是一個(gè)空數(shù)組

outputs 只有一個(gè)元素、且是空指針

因?yàn)?outputs 中的 tensor 指針都是空的，所以需要?jiǎng)?chuàng)建一個(gè) EagerLocalTensorImpl?對象，其 one::TensorStorage 成員變量是空指針。

因?yàn)?output_eager_blob_objects 中的元素尚未初始化，會(huì)調(diào)用 tensor_impl->InitEagerBlobObject?

進(jìn)行初始化。因?yàn)?tensor_storage_ 還是空的，這個(gè)過程會(huì)執(zhí)行如下操作：

創(chuàng)建 vm::TensorStorage 對象

創(chuàng)建 EagerBlobObject 對象

set_eager_blob_object

UpdateTensorStorage

創(chuàng)建 one::TensorStorage 對象

設(shè)置 tensor 存儲(chǔ)釋放的回調(diào)函數(shù)

上述對象的創(chuàng)建，都只是記錄相關(guān)信息，還不涉及 tensor 的存儲(chǔ)分配。

需要注意的是，注冊到 one::TensorStorage 的回調(diào)函數(shù)被賦值給了成員變量 releaser_hook_，這個(gè)函數(shù)會(huì)通過虛擬機(jī)指令釋放 tensor。

3.2 在指令執(zhí)行過程中分配 tensor 存儲(chǔ)

分配 tensor 存儲(chǔ)的過程如下：

vm::Instruction::Compute

vm::InstructionPolicy::ComputeIf

vm::OpCallInstructionPolicy::Compute

OpCallInstructionUtil::Compute

獲取內(nèi)存分配器

OpCallInstructionUtil::AllocateOutputBlobsMemory

blob_object->TryAllocateBlobBodyMemory

allocator->Allocate

在 EagerBlobObject::TryAllocateBlobBodyMemory 中，allocator 分配的存儲(chǔ)地址會(huì)賦值給 dptr，存儲(chǔ)地址 dptr 和 Free 函數(shù)一起構(gòu)造一個(gè)智能指針，并賦值給 vm::TensorStorage 的 blob_dptr_ 變量。

4?

通過虛擬機(jī)指令釋放 Tensor 存儲(chǔ)

在前面的 3.1 節(jié)提到，EagerLocalTensorImpl 在初始化 EagerBlobObject、創(chuàng)建 one::TensorStorage 的同時(shí)，會(huì)設(shè)置一個(gè)釋放 tensor 的回調(diào)函數(shù)，回調(diào)函數(shù)保存在變量 releaser_hook_ 中

，one::TensorStorage 析構(gòu)時(shí)調(diào)用這個(gè)回調(diào)函數(shù)。把這些信息綜合整理一下，one::TensorStorage 析構(gòu)時(shí)會(huì)執(zhí)行如下操作：

vm::InstructionList instruction_list; InstructionsBuilder instructions_builder(&instruction_list); // JUST(Build(&instructions_builder)); if (eager_blob_object->producer_stream().has_value()) { JUST(instructions_builder->ReleaseTensor(eager_blob_object)); } JUST(vm::Run(instructions_builder.mut_instruction_list()));

在 InstructionsBuilder::ReleaseTensor 中，如果有其它 stream 最近使用了 eager_blob_object，會(huì)通過 SoftSyncStreamBetween 進(jìn)行同步。通過這種方式解決存儲(chǔ)的依賴問題。

一般情況下，通過 tensor 的 producer_stream 釋放存儲(chǔ)，根據(jù)這個(gè)對象獲取對應(yīng)的 vm::Stream 對象，并據(jù)此構(gòu)造指令 instruction（包含 eager_blob_object 和 vm_stream），示例代碼對應(yīng)的指令類型是 FastReleaseTensorInstructionPolicy，其 Compute 方法執(zhí)行具體的存儲(chǔ)釋放邏輯，過程如下：

ReleaseTensorInstructionPolicy::Release()

eager_blob_object->DeallocateBlobDataPtr()

tensor_storage_->Release()

tensor_storage_->_Release()

blob_dptr_.reset()

智能指針重置，調(diào)用分配存儲(chǔ)時(shí)指定的 Free 方法

5

reshape 等場景的存儲(chǔ)管理

在 reshape、slice、transpose 等場景中，調(diào)用的 EagerLocalTensorImpl 構(gòu)造函數(shù)的參數(shù)包括 input 的 tensor_storage，所以這個(gè) tensor 的 tensor_storage_ 變量不是空的，在執(zhí)行 InitEagerBlobObject 時(shí)，只創(chuàng)建 EagerBlobObject以提供 shape、stride等信息；但不會(huì)再創(chuàng)建 one::TensorStorage，而是復(fù)用 input 的存儲(chǔ)。

6?

兩個(gè) TensorStorage 類型可以合并嗎？

為什么在 one::TensorStorage 析構(gòu)時(shí)、由它保存的回調(diào)函數(shù)來觸發(fā)釋放 vm::TensorStorage 中的存儲(chǔ)呢？

one::TensorStorage 只多了一個(gè) releaser，這兩個(gè) Storage 類型是否可以合并呢？

在當(dāng)前的設(shè)計(jì)下，這兩個(gè)類型不能合并。因?yàn)?one::TensorStorage::releaser_hook_ 中持有 EagerBlobObject 的智能指針，EagerBlobObject 中也持有 vm::TensorStorage 的智能指針。如果兩個(gè) Storage 類型合并為一個(gè)，就會(huì)出現(xiàn)循環(huán)引用、對象無法析構(gòu)而導(dǎo)致內(nèi)存泄漏。

所以，vm::TensorStorage 只是單純的存儲(chǔ)，可以在多個(gè) tensor 之間共享。EagerBlobObject 既包括存儲(chǔ)、也包括 shape、stride、data_type 等獨(dú)特的對象信息。而 one::TensorStorage 是為了避免循環(huán)引用而引入的、專門負(fù)責(zé)釋放存儲(chǔ)的角色。

7

附錄

GDB 斷點(diǎn)示例

break oneflow::one::MakeLocalTensorFromData break oneflow::one::NaiveInterpret break oneflow::vm::VirtualMachineEngine::DispatchInstruction break oneflow::vm::OpCallInstructionUtil::Compute break oneflow::vm::OpCallInstructionUtil::AllocateOutputBlobsMemory break oneflow::vm::EagerBlobObject::TryAllocateBlobBodyMemory break oneflow::vm::ReleaseTensorInstructionPolicy::Release break oneflow/core/eager/eager_blob_object.cpp:107

參考資料

OneFlow（https://github.com/Oneflow-Inc/oneflow/tree/b51cb72430619f6088e47bbb8b8226f37299573a）

OneFlow源碼解析：Tensor類型體系與Local Tensor

其他人都在看

“ChatGPT們”的淘金時(shí)代

推演語言模型的大小與計(jì)算開銷

狂追ChatGPT：開源社區(qū)的“平替”熱潮

谷歌科學(xué)家：ChatGPT秘密武器的演進(jìn)與局限

比快更快，開源Stable Diffusion刷新作圖速度

OneEmbedding:單卡訓(xùn)練TB級推薦模型不是夢

GLM訓(xùn)練加速：性能最高提升3倍，顯存節(jié)省1/3

歡迎Star、試用OneFlow: github.com/Oneflow-Inc/oneflow/http://github.com/Oneflow-Inc/oneflow/

關(guān)鍵詞：