Gradient Accumulation corect în DDP: no_sync() fără sincronizări inutile

Patru micro-batch-uri ar trebui să producă o singură actualizare a modelului. Într-un job DistributedDataParallel, însă, acumularea implementată doar prin amânarea lui optimizer.step() continuă să sincronizeze gradientele după fiecare backward(). Rezultatul matematic poate fi corect, dar rețeaua efectuează de patru ori mai multe sincronizări decât sunt necesare.

Proiectul de mai jos măsoară direct diferența. Antrenează același CNN pe CIFAR-10, pe două GPU-uri, în două moduri: every, care sincronizează fiecare micro-batch, și no_sync, care comunică numai la capătul ferestrei de acumulare. Un communication hook numără sincronizările DDP, iar un verificator confirmă raportul așteptat: 400 de sincronizări față de 100 pentru fiecare epocă, la același număr de actualizări ale optimizatorului.

La final rămân checkpointul, metricile pe test, matricea de confuzie, perechile de clase confundate cel mai des și un program CLI care clasifică o imagine. Tehnica se poate combina ulterior cu Selective Activation Checkpointing pentru reducerea VRAM, dar cele două optimizări rezolvă probleme diferite: checkpointing-ul reduce memoria activă, iar no_sync() reduce comunicarea dintre procese.

Fixăm experimentul la două GPU-uri și un batch efectiv de 400 de imagini

Laboratorul folosește Ubuntu 24.04 LTS, Python 3.11, PyTorch 2.8.0, torchvision 0.23.0 și pachetele CUDA 12.8 distribuite de proiectul PyTorch. Sunt necesare exact două GPU-uri NVIDIA cu suport BF16 și minimum 8 GB VRAM fiecare, aproximativ 6 GB de spațiu liber și 45–90 de minute pentru cele două rulări, în funcție de GPU.

CIFAR-10 conține 60.000 de imagini color de 32×32 pixeli, împărțite în zece clase: 50.000 pentru training și 10.000 pentru test. Vom separa determinist cele 50.000 de imagini în 40.000 pentru antrenare și 10.000 pentru validare. Pagina oficială CIFAR-10 nu afișează o licență explicită; înaintea reutilizării comerciale trebuie verificate termenii sursei și proveniența datelor, nu presupusă o licență.

Verifică stația înainte de instalare:

cat /etc/os-release
python3 --version
nvidia-smi
df -h .

Exemplu de rezultat acceptabil; modelele GPU, driverul și spațiul disponibil vor diferi:

PRETTY_NAME="Ubuntu 24.04.3 LTS"
Python 3.11.x
NVIDIA-SMI ...   2 GPU-uri vizibile
Filesystem ...   Avail 6G sau mai mult

Creează mediul și instalează versiunile fixate:

mkdir -p ~/ddp-no-sync-lab
cd ~/ddp-no-sync-lab
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
cat > requirements.txt <<'EOF'
--index-url https://download.pytorch.org/whl/cu128
torch==2.8.0
torchvision==0.23.0
EOF
python -m pip install -r requirements.txt
python - <<'PY'
import torch
import torchvision

print("torch:", torch.__version__)
print("torchvision:", torchvision.__version__)
print("CUDA runtime:", torch.version.cuda)
print("GPU count:", torch.cuda.device_count())
print("BF16:", torch.cuda.is_bf16_supported())

assert torch.cuda.device_count() == 2, "Laboratorul cere exact două GPU-uri vizibile"
assert torch.cuda.is_bf16_supported(), "GPU-ul curent nu raportează suport BF16"
PY

Outputul trebuie să confirme torch 2.8.0, torchvision 0.23.0, două GPU-uri și BF16: True. Sufixul versiunii PyTorch poate include identificatorul buildului CUDA.

Fișierele proiectului vor avea această structură:

