{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":3649,"databundleVersionId":46718}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"bf72b9f4-0747-4ca3-a6a7-193e26a784c9","cell_type":"code","source":"!pip install -q torch torchvision matplotlib numpy ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:30:19.996662Z","iopub.execute_input":"2026-05-03T01:30:19.997576Z","iopub.status.idle":"2026-05-03T01:30:23.339554Z","shell.execute_reply.started":"2026-05-03T01:30:19.997546Z","shell.execute_reply":"2026-05-03T01:30:23.338458Z"}},"outputs":[],"execution_count":null},{"id":"d415e100","cell_type":"code","source":"import os\nimport torch\n\nprint(\"torch:\", torch.__version__)\nprint(\"cuda available:\", torch.cuda.is_available())\nprint(\"gpu count:\", torch.cuda.device_count())\nif torch.cuda.is_available():\n print(\"gpu0:\", torch.cuda.get_device_name(0))\nprint(\"on kaggle:\", os.path.exists(\"/kaggle\"))\n\n# Optional dependency check (Kaggle has these preinstalled).\ntry:\n import torchvision, matplotlib, numpy # noqa: F401\n print(\"torchvision:\", torchvision.__version__)\nexcept Exception as e:\n print(\"Dependency import failed:\", repr(e))\n print(\"If you're on Kaggle, ensure the notebook uses a GPU runtime.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:30:23.341279Z","iopub.execute_input":"2026-05-03T01:30:23.341704Z","iopub.status.idle":"2026-05-03T01:30:23.348575Z","shell.execute_reply.started":"2026-05-03T01:30:23.341676Z","shell.execute_reply":"2026-05-03T01:30:23.347723Z"}},"outputs":[],"execution_count":null},{"id":"6f13950e","cell_type":"code","source":"!nvidia-smi -L","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:30:23.349399Z","iopub.execute_input":"2026-05-03T01:30:23.349700Z","iopub.status.idle":"2026-05-03T01:30:23.519342Z","shell.execute_reply.started":"2026-05-03T01:30:23.349678Z","shell.execute_reply":"2026-05-03T01:30:23.518384Z"}},"outputs":[],"execution_count":null},{"id":"9b9f072e","cell_type":"code","source":"%%writefile lab_distributed_training.py\n#!/usr/bin/env python\n# =============================================================================\n# Chapter 3 — Mini-Project Lab Work\n# \"Distributed Training in the Wild\"\n# AI Master's Programme — Distributed Computing for AI\n#\n# Target environment : Kaggle 2 × NVIDIA T4 (16 GB each)\n# Dataset : CIFAR-10 (real, from torchvision)\n# Framework : PyTorch ≥ 2.1\n#\n# Structure\n# ─────────\n# TASK 1 § Data Parallelism — DDP baseline + scaling analysis\n# TASK 2 § Model Parallelism — Pipeline parallelism (2 GPUs)\n# TASK 3 § Hybrid Parallelism — DP + Gradient Accumulation\n# TASK 4 § Profiling & Plotting — torch.profiler + matplotlib\n# TASK 5 § (Bonus) FSDP — Memory sharding\n#\n# Each task has:\n# • A theory reminder (links to lecture slides)\n# • A code scaffold with ### TODO sections for students to fill in\n# • Evaluation criteria\n# • Expected output / validation checklist\n#\n# Launch command (Kaggle terminal):\n# torchrun --nproc_per_node=2 lab_distributed_training.py --task 1\n# =============================================================================\n\nimport os, sys, time, math, argparse, csv, json\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data import DataLoader, Subset\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch import autocast\ntry:\n from torch.cuda.amp import GradScaler\nexcept Exception: # pragma: no cover\n from torch.cuda.amp import GradScaler\nfrom torch.profiler import profile, record_function, ProfilerActivity\n\ntry:\n from torch.distributed.pipeline.sync import Pipe\nexcept Exception: # pragma: no cover\n Pipe = None\n\nimport torchvision\nimport torchvision.transforms as T\ntorchvision.datasets.CIFAR10.url = \"https://data.brainchip.com/dataset-mirror/cifar10/cifar-10-python.tar.gz\"\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nimport numpy as np\n# ─── OUTPUT DIRECTORY ────────────────────────────────────────────────────────\nOUT_DIR = \"/kaggle/working\" if os.path.exists(\"/kaggle\") else \"/tmp/lab_output\"\nos.makedirs(OUT_DIR, exist_ok=True)\n\n# ─── SHARED HYPER-PARAMETERS ─────────────────────────────────────────────────\nHP = dict(\n data_dir = \"/tmp/cifar10\",\n num_classes = 10,\n hidden_dim = 256, # keep small so tasks run fast on T4\n num_layers = 4,\n micro_batch = 128,\n lr = 1e-3,\n epochs = 4,\n backend = \"nccl\",\n)\n\n\n# ─────────────────────────────────────────────────────────────────────────────\n# SHARED UTILITIES (pre-written; students should read but not modify)\n# ─────────────────────────────────────────────────────────────────────────────\n\ndef setup_dist(rank, world_size, backend=\"nccl\"):\n os.environ.setdefault(\"MASTER_ADDR\", \"localhost\")\n os.environ.setdefault(\"MASTER_PORT\", \"12399\")\n dist.init_process_group(backend=backend, rank=rank, world_size=world_size)\n torch.cuda.set_device(rank)\n\ndef teardown_dist():\n dist.destroy_process_group()\n\ndef get_cifar10(train, rank, world_size, batch_size=None):\n \"\"\"CIFAR-10 DataLoader with DistributedSampler.\"\"\"\n bs = batch_size or HP[\"micro_batch\"]\n os.makedirs(HP[\"data_dir\"], exist_ok=True)\n\n tf = T.Compose([\n T.RandomCrop(32, padding=4) if train else T.CenterCrop(32),\n T.RandomHorizontalFlip() if train else T.Lambda(lambda x: x),\n T.ToTensor(),\n T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n ])\n\n is_dist = dist.is_available() and dist.is_initialized() and world_size > 1\n if is_dist:\n if rank == 0:\n torchvision.datasets.CIFAR10(\n root=HP[\"data_dir\"],\n train=train,\n transform=tf,\n download=True,\n )\n dist.barrier()\n download = False\n else:\n download = True\n\n dataset = torchvision.datasets.CIFAR10(\n root=HP[\"data_dir\"],\n train=train,\n transform=tf,\n download=download,\n )\n\n sampler = (\n DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=train)\n if is_dist else None\n )\n\n return DataLoader(\n dataset,\n batch_size=bs,\n sampler=sampler,\n shuffle=(sampler is None and train),\n num_workers=2,\n pin_memory=True,\n )\n\n\n# ─── Simple ViT-mini (same as demo file) ─────────────────────────────────────\n\nclass PatchEmbed(nn.Module):\n def __init__(self, hidden=256):\n super().__init__()\n self.proj = nn.Conv2d(3, hidden, kernel_size=4, stride=4)\n self.cls = nn.Parameter(torch.randn(1, 1, hidden))\n self.pos = nn.Parameter(torch.randn(1, 65, hidden)) # 64 patches + 1 cls\n def forward(self, x):\n B = x.size(0)\n x = self.proj(x).flatten(2).transpose(1, 2)\n x = torch.cat([self.cls.expand(B, -1, -1), x], 1) + self.pos\n return x\n\nclass TBlock(nn.Module):\n def __init__(self, d=256, h=8):\n super().__init__()\n self.n1 = nn.LayerNorm(d)\n self.at = nn.MultiheadAttention(d, h, batch_first=True)\n self.n2 = nn.LayerNorm(d)\n self.ff = nn.Sequential(nn.Linear(d, 4*d), nn.GELU(), nn.Linear(4*d, d))\n def forward(self, x):\n y, _ = self.at(self.n1(x), self.n1(x), self.n1(x))\n x = x + y\n return x + self.ff(self.n2(x))\n\nclass CHead(nn.Module):\n def __init__(self, d=256, nc=10):\n super().__init__()\n self.n = nn.LayerNorm(d)\n self.h = nn.Linear(d, nc)\n def forward(self, x):\n return self.h(self.n(x[:, 0]))\n\ndef build_model(d=None, nl=None):\n d = d or HP[\"hidden_dim\"]\n nl = nl or HP[\"num_layers\"]\n return nn.Sequential(\n PatchEmbed(d),\n *[TBlock(d) for _ in range(nl)],\n CHead(d)\n )\n\n\n# ─── Logging utilities ───────────────────────────────────────────────────────\n\nclass Logger:\n \"\"\"Accumulates and persists training metrics to a CSV file.\"\"\"\n def __init__(self, name):\n self.name = name\n self.rows = []\n self.path = os.path.join(OUT_DIR, f\"{name}_log.csv\")\n\n def log(self, **kwargs):\n self.rows.append(kwargs)\n\n def save(self):\n if not self.rows:\n return\n with open(self.path, \"w\", newline=\"\") as f:\n w = csv.DictWriter(f, fieldnames=self.rows[0].keys())\n w.writeheader()\n w.writerows(self.rows)\n print(f\" [Log] {self.path}\")\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# TASK 1 — DATA PARALLELISM\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef train_one_epoch_task1(model, loader, optimizer, scaler, device,\n rank, world_size, epoch, accum_steps, logger):\n model.train()\n if hasattr(loader, \"sampler\") and hasattr(loader.sampler, \"set_epoch\"):\n loader.sampler.set_epoch(epoch)\n\n optimizer.zero_grad(set_to_none=True)\n total_loss, total_corr, total_samp = 0.0, 0, 0\n t0 = time.perf_counter()\n\n for step, (images, labels) in enumerate(loader):\n images = images.to(device, non_blocking=True)\n labels = labels.to(device, non_blocking=True)\n\n with autocast(device_type=device.type, enabled=(device.type == \"cuda\")):\n logits = model(images)\n loss_raw = F.cross_entropy(logits, labels)\n loss = loss_raw / accum_steps\n\n scaler.scale(loss).backward()\n\n do_step = ((step + 1) % accum_steps == 0) or ((step + 1) == len(loader))\n if do_step:\n scaler.unscale_(optimizer)\n nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n scaler.step(optimizer)\n scaler.update()\n optimizer.zero_grad(set_to_none=True)\n\n bs = images.size(0)\n total_samp += bs\n total_loss += loss_raw.detach().item() * bs\n total_corr += (logits.argmax(dim=1) == labels).sum().item()\n\n elapsed = time.perf_counter() - t0\n\n t_loss_sum = torch.tensor(total_loss, device=device, dtype=torch.float64)\n t_corr = torch.tensor(float(total_corr), device=device, dtype=torch.float64)\n t_samp = torch.tensor(float(total_samp), device=device, dtype=torch.float64)\n t_elapsed = torch.tensor(float(elapsed), device=device, dtype=torch.float64)\n\n if dist.is_available() and dist.is_initialized() and world_size > 1:\n dist.all_reduce(t_loss_sum, op=dist.ReduceOp.SUM)\n dist.all_reduce(t_corr, op=dist.ReduceOp.SUM)\n dist.all_reduce(t_samp, op=dist.ReduceOp.SUM)\n dist.all_reduce(t_elapsed, op=dist.ReduceOp.MAX)\n\n denom = max(t_samp.item(), 1.0)\n avg_loss = (t_loss_sum.item() / denom)\n acc = (t_corr.item() / denom) * 100.0\n throughput = denom / max(t_elapsed.item(), 1e-9)\n\n mem_mb = torch.cuda.max_memory_allocated(device) / 1e6 if device.type == \"cuda\" else 0.0\n\n if rank == 0:\n logger.log(\n epoch=int(epoch),\n loss=float(avg_loss),\n acc=float(acc),\n throughput=float(throughput),\n mem_mb=float(mem_mb),\n )\n print(\n f\" [Task1] Epoch {epoch} loss={avg_loss:.4f} acc={acc:.2f}% \"\n f\"throughput={throughput:.0f} samp/s mem={mem_mb:.0f}MB\"\n )\n\n return avg_loss, acc, throughput, mem_mb\n\n\ndef run_task1(rank, world_size, accum_steps=1):\n print(f\"\\n{'─'*55}\")\n print(f\" TASK 1 — Data Parallelism (rank={rank}/{world_size}, accum={accum_steps})\")\n print(f\"{'─'*55}\")\n\n setup_dist(rank, world_size, backend=HP[\"backend\"])\n device = torch.device(f\"cuda:{rank}\")\n\n model = build_model().to(device)\n ddp_model = DDP(model, device_ids=[rank], bucket_cap_mb=25)\n\n optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=HP[\"lr\"], weight_decay=1e-2)\n scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=HP[\"epochs\"])\n scaler = GradScaler(enabled=True)\n\n logger = Logger(f\"task1_r{rank}_a{accum_steps}\")\n\n train_loader = get_cifar10(True, rank, world_size)\n val_loader = get_cifar10(False, rank, world_size)\n\n for epoch in range(1, HP[\"epochs\"] + 1):\n train_one_epoch_task1(\n ddp_model,\n train_loader, optimizer, scaler,\n device, rank, world_size, epoch, accum_steps, logger\n )\n\n val_loss, val_acc = validate(ddp_model, val_loader, device, rank, world_size)\n if rank == 0:\n logger.rows[-1].update(val_loss=float(val_loss), val_acc=float(val_acc))\n print(f\" [Task1] val_loss={val_loss:.4f} val_acc={val_acc:.2f}%\")\n\n if rank == 0 and (epoch % 2 == 0):\n ckpt = {\n \"epoch\": int(epoch),\n \"model\": ddp_model.module.state_dict(),\n \"optimizer\": optimizer.state_dict(),\n }\n ckpt_path = os.path.join(OUT_DIR, f\"task1_ckpt_epoch{epoch}.pt\")\n torch.save(ckpt, ckpt_path)\n print(f\" [Ckpt] {ckpt_path}\")\n\n scheduler.step()\n\n if rank == 0:\n n_params = sum(p.numel() for p in ddp_model.module.parameters())\n avg_tp = float(np.mean([r[\"throughput\"] for r in logger.rows])) if logger.rows else 0.0\n mfu = compute_mfu(n_params, avg_tp) * 100.0\n print(f\" [Task1] Params={n_params:,} avg_throughput={avg_tp:.0f} samp/s MFU≈{mfu:.2f}%\")\n\n logger.save()\n plot_task1(logger)\n\n teardown_dist()\n\n\n@torch.no_grad()\ndef validate(model, loader, device, rank, world_size):\n model.eval()\n total_loss, total_corr, total_samp = 0.0, 0, 0\n for images, labels in loader:\n images = images.to(device, non_blocking=True)\n labels = labels.to(device, non_blocking=True)\n with autocast(device_type=device.type, enabled=(device.type == \"cuda\")):\n logits = model(images)\n loss = F.cross_entropy(logits, labels)\n\n bs = images.size(0)\n total_samp += bs\n total_loss += loss.detach().item() * bs\n total_corr += (logits.argmax(dim=1) == labels).sum().item()\n\n t_loss_sum = torch.tensor(total_loss, device=device, dtype=torch.float64)\n t_corr = torch.tensor(float(total_corr), device=device, dtype=torch.float64)\n t_samp = torch.tensor(float(total_samp), device=device, dtype=torch.float64)\n if dist.is_available() and dist.is_initialized() and world_size > 1:\n dist.all_reduce(t_loss_sum, op=dist.ReduceOp.SUM)\n dist.all_reduce(t_corr, op=dist.ReduceOp.SUM)\n dist.all_reduce(t_samp, op=dist.ReduceOp.SUM)\n\n denom = max(t_samp.item(), 1.0)\n avg_loss = t_loss_sum.item() / denom\n acc = (t_corr.item() / denom) * 100.0\n return float(avg_loss), float(acc)\n\n\ndef plot_task1(logger: Logger):\n if not logger.rows:\n print(\" [Task1] No data to plot.\")\n return\n\n epochs = [r[\"epoch\"] for r in logger.rows]\n losses = [r[\"loss\"] for r in logger.rows]\n accs = [r[\"acc\"] for r in logger.rows]\n throughputs = [r[\"throughput\"] for r in logger.rows]\n mems = [r[\"mem_mb\"] for r in logger.rows]\n\n fig, axes = plt.subplots(2, 2, figsize=(12, 8))\n fig.suptitle(\"Task 1 — Data Parallelism (2×T4)\", fontsize=13)\n\n ax = axes[0, 0]\n ax.plot(epochs, losses, \"o-\", label=\"train\")\n ax.set_title(\"Train loss\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Loss\")\n ax.grid(True, alpha=0.3)\n ax.legend()\n\n ax = axes[0, 1]\n ax.plot(epochs, accs, \"o-\", label=\"train\")\n if \"val_acc\" in logger.rows[0]:\n ax.plot(epochs, [r.get(\"val_acc\", float(\"nan\")) for r in logger.rows], \"o--\", label=\"val\")\n ax.set_title(\"Accuracy\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Accuracy (%)\")\n ax.grid(True, alpha=0.3)\n ax.legend()\n\n ax = axes[1, 0]\n ax.bar([str(e) for e in epochs], throughputs)\n ax.set_title(\"Throughput\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Samples / sec\")\n ax.grid(True, axis=\"y\", alpha=0.3)\n\n ax = axes[1, 1]\n ax.plot(epochs, mems, \"o-\", label=\"GPU-0 peak\")\n ax.axhline(16000, ls=\"--\", color=\"k\", alpha=0.6, label=\"T4 16 GB\")\n ax.set_title(\"Peak memory\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Memory (MB)\")\n ax.grid(True, alpha=0.3)\n ax.legend()\n\n plt.tight_layout()\n path = os.path.join(OUT_DIR, \"task1_results.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# TASK 2 — MODEL PARALLELISM (Pipeline)\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef build_pipeline_model_task2(num_chunks=4):\n half = HP[\"num_layers\"] // 2\n\n stage0 = nn.Sequential(\n PatchEmbed(HP[\"hidden_dim\"]),\n *[TBlock(HP[\"hidden_dim\"]) for _ in range(half)],\n ).to(\"cuda:0\")\n\n stage1 = nn.Sequential(\n *[TBlock(HP[\"hidden_dim\"]) for _ in range(HP[\"num_layers\"] - half)],\n CHead(HP[\"hidden_dim\"], HP[\"num_classes\"]),\n ).to(\"cuda:1\")\n\n if Pipe is None:\n class _TwoStage(nn.Module):\n def __init__(self, s0, s1):\n super().__init__()\n self.s0 = s0\n self.s1 = s1\n def forward(self, x):\n x = self.s0(x)\n x = x.to(\"cuda:1\")\n return self.s1(x)\n return _TwoStage(stage0, stage1)\n\n return Pipe(nn.Sequential(stage0, stage1), chunks=num_chunks)\n\n\ndef run_task2(num_chunks=4):\n print(f\"\\n{'─'*55}\")\n print(f\" TASK 2 — Pipeline Parallelism (chunks={num_chunks})\")\n print(f\"{'─'*55}\")\n\n if \"RANK\" in os.environ and int(os.environ[\"RANK\"]) != 0:\n print(\" [Task2] Pipeline runs on rank 0 only; skipping.\")\n return\n\n if torch.cuda.device_count() < 2:\n raise RuntimeError(\"Task 2 requires 2 CUDA GPUs (cuda:0 and cuda:1).\")\n\n import torch.distributed.rpc as rpc\n os.environ.setdefault(\"MASTER_ADDR\", \"localhost\")\n os.environ.setdefault(\"MASTER_PORT\", \"12400\")\n\n rpc.init_rpc(\"worker0\", rank=0, world_size=1)\n\n model = build_pipeline_model_task2(num_chunks=num_chunks)\n optimizer = torch.optim.AdamW(model.parameters(), lr=HP[\"lr\"], weight_decay=1e-2)\n\n tf = T.Compose([\n T.RandomCrop(32, padding=4), T.RandomHorizontalFlip(), T.ToTensor(),\n T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n ])\n dataset = torchvision.datasets.CIFAR10(\n root=HP[\"data_dir\"], train=True, transform=tf, download=True)\n\n loader = DataLoader(\n dataset,\n batch_size=HP[\"micro_batch\"],\n shuffle=True,\n num_workers=2,\n pin_memory=True,\n )\n\n val_loader = get_cifar10(False, rank=0, world_size=1)\n\n @torch.no_grad()\n def _validate_pipeline(pipe_model, vloader):\n pipe_model.eval()\n total_loss, total_corr, total_samp = 0.0, 0, 0\n for images, labels in vloader:\n images = images.to(\"cuda:0\", non_blocking=True)\n labels = labels.to(\"cuda:1\", non_blocking=True)\n out = pipe_model(images)\n out = out.local_value() if hasattr(out, \"local_value\") else out\n loss = F.cross_entropy(out, labels)\n\n bs = images.size(0)\n total_samp += bs\n total_loss += loss.detach().item() * bs\n total_corr += (out.argmax(dim=1) == labels).sum().item()\n\n avg_loss = total_loss / max(total_samp, 1)\n acc = (total_corr / max(total_samp, 1)) * 100.0\n return float(avg_loss), float(acc)\n\n gpu0_mem_history, gpu1_mem_history = [], []\n\n for epoch in range(1, HP[\"epochs\"] + 1):\n t0 = time.perf_counter()\n total_loss, total_corr, total_samp = 0.0, 0, 0\n\n model.train()\n for images, labels in loader:\n images = images.to(\"cuda:0\", non_blocking=True)\n labels = labels.to(\"cuda:1\", non_blocking=True)\n\n optimizer.zero_grad(set_to_none=True)\n out = model(images)\n out = out.local_value() if hasattr(out, \"local_value\") else out\n loss = F.cross_entropy(out, labels)\n loss.backward()\n optimizer.step()\n\n bs = images.size(0)\n total_samp += bs\n total_loss += loss.detach().item() * bs\n total_corr += (out.argmax(dim=1) == labels).sum().item()\n\n avg_loss = total_loss / max(total_samp, 1)\n train_acc = (total_corr / max(total_samp, 1)) * 100.0\n\n gpu0_mem = torch.cuda.max_memory_allocated(\"cuda:0\") / 1e6\n gpu1_mem = torch.cuda.max_memory_allocated(\"cuda:1\") / 1e6\n gpu0_mem_history.append(float(gpu0_mem))\n gpu1_mem_history.append(float(gpu1_mem))\n torch.cuda.reset_peak_memory_stats(\"cuda:0\")\n torch.cuda.reset_peak_memory_stats(\"cuda:1\")\n\n val_loss, val_acc = _validate_pipeline(model, val_loader)\n\n elapsed = time.perf_counter() - t0\n print(\n f\" [Task2] Epoch {epoch} loss={avg_loss:.4f} acc={train_acc:.2f}% \"\n f\"val_acc={val_acc:.2f}% GPU0={gpu0_mem:.0f}MB GPU1={gpu1_mem:.0f}MB t={elapsed:.1f}s\"\n )\n\n plot_task2_memory(gpu0_mem_history, gpu1_mem_history)\n plot_task2_bubble(max_chunks=8, used_chunks=num_chunks)\n\n rpc.shutdown()\n\n\ndef plot_task2_memory(gpu0_mem, gpu1_mem):\n epochs = list(range(1, len(gpu0_mem) + 1))\n x = np.arange(len(epochs))\n w = 0.38\n\n fig, ax = plt.subplots(figsize=(8, 4))\n ax.bar(x - w / 2, gpu0_mem, width=w, label=\"GPU-0\")\n ax.bar(x + w / 2, gpu1_mem, width=w, label=\"GPU-1\")\n ax.axhline(16000, ls=\"--\", color=\"k\", alpha=0.6, label=\"T4 16 GB\")\n ax.set_xticks(x)\n ax.set_xticklabels([str(e) for e in epochs])\n ax.set_title(\"Task 2 — Peak memory per epoch\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Peak memory (MB)\")\n ax.grid(True, axis=\"y\", alpha=0.3)\n ax.legend()\n\n plt.tight_layout()\n path = os.path.join(OUT_DIR, \"task2_memory.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n\ndef plot_task2_bubble(max_chunks=8, used_chunks=None):\n K = 2\n F_vals = list(range(1, max_chunks + 1))\n bubbles = [(K - 1) / (K - 1 + F) for F in F_vals]\n\n fig, ax = plt.subplots(figsize=(7, 4))\n ax.plot(F_vals, [b * 100 for b in bubbles], \"b-o\", markersize=6)\n ax.set_title(f\"Pipeline Bubble Fraction (K={K} stages)\")\n ax.set_xlabel(\"Number of micro-batches (F)\")\n ax.set_ylabel(\"Bubble fraction (%)\")\n if used_chunks is not None:\n ax.axvline(used_chunks, color=\"k\", ls=\"--\", alpha=0.6)\n ax.annotate(\n f\"used F={used_chunks}\",\n xy=(used_chunks, (K - 1) / (K - 1 + used_chunks) * 100),\n xytext=(used_chunks + 0.2, 60),\n arrowprops=dict(arrowstyle=\"->\", lw=1.0),\n )\n ax.grid(True, alpha=0.3)\n plt.tight_layout()\n path = os.path.join(OUT_DIR, \"task2_bubble.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# TASK 3 — HYBRID PARALLELISM\n# ═══════════════════════════════════════════════════════════════════════════════\n\nclass TopKSparseHook:\n def __init__(self, k_fraction=0.1):\n self.k_fraction = k_fraction\n self.total_elems = 0\n self.kept_elems = 0\n\n def hook(self, process_group, bucket):\n tensor = bucket.buffer()\n numel = tensor.numel()\n\n k = max(1, int(numel * self.k_fraction))\n with torch.no_grad():\n topk = torch.topk(tensor.abs(), k, largest=True, sorted=False)\n threshold = topk.values.min()\n mask = tensor.abs() >= threshold\n tensor.mul_(mask)\n\n self.total_elems += numel\n self.kept_elems += int(mask.sum().item())\n\n fut = dist.all_reduce(tensor, group=process_group, async_op=True).get_future()\n\n def _then(f):\n (t,) = f.value()\n if dist.is_available() and dist.is_initialized():\n ws = dist.get_world_size(group=process_group) if process_group is not None else dist.get_world_size()\n else:\n ws = 1\n t.div_(ws)\n return t\n\n return fut.then(_then)\n\n @property\n def compression_ratio(self):\n if self.total_elems == 0:\n return 0.0\n return self.kept_elems / self.total_elems\n\n\ndef run_task3(rank, world_size, accum_steps=4, use_compression=False):\n effective_batch = HP[\"micro_batch\"] * accum_steps * world_size\n print(f\"\\n{'─'*55}\")\n print(f\" TASK 3 — Hybrid Parallelism (rank={rank}/{world_size})\")\n print(f\" Effective batch = {HP['micro_batch']} × {accum_steps} × {world_size} = {effective_batch}\")\n print(f\"{'─'*55}\")\n\n setup_dist(rank, world_size, backend=HP[\"backend\"])\n device = torch.device(f\"cuda:{rank}\")\n\n model = build_model().to(device)\n ddp_model = DDP(model, device_ids=[rank], bucket_cap_mb=25)\n\n hook_obj = None\n if use_compression:\n hook_obj = TopKSparseHook(k_fraction=0.1)\n ddp_model.register_comm_hook(None, hook_obj.hook)\n\n optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=HP[\"lr\"], weight_decay=1e-2)\n scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=HP[\"epochs\"])\n scaler = GradScaler(enabled=True)\n train_loader = get_cifar10(True, rank, world_size)\n val_loader = get_cifar10(False, rank, world_size)\n\n logger = Logger(f\"task3_r{rank}_a{accum_steps}_comp{use_compression}\")\n\n for epoch in range(1, HP[\"epochs\"] + 1):\n if hook_obj is not None:\n hook_obj.total_elems = 0\n hook_obj.kept_elems = 0\n\n train_one_epoch_task1(\n ddp_model, train_loader, optimizer, scaler,\n device, rank, world_size, epoch, accum_steps, logger\n )\n val_loss, val_acc = validate(ddp_model, val_loader, device, rank, world_size)\n\n if rank == 0:\n logger.rows[-1].update(\n val_loss=float(val_loss),\n val_acc=float(val_acc),\n effective_batch=int(effective_batch),\n compression_ratio=float(hook_obj.compression_ratio if hook_obj is not None else 0.0),\n )\n if use_compression:\n print(f\" [Compression] kept/total={hook_obj.compression_ratio:.3f}\")\n\n scheduler.step()\n\n if rank == 0:\n logger.save()\n plot_task3(logger, effective_batch, use_compression)\n\n teardown_dist()\n\n\ndef plot_task3(logger: Logger, effective_batch: int, compressed: bool):\n if not logger.rows:\n return\n\n epochs = [r[\"epoch\"] for r in logger.rows]\n losses = [r[\"loss\"] for r in logger.rows]\n accs = [r[\"acc\"] for r in logger.rows]\n val_accs = [r.get(\"val_acc\", float(\"nan\")) for r in logger.rows]\n\n fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n title = f\"Task 3 — Hybrid (effective_batch={effective_batch})\"\n if compressed:\n title += \" + TopK\"\n fig.suptitle(title, fontsize=12)\n\n ax = axes[0]\n ax.plot(epochs, losses, \"o-\", label=\"train\")\n ax.set_title(\"Loss\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Loss\")\n ax.grid(True, alpha=0.3)\n ax.legend()\n\n ax = axes[1]\n ax.plot(epochs, accs, \"o-\", label=\"train\")\n ax.plot(epochs, val_accs, \"o--\", label=\"val\")\n ax.set_title(\"Accuracy\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Accuracy (%)\")\n ax.grid(True, alpha=0.3)\n ax.legend()\n\n plt.tight_layout()\n suffix = \"_topk\" if compressed else \"\"\n path = os.path.join(OUT_DIR, f\"task3_eb{effective_batch}{suffix}.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n\ndef effective_batch_sweep_plot(results: dict):\n if not results:\n return\n\n xs = sorted(results.keys())\n ys = [results[x] * 100.0 for x in xs]\n best_x = max(xs, key=lambda x: results[x])\n best_y = results[best_x] * 100.0\n\n fig, ax = plt.subplots(figsize=(7, 4))\n ax.plot(xs, ys, \"o-\")\n ax.set_xscale(\"log\")\n ax.set_title(\"Effective batch sweep\")\n ax.set_xlabel(\"Effective batch size (log scale)\")\n ax.set_ylabel(\"Final val accuracy (%)\")\n ax.grid(True, which=\"both\", alpha=0.3)\n ax.annotate(\n f\"best={best_x} ({best_y:.1f}%)\",\n xy=(best_x, best_y),\n xytext=(best_x * 1.15, best_y - 3),\n arrowprops=dict(arrowstyle=\"->\", lw=1.0),\n )\n\n plt.tight_layout()\n path = os.path.join(OUT_DIR, \"task3_effective_batch_sweep.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# TASK 4 — PROFILING AND MONITORING\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef run_task4_profiler(rank, world_size):\n print(f\"\\n{'─'*55}\")\n print(f\" TASK 4 — Profiling (rank={rank}/{world_size})\")\n print(f\"{'─'*55}\")\n\n setup_dist(rank, world_size)\n device = torch.device(f\"cuda:{rank}\")\n\n model = build_model().to(device)\n ddp_model = DDP(model, device_ids=[rank])\n optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=HP[\"lr\"])\n scaler = GradScaler(enabled=True)\n loader = get_cifar10(True, rank, world_size)\n\n prof_dir = os.path.join(OUT_DIR, \"prof_trace\")\n os.makedirs(prof_dir, exist_ok=True)\n\n import gzip\n import shutil\n import torch.profiler as tprof\n\n schedule = tprof.schedule(wait=1, warmup=1, active=3, repeat=1)\n\n with profile(\n activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],\n schedule=schedule,\n on_trace_ready=None,\n record_shapes=True,\n profile_memory=True,\n ) as prof:\n ddp_model.train()\n it = iter(loader)\n for _step in range(6):\n images, labels = next(it)\n images = images.to(device, non_blocking=True)\n labels = labels.to(device, non_blocking=True)\n\n optimizer.zero_grad(set_to_none=True)\n with autocast(device_type=\"cuda\", enabled=True):\n with record_function(\"model_forward\"):\n logits = ddp_model(images)\n with record_function(\"loss\"):\n loss = F.cross_entropy(logits, labels)\n with record_function(\"backward\"):\n scaler.scale(loss).backward()\n with record_function(\"optimizer\"):\n scaler.step(optimizer)\n scaler.update()\n\n prof.step()\n\n if rank == 0:\n chrome_path = os.path.join(prof_dir, f\"trace_rank{rank}.json\")\n prof.export_chrome_trace(chrome_path)\n gz_path = chrome_path + \".gz\"\n with open(chrome_path, \"rb\") as f_in, gzip.open(gz_path, \"wb\") as f_out:\n shutil.copyfileobj(f_in, f_out)\n print(f\" [Trace] {gz_path}\")\n print(prof.key_averages().table(sort_by=\"cuda_time_total\", row_limit=20))\n\n plot_task4_operator_breakdown(prof, rank)\n\n teardown_dist()\n\n\ndef plot_task4_operator_breakdown(prof, rank):\n if rank != 0:\n return\n avgs = prof.key_averages()\n top = sorted(avgs, key=lambda x: x.cuda_time_total, reverse=True)[:10]\n\n names = [e.key for e in reversed(top)]\n times_ms = [e.cuda_time_total / 1000.0 for e in reversed(top)]\n\n def _cat(name: str) -> str:\n n = name.lower()\n if \"attention\" in n or \"scaled_dot_product\" in n:\n return \"attn\"\n if \"allreduce\" in n or \"nccl\" in n or \"all_gather\" in n or \"reduce\" in n:\n return \"comms\"\n if \"layer_norm\" in n or \"norm\" in n:\n return \"norm\"\n if \"matmul\" in n or \"mm\" in n or \"addmm\" in n or \"linear\" in n or \"gemm\" in n:\n return \"linear\"\n return \"other\"\n\n color_map = {\n \"attn\": \"tab:blue\",\n \"linear\": \"tab:orange\",\n \"norm\": \"tab:green\",\n \"comms\": \"tab:red\",\n \"other\": \"tab:gray\",\n }\n colors = [color_map[_cat(n)] for n in names]\n\n fig, ax = plt.subplots(figsize=(10, 5))\n y = np.arange(len(names))\n ax.barh(y, times_ms, color=colors)\n ax.set_yticks(y)\n ax.set_yticklabels(names, fontsize=8)\n ax.set_title(\"Task 4 — Top-10 ops by CUDA time\")\n ax.set_xlabel(\"CUDA time total (ms)\")\n ax.grid(True, axis=\"x\", alpha=0.3)\n plt.tight_layout()\n path = os.path.join(OUT_DIR, \"task4_op_breakdown.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n\ndef compute_mfu(n_params, avg_throughput_samp_per_s,\n peak_flops=65e12, flops_factor=6):\n flops_per_sample = flops_factor * n_params\n mfu = avg_throughput_samp_per_s * flops_per_sample / peak_flops\n return mfu\n\n\ndef plot_task4_mfu_comparison(mfu_dict: dict):\n if not mfu_dict:\n return\n\n labels = list(mfu_dict.keys())\n vals = [float(mfu_dict[k]) for k in labels]\n\n fig, ax = plt.subplots(figsize=(8, 4))\n ax.bar(labels, vals)\n ax.set_ylabel(\"MFU (%)\")\n ax.set_title(\"MFU comparison\")\n ax.grid(True, axis=\"y\", alpha=0.3)\n\n for name, v in [(\"PaLM-540B\", 46.2), (\"BLOOM-176B\", 33.0), (\"LLaMA-3-405B\", 41.0)]:\n ax.axhline(v, ls=\"--\", alpha=0.5)\n ax.text(0.01, v + 0.5, name, transform=ax.get_yaxis_transform(), fontsize=8)\n\n plt.tight_layout()\n path = os.path.join(OUT_DIR, \"task4_mfu.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# TASK 5 (BONUS) — PyTorch FSDP\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef run_task5_fsdp(rank, world_size):\n print(f\"\\n{'─'*55}\")\n print(f\" TASK 5 (Bonus) — FSDP (rank={rank}/{world_size})\")\n print(f\"{'─'*55}\")\n\n setup_dist(rank, world_size, backend=HP[\"backend\"])\n device = torch.device(f\"cuda:{rank}\")\n\n from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n from torch.distributed.fsdp import MixedPrecision\n from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy\n import functools\n\n cap_major, _cap_minor = torch.cuda.get_device_capability(rank)\n mp_dtype = torch.bfloat16 if cap_major >= 8 else torch.float16\n\n def _train(model, optimizer, scaler, train_loader, val_loader, tag):\n logger = Logger(f\"task5_{tag}_r{rank}\")\n mem_hist = []\n for epoch in range(1, HP[\"epochs\"] + 1):\n torch.cuda.reset_peak_memory_stats(device)\n train_one_epoch_task1(\n model, train_loader, optimizer, scaler,\n device, rank, world_size, epoch, accum_steps=1, logger=logger\n )\n val_loss, val_acc = validate(model, val_loader, device, rank, world_size)\n if rank == 0:\n logger.rows[-1].update(val_loss=float(val_loss), val_acc=float(val_acc))\n mem_hist.append(torch.cuda.max_memory_allocated(device) / 1e6)\n return logger, mem_hist\n\n train_loader = get_cifar10(True, rank, world_size)\n val_loader = get_cifar10(False, rank, world_size)\n\n ddp_base = DDP(build_model().to(device), device_ids=[rank], bucket_cap_mb=25)\n ddp_opt = torch.optim.AdamW(ddp_base.parameters(), lr=HP[\"lr\"], weight_decay=1e-2)\n ddp_scaler = GradScaler(enabled=True)\n ddp_logger, ddp_mem = _train(ddp_base, ddp_opt, ddp_scaler, train_loader, val_loader, tag=\"ddp\")\n\n auto_wrap = functools.partial(size_based_auto_wrap_policy, min_num_params=100_000)\n fsdp_model = FSDP(\n build_model().to(device),\n auto_wrap_policy=auto_wrap,\n mixed_precision=MixedPrecision(param_dtype=mp_dtype, reduce_dtype=mp_dtype, buffer_dtype=mp_dtype),\n device_id=rank,\n )\n fsdp_opt = torch.optim.AdamW(fsdp_model.parameters(), lr=HP[\"lr\"], weight_decay=1e-2)\n fsdp_scaler = GradScaler(enabled=(mp_dtype == torch.float16))\n fsdp_logger, fsdp_mem = _train(fsdp_model, fsdp_opt, fsdp_scaler, train_loader, val_loader, tag=\"fsdp\")\n\n if rank == 0:\n ddp_logger.save()\n fsdp_logger.save()\n\n epochs = list(range(1, HP[\"epochs\"] + 1))\n x = np.arange(len(epochs))\n w = 0.38\n fig, ax = plt.subplots(figsize=(8, 4))\n ax.bar(x - w / 2, ddp_mem, width=w, label=\"DDP\")\n ax.bar(x + w / 2, fsdp_mem, width=w, label=\"FSDP\")\n ax.axhline(16000, ls=\"--\", color=\"k\", alpha=0.6, label=\"T4 16 GB\")\n ax.set_xticks(x)\n ax.set_xticklabels([str(e) for e in epochs])\n ax.set_title(\"Task 5 — Peak memory: DDP vs FSDP\")\n ax.set_xlabel(\"Epoch\")\n ax.set_ylabel(\"Peak memory (MB)\")\n ax.grid(True, axis=\"y\", alpha=0.3)\n ax.legend()\n plt.tight_layout()\n path = os.path.join(OUT_DIR, \"task5_ddp_vs_fsdp_memory.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\")\n plt.close()\n print(f\" [Plot] {path}\")\n\n teardown_dist()\n\n\nSTRATEGY_TABLE_TEMPLATE = \"\"\"\n╔══════════════════════╦═══════════╦══════════════╦═══════════════════╦═══════════╗\n║ Strategy ║ Val Acc ║ Throughput ║ Peak GPU Mem (MB) ║ MFU (%) ║\n╠══════════════════════╬═══════════╬══════════════╬═══════════════════╬═══════════╣\n║ DDP (baseline) ║ ? ║ ? ║ ? ║ ? ║\n║ DDP + Grad Accum ║ ? ║ ? ║ ? ║ ? ║\n║ Pipeline (chunks=4) ║ ? ║ ? ║ ? ║ ? ║\n║ Hybrid (DDP+GA) ║ ? ║ ? ║ ? ║ ? ║\n║ FSDP (bonus) ║ ? ║ ? ║ ? ║ ? ║\n╚══════════════════════╩═══════════╩══════════════╩═══════════════════╩═══════════╝\n\"\"\"\n\n\ndef spawn_fn(fn, world_size, extra_args=()):\n if \"RANK\" in os.environ:\n rank = int(os.environ[\"RANK\"])\n fn(rank, world_size, *extra_args)\n else:\n mp.spawn(fn, args=(world_size, *extra_args), nprocs=world_size, join=True)\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Chapter-3 Lab Work\")\n parser.add_argument(\"--task\", type=int, default=1, choices=[1,2,3,4,5],\n help=\"Which task to run (1-5)\")\n parser.add_argument(\"--accum_steps\", type=int, default=1)\n parser.add_argument(\"--chunks\", type=int, default=4,\n help=\"Pipeline micro-batch chunks (Task 2)\")\n parser.add_argument(\"--compress\", action=\"store_true\",\n help=\"Enable Top-K gradient compression (Task 3)\")\n args = parser.parse_args()\n\n world_size = torch.cuda.device_count()\n if world_size == 0:\n print(\"⚠ No CUDA GPUs found. Enable 2×T4 in Kaggle settings.\")\n return\n print(f\"\\n GPUs: {world_size} × {torch.cuda.get_device_name(0)}\")\n print(f\" Task: {args.task}\\n\")\n\n if args.task == 1:\n spawn_fn(run_task1, world_size, extra_args=(args.accum_steps,))\n\n elif args.task == 2:\n run_task2(num_chunks=args.chunks)\n\n elif args.task == 3:\n spawn_fn(run_task3, world_size,\n extra_args=(args.accum_steps, args.compress))\n\n elif args.task == 4:\n spawn_fn(run_task4_profiler, world_size)\n\n elif args.task == 5:\n spawn_fn(run_task5_fsdp, world_size)\n\n print(f\"\\n ✓ Task {args.task} done. Results in {OUT_DIR}\")\n print(STRATEGY_TABLE_TEMPLATE)\n\n\nif __name__ == \"__main__\":\n main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:30:23.522251Z","iopub.execute_input":"2026-05-03T01:30:23.522498Z","iopub.status.idle":"2026-05-03T01:30:23.544425Z","shell.execute_reply.started":"2026-05-03T01:30:23.522472Z","shell.execute_reply":"2026-05-03T01:30:23.543816Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"id":"da48d178","cell_type":"code","source":"# Task 1 — DDP baseline (accum_steps=1)\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 1 --accum_steps 1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:30:23.545462Z","iopub.execute_input":"2026-05-03T01:30:23.545853Z","iopub.status.idle":"2026-05-03T01:31:39.214791Z","shell.execute_reply.started":"2026-05-03T01:30:23.545808Z","shell.execute_reply":"2026-05-03T01:31:39.214008Z"}},"outputs":[],"execution_count":null},{"id":"d11fe149","cell_type":"code","source":"# Task 1 — DDP with gradient accumulation (accum_steps=4)\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 1 --accum_steps 4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:31:39.216271Z","iopub.execute_input":"2026-05-03T01:31:39.217037Z","iopub.status.idle":"2026-05-03T01:32:44.667387Z","shell.execute_reply.started":"2026-05-03T01:31:39.217006Z","shell.execute_reply":"2026-05-03T01:32:44.666570Z"}},"outputs":[],"execution_count":null},{"id":"b9c64dbb","cell_type":"code","source":"# Task 2 — Pipeline Parallelism (chunks=4)\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 2 --chunks 4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:32:44.668875Z","iopub.execute_input":"2026-05-03T01:32:44.669258Z","iopub.status.idle":"2026-05-03T01:34:38.030394Z","shell.execute_reply.started":"2026-05-03T01:32:44.669227Z","shell.execute_reply":"2026-05-03T01:34:38.029536Z"},"scrolled":true},"outputs":[],"execution_count":null},{"id":"eacb237a","cell_type":"code","source":"# Task 2 — Pipeline Parallelism (chunks=8)\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 2 --chunks 8","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:34:38.031813Z","iopub.execute_input":"2026-05-03T01:34:38.032252Z","iopub.status.idle":"2026-05-03T01:36:32.603616Z","shell.execute_reply.started":"2026-05-03T01:34:38.032220Z","shell.execute_reply":"2026-05-03T01:36:32.602885Z"}},"outputs":[],"execution_count":null},{"id":"e3e9ba49","cell_type":"code","source":"# Task 3 — Hybrid (DDP + grad accumulation, no compression)\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 3 --accum_steps 4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:36:32.605170Z","iopub.execute_input":"2026-05-03T01:36:32.605514Z","iopub.status.idle":"2026-05-03T01:37:37.273216Z","shell.execute_reply.started":"2026-05-03T01:36:32.605474Z","shell.execute_reply":"2026-05-03T01:37:37.272268Z"}},"outputs":[],"execution_count":null},{"id":"42d27499","cell_type":"code","source":"# Task 3 — Hybrid + Top-K gradient compression\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 3 --accum_steps 4 --compress","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:37:37.276146Z","iopub.execute_input":"2026-05-03T01:37:37.276706Z","iopub.status.idle":"2026-05-03T01:38:42.581814Z","shell.execute_reply.started":"2026-05-03T01:37:37.276675Z","shell.execute_reply":"2026-05-03T01:38:42.580814Z"}},"outputs":[],"execution_count":null},{"id":"3a623b54","cell_type":"code","source":"# Task 4 — Profiler (+ trace export + plots)\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:38:42.583206Z","iopub.execute_input":"2026-05-03T01:38:42.583537Z","iopub.status.idle":"2026-05-03T01:38:56.603553Z","shell.execute_reply.started":"2026-05-03T01:38:42.583508Z","shell.execute_reply":"2026-05-03T01:38:56.602506Z"}},"outputs":[],"execution_count":null},{"id":"3e7fba64","cell_type":"code","source":"# Task 5 (Bonus) — FSDP (may take longer / use more memory)\n!torchrun --nproc_per_node=2 lab_distributed_training.py --task 5","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:38:56.605047Z","iopub.execute_input":"2026-05-03T01:38:56.605412Z","iopub.status.idle":"2026-05-03T01:41:13.390929Z","shell.execute_reply.started":"2026-05-03T01:38:56.605354Z","shell.execute_reply":"2026-05-03T01:41:13.390129Z"}},"outputs":[],"execution_count":null},{"id":"6c88ca77","cell_type":"code","source":"# Show artifacts saved by the lab script\nimport os, glob\nfrom IPython.display import display, Image\n\nout_dir = \"/kaggle/working\" if os.path.exists(\"/kaggle\") else \"/tmp/lab_output\"\nprint(\"OUT_DIR:\", out_dir)\nprint(\"Files:\")\nfor p in sorted(glob.glob(os.path.join(out_dir, \"*\"))):\n print(\" -\", os.path.basename(p))\n\npngs = sorted(glob.glob(os.path.join(out_dir, \"task*.png\")))\nif not pngs:\n print(\"No task plots found yet. Run one of the tasks above.\")\nelse:\n for p in pngs:\n print(\"\\n\", os.path.basename(p))\n display(Image(filename=p))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:41:13.392488Z","iopub.execute_input":"2026-05-03T01:41:13.392854Z","iopub.status.idle":"2026-05-03T01:41:13.423035Z","shell.execute_reply.started":"2026-05-03T01:41:13.392807Z","shell.execute_reply":"2026-05-03T01:41:13.422458Z"}},"outputs":[],"execution_count":null},{"id":"432e71da","cell_type":"code","source":"%%writefile chapter3_kaggle_2xT4_demo.py\n#!/usr/bin/env python\n# =============================================================================\n# Chapter 3 — Distributed Training Algorithms & Strategies\n# Complete Hands-On Demo · Kaggle 2 × NVIDIA T4 (16 GB each)\n#\n# ╔══════════════════════════════════════════════════════════════════╗\n# ║ FIXED for PyTorch 2.x — three breaking changes corrected: ║\n# ║ 1. torch.distributed.pipeline (Pipe) REMOVED in PyTorch 2.0 ║\n# ║ → replaced with a manual 2-stage .to(\"cuda:N\") pipeline ║\n# ║ 2. torch.cuda.amp.GradScaler deprecated in PyTorch 2.x ║\n# ║ → replaced with torch.amp.GradScaler(\"cuda\", ...) ║\n# ║ 3. torch.distributed.rpc (required by old Pipe) removed ║\n# ║ → no RPC needed in the manual pipeline ║\n# ╚══════════════════════════════════════════════════════════════════╝\n#\n# Topics covered (mirrors the lecture slides):\n# §3.1 Data Parallelism — DDP + Ring-AllReduce + Gradient Accumulation\n# §3.2 Model Parallelism — Manual 2-GPU Pipeline (GPipe-style)\n# — Tensor-Parallel sketch (column-parallel linear)\n# §3.3 Hybrid Parallelism — DDP + Gradient Accumulation + PowerSGD\n# §3.4 Fault Tolerance — Checkpointing, torch.profiler, MFU estimation\n#\n# Real dataset : CIFAR-10 (auto-downloaded via torchvision)\n#\n# HOW TO RUN IN KAGGLE\n# ─────────────────────────────────────────────────────────────────────\n# Settings → Accelerator → \"2 x GPU T4\" (REQUIRED)\n#\n# Cell 0 – verify\n# !python -c \"import torch; print(torch.__version__, torch.cuda.device_count(), 'GPUs')\"\n#\n# Cell 1 – §3.1 Data Parallelism (DDP)\n# !torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode ddp\n#\n# Cell 2 – §3.2 Pipeline Parallelism (manual, no Pipe module needed)\n# !torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode pipeline\n#\n# Cell 3 – §3.3 Hybrid (DDP + Gradient Accumulation + PowerSGD)\n# !torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode hybrid\n#\n# Cell 4 – §3.2 Tensor-Parallel correctness check\n# !torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode tensor_parallel\n#\n# Cell 5 – §3.4 Profiling\n# !torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode profile\n#\n# Cell 6 – Everything in one run (3 epochs)\n# !torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode all --epochs 3\n# =============================================================================\n\nimport os, sys, time, argparse, datetime\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch import autocast\n\n# ── FIX 2: correct GradScaler import for PyTorch 2.x ────────────────────────\n# WRONG (crashes in PyTorch 2.x): torch.cuda.amp → GradScaler was moved\n# RIGHT:\nfrom torch.amp import GradScaler\n\nfrom torch.profiler import profile, record_function, ProfilerActivity\n\nimport torchvision\nimport torchvision.transforms as T\n\nimport matplotlib\nmatplotlib.use(\"Agg\") # headless – Kaggle has no display\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 0. GLOBAL CONFIGURATION\n# ═══════════════════════════════════════════════════════════════════════════════\n\nCFG = dict(\n # ── Dataset ───────────────────────────────────────────────────────────────\n data_dir = \"/tmp/cifar10\",\n num_classes = 10,\n img_size = 32,\n\n # ── Model ─────────────────────────────────────────────────────────────────\n hidden_dim = 512, # Transformer width\n num_layers = 6, # total Transformer blocks (3+3 for pipeline)\n\n # ── Training ──────────────────────────────────────────────────────────────\n epochs = 5,\n micro_batch = 64, # per-GPU batch size\n accum_steps = 4, # gradient-accumulation steps\n # effective_batch = micro_batch x accum_steps x world_size = 64x4x2 = 512\n lr = 3e-4,\n weight_decay = 1e-2,\n\n # ── Communication ─────────────────────────────────────────────────────────\n backend = \"nccl\", # fastest on T4 (GPU-Direct RDMA via NCCL)\n bucket_mb = 25, # DDP AllReduce bucket size (MB)\n\n # ── Mixed precision ───────────────────────────────────────────────────────\n use_amp = True, # FP16 autocast (T4 Tensor Cores)\n\n # ── Pipeline ──────────────────────────────────────────────────────────────\n pipe_chunks = 4, # GPipe micro-batches per global batch\n # bubble fraction = (K-1)/(K-1+F) with K=2, F=4 → 1/5 = 20%\n\n # ── Checkpointing ─────────────────────────────────────────────────────────\n ckpt_dir = \"/tmp/ckpts\",\n save_every = 2,\n\n # ── Profiling ─────────────────────────────────────────────────────────────\n profile_steps = 3,\n profile_dir = \"/tmp/prof\",\n\n # ── Output ────────────────────────────────────────────────────────────────\n plot_dir = \"/kaggle/working\" if os.path.exists(\"/kaggle\") else \"/tmp\",\n)\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 1. DATASET HELPERS\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef get_cifar10(train: bool, rank: int, world_size: int, batch_size: int = None):\n bs = batch_size or CFG[\"micro_batch\"]\n tf = T.Compose([\n T.RandomCrop(32, padding=4) if train else T.CenterCrop(32),\n T.RandomHorizontalFlip() if train else T.Lambda(lambda x: x),\n T.ToTensor(),\n T.Normalize((0.4914, 0.4822, 0.4465),\n (0.2470, 0.2435, 0.2616)),\n ])\n\n if dist.is_available() and dist.is_initialized() and world_size > 1:\n if rank == 0:\n # download once\n torchvision.datasets.CIFAR10(\n root=CFG[\"data_dir\"],\n train=train,\n transform=tf,\n download=True,\n )\n dist.barrier()\n\n dataset = torchvision.datasets.CIFAR10(\n root=CFG[\"data_dir\"],\n train=train,\n transform=tf,\n download=False,\n )\n\n sampler = DistributedSampler(\n dataset, num_replicas=world_size, rank=rank,\n shuffle=train, drop_last=train,\n ) if world_size > 1 else None\n\n return DataLoader(\n dataset,\n batch_size=bs,\n sampler=sampler,\n shuffle=(sampler is None and train),\n num_workers=2,\n pin_memory=True,\n persistent_workers=True,\n )\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 2. MODEL — ViT-mini\n# PatchEmbedding → N x TransformerBlock → ClassificationHead\n# Kept as nn.Sequential so it can be trivially split for pipeline.\n# ═══════════════════════════════════════════════════════════════════════════════\n\nclass PatchEmbedding(nn.Module):\n \"\"\"\n Tokenise a 32x32 image into 64 patch tokens + 1 [CLS] token.\n patch_size=4 → (32/4)^2 = 64 patches, each projected to hidden_dim.\n \"\"\"\n def __init__(self, img_size=32, patch_size=4, in_ch=3, hidden_dim=512):\n super().__init__()\n self.n_patches = (img_size // patch_size) ** 2\n self.proj = nn.Conv2d(in_ch, hidden_dim,\n kernel_size=patch_size, stride=patch_size)\n self.cls_token = nn.Parameter(torch.randn(1, 1, hidden_dim))\n self.pos_embed = nn.Parameter(torch.randn(1, self.n_patches + 1, hidden_dim))\n\n def forward(self, x):\n B = x.size(0)\n x = self.proj(x).flatten(2).transpose(1, 2) # (B, 64, D)\n cls = self.cls_token.expand(B, -1, -1) # (B, 1, D)\n x = torch.cat([cls, x], dim=1) + self.pos_embed # (B, 65, D)\n return x\n\n\nclass TransformerBlock(nn.Module):\n \"\"\"\n Pre-LayerNorm Transformer block:\n LN → MultiHeadAttention → residual\n LN → FFN (4x expansion, GELU) → residual\n \"\"\"\n def __init__(self, hidden_dim=512, n_heads=8, dropout=0.1):\n super().__init__()\n self.norm1 = nn.LayerNorm(hidden_dim)\n self.attn = nn.MultiheadAttention(\n hidden_dim, n_heads, dropout=dropout, batch_first=True)\n self.norm2 = nn.LayerNorm(hidden_dim)\n self.ffn = nn.Sequential(\n nn.Linear(hidden_dim, 4 * hidden_dim), nn.GELU(), nn.Dropout(dropout),\n nn.Linear(4 * hidden_dim, hidden_dim), nn.Dropout(dropout),\n )\n\n def forward(self, x):\n y, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))\n x = x + y\n return x + self.ffn(self.norm2(x))\n\n\nclass ClassificationHead(nn.Module):\n \"\"\"Pool the [CLS] token and project to num_classes logits.\"\"\"\n def __init__(self, hidden_dim=512, num_classes=10):\n super().__init__()\n self.norm = nn.LayerNorm(hidden_dim)\n self.head = nn.Linear(hidden_dim, num_classes)\n\n def forward(self, x):\n return self.head(self.norm(x[:, 0])) # x[:,0] = CLS token\n\n\ndef build_model(hidden_dim=None, num_layers=None, num_classes=None):\n \"\"\"Build the full ViT-mini as a flat nn.Sequential.\"\"\"\n D = hidden_dim or CFG[\"hidden_dim\"]\n NL = num_layers or CFG[\"num_layers\"]\n NC = num_classes or CFG[\"num_classes\"]\n return nn.Sequential(\n PatchEmbedding(hidden_dim=D),\n *[TransformerBlock(hidden_dim=D) for _ in range(NL)],\n ClassificationHead(hidden_dim=D, num_classes=NC),\n )\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 3. PROCESS-GROUP SETUP / TEARDOWN\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef setup_dist(rank: int, world_size: int, backend: str = \"nccl\"):\n \"\"\"\n Initialise the default NCCL process group.\n torchrun sets MASTER_ADDR/PORT automatically; mp.spawn does not.\n \"\"\"\n os.environ.setdefault(\"MASTER_ADDR\", \"localhost\")\n os.environ.setdefault(\"MASTER_PORT\", \"12355\")\n dist.init_process_group(\n backend = backend,\n rank = rank,\n world_size = world_size,\n timeout = datetime.timedelta(minutes=10),\n )\n torch.cuda.set_device(rank)\n\n\ndef teardown_dist():\n dist.destroy_process_group()\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 4. METRICS TRACKER\n# ═══════════════════════════════════════════════════════════════════════════════\n\nclass MetricsTracker:\n \"\"\"Accumulate per-epoch metrics for later plotting.\"\"\"\n def __init__(self):\n self.reset()\n\n def reset(self):\n self.train_losses, self.val_losses = [], []\n self.train_accs, self.val_accs = [], []\n self.throughputs = [] # samples / second\n self.gpu_mem_mb = [] # peak allocated MB\n self.step_times = []\n\n def log_train(self, loss, acc, throughput, mem_mb, step_time):\n self.train_losses.append(loss); self.train_accs.append(acc)\n self.throughputs.append(throughput)\n self.gpu_mem_mb.append(mem_mb); self.step_times.append(step_time)\n\n def log_val(self, loss, acc):\n self.val_losses.append(loss); self.val_accs.append(acc)\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 5. §3.1 DATA PARALLELISM — DDP\n#\n# Concept (slides 4-9):\n# • Every GPU holds a FULL model copy.\n# • Each GPU trains on a DIFFERENT mini-batch (DistributedSampler).\n# • DDP hooks into backward(): AllReduce fires per gradient-bucket as\n# gradients become available → overlaps comm with backward (slide 20).\n# • Ring-AllReduce (slides 6-7): bandwidth/GPU = 2x|gradient|, optimal.\n# • Gradient accumulation (slide 9): effective_batch = micro*accum*world.\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef train_epoch_ddp(model, loader, optimizer, scaler,\n device, rank, world_size, epoch, metrics):\n \"\"\"\n One training epoch: DDP + AMP + gradient accumulation.\n\n Implementation notes:\n (1) set_epoch(epoch) — reshuffles data uniquely each epoch.\n (2) autocast(\"cuda\") — FP16 matmuls on T4 Tensor Cores (~2x faster).\n (3) loss/accum_steps — gradient magnitude = full-batch gradient.\n (4) scaler.scale(loss).backward() — prevents FP16 gradient underflow.\n (5) DDP fires AllReduce per bucket automatically during backward.\n (6) AllReduce scalars — aggregate global loss/acc across both GPUs.\n \"\"\"\n model.train()\n if hasattr(loader.sampler, \"set_epoch\"):\n loader.sampler.set_epoch(epoch) # (1)\n\n optimizer.zero_grad(set_to_none=True)\n total_loss, total_corr, total_samp = 0.0, 0, 0\n t0 = time.perf_counter()\n\n for step, (images, labels) in enumerate(loader):\n images = images.to(device, non_blocking=True)\n labels = labels.to(device, non_blocking=True)\n\n with autocast(device_type=\"cuda\", enabled=CFG[\"use_amp\"]): # (2)\n logits = model(images)\n loss = F.cross_entropy(logits, labels) / CFG[\"accum_steps\"] # (3)\n\n scaler.scale(loss).backward() # (4) + (5) DDP AllReduce per bucket\n\n total_loss += loss.item() * CFG[\"accum_steps\"]\n total_corr += (logits.detach().argmax(1) == labels).sum().item()\n total_samp += images.size(0)\n\n if (step + 1) % CFG[\"accum_steps\"] == 0:\n scaler.unscale_(optimizer)\n nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n scaler.step(optimizer)\n scaler.update()\n optimizer.zero_grad(set_to_none=True)\n\n # (6) AllReduce global metrics\n t_loss = torch.tensor(total_loss / len(loader), device=device)\n t_corr = torch.tensor(float(total_corr), device=device)\n t_samp = torch.tensor(float(total_samp), device=device)\n dist.all_reduce(t_loss, op=dist.ReduceOp.SUM); t_loss /= world_size\n dist.all_reduce(t_corr, op=dist.ReduceOp.SUM)\n dist.all_reduce(t_samp, op=dist.ReduceOp.SUM)\n\n elapsed = time.perf_counter() - t0\n throughput = t_samp.item() / elapsed\n mem_mb = torch.cuda.max_memory_allocated(device) / 1e6\n\n if rank == 0:\n metrics.log_train(t_loss.item(), t_corr.item() / t_samp.item(),\n throughput, mem_mb, elapsed)\n print(f\" [Train] Epoch {epoch:>2} loss={t_loss:.4f} \"\n f\"acc={100*t_corr.item()/t_samp.item():.1f}% \"\n f\"thr={throughput:.0f} samp/s mem={mem_mb:.0f}MB\")\n\n\n@torch.no_grad()\ndef val_epoch(model, loader, device, rank, world_size, metrics):\n \"\"\"Validation with AllReduce for global accuracy across both GPUs.\"\"\"\n model.eval()\n total_loss, total_corr, total_samp = 0.0, 0, 0\n for images, labels in loader:\n images = images.to(device, non_blocking=True)\n labels = labels.to(device, non_blocking=True)\n with autocast(device_type=\"cuda\", enabled=CFG[\"use_amp\"]):\n logits = model(images)\n total_loss += F.cross_entropy(logits, labels).item()\n total_corr += (logits.argmax(1) == labels).sum().item()\n total_samp += images.size(0)\n\n t_loss = torch.tensor(total_loss / len(loader), device=device)\n t_corr = torch.tensor(float(total_corr), device=device)\n t_samp = torch.tensor(float(total_samp), device=device)\n dist.all_reduce(t_loss, op=dist.ReduceOp.SUM); t_loss /= world_size\n dist.all_reduce(t_corr, op=dist.ReduceOp.SUM)\n dist.all_reduce(t_samp, op=dist.ReduceOp.SUM)\n\n if rank == 0:\n metrics.log_val(t_loss.item(), t_corr.item() / t_samp.item())\n print(f\" [ Val ] loss={t_loss:.4f} \"\n f\"acc={100*t_corr.item()/t_samp.item():.1f}%\")\n\n\ndef run_ddp(rank: int, world_size: int):\n \"\"\"\n §3.1 Data Parallelism entry-point.\n Full model on each GPU; DDP syncs gradients via Ring-AllReduce.\n \"\"\"\n print(f\"\\n{'='*60}\")\n print(f\" §3.1 DATA PARALLELISM — DDP (rank={rank}/{world_size})\")\n print(f\"{'='*60}\")\n\n setup_dist(rank, world_size, CFG[\"backend\"])\n device = torch.device(f\"cuda:{rank}\")\n\n model = build_model().to(device)\n ddp_model = DDP(model, device_ids=[rank],\n bucket_cap_mb=CFG[\"bucket_mb\"],\n find_unused_parameters=False)\n\n optimizer = torch.optim.AdamW(ddp_model.parameters(),\n lr=CFG[\"lr\"], weight_decay=CFG[\"weight_decay\"])\n scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(\n optimizer, T_max=CFG[\"epochs\"])\n\n # FIX 2: GradScaler(\"cuda\", ...) is the correct PyTorch 2.x API\n scaler = GradScaler(\"cuda\", enabled=CFG[\"use_amp\"])\n\n train_loader = get_cifar10(True, rank, world_size)\n val_loader = get_cifar10(False, rank, world_size)\n os.makedirs(CFG[\"ckpt_dir\"], exist_ok=True)\n metrics = MetricsTracker()\n\n for epoch in range(1, CFG[\"epochs\"] + 1):\n train_epoch_ddp(ddp_model, train_loader, optimizer, scaler,\n device, rank, world_size, epoch, metrics)\n val_epoch(ddp_model, val_loader, device, rank, world_size, metrics)\n scheduler.step()\n\n # §3.4 Checkpointing: only rank-0 writes to avoid concurrent writes.\n # .module unwraps DDP to get the raw model state_dict.\n if rank == 0 and epoch % CFG[\"save_every\"] == 0:\n ckpt = os.path.join(CFG[\"ckpt_dir\"], f\"ddp_epoch{epoch}.pt\")\n torch.save({\n \"epoch\" : epoch,\n \"model_state\" : ddp_model.module.state_dict(),\n \"optim_state\" : optimizer.state_dict(),\n \"scaler_state\": scaler.state_dict(),\n \"rng\" : torch.get_rng_state(), # for exact resume\n }, ckpt)\n print(f\" [Ckpt] Saved → {ckpt}\")\n\n # §3.4 MFU Estimation (slide 17)\n # MFU = (throughput x FLOPs_per_sample) / GPU_peak_FLOP/s\n if rank == 0:\n n_params = sum(p.numel() for p in model.parameters())\n flops_per_sample = 6 * n_params # rule-of-thumb: 6N (slide 17)\n avg_thr = float(np.mean(metrics.throughputs))\n mfu = avg_thr * flops_per_sample / 65e12 # T4 peak FP16\n print(f\"\\n [MFU] Params={n_params/1e6:.1f}M \"\n f\"Throughput={avg_thr:.0f} samp/s MFU={100*mfu:.1f}%\")\n _plot_ddp_metrics(metrics, \"DDP\")\n\n teardown_dist()\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 6. §3.2 PIPELINE PARALLELISM — MANUAL 2-GPU IMPLEMENTATION\n#\n# FIX 1: torch.distributed.pipeline.sync.Pipe was REMOVED in PyTorch 2.0.\n# torch.distributed.rpc (required by Pipe) was also removed.\n# The manual implementation below does exactly what Pipe did internally:\n# • Split the model into sequential stages placed on different GPUs.\n# • Transfer activations between stages using tensor.to(\"cuda:N\").\n# • PyTorch autograd handles cross-device gradient flow automatically\n# via the CopyBackwards node created by .to().\n# No Pipe, no RPC, no external library needed.\n#\n# Concept (slides 12-13 — GPipe):\n# Stage-0 (cuda:0): PatchEmbedding + first 3 TransformerBlocks\n# Stage-1 (cuda:1): last 3 TransformerBlocks + ClassificationHead\n#\n# For each global batch:\n# Split into F micro-batches (pipe_chunks).\n# For each micro-batch:\n# act = stage0(x) # cuda:0\n# act = act.to(\"cuda:1\") # transfer + autograd edge\n# logits = stage1(act) # cuda:1\n# loss.backward() # gradients flow back through autograd edge\n#\n# Bubble fraction = (K-1) / (K-1+F)\n# K=2, F=4 → bubble = 20%\n# K=2, F=8 → bubble = 11%\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef build_pipeline_stages():\n \"\"\"\n Build Stage-0 (cuda:0) and Stage-1 (cuda:1).\n\n Stage-0: PatchEmbedding + first num_layers//2 TransformerBlocks\n Stage-1: remaining TransformerBlocks + ClassificationHead\n\n Each GPU stores roughly half the parameters, so neither GPU needs\n the full model weight memory simultaneously.\n \"\"\"\n D = CFG[\"hidden_dim\"]\n half = CFG[\"num_layers\"] // 2 # = 3\n\n stage0 = nn.Sequential(\n PatchEmbedding(hidden_dim=D),\n *[TransformerBlock(hidden_dim=D) for _ in range(half)],\n ).to(\"cuda:0\")\n\n stage1 = nn.Sequential(\n *[TransformerBlock(hidden_dim=D) for _ in range(CFG[\"num_layers\"] - half)],\n ClassificationHead(hidden_dim=D, num_classes=CFG[\"num_classes\"]),\n ).to(\"cuda:1\")\n\n return stage0, stage1\n\n\ndef gpipe_forward(stage0, stage1, x_global, labels_global, n_chunks, backward=True):\n x_chunks = x_global.chunk(n_chunks, dim=0)\n y_chunks = labels_global.chunk(n_chunks, dim=0)\n\n total_loss = 0.0\n total_corr = 0\n\n for x_mb, y_mb in zip(x_chunks, y_chunks):\n x_mb = x_mb.to(\"cuda:0\", non_blocking=True)\n y_mb = y_mb.to(\"cuda:1\", non_blocking=True)\n\n act = stage0(x_mb)\n act = act.to(\"cuda:1\")\n logits = stage1(act)\n\n loss = F.cross_entropy(logits, y_mb)\n\n if backward:\n (loss / n_chunks).backward()\n\n total_loss += loss.item()\n with torch.no_grad():\n total_corr += (logits.argmax(1) == y_mb).sum().item()\n\n return total_loss, total_corr\n\n\ndef run_pipeline(rank: int, world_size: int):\n \"\"\"\n §3.2 Manual Pipeline Parallelism demo.\n\n Runs as a single process that directly drives both GPUs.\n rank != 0 returns immediately (no work to do).\n\n In production (PiPPy / Megatron-LM) each stage lives in its own process\n and activations travel via dist.send/recv. Single-process is functionally\n identical and transparent for teaching.\n \"\"\"\n if rank != 0:\n return # rank-1 has nothing to do in single-process pipeline\n\n print(f\"\\n{'='*60}\")\n print(f\" §3.2 MODEL PARALLELISM — MANUAL PIPELINE (2 GPUs, no Pipe)\")\n n_chunks = CFG[\"pipe_chunks\"]\n bubble = 1 / (1 + n_chunks) # (K-1)/(K-1+F) with K=2\n print(f\" GPipe: chunks={n_chunks} theoretical bubble={100*bubble:.0f}%\")\n print(f\"{'='*60}\")\n\n stage0, stage1 = build_pipeline_stages()\n n0 = sum(p.numel() for p in stage0.parameters())\n n1 = sum(p.numel() for p in stage1.parameters())\n print(f\" Stage-0 (cuda:0): {n0/1e6:.1f}M params\")\n print(f\" Stage-1 (cuda:1): {n1/1e6:.1f}M params\")\n\n all_params = list(stage0.parameters()) + list(stage1.parameters())\n optimizer = torch.optim.AdamW(all_params, lr=CFG[\"lr\"],\n weight_decay=CFG[\"weight_decay\"])\n\n tf = T.Compose([\n T.RandomCrop(32, padding=4), T.RandomHorizontalFlip(), T.ToTensor(),\n T.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),\n ])\n batch_size = CFG[\"micro_batch\"] * n_chunks\n train_set = torchvision.datasets.CIFAR10(\n root=CFG[\"data_dir\"], train=True, transform=tf, download=True)\n val_set = torchvision.datasets.CIFAR10(\n root=CFG[\"data_dir\"], train=False, transform=tf, download=True)\n train_loader = DataLoader(train_set, batch_size=batch_size,\n shuffle=True, num_workers=2, pin_memory=True)\n val_loader = DataLoader(val_set, batch_size=batch_size,\n shuffle=False, num_workers=2, pin_memory=True)\n\n metrics = MetricsTracker()\n\n for epoch in range(1, CFG[\"epochs\"] + 1):\n stage0.train(); stage1.train()\n t0 = time.perf_counter()\n total_loss, total_corr, total_samp = 0.0, 0, 0\n\n for images, labels in train_loader:\n optimizer.zero_grad(set_to_none=True)\n loss, corr = gpipe_forward(stage0, stage1, images, labels, n_chunks, backward=True)\n nn.utils.clip_grad_norm_(all_params, 1.0)\n optimizer.step()\n total_loss += loss; total_corr += corr; total_samp += images.size(0)\n\n elapsed = time.perf_counter() - t0\n throughput = total_samp / elapsed\n mem0 = torch.cuda.max_memory_allocated(\"cuda:0\") / 1e6\n mem1 = torch.cuda.max_memory_allocated(\"cuda:1\") / 1e6\n torch.cuda.reset_peak_memory_stats(\"cuda:0\")\n torch.cuda.reset_peak_memory_stats(\"cuda:1\")\n\n stage0.eval(); stage1.eval()\n val_corr, val_samp = 0, 0\n with torch.no_grad():\n for vimg, vlbl in val_loader:\n _, c = gpipe_forward(stage0, stage1, vimg, vlbl, n_chunks, backward=False)\n val_corr += c\n val_samp += vimg.size(0)\n\n train_acc = total_corr / max(total_samp, 1)\n val_acc = val_corr / max(val_samp, 1)\n metrics.log_train(total_loss / len(train_loader), train_acc,\n throughput, mem0 + mem1, elapsed)\n metrics.log_val(0.0, val_acc)\n print(f\" [Pipe] Epoch {epoch:>2} \"\n f\"train={100*train_acc:.1f}% val={100*val_acc:.1f}% \"\n f\"GPU0={mem0:.0f}MB GPU1={mem1:.0f}MB thr={throughput:.0f} samp/s\")\n\n _plot_pipeline_memory(metrics)\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 7. §3.2 TENSOR PARALLELISM — column-parallel linear (Megatron-LM style)\n#\n# Concept (slide 11):\n# Y = X * W split column-wise:\n# GPU-0: W[:, :out/2] → Y_0 = X * W_0 (no comm)\n# GPU-1: W[:,out/2: ] → Y_1 = X * W_1 (no comm)\n# AllGather → Y = [Y_0 | Y_1] (1 collective per layer-forward)\n#\n# Memory saving: each GPU stores (in_f x out_f / world_size) params.\n# Comm cost: 1 AllGather per layer forward (proportional to activation size).\n# ═══════════════════════════════════════════════════════════════════════════════\n\nclass ColumnParallelLinear(nn.Module):\n \"\"\"\n Tensor-Parallel column-parallel linear (Megatron-LM approach).\n\n Each rank stores W[:, col_start:col_end]:\n rank-0: columns [0 ... out/2-1]\n rank-1: columns [out/2 ... out-1]\n\n Forward:\n (1) y_local = x @ W_local — pure local matmul, zero communication\n (2) AllGather(y_local) — one collective, all ranks get full Y\n \"\"\"\n def __init__(self, in_features, out_features, rank, world_size):\n super().__init__()\n assert out_features % world_size == 0, \\\n \"out_features must be divisible by world_size\"\n self.rank = rank\n self.world_size = world_size\n self.out_local = out_features // world_size\n\n self.weight = nn.Parameter(\n torch.randn(in_features, self.out_local) * (in_features ** -0.5))\n self.bias = nn.Parameter(torch.zeros(out_features))\n\n def forward(self, x):\n y_local = F.linear(x, self.weight.T) # (B, out_local)\n gathered = [torch.empty_like(y_local) for _ in range(self.world_size)]\n dist.all_gather(gathered, y_local)\n return torch.cat(gathered, dim=-1) + self.bias # (B, out_full)\n\n\ndef demo_tensor_parallel(rank: int, world_size: int):\n \"\"\"\n Verify ColumnParallelLinear against a reference single-GPU matmul.\n Expected max absolute error < 1e-3 in FP32.\n \"\"\"\n setup_dist(rank, world_size)\n device = torch.device(f\"cuda:{rank}\")\n\n B, IN, OUT = 16, 512, 1024\n torch.manual_seed(42)\n\n W_full = torch.randn(IN, OUT, device=\"cuda:0\") * (IN ** -0.5)\n\n # Replicate input across ranks (Megatron: X is identical on all GPUs)\n X = torch.randn(B, IN, device=device)\n dist.broadcast(X, src=0)\n\n # Load each rank's column shard from W_full\n tp = ColumnParallelLinear(IN, OUT, rank, world_size).to(device)\n with torch.no_grad():\n c_start = rank * (OUT // world_size)\n c_end = c_start + (OUT // world_size)\n tp.weight.data = W_full[:, c_start:c_end].T.to(device)\n\n y_tp = tp(X) # tensor-parallel result\n y_ref = F.linear(X.to(\"cuda:0\"), W_full.T) # reference (single GPU)\n\n if rank == 0:\n err = (y_tp.cpu() - y_ref.cpu()).abs().max().item()\n status = \"PASS\" if err < 1e-3 else \"FAIL\"\n print(f\"\\n [TP] max|y_tp - y_ref| = {err:.2e} [{status}]\")\n print(\" (< 1e-3 expected; small rounding differences in FP32 are normal)\")\n\n teardown_dist()\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 8. §3.3 HYBRID PARALLELISM\n#\n# Concept (slides 9, 15, 19):\n# On 2xT4 the accessible hybrid is:\n# DDP (world_size=2) x Gradient Accumulation (accum_steps=4)\n# effective_batch = 64 x 4 x 2 = 512\n#\n# PowerSGD (slide 19): gradient matrices compressed to rank-r factors.\n# AllReduce bandwidth reduces from d^2 to 2*r*d (10-100x savings).\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef run_hybrid(rank: int, world_size: int):\n \"\"\"§3.3 Hybrid: DDP + gradient accumulation + optional PowerSGD.\"\"\"\n eff = CFG[\"micro_batch\"] * CFG[\"accum_steps\"] * world_size\n print(f\"\\n{'='*60}\")\n print(f\" §3.3 HYBRID PARALLELISM (rank={rank}/{world_size})\")\n print(f\" Effective batch = {CFG['micro_batch']} x \"\n f\"{CFG['accum_steps']} x {world_size} = {eff}\")\n print(f\"{'='*60}\")\n\n setup_dist(rank, world_size, CFG[\"backend\"])\n device = torch.device(f\"cuda:{rank}\")\n model = build_model().to(device)\n ddp_model = DDP(model, device_ids=[rank], bucket_cap_mb=CFG[\"bucket_mb\"])\n\n try:\n from torch.distributed.algorithms.ddp_comm_hooks import (\n powerSGD_hook as psgd)\n state = psgd.PowerSGDState(\n process_group=None, matrix_approximation_rank=4, warm_start=True)\n ddp_model.register_comm_hook(state, psgd.powerSGD_hook)\n if rank == 0:\n print(\" [PowerSGD] rank-4 gradient compression hook registered\")\n except Exception as e:\n if rank == 0:\n print(f\" [PowerSGD] unavailable ({e}); using plain AllReduce\")\n\n optimizer = torch.optim.AdamW(ddp_model.parameters(),\n lr=CFG[\"lr\"], weight_decay=CFG[\"weight_decay\"])\n # FIX 2: correct PyTorch 2.x GradScaler API\n scaler = GradScaler(\"cuda\", enabled=CFG[\"use_amp\"])\n\n train_loader = get_cifar10(True, rank, world_size)\n val_loader = get_cifar10(False, rank, world_size)\n metrics = MetricsTracker()\n\n for epoch in range(1, CFG[\"epochs\"] + 1):\n train_epoch_ddp(ddp_model, train_loader, optimizer, scaler,\n device, rank, world_size, epoch, metrics)\n val_epoch(ddp_model, val_loader, device, rank, world_size, metrics)\n\n if rank == 0:\n _plot_ddp_metrics(metrics, \"Hybrid\")\n\n teardown_dist()\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 9. §3.4 PROFILING WITH torch.profiler\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef run_profiled_training(rank: int, world_size: int):\n \"\"\"\n §3.4 torch.profiler demo.\n Schedule: wait=1 (JIT warmup), warmup=1, active=3 recorded steps.\n Output: Chrome JSON trace viewable at https://ui.perfetto.dev\n \"\"\"\n print(f\"\\n{'='*60}\")\n print(f\" §3.4 PROFILING (rank={rank}/{world_size})\")\n print(f\"{'='*60}\")\n\n setup_dist(rank, world_size, CFG[\"backend\"])\n device = torch.device(f\"cuda:{rank}\")\n model = build_model().to(device)\n ddp_model = DDP(model, device_ids=[rank])\n optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=CFG[\"lr\"])\n # FIX 2: correct GradScaler API\n scaler = GradScaler(\"cuda\", enabled=CFG[\"use_amp\"])\n loader = get_cifar10(True, rank, world_size)\n\n os.makedirs(CFG[\"profile_dir\"], exist_ok=True)\n\n schedule = torch.profiler.schedule(wait=1, warmup=1,\n active=CFG[\"profile_steps\"])\n with profile(\n activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA],\n schedule = schedule,\n on_trace_ready = torch.profiler.tensorboard_trace_handler(\n CFG[\"profile_dir\"]),\n record_shapes = True,\n profile_memory = True,\n with_stack = False,\n ) as prof:\n ddp_model.train()\n loader_it = iter(loader)\n for _ in range(6): # wait(1) + warmup(1) + active(3) + 1 extra\n images, labels = next(loader_it)\n images = images.to(device, non_blocking=True)\n labels = labels.to(device, non_blocking=True)\n optimizer.zero_grad(set_to_none=True)\n with autocast(device_type=\"cuda\", enabled=CFG[\"use_amp\"]):\n with record_function(\"forward\"):\n logits = ddp_model(images)\n with record_function(\"loss\"):\n loss = F.cross_entropy(logits, labels)\n with record_function(\"backward\"):\n scaler.scale(loss).backward()\n with record_function(\"optimizer\"):\n scaler.step(optimizer); scaler.update()\n prof.step()\n\n if rank == 0:\n print(\"\\n Top-20 ops by CUDA time:\")\n print(prof.key_averages().table(sort_by=\"cuda_time_total\", row_limit=20))\n print(f\"\\n [Prof] Trace saved to {CFG['profile_dir']}\")\n print(\" Upload the .json.gz to https://ui.perfetto.dev\")\n _plot_profiler_ops(prof)\n\n teardown_dist()\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 10. VISUALISATION HELPERS\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef _plot_ddp_metrics(metrics: MetricsTracker, label: str = \"DDP\"):\n \"\"\"4-panel dashboard: loss · accuracy · throughput · GPU memory.\"\"\"\n epochs = range(1, len(metrics.train_losses) + 1)\n fig, axes = plt.subplots(2, 2, figsize=(12, 8))\n fig.suptitle(f\"Training Dashboard — {label} · Kaggle 2xT4\", fontsize=13)\n\n ax = axes[0, 0]\n ax.plot(epochs, metrics.train_losses, \"b-o\", ms=5, label=\"Train\")\n if metrics.val_losses:\n ax.plot(epochs, metrics.val_losses, \"r--s\", ms=5, label=\"Val\")\n ax.set(title=\"Loss\", xlabel=\"Epoch\", ylabel=\"Cross-Entropy\"); ax.legend(); ax.grid(alpha=0.3)\n\n ax = axes[0, 1]\n ax.plot(epochs, [a*100 for a in metrics.train_accs], \"b-o\", ms=5, label=\"Train\")\n if metrics.val_accs:\n ax.plot(epochs, [a*100 for a in metrics.val_accs], \"r--s\", ms=5, label=\"Val\")\n ax.set(title=\"Accuracy\", xlabel=\"Epoch\", ylabel=\"%\"); ax.legend(); ax.grid(alpha=0.3)\n\n ax = axes[1, 0]\n ax.bar(epochs, metrics.throughputs, color=\"steelblue\", alpha=0.8)\n mean_thr = float(np.mean(metrics.throughputs))\n ax.axhline(mean_thr, color=\"red\", ls=\"--\", label=f\"Mean={mean_thr:.0f}\")\n ax.set(title=\"Throughput (samples/s)\", xlabel=\"Epoch\", ylabel=\"Samples/sec\")\n ax.legend(); ax.grid(alpha=0.3, axis=\"y\")\n\n ax = axes[1, 1]\n ax.plot(epochs, metrics.gpu_mem_mb, \"g-^\", ms=6)\n ax.axhline(16000, color=\"red\", ls=\":\", label=\"T4 limit (16 GB)\")\n ax.set(title=\"Peak GPU Memory\", xlabel=\"Epoch\", ylabel=\"MB\")\n ax.legend(); ax.grid(alpha=0.3)\n\n plt.tight_layout()\n path = os.path.join(CFG[\"plot_dir\"],\n f\"dashboard_{label.lower().replace(' ','_')}.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\"); plt.close()\n print(f\" [Plot] Dashboard → {path}\")\n\n\ndef _plot_pipeline_memory(metrics: MetricsTracker):\n \"\"\"Accuracy + combined pipeline GPU memory across epochs.\"\"\"\n epochs = range(1, len(metrics.train_losses) + 1)\n fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n fig.suptitle(\"Pipeline Parallelism · Kaggle 2xT4\", fontsize=13)\n\n ax = axes[0]\n ax.plot(epochs, [a*100 for a in metrics.train_accs], \"b-o\", ms=5, label=\"Train\")\n if metrics.val_accs:\n ax.plot(epochs, [a*100 for a in metrics.val_accs], \"r--s\", ms=5, label=\"Val\")\n ax.set(title=\"Accuracy per Epoch\", xlabel=\"Epoch\", ylabel=\"%\"); ax.legend(); ax.grid(alpha=0.3)\n\n ax = axes[1]\n ax.plot(epochs, metrics.gpu_mem_mb, \"m-o\", ms=5, label=\"GPU0+GPU1 total\")\n ax.axhline(16000, color=\"r\", ls=\":\", label=\"Single T4 (16 GB)\")\n ax.axhline(32000, color=\"g\", ls=\"--\", label=\"2xT4 total (32 GB)\")\n ax.set(title=\"Combined GPU Memory\", xlabel=\"Epoch\", ylabel=\"MB\")\n ax.legend(); ax.grid(alpha=0.3)\n\n plt.tight_layout()\n path = os.path.join(CFG[\"plot_dir\"], \"pipeline_memory.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\"); plt.close()\n print(f\" [Plot] Pipeline memory → {path}\")\n\n\ndef _plot_profiler_ops(prof):\n \"\"\"Horizontal bar chart of top-10 ops by CUDA time.\"\"\"\n avgs = sorted(prof.key_averages(),\n key=lambda a: a.cuda_time_total, reverse=True)[:10]\n names = [a.key[:40] for a in avgs]\n times = [a.cuda_time_total / 1e3 for a in avgs] # us → ms\n\n fig, ax = plt.subplots(figsize=(9, 5))\n bars = ax.barh(names[::-1], times[::-1], color=\"steelblue\", alpha=0.85)\n ax.bar_label(bars, fmt=\"%.2f ms\", padding=3, fontsize=8)\n ax.set(title=\"Top-10 CUDA Ops (profiler)\", xlabel=\"CUDA time (ms)\")\n ax.grid(alpha=0.3, axis=\"x\")\n plt.tight_layout()\n path = os.path.join(CFG[\"plot_dir\"], \"profiler_ops.png\")\n plt.savefig(path, dpi=150, bbox_inches=\"tight\"); plt.close()\n print(f\" [Plot] Profiler ops → {path}\")\n\n\n# ═══════════════════════════════════════════════════════════════════════════════\n# 11. MAIN\n# ═══════════════════════════════════════════════════════════════════════════════\n\ndef spawn_or_torchrun(fn, world_size: int):\n \"\"\"\n torchrun sets RANK env var → we are already the target process.\n plain python → use mp.spawn to fork world_size processes.\n \"\"\"\n if \"RANK\" in os.environ:\n fn(int(os.environ[\"RANK\"]), world_size)\n else:\n mp.spawn(fn, args=(world_size,), nprocs=world_size, join=True)\n\n\ndef main():\n parser = argparse.ArgumentParser(description=\"Chapter-3 Distributed Demo\")\n parser.add_argument(\"--mode\", default=\"ddp\",\n choices=[\"ddp\", \"pipeline\", \"hybrid\", \"tensor_parallel\", \"profile\", \"all\"])\n parser.add_argument(\"--epochs\", type=int, default=CFG[\"epochs\"])\n parser.add_argument(\"--micro_batch\", type=int, default=CFG[\"micro_batch\"])\n parser.add_argument(\"--accum_steps\", type=int, default=CFG[\"accum_steps\"])\n parser.add_argument(\"--pipe_chunks\", type=int, default=CFG[\"pipe_chunks\"])\n args = parser.parse_args()\n\n CFG[\"epochs\"] = args.epochs\n CFG[\"micro_batch\"] = args.micro_batch\n CFG[\"accum_steps\"] = args.accum_steps\n CFG[\"pipe_chunks\"] = args.pipe_chunks\n\n world_size = torch.cuda.device_count()\n if world_size == 0:\n print(\"No CUDA GPUs detected. \"\n \"Enable '2 x GPU T4' in Kaggle Settings → Accelerator.\")\n sys.exit(1)\n\n print(f\"\\n PyTorch : {torch.__version__}\")\n print(f\" GPUs : {world_size} x {torch.cuda.get_device_name(0)}\")\n print(f\" Mode : {args.mode}\")\n\n for d in (CFG[\"plot_dir\"], CFG[\"ckpt_dir\"], CFG[\"data_dir\"]):\n os.makedirs(d, exist_ok=True)\n\n mode = args.mode\n\n if mode in (\"ddp\", \"all\"):\n spawn_or_torchrun(run_ddp, world_size)\n\n if mode in (\"pipeline\", \"all\") and world_size >= 2:\n # run_pipeline drives both GPUs from rank-0; rank-1 returns immediately.\n spawn_or_torchrun(run_pipeline, world_size)\n\n if mode in (\"hybrid\", \"all\"):\n spawn_or_torchrun(run_hybrid, world_size)\n\n if mode in (\"tensor_parallel\", \"all\") and world_size >= 2:\n spawn_or_torchrun(demo_tensor_parallel, world_size)\n\n if mode in (\"profile\", \"all\"):\n spawn_or_torchrun(run_profiled_training, world_size)\n\n print(f\"\\n Done. Plots saved to {CFG['plot_dir']}\")\n\n\nif __name__ == \"__main__\":\n main()\n","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2026-05-03T01:41:13.424468Z","iopub.execute_input":"2026-05-03T01:41:13.424702Z","iopub.status.idle":"2026-05-03T01:41:13.448154Z","shell.execute_reply.started":"2026-05-03T01:41:13.424680Z","shell.execute_reply":"2026-05-03T01:41:13.447507Z"}},"outputs":[],"execution_count":null},{"id":"8337bda4","cell_type":"code","source":"# Patch the chapter3 demo so `--mode all` is stable under `torchrun`\n# Fixes common causes of NCCL ncclRemoteError / connection refused:\n# 1) pipeline runs only on rank0 but uses both GPUs -> rank1 must wait\n# 2) rank teardown desync between modes -> barrier before destroy\n# 3) tensor-parallel demo must not touch cuda:0 on rank1 (single GPU per proc)\nfrom pathlib import Path\nimport re\n\npath = Path(\"chapter3_kaggle_2xT4_demo.py\")\ntext = path.read_text(encoding=\"utf-8\")\nchanged = False\n\ndef _apply(name: str, fn):\n global text, changed\n new_text = fn(text)\n if new_text != text:\n print(f\"[patch] applied: {name}\")\n text = new_text\n changed = True\n else:\n print(f\"[patch] skipped: {name} (already applied or pattern not found)\")\n\n# --- (1) Pipeline rank synchronization (rank1 waits for rank0 completion) ---\ndef patch_pipeline(t: str) -> str:\n needle = \" if rank != 0:\\n return # rank-1 has nothing to do in single-process pipeline\\n\"\n if needle not in t:\n return t\n replacement = (\n \" os.makedirs(CFG[\\\"ckpt_dir\\\"], exist_ok=True)\\n\"\n \" run_id = os.environ.get(\\\"TORCHELASTIC_RUN_ID\\\") or os.environ.get(\\\"MASTER_PORT\\\") or \\\"default\\\"\\n\"\n \" done_file = os.path.join(CFG[\\\"ckpt_dir\\\"], f\\\".pipeline_done_{run_id}\\\")\\n\\n\"\n \" if rank != 0:\\n\"\n \" t_start = time.time()\\n\"\n \" while not os.path.exists(done_file):\\n\"\n \" if time.time() - t_start > 6 * 60 * 60:\\n\"\n \" raise TimeoutError(f\\\"Pipeline wait timed out: {done_file}\\\")\\n\"\n \" time.sleep(1.0)\\n\"\n \" return\\n\\n\"\n \" # rank 0 runs the pipeline and then signals completion\\n\"\n \" try:\\n\"\n \" if os.path.exists(done_file):\\n\"\n \" os.remove(done_file)\\n\"\n \" except OSError:\\n\"\n \" pass\\n\"\n )\n t = t.replace(needle, \"\".join(replacement))\n # Insert completion signal after plotting, if not already present\n if \"# Signal other ranks they can proceed\" not in t:\n anchor = \" _plot_pipeline_memory(metrics)\\n\"\n if anchor in t:\n t = t.replace(\n anchor,\n anchor\n + \"\\n\"\n + \" # Signal other ranks they can proceed to the next mode.\\n\"\n + \" with open(done_file, \\\"w\\\", encoding=\\\"utf-8\\\") as f:\\n\"\n + \" f.write(\\\"done\\\\n\\\")\\n\",\n 1,\n )\n return t\n\n_apply(\"pipeline rank sync\", patch_pipeline)\n\n# --- (2) Make teardown synchronized (barrier before destroy) ---\ndef patch_teardown(t: str) -> str:\n pattern = r\"def teardown_dist\\(\\):\\n\\s+dist\\.destroy_process_group\\(\\)\"\n repl = (\n \"def teardown_dist():\\n\"\n \" if dist.is_available() and dist.is_initialized():\\n\"\n \" dist.barrier()\\n\"\n \" dist.destroy_process_group()\",\n )\n return re.sub(pattern, \"\".join(repl), t)\n\n_apply(\"teardown barrier\", patch_teardown)\n\n# --- (3) Tensor-parallel demo: keep each rank on its local GPU + broadcast W_full ---\ndef patch_tensor_parallel(t: str) -> str:\n # Replace any hardcoded cuda:0 weight allocation with local device\n t = t.replace(\n \"W_full = torch.randn(IN, OUT, device=\\\"cuda:0\\\") * (IN ** -0.5)\",\n \"W_full = torch.randn(IN, OUT, device=device) * (IN ** -0.5)\",\n )\n # Insert broadcast of W_full from rank0 right after the W_full assignment\n if \"dist.broadcast(W_full, src=0)\" not in t:\n pat = r\"(?m)^(\\s*W_full\\s*=\\s*torch\\.randn\\(IN,\\s*OUT,\\s*device=device\\)\\s*\\*\\s*\\(IN\\s*\\*\\*\\s*-0\\.5\\)\\s*\\n)\"\n t = re.sub(pat, r\"\\1\\1\".replace(\"\\1\", \"\"), t, count=0)\n return t","metadata":{"trusted":true,"jupyter":{"source_hidden":true},"execution":{"iopub.status.busy":"2026-05-03T01:41:13.448993Z","iopub.execute_input":"2026-05-03T01:41:13.449474Z","iopub.status.idle":"2026-05-03T01:41:13.465224Z","shell.execute_reply.started":"2026-05-03T01:41:13.449439Z","shell.execute_reply":"2026-05-03T01:41:13.464465Z"}},"outputs":[],"execution_count":null},{"id":"c6175173","cell_type":"code","source":"!torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode ddp","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:41:13.466214Z","iopub.execute_input":"2026-05-03T01:41:13.466804Z","iopub.status.idle":"2026-05-03T01:43:54.149154Z","shell.execute_reply.started":"2026-05-03T01:41:13.466772Z","shell.execute_reply":"2026-05-03T01:43:54.148152Z"}},"outputs":[],"execution_count":null},{"id":"7caffcc2","cell_type":"code","source":"# Demo: all modes back-to-back (now safe with the patch cell above)\n%env OMP_NUM_THREADS=1\n!torchrun --nproc_per_node=2 chapter3_kaggle_2xT4_demo.py --mode all --epochs 3","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-03T01:43:54.151384Z","iopub.execute_input":"2026-05-03T01:43:54.151632Z","iopub.status.idle":"2026-05-03T01:46:35.030881Z","shell.execute_reply.started":"2026-05-03T01:43:54.151606Z","shell.execute_reply":"2026-05-03T01:46:35.029817Z"}},"outputs":[],"execution_count":null}]}