Jigsaw: Domain and Tensor Parallelism for High-Resolution Inp...
summary
Jigsaw shards both model weights and input data so GPUs hold disjoint weight and data slices and avoid redundant I O, implementing distributed matrix multiplication kernels in PyTorch to override autograd for linear layers and attention. It trains on ~400 MB samples, runs a 1B-parameter model on one GPU and a 4B-parameter model across 8 GPUs, overlaps comms with compute, avoids all-gathers, and composes with DDP FSDP Megatron style approaches.
AI-written summary · about
chapters (8)
- 0:00Jigsaw introduction
- 1:47Limitations of current parallelism
- 3:10Tensor parallelism data issues
- 4:32Domain parallelism memory fix
- 5:57Linear layer foundation
- 7:18Jigsaw for attention
- 8:23Jigsaw results and scalability
- 9:24Jigsaw integration summary
related talks
13mLightning Talk: Exploring PiPPY, Tensor Parallel and Torchserve for Large...
24mTraining large scale models using PyTorch
26mEfficiently serving LLMs at scale.
55mLarge Scale Distributed AI on the JVM
34mDeep Learning at Scale
24mKeynote: Next-generation frameworks for Large-scale AI
21mLessons Learned in WatsonX Training: Scaling Cloud-Native...
01h 03mScaling up machine learning without tears (and what do programming languages ha...