| """ |
| SCUGNIZZ - Hugging Face Jobs edition |
| |
| NOTE: |
| - Configurato per training su Hugging Face Jobs. |
| - Usa FineWeb in streaming. |
| - Parametri modello aumentati (12L / 768D / 12H). |
| - TARGET_TOKENS rappresenta un obiettivo logico di training. |
| - Per usare l'intero FineWeb è consigliabile eliminare il memmap e |
| passare a un DataLoader streaming. Questa versione mantiene la |
| struttura originale per ridurre le modifiche. |
| """ |
|
|
| !pip -q install datasets transformers huggingface_hub |
|
|
| import os |
| import math |
| import time |
| import random |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from contextlib import nullcontext |
| from datasets import load_dataset |
| from transformers import GPT2TokenizerFast |
|
|
| seed = 42 |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
|
|
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(seed) |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| print("Device:", device) |
|
|
| if torch.cuda.is_available(): |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
|
|
| |
| |
| |
|
|
| TARGET_TOKENS = 10_000_000_000 |
|
|
| DATA_FILE = "fineweb_full_uint32.dat" |
| CKPT_FILE = "pcs_fineweb_checkpoint_last.pt" |
| BEST_FILE = "pcs_fineweb_checkpoint_best.pt" |
| FINAL_FILE = "pcs_fineweb_final.pt" |
|
|
| batch_size = 16 |
| block_size = 1024 |
|
|
| |
| TOKENS_PER_STEP = batch_size * block_size |
| TRAIN_TOKENS = int(TARGET_TOKENS * 0.9) |
| max_iters = (TRAIN_TOKENS + TOKENS_PER_STEP - 1) // TOKENS_PER_STEP |
|
|
| print(f"Training tokens : {TRAIN_TOKENS:,}") |
| print(f"Token/step : {TOKENS_PER_STEP:,}") |
| print(f"Iterazioni : {max_iters:,} (~1 epoca)") |
| eval_interval = 500 |
| save_interval = 1000 |
| eval_iters = 20 |
|
|
| learning_rate = 3e-4 |
| min_lr = 3e-5 |
| warmup_iters = 1000 |
| weight_decay = 0.1 |
| beta1 = 0.9 |
| beta2 = 0.95 |
| grad_clip = 1.0 |
|
|
| n_embd = 768 |
| n_head = 12 |
| n_layer = 12 |
| dropout = 0.1 |
| bias = False |
|
|
| pcs_a = 0.8309193524478643 |
| pcs_b = 0.0 |
|
|
| tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") |
| tokenizer.pad_token = tokenizer.eos_token |
| vocab_size = tokenizer.vocab_size |
| eos_id = tokenizer.eos_token_id |
|
|
| print("Vocab size:", vocab_size) |
|
|
| if not os.path.exists(DATA_FILE): |
| print("Creating FineWeb memmap...") |
|
|
| arr = np.memmap( |
| DATA_FILE, |
| dtype=np.uint32, |
| mode="w+", |
| shape=(TARGET_TOKENS,) |
| ) |
|
|
| ds = load_dataset( |
| "HuggingFaceFW/fineweb", |
| name="CC-MAIN-2024-10", |
| split="train", |
| streaming=True |
| ) |
|
|
| pos = 0 |
| last_print_million = -1 |
|
|
| for row in ds: |
| txt = row["text"] |
|
|
| if txt and len(txt) > 100: |
| ids = tokenizer.encode(txt + tokenizer.eos_token) |
| n = min(len(ids), TARGET_TOKENS - pos) |
|
|
| if n > 0: |
| arr[pos:pos+n] = np.array(ids[:n], dtype=np.uint32) |
| pos += n |
|
|
| cur_million = pos // 1_000_000 |
| if cur_million != last_print_million: |
| print("Saved tokens:", pos) |
| last_print_million = cur_million |
|
|
| if pos >= TARGET_TOKENS: |
| break |
|
|
| arr.flush() |
| print("Memmap created. Tokens written:", pos) |
|
|
| else: |
| print("Memmap already exists:", DATA_FILE) |
|
|
| data = np.memmap( |
| DATA_FILE, |
| dtype=np.uint32, |
| mode="r", |
| shape=(TARGET_TOKENS,) |
| ) |
|
|
| split_idx = int(0.9 * TARGET_TOKENS) |
| train_len = split_idx |
| val_len = TARGET_TOKENS - split_idx |
|
|
| print("Train tokens:", train_len) |
| print("Val tokens:", val_len) |
|
|
| def get_batch(split_name): |
| if split_name == "train": |
| lo = 0 |
| hi = train_len - block_size - 1 |
| else: |
| lo = train_len |
| hi = TARGET_TOKENS - block_size - 1 |
|
|
| ix = np.random.randint(lo, hi, size=(batch_size,)) |
| x = np.stack([data[i:i+block_size] for i in ix]) |
| y = np.stack([data[i+1:i+block_size+1] for i in ix]) |
|
|
| x = torch.tensor(x, dtype=torch.long, device=device) |
| y = torch.tensor(y, dtype=torch.long, device=device) |
| return x, y |
|
|
| def get_lr(it): |
| if it < warmup_iters: |
| return learning_rate * (it + 1) / warmup_iters |
|
|
| if it > max_iters: |
| return min_lr |
|
|
| decay_ratio = (it - warmup_iters) / (max_iters - warmup_iters) |
| coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) |
| return min_lr + coeff * (learning_rate - min_lr) |
|
|
| class PCS(nn.Module): |
| def __init__(self, a=pcs_a, b=pcs_b): |
| super().__init__() |
| self.a = a |
| self.b = b |
|
|
| def forward(self, x): |
| return x * torch.sin(self.a * x) + self.b * torch.cos(x) |
|
|
| class CausalSelfAttention(nn.Module): |
| def __init__(self, n_embd, n_head, block_size, dropout, bias=False): |
| super().__init__() |
| assert n_embd % n_head == 0 |
|
|
| self.n_head = n_head |
| self.head_dim = n_embd // n_head |
|
|
| self.q_proj = nn.Linear(n_embd, n_embd, bias=bias) |
| self.k_proj = nn.Linear(n_embd, n_embd, bias=bias) |
| self.v_proj = nn.Linear(n_embd, n_embd, bias=bias) |
| self.out_proj = nn.Linear(n_embd, n_embd, bias=bias) |
|
|
| self.attn_dropout = nn.Dropout(dropout) |
| self.resid_dropout = nn.Dropout(dropout) |
|
|
| mask = torch.tril(torch.ones(block_size, block_size)) |
| self.register_buffer("mask", mask.view(1, 1, block_size, block_size)) |
|
|
| def forward(self, x): |
| B, T, C = x.shape |
|
|
| q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| k = self.k_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| v = self.v_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
|
|
| att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| att = att.masked_fill(self.mask[:, :, :T, :T] == 0, float("-inf")) |
| att = F.softmax(att, dim=-1) |
| att = self.attn_dropout(att) |
|
|
| y = att @ v |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| y = self.resid_dropout(self.out_proj(y)) |
| return y |
|
|
| class MLP(nn.Module): |
| def __init__(self, n_embd, dropout, bias=False): |
| super().__init__() |
| self.fc1 = nn.Linear(n_embd, 4 * n_embd, bias=bias) |
| self.act = PCS() |
| self.fc2 = nn.Linear(4 * n_embd, n_embd, bias=bias) |
| self.dropout = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| x = self.fc1(x) |
| x = self.act(x) |
| x = self.fc2(x) |
| x = self.dropout(x) |
| return x |
|
|
| class Block(nn.Module): |
| def __init__(self, n_embd, n_head, block_size, dropout, bias=False): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(n_embd) |
| self.attn = CausalSelfAttention(n_embd, n_head, block_size, dropout, bias=bias) |
| self.ln2 = nn.LayerNorm(n_embd) |
| self.mlp = MLP(n_embd, dropout, bias=bias) |
|
|
| def forward(self, x): |
| x = x + self.attn(self.ln1(x)) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
| class GPT(nn.Module): |
| def __init__(self): |
| super().__init__() |
|
|
| self.tok_emb = nn.Embedding(vocab_size, n_embd) |
| self.pos_emb = nn.Embedding(block_size, n_embd) |
| self.drop = nn.Dropout(dropout) |
|
|
| self.blocks = nn.ModuleList([ |
| Block(n_embd, n_head, block_size, dropout, bias=bias) |
| for _ in range(n_layer) |
| ]) |
|
|
| self.ln_f = nn.LayerNorm(n_embd) |
| self.lm_head = nn.Linear(n_embd, vocab_size, bias=False) |
|
|
| self.tok_emb.weight = self.lm_head.weight |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| if module.bias is not None: |
| nn.init.zeros_(module.bias) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
| def forward(self, idx, targets=None): |
| B, T = idx.shape |
| assert T <= block_size |
|
|
| pos = torch.arange(0, T, device=idx.device, dtype=torch.long) |
| x = self.tok_emb(idx) + self.pos_emb(pos) |
| x = self.drop(x) |
|
|
| for block in self.blocks: |
| x = block(x) |
|
|
| x = self.ln_f(x) |
| logits = self.lm_head(x) |
|
|
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy( |
| logits.reshape(B * T, logits.size(-1)), |
| targets.reshape(B * T) |
| ) |
|
|
| return logits, loss |
|
|
| model = GPT().to(device) |
| print("Parameters (M):", sum(p.numel() for p in model.parameters()) / 1e6) |
|
|
| optimizer = torch.optim.AdamW( |
| model.parameters(), |
| lr=learning_rate, |
| betas=(beta1, beta2), |
| weight_decay=weight_decay |
| ) |
|
|
| use_amp = (device == "cuda") |
| scaler = torch.cuda.amp.GradScaler(enabled=use_amp) |
|
|
| start_iter = 0 |
| best_val = float("inf") |
|
|
| if os.path.exists(CKPT_FILE): |
| print("Loading checkpoint...") |
| ckpt = torch.load(CKPT_FILE, map_location=device) |
| model.load_state_dict(ckpt["model"]) |
| optimizer.load_state_dict(ckpt["optimizer"]) |
| start_iter = ckpt["iter"] + 1 |
| best_val = ckpt.get("best_val", float("inf")) |
| print("Resume from iter:", start_iter) |
| print("Best val:", best_val) |
|
|
| @torch.no_grad() |
| def estimate_loss(): |
| out = {} |
| model.eval() |
|
|
| for split in ["train", "val"]: |
| losses = torch.zeros(eval_iters) |
|
|
| for _ in range(eval_iters): |
| x, y = get_batch(split) |
| ctx = torch.cuda.amp.autocast() if use_amp else nullcontext() |
|
|
| with ctx: |
| _, loss = model(x, y) |
|
|
| losses[_] = loss.item() |
|
|
| out[split] = losses.mean().item() |
|
|
| model.train() |
| return out |
|
|
| print("Starting training...") |
| t0 = time.time() |
|
|
| for it in range(start_iter, max_iters + 1): |
| lr = get_lr(it) |
| for param_group in optimizer.param_groups: |
| param_group["lr"] = lr |
|
|
| xb, yb = get_batch("train") |
|
|
| ctx = torch.cuda.amp.autocast() if use_amp else nullcontext() |
| with ctx: |
| _, loss = model(xb, yb) |
|
|
| optimizer.zero_grad(set_to_none=True) |
|
|
| scaler.scale(loss).backward() |
| scaler.unscale_(optimizer) |
| torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) |
| scaler.step(optimizer) |
| scaler.update() |
|
|
| if it % eval_interval == 0: |
| losses = estimate_loss() |
| train_loss = losses["train"] |
| val_loss = losses["val"] |
| ppl = math.exp(val_loss) |
|
|
| print("Iter", f"{it:06d}", "|", |
| "LR", f"{lr:.6e}", "|", |
| "Train", f"{train_loss:.4f}", "|", |
| "Val", f"{val_loss:.4f}", "|", |
| "PPL", f"{ppl:.2f}") |
|
|
| if val_loss < best_val: |
| best_val = val_loss |
| torch.save( |
| { |
| "iter": it, |
| "model": model.state_dict(), |
| "optimizer": optimizer.state_dict(), |
| "best_val": best_val |
| }, |
| BEST_FILE |
| ) |
| print("New best checkpoint saved:", BEST_FILE) |
|
|
| if it % save_interval == 0 and it > 0: |
| torch.save( |
| { |
| "iter": it, |
| "model": model.state_dict(), |
| "optimizer": optimizer.state_dict(), |
| "best_val": best_val |
| }, |
| CKPT_FILE |
| ) |
| print("Checkpoint saved:", CKPT_FILE) |
|
|
| elapsed = (time.time() - t0) / 60 |
| print("Training finished in", round(elapsed, 2), "minutes") |
|
|
| torch.save(model.state_dict(), FINAL_FILE) |
| print("Final model saved:", FINAL_FILE) |
|
|
| @torch.no_grad() |
| def generate(prompt, max_new_tokens=150, temperature=0.8, top_k=50, top_p=0.95, repetition_penalty=1.10): |
| model.eval() |
|
|
| ids = tokenizer.encode(prompt) |
| x = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0) |
|
|
| for _ in range(max_new_tokens): |
| x_cond = x[:, -block_size:] |
|
|
| ctx = torch.cuda.amp.autocast() if use_amp else nullcontext() |
| with ctx: |
| logits, _ = model(x_cond) |
|
|
| logits = logits[:, -1, :] |
|
|
| if repetition_penalty != 1.0: |
| used_tokens = torch.unique(x[0]) |
| for token_id in used_tokens: |
| token_id = token_id.item() |
| if logits[0, token_id] < 0: |
| logits[0, token_id] *= repetition_penalty |
| else: |
| logits[0, token_id] /= repetition_penalty |
|
|
| logits = logits / temperature |
|
|
| if top_k is not None and top_k > 0: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = float("-inf") |
|
|
| if top_p is not None and top_p < 1.0: |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| sorted_probs = F.softmax(sorted_logits, dim=-1) |
| cumulative_probs = torch.cumsum(sorted_probs, dim=-1) |
|
|
| sorted_indices_to_remove = cumulative_probs > top_p |
| sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone() |
| sorted_indices_to_remove[:, 0] = False |
|
|
| indices_to_remove = sorted_indices_to_remove.scatter( |
| 1, sorted_indices, sorted_indices_to_remove |
| ) |
| logits = logits.masked_fill(indices_to_remove, float("-inf")) |
|
|
| probs = F.softmax(logits, dim=-1) |
| next_id = torch.multinomial(probs, num_samples=1) |
| x = torch.cat((x, next_id), dim=1) |
|
|
| if next_id.item() == eos_id: |
| break |
|
|
| return tokenizer.decode(x[0].tolist()) |
|
|
| print("============================================================") |
| print("Quick generation test") |
| print("============================================================") |
| print(generate("Artificial intelligence is", max_new_tokens=120)) |
|
|
| print("============================================================") |
| print("PCS GPT - Chat") |
| print("Scrivi exit per uscire.") |
| print("============================================================") |
|
|
| while True: |
| user_in = input("Tu: ").strip() |
|
|
| if user_in.lower() == "exit": |
| break |
|
|
| out = generate( |
| prompt=user_in, |
| max_new_tokens=120, |
| temperature=0.8, |
| top_k=40, |
| top_p=0.9, |
| repetition_penalty=1.12 |
| ) |
|
|
| print("PCS:") |
| print(out) |
