train_model.py

"""
SupraLabs / Supra-Mini-v6  ~1.4M params, Llama-arch
5B-Token pretraining on FineWeb-Edu + Cosmopedia (70/30)
Target: RTX 5060 Ti 16GB, bf16
"""

import os, math, numpy as np, torch
os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

from datasets import load_dataset, interleave_datasets
from tokenizers import ByteLevelBPETokenizer
from transformers import (
    LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast,
    Trainer, TrainingArguments,
)
from torch.utils.data import Dataset
from tqdm import tqdm

# --------------------------------------------------------------- Tokenizer
print("[*] Loading tokenizer...")
fast_tok = ByteLevelBPETokenizer(
    "./custom_llama_tokenizer-vocab.json",
    "./custom_llama_tokenizer-merges.txt",
)
tokenizer = PreTrainedTokenizerFast(
    tokenizer_object=fast_tok,
    bos_token="<s>", eos_token="</s>",
    unk_token="<unk>", pad_token="<pad>",
)
assert len(tokenizer.get_vocab()) < 2**16, "vocab too large for uint16"

# --------------------------------------------------------------- Data
TOKEN_BIN     = "./tokens_mix_5B.bin"
TARGET_TOKENS = 5_000_000_000
SEQ_LEN       = 1024
N_VAL_TOKENS  = 5_000_000              # 5M val (1024-aligned)
FLUSH_EVERY   = 4_000_000
BATCH_TEXTS   = 2000

def build_mixed_token_bin(path=TOKEN_BIN, target_tokens=TARGET_TOKENS):
    if os.path.exists(path) and os.path.getsize(path) >= target_tokens * 2:
        print(f"[=] Reusing {path}")
        return
    ds_fw   = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT",
                           split="train", streaming=True)
    ds_cosm = load_dataset("HuggingFaceTB/smollm-corpus", "cosmopedia-v2",
                           split="train", streaming=True)
    mixed = interleave_datasets([ds_fw, ds_cosm], probabilities=[0.70, 0.30],
                                seed=42, stopping_strategy="all_exhausted")
    mm = np.memmap(path, dtype=np.uint16, mode="w+", shape=(target_tokens,))
    eos = tokenizer.eos_token_id
    written, buf, texts = 0, [], []
    pbar = tqdm(total=target_tokens, desc="tok", unit="tok")

    def flush():
        nonlocal written, buf
        if not buf: return False
        n = min(len(buf), target_tokens - written)
        mm[written:written+n] = np.asarray(buf[:n], dtype=np.uint16)
        written += n; pbar.update(n); del buf[:n]
        return written >= target_tokens

    for ex in mixed:
        texts.append(ex["text"])
        if len(texts) >= BATCH_TEXTS:
            for e in fast_tok.encode_batch(texts):
                buf.extend(e.ids); buf.append(eos)
            texts.clear()
            if len(buf) >= FLUSH_EVERY and flush(): break
    if written < target_tokens and texts:
        for e in fast_tok.encode_batch(texts):
            buf.extend(e.ids); buf.append(eos)
    flush()
    pbar.close(); mm.flush(); del mm

class MemmapDataset(Dataset):
    def __init__(self, path, offset_tokens, length_tokens, seq_len):
        self.path = path
        self.seq_len = seq_len
        self.offset = offset_tokens
        self.n = length_tokens // seq_len
        self._d = None
    @property
    def d(self):
        if self._d is None:
            self._d = np.memmap(self.path, dtype=np.uint16, mode="r",
                                offset=self.offset * 2,
                                shape=(self.n * self.seq_len,))
        return self._d
    def __len__(self): return self.n
    def __getitem__(self, i):
        s = i * self.seq_len
        ids = torch.from_numpy(np.asarray(self.d[s:s+self.seq_len], dtype=np.int64))
        return {"input_ids": ids, "labels": ids.clone()}

def collate(batch):
    return {"input_ids": torch.stack([b["input_ids"] for b in batch]),
            "labels":    torch.stack([b["labels"]    for b in batch])}

print(f"[*] Building token bin ({TARGET_TOKENS:,} toks)...")
build_mixed_token_bin()
N_TRAIN = TARGET_TOKENS - N_VAL_TOKENS
train_dataset = MemmapDataset(TOKEN_BIN, 0,        N_TRAIN,      SEQ_LEN)
val_dataset   = MemmapDataset(TOKEN_BIN, N_TRAIN,  N_VAL_TOKENS, SEQ_LEN)
print(f"[+] Train: {len(train_dataset):,} chunks | Val: {len(val_dataset):,} chunks")

