混合式分片資料平行處理

分割資料平行程度是一種節省記憶體的分散式訓練技術，可將模型狀態 (模型參數、漸層和最佳化器狀態) 分割到不同裝置之間。這有助於您使用釋放的 GPU 記憶體配合較大型號或增加批次大小。SMP 程式庫提供使用 PyTorch 完全分片資料並行 (FSDP) 來執行分片資料平行處理的功能。 PyTorch FSDP 默認情況下，在所使用的整個 GPU 集中進行碎片。在 SMP v2 中，該庫通過擴展 PyTorch 混合分片（HYBRID_SHARD）來提供 PyTorch FSDP 之上的分片數據並行性，這是 FSDP 提供的分片策略之一：，，，，。 PyTorch FULL_SHARD SHARD_GRAD_OP HYBRID_SHARD _HYBRID_SHARD_ZERO2以這種方式擴展混合分片有助於實現 scale-aware-sharding 如 FSDP 的巨型模型訓練的近線性縮放博客中所述。 AWS PyTorch

SMP 程式庫可讓您_HYBRID_SHARD_ZERO2在任何可設定數量的 GPU 上輕鬆使用，HYBRID_SHARD並延伸支援跨單一節點 (HYBRID_SHARD) 或所有 GPU () 分割的原生 PyTorch FSDP。FULL_SHARD PyTorch FSDP 呼叫可以保持原狀，而且您只需要將hybrid_shard_degree引數新增至 SMP 組態，如下列程式碼範例所示。您不需要在模型周圍變更 PyTorch FSDP 包裝函式中的sharding_strategy引數值。 PyTorch 您可以ShardingStrategy.HYBRID_SHARD作為值傳遞。或者，SMP 程式庫會覆寫指令碼中的策略，ShardingStrategy.HYBRID_SHARD如果您指定的值等於或大於 2 的hybrid_shard_degree參數，則會將其設定為。

下列程式碼片段說明如何在遵循中介紹的兩個步驟程序時，將 SMP 初始化模組新增torch.sagemaker.init()至訓練指令碼，並以 JSON 格式設定 SMP 組態字典，以供訓練工作啟動器使用。開始使用 SageMaker 模型平行程式庫 v2您不需要對 PyTorch 模型或 PyTorch FSDP 組態進行任何變更。如需 hybrid_shard_degree 參數的詳細資訊，請參閱 SMP v2 核心功能組態參數。

SMP 組態設定字典

{ "hybrid_shard_degree": 16 }

在訓練指令碼中

import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
    model,
    ...
)

# Optimizer needs to be created after FSDP wrapper
optimizer = ...