Distributed Training

@gridweave.train(gpus=N) runs your function as a distributed PyTorch job across N GPUs. You write an ordinary training loop; Ray Train's TorchTrainer runs it once per rank with the process group already set up (RANK, WORLD_SIZE, the right backend), so torch DDP, HuggingFace Trainer, or accelerate just work. Usage examples are in the SDK README; this note is what happens underneath.

How the GPUs are chosen

Training needs N whole GPUs (no fractions), and they have to be able to talk to each other:

• Same vendor only — all NVIDIA or all AMD, because the collective backend differs: NCCL for NVIDIA, gloo for AMD.
• Same node preferred — GPUs on one machine talk over NVLink/PCIe. If N doesn't fit on a single node, it falls back to N spread across nodes, and the ranks sync over the network (still same-vendor).

The gateway's VRAM ledger reserves the slots — the same ledger that schedules regular jobs — so a training job competes for GPUs with everything else.

Running it

It rides the normal job path: POST /v1/jobs/train to submit, then the same /v1/jobs/{id}/status · /logs · /result endpoints as any job. Logs are rank-prefixed (rank 0's are flushed to S3), and gridweave.run(train_fn) submits and blocks for the result like a regular job.

Billing

A training job is billed as wall-time × N GPUs — each reserved GPU costs its hourly rate for as long as the job holds it — through the same credit/metering path as jobs (see job_gateway).

Worker requirements

Ray workers ship the training stack (torch, transformers, datasets, accelerate) in the ray-worker image, so most fine-tuning code runs unmodified.