Distribuerad träning

Tre sätt att dela upp träning över flera GPU:er — de löser olika problem:

Regel: börja med dataparallellism (DDP). Räcker inte minnet: lägg på ZeRO/FSDP. Ryms modellen fortfarande inte: tensorparallellism inom en nod, pipeline mellan noder.

ZeRO i tre steg (Rajbhandari m.fl. 2019) — varje steg delar upp en sak till över alla N GPU:er:

Med 8 GPU:er och ZeRO-3 blir det ~11 GB per kort i stället för 84.

Effektiv batchstorlek = batch_per_gpu × antal_gpu × accumulation_steps. Den ska hållas konstant när du ändrar uppsättning, annars ändrar du två saker samtidigt och resultatet går inte att jämföra.

Gradient accumulation simulerar stor batch på lite minne: kör k mikrobatchar, ackumulera gradienterna, uppdatera en gång. Matematiskt nästan ekvivalent med en k gånger större batch (skillnaden ligger i batchnorm-statistik, om sådan används).

Vanligaste praktiska felet: att inte skala inlärningstakten när effektiv batch växer, eller att glömma att bara rank 0 ska logga och spara checkpoints — annars skriver alla över varandra.

# FSDP i PyTorch (motsvarar ZeRO-3)
import torch, torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.utils.data.distributed import DistributedSampler

dist.init_process_group("nccl")
rank, world = dist.get_rank(), dist.get_world_size()
torch.cuda.set_device(rank)

modell = FSDP(modell.cuda(), device_id=rank)
sampler = DistributedSampler(dataset, num_replicas=world, rank=rank, shuffle=True)
loader = torch.utils.data.DataLoader(dataset, batch_size=4, sampler=sampler)

ACC = 4          # effektiv batch = 4 * world * 4
for epok in range(epoker):
    sampler.set_epoch(epok)                   # annars samma ordning varje epok
    for i, (xb, yb) in enumerate(loader):
        loss = kriterium(modell(xb.cuda()), yb.cuda()) / ACC
        loss.backward()
        if (i + 1) % ACC == 0:
            opt.step(); opt.zero_grad()
    if rank == 0:
        spara_checkpoint(modell, epok)        # bara en process skriver