DeepSeed Training Dashboard

!nvidia-smi
!pip install deepspeed transformers datasets pynvml accelerate -q

import os
import json
import time
import math
import numpy as np
from collections import defaultdict

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

import deepspeed
from transformers import BertForSequenceClassification, BertTokenizer
from datasets import load_dataset
from sklearn.metrics import accuracy_score, f1_score

import pynvml
pynvml.nvmlInit()
gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(0)
gpu_name = pynvml.nvmlDeviceGetName(gpu_handle)
gpu_mem_total = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle).total / 1e9

print(f"GPU: {gpu_name}")
print(f"GPU Memory: {gpu_mem_total:.1f} GB")
print(f"PyTorch: {torch.__version__}")
print(f"DeepSpeed: {deepspeed.__version__}")
print(f"CUDA: {torch.version.cuda}")

DS_CONFIG = {
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 4,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 2e-5,
            "betas": [0.9, 0.999],
            "eps": 1e-8,
            "weight_decay": 0.01
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 2e-5,
            "warmup_num_steps": 150,
            "total_num_steps": 2400
        }
    },
    "fp16": {
        "enabled": True,
        "loss_scale": 0,
        "initial_scale_power": 16,
        "loss_scale_window": 1000
    },
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True
        },
        "allgather_partitions": True,
        "allgather_bucket_size": 2e8,
        "overlap_comm": True,
        "reduce_scatter": True,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": True
    },
    "gradient_clipping": 1.0,
    "wall_clock_breakdown": True,
    "steps_per_print": 50
}

with open("ds_config.json", "w") as f:
    json.dump(DS_CONFIG, f, indent=2)

model_name = "bert-large-uncased"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2)

total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

print(f"Model: {model_name}")
print(f"Total parameters: {total_params:,} ({total_params/1e6:.0f}M)")
print(f"Layers: {model.config.num_hidden_layers}")       # 24
print(f"Hidden size: {model.config.hidden_size}")          # 1024
print(f"Attention heads: {model.config.num_attention_heads}") # 16

MAX_SEQ_LEN = 256

dataset = load_dataset("imdb")

def tokenize_fn(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=MAX_SEQ_LEN,
    )

tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=["text"])
tokenized = tokenized.rename_column("label", "labels")
tokenized.set_format("torch")

train_dataset = tokenized["train"]   # 25,000 samples
val_dataset = tokenized["test"]      # 25,000 samples

print(f"Train samples: {len(train_dataset)}")
print(f"Val samples: {len(val_dataset)}")
print(f"Max sequence length: {MAX_SEQ_LEN}")

class MetricsLogger:
    """Collects real GPU metrics, timing, and training stats per step."""

    def __init__(self, gpu_handle):
        self.gpu_handle = gpu_handle
        self.steps = []
        self.evaluations = []

        # CUDA events for sub-millisecond timing
        self.evt_start = torch.cuda.Event(enable_timing=True)
        self.evt_fwd_end = torch.cuda.Event(enable_timing=True)
        self.evt_bwd_end = torch.cuda.Event(enable_timing=True)
        self.evt_opt_end = torch.cuda.Event(enable_timing=True)

    def get_gpu_metrics(self):
        util = pynvml.nvmlDeviceGetUtilizationRates(self.gpu_handle)
        mem = pynvml.nvmlDeviceGetMemoryInfo(self.gpu_handle)
        return {
            "gpu_util_pct": util.gpu,
            "gpu_mem_used_gb": round(mem.used / 1e9, 3),
        }

    def log_step(self, step, epoch, loss, lr, grad_norm, batch_size, layer_times=None):
        torch.cuda.synchronize()
        forward_ms = self.evt_start.elapsed_time(self.evt_fwd_end)
        backward_ms = self.evt_fwd_end.elapsed_time(self.evt_bwd_end)
        optimizer_ms = self.evt_bwd_end.elapsed_time(self.evt_opt_end)
        gpu = self.get_gpu_metrics()
        # ... builds record dict with all fields
        self.steps.append(record)
        return record

