Write Code from Examples¶
Writing research code from scratch is rarely necessary. Mila provides two starting points: a collection of minimal examples that each illustrate one concept in isolation, and the Research Project Template, a complete project structure ready to be used for a new research project. Both are designed to run on the Mila cluster as-is.
Before you begin¶
-
Obtain a Mila account, enable cluster access and MFA, install
uvandmilatools, configure SSH access and connect to the cluster for the first time.
What this page covers¶
- Choose between the minimal examples and the Research Project Template
- Find the minimal example matching a specific need
- Start a new project from the Research Project Template
Choose a starting point¶
| Need | Starting point |
|---|---|
| Understand one concept (multi-GPU, checkpointing, etc.) | Minimal examples |
| Adapt an existing script to run on the cluster | Minimal examples |
| Start a new research project with a full, tested codebase | Research Project Template |
Minimal examples¶
The minimal examples are small, self-contained
projects. Each one contains a job.sh Slurm script, launched with
sbatch job.sh, and a main.py Python script. Many examples are displayed as
a difference with respect to a simpler base example, which highlights exactly
what changes are required to add a feature.
-
Set up an environment with PyTorch, JAX or Flash Attention using
uv. -
Scale a training script from a single GPU to multiple GPUs and nodes.
-
Add checkpointing, experiment tracking, job arrays and hyperparameter search.
-
Combine these concepts, for instance to train on ImageNet across multiple nodes.
Read the examples in order
Each distributed training example builds on the previous one. Start with the single GPU job before moving on to the multi-GPU and multi-node examples.
Add metric logging from the start
The WandB setup example shows a training script that already logs metrics. Adding logging while writing the code makes the results available in Phase 3. See Monitor and Optimize Experiments.
Research Project Template¶
The Research Project Template is a starting point for new machine learning research projects. It provides a project structure that already integrates the tools commonly used at Mila:
- Hydra to configure experiments.
- PyTorch Lightning and JAX to write training code.
- Weights & Biases to track experiments.
- pytest to test the code.
Instead of assembling these components one by one, the template provides them already configured and working together, so research can start on a tested codebase. Follow the instructions in the Research Project Template documentation to create a new project from it.
Other tools
The Toolbox section describes more tools that can be used alongside the template or the minimal examples.
Next step¶
-
Submit and manage jobs on the cluster with Slurm.