dim: int = 4096

n_layers: int = 32

n_heads: int = 32

n_kv_heads: Optional[int] = None

vocab_size: int = -1 # defined later by tokenizer

multiple_of: int = 256 # make SwiGLU hidden layer size multiple of large power of 2

norm_eps: float = 1e-5

max_seq_len: int = 2048

def __init__(self, dim: int, eps: float):

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)

freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

t = torch.arange(end, device=freqs.device) # type: ignore

freqs = torch.outer(t, freqs).float() # type: ignore

freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64

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)]

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)

"""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)

def __init__(self, args: ModelArgs):

super().__init__()

self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads

model_parallel_size = 1

self.n_local_heads = args.n_heads // model_parallel_size

self.n_local_kv_heads = self.n_kv_heads // model_parallel_size

self.n_rep = self.n_local_heads // self.n_local_kv_heads

self.head_dim = args.dim // args.n_heads

self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)

self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)

self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)

self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)

self.attn_dropout = nn.Dropout(args.dropout)

self.resid_dropout = nn.Dropout(args.dropout)

self.dropout = args.dropout

# use flash attention or a manual implementation?

self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')

if not self.flash:

print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")

mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))

mask = torch.triu(mask, diagonal=1)

self.register_buffer("mask", mask)

def forward(

self,

x: torch.Tensor,

freqs_cis: torch.Tensor,

):

bsz, seqlen, _ = x.shape

# QKV

xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)

xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

# RoPE relative positional embeddings

xq, xk = apply_rotary_emb(xq, xk, freqs_cis)

# grouped multiquery attention: expand out keys and values

xk = repeat_kv(xk, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)

xv = repeat_kv(xv, self.n_rep) # (bs, seqlen, n_local_heads, head_dim)

# make heads into a batch dimension

xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim)

xk = xk.transpose(1, 2)

xv = xv.transpose(1, 2)

# flash implementation

if self.flash:

output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)

else:

# manual implementation

scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)

scores = scores + self.mask[:, :, :seqlen, :seqlen] # (bs, n_local_heads, seqlen, cache_len + seqlen)

scores = F.softmax(scores.float(), dim=-1).type_as(xq)

scores = self.attn_dropout(scores)

output = torch.matmul(scores, xv) # (bs, n_local_heads, seqlen, head_dim)

# restore time as batch dimension and concat heads

output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)

# final projection into the residual stream

output = self.wo(output)

output = self.resid_dropout(output)

def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):

super().__init__()

hidden_dim = int(2 * hidden_dim / 3)

hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

self.w1 = nn.Linear(dim, hidden_dim, bias=False)

self.w2 = nn.Linear(hidden_dim, dim, bias=False)

self.w3 = nn.Linear(dim, hidden_dim, bias=False)

self.dropout = nn.Dropout(dropout)

def forward(self, x):

def __init__(self, layer_id: int, args: ModelArgs):

super().__init__()

self.n_heads = args.n_heads

self.dim = args.dim

self.head_dim = args.dim // args.n_heads

self.attention = Attention(args)

self.feed_forward = FeedForward(

dim=args.dim,

hidden_dim=4 * args.dim,

multiple_of=args.multiple_of,

dropout=args.dropout,

)

self.layer_id = layer_id

self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)

self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)

def forward(self, x, freqs_cis):

h = x + self.attention.forward(self.attention_norm(x), freqs_cis)

out = h + self.feed_forward.forward(self.ffn_norm(h))

def __init__(self, params: ModelArgs):

super().__init__()

self.params = params

self.vocab_size = params.vocab_size

self.n_layers = params.n_layers

self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)

self.dropout = nn.Dropout(params.dropout)

self.layers = torch.nn.ModuleList()

for layer_id in range(params.n_layers):

self.layers.append(TransformerBlock(layer_id, params))

self.norm = RMSNorm(params.dim, eps=params.norm_eps)

self.output = nn.Linear(params.dim, params.vocab_size, bias=False)

# share the unembedding parameters with the embedding parameters

self.tok_embeddings.weight = self.output.weight # https://paperswithcode.com/method/weight-tying

# some useful precompute for the RoPE relative positional embeddings. TODO why * 2 here? confuse

freqs_cis = precompute_freqs_cis(self.params.dim // self.params.n_heads, self.params.max_seq_len * 2)

self.register_buffer("freqs_cis", freqs_cis, persistent=False)

# init all weights

self.apply(self._init_weights)

# apply special scaled init to the residual projections, per GPT-2 paper

for pn, p in self.named_parameters():

if pn.endswith('w3.weight') or pn.endswith('wo.weight'):

torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * params.n_layers))

def _init_weights(self, module):

if isinstance(module, nn.Linear):

torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)

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)

def forward(self, tokens, targets=None):

_bsz, seqlen = tokens.shape

h = self.tok_embeddings(tokens)

h = self.dropout(h)

freqs_cis = self.freqs_cis[:seqlen]

for layer in self.layers:

h = layer(h, freqs_cis)

h = self.norm(h)

if targets is not None:

# if we are given some desired targets also calculate the loss

logits = self.output(h)

loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)

else:

# inference-time mini-optimization: only forward the output on the very last position

logits = self.output(h[:, [-1], :]) # note: using list [-1] to preserve the time dim

loss = None

return logits, loss

def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):

# 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)

print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters")

print(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'

extra_args = dict(fused=True) if use_fused else dict()

optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args)

print(f"using fused AdamW: {use_fused}")

return optimizer

def estimate_mfu(self, fwdbwd_per_iter, dt):

""" estimate model flops utilization (MFU) in units of A100 bfloat16 peak FLOPS """

# first estimate the number of flops we do per iteration.

# see PaLM paper Appendix B as ref: https://arxiv.org/abs/2204.02311

N = sum(p.numel() for p in self.parameters())

cfg = self.params

L, H, Q, T = cfg.n_layers, cfg.n_heads, cfg.dim//cfg.n_heads, cfg.max_seq_len

flops_per_token = 6*N + 12*L*H*Q*T

flops_per_fwdbwd = flops_per_token * T

flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter

# express our flops throughput as ratio of A100 bfloat16 peak flops

flops_achieved = flops_per_iter * (1.0/dt) # per second

flops_promised = 312e12 # A100 GPU bfloat16 peak flops is 312 TFLOPS

mfu = flops_achieved / flops_promised

return mfu

@torch.inference_mode()

def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):

"""

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.params.max_seq_len else idx[:, -self.params.max_seq_len:]

