1. 管道并行

DeepSpeed v0.3增(zeng)加了(le)對管道并(bing)行的(de)新(xin)支持。管道并(bing)行將(jiang)模(mo)型(xing)的(de)層劃(hua)分(fen)為(wei)階段，可(ke)以并(bing)行處理，從而(er)提(ti)(ti)高深度學習訓練的(de)內(nei)存和(he)計算效率。DeepSpeed的(de)訓練引擎提(ti)(ti)供了(le)**混合(he)數據和(he)管道并(bing)行**，并(bing)可(ke)進一步與**模(mo)型(xing)并(bing)行**（如Megatron-LM）結合(he)使用

DeepSpeed使用(yong)梯(ti)度(du)(du)累積(ji)來提(ti)取管(guan)道并行(xing)(xing)性。每(mei)個(ge)(ge)訓練數據批(pi)次(ci)(ci)被(bei)劃分為(wei)可以由管(guan)道階段(duan)并行(xing)(xing)處(chu)理的微(wei)批(pi)次(ci)(ci)。一旦一個(ge)(ge)階段(duan)完成(cheng)(cheng)微(wei)批(pi)次(ci)(ci)的正向(xiang)傳遞(di)(di)(di)(di)，激活內存就(jiu)會(hui)傳遞(di)(di)(di)(di)到管(guan)道中的下一個(ge)(ge)階段(duan)。同樣(yang)地，當下一個(ge)(ge)階段(duan)在微(wei)批(pi)次(ci)(ci)上完成(cheng)(cheng)反向(xiang)傳遞(di)(di)(di)(di)時，關于激活的梯(ti)度(du)(du)就(jiu)會(hui)通過管(guan)道向(xiang)后傳遞(di)(di)(di)(di)。每(mei)個(ge)(ge)反向(xiang)傳遞(di)(di)(di)(di)都會(hui)本地累積(ji)梯(ti)度(du)(du)。接下來，所有(you)數據并行(xing)(xing)組都會(hui)并行(xing)(xing)執(zhi)行(xing)(xing)梯(ti)度(du)(du)的縮減(jian)操作。最后，優化器會(hui)更新模(mo)型權重。

2.Pipline 并行開始

DeepSpeed致力于加速和簡化(hua)管(guan)道并(bing)行訓練(lian)過程。本節(jie)提供了(le)使用混合數據(ju)和管(guan)道并(bing)行訓練(lian)torchvision的(de)AlexNet模型(xing)的(de)第一(yi)步。

2.1 表示管道模型

管(guan)(guan)道并行需要(yao)將模型(xing)表示(shi)為(wei)層序(xu)列(lie)。在正向(xiang)傳遞中，每個層都使用前一層的輸出。事實上(shang)，對于管(guan)(guan)道并行模型(xing)，不需要(yao)指定(ding)forward（）！管(guan)(guan)道并行模型(xing)的正向(xiang)傳遞隱式采用以下形式：

def forward(self, inputs):
    x = inputs
    for layer in self.layers:
        x = layer(x)
    return x

PyTorch的(de)`torch.nn.Sequential`是(shi)表達管(guan)道(dao)并(bing)行(xing)模型的(de)方便容器，并(bing)且可(ke)以在無需修改(gai)的(de)情況下由DeepSpeed并(bing)行(xing)化。

net = nn.Sequential(
    nn.Linear(in_features, hidden_dim),
    nn.ReLU(inplace=True),
    nn.Linear(hidden_dim, out_features)
)
from deepspeed.pipe import PipelineModule
net = PipelineModule(layers=net, num_stages=2)

`PipelineModule`使(shi)用(yong)其`layers`參數作為(wei)構成(cheng)模型的層序(xu)列。初始化后(hou)，net被(bei)劃分(fen)為(wei)兩個(ge)管道階段，其層被(bei)移動到相應的GPU上。如果存(cun)在多(duo)于(yu)兩個(ge)GPU，則DeepSpeed還將使(shi)用(yong)**混合數據并(bing)行(xing)。**

Note:The total number of GPUs must be divisible by the number of pipeline stages.

2.2 輸入和輸出

在`torch.nn.Sequential`之后，每個層的(de)輸(shu)入和(he)輸(shu)出必須(xu)是單個`torch.Tensor`或一個元組(zu)的(de)張量(liang)。實際上，某(mou)些模型(xing)可能需要修(xiu)改(gai)它們的(de)正(zheng)向傳(chuan)遞以打(da)包和(he)解(jie)包`forward()`的(de)參數(shu)。考慮一下Transformer block的(de)簡化實現：

class TransformerBlock(nn.Module)
    ...
    def forward(self, hidden, mask):
        output = self.compute(hidden, mask)
        return output
    ...

stack = [ TransformerBlock() for _ in range(num_layers) ]

需(xu)要對(dui)TransformerBlock進行兩個(ge)修(xiu)改:

