%%writefile train_jirack_bitnet_1b_lambda_warmup.py

# =============================================================================
# COPYRIGHT © 2025-2026  Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
# CMS Manhattan JiRack Technology — PATENT PENDING
#
# This code is proprietary. 
# Personal and non-commercial research use is allowed.
# Any commercial use, derivative works for profit, or distribution 
# requires a paid license and 5% royalty.
#
# Unauthorized commercial use is strictly prohibited.
# Contact: grabko@cmsmanhattan.com
# =============================================================================
# train_jirack_bitnet_1b_lambda_warmup.py
# =============================================================================
# QAT training with lambda warmup — for the SubLN-equipped, bitnet.cpp-
# compatible architecture (JiRackTernaryPyTorch_1b_bitnet.py).
#
# Same warmup machinery as the 10B script, re-pointed at the 1B model.
# Differences from the 10B version, all deliberate:
#
#   [B-1] import from JiRackTernaryPyTorch_1b_bitnet (SubLN architecture).
#   [B-2] BASE_CHECKPOINT points at the output of
#         load_old_checkpoint_into_bitnet_arch_1b.py — the migrated
#         checkpoint where SubLN weights are freshly initialized at 1.0.
#   [B-3] BATCH_SIZE raised 2 -> 8 and GRAD_ACCUM lowered 5 -> 2. The 1B
#         model is ~10x smaller, so on a 96GB card there is no reason to
#         crawl at micro-batch 2. Effective batch stays 16 (was 10).
#   [B-4] LAMBDA_WARMUP_STEPS 5000 -> 2000 micro-batches. The 1B weights
#         are ALREADY fully ternary-trained (the original had no lambda at
#         all — it ran quantized from step one), so this warmup only has to
#         settle the new SubLN weights, not re-ternarize from scratch.
#
# Everything else — the sigmoid lambda schedule, resume-by-global-step,
# optimizer state save/restore, lambda-per-accumulation-window, lambda-
# tagged validation — is unchanged.
#
# WHAT TO WATCH, 1B-specific: this run changes TWO things at once, SubLN
# appearing and activation quantization dropping its mean-centering ([C-3]
# in the architecture file). So expect a real loss spike as lambda crosses
# the steep middle of the sigmoid. Judge by val loss @ a given lambda
# trending down shard over shard. If it has not recovered by the end of
# warmup, raise LAMBDA_WARMUP_STEPS (3000-4000) rather than touching LR —
# the sigmoid midpoint is the sensitive part, not the learning rate.
# =============================================================================

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
import glob
import re
import gc
import math
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from tqdm import tqdm
from transformers.optimization import Adafactor

# [B-1] SubLN / bitnet.cpp-compatible model class
from JiRackTernaryPyTorch_1b_bitnet import JiRackTransformer1B, JiRackConfig1B

# ========================= НАСТРОЙКИ =========================
DATA_DIR = "/content/prepared_sft_data"
OUTPUT_DIR = "/content/JRock_1B_Bitnet_Checkpoints"
# [B-2] output of load_old_checkpoint_into_bitnet_arch_1b.py
BASE_CHECKPOINT = "/content/jirack_1b_bitnet_init.pt"

BATCH_SIZE = 8          # [B-3] 1B fits comfortably; raise further if VRAM allows
GRAD_ACCUM = 2          # [B-3] effective batch = 16
LR = 4e-5
VAL_RATIO = 0.05

# === Lambda Warmup Settings ===
# [B-4] 2000 micro-batches = 1000 optimizer steps @ GRAD_ACCUM=2.
# Shorter than the 10B run on purpose: the 1B weights are already ternary,
# this warmup only has to settle the new SubLN. If val loss has not
# recovered by the end of warmup, raise this (3000-4000), not LR.
LAMBDA_WARMUP_STEPS = 2000      # in MICRO-batches
MAX_LAMBDA = 1.0
SAVE_OPTIMIZER = True
AUTOSAVE_EVERY = 1000
# =============================================================

