パイプライン並列処理とテンソル並列処理を組み合わせて使用する場合のランキングメカニズム

このセクションでは、モデル並列処理のランキングメカニズムがテンソル並列処理とどのように作用するかについて説明します。これは、Model SageMaker Parallelism Library の主な機能 のランキングの基本の延長です。テンソル並列処理では、ライブラリは 3 種類のランキングとプロセスグループ を導入しますAPIs。テンソル並列ランクsmp.tp_rank()の場合は 、パイプライン並列ランクsmp.pp_rank()の場合は 、縮小データ並列ランクsmp.rdp_rank()の場合は です。対応する通信プロセスグループは、テンソル並列グループ (TP_GROUP)、パイプライン並列グループ (PP_GROUP)、および縮小データ並列グループ(RDP_GROUP) です。これらのグループは以下のように定義されます。

• テンソル並列グループ (TP_GROUP) は、データ並列グループの均等に割り切れるサブセットであり、それに対してモジュールのテンソル並列分散が行われます。パイプラインの並列度が 1 の場合、TP_GROUP はモデル並列グループ (MP_GROUP) と同じです。

• パイプライン並列グループ (PP_GROUP) は、パイプライン並列処理が行われるプロセスのグループです。テンソル並列度が 1 のとき、PP_GROUP は MP_GROUP と同じです。

• 縮小データ並列グループ (RDP_GROUP) は、同じパイプライン並列パーティションと、同じテンソル並列パーティションの両方を保持し、それらの間でデータ並列処理を実行する一連のプロセスです。これは、データ並列グループ、DP_GROUP 全体のサブセットであるため、削減データ並列グループと呼ばれます。TP_GROUP 内部に分布するモデルパラメータについては、勾配 allreduce オペレーションは縮小データ並列グループに対してのみ実行され、分布していないパラメータでは、勾配 allreduce は DP_GROUP 全体にわたって行われます。

• モデル並列グループ (MP_GROUP) とは、モデル全体をまとめて格納するプロセスのグループを指します。これは、現在のプロセスの TP_GROUP にあるすべてのランクの PP_GROUP の統合で構成されています。テンソル並列度が 1 であるとき、MP_GROUP は PP_GROUP と等しくなります。また、前の smdistributed リリースから、MP_GROUP の既存の定義と一致しています。なお、現在の TP_GROUP は現在の DP_GROUP と現在の MP_GROUP の両方のサブセットです。

SageMaker モデル並列処理ライブラリAPIsの通信プロセスの詳細については、SageMaker Python SDKドキュメントの「共通API」とPyTorch「固有APIs」を参照してください。

例えば、テンソル並列処理の度合いが 2GPUs、パイプライン並列処理の度合いが 2、データ並列処理の度合いが 4 である 8 の単一ノードのプロセスグループを考えてみましょう。前の図の中央上部には、4 レイヤーモデルの例が示されています。図の左下と右下の部分は、パイプライン並列処理とテンソル並列処理の両方GPUsを使用して 4 層に分散された 4 層モデルを示しています。テンソル並列処理は、中央の 2 つのレイヤーに使用されます。これらの 2 つの図は、異なるグループ境界線を示す単純なコピーです。パーティション化されたモデルは、0～3 と GPUs 4～7 のデータ並列処理用にレプリケートされます。左下の図は、MP_GROUP、PP_GROUP、および TP_GROUP の定義を示しています。右下の図はRDP_GROUP、同じ セットWORLDの DP_GROUP、、および を示していますGPUs。同じ色を持つレイヤーとレイヤースライスの勾配は、データ並列処理に対して一緒に allreduce されます。例えば、最初のレイヤー (水色) では、DP_GROUP 全体にわたって allreduce オペレーションが実行されます。一方、2 つ目のレイヤーの濃いオレンジのスライスでは、そのプロセスの RDP_GROUP 内でのみ allreduce オペレーションが行われます。太字の濃い赤色の矢印は、TP_GROUP 全体のバッチのあるテンソルを表しています。

GPU0: pp_rank 0, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 0
GPU1: pp_rank 1, tp_rank 0, rdp_rank 0, dp_rank 0, mp_rank 1
GPU2: pp_rank 0, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 2
GPU3: pp_rank 1, tp_rank 1, rdp_rank 0, dp_rank 1, mp_rank 3
GPU4: pp_rank 0, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 0
GPU5: pp_rank 1, tp_rank 0, rdp_rank 1, dp_rank 2, mp_rank 1
GPU6: pp_rank 0, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 2
GPU7: pp_rank 1, tp_rank 1, rdp_rank 1, dp_rank 3, mp_rank 3

この例では、パイプラインの並列処理はGPUペア (0,1)、 (2,3)、 (4,5)、および (6,7) で行われます。さらに、データ並列処理 (allreduce) は GPUs 0、2、4、6、および GPUs1、3、5、7 で個別に行われます。テンソル並列処理DP_GROUPは、GPUペア (0,2)、 (1,3)、 (4,6)、および (5,7) の のサブセットで行われます。