From ccdbdfd4ceaaad811873da9e93bd24e2161465d5 Mon Sep 17 00:00:00 2001 From: Aleksa Gordic Date: Sat, 3 Aug 2024 21:49:45 +0200 Subject: [PATCH] Going towards pure llama 3 file - fixed attn --- train_gpt2.py | 41 +++++++++++++++-------------------------- 1 file changed, 15 insertions(+), 26 deletions(-) diff --git a/train_gpt2.py b/train_gpt2.py index 454b162..99a6356 100644 --- a/train_gpt2.py +++ b/train_gpt2.py @@ -1,19 +1,18 @@ """ -Reference code for GPT-2 training and inference. +Reference code for LLaMA-3.1 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 +# 1) https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/tokenizer.py +# 2) https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py +# 3) https://github.com/meta-llama/llama3/blob/11817d47e1ba7a4959b025eb1ca308572e0e3963/llama/generation.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 +python train_llama3.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 +torchrun --standalone --nproc_per_node=4 train_llama3.py --write_tensors=0 --num_iterations=50 --sequence_length=1024 --compile=1 --tensorcores=1 --dtype=bfloat16 """ import os @@ -48,11 +47,6 @@ from llmc_py.utils import repeat_kv, sample_top_p, RMSNorm # ----------------------------------------------------------------------------- # 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 @@ -61,21 +55,17 @@ class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 - self.is_llama = config.is_llama - if not self.is_llama: - assert config.n_head == config.n_kv_head, "GQA is only available for LLaMA" self.n_head = config.n_head self.n_kv_head = config.n_kv_head self.n_rep = self.n_head // self.n_kv_head self.hd = config.n_embd // config.n_head - self.c_attn = nn.Linear(config.n_embd, (config.n_head + 2 * config.n_kv_head) * self.hd, bias=not self.is_llama) # key, query, value projections - self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=not self.is_llama) # output projection - - if not self.is_llama: - self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 + self.c_attn = nn.Linear(config.n_embd, (config.n_head + 2 * config.n_kv_head) * self.hd, bias=False) # key, query, value projections + self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False) # output projection + self.c_proj.LLMC_RESIDUAL_SCALE_FLAG = 1 + # static KV cache self.cache_k = torch.zeros((config.max_gen_batch_size, config.block_size, config.n_kv_head, self.hd)) self.cache_v = torch.zeros((config.max_gen_batch_size, config.block_size, config.n_kv_head, self.hd)) @@ -85,12 +75,12 @@ class CausalSelfAttention(nn.Module): def forward(self, x, freqs_cis=None, start_pos=None): 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_head * self.hd, self.n_kv_head * self.hd, self.n_kv_head * self.hd], dim=-1) q, k, v = map(lambda t: t.view(B, T, -1, self.hd), (q, k, v)) # (B, T, NH, HD) - if self.is_llama: - q, k = apply_rotary_emb(q, k, freqs_cis=freqs_cis) # rotate QK (rope) + q, k = apply_rotary_emb(q, k, freqs_cis=freqs_cis) # rotate QK (rope) if start_pos >= 0: # kv-caching (which we can disable by setting start_pos = -1) self.cache_k[:B, start_pos : start_pos + T] = k @@ -98,9 +88,8 @@ class CausalSelfAttention(nn.Module): k = self.cache_k[:B, : start_pos + T] v = self.cache_v[:B, : start_pos + T] - if self.is_llama: - k = repeat_kv(k, self.n_rep) # GQA - v = repeat_kv(v, self.n_rep) + k = repeat_kv(k, self.n_rep) # GQA + v = repeat_kv(v, self.n_rep) q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v)) # (B, NH, T, HD) @@ -768,7 +757,7 @@ if __name__ == "__main__": # default settings will overfit a tiny batch of data # and save model weights and debug state to disk on the first iteration parser = argparse.ArgumentParser() - parser.add_argument("--llama3", type=int, default=0, help="use llama3 model") + parser.add_argument("--llama3", type=int, default=1, help="use llama3 model") parser.add_argument("--llama3_ckpt_dir", type=str, default=None, help="path to llama3 model checkpoint") parser.add_argument("--llama3_tokenizer_path", type=str, default=None, help="path to llama3 tokenizer") # file system input / output