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
- 1distributed.run(fn) starts one process per card the container was given, gives each its rank and joins them into a process group.
- 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.
- 3Rank 0's return value is the job's result; the other ranks return nothing.
- 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
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)]))
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
Train four models in parallel on L40S
Prep on CPU, fan out training to one GPU container per config, evaluate, package: a pipeline in plain Python.
Process ten thousand inputs with .map()
One container per item, sub-second warm starts, scale back to zero when the queue drains.
Open a JupyterLab notebook on a GPU
Your volumes and secrets mounted, live CPU, memory and GPU meters, the kernel on the cluster.