本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。
混合碎片資料平行處理
陰影資料平行處理是一種節省記憶體的分散式訓練技術,可將模型的狀態 (模型參數、梯度和最佳化工具狀態) 分割到各個裝置。這可協助您適應較大的模型,或使用釋放的GPU記憶體來增加批次大小。SMP 該程式庫提供使用 PyTorch 全碎片資料平行 (FSDP) PyTorch FSDP 執行碎片資料平行處理的功能。預設碎片涵蓋整個GPUs使用集合。在 SMP v2 中,程式庫透過擴展 PyTorch 混合碎片 (HYBRID_SHARD
) 來在 PyTorch FSDP之上提供此碎片資料平行處理,這是 提供的碎片策略 PyTorch FSDPFULL_SHARD
、SHARD_GRAD_OP
、HYBRID_SHARD
、、_HYBRID_SHARD_ZERO2
。以此方式擴展混合分割有助於實作 scale-aware-sharding 上巨型模型訓練的部落格近線擴展 AWS
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 FSDPhybrid_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 = ...