1. 參數必須被收集到一個元組中。
2. 還必須(xu)從`forward()`返回(hui)`mask`以(yi)傳遞給下一層。

這些修改可(ke)以通過一個短的子類(lei)來完(wan)成：

class TransformerBlockPipe(TransformerBlock)
    def forward(self, inputs):
        hidden, mask = inputs
        output = super().forward(hidden, mask)
        return (output, mask)
stack = [ TransformerBlockPipe() for _ in range(num_layers) ]

2.3 訓練循環

管道并(bing)(bing)行交替執(zhi)行前向傳遞和(he)反(fan)(fan)向傳遞，因此訓練循環不能分為`forward()`、`backward`和(he)`step()`等單(dan)獨(du)的階(jie)段(duan)。相反(fan)(fan)，DeepSpeed的管道引(yin)擎(qing)提供了`train_batch()`方法(fa)，該(gai)方法(fa)推進管道引(yin)擎(qing)直(zhi)到消耗(hao)下(xia)一批訓練數據并(bing)(bing)更新模型(xing)權重(zhong)。 

2.4 數據處理

通常(chang)，數據并(bing)行(xing)訓練會(hui)使每(mei)個工作器在(zai)每(mei)個批次開始時獨(du)立執行(xing)IO。但是，在(zai)管道并(bing)行(xing)環(huan)境中，只有(you)第(di)一(yi)個階段使用輸入數據，只有(you)最后一(yi)個階段使用標簽進行(xing)損失計算。

`注意：管道引擎希望數據加載器返(fan)回兩個(ge)(ge)項(xiang)目的元組(zu)。第一個(ge)(ge)返(fan)回項(xiang)是輸(shu)入(ru)批次數據，第二(er)個(ge)(ge)項(xiang)目是用(yong)于計算(suan)損失的數據。與(yu)之(zhi)前一樣，輸(shu)入(ru)和標(biao)簽應(ying)為(wei)torch.Tensor類型或張(zhang)量的元組(zu)。`

出于(yu)方便起見，DeepSpeed管(guan)(guan)道(dao)引擎可以在向`deepspeed.initialize()`提供(gong)數據集時構建分(fen)布式數據加(jia)載(zai)器(qi)。DeepSpeed處理了數據加(jia)載(zai)的其余(yu)復雜性(xing)，因此(ci)管(guan)(guan)道(dao)訓(xun)練循環變為：

engine, _, _, _ = deepspeed.initialize(
    args=args,
    model=net,
    model_parameters=[p for p in net.parameters() if p.requires_grad],
    training_data=cifar_trainset())

for step in range(args.steps):
    loss = engine.train_batch()

當(dang)然，DeepSpeed將與(yu)您希(xi)望使用的任何數據(ju)加載器一起工作(zuo)。數據(ju)加載器應由(you)管道中(zhong)的第(di)一個和最后(hou)一個階段構(gou)建。每個工作(zuo)器應加載`engine.train_micro_batch_size_per_gpu()`大小的微批次，并且每個`train_batch()`將查(cha)詢(xun)`engine.gradient_accumulation_steps()`次。

`請注意！管道引擎從迭代器(qi)中(zhong)提取(qu)數(shu)據(ju)而不是對其進行迭代。在訓練批(pi)次(ci)中(zhong)間(jian)，數(shu)據(ju)流不能為空。每次(ci)調用train_batch()將從數(shu)據(ju)迭代器(qi)中(zhong)拉取(qu)engine.gradient_accumulation_steps()個微批(pi)次(ci)數(shu)據(ju)。`

DeepSpeed提供(gong)了一個方便(bian)的(de)類`deepspeed.utils.RepeatingLoader`，它(ta)簡單地(di)包裝可迭(die)代對象（例如數據(ju)加載器(qi)），并在到達結尾(wei)時(shi)重(zhong)新啟動它(ta)：

train_loader = deepspeed.utils.RepeatingLoader(train_loader)
train_iter = iter(train_loader)
for step in range(args.steps):
    loss = engine.train_batch(data_iter=trainiter)

3. 進一步建議

3.1 管道模塊的負載平衡

管道并行(xing)訓(xun)練的性能強烈依賴(lai)于(yu)負載均(jun)衡。DeepSpeed提(ti)供(gong)了幾種在GPU之間分配模型的機(ji)制。可以使用`PipelineModule`的`partition_method`關鍵字參(can)數設置(zhi)這些策(ce)略。以下是DeepSpeed當前提(ti)供(gong)的分區方法：

- `partition_method="parameters" `(默認)將每個管道階段上可訓練參數的數量平衡起來。這在內存受限的環境和當層的大小與計算時間成比例時特別有用。 
- `partition_method="type:[regex]"`平衡類名與[regex]匹配的層。正則表達式不區分大小寫。例如，`partition_method="type:transformer"`將平衡每個階段中transformer層的數量。 
- `partition_method="uniform" `平(ping)衡每(mei)個階段中的層數。

