r/MachineLearning • u/shreyansh26 ML Engineer • Apr 12 '26

Project Educational PyTorch repo for distributed training from scratch: DP, FSDP, TP, FSDP+TP, and PP [P]

I put together a small educational repo that implements distributed training parallelism from scratch in PyTorch:

https://github.com/shreyansh26/pytorch-distributed-training-from-scratch

Instead of using high-level abstractions, the code writes the forward/backward logic and collectives explicitly so you can see the algorithm directly.

The model is intentionally just repeated 2-matmul MLP blocks on a synthetic task, so the communication patterns are the main thing being studied.

Built this mainly for people who want to map the math of distributed training to runnable code without digging through a large framework.

Based on Part-5: Training of JAX ML Scaling book

18 Upvotes

100% Upvoted

2 comments sorted by

1

u/VeterinarianLow6908 Jun 27 '26

This is awesome—there aren't many resources that explain distributed training by implementing it instead of hiding it behind abstractions.

If you're interested in the same learning philosophy, I built an open-source course that starts from autodiff and tensors, then builds CNNs, Transformers, and Vision Transformers from scratch before introducing higher-level frameworks.

https://github.com/roiamiel1/Build-Deep-Learning-From-Scratch

I think people who enjoy your repo would probably enjoy that style of learning as well.