ddp-no-sync-lab/
├── requirements.txt       # versiunile Python și indexul wheel-urilor CUDA
├── model.py               # CNN-ul folosit identic la training și inferență
├── train.py               # date, baseline, DDP, no_sync, evaluare și artefacte
├── infer.py               # clasificarea unei imagini cu checkpointul salvat
├── verify.py              # verifică numărul sincronizărilor și rapoartele
├── data/                  # arhiva CIFAR-10 descărcată de torchvision
└── artifacts/
    ├── every/             # rularea care sincronizează fiecare micro-batch
    └── no_sync/           # rularea optimizată și imaginea pentru inferență

Pasul 1 — Definim un CNN suficient de mic pentru un singur bucket DDP

Scopul este să izolăm comportamentul acumulării, nu să obținem un record pe CIFAR-10. Modelul are mai puțin de 25 MB de gradienți, iar DistributedDataParallel va fi configurat cu bucket_cap_mb=25, astfel încât communication hook-ul să observe o sincronizare pe fiecare backward sincronizat.

Fișierul model.py:

from torch import nn


CLASSES = (
    "airplane",
    "automobile",
    "bird",
    "cat",
    "deer",
    "dog",
    "frog",
    "horse",
    "ship",
    "truck",
)


class CifarCNN(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.Conv2d(128, 128, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(256, len(CLASSES))

    def forward(self, inputs):
        features = self.features(inputs)
        return self.classifier(features.flatten(1))

AdaptiveAvgPool2d(1) elimină dependența clasificatorului de dimensiunea spațială finală. Normalizările BatchNorm rămân locale fiecărei replici; proiectul nu transformă modelul în SyncBatchNorm, pentru a nu introduce alte colective în măsurătoare.

Verifică modelul:

python - <<'PY'
import torch
from model import CifarCNN

model = CifarCNN()
output = model(torch.zeros(2, 3, 32, 32))
parameters = sum(p.numel() for p in model.parameters())
print("output:", tuple(output.shape))
print("parameters:", parameters)
assert output.shape == (2, 10)
PY

Trebuie să vezi:

output: (2, 10)
parameters: 445834

Punct de control: dacă forma nu este (2, 10), verifică dacă model.py a fost copiat integral și dacă rulezi comanda din ~/ddp-no-sync-lab.

Pasul 2 — Legăm fereastra de acumulare de contextul no_sync()

Aici se află diferența care contează. Contextul ddp.no_sync() trebuie să cuprindă atât forward-ul, cât și backward-ul micro-batch-urilor intermediare. Documentația PyTorch avertizează explicit că includerea numai a lui backward() nu împiedică sincronizarea, deoarece forward-ul DDP pregătește mecanismele folosite la reducerea gradientelor.

Fișierul complet train.py:

import argparse
import json
import os
import random
from collections import Counter
from contextlib import nullcontext
from pathlib import Path

import torch
import torch.distributed as dist
from torch import nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Subset
from torch.utils.data.distributed import DistributedSampler
from torchvision import datasets, transforms

from model import CLASSES, CifarCNN


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=("every", "no_sync"), required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--batch-size", type=int, default=50)
    parser.add_argument("--accum-steps", type=int, default=4)
    parser.add_argument("--seed", type=int, default=47)
    parser.add_argument("--expected-world-size", type=int, default=2)
    return parser.parse_args()


def seed_everything(seed):
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.use_deterministic_algorithms(True)


def build_loaders(args, rank, world_size):
    data_root = Path("data")
    train_transform = transforms.Compose(
        [
            transforms.RandomCrop(32, padding=4),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(
                (0.4914, 0.4822, 0.4465),
                (0.2470, 0.2435, 0.2616),
            ),
        ]
    )
    eval_transform = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(
                (0.4914, 0.4822, 0.4465),
                (0.2470, 0.2435, 0.2616),
            ),
        ]
    )

    if rank == 0:
        datasets.CIFAR10(data_root, train=True, download=True)
        datasets.CIFAR10(data_root, train=False, download=True)
    dist.barrier()

    train_source = datasets.CIFAR10(
        data_root, train=True, transform=train_transform, download=False
    )
    validation_source = datasets.CIFAR10(
        data_root, train=True, transform=eval_transform, download=False
    )
    test_source = datasets.CIFAR10(
        data_root, train=False, transform=eval_transform, download=False
    )

    generator = torch.Generator().manual_seed(args.seed)
    permutation = torch.randperm(50000, generator=generator).tolist()
    train_indices = permutation[:40000]
    validation_indices = permutation[40000:]

    train_set = Subset(train_source, train_indices)
    validation_set = Subset(validation_source, validation_indices)

    train_sampler = DistributedSampler(
        train_set,
        num_replicas=world_size,
        rank=rank,
        shuffle=True,
        seed=args.seed,
        drop_last=False,
    )
    validation_sampler = DistributedSampler(
        validation_set,
        num_replicas=world_size,
        rank=rank,
        shuffle=False,
        drop_last=False,
    )
    test_sampler = DistributedSampler(
        test_source,
        num_replicas=world_size,
        rank=rank,
        shuffle=False,
        drop_last=False,
    )

    common = {
        "batch_size": args.batch_size,
        "num_workers": 0,
        "pin_memory": True,
    }
    train_loader = DataLoader(
        train_set, sampler=train_sampler, drop_last=False, **common
    )
    validation_loader = DataLoader(
        validation_set, sampler=validation_sampler, drop_last=False, **common
    )
    test_loader = DataLoader(
        test_source, sampler=test_sampler, drop_last=False, **common
    )

    train_targets = [train_source.targets[index] for index in train_indices]
    majority_class = Counter(train_targets).most_common(1)[0][0]
    return (
        train_loader,
        validation_loader,
        test_loader,
        train_sampler,
        majority_class,
    )


