train_toolace_lora_v2.py

# ==============================================================================
# JiRack ToolACE LoRA SFT + merge (single script)
# COPYRIGHT (c) 2026 Konstantin Vladimirovich Grabko.
#
# Trains tool-calling on your converted ToolACE dataset
# (toolace_sft_jirack_precision_8b.jsonl, verified 100% parse rate) using LoRA
# adapters injected directly into JiRackTransformer's nn.Linear layers --
# no PEFT/HF-model needed, works with your custom class and .pt checkpoint.
#
# What it does:
#   1. Loads JiRackTransformer + your expanded .pt (same as chat scripts).
#   2. Freezes everything; injects LoRA (A/B low-rank pairs) into every
#      nn.Linear except the LM head (out_features == vocab_size).
#   3. ALSO unfreezes the embedding rows of your custom special tokens
#      (additional_special_tokens) -- those rows are untrained noise right now;
#      the model can't emit <|tool_call_start|> etc. until they're trained.
#      A gradient hook zeroes grads for all other rows, so the base vocab
#      embeddings stay untouched.
#   4. Trains with assistant-only loss masking (user/system/tool-result tokens
#      get label -100), bf16 autocast, gradient accumulation.
#   5. Saves: (a) the LoRA adapter alone (small, reusable), and
#             (b) OPTIONAL merged checkpoint: W' = W + (alpha/r)*B@A folded into
#                 the original weights, wrappers removed -> state_dict has
#                 EXACTLY the same keys as the input .pt. Drop-in replacement
#                 for chat_jirack_8b_tools.py's MODEL_PATH.
#
# Hardware notes:
#   * GPU strongly recommended. CPU works but will be very slow; remember:
#       export MKL_ENABLE_INSTRUCTIONS=AVX
#       export MKL_DEBUG_CPU_TYPE=5
#   * Memory: base weights sit in bf16 frozen (no optimizer state for them).
#     Optimizer state only for LoRA params (+2 embedding matrices' rows).
#     8B on a 48GB GPU fits comfortably at MAX_LEN=2048, BATCH=1, accum=16.
#   * Ternary: training runs at set_lambda(0.0) (full precision). If you want
#     quant-aware finetuning instead, raise LAMBDA below -- but for teaching
#     tool-call FORMAT, full precision is the right call.
#
# For other sizes: edit the import + 3 paths below (16B/2B/36B analogous).
# ==============================================================================

import json
import math
import os
import random
import sys
import time

import torch
import torch.nn as nn
from transformers import AutoTokenizer
from transformers.optimization import Adafactor

try:
    import bitsandbytes as bnb
    _HAS_BNB = True
except ImportError:
    _HAS_BNB = False

sys.path.append(os.getcwd())
from JiRackPrecision_8b import JiRackTransformer, JiRackConfig

# ========================= EDIT THESE =========================
MODEL_PATH    = "/mnt/nfs_clientshare/JiRackPrecision_8b/ds8b_expanded.pt"
#TOKENIZER_DIR = "."
TOKENIZER_DIR = "/mnt/nfs_clientshare/JiPrecision_Tokenizer/ji_precision_tokenizer"
DATASET_PATH  = "/mnt/nfs_clientshare/JiRackPrecision_8b/toolace_sft_jirack_precision_8b.jsonl"
ADAPTER_OUT   = "/mnt/nfs_clientshare/JiRackPrecision_8b/toolace_lora_adapter.pt"
MERGED_OUT    = "/mnt/nfs_clientshare/JiRackPrecision_8b/ds8b_expanded_toolace.pt"

# LoRA
LORA_R        = 16
LORA_ALPHA    = 32
LORA_DROPOUT  = 0.05

# Training
EPOCHS        = 2
LR            = 2e-4          # LoRA params
EMBED_LR      = 5e-5          # new-token embedding rows (gentler)
BATCH_SIZE    = 1
GRAD_ACCUM    = 16
MAX_LEN       = 2048          # truncate long conversations
WARMUP_STEPS  = 50
SEED          = 42
LAMBDA        = 0.0           # 0.0 = full-precision training (recommended here)
SAVE_EVERY    = 500           # optimizer steps between adapter checkpoints
MERGE_AT_END  = True          # write MERGED_OUT after training
OPTIMIZER     = "adafactor"   # "adamw" or "adafactor"
                               # adafactor: no momentum buffer, ~2 bytes/param
                               # optimizer state vs AdamW's ~8 bytes/param --
                               # matters a lot on tight VRAM (e.g. 36B on 96GB).
