wehub-resource-sync
336 lines
13 KiB
Python
336 lines
13 KiB
Python
# 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()
|