File size: 2,396 Bytes
d5743c4 87f48ad d5743c4 b363e2c 87f48ad b363e2c d5743c4 87f48ad b363e2c 87f48ad d5743c4 b363e2c d5743c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """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]
# char-level pseudo-tokenization for demo (real run uses train_tokenizer.py + proper encode)
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()
|