混合碎片資料平行處理

陰影資料平行處理是一種節省記憶體的分散式訓練技術，可將模型的狀態 （模型參數、梯度和最佳化工具狀態） 分割到各個裝置。這可協助您適應較大的模型，或使用釋放的GPU記憶體來增加批次大小。SMP 該程式庫提供使用 PyTorch 全碎片資料平行 （FSDP） PyTorch FSDP 執行碎片資料平行處理的功能。預設碎片涵蓋整個GPUs使用集合。在 SMP v2 中，程式庫透過擴展 PyTorch 混合碎片 （HYBRID_SHARD） 來在 PyTorch FSDP之上提供此碎片資料平行處理，這是 提供的碎片策略 PyTorch FSDP之一：FULL_SHARD、SHARD_GRAD_OP、HYBRID_SHARD、、_HYBRID_SHARD_ZERO2。以此方式擴展混合分割有助於實作 scale-aware-sharding 上巨型模型訓練的部落格近線擴展 AWS中所述FSDP。 PyTorch

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

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

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 = ...