class LayerProfiler:
    """
    Registers forward/backward hooks on all BERT encoder layers,
    embeddings, and classifier head to measure real per-layer timing.
    """

    def __init__(self, model):
        self.timings = {}
        self._events = {}
        self._hooks = []
        self._setup(model)

    def _setup(self, model):
        base = model.module if hasattr(model, 'module') else model

        # Register hooks on all 26 layers
        self._register(base.bert.embeddings, "embeddings")
        for i, layer in enumerate(base.bert.encoder.layer):
            self._register(layer, f"encoder.{i}")
        self._register(base.classifier, "classifier")

    def _register(self, module, name):
        # Creates CUDA event pairs and registers
        # forward_pre_hook, forward_hook, backward_hook
        # on each module to record timing events
        ...

    def collect(self):
        """Call after torch.cuda.synchronize() to read timings."""
        result = {}
        for name, evts in self._events.items():
            fwd_ms = evts["fwd_start"].elapsed_time(evts["fwd_end"])
            result[name] = {"fwd_ms": round(fwd_ms, 3)}
        return result

# Single-GPU Colab environment setup
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29500"
os.environ["RANK"] = "0"
os.environ["LOCAL_RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"

train_loader = DataLoader(
    train_dataset,
    batch_size=DS_CONFIG["train_micro_batch_size_per_gpu"],  # 4
    shuffle=True,
    num_workers=2,
    pin_memory=True,
    drop_last=True,
)

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config=DS_CONFIG,
    model_parameters=model.parameters(),
)

# Attach per-layer profiler
layer_profiler = LayerProfiler(model_engine)

print(f"ZeRO Stage: {DS_CONFIG['zero_optimization']['stage']}")
print(f"FP16: {DS_CONFIG['fp16']['enabled']}")
print(f"CPU Offload: {DS_CONFIG['zero_optimization']['offload_optimizer']['device']}")
print(f"Effective batch: {4 * 4}")  # micro_batch * grad_accum = 16

NUM_EPOCHS = 3
LOG_EVERY = 5        # log metrics every 5 steps
EVAL_EVERY = 200    # evaluate every 200 steps
GRAD_ACCUM = 4

for epoch in range(NUM_EPOCHS):
    model_engine.train()

    for batch_idx, batch in enumerate(train_loader):
        # ---- Data loading timing ----
        metrics_logger.start_data_load()
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)
        metrics_logger.end_data_load()

        # ---- Forward pass ----
        metrics_logger.start_step()        # records CUDA event
        outputs = model_engine(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels,
        )
        loss = outputs.loss
        metrics_logger.end_forward()        # records CUDA event

        # ---- Backward pass ----
        model_engine.backward(loss)          # DeepSpeed handles FP16 scaling
        metrics_logger.end_backward()       # records CUDA event

        # ---- Optimizer step ----
        model_engine.step()                  # DeepSpeed handles grad accum
        metrics_logger.end_optimizer()      # records CUDA event

        # ---- Log every 5 steps ----
        if batch_idx % LOG_EVERY == 0:
            torch.cuda.synchronize()
            layer_times = layer_profiler.collect()

            # Compute gradient L2 norm
            grad_norm = 0.0
            for p in model_engine.module.parameters():
                if p.grad is not None:
                    grad_norm += p.grad.data.float().norm(2).item() ** 2
            grad_norm = grad_norm ** 0.5

            metrics_logger.log_step(
                step=global_step, epoch=..., loss=loss.item(),
                lr=lr, grad_norm=grad_norm, batch_size=4,
                layer_times=layer_times,
            )

        # ---- Evaluate every 200 steps ----
        if global_step % EVAL_EVERY == 0:
            val_loss, val_acc, val_f1 = evaluate(model_engine, val_loader, device)
            metrics_logger.log_eval(step=global_step, ...)