# --------------------------------------------------------------- Model
config = LlamaConfig(
    vocab_size=len(tokenizer.get_vocab()),
    hidden_size=128, intermediate_size=256,
    num_hidden_layers=6,
    num_attention_heads=4, num_key_value_heads=2,
    max_position_embeddings=SEQ_LEN,
    rope_theta=10000.0, rms_norm_eps=1e-5,
    tie_word_embeddings=True, initializer_range=0.02,
    attention_bias=False, mlp_bias=False,
    pad_token_id=tokenizer.pad_token_id,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    attn_implementation="sdpa",
)
model = LlamaForCausalLM(config)

n_layers = config.num_hidden_layers
for name, p in model.named_parameters():
    if name.endswith("o_proj.weight") or name.endswith("down_proj.weight"):
        torch.nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * n_layers))

print(f"[*] Params: {model.num_parameters():,}")
print(f"[*] Start loss (ln vocab): {math.log(len(tokenizer.get_vocab())):.3f}")

# --------------------------------------------------------------- WSD scheduler
def get_wsd_lambda(num_training_steps, warmup_steps=500, decay_frac=0.20):
    dec_start = int(num_training_steps * (1.0 - decay_frac))
    def fn(step):
        if step < warmup_steps:
            return step / max(1, warmup_steps)
        if step < dec_start:
            return 1.0
        prog = (step - dec_start) / max(1, num_training_steps - dec_start)
        return 1.0 - math.sqrt(prog)
    return fn

class WSDTrainer(Trainer):
    def create_optimizer(self):
        decay, nodecay, embed = [], [], []
        for n, p in self.model.named_parameters():
            if not p.requires_grad: continue
            if "embed_tokens" in n: embed.append(p)
            elif p.ndim >= 2: decay.append(p)
            else: nodecay.append(p)
        self.optimizer = torch.optim.AdamW(
            [{"params": decay,   "weight_decay": 0.1},
             {"params": nodecay, "weight_decay": 0.0},
             {"params": embed,   "weight_decay": 0.0,
              "lr": self.args.learning_rate * 0.3}],
            lr=self.args.learning_rate, betas=(0.9, 0.95), eps=1e-8,
            fused=True)
        return self.optimizer

    def compute_loss(self, model, inputs, return_outputs=False,
                     num_items_in_batch=None):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs.logits[..., :-1, :].contiguous()
        labels = labels[..., 1:].contiguous()
        flat_logits = logits.view(-1, logits.size(-1))
        flat_labels = labels.view(-1)

        if num_items_in_batch is not None:
            ce = torch.nn.functional.cross_entropy(
                flat_logits, flat_labels, ignore_index=-100, reduction="sum")
            lse = torch.logsumexp(flat_logits, dim=-1)
            z_loss = 1e-4 * (lse ** 2).sum()
            loss = (ce + z_loss) / num_items_in_batch
        else:
            ce = torch.nn.functional.cross_entropy(
                flat_logits, flat_labels, ignore_index=-100)
            lse = torch.logsumexp(flat_logits, dim=-1)
            z_loss = 1e-4 * (lse ** 2).mean()
            loss = ce + z_loss
        return (loss, outputs) if return_outputs else loss

    def create_scheduler(self, num_training_steps, optimizer=None):
        opt = optimizer or self.optimizer
        self.lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
            opt, get_wsd_lambda(num_training_steps, warmup_steps=500))
        return self.lr_scheduler

# --------------------------------------------------------------- Training
args = TrainingArguments(
    output_dir="./Supra-Mini-v6-5B",
    num_train_epochs=1,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=8,            # global batch = 256 seqs × 1024 = 262k tok
    learning_rate=6e-4,
    weight_decay=0.1,
    adam_beta1=0.9, adam_beta2=0.95, adam_epsilon=1e-8,
    max_grad_norm=1.0,
    bf16=True, fp16=False, tf32=True,
    torch_compile=True,
    logging_steps=50,
    eval_strategy="steps", eval_steps=1000,
    per_device_eval_batch_size=64,
    save_strategy="steps", save_steps=5000, save_total_limit=3,
    dataloader_num_workers=8,
    dataloader_pin_memory=True,
    dataloader_persistent_workers=True,
    report_to="none",
    gradient_checkpointing=False,
    seed=42,
)

trainer = WSDTrainer(
    model=model, args=args,
    train_dataset=train_dataset, eval_dataset=val_dataset,
    data_collator=collate,
)
trainer.train()
trainer.save_model("./Supra-Mini-v6-5B-FINAL")
tokenizer.save_pretrained("./Supra-Mini-v6-5B-FINAL")
print("[*] done.")