FREEZE_8BIT   = False         # cast frozen base Linear weights to int8
                               # (bitsandbytes) to shrink the frozen backbone's
                               # footprint ~2x. LoRA/embeddings stay full bf16.
                               # Needs `pip install bitsandbytes`.
# ================================================================


# ------------------------------ LoRA machinery ------------------------------

class LoRALinear(nn.Module):
    """Wraps a frozen nn.Linear; adds trainable low-rank A/B path."""

    def __init__(self, base: nn.Linear, r: int, alpha: int, dropout: float):
        super().__init__()
        self.base = base
        self.r = r
        self.scale = alpha / r
        self.lora_A = nn.Parameter(torch.zeros(r, base.in_features))
        self.lora_B = nn.Parameter(torch.zeros(base.out_features, r))
        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
        # B starts at zero -> identity behavior at step 0
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()

    def forward(self, x):
        out = self.base(x)
        lx = self.dropout(x).to(self.lora_A.dtype)
        out = out + (lx @ self.lora_A.T @ self.lora_B.T) * self.scale
        return out

    @torch.no_grad()
    def merge_into_base(self):
        """Fold LoRA delta into the base weight. If the base is an int8
        bitsandbytes layer (FREEZE_8BIT=True), merging happens in fp32 and
        the result is a plain bf16 nn.Linear -- you get a full-precision
        merged checkpoint either way, the 8bit trick only ever affected
        training-time memory, never the final merged weights."""
        delta = (self.lora_B.float() @ self.lora_A.float()) * self.scale
        if _HAS_BNB and isinstance(self.base, bnb.nn.Linear8bitLt):
            w = self.base.weight
            dequant = bnb.functional.dequantize_4bit(w.data, w.quant_state) \
                if hasattr(w, "quant_state") and w.quant_state is not None \
                else self.base(torch.eye(self.base.in_features,
                                          dtype=torch.bfloat16,
                                          device=delta.device)).T  # fallback dequant via identity matmul
            merged_w = dequant.float() + delta
            new_linear = nn.Linear(self.base.in_features, self.base.out_features,
                                    bias=self.base.bias is not None)
            new_linear.weight.data = merged_w.to(torch.bfloat16)
            if self.base.bias is not None:
                new_linear.bias.data = self.base.bias.data.to(torch.bfloat16)
            self.base = new_linear
        else:
            self.base.weight.data += delta.to(self.base.weight.dtype)


def _to_8bit_linear(linear):
    """Swap a frozen nn.Linear for bitsandbytes' int8 version (weights only --
    this is inference-style quantization for a FROZEN layer, not QLoRA's
    double-quant, but good enough to roughly halve backbone memory)."""
    q = bnb.nn.Linear8bitLt(linear.in_features, linear.out_features,
                             bias=linear.bias is not None, has_fp16_weights=False)
    q.load_state_dict(linear.state_dict())
    for p in q.parameters():
        p.requires_grad = False
    return q


def inject_lora(model, vocab_size):
    """Replace every nn.Linear (except the vocab-sized head) with LoRALinear.
    If FREEZE_8BIT is set, the wrapped base layer is first cast to int8
    (bitsandbytes) to shrink the frozen backbone's memory footprint --
    LoRA's own A/B matrices always stay full precision regardless."""
    if FREEZE_8BIT and not _HAS_BNB:
        sys.exit("❌ FREEZE_8BIT=True but bitsandbytes isn't installed. "
                  "Run: pip install bitsandbytes")
    wrapped = []
    for parent_name, parent in list(model.named_modules()):
        for child_name, child in list(parent.named_children()):
            if isinstance(child, nn.Linear) and child.out_features != vocab_size:
                base = _to_8bit_linear(child) if FREEZE_8BIT else child
                setattr(parent, child_name,
                        LoRALinear(base, LORA_R, LORA_ALPHA, LORA_DROPOUT))
                full = f"{parent_name}.{child_name}" if parent_name else child_name
                wrapped.append(full)
    return wrapped


def merge_and_unwrap(model):
    """Fold LoRA into base weights and restore original nn.Linear modules,
    so state_dict() keys match the original checkpoint exactly."""
    for parent_name, parent in list(model.named_modules()):
        for child_name, child in list(parent.named_children()):
            if isinstance(child, LoRALinear):
                child.merge_into_base()
                setattr(parent, child_name, child.base)


# ------------------------------ Dataset ------------------------------

def load_dataset(path):
    convs = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            msgs = obj.get("messages", obj)
            if isinstance(msgs, list) and any(m.get("role") == "assistant" for m in msgs):
                convs.append(msgs)
    return convs


