混合碎片資料平行化

碎片資料平行化是節省記憶體的分散式訓練技術，可將模型狀態 (模型參數、漸層與最佳化工具狀態) 分割到不同的裝置。這可協助您適應較大的模型，或使用釋放的 GPU 記憶體增加批次大小。SMP 程式庫提供使用 PyTorch 全碎片資料平行化 (FSDP) 執行碎片資料平行化的功能。根據預設，PyTorch FSDP 會跨整組正在使用的 GPU 碎片。在 SMP v2 中，程式庫透過擴展 PyTorch 混合碎片 (HYBRID_SHARD)，在 PyTorch FSDP 之上提供此碎片資料平行化，這是 PyTorch FSDP 提供的碎片策略之一：FULL_SHARD、SHARD_GRAD_OP、HYBRID_SHARD、_HYBRID_SHARD_ZERO2。以此方式擴展混合碎片有助於實作規模感知碎片，如適用於 PyTorch FSDP 的部落格文章AWS上巨型模型訓練的近線性擴展中所述。

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

下列程式碼片段示範如何在訓練指令碼中新增 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 = ...