Composable, production-ready building blocks for large-scale AI training:
mix and match parallelism strategies like Lego bricks.
Training a large language model today means orchestrating thousands of GPUs across a massive design space of parallelism strategies, memory optimizations, and hardware configurations. One wrong choice can waste millions of dollars in compute. LegoAI is our research program at DASlab that tackles this challenge head-on by building composable, modular tools — just like Lego bricks, you assemble only the pieces you need.
The program currently includes two systems: TorchTitan, a production-grade distributed training framework that makes it easy to combine different parallelism techniques in any configuration, and TorchSim, a fast simulator that lets you predict how a training run will perform — without ever touching a GPU cluster.
Together, these tools dramatically reduce the barrier to training frontier AI models at scale, enabling researchers to explore, iterate, and innovate faster.
Modern AI systems are too rigid. We build general-purpose composable primitives that snap together seamlessly — giving researchers the flexibility to explore the full design space without re-engineering the stack.
TorchTitan delivers dramatic speedups across model scales,
measured on NVIDIA H100 GPUs training Llama 3.1 models.
Training the next generation of LLMs requires combining multiple parallelism techniques simultaneously. TorchTitan is an open-source, PyTorch-native framework that makes this composable and easy.
Compose data, tensor, pipeline, and context parallelism in any combination.
Hardware-software co-design for maximum throughput on modern accelerators.
Efficient fault-tolerant checkpointing that reduces overhead by up to 15×.
Built-in debugging tools to diagnose distributed training failures at scale.
Existing training systems are often monolithic — once you pick a configuration, you are locked in. TorchTitan was designed from the ground up around 4D parallelism: the seamless combination of Data Parallel (DP), Tensor Parallel (TP), Pipeline Parallel (PP), and Context Parallel (CP) in any configuration you choose. This composability means that the same codebase that trains an 8-billion-parameter model on 8 GPUs can scale — without modification — to a 405-billion-parameter model on 512 GPUs.
TorchTitan also integrates hardware-software co-designed optimizations, including Float8 mixed-precision training and Asynchronous Tensor Parallel (AsyncTP), which overlaps computation and communication to maximize hardware utilization. Production features like distributed checkpointing and a built-in flight recorder for debugging make TorchTitan ready for real workloads — not just research demos.
TorchTitan exposes a vast configuration space. TorchSim lets you explore it entirely on a commodity CPU — no GPU cluster required.
Parallelism strategies, memory optimizations, precision settings — empirically benchmarking thousands of configurations would cost enormous amounts of GPU time and money before a single model trains.
TorchSim simulates a full distributed training run — runtime and memory — by mirroring the real PyTorch execution engine with FakeTensors and FakeProcessGroups, all on an ordinary laptop.
Beyond a single number, TorchSim breaks down runtime into exposed vs. overlapped communication, and memory into parameters, activations, gradients, and optimizer states — at per-module granularity.
TorchSim is built around the insight that you only need tensor metadata —
shapes, dtypes, and device placement — to faithfully simulate a training step.
By running under PyTorch's FakeTensorMode, TorchSim executes your
exact training script (forward pass, backward pass, optimizer step) without any
actual data. It intercepts every operator dispatch and synchronization primitive,
building a precise timeline of compute and communication that it simulates using
hardware-aware cost models. The result: evaluate thousands of training configurations
in minutes and arrive at your GPU cluster with confidence.
