Lightning Talk: Jigsaw: Domain and Tensor Parallelism for High-Resolution Inp... Deifilia Kieckhefen

This talk explores Jigsaw, a PyTorch library designed to enhance domain and tensor parallelism for high-resolution input training in distributed neural network frameworks. The speaker, Deifilia Kieckhefen from Karlsruhe Institute of Technology, discusses how Jigsaw efficiently shards both model weights and input data across parallel processes, addressing common bottlenecks found in complex model architectures. By parallelizing activations, convolutions, linear layers, and attention with a distributed matrix multiplication backend, Jigsaw enables improved performance for multi-billion-parameter models. The session demonstrates Jigsaw's usability across various model architectures and compares its scalability to established methods like DDP, FSDP, and Megatron-LM.