os.makedirs(OUTPUT_DIR, exist_ok=True)

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

print("🚀 Loading JiRack 1B Ternary (SubLN / bitnet.cpp) + Lambda Warmup...")

config = JiRackConfig1B()
# [B-3] gradient checkpointing OFF for 1B — it trades ~30% speed for memory
# you do not need at this size. Flip to True if you raise BATCH_SIZE a lot.
USE_GRAD_CHECKPOINT = False

model = JiRackTransformer1B(config, use_checkpoint=USE_GRAD_CHECKPOINT)
model.to("cuda")

for block in model.blocks:
    block.use_checkpoint = USE_GRAD_CHECKPOINT
print(f"✅ Gradient checkpointing: {'enabled' if USE_GRAD_CHECKPOINT else 'disabled'}")

optimizer = Adafactor(
    model.parameters(),
    lr=LR,
    eps=(1e-30, 1e-3),
    clip_threshold=1.0,
    decay_rate=-0.8,
    weight_decay=0.0001,
    scale_parameter=False,
    relative_step=False,
    warmup_init=False,
)

criterion = nn.CrossEntropyLoss(ignore_index=-100)


# ==================== LAMBDA SCHEDULE ====================
def lambda_schedule(step: int) -> float:
    """Sigmoid ramp over LAMBDA_WARMUP_STEPS micro-batches."""
    if step >= LAMBDA_WARMUP_STEPS:
        return MAX_LAMBDA
    return MAX_LAMBDA / (
        1.0 + math.exp(-10 * (step - LAMBDA_WARMUP_STEPS / 2) / LAMBDA_WARMUP_STEPS)
    )


def invert_lambda_schedule(lam: float) -> int:
    """Reconstruct global_step from a saved lambda (legacy checkpoints
    that stored only the model state_dict). Inverse of the sigmoid above."""
    if lam >= MAX_LAMBDA * 0.999:
        return LAMBDA_WARMUP_STEPS
    if lam <= 1e-6:
        return 0
    p = lam / MAX_LAMBDA
    x = -math.log(1.0 / p - 1.0)  # logit
    return int(round(x * LAMBDA_WARMUP_STEPS / 10 + LAMBDA_WARMUP_STEPS / 2))


# ==================== CHECKPOINT LOAD ====================
def load_any_checkpoint(path, model, optimizer):
    """Handles both new-format dicts and legacy plain state_dicts.
    Returns restored global_step.

    NOTE: on the very first load (BASE_CHECKPOINT = migration output),
    the state_dict has EVERY tensor including attn_sub_norm/ffn_sub_norm
    (the migration script already wrote them, at value 1.0), so this is
    a normal strict-ish load — no special-casing needed here for SubLN."""
    ckpt = torch.load(path, map_location="cpu", weights_only=True)

    if isinstance(ckpt, dict) and "model" in ckpt:
        missing, unexpected = model.load_state_dict(ckpt["model"], strict=False)
        assert not unexpected, unexpected
        assert all(k.endswith("lambda_") for k in missing), missing
        if SAVE_OPTIMIZER and "optimizer" in ckpt and ckpt["optimizer"] is not None:
            try:
                optimizer.load_state_dict(ckpt["optimizer"])
                print("✅ Optimizer state restored")
            except Exception as e:
                print(f"⚠️ Optimizer state not restored ({e}); continuing fresh")
        step = int(ckpt.get("global_step", 0))
        print(f"✅ Resumed (new format): global_step={step}, "
              f"lambda={ckpt.get('lambda', 'n/a')}")
        return step

    # legacy: plain state_dict (e.g. straight from load_from_old_checkpoint
    # + torch.save, no wrapper dict) — missing keys should only be lambda_
    # buffers here, since the migration script already fills SubLN weights.
    missing, unexpected = model.load_state_dict(ckpt, strict=False)
    assert not unexpected, unexpected
    assert all(k.endswith("lambda_") for k in missing), missing
    lam = model.get_lambda()
    step = invert_lambda_schedule(lam)
    print(f"✅ Resumed (legacy format): lambda={lam:.4f} -> "
          f"reconstructed global_step={step}")
    return step


