在深入探討 PyTorch 中的 Storage 類以及其在參數 offload 場景中的應用之前，讓我們首先了解一下 PyTorch 和它的基礎組件。PyTorch 是一個廣泛使用的開源機器學習庫，它不僅提供了強大的計算圖功能和自動梯度計算，還允許開發者直接操作底層數據結構，這其中就包括 Storage。
1. 什麼是 torch.Storage?
在 PyTorch 中，Storage 是一種容納數據的一維數組，它可以看作是一個底層的內存塊，其中存儲着特定類型的數據。與 Tensor 的關係非常緊密，實際上，每個 Tensor 都有一個與之關聯的 Storage 對象。Tensor 提供了一個高維視圖來操作存儲在 Storage 中的數據。
Storage 的一個關鍵特性是它的數據排列是連續的，這使得數據可以迅速地在設備之間傳輸，例如從 CPU 到 GPU，省去了頻繁索引的操作。此外，Storage 可以存在於不同的設備上，如 CPU 或 CUDA（GPU）。
使用 storage 實現 offload 參數場景大致有如下：
-
模型訓練時的內存優化：
在深度學習模型訓練過程中，特別是當使用的模型非常大，以至於單個 GPU 顯存不足時，可以使用 offload 技術將部分數據暫時存儲到 CPU 內存中，從而釋放 GPU 顯存用於計算。 -
數據預處理：
在進行大規模數據處理時，可以將不活躍的數據段 offload 到 CPU，以保持 GPU 資源用於執行高優先級的任務。 -
長期數據存儲：
對於不需要頻繁訪問的大量數據，可以將其 offload 到 CPU 或其他存儲系統，以減少昂貴的 GPU 存儲資源的佔用。
2. 理解 Storage
2.1 簡單例子
import torch
x = torch.arange(3, dtype=torch.float32).cuda()
print(x.storage())
輸出結果如下，可以看到打印出來的結果符合預期，有三個浮點數，storage 的類型是 torch.storage.TypedStorage
0.0
1.0
2.0
[torch.storage.TypedStorage(dtype=torch.float32, device=cuda:0) of size 3]
更一般地，我們還能打印看看無類型的 storage 是什麼樣的
x_storage = x.storage()._untyped_storage
print(x_storage)
輸出結果如下，可以看到總共有 12 個整數，這是因爲前面我們使用的數據類型是 float32，也就是說每個數由 4 個字節（bytes）表示。因爲 變量 x 總共有 3 個數，所有它的 storage 總共有 12 個字節。
0
0
0
0
0
0
128
63
0
0
0
64
[torch.storage.UntypedStorage(device=cuda:0) of size 12]
這些值實際上是浮點數0、1、2在內存中的字節級表示。需要注意的是，上面輸出結果並不是隨機值，而是這些浮點數在 IEEE 754 標準下的二進制表達。我們可以逐個解釋這些值如何來的。
2.2 浮點數的 IEEE 754 表示
對於類型 float32（即單精度浮點數），每個數字佔用 4 個字節（32位），具體編碼方式爲：
- 1 位符號位（最高位）
- 8 位指數位
- 23 位尾數位
在解釋這些值之前，我們先了解一下計算機中的 小端序（Little Endian） 存儲方式：在這種存儲方式中，低位字節存放在內存的低地址端，高位字節存放在高地址端。
以Tensor[0., 1., 2.] 爲例，我們來看看這些值在內存中是如何表示的：
-
數字 0 的浮點表示：
- 符號位：0
- 指數位：全0（偏移量爲127，因此全0表示指數-127）
- 尾數位：全0
- 二進制表示：
00000000 00000000 00000000 00000000 - 十六進制表示：
00 00 00 00 - 小端序下的字節表示：
00 00 00 00 - 上面結果轉化成十進制表示：
0 0 0 0
-
數字 1 的浮點表示：
- 符號位：0
- 指數位：127（偏移後爲0，
01111111） - 尾數位：全0（因爲1.0的尾數部分無需額外存儲）
- 二進制表示：
001111111 00000000000000000000000 - 十六進制表示：
3F 80 00 00 - 小端序下的字節表示：
00 00 80 3F - 上面結果轉化成十進制表示：
0 0 128 63(80十六進制轉十進制是128，3F轉十進制是63)
-
數字 2 的浮點表示：
- 符號位：0
- 指數位：128（偏移後爲1，
10000000） - 尾數位：全0（因爲2.0的尾數部分也無需額外存儲）
- 二進制表示：
010000000 00000000000000000000000 - 十六進制表示：
40 00 00 00 - 小端序下的字節表示：
00 00 00 40
3. 使用 Storage 實現參數 offload 到 cpu
前面例子中的變量x在 cuda上，爲了實現 offload，我們需要在 cpu 上創建一個 storage，如下：
offload_storage = torch.UntypedStorage(x.nbytes).pin_memory(x.device)
print(offload_storage.device)
print(offload_storage)
輸出結果如下,可以看到offload_storage是在 cpu 上，目前其上面的值都是一些隨機值。
cpu
208
238
22
7
0
0
0
0
208
66
20
6
[torch.storage.UntypedStorage(device=cpu) of size 12]
接下來我們需要把 x offload 到 cpu 上，只需要對 storage 做 copy 操作即可，代碼如下：
offload_storage.copy_(x_storage)
print(offload_storage.device)
print(offload_storage)
輸出結果如下：
cpu
0
0
0
0
0
0
128
63
0
0
0
64
[torch.storage.UntypedStorage(device=cpu) of size 12]
可以看到x的值被成功拷貝到 cpu 上，但是這離實現 offload 還有一步之遙，我們接下來繼續看一個簡單的 offload 例子。
4. gpu 參數 和 cpu 參數互換
我們接着將探討如何利用 Storage 實現 GPU 和 CPU 之間的數據互換，這對於處理大型數據集或進行復雜的數據處理任務時尤其有用。
假設我們有以下設置：
- 一個 CUDA
Tensor用於當前計算。 - 多個 CPU
Storage用於存儲額外的數據集，這些數據集可能在不同時間被需求到 GPU。
4.1 初始化環境
首先，我們定義一個在 CUDA 上的 Tensor 和多個在 CPU 上的 Storage，準備用於數據交換：
import torch
# 定義 CUDA Tensors (用於當前計算)
current_data = torch.tensor([0.0, 1.0], device='cuda')
# 定義 CPU Storages (用於存儲額外數據)
extra_data1 = torch.FloatTensor([2.0, 3.0]).storage().pin_memory()
extra_data2 = torch.FloatTensor([4.0, 5.0]).storage().pin_memory()
extra_data3 = torch.FloatTensor([6.0, 7.0]).storage().pin_memory()
print("Initial CUDA Tensor (Current Data):")
print(current_data)
print("\nInitial CPU Storages (Extra Data):")
print("Extra Data 1:", list(extra_data1))
print("Extra Data 2:", list(extra_data2))
print("Extra Data 3:", list(extra_data3))
輸出結果爲：
Initial CUDA Tensor (Current Data):
tensor([0., 1.], device='cuda:0')
Initial CPU Storages (Extra Data):
Extra Data 1: [2.0, 3.0]
Extra Data 2: [4.0, 5.0]
Extra Data 3: [6.0, 7.0]
4.2 使用緩衝區進行數據交換
接下來，我們將根據需要將 CPU 上的數據加載到 CUDA Tensor 中，同時將當前 CUDA Tensor 的數據存儲回某個 CPU Storage，這可以申請一個 buffer 來作爲中間變量，反正數據丟失。
# 緩衝區定義
cpu_buffer = torch.FloatTensor(current_data.size()).storage().pin_memory() # CPU buffer storage
# 場景1：將 current_data 保存到 extra_data1，從 extra_data1 加載新數據到 current_data
cpu_buffer.copy_(current_data.storage()) # Save current GPU data to CPU buffer
current_data.storage().copy_(extra_data1) # Move from CUDA buffer to current_data
extra_data1.copy_(cpu_buffer) # Move from CPU buffer to extra_data1 Storage
print("\nAfter Data Exchange Scenario 1:")
print(f"Updated Current Data on {current_data.device}:", current_data)
print(f"Updated Extra Data 1 on {extra_data1.device}:", list(extra_data1))
print("Extra Data 2:", list(extra_data2))
print("Extra Data 3:", list(extra_data3))
輸出結果
After Data Exchange Scenario 1:
Updated Current Data on cuda:0: tensor([0., 1.], device='cuda:0')
Updated Extra Data 1 on cpu: [2.0, 3.0]
Extra Data 2: [4.0, 5.0]
Extra Data 3: [6.0, 7.0]
此示例清晰地展示瞭如何利用 PyTorch 的 Storage 類來有效管理內存資源，並通過使用 CPU 和 CUDA 緩衝區動態切換數據來優化應用性能。這種方法尤其適用於需要頻繁在不同計算設備之間遷移數據的場景，從而保證計算效率和響應速度。
儘管可以通過 PyTorch 的 to('cpu') 或 to('cuda') 方法簡單地在設備間遷移數據，使用 Storage 提供了更細粒度的控制。這在處理需要大量連續物理存儲空間的複雜模型時顯得尤爲重要。
例如在混合專家模型（MoE）中，系統需要根據不同的請求動態調用不同的專家（模型）。每個專家可能包含的是由多層感知機 (MLP) 或更復雜結構組成的模型，其中每層的參數在內存中通常是不連續的。這種不連續性可能導致在將參數 offload 到 CPU 或重新加載到 GPU 時，因頻繁的內存訪問和索引操作而增加通信開銷。