def evaluate(model_engine, val_loader, device):
    """Run evaluation and return real loss, accuracy, F1."""
    model_engine.eval()
    total_loss = 0
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for batch in val_loader:
            outputs = model_engine(
                input_ids=batch["input_ids"].to(device),
                attention_mask=batch["attention_mask"].to(device),
                labels=batch["labels"].to(device),
            )
            total_loss += outputs.loss.item()
            preds = torch.argmax(outputs.logits, dim=-1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(batch["labels"].numpy())

    model_engine.train()
    acc = accuracy_score(all_labels, all_preds)
    f1 = f1_score(all_labels, all_preds, average="binary")
    return total_loss / n_batches, acc, f1

# Compute ZeRO memory breakdown
param_bytes = total_params * 2     # FP16 = 2 bytes/param = 0.67 GB
grad_bytes = total_params * 2      # FP16 = 2 bytes/param = 0.67 GB
opt_bytes = total_params * 12      # Adam FP32: copy + momentum + variance = 4.02 GB (CPU)
act_bytes = micro_batch * seq_len * hidden * n_layers * 2  # ~0.20 GB

output = {
    "meta": {
        "model": "bert-large-uncased",
        "params": 335141890,
        "gpu": gpu_name,
        "gpu_memory_total_gb": 15.8,
        "zero_stage": 2,
        "fp16": True,
        "batch_sizes": {"micro": 4, "grad_accum": 4, "effective": 16},
        "epochs": 3,
        "total_time_minutes": 42.3,
        "final_accuracy": final_acc,
        "final_f1": final_f1,
    },
    "steps": metrics_logger.steps,            # 460 records
    "evaluations": metrics_logger.evaluations,  # 11 records
    "zero_memory": {
        "parameters_gb": 0.670,
        "gradients_gb": 0.670,
        "optimizer_cpu_gb": 4.022,
        "activations_gb": 0.201,
    },
}

with open("training_metrics.json", "w") as f:
    json.dump(output, f, indent=2)

# Download from Colab
from google.colab import files
files.download("training_metrics.json")

# DeepSpeed calls this internally during deepspeed.initialize():
torch.distributed.init_process_group(
    backend="nccl",          # NVIDIA Collective Communications Library
    init_method="env://",     # Reads MASTER_ADDR, MASTER_PORT, RANK, WORLD_SIZE
    world_size=1,             # Single GPU in our Colab setup
    rank=0,                   # This process's rank
)

# Process group topology for multi-GPU:
#   Rank 0 (GPU 0) ──NCCL──> Rank 1 (GPU 1)
#        │                        │
#        └──────NCCL──────────────┘
#
# NCCL uses:
#   - NVLink (if available): 600 GB/s bidirectional (A100)
#   - PCIe Gen4: 32 GB/s per direction
#   - InfiniBand: 200 Gb/s (multi-node)

# The actual AllReduce algorithm (Ring AllReduce):
# Step 1: Reduce-Scatter - each rank gets 1/N of the reduced gradient
# Step 2: All-Gather     - each rank broadcasts its chunk to all others
# Total bytes transferred: 2 * (N-1)/N * model_size
# For BERT-Large FP16: 2 * (N-1)/N * 670 MB

# DDP groups parameters into "buckets" for AllReduce efficiency.
# Default bucket size: 25MB. DeepSpeed uses 200MB buckets.

# How overlapping works:
#
# Timeline (backward pass):
# ┌──────────────────────────────────────────────────┐
# │ Layer 24 grad  │ Layer 23 grad  │ Layer 22 grad  │ ... (backward)
# └──────────────────────────────────────────────────┘
# ┌────────────────┐                                    
# │ AllReduce Bkt 1 │ (layers 24-20, fires as soon as full)
# └────────────────┘
#                  ┌────────────────┐
#                  │ AllReduce Bkt 2 │ (layers 19-15)
#                  └────────────────┘

# DeepSpeed ZeRO-2 uses ReduceScatter instead of AllReduce:
DS_CONFIG["zero_optimization"] = {
    "stage": 2,
    "reduce_scatter": True,         # Each rank gets 1/N of reduced grads
    "reduce_bucket_size": 2e8,     # 200MB buckets
    "overlap_comm": True,          # Overlap with backward compute
    "contiguous_gradients": True,  # Pack gradients contiguously for faster reduce
}

# How you'd replicate this EXACT run with PyTorch FSDP instead of DeepSpeed:

from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    ShardingStrategy,
    MixedPrecision,
    CPUOffload,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers.models.bert.modeling_bert import BertLayer

# 1. Define sharding policy: wrap each BertLayer as a FSDP unit
auto_wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={BertLayer},  # Each of the 24 encoder layers
)

