| """Tiny Llama pretrain demo (CPU/4GB friendly). For 1B+ use cloud + DeepSpeed.""" |
| import json |
| from pathlib import Path |
| import torch |
| from transformers import LlamaConfig, LlamaForCausalLM, Trainer, TrainingArguments |
| from torch.utils.data import Dataset |
|
|
| class JsonlDS(Dataset): |
| def __init__(self, path, tok_len=128): |
| self.rows = [json.loads(l) for l in open(path, encoding="utf-8")] |
| self.tok_len = tok_len |
| def __len__(self): return len(self.rows) |
| def __getitem__(self, i): |
| import random |
| r = self.rows[i] |
| |
| s = (r.get("prompt", "") + "\n" + r.get("think", "") + "\n" + r.get("answer", ""))[:2000] |
| ids = [min(ord(c), 30000) for c in s][:self.tok_len] |
| ids += [0] * (self.tok_len - len(ids)) |
| return {"input_ids": torch.tensor(ids, dtype=torch.long), "labels": torch.tensor(ids, dtype=torch.long)} |
|
|
| def main(): |
| import argparse |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--epochs", type=int, default=3) |
| ap.add_argument("--out", default="outputs/tiny-50M") |
| ap.add_argument("--cpu", action="store_true", help="force CPU") |
| a = ap.parse_args() |
| cfg = LlamaConfig(hidden_size=512, num_hidden_layers=8, num_attention_heads=8, |
| num_key_value_heads=4, intermediate_size=1376, vocab_size=32000, max_position_embeddings=2048) |
| model = LlamaForCausalLM(cfg) |
| ds = JsonlDS("data/processed/train.jsonl") |
| use_cuda = torch.cuda.is_available() and not a.cpu |
| import os |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") |
| args = TrainingArguments(output_dir=a.out, per_device_train_batch_size=1, |
| gradient_accumulation_steps=32, |
| num_train_epochs=a.epochs, logging_steps=5, save_steps=200, |
| save_total_limit=1, fp16=use_cuda, use_cpu=a.cpu, |
| gradient_checkpointing=True, report_to="none", |
| optim="adafactor", |
| dataloader_num_workers=0, dataloader_pin_memory=False, |
| max_grad_norm=0.0) |
| Trainer(model=model, args=args, train_dataset=ds).train() |
| model.save_pretrained(a.out) |
| print(f"saved -> {a.out}") |
|
|
| if __name__ == "__main__": |
| main() |
|
|
