When we train large language models, there are a lot of systems challenges and different sharding schemes one can use. While there are many great resources on scaling LLMs out there (https://huggingface.co/spaces/nanotron/ultrascale-playbook or https://jax-ml.github.io/scaling-book/), I felt like there was still a gap when it comes to visualising different forms of parallelism and building intuition around overlaps and execution order for a distributed training run

The idea is to make it easier to visualise FSDP/Tensor Parallel/Expert Parallel/Context parallel and reason about it - you can drag and drop compute kernels and collectives to create DDP/TP/FSDP/EP/CP traces based on real torchtitan profiles.