def reduce_metrics(loss_sum, correct, count, device):
    values = torch.tensor(
        [loss_sum, float(correct), float(count)],
        dtype=torch.float64,
        device=device,
    )
    dist.all_reduce(values, op=dist.ReduceOp.SUM)
    return values[0].item() / values[2].item(), values[1].item() / values[2].item()


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    loss_sum = 0.0
    correct = 0
    count = 0
    confusion = torch.zeros(
        len(CLASSES), len(CLASSES), dtype=torch.int64, device=device
    )

    for inputs, targets in loader:
        inputs = inputs.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            logits = model(inputs)
            loss = criterion(logits, targets)

        predictions = logits.argmax(dim=1)
        loss_sum += loss.item() * targets.size(0)
        correct += predictions.eq(targets).sum().item()
        count += targets.size(0)

        flat = targets * len(CLASSES) + predictions
        confusion += torch.bincount(
            flat, minlength=len(CLASSES) ** 2
        ).reshape(len(CLASSES), len(CLASSES))

    loss, accuracy = reduce_metrics(loss_sum, correct, count, device)
    dist.all_reduce(confusion, op=dist.ReduceOp.SUM)
    return loss, accuracy, confusion


def counting_allreduce_hook(state, bucket):
    state["calls"] += 1
    tensor = bucket.buffer()
    tensor.div_(state["world_size"])
    return dist.all_reduce(tensor, group=state["group"], async_op=True).get_future().then(
        lambda future: future.value()[0]
    )


def train_epoch(
    ddp_model,
    loader,
    sampler,
    optimizer,
    criterion,
    device,
    epoch,
    mode,
    accum_steps,
    hook_state,
):
    ddp_model.train()
    sampler.set_epoch(epoch)
    optimizer.zero_grad(set_to_none=True)

    loss_sum = 0.0
    correct = 0
    count = 0
    optimizer_steps = 0
    calls_before = hook_state["calls"]

    if len(loader) % accum_steps != 0:
        raise RuntimeError(
            f"{len(loader)} micro-batch-uri nu se împart exact la accum_steps={accum_steps}"
        )

    for micro_step, (inputs, targets) in enumerate(loader, start=1):
        inputs = inputs.to(device, non_blocking=True)
        targets = targets.to(device, non_blocking=True)
        update_now = micro_step % accum_steps == 0
        suppress_sync = mode == "no_sync" and not update_now
        sync_context = ddp_model.no_sync() if suppress_sync else nullcontext()

        with sync_context:
            with torch.autocast("cuda", dtype=torch.bfloat16):
                logits = ddp_model(inputs)
                raw_loss = criterion(logits, targets)
                scaled_loss = raw_loss / accum_steps
            scaled_loss.backward()

        predictions = logits.detach().argmax(dim=1)
        loss_sum += raw_loss.detach().item() * targets.size(0)
        correct += predictions.eq(targets).sum().item()
        count += targets.size(0)

        if update_now:
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            optimizer_steps += 1

    train_loss, train_accuracy = reduce_metrics(
        loss_sum, correct, count, device
    )
    sync_calls = hook_state["calls"] - calls_before
    return train_loss, train_accuracy, optimizer_steps, sync_calls