# forward the model to get the logits for the index in the sequence

logits, _ = self(idx_cond)

logits = logits[:, -1, :] # crop to just the final time step

if temperature == 0.0:

# "sample" the single most likely index

_, idx_next = torch.topk(logits, k=1, dim=-1)

else:

# pluck the logits at the final step and scale by desired temperature

logits = logits / 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)

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

def export(self, filepath='model.bin'):

"""export the model weights in fp32 into .bin file to be read from C"""

f = open(filepath, 'wb')

def serialize(t):

d = t.detach().cpu().view(-1).numpy().astype(np.float32)

b = struct.pack(f'{len(d)}f', *d)

f.write(b)

# first write out the header

hidden_dim = self.layers[0].feed_forward.w1.weight.shape[0]

p = self.params

n_kv_heads = p.n_heads if p.n_kv_heads is None else p.n_kv_heads

header = struct.pack('iiiiiii', p.dim, hidden_dim, p.n_layers, p.n_heads,

n_kv_heads, p.vocab_size, p.max_seq_len)

f.write(header)

# next write out the embedding weights

serialize(self.tok_embeddings.weight)

# now all the layers

# attention weights

for layer in self.layers:

serialize(layer.attention_norm.weight)

for layer in self.layers:

serialize(layer.attention.wq.weight)

for layer in self.layers:

serialize(layer.attention.wk.weight)

for layer in self.layers:

serialize(layer.attention.wv.weight)

for layer in self.layers:

serialize(layer.attention.wo.weight)

# ffn weights

for layer in self.layers:

serialize(layer.ffn_norm.weight)

for layer in self.layers:

serialize(layer.feed_forward.w1.weight)

for layer in self.layers:

serialize(layer.feed_forward.w2.weight)

for layer in self.layers:

serialize(layer.feed_forward.w3.weight)

# final rmsnorm

serialize(self.norm.weight)

# note: no need to write final classifier weights due to weight sharing

# freqs_cis

serialize(self.freqs_cis.real[:p.max_seq_len])

serialize(self.freqs_cis.imag[:p.max_seq_len])

# write to binary file

f.close()