"""A small GPT-style DDP run, instrumented with the Warpscale Python SDK.
Synthetic tokens only — no dataset download — so the run is reproducible on any
node with two GPUs.
"""
import argparse
import os
import time
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, Dataset, DistributedSampler
import warpscale
class SyntheticTokens(Dataset):
"""Random token ids built on the CPU worker so every batch costs a real
host-to-device copy, which is what the data-transfer features measure."""
def __init__(self, n_samples, seq_len, vocab_size, seed=0):
self.n_samples = n_samples
self.seq_len = seq_len
self.vocab_size = vocab_size
self.seed = seed
def __len__(self):
return self.n_samples
def __getitem__(self, idx):
g = torch.Generator().manual_seed(self.seed + idx)
ids = torch.randint(0, self.vocab_size, (self.seq_len + 1,), generator=g)
return ids[:-1], ids[1:]
class Block(nn.Module):
def __init__(self, d_model, n_head):
super().__init__()
self.n_head = n_head
self.ln1 = nn.LayerNorm(d_model)
self.qkv = nn.Linear(d_model, 3 * d_model)
self.proj = nn.Linear(d_model, d_model)
self.ln2 = nn.LayerNorm(d_model)
self.mlp = nn.Sequential(
nn.Linear(d_model, 4 * d_model),
nn.GELU(),
nn.Linear(4 * d_model, d_model),
)
def forward(self, x):
b, t, c = x.shape
q, k, v = self.qkv(self.ln1(x)).split(c, dim=2)
q, k, v = (
z.view(b, t, self.n_head, c // self.n_head).transpose(1, 2) for z in (q, k, v)
)
a = F.scaled_dot_product_attention(q, k, v, is_causal=True)
x = x + self.proj(a.transpose(1, 2).reshape(b, t, c))
return x + self.mlp(self.ln2(x))
class TinyGPT(nn.Module):
def __init__(self, vocab_size, seq_len, d_model, n_layer, n_head):
super().__init__()
self.tok = nn.Embedding(vocab_size, d_model)
self.pos = nn.Embedding(seq_len, d_model)
self.blocks = nn.ModuleList(Block(d_model, n_head) for _ in range(n_layer))
self.ln_f = nn.LayerNorm(d_model)
self.head = nn.Linear(d_model, vocab_size, bias=False)
def forward(self, idx):
pos = torch.arange(idx.shape[1], device=idx.device)
x = self.tok(idx) + self.pos(pos)
for block in self.blocks:
x = block(x)
return self.head(self.ln_f(x))
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--epochs", type=int, default=3)
p.add_argument("--steps-per-epoch", type=int, default=800)
p.add_argument("--batch-size", type=int, default=8, help="sequences per rank")
p.add_argument("--seq-len", type=int, default=1024)
p.add_argument("--vocab-size", type=int, default=32000)
p.add_argument("--d-model", type=int, default=768)
p.add_argument("--n-layer", type=int, default=12)
p.add_argument("--n-head", type=int, default=12)
p.add_argument("--num-workers", type=int, default=4)
p.add_argument("--log-every", type=int, default=50)
p.add_argument("--ckpt-dir", default="./ckpt")
return p.parse_args()
def save_checkpoint(model, optimizer, ckpt_dir, epoch):
"""Write one full checkpoint, overwriting the last so repeated epochs keep
exercising checkpoint I/O without growing on disk."""
os.makedirs(ckpt_dir, exist_ok=True)
path = os.path.join(ckpt_dir, "latest.pt")
tmp = path + ".tmp"
torch.save(
{
"epoch": epoch,
"model": model.module.state_dict(),
"optimizer": optimizer.state_dict(),
},
tmp,
)
os.replace(tmp, path)
return path
def main():
args = parse_args()
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
dist.init_process_group(backend="nccl", device_id=device)
rank = dist.get_rank()
world_size = dist.get_world_size()
torch.manual_seed(1234 + rank)
model = TinyGPT(
args.vocab_size, args.seq_len, args.d_model, args.n_layer, args.n_head
).to(device)
param_count = sum(p.numel() for p in model.parameters())
ddp = DistributedDataParallel(model, device_ids=[local_rank])
optimizer = torch.optim.AdamW(ddp.parameters(), lr=3e-4, fused=True)
# Elects rank 0 as the sole emitter, so it must follow init_process_group and
# precede the DataLoader whose workers fork.
warpscale.init(
total_steps=args.steps_per_epoch * args.epochs,
total_epochs=args.epochs,
global_batch_size=args.batch_size * world_size,
precision="bf16",
framework="ddp",
param_count=param_count,
flops_per_token=6 * param_count,
)
dataset = SyntheticTokens(
args.steps_per_epoch * args.batch_size * world_size, args.seq_len, args.vocab_size
)
sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, drop_last=True)
loader = DataLoader(
dataset,
batch_size=args.batch_size,
sampler=sampler,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True,
persistent_workers=args.num_workers > 0,
)
if rank == 0:
print(
f"[train] world_size={world_size} params={param_count/1e6:.1f}M "
f"global_batch={args.batch_size * world_size} seq_len={args.seq_len} "
f"run_id={os.environ.get('WS_RUN_ID', '<unset>')}",
flush=True,
)
started = time.perf_counter()
for epoch in range(args.epochs):
sampler.set_epoch(epoch)
epoch_started = time.perf_counter()
for step, (x, y) in enumerate(loader):
x = x.to(device, non_blocking=True)
y = y.to(device, non_blocking=True)
with torch.autocast("cuda", dtype=torch.bfloat16):
loss = F.cross_entropy(ddp(x).flatten(0, 1), y.flatten())
loss.backward()
optimizer.step()
optimizer.zero_grad(set_to_none=True)
warpscale.step()
if rank == 0 and (step + 1) % args.log_every == 0:
elapsed = time.perf_counter() - epoch_started
print(
f"[train] epoch {epoch} step {step + 1}/{args.steps_per_epoch} "
f"loss {loss.item():.3f} {(step + 1) / elapsed:.2f} steps/s",
flush=True,
)
warpscale.epoch()
warpscale.checkpoint_started()
if rank == 0:
path = save_checkpoint(ddp, optimizer, args.ckpt_dir, epoch)
print(f"[train] epoch {epoch} checkpoint -> {path}", flush=True)
dist.barrier()
warpscale.checkpoint_persisted()
if rank == 0:
print(f"[train] done in {time.perf_counter() - started:.1f}s", flush=True)
dist.destroy_process_group()
if __name__ == "__main__":
main()