def top_confusions(confusion, limit=5):
    matrix = confusion.cpu().clone()
    matrix.fill_diagonal_(0)
    entries = []
    for flat_index in torch.argsort(matrix.flatten(), descending=True)[:limit]:
        true_index = int(flat_index // len(CLASSES))
        predicted_index = int(flat_index % len(CLASSES))
        entries.append(
            {
                "true": CLASSES[true_index],
                "predicted": CLASSES[predicted_index],
                "count": int(matrix[true_index, predicted_index]),
            }
        )
    return entries


def main():
    args = parse_args()
    local_rank = int(os.environ["LOCAL_RANK"])
    rank = int(os.environ["RANK"])
    world_size = int(os.environ["WORLD_SIZE"])

    if world_size != args.expected_world_size:
        raise ValueError(
            f"expected world_size={args.expected_world_size}, got {world_size}"
        )

    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)
    dist.init_process_group(backend="nccl")
    seed_everything(args.seed + rank)

    args.output.mkdir(parents=True, exist_ok=True)
    (
        train_loader,
        validation_loader,
        test_loader,
        train_sampler,
        majority_class,
    ) = build_loaders(args, rank, world_size)

    model = CifarCNN().to(device)
    ddp_model = DDP(
        model,
        device_ids=[local_rank],
        output_device=local_rank,
        bucket_cap_mb=25,
        broadcast_buffers=True,
    )
    hook_state = {
        "calls": 0,
        "group": dist.group.WORLD,
        "world_size": world_size,
    }
    ddp_model.register_comm_hook(hook_state, counting_allreduce_hook)

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(
        ddp_model.parameters(), lr=3e-3, weight_decay=1e-4
    )
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer, T_max=args.epochs
    )

    history = []
    for epoch in range(args.epochs):
        train_loss, train_accuracy, steps, syncs = train_epoch(
            ddp_model,
            train_loader,
            train_sampler,
            optimizer,
            criterion,
            device,
            epoch,
            args.mode,
            args.accum_steps,
            hook_state,
        )
        validation_loss, validation_accuracy, _ = evaluate(
            ddp_model, validation_loader, criterion, device
        )
        scheduler.step()

        row = {
            "epoch": epoch + 1,
            "mode": args.mode,
            "optimizer_steps": steps,
            "gradient_syncs": syncs,
            "train_loss": train_loss,
            "train_accuracy": train_accuracy,
            "validation_loss": validation_loss,
            "validation_accuracy": validation_accuracy,
        }
        history.append(row)
        if rank == 0:
            print(json.dumps(row), flush=True)

    test_loss, test_accuracy, confusion = evaluate(
        ddp_model, test_loader, criterion, device
    )

    if rank == 0:
        checkpoint = {
            "model": ddp_model.module.state_dict(),
            "classes": CLASSES,
            "seed": args.seed,
        }
        torch.save(checkpoint, args.output / "model.pt")
        (args.output / "history.json").write_text(
            json.dumps(history, indent=2), encoding="utf-8"
        )

        baseline_accuracy = float(confusion[majority_class].sum().item()) / float(
            confusion.sum().item()
        )
        report = {
            "mode": args.mode,
            "baseline_class": CLASSES[majority_class],
            "baseline_accuracy": baseline_accuracy,
            "model_test_loss": test_loss,
            "model_test_accuracy": test_accuracy,
            "confusion_matrix": confusion.cpu().tolist(),
            "top_confusions": top_confusions(confusion),
        }
        (args.output / "report.json").write_text(
            json.dumps(report, indent=2), encoding="utf-8"
        )

        raw_test = datasets.CIFAR10(
            "data", train=False, download=False
        )
        raw_test[0][0].save(args.output / "sample_test.png")
        print(json.dumps(report), flush=True)

    dist.barrier()
    dist.destroy_process_group()