3.2 內存高效的模型構建

將一(yi)個(ge)(ge) `Sequential `容器構(gou)建并(bing)提供給(gei) `PipelineModule` 是(shi)指定管(guan)道并(bing)行(xing)模(mo)型(xing)的(de)一(yi)種方便的(de)方式(shi)。然而，對于大(da)型(xing)模(mo)型(xing)，這(zhe)種方法會遇(yu)到可擴展性(xing)問題，因(yin)為每個(ge)(ge) worker 都會在 CPU 內存中(zhong)復制整(zheng)個(ge)(ge)模(mo)型(xing)。例(li)如(ru)，一(yi)個(ge)(ge)具有 16 個(ge)(ge) GPU 的(de)機器必(bi)須擁有與模(mo)型(xing)大(da)小的(de) 16 倍相同的(de)本(ben)地(di) CPU 內存。

DeepSpeed 提(ti)供了一(yi)個(ge)` LayerSpec `類，它延遲模(mo)塊的(de)構建，直到模(mo)型層(ceng)次(ci)已經(jing)被分(fen)配(pei)(pei)給工作節點。然(ran)后每個(ge) worker 只(zhi)分(fen)配(pei)(pei)它所分(fen)配(pei)(pei)的(de)層(ceng)次(ci)。因此，與前(qian)一(yi)段中(zhong)的(de)示例進行(xing)比較(jiao)，使(shi)用(yong) `LayerSpec`，具有 16 個(ge) GPU 的(de)機器將需要在其 CPU 內存(cun)上(shang)分(fen)配(pei)(pei)總共 1x 模(mo)型大小(xiao)而不(bu)是 16 倍。

某些模型(xing)不(bu)能完全表示為(wei)管(guan)道(dao)(dao)并行(xing)模型(xing)，因(yin)為(wei)一(yi)些層在(zai)(zai)(zai)管(guan)道(dao)(dao)內被重用。例如(ru)，基于 Transformer 的(de)語言模型(xing)通常在(zai)(zai)(zai)管(guan)道(dao)(dao)早期使用嵌入(ru)層將詞匯(hui)映射到隱藏(zang)狀(zhuang)態，然后在(zai)(zai)(zai)管(guan)道(dao)(dao)末(mo)尾使用嵌入(ru)將隱藏(zang)狀(zhuang)態映射回詞匯(hui)。如(ru)果該(gai)模型(xing)受到純管(guan)道(dao)(dao)并行(xing)性的(de)限制，則這種嵌入(ru)重用將禁止管(guan)道(dao)(dao)并行(xing)性。

3.3 Tied Layers

某(mou)些(xie)模(mo)型(xing)不能(neng)完全表(biao)示為管(guan)(guan)(guan)道(dao)(dao)并(bing)(bing)行(xing)模(mo)型(xing)，因為一些(xie)層在(zai)(zai)(zai)管(guan)(guan)(guan)道(dao)(dao)內被重(zhong)用。例如，基于 Transformer 的(de)語言模(mo)型(xing)通常在(zai)(zai)(zai)管(guan)(guan)(guan)道(dao)(dao)早期(qi)使用嵌(qian)入層將詞匯(hui)映射到隱藏(zang)狀態，然后(hou)在(zai)(zai)(zai)管(guan)(guan)(guan)道(dao)(dao)末尾使用嵌(qian)入將隱藏(zang)狀態映射回詞匯(hui)。如果(guo)該(gai)模(mo)型(xing)受到純管(guan)(guan)(guan)道(dao)(dao)并(bing)(bing)行(xing)性(xing)的(de)限(xian)制(zhi)，則這種嵌(qian)入重(zhong)用將禁止(zhi)管(guan)(guan)(guan)道(dao)(dao)并(bing)(bing)行(xing)性(xing)。

DeepSpeed 提供了(le)` TiedLayerSpec`，它是 `LayerSpec` 的擴(kuo)展。`TiedLayerSpec `需要(yao)一個額外的參數：`key`。每次重(zhong)用(yong)層(ceng)都使用(yong) `TiedLayerSpec `進行指定(ding)，`key `字段用(yong)于(yu)識別層(ceng)何處被(bei)重(zhong)用(yong)。

被捆綁(bang)的(de)(de)層會在(zai)(zai)每個擁(yong)有“重用”實例的(de)(de)管道(dao)階(jie)段中復制。然后訓練(lian)將像平常一樣(yang)進行，但在(zai)(zai)所有反(fan)向(xiang)傳遞完成(cheng)后，會添(tian)加一個額外的(de)(de)捆綁(bang)梯(ti)度的(de)(de)全(quan)局歸(gui)約操作。這個全(quan)局歸(gui)約確保了(le)捆綁(bang)層的(de)(de)權重在(zai)(zai)管道(dao)階(jie)段之(zhi)間保持(chi)同步。