Upload 2 files

code/model.py

@@ -0,0 +1,94 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+        
+    def forward(self, x):
+        return (x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) * self.w).type_as(x)
+
+class ResonanceLayerKaggle(nn.Module):
+    def __init__(self, d_model=1024, n_heads=16):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        self.to_amp = nn.Linear(d_model, d_model, bias=False)
+        self.to_freq = nn.Linear(d_model, n_heads, bias=False)
+        self.to_phase = nn.Linear(d_model, n_heads, bias=False)
+        self.to_v = nn.Linear(d_model, d_model, bias=False)
+        self.to_out = nn.Linear(d_model, d_model, bias=False)
+        self.norm = RMSNorm(d_model)
+
+    def forward(self, x):
+        B, T, C = x.shape
+        h = self.norm(x)
+        A = self.to_amp(h).view(B, T, self.n_heads, self.head_dim)
+        V = self.to_v(h).view(B, T, self.n_heads, self.head_dim)
+        freq = torch.tanh(self.to_freq(h)) * 2.0  # B,T,H
+        phase = torch.tanh(self.to_phase(h)) * math.pi
+
+        # Distance - causal ke liye
+        pos = torch.arange(T, device=x.device).float()
+        chunk = 512  # safe for Kaggle T4
+        out = torch.zeros(B, T, C, device=x.device, dtype=x.dtype)
+        
+        for i in range(0, T, chunk):
+            end = min(i+chunk, T)
+            phase_i = phase[:, i:end, :].permute(0,2,1).unsqueeze(-1) # B,H,chunk,1
+            phase_j = phase.permute(0,2,1).unsqueeze(-2) # B,H,1,T
+            
+            dist = (torch.arange(i, end, device=x.device)[:,None] - pos[None,:]).abs().clamp(min=0)
+            dist = dist.view(1, 1, end-i, T)
+            
+            freq_i = freq[:, i:end, :].permute(0,2,1).unsqueeze(-1)
+            angle = (phase_i - phase_j) + freq_i * dist * 0.05
+            score = torch.cos(angle) / math.sqrt(self.head_dim)  # Normalize to prevent explosion
+            
+            causal_mask = (torch.arange(i, end, device=x.device)[:,None] >= pos[None,:]).float().view(1,1,end-i,T)
+            score = score * causal_mask
+            
+            amp = A.norm(dim=-1).permute(0,2,1) # B,H,T
+            amp_i = amp[:,:,i:end].unsqueeze(-1)
+            amp_j = amp.unsqueeze(-2)
+            score = score * torch.sigmoid(amp_i * amp_j)
+            
+            V_t = V.permute(0,2,1,3) # B,H,T,D
+            Bh = B * self.n_heads
+            s = score.reshape(Bh, end-i, T)
+            v = V_t.reshape(Bh, T, self.head_dim)
+            o = torch.bmm(s, v).view(B, self.n_heads, end-i, self.head_dim).permute(0,2,1,3).reshape(B, end-i, C)
+            out[:, i:end] = o
+
+        return self.to_out(out)
+
+class ViuResonance100M(nn.Module):
+    def __init__(self, vocab_size=64000, d_model=1024, n_layers=12, n_heads=16, dropout=0.1):
+        super().__init__()
+        self.emb = nn.Embedding(vocab_size, d_model)
+        self.layers = nn.ModuleList([ResonanceLayerKaggle(d_model, n_heads) for _ in range(n_layers)])
+        self.norm = RMSNorm(d_model)
+        self.head = nn.Linear(d_model, vocab_size, bias=False)
+        self.head.weight = self.emb.weight
+        self.resid_dropout = nn.Dropout(dropout)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        x = self.emb(idx)
+        for layer in self.layers:
+            x = x + self.resid_dropout(layer(x))
+        logits = self.head(self.norm(x))
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-100)
+        return logits, loss

code/train.py

