"""Reference code for GPT-2 training and inference.Will save the model weights into files, to be read from C as initialization. References:1) the official GPT-2 TensorFlow implementation released by OpenAI:https://github.com/openai/gpt-2/blob/master/src/model.py2) huggingface/transformers PyTorch implementation:https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py Example launches to only benchmark the speed of bfloat16 compiled GPU training:1 GPU:python train_gpt2.py --write_tensors=0 --num_iterations=50 --sequence_length=1024 --compile=1 --tensorcores=1 --dtype=bfloat16you can also turn on flash-attention by appending --flash=14 GPU:torchrun --standalone --nproc_per_node=4 train_gpt2.py --write_tensors=0 --num_iterations=50 --sequence_length=1024 --compile=1 --tensorcores=1 --dtype=bfloat16""" import osimport mathimport globimport structimport inspectfrom contextlib import nullcontextfrom dataclasses import dataclass import numpy as npimport torchimport torch.nn as nnfrom torch.nn import functional as Fimport torch._inductor.config as configfrom torch.nn.parallel import DistributedDataParallel as DDPfrom torch.distributed import init_process_group, destroy_process_groupfrom torch.distributed.optim import ZeroRedundancyOptimizer # -----------------------------------------------------------------------------# PyTorch nn.Module definitions for the GPT-2 model class NewGELU(nn.Module): """Careful there are a few versions of GeLU, this one is the exact one used by OpenAI""" def forward(self, input): return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) # using a global to toggle flash-attentionFLASH = 0 class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 # key, query, value projections for all heads, but in a batch self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) # output projection self.c_proj = nn.Linear(config.n_embd, config.n_embd) self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 # regularization self.n_head = config.n_head self.n_embd = config.n_embd # not really a 'bias', more of a mask, but following the OpenAI/HF naming though self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) .view(1, 1, config.block_size, config.block_size)) def forward(self, x): B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) # calculate query, key, values for all heads in batch and move head forward to be the batch dim qkv = self.c_attn(x) q, k, v = qkv.split(self.n_embd, dim=2) k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) if FLASH: # flashattention y = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: # manual implementation of attention # this materializes the large (T,T) matrix for all the queries and keys att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf')) att = F.softmax(att, dim=-1) y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side # output projection y = self.c_proj(y) return y class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd) self.gelu = NewGELU() self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd) self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 def forward(self, x): x = self.c_fc(x) x = self.gelu(x) x = self.c_proj(x) return x class Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd) self.attn = CausalSelfAttention(config) self.ln_2 = nn.LayerNorm(config.n_embd) self.mlp = MLP(config) def forward(self, x): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x # -----------------------------------------------------------------------------# The main GPT-2 model @dataclassclass GPTConfig: block_size: int = 1024 vocab_size: int = 50257 n_layer: int = 12 n_head: int = 12 n_embd: int = 768 class GPT(nn.Module): def __init__(self, config): super().__init__() self.config = config self.transformer = nn.ModuleDict(dict( wte = nn.Embedding(config.vocab_size, config.n_embd), wpe = nn.Embedding(config.block_size, config.n_embd), h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f = nn.LayerNorm(config.n_embd), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.lm_head.LLMC_SKIP_INIT = 1 # don't init this one, we will tie weights self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying # init all weights, use a torch rng object to be very careful self.init_rng = torch.Generator() self.init_rng.manual_seed(42) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): # apply special scaled init to the residual projections, per GPT-2 paper std = 0.02 if not hasattr(module, 'LLMC_RESIDUAL_SCALE_FLAG') else 0.02/math.sqrt(2 * self.config.n_layer) # we want to skip initializing lm_head, which shares parameters with wte # and wte was already initialized down below during the Embedding init if not hasattr(module, 'LLMC_SKIP_INIT'): torch.nn.init.normal_(module.weight, mean=0.0, std=std, generator=self.init_rng) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02, generator=self.init_rng) def forward(self, idx, targets=None, return_logits=True): device = idx.device b, t = idx.size() assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}" pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t) # forward the GPT model itself tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd) x = tok_emb + pos_emb for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) if targets is not None: # if we are given some desired targets also calculate the loss logits = self.lm_head(x) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) else: # inference-time mini-optimization: only forward the lm_head on the very last position logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim loss = None # there are performance reasons why not returning logits is prudent, if not needed if not return_logits: logits = None return logits, loss @classmethod def from_pretrained(cls, model_type): """Loads pretrained GPT-2 model weights from huggingface""" assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'} from transformers import GPT2LMHeadModel print("loading weights from pretrained gpt: %s" % model_type) # n_layer, n_head and n_embd are determined from model_type config_args = { 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params }[model_type] config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints # create a from-scratch initialized minGPT model config = GPTConfig(**config_args) model = GPT(config) sd = model.state_dict() sd_keys = sd.keys() sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param # init a huggingface/transformers model model_hf = GPT2LMHeadModel.from_pretrained(model_type) sd_hf = model_hf.state_dict() # copy while ensuring all of the parameters are aligned and match in names and shapes sd_keys_hf = sd_hf.keys() sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer) transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight'] # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear # this means that we have to transpose these weights when we import them assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}" for k in sd_keys_hf: if any(k.endswith(w) for w in transposed): # special treatment for the Conv1D weights we need to transpose assert sd_hf[k].shape[::-1] == sd[k].shape with torch.no_grad(): sd[k].copy_(sd_hf[k].t()) else: # vanilla copy over the other parameters assert sd_hf[k].shape == sd[k].shape with torch.no_grad(): sd[k].copy_(sd_hf[k]) return model def configure_optimizers(self, weight_decay, learning_rate, betas, device_type, zero_stage): # start with all of the candidate parameters param_dict = {pn: p for pn, p in self.named_parameters()} # filter out those that do not require grad param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] optim_groups = [ {'params': decay_params, 'weight_decay': weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0} ] num_decay_params = sum(p.numel() for p in decay_params) num_nodecay_params = sum(p.numel() for p in nodecay_params) print0(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") print0(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") # Create AdamW optimizer and use the fused version if it is available fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters use_fused = fused_available and device_type == 'cuda' print0(f"using fused AdamW: {use_fused}") if zero_stage == 1: print0("using ZeroRedundancyOptimizer") optimizer = ZeroRedundancyOptimizer(**optim_groups[0], optimizer_class=torch.optim.AdamW, lr=learning_rate, betas=betas, fused=use_fused) optimizer.add_param_group(optim_groups[1]) else: print0("using regular AdamW") optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=use_fused) return optimizer @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): """ Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete the sequence max_new_tokens times, feeding the predictions back into the model each time. Most likely you'll want to make sure to be in model.eval() mode of operation for this. """ for _ in range(max_new_tokens): # if the sequence context is growing too long we must crop it at block_size idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] # forward the model to get the logits for the index in the sequence logits, _ = self(idx_cond) # pluck the logits at the final step and scale by desired temperature logits = logits[:, -1, :] / temperature # optionally crop the logits to only the top k options if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float('Inf') # apply softmax to convert logits to (normalized) probabilities probs = F.softmax(logits, dim=-1) # sample from the distribution idx_next = torch.multinomial(probs, num_samples=1) # append sampled index to the running sequence and continue idx = torch.cat((idx, idx_next), dim=1) return idx # -----------------------------------------------------------------------------# Our own simple Distributed Data Loader def _peek_data_shard(filename): # only reads the header, returns header data with open(filename, "rb") as f: # first read the header, which is 256 int32 integers (4 bytes each) header = np.frombuffer(f.read(256*4), dtype=np.int32) if header[0] != 20240520: print("ERROR: magic number mismatch in the data .bin file!") print("---> HINT: Are you passing in a correct file with --input_bin?") print("---> HINT: Dataset encoding changed recently, re-run data prepro or refer again to README") print("---> HINT: For example re-run: `python dev/data/tinyshakespeare.py`, then re-try") exit(1) assert header[1] == 1, "unsupported version" ntok = header[2] # number of tokens (claimed) return ntok # for now just return the number of tokens def _load_data_shard(filename): with open(filename, "rb") as f: # first read the header, which is 256 int32 integers (4 bytes each) header = np.frombuffer(f.read(256*4), dtype=np.int32) assert header[0] == 20240520, "magic number mismatch in the data .bin file" assert header[1] == 1, "unsupported version" ntok = header[2] # number of tokens (claimed) # the rest of it are tokens, stored as uint16 tokens = np.frombuffer(f.read(), dtype=np.uint16) assert len(tokens) == ntok, "number of tokens read does not match header?" return tokens class DistributedDataLoader: def __init__(self, filename_pattern, B, T, process_rank, num_processes): self.process_rank = process_rank self.num_processes = num_processes self.B = B self.T = T # glob files that match the pattern self.files = sorted(glob.glob(filename_pattern)) assert len(self.files) > 0, f"did not find any files that match the pattern {filename_pattern}" # load and validate all data shards, count number of tokens in total ntok_total = 0 for fname in self.files: shard_ntok = _peek_data_shard(fname) assert shard_ntok >= num_processes * B * T + 1 ntok_total += shard_ntok self.ntok_total = ntok_total print0(f"DataLoader: total number of tokens: {ntok_total:,} across {len(self.files)} files") # kick things off self.current_shard = None self.reset() def reset(self): # we're being a bit clever here: if we already had shard 0 loaded, # then don't do the work to reload it, just reset the pointer if self.current_shard != 0: self.current_shard = 0 self.tokens = _load_data_shard(self.files[self.current_shard]) self.current_position = self.process_rank * self.B * self.T def advance(self): # advance to next data shard self.current_shard = (self.current_shard + 1) % len(self.files) self.current_position = self.process_rank * self.B * self.T self.tokens = _load_data_shard(self.files[self.current_shard]) def next_batch(self): B = self.B T = self.T buf = self.tokens[self.current_position : self.current_position+B*T+1] buf = torch.tensor(buf.astype(np.int32), dtype=torch.long) x = (buf[:-1]).view(B, T) # inputs y = (buf[1:]).view(B, T) # targets # advance the start pointer in current shard self.current_position += B * T * self.num_processes # if loading the next batch would be out of bounds advance the shard if self.current_position + (B * T * self.num_processes + 1) > len(self.tokens): self.advance() return x, y # -----------------------------------------------------------------------------# Python -> C bridge utilities for saving params/grads/activations to .bin files def write_fp32(tensor, file): t = tensor.detach().cpu().to(torch.float32) b = t.numpy().tobytes() file.write(b) def write_bf16(tensor, file): t = tensor.detach().cpu().to(torch.bfloat16) # numpy doesn't have bf16 datatype so we have to trick it t = t.view(torch.int16) # trick: reinterpret as int16 b = t.numpy().tobytes() file.write(b) def write_tensors(model_tensors, L, file, dtype): # writes the GPT-2 model's weights to a binary file assert dtype in {"float32", "bfloat16"} write_fun = write_fp32 if dtype == "float32" else write_bf16 write_fun(model_tensors["transformer.wte.weight"], file) # (V, C) write_fun(model_tensors["transformer.wpe.weight"], file) # (T, C) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_1.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_1.bias"], file) for i in range(L): # (L, 3C, C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_attn.weight"], file) for i in range(L): # (L, 3C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_attn.bias"], file) for i in range(L): # (L, C, C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_proj.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_proj.bias"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_2.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_2.bias"], file) for i in range(L): # (L, 4C, C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_fc.weight"], file) for i in range(L): # (L, 4C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_fc.bias"], file) for i in range(L): # (L, C, 4C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_proj.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_proj.bias"], file) write_fun(model_tensors["transformer.ln_f.weight"], file) # (C, ) write_fun(model_tensors["transformer.ln_f.bias"], file) # (C, ) @torch.no_grad()def pad_vocab(tensor, multiple=128, value=0): """ The dimension of the vocab size in GPT-2 is 50,257 which is unfortunately a very unfriendly number for a lot of matrix operations on the GPU. So we pad it to the nearest friendlier multiple, e.g. 50,304 if multiple=128 when we export the weights into C land. This is a NOOP algorithmically and is only done to make the tensor operations more efficient. """ assert tensor.ndim == 2 V, C = tensor.shape assert V == 50257, "just being defensive here" # calculate padded vocab size by rounding up to nearest multiple Vp = ((V + multiple - 1) // multiple) * multiple # pad the tensor pad_rows = Vp - V padded = tensor if pad_rows == 0 else F.pad(tensor, (0, 0, 0, pad_rows), value=value) assert padded.shape == (Vp, C) return padded def write_model(model, filename, dtype): # everything we need to instantiate the model # 1) header is: version int, GPTConfig ints, padding to 1024 bytes assert dtype in {"float32", "bfloat16"} # float16 todo maybe later version = { "float32": 3, # 3: all tensors are fp32, padded vocab "bfloat16": 5, # 5: all tensors are bf16, padded vocab }[dtype] header = torch.zeros(256, dtype=torch.int32) header[0] = 20240326 # magic header[1] = version # checkpoint version header[2] = model.config.block_size header[3] = model.config.vocab_size header[4] = model.config.n_layer header[5] = model.config.n_head header[6] = model.config.n_embd # 2) the parameters follow the header params = {name: param.cpu() for name, param in model.named_parameters()} # pad the vocab to a multiple of 128 here at export, for efficiency in C wte = params["transformer.wte.weight"] # (V, C) wte_padded = pad_vocab(wte) # (Vp, C) params["transformer.wte.weight"] = wte_padded # (Vp, C) print(f"padded vocab size from {wte.size(0)} to {wte_padded.size(0)}") header[7] = wte_padded.size(0) # padded vocab size store in header # now write to file with open(filename, "wb") as file: file.write(header.numpy().tobytes()) # header write_tensors(params, model.config.n_layer, file, dtype) # params print(f"wrote {filename}") def write_state(model, x, y, logits, loss, filename): # the state is used for debugging. # it contains information about the input, logits, loss, and the parameter gradients # this can be used for checking the computation correctness in C header = torch.zeros(256, dtype=torch.int32) header[0] = 20240327 # magic header[1] = 2 # run state version = 2 (1 -> 2 for padded vocab changes) header[2] = x.size(0) # batch size of the batch, B header[3] = x.size(1) # temporal extent of the batch, T grads = {name: param.grad.cpu() for name, param in model.named_parameters()} # pad the vocab grads here as well, to mirror write_model wte_grad = grads["transformer.wte.weight"] # (V, C) wte_grad_padded = pad_vocab(wte_grad, value=0) # (Vp, C) # TODO later maybe pad with nan? grads["transformer.wte.weight"] = wte_grad_padded # (Vp, C) print(f"padded vocab size in reference grads from {wte_grad.size(0)} to {wte_grad_padded.size(0)}") with open(filename, "wb") as file: # header file.write(header.numpy().tobytes()) # input x file.write(x.cpu().numpy().astype("int32").tobytes()) # (B, T) # targets y file.write(y.cpu().numpy().astype("int32").tobytes()) # (B, T) # logits (result of the model forward pass)Lines 1–500 of 858. Load more
""" Reference code for GPT-2 training and inference. Will save the model weights into files, to be read from C as initialization. References: 1) the official GPT-2 TensorFlow implementation released by OpenAI: https://github.com/openai/gpt-2/blob/master/src/model.py 2) huggingface/transformers PyTorch implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py Example launches to only benchmark the speed of bfloat16 compiled GPU training: 1 GPU: python train_gpt2.py --write_tensors=0 --num_iterations=50 --sequence_length=1024 --compile=1 --tensorcores=1 --dtype=bfloat16 you can also turn on flash-attention by appending --flash=1 4 GPU: torchrun --standalone --nproc_per_node=4 train_gpt2.py --write_tensors=0 --num_iterations=50 --sequence_length=1024 --compile=1 --tensorcores=1 --dtype=bfloat16 """ import os import math import glob import struct import inspect from contextlib import nullcontext from dataclasses import dataclass import numpy as np import torch import torch.nn as nn from torch.nn import functional as F import torch._inductor.config as config from torch.nn.parallel import DistributedDataParallel as DDP from torch.distributed import init_process_group, destroy_process_group from torch.distributed.optim import ZeroRedundancyOptimizer # ----------------------------------------------------------------------------- # PyTorch nn.Module definitions for the GPT-2 model class NewGELU(nn.Module): """Careful there are a few versions of GeLU, this one is the exact one used by OpenAI""" def forward(self, input): return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * torch.pow(input, 3.0)))) # using a global to toggle flash-attention FLASH = 0 class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 # key, query, value projections for all heads, but in a batch self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) # output projection self.c_proj = nn.Linear(config.n_embd, config.n_embd) self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 # regularization self.n_head = config.n_head self.n_embd = config.n_embd # not really a 'bias', more of a mask, but following the OpenAI/HF naming though self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)) .view(1, 1, config.block_size, config.block_size)) def forward(self, x): B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd) # calculate query, key, values for all heads in batch and move head forward to be the batch dim qkv = self.c_attn(x) q, k, v = qkv.split(self.n_embd, dim=2) k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) if FLASH: # flashattention y = F.scaled_dot_product_attention(q, k, v, is_causal=True) else: # manual implementation of attention # this materializes the large (T,T) matrix for all the queries and keys att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf')) att = F.softmax(att, dim=-1) y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side # output projection y = self.c_proj(y) return y class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd) self.gelu = NewGELU() self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd) self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 def forward(self, x): x = self.c_fc(x) x = self.gelu(x) x = self.c_proj(x) return x class Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd) self.attn = CausalSelfAttention(config) self.ln_2 = nn.LayerNorm(config.n_embd) self.mlp = MLP(config) def forward(self, x): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x # ----------------------------------------------------------------------------- # The main GPT-2 model @dataclass class GPTConfig: block_size: int = 1024 vocab_size: int = 50257 n_layer: int = 12 n_head: int = 12 n_embd: int = 768 class GPT(nn.Module): def __init__(self, config): super().__init__() self.config = config self.transformer = nn.ModuleDict(dict( wte = nn.Embedding(config.vocab_size, config.n_embd), wpe = nn.Embedding(config.block_size, config.n_embd), h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f = nn.LayerNorm(config.n_embd), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.lm_head.LLMC_SKIP_INIT = 1 # don't init this one, we will tie weights self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying # init all weights, use a torch rng object to be very careful self.init_rng = torch.Generator() self.init_rng.manual_seed(42) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): # apply special scaled init to the residual projections, per GPT-2 paper std = 0.02 if not hasattr(module, 'LLMC_RESIDUAL_SCALE_FLAG') else 0.02/math.sqrt(2 * self.config.n_layer) # we want to skip initializing lm_head, which shares parameters with wte # and wte was already initialized down below during the Embedding init if not hasattr(module, 'LLMC_SKIP_INIT'): torch.nn.init.normal_(module.weight, mean=0.0, std=std, generator=self.init_rng) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02, generator=self.init_rng) def forward(self, idx, targets=None, return_logits=True): device = idx.device b, t = idx.size() assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}" pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t) # forward the GPT model itself tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd) pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd) x = tok_emb + pos_emb for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) if targets is not None: # if we are given some desired targets also calculate the loss logits = self.lm_head(x) loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) else: # inference-time mini-optimization: only forward the lm_head on the very last position logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim loss = None # there are performance reasons why not returning logits is prudent, if not needed if not return_logits: logits = None return logits, loss @classmethod def from_pretrained(cls, model_type): """Loads pretrained GPT-2 model weights from huggingface""" assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'} from transformers import GPT2LMHeadModel print("loading weights from pretrained gpt: %s" % model_type) # n_layer, n_head and n_embd are determined from model_type config_args = { 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params }[model_type] config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints # create a from-scratch initialized minGPT model config = GPTConfig(**config_args) model = GPT(config) sd = model.state_dict() sd_keys = sd.keys() sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param # init a huggingface/transformers model model_hf = GPT2LMHeadModel.from_pretrained(model_type) sd_hf = model_hf.state_dict() # copy while ensuring all of the parameters are aligned and match in names and shapes sd_keys_hf = sd_hf.keys() sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer) transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight'] # basically the openai checkpoints use a "Conv1D" module, but we only want to use a vanilla Linear # this means that we have to transpose these weights when we import them assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}" for k in sd_keys_hf: if any(k.endswith(w) for w in transposed): # special treatment for the Conv1D weights we need to transpose assert sd_hf[k].shape[::-1] == sd[k].shape with torch.no_grad(): sd[k].copy_(sd_hf[k].t()) else: # vanilla copy over the other parameters assert sd_hf[k].shape == sd[k].shape with torch.no_grad(): sd[k].copy_(sd_hf[k]) return model def configure_optimizers(self, weight_decay, learning_rate, betas, device_type, zero_stage): # start with all of the candidate parameters param_dict = {pn: p for pn, p in self.named_parameters()} # filter out those that do not require grad param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad} # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no. # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't. decay_params = [p for n, p in param_dict.items() if p.dim() >= 2] nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2] optim_groups = [ {'params': decay_params, 'weight_decay': weight_decay}, {'params': nodecay_params, 'weight_decay': 0.0} ] num_decay_params = sum(p.numel() for p in decay_params) num_nodecay_params = sum(p.numel() for p in nodecay_params) print0(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") print0(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") # Create AdamW optimizer and use the fused version if it is available fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters use_fused = fused_available and device_type == 'cuda' print0(f"using fused AdamW: {use_fused}") if zero_stage == 1: print0("using ZeroRedundancyOptimizer") optimizer = ZeroRedundancyOptimizer(**optim_groups[0], optimizer_class=torch.optim.AdamW, lr=learning_rate, betas=betas, fused=use_fused) optimizer.add_param_group(optim_groups[1]) else: print0("using regular AdamW") optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, fused=use_fused) return optimizer @torch.no_grad() def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None): """ Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete the sequence max_new_tokens times, feeding the predictions back into the model each time. Most likely you'll want to make sure to be in model.eval() mode of operation for this. """ for _ in range(max_new_tokens): # if the sequence context is growing too long we must crop it at block_size idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:] # forward the model to get the logits for the index in the sequence logits, _ = self(idx_cond) # pluck the logits at the final step and scale by desired temperature logits = logits[:, -1, :] / temperature # optionally crop the logits to only the top k options if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) logits[logits < v[:, [-1]]] = -float('Inf') # apply softmax to convert logits to (normalized) probabilities probs = F.softmax(logits, dim=-1) # sample from the distribution idx_next = torch.multinomial(probs, num_samples=1) # append sampled index to the running sequence and continue idx = torch.cat((idx, idx_next), dim=1) return idx # ----------------------------------------------------------------------------- # Our own simple Distributed Data Loader def _peek_data_shard(filename): # only reads the header, returns header data with open(filename, "rb") as f: # first read the header, which is 256 int32 integers (4 bytes each) header = np.frombuffer(f.read(256*4), dtype=np.int32) if header[0] != 20240520: print("ERROR: magic number mismatch in the data .bin file!") print("---> HINT: Are you passing in a correct file with --input_bin?") print("---> HINT: Dataset encoding changed recently, re-run data prepro or refer again to README") print("---> HINT: For example re-run: `python dev/data/tinyshakespeare.py`, then re-try") exit(1) assert header[1] == 1, "unsupported version" ntok = header[2] # number of tokens (claimed) return ntok # for now just return the number of tokens def _load_data_shard(filename): with open(filename, "rb") as f: # first read the header, which is 256 int32 integers (4 bytes each) header = np.frombuffer(f.read(256*4), dtype=np.int32) assert header[0] == 20240520, "magic number mismatch in the data .bin file" assert header[1] == 1, "unsupported version" ntok = header[2] # number of tokens (claimed) # the rest of it are tokens, stored as uint16 tokens = np.frombuffer(f.read(), dtype=np.uint16) assert len(tokens) == ntok, "number of tokens read does not match header?" return tokens class DistributedDataLoader: def __init__(self, filename_pattern, B, T, process_rank, num_processes): self.process_rank = process_rank self.num_processes = num_processes self.B = B self.T = T # glob files that match the pattern self.files = sorted(glob.glob(filename_pattern)) assert len(self.files) > 0, f"did not find any files that match the pattern {filename_pattern}" # load and validate all data shards, count number of tokens in total ntok_total = 0 for fname in self.files: shard_ntok = _peek_data_shard(fname) assert shard_ntok >= num_processes * B * T + 1 ntok_total += shard_ntok self.ntok_total = ntok_total print0(f"DataLoader: total number of tokens: {ntok_total:,} across {len(self.files)} files") # kick things off self.current_shard = None self.reset() def reset(self): # we're being a bit clever here: if we already had shard 0 loaded, # then don't do the work to reload it, just reset the pointer if self.current_shard != 0: self.current_shard = 0 self.tokens = _load_data_shard(self.files[self.current_shard]) self.current_position = self.process_rank * self.B * self.T def advance(self): # advance to next data shard self.current_shard = (self.current_shard + 1) % len(self.files) self.current_position = self.process_rank * self.B * self.T self.tokens = _load_data_shard(self.files[self.current_shard]) def next_batch(self): B = self.B T = self.T buf = self.tokens[self.current_position : self.current_position+B*T+1] buf = torch.tensor(buf.astype(np.int32), dtype=torch.long) x = (buf[:-1]).view(B, T) # inputs y = (buf[1:]).view(B, T) # targets # advance the start pointer in current shard self.current_position += B * T * self.num_processes # if loading the next batch would be out of bounds advance the shard if self.current_position + (B * T * self.num_processes + 1) > len(self.tokens): self.advance() return x, y # ----------------------------------------------------------------------------- # Python -> C bridge utilities for saving params/grads/activations to .bin files def write_fp32(tensor, file): t = tensor.detach().cpu().to(torch.float32) b = t.numpy().tobytes() file.write(b) def write_bf16(tensor, file): t = tensor.detach().cpu().to(torch.bfloat16) # numpy doesn't have bf16 datatype so we have to trick it t = t.view(torch.int16) # trick: reinterpret as int16 b = t.numpy().tobytes() file.write(b) def write_tensors(model_tensors, L, file, dtype): # writes the GPT-2 model's weights to a binary file assert dtype in {"float32", "bfloat16"} write_fun = write_fp32 if dtype == "float32" else write_bf16 write_fun(model_tensors["transformer.wte.weight"], file) # (V, C) write_fun(model_tensors["transformer.wpe.weight"], file) # (T, C) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_1.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_1.bias"], file) for i in range(L): # (L, 3C, C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_attn.weight"], file) for i in range(L): # (L, 3C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_attn.bias"], file) for i in range(L): # (L, C, C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_proj.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.attn.c_proj.bias"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_2.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.ln_2.bias"], file) for i in range(L): # (L, 4C, C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_fc.weight"], file) for i in range(L): # (L, 4C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_fc.bias"], file) for i in range(L): # (L, C, 4C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_proj.weight"], file) for i in range(L): # (L, C) write_fun(model_tensors[f"transformer.h.{i}.mlp.c_proj.bias"], file) write_fun(model_tensors["transformer.ln_f.weight"], file) # (C, ) write_fun(model_tensors["transformer.ln_f.bias"], file) # (C, ) @torch.no_grad() def pad_vocab(tensor, multiple=128, value=0): """ The dimension of the vocab size in GPT-2 is 50,257 which is unfortunately a very unfriendly number for a lot of matrix operations on the GPU. So we pad it to the nearest friendlier multiple, e.g. 50,304 if multiple=128 when we export the weights into C land. This is a NOOP algorithmically and is only done to make the tensor operations more efficient. """ assert tensor.ndim == 2 V, C = tensor.shape assert V == 50257, "just being defensive here" # calculate padded vocab size by rounding up to nearest multiple Vp = ((V + multiple - 1) // multiple) * multiple # pad the tensor pad_rows = Vp - V padded = tensor if pad_rows == 0 else F.pad(tensor, (0, 0, 0, pad_rows), value=value) assert padded.shape == (Vp, C) return padded def write_model(model, filename, dtype): # everything we need to instantiate the model # 1) header is: version int, GPTConfig ints, padding to 1024 bytes assert dtype in {"float32", "bfloat16"} # float16 todo maybe later version = { "float32": 3, # 3: all tensors are fp32, padded vocab "bfloat16": 5, # 5: all tensors are bf16, padded vocab }[dtype] header = torch.zeros(256, dtype=torch.int32) header[0] = 20240326 # magic header[1] = version # checkpoint version header[2] = model.config.block_size header[3] = model.config.vocab_size header[4] = model.config.n_layer header[5] = model.config.n_head header[6] = model.config.n_embd # 2) the parameters follow the header params = {name: param.cpu() for name, param in model.named_parameters()} # pad the vocab to a multiple of 128 here at export, for efficiency in C wte = params["transformer.wte.weight"] # (V, C) wte_padded = pad_vocab(wte) # (Vp, C) params["transformer.wte.weight"] = wte_padded # (Vp, C) print(f"padded vocab size from {wte.size(0)} to {wte_padded.size(0)}") header[7] = wte_padded.size(0) # padded vocab size store in header # now write to file with open(filename, "wb") as file: file.write(header.numpy().tobytes()) # header write_tensors(params, model.config.n_layer, file, dtype) # params print(f"wrote {filename}") def write_state(model, x, y, logits, loss, filename): # the state is used for debugging. # it contains information about the input, logits, loss, and the parameter gradients # this can be used for checking the computation correctness in C header = torch.zeros(256, dtype=torch.int32) header[0] = 20240327 # magic header[1] = 2 # run state version = 2 (1 -> 2 for padded vocab changes) header[2] = x.size(0) # batch size of the batch, B header[3] = x.size(1) # temporal extent of the batch, T grads = {name: param.grad.cpu() for name, param in model.named_parameters()} # pad the vocab grads here as well, to mirror write_model wte_grad = grads["transformer.wte.weight"] # (V, C) wte_grad_padded = pad_vocab(wte_grad, value=0) # (Vp, C) # TODO later maybe pad with nan? grads["transformer.wte.weight"] = wte_grad_padded # (Vp, C) print(f"padded vocab size in reference grads from {wte_grad.size(0)} to {wte_grad_padded.size(0)}") with open(filename, "wb") as file: # header file.write(header.numpy().tobytes()) # input x file.write(x.cpu().numpy().astype("int32").tobytes()) # (B, T) # targets y file.write(y.cpu().numpy().astype("int32").tobytes()) # (B, T) # logits (result of the model forward pass)