← All examples

Fine-tuning

Train on 1, 2, 4 or 8 cards

Two ways to use several cards, both one line. Many trainings side by side is .map() over a single-card job: a sweep finishes in one run's time. One training across several cards is gpu_count on the decorator plus levers9.distributed.run: one process per card, gradients averaged, and the same function for 1, 2, 4 or 8.

Any GPU pool. A multi-card job needs one node with that many cards free at once; on RTX PRO 6000s, four cards measured 3.15x a single card on the same fine-tune.

How it works

  1. 1distributed.run(fn) starts one process per card the container was given, gives each its rank and joins them into a process group.
  2. 2Inside, dist.device(), dist.wrap(model), dist.shard(rows) and dist.all_reduce_mean(loss) are the whole vocabulary; with one card they answer as a world of one and nothing is spawned.
  3. 3Rank 0's return value is the job's result; the other ranks return nothing.
  4. 4Card quirks live in the SDK: the RTX PRO 6000 goes through shared memory and NCCL older than 2.26 is refused before the first collective, with the package to add.

The code

demo/qwen_asr_train/app.py
from levers9 import distributed as dist, job

@job(gpu="RTXPro6000", gpu_count=4, cpu=8, memory="32Gi", image=IMG)
def train(steps: int = 200) -> dict:
    return dist.run(train_rank, {"steps": steps})    # 4 processes, 4 cards, rank 0's result

def train_rank(cfg: dict) -> dict:
    model = dist.wrap(load_model().to(dist.device()))   # DDP, or the model itself on one card
    for step in range(cfg["steps"]):
        loss = model(**batch_for(step, dist.rank(), dist.world_size())).loss
        loss.backward(); opt.step(); opt.zero_grad()
        dist.print0(f"step {step} loss={dist.all_reduce_mean(loss).item():.4f}")
    return evaluate_and_save(model) if dist.is_main() else {}

# the other shape: many single-card runs at once
results = list(cell.map([(8, 1e-4), (16, 1e-4), (32, 1e-4), (64, 2e-4)]))
run it
cd demo/qwen_asr_train
python app.py rtx-par            # four single-card cells at once, then a table by WER
python app.py rtx4               # one training on four cards of one node

Where to go next

Related