if __name__ == "__main__":
    main()

Batch-ul local este 50, sunt două procese, iar patru micro-batch-uri contribuie la o actualizare. Batch-ul efectiv este astfel 50 × 2 × 4 = 400.

Împărțirea la accum_steps înainte de backward() păstrează scara gradientului echivalentă mediei celor patru micro-batch-uri. Cele 40.000 de exemple produc exact 400 de micro-batch-uri pe rang și 100 de actualizări pe epocă; nu există o fereastră finală incompletă care să necesite alt divizor.

set_epoch(epoch) schimbă ordinea deterministă folosită de DistributedSampler. Fără apel, samplerul ar repeta aceeași ordine la fiecare epocă.

Communication hook-ul împarte bufferul la world_size înainte de all_reduce, păstrând comportamentul de mediere al DDP. El nu este o optimizare recomandată pentru aplicația finală, ci instrumentul de audit al acestui laborator.

Verifică mai întâi sintaxa:

python -m py_compile model.py train.py

Lipsa outputului înseamnă că fișierele se compilează. Punct de control: o eroare SyntaxError indică linia la care fișierul a fost copiat incomplet.

Pasul 3 — Măsurăm costul implementării care sincronizează fiecare micro-batch

Prima rulare amână actualizarea parametrilor, dar nu suprimă sincronizările. Ea este controlul experimental: 100 de pași ai optimizatorului și 400 de apeluri ale communication hook-ului în fiecare epocă.

Rulează:

source .venv/bin/activate
torchrun --standalone --nnodes=1 --nproc-per-node=2 
  train.py 
  --mode every 
  --output artifacts/every 
  --epochs 5 
  --batch-size 50 
  --accum-steps 4

torchrun pornește câte un proces pe GPU și furnizează variabilele LOCAL_RANK, RANK și WORLD_SIZE. --standalone creează rendezvous-ul local, iar --nproc-per-node=2 trebuie să corespundă celor două GPU-uri vizibile.

După descărcarea inițială a datasetului, fiecare epocă va afișa un obiect JSON asemănător cu acesta; valorile pierderii și acurateții depind de hardware și de execuție:

{"epoch": 1, "mode": "every", "optimizer_steps": 100, "gradient_syncs": 400, "train_loss": ..., "train_accuracy": ..., "validation_loss": ..., "validation_accuracy": ...}

Trebuie să apară artifacts/every/model.pt, history.json, report.json și sample_test.png.

Punct de control: dacă optimizer_steps nu este 100, verifică simultan --batch-size 50, cele două procese și splitul de 40.000 de imagini; nu modifica doar valoarea așteptată.

Pasul 4 — Suprimăm trei reduceri din patru fără să schimbăm batch-ul efectiv

A doua rulare folosește același seed, aceleași date, aceeași inițializare și același optimizator. Singura diferență este contextul no_sync() aplicat primelor trei micro-batch-uri din fiecare fereastră.

Rulează:

torchrun --standalone --nnodes=1 --nproc-per-node=2 
  train.py 
  --mode no_sync 
  --output artifacts/no_sync 
  --epochs 5 
  --batch-size 50 
  --accum-steps 4

Pe primele trei micro-batch-uri, gradientele se acumulează local în câmpurile .grad. Al patrulea forward-backward rulează în afara contextului, iar DDP sincronizează inclusiv gradientele acumulate anterior. Abia după această reducere se execută optimizer.step() și zero_grad().

Outputul structural corect este:

{"epoch": 1, "mode": "no_sync", "optimizer_steps": 100, "gradient_syncs": 100, "train_loss": ..., "train_accuracy": ..., "validation_loss": ..., "validation_accuracy": ...}

