ゼロから学ぶ分散学習のための教育用PyTorchリポジトリ：DP、FSDP、TP、FSDP+TP、PP [P]

Reddit r/MachineLearning / 2026/4/12

💬 オピニオンDeveloper Stack & InfrastructureIdeas & Deep AnalysisTools & Practical Usage

共有:

要点

著者は、分散学習の並列化をゼロから実装する教育用のPyTorch GitHubリポジトリを共有しており、順伝播／逆伝播のロジックと集合通信（collective communications）を明示的にコード化している。
このリポジトリでは、データ並列（DP）、完全シャーディング・データ並列（FSDP）、テンソル並列（TP）といった複数の戦略に加え、FSDP+TPやパイプライン並列（PP）といった組み合わせも示している。
モデルとタスクは意図的に単純にされている（合成データセット上で、2つの行列積（2-matmul）MLPブロックを繰り返す）ため、焦点が通信パターンとアルゴリズム的な挙動に保たれる。
このリポジトリは、大規模なフレームワークの抽象化に依存せずに、分散学習の数学をそのまま実行可能なコードへ対応づけることを学習者が理解できるように設計されており、JAX ML Scaling bookの「training」セクションから着想を得ている。

PyTorchでゼロから分散学習の並列化を実装する、小さな教育用リポジトリを作りました：

高レベルの抽象化を使うのではなく、コードは順伝播／逆伝播のロジックと集合通信（collectives）を明示的に書いているので、アルゴリズムをそのまま見られます。

モデルは意図的に、合成タスク上での「2つの行列積（2-matmul）」のMLPブロックを繰り返すだけのものです。したがって、主に調べているのは通信パターンです。

これは主に、大規模なフレームワークを掘り下げることなく、分散学習の数式をそのまま実行可能なコードに対応づけたい人向けに作りました。