# 2. Mixed precision: FP16 forward/backward, FP32 reduce
mp_policy = MixedPrecision(
    param_dtype=torch.float16,     # Parameters cast to FP16
    reduce_dtype=torch.float16,    # Gradient reduce in FP16
    buffer_dtype=torch.float16,    # Buffers (LayerNorm stats) in FP16
)

# 3. Wrap model — equivalent to DeepSpeed ZeRO-2 + CPU offload
model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,  # = ZeRO-2
    cpu_offload=CPUOffload(offload_params=False),      # Offload optimizer, not params
    mixed_precision=mp_policy,
    auto_wrap_policy=auto_wrap_policy,
    device_id=torch.cuda.current_device(),
)

# 4. Standard PyTorch optimizer (no DeepSpeed config needed)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)

# 5. Training loop is standard PyTorch:
for batch in train_loader:
    loss = model(**batch).loss
    loss.backward()               # FSDP handles gradient sharding
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    optimizer.zero_grad()

# DeepSpeed's FP16 config from our run:
"fp16": {
    "enabled": True,
    "loss_scale": 0,              # 0 = dynamic loss scaling
    "initial_scale_power": 16,    # Start at 2^16 = 65536
    "loss_scale_window": 1000,    # Double scale every 1000 good steps
}

# Dynamic loss scaling algorithm (what DeepSpeed does internally):
#
# loss_scale = 2^16 = 65536
# good_steps = 0
#
# for each step:
#     scaled_loss = loss * loss_scale
#     scaled_loss.backward()
#     gradients = [p.grad / loss_scale for p in params]  # unscale
#
#     if any gradient is inf or nan:
#         loss_scale /= 2          # halve scale
#         skip optimizer step       # discard this batch
#         good_steps = 0
#     else:
#         optimizer.step()          # apply update
#         good_steps += 1
#         if good_steps >= 1000:
#             loss_scale *= 2       # try larger scale
#             good_steps = 0

# FP16 vs BF16 precision comparison:
# ┌──────────┬───────────┬───────────┬────────────────────┐
# │ Format   │ Sign+Exp  │ Mantissa  │ Range              │
# ├──────────┼───────────┼───────────┼────────────────────┤
# │ FP32     │ 1+8 bits  │ 23 bits   │ ±3.4e38            │
# │ FP16     │ 1+5 bits  │ 10 bits   │ ±65504   (narrow!) │
# │ BF16     │ 1+8 bits  │ 7 bits    │ ±3.4e38  (FP32!)   │
# └──────────┴───────────┴───────────┴────────────────────┘
# BF16 has FP32's range → no loss scaling needed!
# But T4 doesn't support BF16 — requires Ampere (A100) or newer.
# That's why this run uses FP16 with dynamic loss scaling.

# Activation memory is the biggest consumer (~8.5 GB).
# If we were memory-constrained, gradient checkpointing trades compute for memory:

from torch.utils.checkpoint import checkpoint

# Without checkpointing (our run — activations stored for all 24 layers):
# Memory: O(num_layers) = 24 layers of activations
# Speed:  1× forward, 1× backward

# With checkpointing (recompute activations during backward):
# Memory: O(sqrt(num_layers)) = only ~5 layers stored
# Speed:  ~1.3× forward (33% slower — recomputes activations)

# DeepSpeed activation checkpointing config:
DS_CONFIG["activation_checkpointing"] = {
    "partition_activations": True,     # Shard activations across GPUs
    "cpu_checkpointing": True,        # Offload checkpoints to CPU
    "contiguous_memory_optimization": True,
    "number_checkpoints": 24,         # Checkpoint every encoder layer
}

# PyTorch native equivalent:
for layer in model.bert.encoder.layer:
    layer.forward = functools.partial(
        checkpoint, layer.forward, use_reentrant=False
    )

# We DON'T use this in our run because T4 has enough memory.
# Activation memory: ~8.5 GB, total: ~11 GB, T4 total: 15.8 GB
# Headroom: 4.8 GB — no need to trade speed for memory.

