Upload modeling_sage_1b.py with huggingface_hub

modeling_sage_1b.py

@@ -0,0 +1,340 @@
+import json, os, pickle, math, time, sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, processors
+from datasets import load_dataset
+import numpy as np
+import warnings
+warnings.filterwarnings("ignore")
+
+os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+DATA_DIR = os.path.join(BASE_DIR, "data")
+TOKENIZER_DIR = os.path.join(BASE_DIR, "tokenizer")
+MODEL_DIR = os.path.join(BASE_DIR, "model")
+os.makedirs(DATA_DIR, exist_ok=True)
+os.makedirs(TOKENIZER_DIR, exist_ok=True)
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+torch.set_num_threads(4)
+
+VOCAB_SIZE = 50000
+HIDDEN_SIZE = 1536
+NUM_LAYERS = 30
+NUM_HEADS = 12
+HEAD_DIM = HIDDEN_SIZE // NUM_HEADS
+INTERMEDIATE_SIZE = 6144
+MAX_SEQ_LEN = 128
+NUM_SAMPLES = 10000
+TRAIN_BATCH_SIZE = 2
+GRAD_ACCUM_STEPS = 4
+LEARNING_RATE = 4e-4
+NUM_EPOCHS = 3
+WARMUP_STEPS = 50
+
+total_p = (VOCAB_SIZE * HIDDEN_SIZE +
+           NUM_LAYERS * (4 * HIDDEN_SIZE * HIDDEN_SIZE + 3 * HIDDEN_SIZE * INTERMEDIATE_SIZE + 2 * HIDDEN_SIZE) +
+           HIDDEN_SIZE * VOCAB_SIZE)
+print(f"=== Sage 1B ({total_p/1e9:.3f}B params) ===")
+
+# ====== STEP 1: Load English Dataset ======
+print("\n--- Step 1: Loading English text dataset ---")
+dataset = load_dataset("roneneldan/TinyStories", split="train", streaming=True)
+samples = []
+start = time.time()
+for i, example in enumerate(dataset):
+    if i >= NUM_SAMPLES:
+        break
+    text = example.get("text", "").strip()
+    if len(text) >= 100:
+        samples.append(text)
+    if (i+1) % 10000 == 0:
+        print(f"  {i+1}/{NUM_SAMPLES} scanned, {len(samples)} valid ({time.time()-start:.0f}s)")
+
+# Supplement with more if needed
+if len(samples) < 10000:
+    print(f"  Only {len(samples)} valid samples. Trying additional sources...")
+    try:
+        ds2 = load_dataset("wikipedia", "20220301.en", split="train", streaming=True)
+        for i, ex in enumerate(ds2):
+            if len(samples) >= NUM_SAMPLES:
+                break
+            text = ex.get("text", "").strip()
+            if len(text) >= 200:
+                samples.append(text[:2000])
+            if (i+1) % 5000 == 0:
+                print(f"  wiki: {i+1} scanned, {len(samples)} total")
+    except Exception as e:
+        print(f"  Wikipedia supplement failed: {e}")
+
+print(f"Collected {len(samples)} samples in {time.time()-start:.0f}s")
+with open(os.path.join(DATA_DIR, "raw_texts.pkl"), "wb") as f:
+    pickle.dump(samples, f)
+
+# ====== STEP 2: Train BPE Tokenizer ======
+print("\n--- Step 2: Training BPE tokenizer ---")
+tokenizer = Tokenizer(models.BPE())
+tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+tokenizer.decoder = decoders.ByteLevel()
+tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
+trainer = trainers.BpeTrainer(
+    vocab_size=VOCAB_SIZE,
+    special_tokens=["<PAD>", "<UNK>", "<BOS>", "<EOS>"],
+    min_frequency=2,
+)
+tokenizer.train_from_iterator(samples, trainer=trainer)
+tokenizer.save(os.path.join(TOKENIZER_DIR, "tokenizer.json"))
+print(f"Vocabulary size: {tokenizer.get_vocab_size()}")
+
+# ====== STEP 3: Tokenize ======
+print("\n--- Step 3: Tokenizing ---")
+pad_id = tokenizer.token_to_id("<PAD>")
+bos_id = tokenizer.token_to_id("<BOS>")
+eos_id = tokenizer.token_to_id("<EOS>")
+tokenized = []
+for text in samples:
+    ids = tokenizer.encode(text).ids
+    if len(ids) > MAX_SEQ_LEN - 2:
+        ids = ids[:MAX_SEQ_LEN - 2]
+    ids = [bos_id] + ids + [eos_id]
+    if len(ids) < MAX_SEQ_LEN:
+        ids += [pad_id] * (MAX_SEQ_LEN - len(ids))
+    tokenized.append(ids)
+tensor_data = torch.tensor(tokenized, dtype=torch.long)
+torch.save(tensor_data, os.path.join(DATA_DIR, "tokenized.pt"))
+print(f"Tokenized {len(tokenized)} sequences, shape: {tensor_data.shape}")
+
+# ====== STEP 4: Build Model ======
+print("\n--- Step 4: Building Sage 1B model ---")
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len=MAX_SEQ_LEN):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        self.max_seq_len = max_seq_len
+        self._cos = None
+        self._sin = None
+    def get_cos_sin(self, x, seq_len=None):
+        seq_len = seq_len or x.size(1)
+        if self._cos is None or self._cos.size(-2) < seq_len:
+            t = torch.arange(self.max_seq_len, device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1)[None, None]
+            self._cos = emb.cos()
+            self._sin = emb.sin()
+        return self._cos[..., :seq_len, :], self._sin[..., :seq_len, :]
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary(x, cos, sin):
+    return (x * cos) + (rotate_half(x) * sin)
+
+class Attention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
+        self.k_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
+        self.v_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
+        self.o_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False)
+    def forward(self, x, cos, sin, mask):
+        B, T, _ = x.shape
+        q = self.q_proj(x).reshape(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        k = self.k_proj(x).reshape(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        v = self.v_proj(x).reshape(B, T, NUM_HEADS, HEAD_DIM).transpose(1, 2)
+        q, k = apply_rotary(q, cos, sin), apply_rotary(k, cos, sin)
+        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(HEAD_DIM)
+        attn = attn + mask[:, :, :T, :T]
+        attn = F.softmax(attn, dim=-1)
+        return self.o_proj(attn.matmul(v).transpose(1, 2).reshape(B, T, HIDDEN_SIZE))
+
+class FeedForward(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gate = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
+        self.up = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
+        self.down = nn.Linear(INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False)
+    def forward(self, x):
+        return self.down(F.silu(self.gate(x)) * self.up(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn_norm = nn.RMSNorm(HIDDEN_SIZE, eps=1e-6)
+        self.ffn_norm = nn.RMSNorm(HIDDEN_SIZE, eps=1e-6)
+        self.attn = Attention()
+        self.ffn = FeedForward()
+    def forward(self, x, cos, sin, mask):
+        x = x + self.attn(self.attn_norm(x), cos, sin, mask)
+        x = x + self.ffn(self.ffn_norm(x))
+        return x
+
+mask_cache = {}
+def get_causal_mask(T, device):
+    if T not in mask_cache:
+        m = torch.triu(torch.full((T, T), float('-inf'), device=device), diagonal=1)
+        mask_cache[T] = m
+    return mask_cache[T][None, None]
+
+class Sage1B(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(VOCAB_SIZE, HIDDEN_SIZE)
+        self.layers = nn.ModuleList([TransformerBlock() for _ in range(NUM_LAYERS)])
+        self.norm = nn.RMSNorm(HIDDEN_SIZE, eps=1e-6)
+        self.lm_head = nn.Linear(HIDDEN_SIZE, VOCAB_SIZE, bias=False)
+        self.rotary = RotaryEmbedding(HEAD_DIM)
+        self.max_seq_len = MAX_SEQ_LEN
+        self.vocab_size = VOCAB_SIZE
+        self.hidden_size = HIDDEN_SIZE
+
+    def forward(self, input_ids, labels=None):
+        B, T = input_ids.shape
+        x = self.embed_tokens(input_ids) * math.sqrt(HIDDEN_SIZE)
+        cos, sin = self.rotary.get_cos_sin(x, T)
+        mask = get_causal_mask(T, x.device)
+        for layer in self.layers:
+            x = layer(x, cos, sin, mask)
+        x = self.norm(x)
+        logits = self.lm_head(x)
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.view(-1, VOCAB_SIZE), labels.view(-1), ignore_index=0)
+        return loss, logits
+
+    @torch.no_grad()
+    def generate(self, input_ids, max_new_tokens=50, temperature=0.8, top_k=40):
+        self.eval()
+        for _ in range(max_new_tokens):
+            if input_ids.size(1) > MAX_SEQ_LEN:
+                input_ids = input_ids[:, -MAX_SEQ_LEN:]
+            _, logits = self.forward(input_ids)
+            logits = logits[:, -1, :] / temperature
+            if top_k > 0:
+                vals = torch.topk(logits, top_k).values[:, -1:]
+                logits[logits < vals] = float('-inf')
+            probs = F.softmax(logits, dim=-1)
+            nxt = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, nxt], dim=1)
+            if nxt.item() == 3:
+                break
+        return input_ids
+
+model = Sage1B()
+total_params = sum(p.numel() for p in model.parameters())
+print(f"Parameters: {total_params:,} ({total_params/1e9:.3f}B)")
+
+config = {
+    "vocab_size": VOCAB_SIZE, "hidden_size": HIDDEN_SIZE,
+    "num_hidden_layers": NUM_LAYERS, "num_attention_heads": NUM_HEADS,
+    "head_dim": HEAD_DIM, "intermediate_size": INTERMEDIATE_SIZE,
+    "max_position_embeddings": MAX_SEQ_LEN, "model_type": "sage_1b",
+    "total_params": total_params, "torch_dtype": "float32",
+}
+with open(os.path.join(MODEL_DIR, "config.json"), "w") as f:
+    json.dump(config, f, indent=2)
+
+# Copy this file as modeling_sage_1b.py for HF distribution
+with open(os.path.join(MODEL_DIR, "modeling_sage_1b.py"), "w") as f:
+    f.write(open(os.path.abspath(__file__)).read())
+
+# ====== STEP 5: Train ======
+print("\n--- Step 5: Training ---")
+data = torch.load(os.path.join(DATA_DIR, "tokenized.pt"))
+print(f"Training samples: {len(data)}")
+
+class TextDataset(Dataset):
+    def __init__(self, data):
+        self.data = data
+    def __len__(self):
+        return len(self.data)
+    def __getitem__(self, idx):
+        t = self.data[idx]
+        return t[:-1], t[1:]
+
+tds = TextDataset(data)
+loader = DataLoader(tds, batch_size=TRAIN_BATCH_SIZE, shuffle=True, drop_last=True)
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, betas=(0.9, 0.95), weight_decay=0.1)
+
+def get_lr(step):
+    if step < WARMUP_STEPS:
+        return LEARNING_RATE * (step + 1) / WARMUP_STEPS
+    return LEARNING_RATE * (1 - min(step, 10000) / 10000 * 0.9)
+
+best_loss = float('inf')
+global_step = 0
+
+for epoch in range(NUM_EPOCHS):
+    model.train()
+    total_loss = 0
+    n_batches = 0
+    optimizer.zero_grad()
+    epoch_start = time.time()
+
+    for bidx, (inp, tgt) in enumerate(loader):
+        loss, _ = model(inp, labels=tgt)
+        loss = loss / GRAD_ACCUM_STEPS
+        loss.backward()
+
+        if (bidx + 1) % GRAD_ACCUM_STEPS == 0:
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            for pg in optimizer.param_groups:
+                pg['lr'] = get_lr(global_step)
+            optimizer.step()
+            optimizer.zero_grad()
+            global_step += 1
+
+        total_loss += loss.item() * GRAD_ACCUM_STEPS
+        n_batches += 1
+
+        if (bidx + 1) % 200 == 0:
+            elapsed = time.time() - epoch_start
+            avg = total_loss / max(n_batches, 1)
+            lr = optimizer.param_groups[0]['lr']
+            print(f"  E{epoch+1} B{bidx+1}/{len(loader)} | Loss: {avg:.4f} | LR: {lr:.2e} | {elapsed:.0f}s")
+
+    avg = total_loss / max(n_batches, 1)
+    et = time.time() - epoch_start
+    print(f"Epoch {epoch+1} | Avg loss: {avg:.4f} | Time: {et:.0f}s | Steps: {global_step}")
+
+    if avg < best_loss:
+        best_loss = avg
+        sd = model.state_dict()
+        torch.save(sd, os.path.join(MODEL_DIR, "pytorch_model.bin"))
+        torch.save({k: v.half() if v.dtype == torch.float32 else v for k, v in sd.items()},
+                   os.path.join(MODEL_DIR, "pytorch_model_state.bin"))
+        print(f"  Best model saved (loss: {avg:.4f})")
+
+# Final save
+sd = model.state_dict()
+torch.save(sd, os.path.join(MODEL_DIR, "pytorch_model.bin"))
+torch.save({k: v.half() if v.dtype == torch.float32 else v for k, v in sd.items()},
+           os.path.join(MODEL_DIR, "pytorch_model_state.bin"))
+
+# Save tokenizer pickle
+with open(os.path.join(TOKENIZER_DIR, "tokenizer.pkl"), "wb") as f:
+    pickle.dump(tokenizer, f)
+
+# Test generation
+print("\n--- Test generation ---")
+model.eval()
+from tokenizers import Tokenizer as Tk
+test_tokenizer = Tk.from_file(os.path.join(TOKENIZER_DIR, "tokenizer.json"))
+prompt = "Once upon a time"
+tokens = test_tokenizer.encode(prompt).ids
+inp = torch.tensor([[1] + tokens[:20]], dtype=torch.long)
+out = model.generate(inp, max_new_tokens=30, temperature=0.7)
+gen_text = test_tokenizer.decode(out[0].tolist(), skip_special_tokens=True)
+print(f"Prompt: {prompt}")
+print(f"Generated: {gen_text}")
+
+print(f"\n=== DONE ===")
+print(f"Params: {total_params:,} ({total_params/1e9:.3f}B)")
+print(f"Best loss: {best_loss:.4f}")
+print(f"Model: {MODEL_DIR}")