Create train_model.py

train_model.py

@@ -0,0 +1,178 @@
+import os
+os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+print("[*] Loading libraries...")
+import torch
+import math
+import numpy as np
+from datasets import load_dataset
+from tokenizers import ByteLevelBPETokenizer
+from transformers import (
+    LlamaConfig,
+    LlamaForCausalLM,
+    PreTrainedTokenizerFast,
+    Trainer,
+    TrainingArguments,
+)
+from torch.utils.data import Dataset
+from tqdm import tqdm
+
+print("[*] Loading tokenizer...")
+fast_tokenizer = ByteLevelBPETokenizer(
+    "./custom_llama_tokenizer-vocab.json",
+    "./custom_llama_tokenizer-merges.txt"
+)
+tokenizer = PreTrainedTokenizerFast(
+    tokenizer_object=fast_tokenizer,
+    bos_token="<s>",
+    eos_token="</s>",
+    unk_token="<unk>",
+    pad_token="<pad>",
+)
+
+TOKEN_BIN     = "/kaggle/working/tokens.bin"
+TARGET_TOKENS = 500_000_000
+SEQ_LEN       = 256
+BATCH_TEXTS   = 1000
+FLUSH_EVERY   = 1_000_000
+
+
+def build_token_bin(fast_tokenizer, path=TOKEN_BIN, target_tokens=TARGET_TOKENS):
+    if os.path.exists(path) and os.path.getsize(path) >= target_tokens * 2:
+        print(f"[=] Reusing existing token file: {path}")
+        return
+
+    print(f"[*] Streaming + tokenizing {target_tokens:,} tokens → {path}")
+    mm = np.memmap(path, dtype=np.uint16, mode="w+", shape=(target_tokens,))
+
+    dataset = load_dataset(
+        "HuggingFaceFW/fineweb-edu", "sample-10BT",
+        split="train", streaming=True
+    )
+
+    written = 0
+    buf = []
+    texts = []
+    pbar = tqdm(total=target_tokens, desc="[*] Gathering tokens", unit="tok")
+
+    def flush_buf():
+        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 example in dataset:
+        texts.append(example["text"])
+        if len(texts) >= BATCH_TEXTS:
+            encs = fast_tokenizer.encode_batch(texts)
+            texts.clear()
+            for e in encs:
+                buf.extend(e.ids)
+            if len(buf) >= FLUSH_EVERY:
+                if flush_buf():
+                    break
+
+    if written < target_tokens and texts:
+        encs = fast_tokenizer.encode_batch(texts)
+        for e in encs:
+            buf.extend(e.ids)
+    if written < target_tokens:
+        flush_buf()
+
+    pbar.close()
+    mm.flush()
+    del mm
+    print(f"[+] Wrote {written:,} tokens to {path} "
+          f"({os.path.getsize(path)/1e6:.1f} MB)")
+
+
+class MemmapDataset(Dataset):
+    def __init__(self, path, total_tokens, seq_len=SEQ_LEN):
+        self.path     = path
+        self.seq_len  = seq_len
+        self.n_chunks = total_tokens // seq_len
+        self._data    = None  # lazy open (Multiprocessing-safe)
+
+    @property
+    def data(self):
+        if self._data is None:
+            self._data = np.memmap(
+                self.path, dtype=np.uint16, mode="r",
+                shape=(self.n_chunks * self.seq_len,)
+            )
+        return self._data
+
+    def __len__(self):
+        return self.n_chunks
+
+    def __getitem__(self, idx):
+        s   = idx * self.seq_len
+        arr = np.asarray(self.data[s:s + self.seq_len], dtype=np.int64)
+        ids = torch.from_numpy(arr)
+        return {"input_ids": ids, "labels": ids.clone()}
+
+
+def collate_fn(batch):
+    input_ids = torch.stack([b["input_ids"] for b in batch])
+    labels    = torch.stack([b["labels"]    for b in batch])
+    return {"input_ids": input_ids, "labels": labels}
+
+
+print(f"[*] Preparing {TARGET_TOKENS:,} tokens (streaming, memmap-backed)...")
+build_token_bin(fast_tokenizer, TOKEN_BIN, TARGET_TOKENS)
+dataset = MemmapDataset(TOKEN_BIN, TARGET_TOKENS, seq_len=SEQ_LEN)
+print(f"[+] Dataset ready: {len(dataset):,} chunks of {SEQ_LEN} tokens")
+
+print("[*] Setting up model...")
+config = LlamaConfig(
+    vocab_size=len(tokenizer.get_vocab()),
+    hidden_size=48,
+    intermediate_size=96,
+    num_hidden_layers=4,
+    num_attention_heads=4,
+    max_position_embeddings=256,
+    pad_token_id=tokenizer.pad_token_id,
+    bos_token_id=tokenizer.bos_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+)
+model = LlamaForCausalLM(config)
+print(f"[*] Model parameters: {model.num_parameters():,}")
+
+print("[*] Defining training arguments...")
+training_args = TrainingArguments(
+    output_dir="./Supra-Mini-0.1m",
+    num_train_epochs=2,
+    per_device_train_batch_size=1024,
+    gradient_accumulation_steps=1,
+    save_steps=500,
+    save_total_limit=2,
+    logging_steps=100,
+    weight_decay=0.01,
+    fp16=torch.cuda.is_available(),
+    push_to_hub=False,
+    report_to="none",
+    dataloader_num_workers=2,
+    dataloader_pin_memory=True,
+    learning_rate=6e-4,
+    lr_scheduler_type="cosine",
+    warmup_ratio=0.05,
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+    data_collator=collate_fn,
+)
+
+print("[*] Starting training...")
+trainer.train()
+trainer.save_model("./Supra-Mini-0.1m-FINAL")
+tokenizer.save_pretrained("./Supra-Mini-0.1m-FINAL")
+print("[*] Training finished.")