def build_example(tokenizer, messages, max_len):
    """Tokenize a conversation with assistant-only labels.
    Incremental templating: token span of message i = template(msgs[:i+1]) minus
    template(msgs[:i]). Labels = ids inside assistant spans, else -100."""
    ids, labels = [], []
    prev = []
    prev_len = 0
    for m in messages:
        prev.append(m)
        cur = tokenizer.apply_chat_template(prev, tokenize=True,
                                            add_generation_prompt=False)
        span = cur[prev_len:]
        if m["role"] == "assistant":
            labels.extend(span)
        else:
            labels.extend([-100] * len(span))
        ids = cur
        prev_len = len(cur)
        if len(ids) >= max_len:
            break
    ids = ids[:max_len]
    labels = labels[:max_len]
    if all(l == -100 for l in labels):
        return None
    return torch.tensor(ids), torch.tensor(labels)


# ------------------------------ Training ------------------------------

def main():
    random.seed(SEED)
    torch.manual_seed(SEED)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"🚀 Device: {device.upper()}")
    print(f"⚙️  Optimizer={OPTIMIZER}  FREEZE_8BIT={FREEZE_8BIT}  "
          f"(tip: for tight VRAM like 36B on 96GB, use adafactor + FREEZE_8BIT=True)")

    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
    vocab_rows = None

    # --- model ---
    config = JiRackConfig()
    model = JiRackTransformer(config, use_checkpoint=True)  # activation ckpt on
    print(f"📥 Loading {MODEL_PATH} ...")
    ckpt = torch.load(MODEL_PATH, map_location="cpu", weights_only=False)
    sd = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
    missing, unexpected = model.load_state_dict(sd, strict=False)
    real_missing = [k for k in missing if not k.endswith("lambda_")]
    if real_missing:
        print(f"⚠️ Missing keys: {real_missing[:10]}")
    model = model.to(dtype=torch.bfloat16, device=device)
    model.set_lambda(LAMBDA)

    # find embedding module + vocab size
    embed = None
    for mod in model.modules():
        if isinstance(mod, nn.Embedding):
            embed = mod
            break
    if embed is None:
        sys.exit("❌ No nn.Embedding found in model")
    vocab_rows = embed.weight.shape[0]
    print(f"   embedding rows: {vocab_rows}")

    # --- freeze all, inject LoRA ---
    for p in model.parameters():
        p.requires_grad = False
    wrapped = inject_lora(model, vocab_rows)
    model = model.to(device)
    print(f"🧩 LoRA injected into {len(wrapped)} Linear layers (r={LORA_R}, alpha={LORA_ALPHA})")

    lora_params = [p for n, p in model.named_parameters() if "lora_" in n]
    for p in lora_params:
        p.requires_grad = True

    # --- unfreeze ONLY the custom special-token embedding rows ---
    special_ids = sorted(set(tokenizer.additional_special_tokens_ids or []))
    special_ids = [i for i in special_ids if i < vocab_rows]
    embed.weight.requires_grad = True
    row_mask = torch.zeros(vocab_rows, 1, device=device)
    for i in special_ids:
        row_mask[i] = 1.0
    embed.weight.register_hook(lambda g: g * row_mask.to(g.dtype))
    print(f"🎯 Training embedding rows for {len(special_ids)} special tokens "
          f"(ids {special_ids[0]}..{special_ids[-1]} range), base vocab frozen via grad mask.")

    # untied lm_head: train the same rows there too (model can't EMIT a token
    # whose output row is noise, even with good input embeddings)
    head = None
    for mod in model.modules():
        if isinstance(mod, nn.Linear) and mod.out_features == vocab_rows:
            head = mod
            break
    if head is not None and head.weight is not embed.weight:
        head.weight.requires_grad = True
        head.weight.register_hook(lambda g: g * row_mask.to(g.dtype))
        print("🎯 LM head is untied -- training the same rows there as well.")

    n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"   trainable params (incl. masked embeds): {n_train/1e6:.1f}M")

    # --- data ---
    convs = load_dataset(DATASET_PATH)
    print(f"📚 {len(convs)} conversations loaded from {DATASET_PATH}")
    random.shuffle(convs)

    # --- optimizer ---
    groups = [{"params": lora_params, "lr": LR}]
    embed_params = [embed.weight] + ([head.weight] if head is not None and head.weight is not embed.weight else [])
    groups.append({"params": embed_params, "lr": EMBED_LR})

    if OPTIMIZER == "adafactor":
        # relative_step=False + explicit per-group lr so our own cosine
        # schedule (LambdaLR below) still controls the learning rate.
        # No momentum buffer -> optimizer state is ~2 bytes/param instead of
        # AdamW's ~8 bytes/param (no exp_avg, no exp_avg_sq kept at full
        # precision). This is the main lever for fitting 36B on 96GB.
        optim = Adafactor(groups, scale_parameter=False, relative_step=False,
                           warmup_init=False, weight_decay=0.0)
        print("⚙️  Optimizer: Adafactor (relative_step=False, no momentum buffer)")
    elif OPTIMIZER == "adamw":
        optim = torch.optim.AdamW(groups, weight_decay=0.0)
        print("⚙️  Optimizer: AdamW")
    else:
        sys.exit(f"❌ Unknown OPTIMIZER '{OPTIMIZER}' -- use 'adamw' or 'adafactor'")

    total_steps = max(1, (len(convs) * EPOCHS) // (BATCH_SIZE * GRAD_ACCUM))

    def lr_lambda(step):
        if step < WARMUP_STEPS:
            return step / max(1, WARMUP_STEPS)
        prog = (step - WARMUP_STEPS) / max(1, total_steps - WARMUP_STEPS)
        return 0.5 * (1.0 + math.cos(math.pi * min(1.0, prog)))
    sched = torch.optim.lr_scheduler.LambdaLR(optim, lr_lambda)

    loss_fn = nn.CrossEntropyLoss(ignore_index=-100)

    def save_adapter(path):
        state = {n: p.detach().cpu() for n, p in model.named_parameters() if "lora_" in n}
        state["__special_ids__"] = torch.tensor(special_ids)
        state["__embed_rows__"] = embed.weight.detach()[special_ids].cpu()
        if head is not None and head.weight is not embed.weight:
            state["__head_rows__"] = head.weight.detach()[special_ids].cpu()
        torch.save({"lora_r": LORA_R, "lora_alpha": LORA_ALPHA, "state": state}, path)
        print(f"💾 Adapter saved: {path}")

    # --- loop ---
    model.train()
    step, micro, running = 0, 0, 0.0
    t0 = time.time()
    for epoch in range(EPOCHS):
        for conv in convs:
            ex = build_example(tokenizer, conv, MAX_LEN)
            if ex is None:
                continue
            ids, labels = ex
            ids = ids.unsqueeze(0).to(device)
            labels = labels.unsqueeze(0).to(device)

            with torch.autocast(device_type=("cuda" if device == "cuda" else "cpu"),
                                dtype=torch.bfloat16):
                logits = model(ids)
                loss = loss_fn(logits[:, :-1, :].reshape(-1, logits.size(-1)).float(),
                               labels[:, 1:].reshape(-1))
            (loss / GRAD_ACCUM).backward()
            running += loss.item()
            micro += 1

            if micro % GRAD_ACCUM == 0:
                torch.nn.utils.clip_grad_norm_(
                    [p for p in model.parameters() if p.requires_grad], 1.0)
                optim.step()
                sched.step()
                optim.zero_grad(set_to_none=True)
                step += 1
                if step % 10 == 0:
                    avg = running / (10 * GRAD_ACCUM)
                    running = 0.0
                    el = time.time() - t0
                    print(f"epoch {epoch+1} step {step}/{total_steps} "
                          f"loss {avg:.4f} lr {sched.get_last_lr()[0]:.2e} "
                          f"[{el/60:.1f} min]")
                if step % SAVE_EVERY == 0:
                    save_adapter(ADAPTER_OUT)

    save_adapter(ADAPTER_OUT)

    # --- merge ---
    if MERGE_AT_END:
        print("🔀 Merging LoRA into base weights ...")
        model.eval()
        merge_and_unwrap(model)
        merged_sd = {k: v.detach().cpu() for k, v in model.state_dict().items()}
        # drop lambda_ buffers if the original checkpoint didn't carry them
        orig_keys = set(sd.keys())
        merged_sd = {k: v for k, v in merged_sd.items()
                     if k in orig_keys or not k.endswith("lambda_")}
        extra = set(merged_sd.keys()) - orig_keys
        missing2 = orig_keys - set(merged_sd.keys())
        if extra:
            print(f"⚠️ Keys not in original ckpt (kept): {list(extra)[:8]}")
        if missing2:
            print(f"⚠️ Original keys absent in merged (check!): {list(missing2)[:8]}")
        torch.save(merged_sd, MERGED_OUT)
        print(f"✅ Merged checkpoint saved: {MERGED_OUT}")
        print("   Point chat_jirack_8b_tools.py MODEL_PATH at it and test.")


if __name__ == "__main__":
    main()