Live Migrator

A research system for live migration of distributed LLM training jobs.

It lets you take a running Megatron-LM training job and move some of its ranks from one set of GPUs to another mid-training (expected event), without restarting from a checkpoint or losing iteration progress.

What it does

Given a running multi-host job (e.g. 24 GPUs across b0, b1, b2), you can ask the system to vacate one host and move its ranks to a fresh host (b3) while the rest of the job keeps training. The controller (main.py on the rank-0 host) builds new NCCL groups in the background, warms up the joining ranks, then does a two-stage switchover and transfers model + optimizer state to the new ranks — all with the live training loop barely stalling.

For the full phase-by-phase walkthrough and the per-iteration measurements a run produces, see microbenchmark.md.

Three Big Components

• live_migrator/ (in-repo, originally from Megatron-LM) — hosts all the migration logic.
• nccl/ — our patched NCCL. It supports the two-stage, zero-memory-overhead NCCL setup.
• pytorch/ — customized PyTorch source build. Built once per cluster and reused. It wires in the NCCL code changes.

Third-party code

third_party/ holds modified copies of NCCL and PyTorch, vendored (with their .git removed) so the system ships as a single repo. Build artifacts are excluded; each project's own .gitignore is preserved.

• third_party/nccl/ — modified from NVIDIA NCCL v2.27.3-1.
• third_party/pytorch/ — modified from PyTorch 2.8.0a0.

Building the environment

The project runs inside a Docker container defined by .devcontainer/devcontainer.json. The container is based on nvcr.io/nvidia/pytorch and its postCreateCommand compiles and installs the patched NCCL and PyTorch for you.

Adjust the mounts for your machine. Open .devcontainer/devcontainer.json and edit the host paths in runArgs (the -v volume mounts for the repo, logs/, checkpoints/, dataset/, nccl/, and pytorch/, plus the SSH keys and GIT_USER_* env vars) so they point at your own directories.

Running it

Once you are inside the container, on the rank-0 host:

python main.py run

What run executes (which experiment, which model, which TP/PP) is decided by uncommenting one of the function calls at the bottom of main.py — e.g. m24_live_migration(MLLM_experiment_profile="20B_TP4_PP2", ...). See microbenchmark.md for the details of a run and where the logs/measurements land.

Other useful actions: python main.py kill (terminate all processes), python main.py sync (sync files across hosts), python main.py test (check SSH connectivity).

Where to look first

• main.py — controller + CLI.
• main.py:m24_live_migration — example experiment driver.
• microbenchmark.md — detailed migration workflow and measurements.