Numărul de actualizări rămâne 100, dar sincronizările scad de la 400 la 100. Asta nu garantează o accelerare de exact patru ori: forward-ul, backward-ul, încărcarea datelor și actualizarea parametrilor continuă să consume timp. Măsurătoarea demonstrează reducerea colectivelor, nu promite un anumit spor de performanță.

Pentru un antrenament lingvistic mai mare, aceeași structură a ferestrei poate fi aplicată unui proiect de continued pretraining cu replay. no_sync() nu rezolvă însă selecția datelor, uitarea catastrofală sau traversarea limitelor dintre documente; acestea rămân probleme separate, inclusiv în proiectele cu sequence packing și măști de attention corecte.

Punct de control: dacă vezi tot 400 de sincronizări, contextul nu cuprinde forward-ul sau rularea folosește accidental --mode every.

Pasul 5 — Blocăm regresia printr-un verificator al artefactelor

Nu este suficient ca jobul să termine fără excepție. Verificatorul cere cinci epoci pentru fiecare mod, 100 de actualizări per epocă și exact raportul de sincronizare proiectat.

Fișierul verify.py:

import json
from pathlib import Path


def load_json(path):
    if not path.is_file():
        raise FileNotFoundError(f"Lipsește artefactul obligatoriu: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def validate_history(path, mode, expected_syncs):
    history = load_json(path)
    if len(history) != 5:
        raise RuntimeError(f"{path}: expected 5 epochs, got {len(history)}")

    for row in history:
        if row["mode"] != mode:
            raise RuntimeError(
                f"{path}: expected mode={mode}, got {row['mode']}"
            )
        if row["optimizer_steps"] != 100:
            raise RuntimeError(
                f"{path}: expected 100 optimizer steps, "
                f"got {row['optimizer_steps']}"
            )
        if row["gradient_syncs"] != expected_syncs:
            raise RuntimeError(
                f"{path}: expected {expected_syncs} gradient syncs, "
                f"got {row['gradient_syncs']}"
            )
    return history


every = validate_history(
    Path("artifacts/every/history.json"), "every", 400
)
optimized = validate_history(
    Path("artifacts/no_sync/history.json"), "no_sync", 100
)
every_report = load_json(Path("artifacts/every/report.json"))
optimized_report = load_json(Path("artifacts/no_sync/report.json"))

required = (
    Path("artifacts/every/model.pt"),
    Path("artifacts/no_sync/model.pt"),
    Path("artifacts/no_sync/sample_test.png"),
)
for artifact in required:
    if not artifact.is_file() or artifact.stat().st_size == 0:
        raise RuntimeError(f"Artefact absent sau gol: {artifact}")

print("| mod | pași optimizer/epocă | sincronizări/epocă | test accuracy |")
print("|---|---:|---:|---:|")
print(
    f"| every | 100 | 400 | "
    f"{every_report['model_test_accuracy']:.4f} |"
)
print(
    f"| no_sync | 100 | 100 | "
    f"{optimized_report['model_test_accuracy']:.4f} |"
)
print(
    "baseline:",
    optimized_report["baseline_class"],
    f"{optimized_report['baseline_accuracy']:.4f}",
)
print("top confusions:", optimized_report["top_confusions"])
print("AUDIT PASSED: no_sync reduced gradient synchronizations by 4x")

Execută testul de capăt:

python verify.py

Formatul rezultatului trebuie să fie:

| mod | pași optimizer/epocă | sincronizări/epocă | test accuracy |
|---|---:|---:|---:|
| every | 100 | 400 | 0.xxxx |
| no_sync | 100 | 100 | 0.xxxx |
baseline: <clasă> 0.1000
top confusions: [...]
AUDIT PASSED: no_sync reduced gradient synchronizations by 4x

CIFAR-10 are câte 1.000 de imagini din fiecare clasă în setul oficial de test, deci predictorul care alege mereu clasa majoritară din splitul de training obține 0,1000 pe test. Acuratețea modelului și confuziile trebuie citite din raportul generat, nu înlocuite cu valori presupuse.

Comparația dintre every și no_sync nu cere checkpointuri identice bit cu bit. Ordinea și momentul reducerilor pot produce mici diferențe numerice, chiar dacă ambele variante reprezintă același batch efectiv. Criteriul structural este numărul colectivelor, iar criteriul ML este raportarea separată a baseline-ului și a metricii pe test.

Punct de control: mesajul final AUDIT PASSED apare numai dacă ambele istorii și toate artefactele obligatorii sunt coerente.

Pasul 6 — Încărcăm checkpointul optimizat într-un CLI de inferență

Servirea minimă este un program local care primește un fișier PNG sau JPEG, aplică aceeași normalizare și afișează clasa, probabilitatea maximă și toate scorurile softmax.

Fișierul infer.py:

import argparse
import json
from pathlib import Path

import torch
from PIL import Image
from torchvision import transforms

from model import CLASSES, CifarCNN


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--image", type=Path, required=True)
    return parser.parse_args()


def main():
    args = parse_args()
    if not args.checkpoint.is_file():
        raise FileNotFoundError(f"Checkpoint inexistent: {args.checkpoint}")
    if not args.image.is_file():
        raise FileNotFoundError(f"Imagine inexistentă: {args.image}")

    checkpoint = torch.load(
        args.checkpoint, map_location="cpu", weights_only=True
    )
    model = CifarCNN()
    model.load_state_dict(checkpoint["model"])
    model.eval()

    preprocess = transforms.Compose(
        [
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize(
                (0.4914, 0.4822, 0.4465),
                (0.2470, 0.2435, 0.2616),
            ),
        ]
    )
    with Image.open(args.image) as image:
        tensor = preprocess(image.convert("RGB")).unsqueeze(0)

    with torch.inference_mode():
        probabilities = model(tensor).softmax(dim=1)[0]

    predicted = int(probabilities.argmax())
    result = {
        "class": CLASSES[predicted],
        "confidence": float(probabilities[predicted]),
        "probabilities": {
            name: float(probabilities[index])
            for index, name in enumerate(CLASSES)
        },
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()

Testează-l pe imaginea exportată din setul oficial:

python infer.py 
  --checkpoint artifacts/no_sync/model.pt 
  --image artifacts/no_sync/sample_test.png

Outputul va avea această structură, cu valori calculate din checkpointul local:

{
  "class": "<una dintre cele 10 clase>",
  "confidence": 0.xxxx,
  "probabilities": {
    "airplane": 0.xxxx,
    "automobile": 0.xxxx
  }
}

Dicționarul real va conține toate cele zece clase. Punct de control: probabilitățile trebuie să însumeze aproximativ 1; diferențele foarte mici sunt efectul reprezentării în virgulă mobilă.

Defecțiuni care arată unde s-a rupt acumularea distribuită

ValueError: expected world_size=2, got 1

Jobul a fost pornit cu python train.py sau cu un singur proces. Folosește lansatorul complet:

torchrun --standalone --nnodes=1 --nproc-per-node=2 train.py 
  --mode no_sync --output artifacts/no_sync

RuntimeError: 400 micro-batch-uri nu se împart exact la accum_steps=...

Fereastra a fost schimbată la o valoare care nu divide numărul de batch-uri per rang. Revino la configurația auditată:

rm -rf artifacts/every artifacts/no_sync
torchrun --standalone --nnodes=1 --nproc-per-node=2 train.py 
  --mode every --output artifacts/every --batch-size 50 --accum-steps 4

Nu elimina verificarea: o ultimă fereastră incompletă împărțită tot la patru ar micșora incorect gradientul acelei actualizări.

RuntimeError: ... expected 100 gradient syncs, got 400

Forward-ul a ieșit din contextul no_sync() sau rularea a folosit modul greșit. Contextul trebuie să închidă ambele operații:

with ddp_model.no_sync():
    logits = ddp_model(inputs)
    loss = criterion(logits, targets) / accum_steps
    loss.backward()

Acest fragment este forma corectă pentru un micro-batch intermediar, nu pentru ultimul din fereastră.

RuntimeError: CUDA error: out of memory

Unul dintre GPU-uri nu are suficientă memorie liberă sau este folosit de alt proces. Inspectează mai întâi:

nvidia-smi

Oprește numai procesele proprii identificate acolo. Nu reduce arbitrar batch-ul în timpul comparației, deoarece ai schimba numărul așteptat de micro-batch-uri și sincronizări.

torch.distributed.DistBackendError: NCCL error ...

Cele două procese nu pot inițializa ori menține colectivele NCCL. Confirmă că ambele GPU-uri sunt vizibile:

CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nnodes=1 
  --nproc-per-node=2 train.py 
  --mode no_sync --output artifacts/no_sync

Nu dezactiva NCCL și nu înlocui backendul doar pentru a ascunde eroarea; acest laborator măsoară comunicarea GPU.

RuntimeError: Deterministic behavior was enabled ...

PyTorch a întâlnit o operație CUDA fără implementare deterministă în mediul concret. Păstrează mesajul complet și verifică operația indicată. Nu elimina torch.use_deterministic_algorithms(True) dintr-o singură rulare: comparația ar avea politici de reproducibilitate diferite.

Ștergem laboratorul fără să atingem alte modele sau dataseturi

Comenzile următoare elimină numai mediul, CIFAR-10 și artefactele create în directorul proiectului. Verifică locația înainte de rm -rf:

cd ~/ddp-no-sync-lab
pwd
test "$PWD" = "$HOME/ddp-no-sync-lab"
rm -rf artifacts data .venv __pycache__
rm -f model.py train.py infer.py verify.py requirements.txt
cd ..
rmdir ~/ddp-no-sync-lab

Dacă vrei să păstrezi rezultatele, copiază mai întâi artifacts/no_sync/model.pt, report.json și history.json într-un director de experimente. Nu este necesară dezinstalarea driverului NVIDIA sau a toolkitului sistemului: wheel-urile proiectului au fost izolate în .venv.

Lista de control pentru un no_sync() implementat corect

☐ Jobul rulează cu exact două procese și câte un proces pe GPU.
☐ Fiecare rang primește 20.000 de exemple și 400 de micro-batch-uri pe epocă.
☐ Batch-ul efectiv este 400 de imagini.
☐ Pierderea fiecărui micro-batch este împărțită la patru înainte de backward().
☐ Forward-ul și backward-ul intermediar sunt ambele în contextul no_sync().
☐ Ultimul micro-batch al ferestrei rulează în afara contextului.
☐ optimizer.step() și zero_grad() apar numai după micro-batch-ul sincronizat.
☐ Fiecare epocă are 100 de pași ai optimizatorului.
☐ Modul every raportează 400 de sincronizări, iar no_sync raportează 100.
☐ Baseline-ul, acuratețea pe test și confuziile provin din report.json.

Următorul experiment util este profilarea ambelor comenzi cu PyTorch Profiler pentru a măsura timpul colectivelor NCCL, nu doar numărul lor. După aceea, aceeași fereastră poate fi combinată cu checkpointing de activări sau aplicată unui training cu secvențe împachetate. Schimbă o singură componentă la fiecare comparație și păstrează verificarea sincronizărilor ca test de regresie.

Întrebări frecvente

De ce nu este suficient să apelez optimizer.step() o dată la patru batch-uri?

Pentru că DDP sincronizează în mod normal gradientele după fiecare backward. Amânarea actualizării parametrilor produce acumulare, dar nu elimină cele trei sincronizări intermediare.

Pot pune numai backward() în interiorul no_sync()?

Nu. Forward-ul DDP trebuie inclus în același context, deoarece pregătește sincronizarea gradientelor folosită de backward.

De ce tutorialul cere exact două GPU-uri?

Numărul de micro-batch-uri, batch-ul efectiv și valorile așteptate de verificator sunt fixate pentru două ranguri. Schimbarea numărului de GPU-uri cere recalcularea întregului experiment.