Capstone Module 11: Mini Llama-3 100M Param Pre-training — Single H100, 1 Week

# Mini Llama-3 pre-training script
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, IterableDataset
import math
import os
import time
from dataclasses import dataclass
 
# (Llama block imports from Module 10 capstone)
# from llama_block import LlamaBlock, LlamaConfig, RMSNorm
 
@dataclass
class MiniLlamaConfig:
    d_model: int = 768       # 100M param: smaller dim
    n_layers: int = 12
    n_heads: int = 12
    n_kv_heads: int = 4      # GQA
    d_head: int = 64
    d_ff: int = 2048
    rope_base: float = 500000.0
    eps: float = 1e-6
    vocab_size: int = 32000   # TurkTokenizer-tr 32K
    max_seq_len: int = 2048
 
 
class MiniLlama(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        # self.blocks = nn.ModuleList([LlamaBlock(config) for _ in range(config.n_layers)])
        self.final_norm = RMSNorm(config.d_model, config.eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        # Tied embedding
        self.lm_head.weight = self.embedding.weight
    
    def forward(self, input_ids):
        x = self.embedding(input_ids)
        # for block in self.blocks:
        #     x = block(x)
        x = self.final_norm(x)
        return self.lm_head(x)
 
 
class StreamingDataset(IterableDataset):
    def __init__(self, corpus_path, tokenizer, max_seq_len):
        self.corpus_path = corpus_path
        self.tokenizer = tokenizer
        self.max_seq_len = max_seq_len
    
    def __iter__(self):
        buffer = []
        with open(self.corpus_path, 'r', encoding='utf-8') as f:
            for line in f:
                tokens = self.tokenizer.encode(line.strip())
                buffer.extend(tokens)
                buffer.append(self.tokenizer.eos_token_id)
                while len(buffer) >= self.max_seq_len:
                    chunk = buffer[:self.max_seq_len]
                    buffer = buffer[self.max_seq_len:]
                    yield torch.tensor(chunk, dtype=torch.long)
 
 
def train(config, corpus_path, save_dir, total_steps=100000):
    model = MiniLlama(config).cuda().bfloat16()
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")
    
    optimizer = optim.AdamW(
        model.parameters(),
        lr=3e-4,
        betas=(0.9, 0.95),
        eps=1e-5,
        weight_decay=0.1,
    )
    
    # LR scheduler
    warmup = 1000
    def lr_lambda(step):
        if step < warmup:
            return step / warmup
        progress = (step - warmup) / (total_steps - warmup)
        return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress))
    scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
    
    # Dataset
    # tokenizer = load_turktokenizer()
    # dataset = StreamingDataset(corpus_path, tokenizer, config.max_seq_len)
    # loader = DataLoader(dataset, batch_size=32, num_workers=4)
    
    # Training loop
    model.train()
    start = time.time()
    # for step, input_ids in enumerate(loader):
    for step in range(total_steps):
        # input_ids = input_ids.cuda()
        input_ids = torch.randint(0, config.vocab_size, (32, config.max_seq_len), device='cuda')
        labels = input_ids.clone()
        
        logits = model(input_ids[:, :-1])
        loss = nn.functional.cross_entropy(
            logits.reshape(-1, config.vocab_size),
            labels[:, 1:].reshape(-1),
        )
        
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        scheduler.step()
        
        if step % 100 == 0:
            elapsed = time.time() - start
            print(f"Step {step}/{total_steps}: loss={loss.item():.4f}, "
                  f"lr={scheduler.get_last_lr()[0]:.2e}, "
                  f"tokens/sec={(step + 1) * 32 * config.max_seq_len / elapsed:.0f}")
        
        if step % 5000 == 0 and step > 0:
            torch.save(model.state_dict(), f"{save_dir}/checkpoint-{step}.pt")
    
    torch.save(model.state_dict(), f"{save_dir}/final.pt")
    print(f"Total time: {(time.time() - start) / 3600:.2f} hours")
 
 
# Run
config = MiniLlamaConfig()
train(config, "corpus/turkish.txt", "checkpoints/")