wehub-resource-sync
1256 lines
56 KiB
Python
1256 lines
56 KiB
Python
"""
|
|
Reference code for LLaMA-3.1 training and inference.
|
|
Will save the model weights into files, to be read from C as initialization.
|
|
|
|
This code differs from GPT-2 very slightly, there are three main differences:
|
|
1) RoPE: LLaMA uses a different positional encoding scheme called Relative Positional Encoding (RoPE).
|
|
2) GQA: Grouped Query Attention (GQA) is used to reduce the number of attention heads.
|
|
3) SwiGLU: Swish-Gated Linear Unit (SwiGLU) is used as the activation function in the MLP.
|
|
|
|
References:
|
|
# 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:
|
|
TODO: add the actual commands
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import math
|
|
import glob
|
|
import inspect
|
|
from contextlib import nullcontext
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
import time
|
|
from typing import (
|
|
AbstractSet,
|
|
Collection,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Literal,
|
|
Optional,
|
|
Sequence,
|
|
Tuple,
|
|
Union,
|
|
cast,
|
|
)
|
|
|
|
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
|
|
import torch.distributed as dist
|
|
|
|
import tiktoken
|
|
from tiktoken.load import load_tiktoken_bpe
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# PyTorch nn.Module definitions for the LLaMA 3.x model
|
|
|
|
# Used in Grouped Query Attention (GQA), broadcasts the key and value tensors
|
|
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
|
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
|
bs, slen, n_kv_heads, head_dim = x.shape
|
|
if n_rep == 1:
|
|
return x
|
|
return (
|
|
x[:, :, :, None, :]
|
|
.expand(bs, slen, n_kv_heads, n_rep, head_dim)
|
|
.reshape(bs, slen, n_kv_heads * n_rep, head_dim)
|
|
)
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# RoPE related
|
|
|
|
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
|
|
ndim = x.ndim
|
|
assert 0 <= 1 < ndim
|
|
assert freqs_cis.shape == (x.shape[1], x.shape[-1])
|
|
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
|
|
return freqs_cis.view(*shape)
|
|
|
|
def apply_scaling(freqs: torch.Tensor):
|
|
# Values obtained from grid search
|
|
scale_factor = 8
|
|
low_freq_factor = 1
|
|
high_freq_factor = 4
|
|
old_context_len = 8192 # original llama3 length
|
|
|
|
low_freq_wavelen = old_context_len / low_freq_factor
|
|
high_freq_wavelen = old_context_len / high_freq_factor
|
|
new_freqs = []
|
|
for freq in freqs:
|
|
wavelen = 2 * math.pi / freq
|
|
if wavelen < high_freq_wavelen:
|
|
new_freqs.append(freq)
|
|
elif wavelen > low_freq_wavelen:
|
|
new_freqs.append(freq / scale_factor)
|
|
else:
|
|
assert low_freq_wavelen != high_freq_wavelen
|
|
smooth = (old_context_len / wavelen - low_freq_factor) / (
|
|
high_freq_factor - low_freq_factor
|
|
)
|
|
new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)
|
|
return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device)
|
|
|
|
def apply_rotary_emb(
|
|
xq: torch.Tensor,
|
|
xk: torch.Tensor,
|
|
freqs_cis: torch.Tensor,
|
|
) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
|
|
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
|
|
freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
|
|
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
|
|
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
|
|
return xq_out.type_as(xq), xk_out.type_as(xk)
|
|
|
|
def precompute_freqs_cis(
|
|
dim: int, end: int, theta: float = 10000.0, use_scaled: bool = False
|
|
):
|
|
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
|
t = torch.arange(end, device=freqs.device, dtype=torch.float32)
|
|
if use_scaled:
|
|
freqs = apply_scaling(freqs)
|
|
freqs = torch.outer(t, freqs)
|
|
freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
|
|
return freqs_cis
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# LLaMA building blocks
|
|
|
|
# LLaMA reference code explicitly implemented RMSNorm so we copy pasted it
|
|
# (https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py)
|
|
# we could also use nn.RMSNorm, it has slightly different numeric properties, but equivalent
|
|
class RMSNorm(torch.nn.Module):
|
|
def __init__(self, dim: int, eps: float = 1e-6):
|
|
super().__init__()
|
|
self.eps = eps
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
|
|
def _norm(self, x):
|
|
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
|
|
|
def forward(self, x):
|
|
output = self._norm(x.float()).type_as(x)
|
|
return output * self.weight
|
|
|
|
class CausalSelfAttention(nn.Module):
|
|
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
assert config.n_embd % config.n_head == 0
|
|
|
|
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.use_kv = config.use_kv
|
|
self.flash = config.flash
|
|
|
|
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
|
|
|
|
# static KV cache - we could alternatively allocate it outside of the model and just pass it in when needed
|
|
if self.use_kv:
|
|
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))
|
|
|
|
def forward(self, x, freqs_cis=None, start_pos=None, mask=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)
|
|
|
|
q, k = apply_rotary_emb(q, k, freqs_cis=freqs_cis) # rotate QK (rope) <-- 1. difference compared to GPT-2
|
|
|
|
if self.use_kv and not self.training and start_pos >= 0: # use kv-caching during inference
|
|
self.cache_k[:B, start_pos : start_pos + T] = k
|
|
self.cache_v[:B, start_pos : start_pos + T] = v
|
|
k = self.cache_k[:B, : start_pos + T]
|
|
v = self.cache_v[:B, : start_pos + T]
|
|
|
|
k = repeat_kv(k, self.n_rep) # GQA <-- 2. difference compared to GPT-2
|
|
v = repeat_kv(v, self.n_rep)
|
|
|
|
q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v)) # (B, NH, T, HD)
|
|
|
|
if self.flash:
|
|
# flashattention
|
|
# if T == 1 no need to mask, otherwise the function complains
|
|
# scaled_dot_product_attention expects a mask where value of True indicates that the element should take part in attention
|
|
# our mask is the opposite, so we need to invert it
|
|
y = F.scaled_dot_product_attention(q, k, v, mask == 0 if T > 1 else None)
|
|
else:
|
|
# manual implementation of attention
|
|
# this materializes the large (T,T) matrix for all the queries and keys
|
|
scores = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.hd))
|
|
if mask is not None:
|
|
scores.masked_fill_(mask, torch.finfo(scores.dtype).min)
|
|
att = F.softmax(scores.float(), dim=-1).type_as(q)
|
|
y = att @ v # (B, NH, T, T) x (B, NH, T, HD) -> (B, NH, T, HD)
|
|
y = y.transpose(1, 2).contiguous().view(B, T, C)
|
|
y = self.c_proj(y)
|
|
return y
|
|
|
|
class MLP(nn.Module):
|
|
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
hidden_dim = 4 * config.n_embd
|
|
hidden_dim = int(2 * hidden_dim / 3)
|
|
# custom dim factor multiplier
|
|
if config.ffn_dim_multiplier is not None:
|
|
hidden_dim = int(config.ffn_dim_multiplier * hidden_dim)
|
|
hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
|
|
self.c_fc = nn.Linear(config.n_embd, hidden_dim, bias=False)
|
|
self.c_fc2 = nn.Linear(config.n_embd, hidden_dim, bias=False)
|
|
self.c_proj = nn.Linear(hidden_dim, config.n_embd, bias=False)
|
|
|
|
def forward(self, x):
|
|
# SwiGLU self.c_proj(F.silu(self.c_fc2(x)) * self.c_fc(x)) <-- 3. difference compared to GPT-2
|
|
x1 = self.c_fc(x)
|
|
x2 = self.c_fc2(x)
|
|
x2 = F.silu(x2)
|
|
x = x1 * x2
|
|
x = self.c_proj(x)
|
|
return x
|
|
|
|
class Block(nn.Module):
|
|
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.ln_1 = RMSNorm(config.n_embd, config.norm_eps)
|
|
self.attn = CausalSelfAttention(config)
|
|
self.ln_2 = RMSNorm(config.n_embd, config.norm_eps)
|
|
self.mlp = MLP(config)
|
|
|
|
def forward(self, x, freqs_cis=None, start_pos=None, mask=None):
|
|
x = x + self.attn(self.ln_1(x), freqs_cis, start_pos, mask)
|
|
x = x + self.mlp(self.ln_2(x))
|
|
return x
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# The main LLaMA 3.1 model
|
|
|
|
@dataclass
|
|
class LlamaConfig:
|
|
version: str = "3.1"
|
|
block_size: int = 8192
|
|
vocab_size: int = 128256
|
|
n_layer: int = 32
|
|
n_head: int = 32
|
|
n_kv_head: int = 8
|
|
n_embd: int = 4096
|
|
ffn_dim_multiplier: float = 1.3
|
|
multiple_of: int = 1024
|
|
norm_eps: float = 1e-5
|
|
rope_theta: float = 500000.0
|
|
use_scaled_rope: bool = True
|
|
max_gen_batch_size: int = 4
|
|
use_kv: bool = True
|
|
flash: bool = False # use flashattention?
|
|
|
|
def __init__(self, **kwargs):
|
|
for k, v in kwargs.items():
|
|
if hasattr(self, k):
|
|
setattr(self, k, v)
|
|
assert self.n_kv_head <= self.n_head
|
|
assert self.n_head % self.n_kv_head == 0
|
|
assert self.n_embd % self.n_head == 0
|
|
|
|
class LLaMA(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),
|
|
h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
|
|
ln_f = RMSNorm(config.n_embd, config.norm_eps),
|
|
))
|
|
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
|
|
|
|
# init all weights, use a torch rng object to be very careful
|
|
self.init_rng = torch.Generator()
|
|
self.init_rng.manual_seed(42)
|
|
|
|
self.freqs_cis = precompute_freqs_cis(
|
|
config.n_embd // config.n_head,
|
|
config.block_size * 2,
|
|
config.rope_theta,
|
|
config.use_scaled_rope,
|
|
)
|
|
|
|
def forward(self, idx, targets=None, return_logits=True, start_pos=0):
|
|
_, t = idx.size()
|
|
assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
|
|
|
|
# forward the LLaMA model itself
|
|
x = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
|
|
freqs_cis = self.freqs_cis[start_pos:start_pos+t]
|
|
|
|
mask = torch.triu(torch.ones((t, t), device=next(self.parameters()).device, dtype=torch.bool), diagonal=1)
|
|
|
|
for i, block in enumerate(self.transformer.h):
|
|
x = block(x, freqs_cis, start_pos, mask)
|
|
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).float()
|
|
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], :]).float() # 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
|
|
|
|
@staticmethod
|
|
def adapt_llama_state_dict_keys(checkpoint, config: LlamaConfig):
|
|
# Modify key names from Meta's LLaMA to our LLaMA
|
|
# our key names are derived from GPT-2's key names
|
|
checkpoint['transformer.wte.weight'] = checkpoint.pop('tok_embeddings.weight')
|
|
|
|
for i in range(config.n_layer):
|
|
for name in ['attention_norm', 'ffn_norm']:
|
|
old_key = f'layers.{i}.{name}.weight' # e.g. layers.x.attention_norm.weight -> transformer.h.x.ln_1.weight
|
|
new_key = f'transformer.h.{i}.ln_{1 if name == "attention_norm" else 2}.weight'
|
|
checkpoint[new_key] = checkpoint.pop(old_key)
|
|
|
|
for i in range(config.n_layer):
|
|
for name in ['attention.wq', 'attention.wk', 'attention.wv']:
|
|
old_key = f'layers.{i}.{name}.weight'
|
|
new_key = f'transformer.h.{i}.attn.c_attn.weight'
|
|
if name == 'attention.wq':
|
|
checkpoint[new_key] = checkpoint.pop(old_key)
|
|
else: # merge 3 weights into transformer.h.x.attn.c_attn.weight
|
|
checkpoint[new_key] = torch.cat((checkpoint[new_key], checkpoint.pop(old_key)), dim=0)
|
|
old_key = f'layers.{i}.attention.wo.weight'
|
|
new_key = f'transformer.h.{i}.attn.c_proj.weight'
|
|
checkpoint[new_key] = checkpoint.pop(old_key)
|
|
|
|
ffn_map = {'w1': 'c_fc2', 'w2': 'c_proj', 'w3': 'c_fc'}
|
|
for i in range(config.n_layer):
|
|
for name in ['feed_forward.w1', 'feed_forward.w2', 'feed_forward.w3']:
|
|
old_key = f'layers.{i}.{name}.weight'
|
|
new_key = f'transformer.h.{i}.mlp.{ffn_map[name.split(".")[-1]]}.weight'
|
|
checkpoint[new_key] = checkpoint.pop(old_key)
|
|
|
|
checkpoint['transformer.ln_f.weight'] = checkpoint.pop('norm.weight')
|
|
checkpoint['lm_head.weight'] = checkpoint.pop('output.weight')
|
|
|
|
return checkpoint
|
|
|
|
@staticmethod
|
|
def adapt_llama_state_dict_keys_hf(checkpoint, config: LlamaConfig):
|
|
# Modify key names from HuggingFace's LLaMA to our LLaMA
|
|
# our key names are derived from GPT-2's key names
|
|
checkpoint['transformer.wte.weight'] = checkpoint.pop('model.embed_tokens.weight')
|
|
|
|
# We need to unpermute K and V because HF script permuted the original Meta-LLaMA weights
|
|
# see: https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py
|
|
def unpermute(w, n_heads, dim1, dim2):
|
|
return w.view(n_heads, 2, dim1 // n_heads // 2, dim2).transpose(1, 2).reshape(dim1, dim2)
|
|
|
|
for i in range(config.n_layer):
|
|
for name in ['input_layernorm', 'post_attention_layernorm']:
|
|
old_key = f'model.layers.{i}.{name}.weight' # e.g. layers.x.attention_norm.weight -> transformer.h.x.ln_1.weight
|
|
new_key = f'transformer.h.{i}.ln_{1 if name == "input_layernorm" else 2}.weight'
|
|
checkpoint[new_key] = checkpoint.pop(old_key)
|
|
|
|
for i in range(config.n_layer):
|
|
for name in ['self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj']:
|
|
old_key = f'model.layers.{i}.{name}.weight'
|
|
new_key = f'transformer.h.{i}.attn.c_attn.weight'
|
|
if name == 'self_attn.q_proj':
|
|
checkpoint[new_key] = unpermute(checkpoint.pop(old_key), config.n_head, config.n_embd, config.n_embd)
|
|
else: # merge 3 weights into transformer.h.x.attn.c_attn.weight
|
|
tensor = checkpoint.pop(old_key)
|
|
if name == 'self_attn.k_proj':
|
|
tensor = unpermute(tensor, config.n_kv_head, config.n_kv_head * (config.n_embd // config.n_head), config.n_embd)
|
|
checkpoint[new_key] = torch.cat((checkpoint[new_key], tensor), dim=0)
|
|
old_key = f'model.layers.{i}.self_attn.o_proj.weight'
|
|
new_key = f'transformer.h.{i}.attn.c_proj.weight'
|
|
checkpoint[new_key] = checkpoint.pop(old_key)
|
|
|
|
ffn_map = {'gate_proj': 'c_fc2', 'down_proj': 'c_proj', 'up_proj': 'c_fc'}
|
|
for i in range(config.n_layer):
|
|
for name in ['gate_proj', 'down_proj', 'up_proj']:
|
|
old_key = f'model.layers.{i}.mlp.{name}.weight'
|
|
new_key = f'transformer.h.{i}.mlp.{ffn_map[name]}.weight'
|
|
checkpoint[new_key] = checkpoint.pop(old_key)
|
|
|
|
checkpoint['transformer.ln_f.weight'] = checkpoint.pop('model.norm.weight')
|
|
|
|
return checkpoint
|
|
|
|
@classmethod
|
|
def from_pretrained_llama3_hf(cls, model_id):
|
|
"""Loads pretrained LLaMA model weights from HuggingFace"""
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
assert model_id == "meta-llama/Meta-Llama-3.1-8B", "Only the 8B-base model is supported for now"
|
|
model_args = LlamaConfig()
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(model_id)
|
|
checkpoint = LLaMA.adapt_llama_state_dict_keys_hf(model.state_dict(), model_args)
|
|
|
|
original_default_type = torch.get_default_dtype() # save the default type
|
|
torch.set_default_tensor_type(torch.cuda.BFloat16Tensor) # much faster loading
|
|
model = LLaMA(model_args)
|
|
model.load_state_dict(checkpoint, strict=False)
|
|
torch.set_default_tensor_type(torch.tensor([], dtype=original_default_type, device="cpu").type()) # restore default type
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id)
|
|
tokenizer.pad_id = 128004 # this is the pad token id for LLaMA 3.1 base, we need to set this explicitly as our generate func expects it
|
|
tokenizer.stop_tokens = [tokenizer.eos_token_id]
|
|
model.tokenizer = tokenizer
|
|
return model
|
|
|
|
@classmethod
|
|
def from_pretrained_llama3_meta(cls, ckpt_dir, tokenizer_path):
|
|
"""Loads pretrained LLaMA model weights from a checkpoint directory"""
|
|
model_args = LlamaConfig()
|
|
|
|
ckpt_path = sorted(Path(ckpt_dir).glob("*.pth"))[0]
|
|
checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True)
|
|
checkpoint = LLaMA.adapt_llama_state_dict_keys(checkpoint, model_args)
|
|
|
|
original_default_type = torch.get_default_dtype() # save the default type
|
|
torch.set_default_tensor_type(torch.cuda.BFloat16Tensor) # much faster loading
|
|
model = LLaMA(model_args)
|
|
model.load_state_dict(checkpoint, strict=False)
|
|
torch.set_default_tensor_type(torch.tensor([], dtype=original_default_type, device="cpu").type()) # restore default type
|
|
|
|
tokenizer = Tokenizer(model_path=tokenizer_path)
|
|
model.tokenizer = tokenizer
|
|
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.inference_mode()
|
|
def generate(
|
|
self,
|
|
prompt_tokens: List[List[int]],
|
|
max_gen_len: int,
|
|
temperature: float = 0.6,
|
|
top_p: float = 0.9,
|
|
echo: bool = False,
|
|
) -> Tuple[List[List[int]], Optional[List[List[float]]]]:
|
|
"""
|
|
Generate text sequences based on provided prompts using the language generation model.
|
|
|
|
Args:
|
|
prompt_tokens (List[List[int]]): List of tokenized prompts, where each prompt is represented as a list of integers.
|
|
max_gen_len (int): Maximum length of the generated text sequence.
|
|
temperature (float, optional): Temperature value for controlling randomness in sampling. Defaults to 0.6.
|
|
top_p (float, optional): Top-p probability threshold for nucleus sampling. Defaults to 0.9.
|
|
echo (bool, optional): Flag indicating whether to include prompt tokens in the generated output. Defaults to False.
|
|
|
|
Returns:
|
|
Tuple[List[List[int]], Optional[List[List[float]]]]: A tuple containing generated token sequences.
|
|
|
|
Note:
|
|
This method uses the provided prompts as a basis for generating text. It employs nucleus sampling to produce text with controlled randomness.
|
|
|
|
"""
|
|
bsz = len(prompt_tokens)
|
|
assert bsz <= self.config.max_gen_batch_size, f"Batch size {bsz} exceeds the maximum generation batch size {self.config.max_gen_batch_size}"
|
|
device = next(self.parameters()).device
|
|
|
|
min_prompt_len = min(len(t) for t in prompt_tokens)
|
|
max_prompt_len = max(len(t) for t in prompt_tokens)
|
|
assert max_prompt_len <= self.config.block_size, f"Prompt length {max_prompt_len} exceeds the maximum block size {self.config.block_size}"
|
|
total_len = min(self.config.block_size, max_gen_len + max_prompt_len)
|
|
|
|
pad_id = self.tokenizer.pad_id
|
|
tokens = torch.full((bsz, total_len), pad_id, dtype=torch.long, device=device)
|
|
for idx, t in enumerate(prompt_tokens):
|
|
tokens[idx, : len(t)] = torch.tensor(t, dtype=torch.long, device=device)
|
|
|
|
prev_pos = 0
|
|
eos_reached = torch.tensor([False] * bsz, device=device)
|
|
input_text_mask = tokens != pad_id
|
|
|
|
if min_prompt_len == total_len:
|
|
logits, _ = self.forward(tokens, start_pos=prev_pos)
|
|
|
|
stop_tokens = torch.tensor(list(self.tokenizer.stop_tokens)).to(device)
|
|
|
|
for cur_pos in range(min_prompt_len, total_len):
|
|
logits, _ = self.forward(tokens[:, prev_pos:cur_pos], start_pos=prev_pos)
|
|
if temperature > 0:
|
|
probs = torch.softmax(logits[:, -1] / temperature, dim=-1)
|
|
next_token = sample_top_p(probs, top_p)
|
|
else:
|
|
next_token = torch.argmax(logits[:, -1], dim=-1)
|
|
|
|
next_token = next_token.reshape(-1)
|
|
# only replace token if prompt has already been generated
|
|
next_token = torch.where(input_text_mask[:, cur_pos], tokens[:, cur_pos], next_token)
|
|
tokens[:, cur_pos] = next_token
|
|
eos_reached |= ~input_text_mask[:, cur_pos] & torch.isin(next_token, stop_tokens)
|
|
prev_pos = cur_pos
|
|
if all(eos_reached):
|
|
break
|
|
|
|
out_tokens = []
|
|
for i, toks in enumerate(tokens.tolist()):
|
|
# cut to max gen len
|
|
start = 0 if echo else len(prompt_tokens[i])
|
|
toks = toks[start : len(prompt_tokens[i]) + max_gen_len]
|
|
# cut to after eos tok if any
|
|
for stop_token in self.tokenizer.stop_tokens:
|
|
try:
|
|
eos_idx = toks.index(stop_token)
|
|
toks = toks[:eos_idx]
|
|
except ValueError:
|
|
pass
|
|
out_tokens.append(toks)
|
|
return out_tokens
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# sampling utils
|
|
|
|
def sample_top_p(probs, p):
|
|
"""
|
|
Perform top-p (nucleus) sampling on a probability distribution.
|
|
|
|
Args:
|
|
probs (torch.Tensor): Probability distribution tensor.
|
|
p (float): Probability threshold for top-p sampling.
|
|
|
|
Returns:
|
|
torch.Tensor: Sampled token indices.
|
|
|
|
Note:
|
|
Top-p sampling selects the smallest set of tokens whose cumulative probability mass
|
|
exceeds the threshold p. The distribution is renormalized based on the selected tokens.
|
|
"""
|
|
probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
|
|
probs_sum = torch.cumsum(probs_sort, dim=-1)
|
|
mask = probs_sum - probs_sort > p
|
|
probs_sort[mask] = 0.0
|
|
probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
|
|
next_token = torch.multinomial(probs_sort, num_samples=1)
|
|
next_token = torch.gather(probs_idx, -1, next_token)
|
|
return next_token
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Llama 3.1 Tokenizer
|
|
|
|
# The tiktoken tokenizer can handle <=400k chars without
|
|
# pyo3_runtime.PanicException.
|
|
TIKTOKEN_MAX_ENCODE_CHARS = 400_000
|
|
|
|
# https://github.com/openai/tiktoken/issues/195
|
|
# Here we iterate over subsequences and split if we exceed the limit
|
|
# of max consecutive non-whitespace or whitespace characters.
|
|
MAX_NO_WHITESPACES_CHARS = 25_000
|
|
|
|
|
|
class Tokenizer:
|
|
"""
|
|
Tokenizing and encoding/decoding text using the Tiktoken tokenizer.
|
|
"""
|
|
|
|
special_tokens: Dict[str, int]
|
|
|
|
num_reserved_special_tokens = 256
|
|
|
|
pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: E501
|
|
|
|
def __init__(self, model_path: str):
|
|
"""
|
|
Initializes the Tokenizer with a Tiktoken model.
|
|
|
|
Args:
|
|
model_path (str): The path to the Tiktoken model file.
|
|
"""
|
|
assert os.path.isfile(model_path), model_path
|
|
|
|
mergeable_ranks = load_tiktoken_bpe(model_path)
|
|
num_base_tokens = len(mergeable_ranks)
|
|
special_tokens = [
|
|
"<|begin_of_text|>",
|
|
"<|end_of_text|>",
|
|
"<|reserved_special_token_0|>",
|
|
"<|reserved_special_token_1|>",
|
|
"<|finetune_right_pad_id|>",
|
|
"<|step_id|>",
|
|
"<|start_header_id|>",
|
|
"<|end_header_id|>",
|
|
"<|eom_id|>", # end of message
|
|
"<|eot_id|>", # end of turn
|
|
"<|python_tag|>",
|
|
]
|
|
reserved_tokens = [
|
|
f"<|reserved_special_token_{2 + i}|>"
|
|
for i in range(self.num_reserved_special_tokens - len(special_tokens))
|
|
]
|
|
special_tokens = special_tokens + reserved_tokens
|
|
|
|
self.special_tokens = {
|
|
token: num_base_tokens + i for i, token in enumerate(special_tokens)
|
|
}
|
|
self.model = tiktoken.Encoding(
|
|
name=Path(model_path).name,
|
|
pat_str=self.pat_str,
|
|
mergeable_ranks=mergeable_ranks,
|
|
special_tokens=self.special_tokens,
|
|
)
|
|
|
|
self.n_words: int = num_base_tokens + len(special_tokens)
|
|
# BOS / EOS token IDs
|
|
self.bos_id: int = self.special_tokens["<|begin_of_text|>"]
|
|
self.eos_id: int = self.special_tokens["<|end_of_text|>"]
|
|
self.eot_id: int = self.special_tokens["<|eot_id|>"]
|
|
self.eom_id: int = self.special_tokens["<|eom_id|>"]
|
|
self.python_tag_id = self.special_tokens["<|python_tag|>"]
|
|
self.pad_id: int = self.special_tokens["<|finetune_right_pad_id|>"]
|
|
# hardcoded stop tokens for the base model
|
|
self.stop_tokens = [
|
|
self.special_tokens["<|begin_of_text|>"],
|
|
self.special_tokens["<|end_of_text|>"],
|
|
]
|
|
|
|
def encode(
|
|
self,
|
|
s: str,
|
|
*,
|
|
bos: bool,
|
|
eos: bool,
|
|
allowed_special: Optional[Union[Literal["all"], AbstractSet[str]]] = None,
|
|
disallowed_special: Union[Literal["all"], Collection[str]] = (),
|
|
) -> List[int]:
|
|
"""
|
|
Encodes a string into a list of token IDs.
|
|
|
|
Args:
|
|
s (str): The input string to be encoded.
|
|
bos (bool): Whether to prepend the beginning-of-sequence token.
|
|
eos (bool): Whether to append the end-of-sequence token.
|
|
allowed_tokens ("all"|set[str]): allowed special tokens in string
|
|
disallowed_tokens ("all"|set[str]): special tokens that raise an error when in string
|
|
|
|
Returns:
|
|
list[int]: A list of token IDs.
|
|
|
|
By default, setting disallowed_special=() encodes a string by ignoring
|
|
special tokens. Specifically:
|
|
- Setting `disallowed_special` to () will cause all text corresponding
|
|
to special tokens to be encoded as natural text (insteading of raising
|
|
an error).
|
|
- Setting `allowed_special` to "all" will treat all text corresponding
|
|
to special tokens to be encoded as special tokens.
|
|
"""
|
|
if allowed_special is None:
|
|
allowed_special = set()
|
|
assert type(s) is str
|
|
|
|
substrs = (
|
|
substr
|
|
for i in range(0, len(s), TIKTOKEN_MAX_ENCODE_CHARS)
|
|
for substr in self._split_whitespaces_or_nonwhitespaces(
|
|
s[i : i + TIKTOKEN_MAX_ENCODE_CHARS], MAX_NO_WHITESPACES_CHARS
|
|
)
|
|
)
|
|
t: List[int] = []
|
|
for substr in substrs:
|
|
t.extend(
|
|
self.model.encode(
|
|
substr,
|
|
allowed_special=allowed_special,
|
|
disallowed_special=disallowed_special,
|
|
)
|
|
)
|
|
if bos:
|
|
t.insert(0, self.bos_id)
|
|
if eos:
|
|
t.append(self.eos_id)
|
|
return t
|
|
|
|
def decode(self, t: Sequence[int]) -> str:
|
|
# Typecast is safe here. Tiktoken doesn't do anything list-related with the sequence.
|
|
return self.model.decode(cast(List[int], t))
|
|
|
|
@staticmethod
|
|
def _split_whitespaces_or_nonwhitespaces(
|
|
s: str, max_consecutive_slice_len: int
|
|
) -> Iterator[str]:
|
|
"""
|
|
Splits the string `s` so that each substring contains no more than `max_consecutive_slice_len`
|
|
consecutive whitespaces or consecutive non-whitespaces.
|
|
"""
|
|
current_slice_len = 0
|
|
current_slice_is_space = s[0].isspace() if len(s) > 0 else False
|
|
slice_start = 0
|
|
|
|
for i in range(len(s)):
|
|
is_now_space = s[i].isspace()
|
|
|
|
if current_slice_is_space ^ is_now_space:
|
|
current_slice_len = 1
|
|
current_slice_is_space = is_now_space
|
|
else:
|
|
current_slice_len += 1
|
|
if current_slice_len > max_consecutive_slice_len:
|
|
yield s[slice_start:i]
|
|
slice_start = i
|
|
current_slice_len = 1
|
|
yield s[slice_start:]
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# 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] != 20240801:
|
|
print("ERROR: magic number mismatch in the data .bin file!")
|
|
exit(1)
|
|
assert header[1] == 7, "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] == 20240801, "magic number mismatch in the data .bin file"
|
|
assert header[1] == 7, "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.uint32)
|
|
assert len(tokens) == ntok, "number of tokens read does not match header?"
|
|
return tokens
|
|
|
|
class DistributedShardedDataLoader:
|
|
"""
|
|
This DataLoader is both:
|
|
- distributed (works correctly in case of multiple processes in DDP)
|
|
- sharded (supports datasets that are broken up into multiple data shards)
|
|
It is not *permuted*, meaning that it itearates over the data in the order
|
|
of the dataset on disk, so the user should make sure to shuffle their examples
|
|
during the creation of their data shards for best performance.
|
|
"""
|
|
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, 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 LLaMA 3 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)
|
|
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, 3C, C)
|
|
write_fun(model_tensors[f"transformer.h.{i}.attn.c_attn.weight"], 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}.ln_2.weight"], 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, C)
|
|
write_fun(model_tensors[f"transformer.h.{i}.mlp.c_fc2.weight"], file)
|
|
for i in range(L): # (L, C, 4C)
|
|
write_fun(model_tensors[f"transformer.h.{i}.mlp.c_proj.weight"], file)
|
|
write_fun(model_tensors["transformer.ln_f.weight"], file) # (C, )
|
|
write_fun(model_tensors["lm_head.weight"], file) # (V, C)
|
|
|
|
def write_model(model, filename, dtype):
|
|
# everything we need to instantiate the model
|
|
# 1) header is: version int, LLaMAConfig ints, padding to 1024 bytes
|
|
assert dtype in {"float32", "bfloat16"}
|
|
version = {
|
|
"float32": 3, # 3: all tensors are fp32
|
|
"bfloat16": 5, # 5: all tensors are bf16
|
|
}[dtype]
|
|
header = torch.zeros(256, dtype=torch.int32)
|
|
header[0] = 20240803 # 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_kv_head
|
|
header[7] = model.config.n_embd
|
|
header[8] = model.config.ffn_dim_multiplier
|
|
header[9] = model.config.multiple_of
|
|
header[10] = model.config.norm_eps
|
|
header[11] = model.config.rope_theta
|
|
header[12] = model.config.use_scaled_rope
|
|
header[13] = model.config.max_gen_batch_size
|
|
header[14] = int(model.config.version.split('.')[0]) # major version
|
|
header[15] = int(model.config.version.split('.')[1]) # minor version
|
|
# 2) the parameters follow the header
|
|
params = {name: param.cpu() for name, param in model.named_parameters()}
|
|
# 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] = 20240803 # magic
|
|
header[1] = x.size(0) # batch size of the batch, B
|
|
header[2] = x.size(1) # temporal extent of the batch, T
|
|
grads = {name: param.grad.cpu() for name, param in model.named_parameters()}
|
|
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)
|
|
write_fp32(logits.cpu(), file)
|
|
# loss (single float, result of the cross entropy loss)
|
|
write_fp32(loss.cpu(), file)
|
|
# gradients
|
|
write_tensors(grads, model.config.n_layer, file, "float32")
|
|
print(f"wrote {filename}")
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# int main
|
|
|
|
def print0(*args, **kwargs):
|
|
# modified print that only prints from the master process
|
|
# if this is not a distributed run, it's just a print
|
|
if int(os.environ.get("RANK", 0)) == 0:
|
|
print(*args, **kwargs)
|
|
|
|
if __name__ == "__main__":
|
|
print0(f"Running pytorch {torch.version.__version__}")
|
|
|
|
# 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("--use_hf", type=int, default=1, help="use HuggingFace (default) or use Meta's model")
|
|
parser.add_argument("--ckpt_dir", type=str, default=None, help="path to llama3 model checkpoint (needed if use_hf=0)")
|
|
parser.add_argument("--tokenizer_path", type=str, default=None, help="path to llama3 tokenizer (needed if use_hf=0)")
|
|
# file system input / output
|
|
parser.add_argument("--input_bin", type=str, default="dev/data/tinyshakespeare/tiny_shakespeare_val.bin", help="input .bin to train on")
|
|
parser.add_argument("--input_val_bin", type=str, default="", help="input .bin to eval validation loss on")
|
|
parser.add_argument("--output_dir", type=str, default="", help="output directory to which to write logs and checkpoints")
|
|
parser.add_argument("--model", type=str, default="meta-llama/Meta-Llama-3.1-8B", help="chose the llama model")
|
|
# token layout for each step of the optimization
|
|
parser.add_argument("--batch_size", type=int, default=4, help="batch size, in units of #batch dimensions")
|
|
parser.add_argument("--sequence_length", type=int, default=64, help="sequence length")
|
|
parser.add_argument("--total_batch_size", type=int, default=256, help="total desired batch size, in units of #tokens")
|
|
# workload (number of steps)
|
|
parser.add_argument("--num_iterations", type=int, default=10, help="number of iterations to run")
|
|
parser.add_argument("--inference_only", type=int, default=0, help="only run inference")
|
|
# optimization
|
|
parser.add_argument("--learning_rate", type=float, default=1e-5, help="learning rate warmup iterations")
|
|
parser.add_argument("--warmup_iters", type=int, default=0, help="learning rate warmup iterations")
|
|
parser.add_argument("--learning_rate_decay_frac", type=float, default=1.0, help="learning rate warmup iterations")
|
|
parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay")
|
|
parser.add_argument("--grad_clip", type=float, default=1.0, help="maximum gradient magnitude")
|
|
# evaluation
|
|
parser.add_argument("--val_loss_every", type=int, default=0, help="every how mant steps to evaluate val loss?")
|
|
parser.add_argument("--val_max_steps", type=int, default=20, help="how many batches of val to average?")
|
|
parser.add_argument("--sample_every", type=int, default=0, help="how often to sample from the model?")
|
|
# debugging
|
|
parser.add_argument("--overfit_single_batch", type=int, default=1, help="overfit just one batch of data")
|
|
# numerics
|
|
parser.add_argument("--tensorcores", type=int, default=0, help="use tensorcores")
|
|
# memory management
|
|
parser.add_argument("--device", type=str, default="", help="by default we autodetect, or set it here")
|
|
parser.add_argument("--compile", type=int, default=0, help="torch.compile the model")
|
|
parser.add_argument("--dtype", type=str, default="bfloat16", help="float32|float16|bfloat16")
|
|
parser.add_argument("--zero_stage", type=int, default=0, help="zero redundancy optimizer stage (0/1/2/3)")
|
|
# python -> C bridge
|
|
parser.add_argument("--write_tensors", type=int, default=0, help="write tensors to disk")
|
|
args = parser.parse_args()
|
|
|
|
# args error checking and convenience variables
|
|
B, T = args.batch_size, args.sequence_length
|
|
assert 1 <= T <= 8192, "sequence length must be between 1 and 8192"
|
|
assert args.dtype in {"float32", "float16", "bfloat16"}
|
|
assert args.model in {"meta-llama/Meta-Llama-3.1-8B"} # only 8B base model supported for now
|
|
|
|
# create the logging directory if it does not exist
|
|
logfile = None
|
|
if args.output_dir:
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
logfile = os.path.join(args.output_dir, "main.log")
|
|
# create the log file "main.log" inside it, and wipe it clean
|
|
with open(logfile, "w") as f:
|
|
pass
|
|
|
|
# set up DDP (distributed data parallel). torchrun sets this env variable
|
|
ddp = int(os.environ.get('RANK', -1)) != -1 # is this a ddp run?
|
|
if ddp:
|
|
# use of DDP atm demands CUDA, we set the device appropriately according to rank
|
|
assert torch.cuda.is_available(), "for now i think we need CUDA for DDP"
|
|
init_process_group(backend='nccl')
|
|
ddp_rank = int(os.environ['RANK'])
|
|
ddp_local_rank = int(os.environ['LOCAL_RANK'])
|
|
ddp_world_size = int(os.environ['WORLD_SIZE'])
|
|
device = f'cuda:{ddp_local_rank}'
|
|
torch.cuda.set_device(device)
|
|
master_process = ddp_rank == 0 # this process will do logging, checkpointing etc.
|
|
seed_offset = 0 # each process gets the exact same seed
|
|
zero_stage = args.zero_stage
|
|
else:
|
|
ddp_rank = 0
|
|
ddp_local_rank = 0
|
|
zero_stage = 0
|
|
ddp_world_size = 1
|
|
master_process = True
|
|
seed_offset = 0
|
|
# select the device
|
|
if args.device:
|
|
# provided explicitly by the user
|
|
device = args.device
|
|
else:
|
|
# attempt to autodetect the device
|
|
device = "cpu"
|
|
if torch.cuda.is_available():
|
|
device = "cuda"
|
|
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
|
device = "mps"
|
|
device_type = 'cuda' if 'cuda' in device else 'cpu'
|
|
assert device_type in {'cuda'}, "GPU required to run LLaMA 3" # we need to load LLaMA as bf16 on CUDA
|
|
print(f"using device: {device}")
|
|
|
|
# calculate gradient accumulation from the desired total batch size and the current run configuration
|
|
tokens_per_fwdbwd = B * T * ddp_world_size
|
|
assert args.total_batch_size % tokens_per_fwdbwd == 0
|
|
grad_accum_steps = args.total_batch_size // tokens_per_fwdbwd
|
|
print0(f"total desired batch size: {args.total_batch_size}")
|
|
print0(f"=> calculated gradient accumulation steps: {grad_accum_steps}")
|
|
|
|
# set up a context manager following the desired dtype and device
|
|
ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[args.dtype]
|
|
ctx = torch.amp.autocast(device_type=device_type, dtype=ptdtype) if (device_type == "cuda") else nullcontext()
|
|
|
|
# rng / reproducibility
|
|
torch.manual_seed(42)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed(42)
|
|
|
|
# set the torch precision mode to use TensorFloat32 (TF32) for matmuls
|
|
# docs https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html
|
|
if args.tensorcores:
|
|
torch.set_float32_matmul_precision('high')
|
|
|
|
# init the model
|
|
if args.use_hf:
|
|
model = LLaMA.from_pretrained_llama3_hf(args.model)
|
|
else: # use Meta's checkpoint
|
|
assert args.ckpt_dir is not None and os.path.exists(args.ckpt_dir), f"llama3 ckpt dir {args.ckpt_dir} does not exist"
|
|
assert args.tokenizer_path is not None and os.path.exists(args.tokenizer_path), f"llama3 tokenizer path {args.tokenizer_path} does not exist"
|
|
model = LLaMA.from_pretrained_llama3_meta(args.ckpt_dir, args.tokenizer_path)
|
|
|
|
model.train()
|
|
if args.compile:
|
|
if hasattr(config, "coordinate_descent_tuning"):
|
|
config.coordinate_descent_tuning = True # suggested by @Chillee
|
|
print0("compiling the model...")
|
|
model = torch.compile(model)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# Our own version of a simple DistributedDataLoader
|
|
|
|
# load tokens
|
|
train_loader = DistributedShardedDataLoader(args.input_bin, B, T, ddp_rank, ddp_world_size)
|
|
val_loader = None
|
|
if args.input_val_bin:
|
|
val_loader = DistributedShardedDataLoader(args.input_val_bin, B, T, ddp_rank, ddp_world_size)
|
|
|
|
# -------------------------------------------------------------------------
|
|
# PyTorch -> C bridge: save some weights and state for C to load later as reference
|
|
|
|
# do one forward pass to generate ground truth for our C tests
|
|
if master_process and args.write_tensors and (not args.inference_only):
|
|
x, y = train_loader.next_batch()
|
|
x, y = x.to(device), y.to(device)
|
|
logits, loss = model(x, y)
|
|
loss.backward()
|
|
# save model params, in bfloat16
|
|
model_to_size = {"meta-llama/Meta-Llama-3.1-8B": "8B"}
|
|
model_size_str = model_to_size[args.model] # e.g. "8B"
|
|
write_model(model, os.path.join(args.output_dir, f"llama3.1_{model_size_str}_bf16.bin"), dtype="bfloat16")
|
|
# save x, y, logits, loss, and parameter gradients, for debugging C
|
|
# always store these in fp32 to have an accurate reference (?)
|
|
write_state(model, x, y, logits, loss, os.path.join(args.output_dir, f"llama3_{model_size_str}_debug_state.bin"))
|
|
# reset the train_loader for the optimization below
|
|
train_loader.reset()
|
|
|
|
# -------------------------------------------------------------------------
|
|
# main training loop
|
|
|
|
# here we wrap model into DDP container
|
|
if ddp:
|
|
model = DDP(model, device_ids=[ddp_local_rank])
|
|
raw_model = model.module if ddp else model # always contains the "raw" unwrapped model
|
|
|
|
# init the optimizer
|
|
optimizer = raw_model.configure_optimizers(weight_decay=args.weight_decay,
|
|
learning_rate=args.learning_rate, betas=(0.9, 0.95),
|
|
device_type=device, zero_stage=zero_stage)
|
|
|
|
# learning rate decay scheduler (cosine with warmup)
|
|
def get_lr(it):
|
|
min_lr = args.learning_rate * args.learning_rate_decay_frac
|
|
# 1) linear warmup for warmup_iters steps
|
|
if it < args.warmup_iters:
|
|
return args.learning_rate * (it+1) / args.warmup_iters
|
|
# 2) if it > lr_decay_iters, return min learning rate
|
|
if it > args.num_iterations:
|
|
return min_lr
|
|
# 3) in between, use cosine decay down to min learning rate
|
|
decay_ratio = (it - args.warmup_iters) / (args.num_iterations - args.warmup_iters)
|
|
assert 0 <= decay_ratio <= 1
|
|
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff starts at 1 and goes to 0
|
|
return min_lr + coeff * (args.learning_rate - min_lr)
|
|
|
|
if device == "cuda":
|
|
torch.cuda.reset_peak_memory_stats()
|
|
timings = []
|
|
norm = -1.0 # dummy value to print in inference-only mode
|
|
for step in range(args.num_iterations + 1):
|
|
t0 = time.time()
|
|
last_step = (step == args.num_iterations)
|
|
|
|
# once in a while evaluate the validation dataset
|
|
if (args.val_loss_every > 0 \
|
|
and (step % args.val_loss_every == 0 or last_step)) \
|
|
and (val_loader is not None):
|
|
model.eval()
|
|
val_loader.reset()
|
|
with torch.no_grad():
|
|
val_loss = 0.0
|
|
for _ in range(args.val_max_steps):
|
|
x, y = val_loader.next_batch()
|
|
x, y = x.to(device), y.to(device)
|
|
_, loss = model(x, y, return_logits=False)
|
|
val_loss += loss.item()
|
|
val_loss /= args.val_max_steps
|
|
# log to console and to file
|
|
print0(f"val loss {val_loss}")
|
|
if master_process and logfile is not None:
|
|
with open(logfile, "a") as f:
|
|
f.write("s:%d tel:%f\n" % (step, val_loss))
|
|
|
|
# once in a while perform model inference on the master process
|
|
if (args.sample_every > 0 \
|
|
and (step % args.sample_every == 0 or last_step)) \
|
|
and master_process:
|
|
model.eval()
|
|
prompts: List[str] = [
|
|
"Clearly, the meaning of life is",
|
|
"Simply put, the theory of relativity states that",
|
|
"""The repo llm.c on GitHub is""",
|
|
"""Translate English to French:
|
|
|
|
sea otter => loutre de mer
|
|
peppermint => menthe poivrée
|
|
plush girafe => girafe peluche
|
|
cheese =>""",
|
|
]
|
|
if args.use_hf:
|
|
prompt_tokens = [model.tokenizer(x).input_ids for x in prompts]
|
|
else: # Meta
|
|
prompt_tokens = [model.tokenizer.encode(x, bos=True, eos=False) for x in prompts]
|
|
|
|
generation_tokens = model.generate(prompt_tokens, max_gen_len=64, temperature=0.6, top_p=0.9, echo=False)
|
|
results = [{"generation": model.tokenizer.decode(t)} for t in generation_tokens]
|
|
for prompt, result in zip(prompts, results):
|
|
print(prompt, end="")
|
|
print(f"{result['generation']}")
|
|
print("\n==================================\n")
|
|
|
|
# bit confusing: we want to make sure to eval and sample on 0th iteration
|
|
# but also after the very last iteration. so we loop for step <= num_iterations
|
|
# instead of just < num_iterations (one extra due to <=), only to do
|
|
# the validation/sampling one last time, and then we break right here as we're done.
|
|
if last_step:
|
|
break
|
|
|
|
# --------------- TRAINING SECTION BEGIN -----------------
|
|
model.train()
|
|
optimizer.zero_grad(set_to_none=True)
|
|
# if we are trying to overfit a single batch, we reset the loader here
|
|
if args.overfit_single_batch:
|
|
train_loader.reset()
|
|
# micro-batch loop where we do gradient accumulation to reach desired total batch size
|
|
lossf = 0.0 # for getting the mean loss (as simple float) over the accumulation steps
|
|
for micro_step in range(grad_accum_steps):
|
|
# fetch a batch
|
|
x, y = train_loader.next_batch()
|
|
x, y = x.to(device), y.to(device)
|
|
if ddp:
|
|
# we want only the last micro-step to sync grads in a DDP model
|
|
# the official way to do this is with model.no_sync(), but that is a
|
|
# context manager that bloats the code, so we just toggle this variable
|
|
model.require_backward_grad_sync = (micro_step == grad_accum_steps - 1)
|
|
# forward pass
|
|
with ctx:
|
|
_, loss = model(x, y, return_logits=False)
|
|
# we have to scale the loss to account for gradient accumulation,
|
|
# because the gradients just add on each successive backward().
|
|
# addition of gradients corresponds to a SUM in the objective, but
|
|
# instead of a SUM we want MEAN, so we scale the loss here
|
|
loss = loss / grad_accum_steps
|
|
lossf += loss.detach() # keep track of the mean loss
|
|
# backward pass
|
|
if not args.inference_only:
|
|
loss.backward()
|
|
if ddp:
|
|
dist.all_reduce(lossf, op=dist.ReduceOp.AVG)
|
|
lossf = lossf.item()
|
|
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
|
|
# determine and set the learning rate for this iteration
|
|
lr = get_lr(step)
|
|
for param_group in optimizer.param_groups:
|
|
param_group['lr'] = lr
|
|
# step the optimizer
|
|
optimizer.step()
|
|
# --------------- TRAINING SECTION END -------------------
|
|
# everything that follows now is just diagnostics, prints, logging, etc.
|
|
|
|
# wait on the CPU for all device work to end so we get accurate per-iteration timings below
|
|
if device == "mps":
|
|
torch.mps.synchronize()
|
|
elif device == "cuda":
|
|
torch.cuda.synchronize()
|
|
# time and print
|
|
t1 = time.time()
|
|
# the 0th iteration is often an outlier (much slower) => skip logging it
|
|
tokens_per_second = grad_accum_steps * ddp_world_size * B * T / (t1-t0)
|
|
print0(f"step {step+1:4d}/{args.num_iterations} | train loss {lossf:.6f} | norm {norm:.4f} | lr {lr:.2e} | ({(t1-t0)*1000:.2f} ms | {tokens_per_second:.0f} tok/s)")
|
|
# log to logile
|
|
if master_process and logfile is not None:
|
|
with open(logfile, "a") as f:
|
|
f.write("s:%d trl:%f\n" % (step, lossf))
|
|
|
|
# keep track of smooth timings, last 20 iterations
|
|
if step > 0 and step > args.num_iterations - 20:
|
|
timings.append(t1-t0)
|
|
|
|
# print the average of the last 20 timings, to get something smooth-ish
|
|
timings = timings[-20:]
|
|
print0(f"final {len(timings)} iters avg: {np.mean(timings)*1000:.3f}ms")
|
|
print0(f"peak memory consumption: {torch.cuda.max_memory_allocated() // 1024 // 1024} MiB")
|
|
|
|
# -------------------------------------------------------------------------
|
|
# clean up nice
|
|
if ddp:
|
|
destroy_process_group()
|