# Torchtune uses YAML configs + composable Python recipes.
# Equivalent of our DeepSpeed BERT run as a Torchtune config:

# config.yaml
model:
  _component_: torchtune.models.llama3.llama3_8b  # (or custom BERT recipe)

tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  path: /tmp/tokenizer.model

dataset:
  _component_: torchtune.datasets.text_completion_dataset
  source: imdb
  split: train

optimizer:
  _component_: torch.optim.AdamW
  lr: 2e-5
  weight_decay: 0.01

loss:
  _component_: torch.nn.CrossEntropyLoss

training:
  batch_size: 4
  epochs: 3
  gradient_accumulation_steps: 4      # Same as our DeepSpeed config
  max_seq_len: 256
  compile: True                        # torch.compile() — Torchtune integrates this
  enable_activation_checkpointing: False

precision: bf16                         # Torchtune prefers BF16 (Ampere+)

# Run with:
# tune run full_finetune_single_device --config config.yaml

# Inside a Torchtune distributed recipe (simplified):
# This is what runs when you do: tune run full_finetune_distributed

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torchtune.training import get_dtype, set_default_dtype
from torchtune.modules import TransformerDecoder

class FullFinetuneRecipeDistributed:

    def setup(self):
        # 1. Initialize process group
        torch.distributed.init_process_group(backend="nccl")

        # 2. Load model with FSDP wrapping
        with set_default_dtype(torch.bfloat16):
            model = self._setup_model()
        model = FSDP(model, ...)  # Shard across GPUs

        # 3. Compile for speed (PyTorch 2.x)
        if self.cfg.compile:
            model = torch.compile(model)  # Fuses ops, reduces memory

        # 4. Setup optimizer with foreach=True for speed
        optimizer = torch.optim.AdamW(
            model.parameters(),
            foreach=True,  # Batched optimizer — 20% faster than per-param
        )

    def train(self):
        for batch in self.dataloader:
            # Mixed precision context (BF16 on Ampere+)
            with torch.amp.autocast("cuda", dtype=torch.bfloat16):
                loss = model(batch)

            loss.backward()

            # Gradient clipping (same as our DeepSpeed clip=1.0)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

            optimizer.step()
            optimizer.zero_grad(set_to_none=True)  # Saves memory vs .zero_grad()

# Scaling this run to multi-node production training:
# (what changes from our single-GPU Colab setup)

# 1. LAUNCHER — replaces manual env vars
# deepspeed --num_nodes=4 --num_gpus=8 train.py --deepspeed ds_config.json
# or: torchrun --nproc_per_node=8 --nnodes=4 --rdzv_backend=c10d train.py

# 2. CONFIG changes for multi-GPU
DS_CONFIG_PROD = {
    "train_micro_batch_size_per_gpu": 4,
    "gradient_accumulation_steps": 1,    # Less accum — more GPUs compensate
    # Effective batch = 4 × 1 × 32 GPUs = 128

    "zero_optimization": {
        "stage": 3,                          # Full sharding at scale
        "offload_optimizer": {"device": "none"},  # No CPU offload with enough GPUs
        "offload_param": {"device": "none"},
    },
    "fp16": {"enabled": False},
    "bf16": {"enabled": True},             # BF16 on A100/H100 — no loss scaling

    "communication_data_type": "bf16",    # AllReduce in BF16
    "prescale_gradients": True,          # Scale before AllReduce for numerical stability
}

# 3. CHECKPOINTING for fault tolerance
# DeepSpeed:
model_engine.save_checkpoint("checkpoints/", tag=f"step_{global_step}")
# Saves: model params (sharded), optimizer state (sharded), scheduler, rng states
# On resume: automatically re-shards if GPU count changes

# FSDP equivalent:
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT):
    torch.save(model.state_dict(), "checkpoint.pt")

# 4. MONITORING (production equivalent of our MetricsLogger)
# - Weights & Biases: wandb.log({"loss": loss, "gpu_util": util})
# - TensorBoard: writer.add_scalar("loss", loss, step)
# - Prometheus + Grafana: DCGM exporter for GPU metrics
# - Our approach: custom JSON → D3.js (what this dashboard shows)