code/kod/faz2_train_shards.py

"""
Faz 2/3 — Pre-shard veriyle eğitim (gerçek veri yolu).

Pre-tokenize parquet shard'lardan okur (yerel dizin VEYA HF dataset repo), kaynakları
oranlarına göre karıştırır, hibrit modeli WSD ile eğitir. input_ids zaten tokenize
(uint16) + dekontamine → loop'ta SP/decontam YOK (hızlı, GPU-aç-bırakmaz).

Faz 2 smoke (yerel CPU, gerçek shard'lar):
  python kod/faz2_train_shards.py --data kod/data/shards --d_model 256 --n_layer 6 --seq_len 2048 --steps 5 --device cpu
Colab smoke (HF'den stream, GPU):
  python kod/faz2_train_shards.py --data kdirgul/smartcore-v1-data --hf --d_model 768 --n_layer 20 --seq_len 2048 --steps 200 --device cuda --bf16
"""
import os, sys, time, math, argparse, random, glob
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import torch
import torch.nn.functional as F
import pyarrow.parquet as pq

MIX = {"en_fineweb_edu": 0.55, "tr_fineweb2_hq": 0.22, "code_codeparrot": 0.13, "math_openwebmath": 0.10}


def wsd_lr(step, total, peak, floor, warmup_frac=0.02, decay_frac=0.25):
    warm = max(1, int(total * warmup_frac)); dec = int(total * (1 - decay_frac))
    if step < warm: return peak * (step + 1) / warm
    if step < dec:  return peak
    return peak - (peak - floor) * (step - dec) / max(1, total - dec)


class ShardStream:
    """Pre-shard parquet okuyucu (yerel dizin veya HF repo), oranlı karışım."""
    def __init__(self, root, weights, seq_len, hf=False, token=None):
        self.names = list(weights); self.w = [weights[n] for n in self.names]
        self.seq_len = seq_len   # satırlar 2048; daha kısa smoke için kırp
        self.iters = {n: self._src(root, n, hf, token) for n in self.names}

    def _src(self, root, name, hf, token):
        if hf:
            from datasets import load_dataset
            while True:  # epoch döngüsü
                ds = load_dataset(root, data_dir=name, split="train", streaming=True, token=token)
                for rec in ds:
                    yield rec["input_ids"]
        else:
            files = sorted(glob.glob(os.path.join(root, name, "shard_*.parquet")))
            assert files, f"shard bulunamadı: {root}/{name}"
            while True:
                for fp in files:
                    for row in pq.read_table(fp, columns=["input_ids"]).column("input_ids"):
                        yield row.as_py()

    def batch(self, bsz):
        rows = [next(self.iters[random.choices(self.names, weights=self.w, k=1)[0]])[:self.seq_len]
                for _ in range(bsz)]
        return torch.tensor(rows, dtype=torch.long)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="kod/data/shards", help="yerel dizin veya HF repo id")
    ap.add_argument("--hf", action="store_true", help="--data bir HF dataset repo id")
    ap.add_argument("--d_model", type=int, default=256)
    ap.add_argument("--n_layer", type=int, default=6)
    ap.add_argument("--seq_len", type=int, default=2048)
    ap.add_argument("--micro_batch", type=int, default=4)
    ap.add_argument("--steps", type=int, default=5)
    ap.add_argument("--peak_lr", type=float, default=5e-4)
    ap.add_argument("--device", default="cpu")
    ap.add_argument("--bf16", action="store_true")
    args = ap.parse_args()

    dev = torch.device(args.device if (args.device != "cuda" or torch.cuda.is_available()) else "cpu")
    torch.manual_seed(0); random.seed(0)

    from hybrid_mamba3 import make_config, build_hybrid, n_params
    # vocab: model çıktısı; shard'lar bu vocab'a göre tokenize edildi (48000)
    cfg = make_config(d_model=args.d_model, n_layer=args.n_layer, vocab=48000,
                      d_mlp_inner=1500 if args.d_model >= 768 else args.d_model * 2,
                      chunk_size=min(64, args.seq_len))
    attn_every = 6 if args.n_layer >= 8 else 3
    model, attn_idx = build_hybrid(cfg, attn_every=attn_every, n_heads=max(2, args.d_model // 64),
                                   n_kv_heads=max(1, args.d_model // 256), device=dev)
    model.to(dev); model.train()
    print(f"model: d={args.d_model} L={args.n_layer} ({args.n_layer-len(attn_idx)} Mamba + "
          f"{len(attn_idx)} GQA) | params={n_params(model)/1e6:.1f}M | dev={dev}")

    tok = os.environ.get("HF_TOKEN")
    print(f"veri: {args.data} ({'HF stream' if args.hf else 'yerel parquet'}) | mixture {MIX}")
    stream = ShardStream(args.data, MIX, args.seq_len, hf=args.hf, token=tok)

    decay = [p for p in model.parameters() if p.ndim >= 2]
    nod = [p for p in model.parameters() if p.ndim < 2]
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.1},
                             {"params": nod, "weight_decay": 0.0}],
                            lr=args.peak_lr, betas=(0.9, 0.95), eps=1e-8)
    use_bf16 = args.bf16 and dev.type == "cuda"
    print(f"smoke | {args.steps} adım | bf16={use_bf16} | başlangıç loss ~ln(48000)=10.78")

    t0 = time.perf_counter(); seen = 0; loss = None
    for step in range(args.steps):
        batch = stream.batch(args.micro_batch).to(dev)
        for g in opt.param_groups:
            g["lr"] = wsd_lr(step, args.steps, args.peak_lr, args.peak_lr * 0.1)
        opt.zero_grad(set_to_none=True)
        ctx = torch.autocast(device_type="cuda", dtype=torch.bfloat16) if use_bf16 else _null()
        with ctx:
            logits, _ = model(batch)
            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)).float(),
                                   batch[:, 1:].reshape(-1))
        loss.backward()
        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        seen += batch.numel()
        print(f"  step {step:3d} | loss {loss.item():6.3f} | grad_norm {gn:5.2f} | "
              f"lr {opt.param_groups[0]['lr']:.2e} | {seen/(time.perf_counter()-t0)/1e3:.1f}k tok/s")
    print(f"\nsmoke BİTTİ. NaN yok: {not math.isnan(loss.item())} | "
          f"{seen/(time.perf_counter()-t0)/1e3:.1f}k tok/s")


class _null:
    def __enter__(self): return self
    def __exit__(self, *a): return False


if __name__ == "__main__":
    main()