train_dataparallel.py

import time

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from src.data import RandomTextDataset
from src.model import TinyTransformerClassifier
from src.runtime import require_cuda

MAX_STEPS = 100


def main():
    require_cuda(min_devices=2)
    device = torch.device("cuda:0")
    world_size = torch.cuda.device_count()

    dataset = RandomTextDataset(num_samples=5000)
    loader = DataLoader(
        dataset,
        batch_size=64 * world_size,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
    )

    model = TinyTransformerClassifier()

    print("Using GPUs:", list(range(world_size)))
    print(f"effective batch size: {64 * world_size}")

    model = nn.DataParallel(model)
    model = model.to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
    criterion = nn.CrossEntropyLoss()

    model.train()
    start = time.time()

    for step, (input_ids, labels) in enumerate(loader):
        if step >= MAX_STEPS:
            break

        input_ids = input_ids.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)

        logits = model(input_ids)
        loss = criterion(logits, labels)

        loss.backward()
        optimizer.step()

        if step % 20 == 0:
            elapsed = time.time() - start
            print(f"step={step}, loss={loss.item():.4f}, elapsed={elapsed:.2f}s")


if __name__ == "__main__":
    main()