checkpoints = sorted(
    glob.glob(os.path.join(OUTPUT_DIR, "jirack_1b_data_*.pt")),
    key=lambda x: int(re.search(r"data_(\d+)", x).group(1)),
)
all_ckpts = glob.glob(os.path.join(OUTPUT_DIR, "*.pt"))

global_step = 0
if all_ckpts:
    LATEST_CKPT = max(all_ckpts, key=os.path.getmtime)
    print(f"📦 Resuming from (newest on disk): {LATEST_CKPT}")
    global_step = load_any_checkpoint(LATEST_CKPT, model, optimizer)
else:
    print(f"📦 Loading base (migrated) checkpoint: {BASE_CHECKPOINT}")
    assert os.path.exists(BASE_CHECKPOINT), (
        f"{BASE_CHECKPOINT} not found — run "
        f"load_old_checkpoint_into_bitnet_arch_1b.py first"
    )
    global_step = load_any_checkpoint(BASE_CHECKPOINT, model, optimizer)
    # Fresh migration -> starting lambda warmup from scratch, even if the
    # checkpoint happened to carry a lambda value from the old run.
    global_step = 0
    print("🔄 Fresh SubLN migration detected: lambda warmup restarted from step 0 "
          "(SubLN has never seen a nonzero lambda before).")

model.set_lambda(lambda_schedule(global_step))
print(f"🔧 Warmup: {LAMBDA_WARMUP_STEPS} steps | starting at "
      f"step={global_step}, lambda={model.get_lambda():.4f}")


# ========================= DATASET =========================
class ShardDataset(Dataset):
    def __init__(self, data_list):
        self.data = data_list

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]


def collate_fn(batch):
    input_ids = pad_sequence(
        [item["input_ids"] for item in batch], batch_first=True, padding_value=0
    )
    attention_mask = pad_sequence(
        [item.get("attention_mask", torch.ones_like(item["input_ids"]))
         for item in batch],
        batch_first=True, padding_value=0,
    )
    labels = input_ids.clone()
    labels[attention_mask == 0] = -100
    return {"input_ids": input_ids, "labels": labels}


# ========================= CHECKPOINT SAVE =========================
def save_checkpoint(path, model, optimizer, global_step, lam):
    model_sd = {k: v.detach().to(torch.bfloat16).cpu()
                for k, v in model.state_dict().items()}
    ckpt = {
        "model": model_sd,
        "optimizer": optimizer.state_dict() if SAVE_OPTIMIZER else None,
        "global_step": global_step,
        "lambda": lam,
    }
    torch.save(ckpt, path)


# ========================= TRAINING =========================
all_shards = sorted(
    glob.glob(f"{DATA_DIR}/sft_data_*.pt"),
    key=lambda x: int(re.search(r"sft_data_(\d+)", x).group(1)),
)

last_done = -1
if checkpoints:
    m = re.search(r"data_(\d+)", checkpoints[-1])
    last_done = int(m.group(1)) if m else -1

