train_toolace_toolcalling_lora_ultra.py
19.9 KB · 471 lines · python Raw
1 #%%writefile train_toolace_lora_ultra.py
2 # ==============================================================================
3 # JiRack Ultra ToolACE LoRA SFT + merge (single script, all four sizes)
4 # COPYRIGHT (c) 2026 Konstantin Vladimirovich Grabko.
5 #
6 # Adapted from the JiRackPrecision_8b ToolACE LoRA script. One file covers
7 # Ultra 1B / 7B / 14B / 32B -- set SIZE below, everything else follows from
8 # the SIZES table.
9 #
10 # What it does (unchanged from the 8B original):
11 # 1. Loads JiRackTransformer + your .pt checkpoint.
12 # 2. Freezes everything; injects LoRA (A/B low-rank pairs) into every
13 # nn.Linear except the LM head (out_features == vocab_size).
14 # 3. ALSO unfreezes the embedding rows of the JiRack special tokens --
15 # those rows are untrained padded slots right now; the model can't emit
16 # <|tool_call_start|> etc. until they're trained. A gradient hook zeroes
17 # grads for all other rows, so the base vocab embeddings stay untouched.
18 # 4. Trains with assistant-only loss masking, bf16 autocast, grad accum.
19 # 5. Saves the LoRA adapter alone + OPTIONAL merged checkpoint whose
20 # state_dict keys match the input .pt exactly.
21 #
22 # ============================ ULTRA-SPECIFIC CHANGES ==========================
23 # [U-A] BitLinear IS an nn.Linear subclass in every Ultra file, so isinstance()
24 # picks it up and LoRA wraps it -- which is what we want. But it also
25 # means the wrapped base runs BitLinear.forward, i.e. the quantization
26 # math, on every call. At LAMBDA=0.0 the result is mathematically
27 # identical to plain F.linear (lam=0 => w_effective=w, x_effective=x),
28 # but the BitLinear fast path only triggers in EVAL mode, so during
29 # training you pay for the quant math with no effect. Tolerable; if you
30 # want it gone, set LAMBDA=0.0 and patch BitLinear's fast-path condition
31 # to also fire while training.
32 # [U-B] FREEZE_8BIT IS DISABLED for Ultra. The 8B original swapped frozen
33 # nn.Linear for bnb.nn.Linear8bitLt -- on Ultra that would REPLACE your
34 # BitLinear modules with plain bnb layers, destroying the lambda_ buffers
35 # and the ternary path, and the merged state_dict keys would no longer
36 # match your checkpoint. Also, the original's merge_into_base() called
37 # bnb.functional.dequantize_4bit on an 8-bit layer, which is the wrong
38 # function anyway. If you need the memory, use adafactor + shorter
39 # MAX_LEN, or shard -- not this.
40 # [U-C] Tokenizer defaults to CMSManhattan/JiRackPrecisionTokenizer (the
41 # published JiRack tokenizer), with a hard assert that it FITS the padded
42 # matrix. Never resize: 151,779 < 151,936 (1B) and < 152,064 (7/14/32B),
43 # so a resize would SHRINK and corrupt the embedding matrix.
44 # [U-D] Per-size memory defaults in the SIZES table (MAX_LEN, GRAD_ACCUM).
45 # ==============================================================================
46
47 import json
48 import math
49 import os
50 import random
51 import sys
52 import time
53
54 import torch
55 import torch.nn as nn
56 from transformers import AutoTokenizer
57 from transformers.optimization import Adafactor
58
59 sys.path.append(os.getcwd())
60
61 # ========================= PICK YOUR SIZE =========================
62 SIZE = "1b" # "1b" | "7b" | "14b" | "32b"
63 # ==================================================================
64
65 # NOTE the module names -- they are NOT uniform in your repo:
66 # 1B -> JiRackTernaryUltra_1b.py
67 # 7B -> JiRackTernaryUltra7b.py <-- no underscore before "7b"!
68 # 14B -> JiRackTernaryUltra_14b.py
69 # 32B -> JiRackTernaryUltra_32b.py
70 # If you rename any of them, fix the "module" field below.
71 SIZES = {
72 "1b": {
73 "module": "JiRackTernaryUltra_1b",
74 "vocab": 151936,
75 "model_path": "/mnt/nfs_clientshare/JiRackUltra_1b/model.pt",
76 "adapter": "/mnt/nfs_clientshare/JiRackUltra_1b/toolace_lora_adapter.pt",
77 "merged": "/mnt/nfs_clientshare/JiRackUltra_1b/ultra1b_toolace.pt",
78 "max_len": 2048,
79 "grad_accum": 8,
80 },
81 "7b": {
82 "module": "JiRackTernaryUltra7b",
83 "vocab": 152064,
84 "model_path": "/mnt/nfs_clientshare/JiRackUltra_7b/model.pt",
85 "adapter": "/mnt/nfs_clientshare/JiRackUltra_7b/toolace_lora_adapter.pt",
86 "merged": "/mnt/nfs_clientshare/JiRackUltra_7b/ultra7b_toolace.pt",
87 "max_len": 2048,
88 "grad_accum": 16,
89 },
90 "14b": {
91 "module": "JiRackTernaryUltra_14b",
92 "vocab": 152064,
93 "model_path": "/mnt/nfs_clientshare/JiRackUltra_14b/model.pt",
94 "adapter": "/mnt/nfs_clientshare/JiRackUltra_14b/toolace_lora_adapter.pt",
95 "merged": "/mnt/nfs_clientshare/JiRackUltra_14b/ultra14b_toolace.pt",
96 "max_len": 1024, # [U-D] halve the context to fit
97 "grad_accum": 16,
98 },
99 "32b": {
100 "module": "JiRackTernaryUltra_32b",
101 "vocab": 152064,
102 "model_path": "/mnt/nfs_clientshare/JiRackUltra_32b/model.pt",
103 "adapter": "/mnt/nfs_clientshare/JiRackUltra_32b/toolace_lora_adapter.pt",
104 "merged": "/mnt/nfs_clientshare/JiRackUltra_32b/ultra32b_toolace.pt",
105 "max_len": 1024,
106 "grad_accum": 32,
107 },
108 }
109
110 if SIZE not in SIZES:
111 sys.exit(f"❌ SIZE must be one of {list(SIZES)}, got '{SIZE}'")
112 CFG = SIZES[SIZE]
113
114 _mod = __import__(CFG["module"], fromlist=["JiRackTransformer", "JiRackConfig"])
115 JiRackTransformer = _mod.JiRackTransformer
116 JiRackConfig = _mod.JiRackConfig
117
118 # ========================= EDIT THESE =========================
119 MODEL_PATH = CFG["model_path"]
120 TOKENIZER_DIR = "CMSManhattan/JiRackPrecisionTokenizer" # [U-C] HF repo or local dir
121 DATASET_PATH = "/mnt/nfs_clientshare/datasets/toolace_sft_jirack_precision.jsonl"
122 ADAPTER_OUT = CFG["adapter"]
123 MERGED_OUT = CFG["merged"]
124
125 # LoRA
126 LORA_R = 16
127 LORA_ALPHA = 32
128 LORA_DROPOUT = 0.05
129
130 # Training
131 EPOCHS = 2
132 LR = 2e-4 # LoRA params
133 EMBED_LR = 5e-5 # new-token embedding rows (gentler)
134 BATCH_SIZE = 1
135 GRAD_ACCUM = CFG["grad_accum"]
136 MAX_LEN = CFG["max_len"]
137 WARMUP_STEPS = 50
138 SEED = 42
139 LAMBDA = 0.0 # 0.0 = full-precision training (recommended:
140 # you're teaching tool-call FORMAT, not
141 # doing QAT -- run the ternarization QAT
142 # scripts separately, AFTER this merge)
143 SAVE_EVERY = 500 # optimizer steps between adapter checkpoints
144 MERGE_AT_END = True
145 OPTIMIZER = "adafactor" # "adamw" or "adafactor"
146 # adafactor: ~2 bytes/param optimizer state
147 # vs AdamW's ~8 -- matters at 14B/32B.
148 # [U-B] FREEZE_8BIT removed on purpose -- see header.
149 # ================================================================
150
151
152 # ------------------------------ LoRA machinery ------------------------------
153
154 class LoRALinear(nn.Module):
155 """Wraps a frozen nn.Linear (or BitLinear); adds trainable low-rank A/B."""
156
157 def __init__(self, base: nn.Linear, r: int, alpha: int, dropout: float):
158 super().__init__()
159 self.base = base
160 self.r = r
161 self.scale = alpha / r
162 self.lora_A = nn.Parameter(torch.zeros(r, base.in_features))
163 self.lora_B = nn.Parameter(torch.zeros(base.out_features, r))
164 nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
165 # B starts at zero -> identity behavior at step 0
166 self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
167
168 def forward(self, x):
169 out = self.base(x) # [U-A] BitLinear.forward
170 lx = self.dropout(x).to(self.lora_A.dtype)
171 out = out + (lx @ self.lora_A.T @ self.lora_B.T) * self.scale
172 return out
173
174 @torch.no_grad()
175 def merge_into_base(self):
176 """Fold the LoRA delta into the base weight, in place. The base stays
177 the SAME module object (BitLinear stays BitLinear), so lambda_ buffers
178 and state_dict keys survive untouched."""
179 delta = (self.lora_B.float() @ self.lora_A.float()) * self.scale
180 self.base.weight.data += delta.to(self.base.weight.dtype)
181
182
183 def inject_lora(model, vocab_size):
184 """Replace every nn.Linear (except the vocab-sized head) with LoRALinear.
185 BitLinear subclasses nn.Linear, so the whole ternary backbone gets wrapped
186 -- intended. Embeddings are nn.Embedding, not nn.Linear, so they're skipped
187 here and handled separately by the row-mask logic."""
188 wrapped = []
189 for parent_name, parent in list(model.named_modules()):
190 for child_name, child in list(parent.named_children()):
191 if isinstance(child, LoRALinear):
192 continue
193 if isinstance(child, nn.Linear) and child.out_features != vocab_size:
194 setattr(parent, child_name,
195 LoRALinear(child, LORA_R, LORA_ALPHA, LORA_DROPOUT))
196 full = f"{parent_name}.{child_name}" if parent_name else child_name
197 wrapped.append(full)
198 return wrapped
199
200
201 def merge_and_unwrap(model):
202 """Fold LoRA into base weights and restore the original modules, so
203 state_dict() keys match the original checkpoint exactly."""
204 for parent_name, parent in list(model.named_modules()):
205 for child_name, child in list(parent.named_children()):
206 if isinstance(child, LoRALinear):
207 child.merge_into_base()
208 setattr(parent, child_name, child.base)
209
210
211 # ------------------------------ Dataset ------------------------------
212
213 def load_dataset(path):
214 convs = []
215 with open(path) as f:
216 for line in f:
217 line = line.strip()
218 if not line:
219 continue
220 obj = json.loads(line)
221 msgs = obj.get("messages", obj)
222 if isinstance(msgs, list) and any(m.get("role") == "assistant" for m in msgs):
223 convs.append(msgs)
224 return convs
225
226
227 def build_example(tokenizer, messages, max_len):
228 """Tokenize a conversation with assistant-only labels.
229 Incremental templating: token span of message i = template(msgs[:i+1]) minus
230 template(msgs[:i]). Labels = ids inside assistant spans, else -100."""
231 ids, labels = [], []
232 prev = []
233 prev_len = 0
234 for m in messages:
235 prev.append(m)
236 cur = tokenizer.apply_chat_template(prev, tokenize=True,
237 add_generation_prompt=False)
238 span = cur[prev_len:]
239 if m["role"] == "assistant":
240 labels.extend(span)
241 else:
242 labels.extend([-100] * len(span))
243 ids = cur
244 prev_len = len(cur)
245 if len(ids) >= max_len:
246 break
247 ids = ids[:max_len]
248 labels = labels[:max_len]
249 if all(l == -100 for l in labels):
250 return None
251 return torch.tensor(ids), torch.tensor(labels)
252
253
254 # ------------------------------ Training ------------------------------
255
256 def main():
257 random.seed(SEED)
258 torch.manual_seed(SEED)
259 device = "cuda" if torch.cuda.is_available() else "cpu"
260 print(f"🚀 JiRack Ultra {SIZE.upper()} ToolACE LoRA | Device: {device.upper()}")
261 print(f"⚙️ module={CFG['module']} optimizer={OPTIMIZER} "
262 f"MAX_LEN={MAX_LEN} GRAD_ACCUM={GRAD_ACCUM}")
263
264 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_DIR)
265
266 # --- model ---
267 config = JiRackConfig()
268 # [U-C] the config's vocab must match what this script expects for this size
269 assert config.vocab_size == CFG["vocab"], (
270 f"{CFG['module']}.JiRackConfig has vocab_size={config.vocab_size} but "
271 f"SIZE='{SIZE}' expects {CFG['vocab']} -- wrong module for this size?"
272 )
273 # [U-C] tokenizer must FIT the padded matrix; never resize
274 assert len(tokenizer) <= config.vocab_size, (
275 f"tokenizer ({len(tokenizer)}) > padded matrix ({config.vocab_size}) — "
276 f"do NOT resize_token_embeddings, fix the tokenizer instead"
277 )
278 print(f"✅ Tokenizer fits: {len(tokenizer)} <= {config.vocab_size}")
279
280 model = JiRackTransformer(config, use_checkpoint=True) # activation ckpt on
281 print(f"📥 Loading {MODEL_PATH} ...")
282 ckpt = torch.load(MODEL_PATH, map_location="cpu", weights_only=False)
283 sd = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
284 missing, unexpected = model.load_state_dict(sd, strict=False)
285 real_missing = [k for k in missing if not k.endswith("lambda_")]
286 if real_missing:
287 print(f"⚠️ Missing keys: {real_missing[:10]}")
288 if unexpected:
289 print(f"⚠️ Unexpected keys: {list(unexpected)[:10]}")
290 model = model.to(dtype=torch.bfloat16, device=device)
291 model.set_lambda(LAMBDA)
292
293 # find embedding module + vocab size
294 embed = None
295 for mod in model.modules():
296 if isinstance(mod, nn.Embedding):
297 embed = mod
298 break
299 if embed is None:
300 sys.exit("❌ No nn.Embedding found in model")
301 vocab_rows = embed.weight.shape[0]
302 print(f" embedding rows: {vocab_rows}")
303
304 # --- find the LM head BEFORE wrapping (after injection it'd be hidden) ---
305 head = None
306 for mod in model.modules():
307 if isinstance(mod, nn.Linear) and mod.out_features == vocab_rows:
308 head = mod
309 break
310
311 # --- freeze all, inject LoRA ---
312 for p in model.parameters():
313 p.requires_grad = False
314 wrapped = inject_lora(model, vocab_rows)
315 model = model.to(device)
316 print(f"🧩 LoRA injected into {len(wrapped)} Linear layers "
317 f"(r={LORA_R}, alpha={LORA_ALPHA})")
318
319 lora_params = [p for n, p in model.named_parameters() if "lora_" in n]
320 for p in lora_params:
321 p.requires_grad = True
322
323 # --- unfreeze ONLY the JiRack special-token embedding rows ---
324 special_ids = sorted(set(tokenizer.additional_special_tokens_ids or []))
325 special_ids = [i for i in special_ids if i < vocab_rows]
326 if not special_ids:
327 print("⚠️ No additional_special_tokens found in the tokenizer -- "
328 "training LoRA only, no embedding rows. If you expected the "
329 "JiRack tool-call/robotics tags here, check the tokenizer repo.")
330 embed.weight.requires_grad = True
331 row_mask = torch.zeros(vocab_rows, 1, device=device)
332 for i in special_ids:
333 row_mask[i] = 1.0
334 embed.weight.register_hook(lambda g: g * row_mask.to(g.dtype))
335 if special_ids:
336 print(f"🎯 Training embedding rows for {len(special_ids)} special tokens "
337 f"(ids {special_ids[0]}..{special_ids[-1]}), base vocab frozen "
338 f"via grad mask.")
339
340 # untied lm_head: train the same rows there too (the model can't EMIT a
341 # token whose output row is noise, even with good input embeddings)
342 if head is not None and head.weight is not embed.weight:
343 head.weight.requires_grad = True
344 head.weight.register_hook(lambda g: g * row_mask.to(g.dtype))
345 print("🎯 LM head is untied -- training the same rows there as well.")
346 elif head is None:
347 print("⚠️ No vocab-sized Linear found -- lm_head not trained.")
348
349 n_train = sum(p.numel() for p in model.parameters() if p.requires_grad)
350 print(f" trainable params (incl. masked embeds): {n_train/1e6:.1f}M")
351
352 # --- data ---
353 convs = load_dataset(DATASET_PATH)
354 print(f"📚 {len(convs)} conversations loaded from {DATASET_PATH}")
355 random.shuffle(convs)
356
357 # --- optimizer ---
358 groups = [{"params": lora_params, "lr": LR}]
359 embed_params = [embed.weight]
360 if head is not None and head.weight is not embed.weight:
361 embed_params.append(head.weight)
362 groups.append({"params": embed_params, "lr": EMBED_LR})
363
364 if OPTIMIZER == "adafactor":
365 # relative_step=False + explicit per-group lr so our own cosine
366 # schedule (LambdaLR below) still controls the learning rate.
367 optim = Adafactor(groups, scale_parameter=False, relative_step=False,
368 warmup_init=False, weight_decay=0.0)
369 print("⚙️ Optimizer: Adafactor (relative_step=False, no momentum buffer)")
370 elif OPTIMIZER == "adamw":
371 optim = torch.optim.AdamW(groups, weight_decay=0.0)
372 print("⚙️ Optimizer: AdamW")
373 else:
374 sys.exit(f"❌ Unknown OPTIMIZER '{OPTIMIZER}' -- use 'adamw' or 'adafactor'")
375
376 total_steps = max(1, (len(convs) * EPOCHS) // (BATCH_SIZE * GRAD_ACCUM))
377
378 def lr_lambda(step):
379 if step < WARMUP_STEPS:
380 return step / max(1, WARMUP_STEPS)
381 prog = (step - WARMUP_STEPS) / max(1, total_steps - WARMUP_STEPS)
382 return 0.5 * (1.0 + math.cos(math.pi * min(1.0, prog)))
383 sched = torch.optim.lr_scheduler.LambdaLR(optim, lr_lambda)
384
385 loss_fn = nn.CrossEntropyLoss(ignore_index=-100)
386
387 def save_adapter(path):
388 state = {n: p.detach().cpu() for n, p in model.named_parameters()
389 if "lora_" in n}
390 state["__special_ids__"] = torch.tensor(special_ids)
391 if special_ids:
392 state["__embed_rows__"] = embed.weight.detach()[special_ids].cpu()
393 if head is not None and head.weight is not embed.weight:
394 state["__head_rows__"] = head.weight.detach()[special_ids].cpu()
395 torch.save({"size": SIZE, "lora_r": LORA_R, "lora_alpha": LORA_ALPHA,
396 "state": state}, path)
397 print(f"💾 Adapter saved: {path}")
398
399 # --- loop ---
400 model.train()
401 step, micro, running = 0, 0, 0.0
402 t0 = time.time()
403 for epoch in range(EPOCHS):
404 for conv in convs:
405 ex = build_example(tokenizer, conv, MAX_LEN)
406 if ex is None:
407 continue
408 ids, labels = ex
409 ids = ids.unsqueeze(0).to(device)
410 labels = labels.unsqueeze(0).to(device)
411
412 with torch.autocast(device_type=("cuda" if device == "cuda" else "cpu"),
413 dtype=torch.bfloat16):
414 logits = model(ids)
415 loss = loss_fn(
416 logits[:, :-1, :].reshape(-1, logits.size(-1)).float(),
417 labels[:, 1:].reshape(-1))
418
419 if torch.isnan(loss) or torch.isinf(loss):
420 print(f"⚠️ NaN/Inf loss at micro-step {micro} — example skipped")
421 optim.zero_grad(set_to_none=True)
422 micro += 1
423 continue
424
425 (loss / GRAD_ACCUM).backward()
426 running += loss.item()
427 micro += 1
428
429 if micro % GRAD_ACCUM == 0:
430 torch.nn.utils.clip_grad_norm_(
431 [p for p in model.parameters() if p.requires_grad], 1.0)
432 optim.step()
433 sched.step()
434 optim.zero_grad(set_to_none=True)
435 step += 1
436 if step % 10 == 0:
437 avg = running / (10 * GRAD_ACCUM)
438 running = 0.0
439 el = time.time() - t0
440 print(f"epoch {epoch+1} step {step}/{total_steps} "
441 f"loss {avg:.4f} lr {sched.get_last_lr()[0]:.2e} "
442 f"[{el/60:.1f} min]")
443 if step % SAVE_EVERY == 0:
444 save_adapter(ADAPTER_OUT)
445
446 save_adapter(ADAPTER_OUT)
447
448 # --- merge ---
449 if MERGE_AT_END:
450 print("🔀 Merging LoRA into base weights ...")
451 model.eval()
452 merge_and_unwrap(model)
453 merged_sd = {k: v.detach().cpu() for k, v in model.state_dict().items()}
454 # drop lambda_ buffers if the original checkpoint didn't carry them
455 orig_keys = set(sd.keys())
456 merged_sd = {k: v for k, v in merged_sd.items()
457 if k in orig_keys or not k.endswith("lambda_")}
458 extra = set(merged_sd.keys()) - orig_keys
459 missing2 = orig_keys - set(merged_sd.keys())
460 if extra:
461 print(f"⚠️ Keys not in original ckpt (kept): {list(extra)[:8]}")
462 if missing2:
463 print(f"⚠️ Original keys absent in merged (check!): {list(missing2)[:8]}")
464 torch.save(merged_sd, MERGED_OUT)
465 print(f"✅ Merged checkpoint saved: {MERGED_OUT}")
466 print(f" Next: point your chat script at it, verify tool tags are "
467 f"emitted, THEN run the ternarization QAT script for {SIZE}.")
468
469
470 if __name__ == "__main__":
471 main()