最佳化工具狀態碎片

最佳化工具狀態碎片是有用的節省記憶體技術，可跨資料平行裝置群組碎片化最佳化工具狀態 (描述最佳化工具狀態的權重集)。每當您使用有狀態優化器（例如 Adam）或優化器（存儲參數FP16和FP32副本）時，都可以使用FP16優化器狀態分片。

注意

最佳化程式狀態分割適用 PyTorch 於 SageMaker 模型平行程式庫 v1.6.0 及更新版本。

如何使用最佳化工具狀態碎片

您可以透過在 modelparallel 組態 設定 "shard_optimizer_state": True 來開啟最佳化工具狀態碎片。

當開啟此功能時，程式庫會基於資料平行處理程度分割模型參數集。對應於第 i 個分區的漸層僅會於第 i 個資料平行等級縮減。在第一次呼叫 smp.step 裝飾項目函式結束時，由 smp.DistributedOptimizer 包裝的最佳化工具會重新定義其參數，以便僅限對應目前資料平行等級分區的參數。重新定義的參數稱為虛擬參數，並與原始參數共用基礎儲存。在第一次呼叫 optimizer.step 期間，會基於這些重新定義的參數建立最佳化工具狀態，這些參數會因為原始分割而進行碎片處理。最佳化工具更新之後， AllGather作業 (做為optimizer.step呼叫的一部分) 會跨資料 parallel 等級執行，以達到一致的參數狀態。

提示

當資料平行處理的程度大於 1 且模型具十億個以上的參數時，最佳化工具狀態碎片很有幫助。

資料平行處理的程度由 (processes_per_host * instance_count / pipeline_parallel_degree) 計算，而 smp.dp_size() 函式會在背景處理大小。

設定 SageMaker PyTorch 估算器

mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True,
        "shard_optimizer_state": True
    }
}

調整您的 PyTorch 訓練指令碼

請參閱 Tensor 平行處理原則結合管道平行處理區段中的調整 PyTorch 訓練指令碼。此指令碼不需要其他修改。