Skip to content

microbenchmark.md

Latest commit

 

History

History
45 lines (27 loc) · 4.36 KB

microbenchmark.md

File metadata and controls

45 lines (27 loc) · 4.36 KB
Raw

Live migration microbenchmark

This document describes, in detail, what happens during a live-migration run (python main.py run) — the phase-by-phase workflow the controller drives, and the per-iteration measurements it produces. For a high-level overview and setup instructions, see the main README.

What a run does

Concretely, given a running 24-GPU job on b0 (ranks 0..7), b1 (ranks 8..15), b2 (ranks 16..23), you can ask the system to vacate b1 and move ranks 8..15 to b3 — while b0, b1, and b2 keep training. The controller (= main.py on b0) drives the workflow:

Phase 0: normal training. Launch the job, all 24 ranks (b0, b1, b2) bring up their NCCL groups and start training normally. The controller then waits a fixed amount of time (currently 80 seconds, simulating "now we want to migrate"). Nothing migration-related happens in this window.

Phase 1: migration.

  1. Pick roles. The controller tags each rank with a role:

    • ranks on b1 (leaving) → MIGRATE_OUT
    • new processes on b3 (joining) → MIGRATE_IN
    • ranks on b0 / b2 (staying) → MIGRATION_HELPER

  2. Send the migration signal. The controller sends a "migration plan" message to every currently-running rank (the helpers and migrate-out). The fresh migrate-in processes on b3 are launched separately and pick up their role from the launch command.

  3. Wait a few rounds, then start migrating. Receiving the signal does not begin the migration — it just schedules it. Helpers and migrate-out ranks keep doing normal training iterations for a small number of rounds (a fixed delay), and only after that delay does the actual migration work below begin. This is just a launch-time buffer so the fresh migrate-in processes have time to come up.

  4. Build new NCCL groups in the background. Once the delay ends, every rank starts initializing the new comm groups one by one. To avoid carrying the memory cost of "old groups + new groups" at the same time, init is done as a lightweight first-touch — the heavy buffers/connections are deferred until stage 2. As each group finishes, the rank sends a "group complete" signal to the controller. The controller waits for everyone, then broadcasts an "allowance" so they all start the next group together. While this is happening, helpers and migrate-out ranks keep running training iterations in parallel, so the live training loop barely stalls.

  5. Migrate-in compute warmup. Once all groups are built (stage 1), the migrate-in ranks run a few warmup iterations to get their per-iteration state caught up to the rest of the cluster.

  6. Stage 2: switchover + state transfer. When both conditions hold — all groups initialized and migrate-in's warmup done — every rank does the second stage of group instantiation, which finalizes the new groups and swaps them in to replace the old ones. Immediately after, migrate-out ships model parameters and optimizer state across to migrate-in over a dedicated state channel.

  7. Migrate-in takes over. Migrate-in now holds the same model + optimizer state as migrate-out, on the new groups. Migrate-out exits; migrate-in resumes the training loop in its place. The job is now running on b0, b2, and b3, stepping iterations normally.

Phase 2: experiment end. When training finishes, every rank signals the controller. The controller then runs a few more iterations involving only the migration-helpers and the new migrate-in ranks (so we get a clean post-migration measurement), and finally kills all processes and exits.

Choosing what to run

What run actually 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", ...).

Measurements and logs

Logs go to logs/<experiment_profile>/port<MASTER_PORT>_rank<RANK>.log.

Per-iteration timing and memory measurements are written to logs/<experiment_profile>/migrator_training_data_port<MASTER_PORT>_rank<RANK>.txt — a TSV with columns iteration, training_elapsed_time (seconds), elapsed_time (seconds, includes non-training overhead), and memory_used (MB). One row per completed iteration, per rank. The port in the filename lets you separate the original training job from the post-migration training job (they run on different MASTER_PORTs).