# ============================================================================== # 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()