@@ -0,0 +1,390 @@
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+
+import json, math, time, random, sys, contextlib
+import numpy as np
+import torch
+
+torch.set_float32_matmul_precision('high')
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True
+
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from huggingface_hub import HfApi, hf_hub_download, create_repo, login
+from datetime import timedelta
+
+def setup_ddp():
+    dist.init_process_group(backend="nccl", timeout=timedelta(minutes=30))
+    rank = dist.get_rank()
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = dist.get_world_size()
+    torch.cuda.set_device(local_rank)
+    return rank, local_rank, world_size
+
+rank, local_rank, world_size = setup_ddp()
+DEVICE = f"cuda:{local_rank}"
+IS_MAIN = rank == 0
+
+def log(msg):
+    if IS_MAIN:
+        print(msg, flush=True)
+
+def get_hf_token():
+    token = os.environ.get("HF_TOKEN")
+    if token:
+        return token
+    try:
+        from kaggle_secrets import UserSecretsClient
+        token = UserSecretsClient().get_secret("HF_TOKEN")
+        if token:
+            return token
+    except Exception:
+        pass
+    raise ValueError("HF_TOKEN environment variable or Kaggle Secret is not set!")
+
+HF_TOKEN = get_hf_token()
+login(token=HF_TOKEN)
+api = HfApi(token=HF_TOKEN)
+
+from transformers import AutoTokenizer
+
+DATA_REPO = "ViuAI/viuai-500m-data"
+MODEL_REPO = "ViuAI/ViuRec"
+TOKENIZER_SUBFOLDER = "tokenizer"
+CODE_SUBFOLDER = "code"
+CKPT_SUBFOLDER = "checkpoints"
+CKPT_REPO = MODEL_REPO
+
+def load_tokenizer_robust():
+    paths_to_try = [
+        "/kaggle/working/ViuRec/tokenizer",
+        "/kaggle/input/viuai-500m-tokenizer",
+        "./viuai-500m-tokenizer"
+    ]
+    for p in paths_to_try:
+        if os.path.exists(os.path.join(p, "vocab.json")):
+            try:
+                tok = AutoTokenizer.from_pretrained(p)
+                if IS_MAIN:
+                    print(f"✅ Tokenizer loaded from {p}")
+                return tok
+            except Exception:
+                pass
+    try:
+        tok = AutoTokenizer.from_pretrained(MODEL_REPO, subfolder=TOKENIZER_SUBFOLDER, token=HF_TOKEN)
+        if IS_MAIN:
+            print(f"✅ Tokenizer loaded from HF ({MODEL_REPO}/{TOKENIZER_SUBFOLDER})")
+        return tok
+    except Exception as e:
+        raise RuntimeError(f"Could not load tokenizer: {e}")
+
+tokenizer = load_tokenizer_robust()
+VOCAB_SIZE = len(tokenizer)
+
+if IS_MAIN:
+    create_repo(CKPT_REPO, repo_type="model", token=HF_TOKEN, exist_ok=True, private=True)
+dist.barrier()
+
+if os.path.exists("/kaggle/working"):
+    WORK = "/kaggle/working"
+elif os.path.exists("/root/working"):
+    WORK = "/root/working"
+else:
+    WORK = "./working"
+os.makedirs(WORK, exist_ok=True)
+
+sys.path.insert(0, WORK)
+sys.path.insert(0, os.path.join(WORK, CODE_SUBFOLDER))
+hf_hub_download(MODEL_REPO, f"{CODE_SUBFOLDER}/model.py", local_dir=WORK, token=HF_TOKEN)
+from model import ViuResonance100M
+
+def with_retry(fn, *args, retries=4, delay=5, **kwargs):
+    for attempt in range(1, retries + 1):
+        try:
+            return fn(*args, **kwargs)
+        except Exception as e:
+            if "404" in str(e) or "EntryNotFoundError" in type(e).__name__:
+                raise
+            if attempt == retries:
+                raise
+            log(f"  [retry {attempt}/{retries}] {type(e).__name__}: {e}")
+            time.sleep(delay)
+            delay *= 2
+
+class ShardPool:
+    def __init__(self, shard_names, repo_id, token, local_dir, pool_size=3, seed=0):
+        self.all_shards = list(shard_names)
+        self.rng = random.Random(seed)
+        self.rng.shuffle(self.all_shards)
+        self.cursor = 0
+        self.repo_id, self.token, self.local_dir, self.pool_size = repo_id, token, local_dir, pool_size
+        os.makedirs(local_dir, exist_ok=True)
+        self.pool = []
+        self.fill()
+
+    def _next_name(self):
+        if self.cursor >= len(self.all_shards):
+            self.rng.shuffle(self.all_shards)
+            self.cursor = 0
+        name = self.all_shards[self.cursor]
+        self.cursor += 1
+        return name
+
+    def fill(self):
+        while len(self.pool) < self.pool_size:
+            name = self._next_name()
+            path = with_retry(hf_hub_download, self.repo_id, name, repo_type="dataset",
+                               local_dir=self.local_dir, token=self.token)
+            mmap_arr = np.memmap(path, dtype=np.uint16, mode="r")
+            if len(mmap_arr) < 2049: # Context length (2048) + 1
+                if IS_MAIN:
+                    print(f"⚠️ Skipping shard {name} as it is too short ({len(mmap_arr)} tokens).")
+                continue
+            self.pool.append((path, mmap_arr))
+
+    def rotate_one(self):
+        old_path, old_mmap = self.pool.pop(0)
+        del old_mmap
+        try: os.remove(old_path)
+        except OSError: pass
+        self.fill()
+
+    def sample_batch(self, batch_size, ctx_len):
+        xs, ys = [], []
+        while len(xs) < batch_size:
+            _, tokens = self.rng.choice(self.pool)
+            max_start = len(tokens) - ctx_len - 1
+            if max_start <= 0:
+                continue
+            i = self.rng.randint(0, max_start)
+            xs.append(np.array(tokens[i:i+ctx_len], dtype=np.int64))
+            ys.append(np.array(tokens[i+1:i+1+ctx_len], dtype=np.int64))
+        return np.stack(xs), np.stack(ys)
+
+def load_val_tokens(val_names, repo_id, token, local_dir):
+    os.makedirs(local_dir, exist_ok=True)
+    arrays = []
+    for name in val_names:
+        path = with_retry(hf_hub_download, repo_id, name, repo_type="dataset", local_dir=local_dir, token=token)
+        arrays.append(np.memmap(path, dtype=np.uint16, mode="r"))
+    return np.concatenate(arrays)
+
+def get_batch(source, batch_size, ctx_len):
+    if isinstance(source, ShardPool):
+        x_np, y_np = source.sample_batch(batch_size, ctx_len)
+    else:
+        i = np.random.randint(0, len(source) - ctx_len - 1, size=batch_size)
+        x_np = np.stack([np.array(source[j:j+ctx_len], dtype=np.int64) for j in i])
+        y_np = np.stack([np.array(source[j+1:j+1+ctx_len], dtype=np.int64) for j in i])
+    x = torch.from_numpy(x_np).to(DEVICE, non_blocking=True)
+    y = torch.from_numpy(y_np).to(DEVICE, non_blocking=True)
+    return x, y
+
+def extract_clean_state_dict(model):
+    raw_model = model.module if hasattr(model, "module") else model
+    state_dict = raw_model.state_dict()
+    return {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
+
+def save_checkpoint(model, optimizer, step, val_loss, filename):
+    if not IS_MAIN:
+        return
+    path = f"{WORK}/{filename}"
+    state_dict = extract_clean_state_dict(model)
+    data = {"model": state_dict, "step": step, "val_loss": val_loss}
+    
+    if filename == "ckpt_latest.pt" and optimizer is not None:
+        data["optimizer"] = optimizer.state_dict()
+        
+    torch.save(data, path)
+    with_retry(api.upload_file, path_or_fileobj=path, path_in_repo=f"{CKPT_SUBFOLDER}/{filename}",
+               repo_id=CKPT_REPO, repo_type="model", token=HF_TOKEN)
+    log(f"  [{filename} saved @ step {step}, val_loss={val_loss:.4f}]")
+
+def load_checkpoint(model, optimizer):
+    try:
+        try:
+            path = with_retry(hf_hub_download, CKPT_REPO, f"{CKPT_SUBFOLDER}/ckpt_latest.pt", repo_type="model",
+                               local_dir=WORK, token=HF_TOKEN, retries=1)
+        except Exception as e:
+            if "404" in str(e) or "EntryNotFoundError" in type(e).__name__:
+                log("No checkpoint found on HF — starting from scratch.")
+                return 0
+            raise
+        
+        ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
+        target = model.module._orig_mod if hasattr(model.module, '_orig_mod') else model.module
+        target.load_state_dict(ckpt["model"])
+        if optimizer is not None and "optimizer" in ckpt:
+            optimizer.load_state_dict(ckpt["optimizer"])
+        log(f"Resumed from step {ckpt['step']} (val_loss={ckpt['val_loss']:.4f})")
+        return ckpt["step"]
+    except Exception as e:
+        log(f"Failed to load checkpoint — starting fresh. ({type(e).__name__}: {e})")
+        return 0
+
+def get_lr(step, warmup, total, max_lr, min_lr):
+    if step < warmup:
+        return max_lr * (step + 1) / warmup
+    if step > total:
+        return min_lr
+    ratio = (step - warmup) / (total - warmup)
+    return min_lr + 0.5 * (1 + math.cos(math.pi * ratio)) * (max_lr - min_lr)
+
+# === GPU CONFIG ===
+mem_gb = torch.cuda.get_device_properties(DEVICE).total_memory / (1024**3)
+if mem_gb > 100:     # H200
+    MICRO_BATCH = 32
+    GRAD_ACCUM = 4
+elif mem_gb > 70:    # A100
+    MICRO_BATCH = 32
+    GRAD_ACCUM = 4
+else:                # T4 or RTX (15GB/24GB)
+    MICRO_BATCH = 2
+    GRAD_ACCUM = 16
+
+if IS_MAIN:
+    print(f"🖥️ Detected GPU with {mem_gb:.1f}GB VRAM -> Auto-Configured: MICRO_BATCH={MICRO_BATCH}, GRAD_ACCUM={GRAD_ACCUM}")
+# ==================
+
+CONTEXT_LEN   = 2048
+MAX_LR        = 3e-4
+MIN_LR        = 3e-5
+WARMUP_STEPS  = 1000
+TOTAL_STEPS   = 40_000
+WEIGHT_DECAY  = 0.1
+GRAD_CLIP     = 1.0
+MAX_STEPS     = int(os.environ.get("MAX_STEPS", TOTAL_STEPS))
+EVAL_EVERY    = 500
+SAVE_EVERY    = 2000
+ROTATE_EVERY  = 500
+
+torch.backends.cudnn.benchmark = False
+
+# Model 100M init
+model = ViuResonance100M(vocab_size=VOCAB_SIZE, d_model=1024, n_layers=12, n_heads=16).to(DEVICE)
+
+if IS_MAIN:
+    log(f"Model params: {sum(p.numel() for p in model.parameters())/1e6:.1f}M | world_size={world_size}")
+    log("Compiling model (if it fails, will fallback to eager)...")
+
+try:
+    model = torch.compile(model)
+except Exception as e:
+    if IS_MAIN:
+        log(f"⚠️ torch.compile failed, falling back to eager mode: {e}")
+model = DDP(model, device_ids=[local_rank])
+
+optimizer = torch.optim.AdamW(model.parameters(), lr=MAX_LR, weight_decay=WEIGHT_DECAY, betas=(0.9, 0.95), fused=True)
+
+start_step = load_checkpoint(model, optimizer)
+start_step_t = torch.tensor([start_step], device=DEVICE)
+dist.broadcast(start_step_t, src=0)
+start_step = start_step_t.item()
+
+files = api.list_repo_files(DATA_REPO, repo_type="dataset")
+train_shards = sorted(f for f in files if f.startswith("shards/") and f.endswith(".bin") and "_val" not in f)
+val_shards = sorted(f for f in files if f.startswith("shards/") and f.endswith("_val.bin"))
+
+shard_pool = ShardPool(train_shards, DATA_REPO, HF_TOKEN, f"{WORK}/train_shards_r{rank}", seed=rank + start_step)
+val_tokens = load_val_tokens(val_shards, DATA_REPO, HF_TOKEN, f"{WORK}/val_shards") if IS_MAIN else None
+
+@torch.no_grad()
+def evaluate():
+    if not IS_MAIN:
+        return None
+    model.eval()
+    losses = []
+    for _ in range(20):
+        x, y = get_batch(val_tokens, 16, CONTEXT_LEN)
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
+            unwrapped = model.module._orig_mod if hasattr(model.module, '_orig_mod') else model.module
+            _, loss = unwrapped(x, y)
+        losses.append(loss.item())
+    model.train()
+    return sum(losses) / len(losses)
+
+model.train()
+val_loss = None
+best_val_loss = None
+ema_tok_s = None
+global_t0 = time.time()
+t0 = time.time()
+end_step = min(start_step + MAX_STEPS, TOTAL_STEPS)
+
+for step in range(start_step, end_step):
+    lr = get_lr(step, WARMUP_STEPS, TOTAL_STEPS, MAX_LR, MIN_LR)
+    for g in optimizer.param_groups:
+        g["lr"] = lr
+
+    optimizer.zero_grad(set_to_none=True)
+    accum_loss = 0.0
+    for micro_i in range(GRAD_ACCUM):
+        x, y = get_batch(shard_pool, MICRO_BATCH, CONTEXT_LEN)
+        sync_ctx = model.no_sync() if micro_i < GRAD_ACCUM - 1 else contextlib.nullcontext()
+        with sync_ctx:
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
+                _, loss = model(x, y)
+                loss = loss / GRAD_ACCUM
+            loss.backward()
+        accum_loss += loss.item()
+        del x, y, loss
+
+    torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
+
+    if not math.isfinite(accum_loss):
+        if IS_MAIN:
+            log(f"⚠️ NaN/Inf loss at step {step}! Skipping optimizer step.")
+        optimizer.zero_grad(set_to_none=True)
+        continue
+
+    optimizer.step()
+
+    if IS_MAIN and step % 10 == 0:
+        if step == start_step:
+            t0 = time.time()
+            mem_alloc = torch.cuda.memory_allocated(DEVICE) / 1e9
+            mem_reserved = torch.cuda.memory_reserved(DEVICE) / 1e9
+            log(f"step {step:5d}/{end_step} | loss {accum_loss:.4f} | lr {lr:.2e} | Calibrating speed... | "
+                f"mem alloc={mem_alloc:.2f}GB reserved={mem_reserved:.2f}GB")
+        else:
+            dt = time.time() - t0
+            tok_s = (MICRO_BATCH * GRAD_ACCUM * CONTEXT_LEN * world_size * 10) / dt if dt > 0 else 0
+            ema_tok_s = tok_s if ema_tok_s is None else 0.7 * ema_tok_s + 0.3 * tok_s
+            t0 = time.time()
+            mem_alloc = torch.cuda.memory_allocated(DEVICE) / 1e9
+            mem_reserved = torch.cuda.memory_reserved(DEVICE) / 1e9
+            steps_done = step - start_step
+            if steps_done > 0:
+                elapsed = time.time() - global_t0
+                eta_s = elapsed / steps_done * (end_step - step)
+                eta_str = f"{eta_s/60:.0f}min" if eta_s < 3600 else f"{eta_s/3600:.1f}hr"
+            else:
+                eta_str = "..."
+            log(f"step {step:5d}/{end_step} | loss {accum_loss:.4f} | lr {lr:.2e} | {ema_tok_s:.0f} tok/s | "
+                f"ETA {eta_str} | mem alloc={mem_alloc:.2f}GB reserved={mem_reserved:.2f}GB")
+
+    if step > start_step and step % ROTATE_EVERY == 0:
+        shard_pool.rotate_one()
+
+    if step > start_step and step % EVAL_EVERY == 0:
+        val_loss = evaluate()
+        if IS_MAIN and val_loss is not None:
+            log(f"step {step} | val_loss {val_loss:.4f}")
+            if best_val_loss is None or val_loss < best_val_loss:
+                best_val_loss = val_loss
+
+    if step > start_step and step % SAVE_EVERY == 0:
+        dist.barrier()
+        if IS_MAIN:
+            save_checkpoint(model.module._orig_mod if hasattr(model.module, '_orig_mod') else model.module, 
+                            optimizer, step, val_loss if val_loss is not None else -1, "ckpt_latest.pt")
+
+dist.barrier()
+final_val = evaluate()
+if IS_MAIN:
+    save_checkpoint(model.module._orig_mod if hasattr(model.module, '_orig_mod') else model.module, 
+                    optimizer, end_step, final_val if final_val is not None else -1, "ckpt_latest.pt")
+    log(f"\n100M Training complete! Final val_loss: {final_val if final_val is not None else -1:.4f}")
+
+dist.destroy_process_group()