大型模型訓練 GPU 記憶體最佳化技術

# 梯度累積範例程式碼
def train_with_gradient_accumulation(model, device, batch_size, gradient_accumulation_steps):
    # 初始化梯度累積步數
    gradient_accumulation_steps = gradient_accumulation_steps
    
    # 遍歷資料集
    for batch_idx, batch in enumerate(dataset):
        # 將資料載入GPU
        inputs = batch.to(device)
        
        # 前向傳播
        outputs = model(inputs)
        loss = loss_fn(outputs, labels)
        
        # 反向傳播
        loss.backward()
        
        # 梯度累積
        if (batch_idx + 1) % gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()

#### 內容解密：
1. `train_with_gradient_accumulation` 函式接受模型、裝置、批次大小和梯度累積步數作為引數。
2. 在遍歷資料集的過程中，將資料載入GPU並進行前向傳播和反向傳播。
3. 使用 `loss.backward()` 計算梯度。
4. 每隔 `gradient_accumulation_steps` 個批次，就執行一次梯度更新和梯度清零操作。

### GPU與CPU之間的記憶體交換

記憶體交換方法非常簡單：它在CPU和GPU之間來回複製啟動值。如果你不熟悉深度學習術語，可以將啟動值視為神經網路每個節點的計算輸出。這個想法是隻保留當前計算步驟所需的資料在GPU中，並將計算結果交換到CPU記憶體中供未來步驟使用。

在此基礎上，一種新的繼電器式執行技術稱為L2L（層到層），它只保留執行層和傳輸緩衝區在GPU上。整個模型和最佳化器（持有狀態）都儲存在CPU空間中。L2L可以大大提高GPU吞吐量，並允許我們在經濟實惠的裝置上開發大型模型。

```python
# L2L範例程式碼（概念性，非實際實作）
def l2l_execution(model, device):
    # 將模型和最佳化器儲存在CPU中
    model_cpu = model.to('cpu')
    optimizer_cpu = optimizer.state_dict()
    
    # 只將執行層載入GPU
    executing_layers = get_executing_layers(model, current_step)
    executing_layers_gpu = executing_layers.to(device)
    
    # 進行前向傳播和反向傳播
    outputs = executing_layers_gpu(inputs)
    loss = loss_fn(outputs, labels)
    loss.backward()
    
    # 更新模型引數
    update_model_parameters(model_cpu, executing_layers_gpu.gradients)

#### 內容解密：
1. `l2l_execution` 函式將模型和最佳化器儲存在CPU中。
2. 只將當前執行層載入GPU。
3. 在GPU上進行前向傳播和反向傳播。
4. 將梯度更新回CPU中的模型引數。

### 管道模型平行

在第4.2節中，我們討論了最常用的分散式訓練方法：資料平行。這種方法在每個裝置上保留整個模型的副本，並將資料分配到多個裝置上。然後，它彙總梯度並在每個訓練步驟中更新模型。只要整個模型可以載入一個GPU，資料平行的方法就非常有效。然而，在本文中，我們將看到，並非總是能夠做到這一點。這時，管道平行就可以派上用場。

#### 模型平行

模型平行的想法是將神經網路分成較小的子網路，並在不同的GPU上執行每個子網路。圖4.6展示了模型平行的方法。

```python
# 模型平行範例程式碼（虛擬PyTorch程式碼）
class ModelParallel(nn.Module):
    def __init__(self):
        super(ModelParallel, self).__init__()
        self.subnet1 = Subnet1().to('cuda:0')
        self.subnet2 = Subnet2().to('cuda:1')
    
    def forward(self, x):
        x = self.subnet1(x.to('cuda:0'))
        x = self.subnet2(x.to('cuda:1'))
        return x

#### 內容解密：
1. `ModelParallel` 類別繼承自 `nn.Module`。
2. 將子網路分配到不同的GPU上。
3. 在 `forward` 方法中，將輸入資料傳遞到不同的子網路，並在不同的GPU上進行計算。

## 4.4 訓練無法載入單一GPU的大型模型

在處理大型深度學習模型時，單一GPU的記憶體限制成為一大挑戰。為瞭解決這個問題，我們可以採用模型平行（Model Parallelism）或管線平行（Pipeline Parallelism）的方法。

### 模型平行

模型平行是一種將大型模型分割成多個子網路，並將這些子網路分配到不同的GPU上進行訓練的方法。以下是一個簡單的PyTorch範例，展示如何實作模型平行：

```python
gpu1 = 0
gpu2 = 1

class LargeModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 初始化網路為兩個子網路
        self.subnet1 = ...
        self.subnet2 = ...
        # 將子網路1和2放到不同的GPU上
        self.subnet1.cuda(gpu1)
        self.subnet2.cuda(gpu2)

    def forward(self, x):
        # 將資料載入GPU 1並計運算元網路1的輸出，此時GPU 2閒置
        x = x.cuda(gpu1)
        x = self.subnet1(x)
        # 將子網路1的輸出移到GPU 2並計運算元網路2的輸出，此時GPU 1閒置
        x = x.cuda(gpu2)
        x = self.subnet2(x)
        return x

# 使用PyTorch GPipe實作管線平行訓練Transformer模型的範例程式碼
from torch.distributed.pipeline.sync import Pipe

# 定義模型
class TransformerModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 初始化Transformer模型
        self.layers = nn.Sequential(
            # ...
        )

    def forward(self, x):
        return self.layers(x)

# 將模型分割到不同的GPU上
model = TransformerModel()
model = Pipe(model, chunks=4)  # 將輸入批次劃分為4個微批次

# 進行訓練
for batch in dataset:
    input_ids = batch['input_ids'].cuda(0)
    labels = batch['labels'].cuda(0)
    output = model(input_ids)
    loss = criterion(output, labels)
    loss.backward()
    optimizer.step()

rpc.init_rpc(
    name="worker",
    rank=0,
    world_size=1,
    # 其他初始化引數
)

num_gpus = 2
partition_len = ((nlayers - 1) // num_gpus) + 1
for i in range(nlayers):
    transformer_block = TransformerEncoderLayer(emsize, nhead, nhid, dropout)
    # 載入到不同的GPU
    device = i // (partition_len)
    tmp_list.append(transformer_block.to(device))

chunks = 8  # 設定微批次數量
model = Pipe(torch.nn.Sequential(*module_list), chunks=chunks)

def train():
    model.train()
    for batch, i in enumerate(range(0, nbatches, bptt)):
        data, targets = get_batch(train_data, i)
        optimizer.zero_grad()
        output = model(data).local_value()
        loss = criterion(output.view(-1, ntokens), targets.cuda(1))
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optimizer.step()