for shard_path in all_shards:
    shard_name = os.path.basename(shard_path)
    shard_num = int(re.search(r"sft_data_(\d+)", shard_name).group(1))

    if shard_num <= last_done:
        print(f"⏭ Skipping already processed: {shard_name}")
        continue

    print(f"\n🔥 Starting shard: {shard_name}")

    raw_shard_data = torch.load(shard_path, map_location="cpu", weights_only=False)
    val_size = int(len(raw_shard_data) * VAL_RATIO)
    train_size = len(raw_shard_data) - val_size

    train_data, val_data = torch.utils.data.random_split(
        raw_shard_data, [train_size, val_size],
        generator=torch.Generator().manual_seed(42),
    )

    train_loader = DataLoader(
        ShardDataset(train_data), batch_size=BATCH_SIZE, shuffle=True,
        collate_fn=collate_fn, pin_memory=True,
    )

    model.train()
    pbar = tqdm(train_loader, desc=f"Shard {shard_num}", dynamic_ncols=True)
    optimizer.zero_grad()

    lambda_value = lambda_schedule(global_step)
    model.set_lambda(lambda_value)

    for step, batch in enumerate(pbar):
        if step % GRAD_ACCUM == 0:
            lambda_value = lambda_schedule(global_step)
            model.set_lambda(lambda_value)

        input_ids = batch["input_ids"].to("cuda", non_blocking=True)
        labels = batch["labels"].to("cuda", non_blocking=True)

        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            logits = model(input_ids)
            if isinstance(logits, tuple):
                logits = logits[0]
            loss = criterion(
                logits[..., :-1, :].reshape(-1, config.vocab_size),
                labels[..., 1:].reshape(-1),
            )
            loss = loss / GRAD_ACCUM

        if torch.isnan(loss) or torch.isinf(loss):
            print(f"\n⚠️ NaN/Inf loss at step {global_step} "
                  f"(lambda={lambda_value:.4f}) — window dropped")
            optimizer.zero_grad(set_to_none=True)
            torch.cuda.empty_cache()
            global_step += 1
            continue

        loss.backward()

        if (step + 1) % GRAD_ACCUM == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()

        if step % 10 == 0:
            pbar.set_postfix({
                "loss": f"{loss.item() * GRAD_ACCUM:.4f}",
                "lambda": f"{lambda_value:.4f}",
                "gstep": global_step,
            })

        if (AUTOSAVE_EVERY > 0 and global_step > 0
                and global_step % AUTOSAVE_EVERY == 0
                and (step + 1) % GRAD_ACCUM == 0):
            autosave_path = os.path.join(OUTPUT_DIR, "autosave_latest.pt")
            tmp_path = autosave_path + ".tmp"
            save_checkpoint(tmp_path, model, optimizer, global_step, lambda_value)
            os.replace(tmp_path, autosave_path)
            pbar.write(f"💾 autosave @ gstep={global_step}, "
                       f"lambda={lambda_value:.4f}")

        global_step += 1

    # ==================== Validation ====================
    print("🧪 Validating...")
    model.eval()
    total_val_loss = 0.0
    val_steps = 0

    val_loader = DataLoader(
        ShardDataset(val_data), batch_size=BATCH_SIZE, shuffle=False,
        collate_fn=collate_fn, pin_memory=True,
    )

    with torch.no_grad():
        for batch in tqdm(val_loader, desc="Validating", leave=False):
            input_ids = batch["input_ids"].to("cuda", non_blocking=True)
            labels = batch["labels"].to("cuda", non_blocking=True)
            with torch.amp.autocast("cuda", dtype=torch.bfloat16):
                logits = model(input_ids)
                if isinstance(logits, tuple):
                    logits = logits[0]
                v_loss = criterion(
                    logits[..., :-1, :].reshape(-1, config.vocab_size),
                    labels[..., 1:].reshape(-1),
                )
                if not (torch.isnan(v_loss) or torch.isinf(v_loss)):
                    total_val_loss += v_loss.item()
                    val_steps += 1

    avg_val_loss = total_val_loss / val_steps if val_steps > 0 else float("inf")
    print(f"📊 Shard {shard_num} — Val Loss: {avg_val_loss:.4f} "
          f"@ lambda={lambda_value:.4f} (gstep={global_step})")

    save_path = os.path.join(OUTPUT_DIR, f"jirack_1b_data_{shard_num}.pt")
    save_checkpoint(save_path, model, optimizer, global_step, lambda_value)
    print(f"💾 Saved: {save_path}")

    del raw_shard_data
    torch.cuda.empty_cache()
    gc.collect()

print("🏁 Training finished with Lambda Warmup (SubLN / bitnet.cpp arch)!")