# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team """ DeepSpeed ZeRO-3 training example with allgather overlap. Trains a GPT-2-style transformer on synthetic data for demonstration. Designed for single-node 8x AMD GPU setup. """ import argparse import math import os import time import torch import torch.nn as nn import deepspeed from deepspeed import comm as dist from deepspeed.accelerator import get_accelerator from torch.utils.data import Dataset, DataLoader # --------------------------------------------------------------------------- # Model: minimal GPT-2-style transformer # --------------------------------------------------------------------------- class CausalSelfAttention(nn.Module): def __init__(self, hidden_size, num_heads, max_seq_len, dropout=0.1): super().__init__() assert hidden_size % num_heads == 0 self.num_heads = num_heads self.head_dim = hidden_size // num_heads self.qkv = nn.Linear(hidden_size, 3 * hidden_size) self.proj = nn.Linear(hidden_size, hidden_size) self.attn_drop = nn.Dropout(dropout) self.proj_drop = nn.Dropout(dropout) self.register_buffer( "causal_mask", torch.tril(torch.ones(max_seq_len, max_seq_len)).view(1, 1, max_seq_len, max_seq_len), ) def forward(self, x): B, T, C = x.size() q, k, v = self.qkv(x).split(C, dim=-1) q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) scale = 1.0 / math.sqrt(self.head_dim) attn = (q @ k.transpose(-2, -1)) * scale attn = attn.masked_fill(self.causal_mask[:, :, :T, :T] == 0, float("-inf")) attn = torch.softmax(attn, dim=-1) attn = self.attn_drop(attn) out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) return self.proj_drop(self.proj(out)) class TransformerBlock(nn.Module): def __init__(self, hidden_size, num_heads, max_seq_len, dropout=0.1): super().__init__() self.ln1 = nn.LayerNorm(hidden_size) self.attn = CausalSelfAttention(hidden_size, num_heads, max_seq_len, dropout) self.ln2 = nn.LayerNorm(hidden_size) self.mlp = nn.Sequential( nn.Linear(hidden_size, 4 * hidden_size), nn.GELU(), nn.Linear(4 * hidden_size, hidden_size), nn.Dropout(dropout), ) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x class GPT2Model(nn.Module): def __init__(self, vocab_size, hidden_size, num_layers, num_heads, max_seq_len, dropout=0.1): super().__init__() self.tok_emb = nn.Embedding(vocab_size, hidden_size) self.pos_emb = nn.Embedding(max_seq_len, hidden_size) self.drop = nn.Dropout(dropout) self.blocks = nn.Sequential( *[TransformerBlock(hidden_size, num_heads, max_seq_len, dropout) for _ in range(num_layers)]) self.ln_f = nn.LayerNorm(hidden_size) self.head = nn.Linear(hidden_size, vocab_size, bias=False) def forward(self, input_ids, labels=None): B, T = input_ids.size() pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0) x = self.drop(self.tok_emb(input_ids) + self.pos_emb(pos)) x = self.blocks(x) x = self.ln_f(x) logits = self.head(x) loss = None if labels is not None: loss = nn.functional.cross_entropy( logits.view(-1, logits.size(-1)), labels.view(-1), ) return loss, logits # --------------------------------------------------------------------------- # Synthetic dataset # --------------------------------------------------------------------------- class SyntheticTextDataset(Dataset): """Generates synthetic token sequences for perf/correctness testing.""" def __init__(self, vocab_size, seq_len, num_samples, seed=42, mode="random"): self.vocab_size = vocab_size self.seq_len = seq_len self.num_samples = num_samples self.seed = seed self.mode = mode def __len__(self): return self.num_samples def __getitem__(self, idx): if self.mode == "random": g = torch.Generator() g.manual_seed(self.seed + idx) tokens = torch.randint(0, self.vocab_size, (self.seq_len + 1, ), generator=g) elif self.mode == "arange": start = (self.seed + idx) % self.vocab_size tokens = (torch.arange(self.seq_len + 1, dtype=torch.long) + start) % self.vocab_size elif self.mode == "repeat": v = (self.seed + idx) % self.vocab_size tokens = torch.full((self.seq_len + 1, ), v, dtype=torch.long) else: raise ValueError(f"Unsupported data mode: {self.mode}") return tokens[:-1], tokens[1:] class WikitextDataset(Dataset): """Real text dataset from HuggingFace wikitext-2 / wikitext-103.""" def __init__(self, vocab_size, seq_len, num_samples, split="train", dataset_name="wikitext-2-raw-v1"): from datasets import load_dataset from transformers import GPT2TokenizerFast tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") raw = load_dataset("wikitext", dataset_name, split=split) text = "\n\n".join([t for t in raw["text"] if t.strip()]) all_ids = tokenizer.encode(text) self.seq_len = seq_len self.samples = [] for i in range(0, len(all_ids) - seq_len - 1, seq_len): self.samples.append(torch.tensor(all_ids[i:i + seq_len + 1], dtype=torch.long)) if len(self.samples) >= num_samples: break def __len__(self): return len(self.samples) def __getitem__(self, idx): tokens = self.samples[idx] return tokens[:-1], tokens[1:] # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def parse_args(): parser = argparse.ArgumentParser(description="DeepSpeed ZeRO-3 training with allgather overlap") parser.add_argument("--vocab_size", type=int, default=50257) parser.add_argument("--hidden_size", type=int, default=4096) parser.add_argument("--num_layers", type=int, default=48) parser.add_argument("--num_heads", type=int, default=32) parser.add_argument("--max_seq_len", type=int, default=2048) parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--num_samples", type=int, default=10000) parser.add_argument("--train_steps", type=int, default=2000) parser.add_argument("--data_mode", type=str, default="random", choices=["random", "arange", "repeat", "wikitext2", "wikitext103"], help="Data mode. random/arange/repeat are synthetic; wikitext2/wikitext103 use real text.") parser.add_argument("--local_rank", type=int, default=-1) parser = deepspeed.add_config_arguments(parser) return parser.parse_args() def main(): args = parse_args() ds_config_path = args.deepspeed_config if ds_config_path and not os.path.isfile(ds_config_path): script_dir = os.path.dirname(os.path.abspath(__file__)) ds_config_path = os.path.join(script_dir, ds_config_path) args.deepspeed_config = ds_config_path deepspeed.init_distributed(dist_backend="cpu:gloo,cuda:nccl") local_rank = args.local_rank get_accelerator().set_device(local_rank) torch.manual_seed(42) get_accelerator().manual_seed_all(42) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False with deepspeed.zero.Init(config_dict_or_path=ds_config_path): model = GPT2Model( vocab_size=args.vocab_size, hidden_size=args.hidden_size, num_layers=args.num_layers, num_heads=args.num_heads, max_seq_len=args.max_seq_len, dropout=0.0, ) total_params = sum(p.numel() for p in model.parameters()) num_gpus = dist.get_world_size() if local_rank == 0: print(f"Model parameters: {total_params / 1e6:.1f}M") print(f"GPUs: {num_gpus}") # FLOPs per token (forward + backward): 6*params + 12*L*H*S # Reference: "Efficient Large-Scale Language Model Training on GPU Clusters # Using Megatron-LM" (Narayanan et al., 2021) flops_per_token = 6 * total_params + 12 * args.num_layers * args.hidden_size * args.max_seq_len if args.data_mode in ("wikitext2", "wikitext103"): ds_name = "wikitext-2-raw-v1" if args.data_mode == "wikitext2" else "wikitext-103-raw-v1" dataset = WikitextDataset(args.vocab_size, args.max_seq_len, args.num_samples, dataset_name=ds_name) else: dataset = SyntheticTextDataset(args.vocab_size, args.max_seq_len, args.num_samples, mode=args.data_mode) if local_rank == 0: if args.data_mode == "random": print(f"Data mode: random (expected CE floor ~ ln(vocab) = {math.log(args.vocab_size):.4f})") elif args.data_mode in ("wikitext2", "wikitext103"): print(f"Data mode: {args.data_mode} (real text, {len(dataset)} samples)") else: print(f"Data mode: {args.data_mode} (learnable pattern, loss should decrease)") model_engine, optimizer, _, lr_scheduler = deepspeed.initialize( args=args, model=model, ) sampler = torch.utils.data.distributed.DistributedSampler( dataset, shuffle=False, seed=42, ) train_loader = DataLoader( dataset, batch_size=model_engine.train_micro_batch_size_per_gpu(), sampler=sampler, num_workers=0, pin_memory=True, ) device = model_engine.device global_batch = model_engine.train_batch_size() tokens_per_step = global_batch * args.max_seq_len warmup_steps = min(50, args.train_steps // 10) step = 0 step_times = [] t_start = time.time() t_steady = None while step < args.train_steps: for batch in train_loader: if step >= args.train_steps: break get_accelerator().synchronize() t_step_start = time.time() input_ids = batch[0].to(device) labels = batch[1].to(device) loss, _ = model_engine(input_ids, labels=labels) model_engine.backward(loss) model_engine.step() get_accelerator().synchronize() step_time_ms = (time.time() - t_step_start) * 1000 if step == warmup_steps: t_steady = time.time() if step >= warmup_steps: step_times.append(step_time_ms) if step % 10 == 0 and local_rank == 0: if step_times: import numpy as np recent = np.array(step_times[-20:]) avg_ms = recent.mean() cur_samples_per_sec = global_batch / (avg_ms / 1000) cur_tokens_per_sec = cur_samples_per_sec * args.max_seq_len cur_tflops_per_gpu = cur_tokens_per_sec * flops_per_token / 1e12 / num_gpus else: avg_ms = step_time_ms cur_tflops_per_gpu = 0.0 cur_samples_per_sec = 0.0 print(f"step {step:5d} | loss {loss.item():.4f} | " f"lr {lr_scheduler.get_last_lr()[0]:.6f} | " f"{cur_samples_per_sec:.1f} samples/s | " f"{cur_tflops_per_gpu:.2f} TFLOPS/GPU | " f"step {avg_ms:.1f} ms") step += 1 if local_rank == 0: import numpy as np total_time = time.time() - t_start st = np.array(step_times) steady_steps = len(st) steady_time = time.time() - t_steady if t_steady else total_time steady_samples_per_sec = steady_steps * global_batch / steady_time steady_tokens_per_sec = steady_samples_per_sec * args.max_seq_len steady_tflops = steady_tokens_per_sec * flops_per_token / 1e12 steady_tflops_per_gpu = steady_tflops / num_gpus print(f"\n{'=' * 70}") print(f"Training complete: {args.train_steps} steps in {total_time:.1f}s") print(f" (warmup={warmup_steps} steps skipped, measured {steady_steps} steps)") print(f"{'=' * 70}") print(f" Throughput : {steady_samples_per_sec:.1f} samples/s") print(f" TFLOPS : {steady_tflops:.1f} (total) | {steady_tflops_per_gpu:.2f} (per GPU)") print(f" Step time (ms) : avg {st.mean():.1f} | p50 {np.median(st):.1f} | " f"p99 {np.percentile(st, 99):.1f} | min {st.min():.1f} | max {st.max():.1f}") print(f"{'=' * 70}") if __name__ == "__main__": main()