| """Stage 2 rational SFT with checkpoint-safe active-vocabulary loss slicing. |
| |
| The Stage-1 embedding and LM-head tensors are never resized or rewritten here. Inputs retain |
| original tokenizer IDs; only the output projection and CE classes are sliced to ACTIVE_TOKEN_IDS. |
| """ |
| from __future__ import annotations |
|
|
| import gc |
| import itertools |
| import json |
| import math |
| import os |
| import shutil |
| import time |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Dict, Iterable, List, Optional |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import DataLoader |
|
|
| from .elevator_pretrainer import attach_lora, make_amp, model_device |
| from .tokenizer_surgeon import (build_active_vocab_ids, render_prompt, render_target, |
| summarize_active_vocab) |
|
|
|
|
| def _memory_snapshot() -> str: |
| rss = None |
| try: |
| import psutil |
| rss = psutil.Process(os.getpid()).memory_info().rss |
| except Exception: |
| pass |
| cgroup = None |
| for candidate in ("/sys/fs/cgroup/memory.current", "/sys/fs/cgroup/memory/usage_in_bytes"): |
| try: |
| cgroup = int(Path(candidate).read_text().strip()) |
| break |
| except Exception: |
| continue |
| fields = [] |
| if rss is not None: |
| fields.append(f"rss={rss / (1024 ** 2):.1f}MiB") |
| if cgroup is not None: |
| fields.append(f"cgroup={cgroup / (1024 ** 2):.1f}MiB") |
| return " ".join(fields) if fields else "memory=n/a" |
|
|
|
|
| def _diag(label: str, started: Optional[float] = None, detail: str = "") -> None: |
| elapsed = f" dt={time.perf_counter() - started:.3f}s" if started is not None else "" |
| extra = f" {detail}" if detail else "" |
| print(f"[SFT-DIAG] {datetime.now().astimezone().isoformat(timespec='seconds')} " |
| f"{label}{elapsed} {_memory_snapshot()}{extra}", flush=True) |
|
|
|
|
| |
| def build_examples(tokenizer, path: Path, max_len: int, im_end_id: int, |
| max_thinking_tokens: Optional[int] = None) -> List[Dict]: |
| """Tokenise every valid SFT record; trim only prompt context, never discard an example.""" |
| examples: List[Dict] = [] |
| modes: Dict[str, int] = {} |
| prompt_trimmed = 0 |
| long_targets = 0 |
| dataset_token_ids = set() |
| with open(path, "r", encoding="utf-8") as handle: |
| for line_no, line in enumerate(handle, 1): |
| if not line.strip(): |
| continue |
| rec = json.loads(line) |
| thinking = rec["thinking"] |
| original_prompt = render_prompt(rec["system"], rec["user"], rec["mode"]) |
| original_target = render_target(thinking, rec["assistant"], rec["mode"]) |
| full_p_ids = tokenizer(original_prompt, add_special_tokens=False)["input_ids"] |
| full_t_ids = tokenizer(original_target, add_special_tokens=False)["input_ids"] |
| |
| dataset_token_ids.update(int(x) for x in full_p_ids) |
| dataset_token_ids.update(int(x) for x in full_t_ids) |
| if max_thinking_tokens is not None and rec["mode"] != "non_thinking": |
| thinking_ids = tokenizer(thinking, add_special_tokens=False)["input_ids"] |
| if len(thinking_ids) > max_thinking_tokens: |
| thinking = tokenizer.decode(thinking_ids[:max_thinking_tokens], skip_special_tokens=False) |
| capped_target = render_target(thinking, rec["assistant"], rec["mode"]) |
| full_t_ids = tokenizer(capped_target, add_special_tokens=False)["input_ids"] |
| dataset_token_ids.update(int(x) for x in full_t_ids) |
| if not full_t_ids or int(full_t_ids[-1]) != int(im_end_id): |
| full_t_ids.append(int(im_end_id)) |
| dataset_token_ids.add(int(im_end_id)) |
|
|
| p_ids, t_ids = list(full_p_ids), list(full_t_ids) |
| if len(p_ids) + len(t_ids) > max_len: |
| prompt_budget = max(0, max_len - len(t_ids)) |
| p_ids = p_ids[-prompt_budget:] if prompt_budget else [] |
| prompt_trimmed += 1 |
| if len(t_ids) > max_len: |
| |
| long_targets += 1 |
| examples.append({"input_ids": p_ids + t_ids, |
| "labels": [-100] * len(p_ids) + t_ids, |
| "source_line": line_no}) |
| modes[rec["mode"]] = modes.get(rec["mode"], 0) + 1 |
|
|
| |
| if examples: |
| examples[0]["_dataset_token_ids"] = sorted(dataset_token_ids) |
| print(f" [sft-data] {len(examples)} examples retained ({', '.join(f'{k}={v}' for k, v in sorted(modes.items()))}); " |
| f"prompt-context-trimmed={prompt_trimmed}, long-targets-kept={long_targets}", flush=True) |
| return examples |
|
|
|
|
| def _collate_sft(rows: List[Dict], pad_id: int) -> Dict[str, torch.Tensor]: |
| longest = max(len(row["input_ids"]) for row in rows) |
| ids = torch.full((len(rows), longest), int(pad_id), dtype=torch.long) |
| labels = torch.full((len(rows), longest), -100, dtype=torch.long) |
| mask = torch.zeros((len(rows), longest), dtype=torch.long) |
| for i, row in enumerate(rows): |
| n = len(row["input_ids"]) |
| ids[i, :n] = torch.as_tensor(row["input_ids"], dtype=torch.long) |
| labels[i, :n] = torch.as_tensor(row["labels"], dtype=torch.long) |
| mask[i, :n] = 1 |
| return {"input_ids": ids, "attention_mask": mask, "labels": labels} |
|
|
|
|
| def _epoch_batch_stream(examples: List[Dict], epochs: int, pad_id: int, device): |
| |
| |
| generator = torch.Generator() |
| generator.manual_seed(11) |
| loader = DataLoader(examples, batch_size=1, shuffle=True, generator=generator, |
| num_workers=0, pin_memory=False, |
| collate_fn=lambda rows: _collate_sft(rows, pad_id)) |
| def stream(): |
| for _epoch in range(epochs): |
| for batch in loader: |
| yield {key: value.to(device, non_blocking=False) for key, value in batch.items()} |
| return stream() |
|
|
|
|
| |
| def _base_model_and_head(model): |
| base = model.get_base_model() if hasattr(model, "get_base_model") else model |
| transformer = getattr(base, "model", None) |
| if transformer is None: |
| raise RuntimeError("active-vocab slicing needs a causal model exposing .model (transformer backbone)") |
| head = base.get_output_embeddings() if hasattr(base, "get_output_embeddings") else getattr(base, "lm_head", None) |
| |
| |
| if head is not None and not hasattr(head, "weight"): |
| saved = getattr(head, "modules_to_save", None) |
| active = getattr(head, "active_adapter", None) |
| active = active[0] if isinstance(active, (list, tuple)) and active else active |
| if saved is not None and active in saved: |
| head = saved[active] |
| elif getattr(head, "original_module", None) is not None: |
| head = head.original_module |
| if head is None or not hasattr(head, "weight"): |
| raise RuntimeError("active-vocab slicing could not locate the output lm_head weight") |
| return base, transformer, head |
|
|
|
|
| def make_stop_weighted_loss(im_end_id: int, stop_weight: float, active_token_ids: List[int], |
| diagnostic=_diag, ring=None): |
| """Use only active lm_head rows and map original IDs to compact CE classes.""" |
| active_ids_cpu = torch.as_tensor(active_token_ids, dtype=torch.long, device="cpu") |
| cached_head = None |
| cached_bias = None |
| vocab_map = None |
| compact_stop_id = None |
| output_rows = None |
|
|
| def loss_fn(model, batch): |
| nonlocal cached_head, cached_bias, vocab_map, compact_stop_id, output_rows |
| base, transformer, head = _base_model_and_head(model) |
| ring_handle = ring or getattr(base, "_rss_ring", None) or getattr(model, "_rss_ring", None) |
| if cached_head is None: |
| started = time.perf_counter() |
| diagnostic("HEAD-SLICE-BEGIN", detail=f"active_ids={active_ids_cpu.numel()}") |
| if ring_handle is not None and "head" not in getattr(ring_handle, "loaded", set()): |
| ring_handle._ensure("head") |
| weight = head.weight |
| output_rows = int(weight.shape[0]) |
| if int(active_ids_cpu[-1]) >= output_rows: |
| raise RuntimeError(f"active token ID {int(active_ids_cpu[-1])} exceeds lm_head rows {output_rows}") |
| index = active_ids_cpu.to(weight.device) |
| with torch.no_grad(): |
| cached_head = weight.index_select(0, index).detach() |
| bias = getattr(head, "bias", None) |
| cached_bias = bias.index_select(0, index).detach() if bias is not None else None |
| vocab_map = torch.full((output_rows,), -1, dtype=torch.long, device=weight.device) |
| vocab_map[index] = torch.arange(index.numel(), dtype=torch.long, device=weight.device) |
| compact_stop_id = int(vocab_map[int(im_end_id)].item()) if int(im_end_id) < output_rows else -1 |
| if ring_handle is not None and getattr(ring_handle, "evict_ends", False): |
| ring_handle._evict("head") |
| if compact_stop_id < 0: |
| raise RuntimeError(f"stop token {im_end_id} is not in ACTIVE_TOKEN_IDS") |
| diagnostic("HEAD-SLICE-DONE", started, detail=f"head_rows={output_rows} active={len(index)}") |
|
|
| t0 = time.perf_counter() |
| diagnostic("FORWARD-BEGIN") |
| outputs = transformer(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], |
| use_cache=False, return_dict=True) |
| hidden = outputs.last_hidden_state if hasattr(outputs, "last_hidden_state") else outputs[0] |
| diagnostic("FORWARD-DONE", t0, detail=f"hidden={tuple(hidden.shape)}") |
|
|
| t0 = time.perf_counter() |
| diagnostic("LOSS-BEGIN") |
| sliced_logits = F.linear(hidden[:, :-1, :], cached_head, cached_bias) |
| labels = batch["labels"][:, 1:].to(device=sliced_logits.device, dtype=torch.long) |
| valid = labels.ne(-100) |
| safe_labels = labels.clamp(min=0, max=output_rows - 1) |
| compact_labels = vocab_map[safe_labels] |
| missing = valid & compact_labels.lt(0) |
| if bool(missing.any()): |
| bad_ids = torch.unique(labels[missing]).detach().cpu().tolist()[:12] |
| raise RuntimeError(f"SFT labels missing from ACTIVE_TOKEN_IDS (first IDs: {bad_ids})") |
| compact_labels = compact_labels.masked_fill(~valid, -100) |
| logits_for_ce = sliced_logits.float() if sliced_logits.dtype in (torch.float16, torch.bfloat16) else sliced_logits |
| per_token = F.cross_entropy(logits_for_ce.reshape(-1, logits_for_ce.shape[-1]), |
| compact_labels.reshape(-1), ignore_index=-100, |
| reduction="none").view_as(compact_labels) |
| weights = valid.to(dtype=per_token.dtype) |
| is_stop = compact_labels.eq(compact_stop_id) & valid |
| weights.mul_(1.0 + (float(stop_weight) - 1.0) * is_stop.to(dtype=weights.dtype)) |
| loss = (per_token * weights).sum() / weights.sum().clamp_min(1.0) |
| diagnostic("LOSS-DONE", t0, detail=f"sliced_logits={tuple(sliced_logits.shape)} loss={float(loss.detach()):.5f}") |
| return loss |
| return loss_fn |
|
|
|
|
| |
| def _train_sft_loop(model, params: List[torch.nn.Parameter], batches: Iterable[Dict], steps: int, |
| lr: float, accum: int, amp, tag: str = "sft") -> float: |
| _diag("OPTIMIZER-INIT-BEGIN", detail=f"trainable_params={sum(p.numel() for p in params):,}") |
| started = time.perf_counter() |
| opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.0, betas=(0.9, 0.95)) |
| warm = max(1, int(0.1 * steps)) |
|
|
| def lr_at(step: int) -> float: |
| if step < warm: |
| return (step + 1) / warm |
| progress = (step - warm) / max(1, steps - warm) |
| return 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * progress)) |
|
|
| scheduler = torch.optim.lr_scheduler.LambdaLR(opt, lr_at) |
| _diag("OPTIMIZER-INIT-DONE", started) |
| model.train() |
| batch_iter = iter(batches) |
| history: List[float] = [] |
| for step in range(steps): |
| group = list(itertools.islice(batch_iter, accum)) |
| if not group: |
| break |
| opt.zero_grad(set_to_none=True) |
| step_loss = 0.0 |
| for micro_idx, batch in enumerate(group, 1): |
| _diag("MICROBATCH-BEGIN", detail=f"step={step + 1}/{steps} micro={micro_idx}/{len(group)}") |
| with amp(): |
| loss = batch["_loss_fn"](model, batch) / len(group) |
| t0 = time.perf_counter() |
| _diag("BACKWARD-BEGIN", detail=f"step={step + 1} micro={micro_idx}") |
| loss.backward() |
| _diag("BACKWARD-DONE", t0, detail=f"step={step + 1} micro={micro_idx}") |
| step_loss += float(loss.detach()) |
| t0 = time.perf_counter() |
| _diag("GRAD-CLIP-BEGIN", detail=f"step={step + 1}") |
| torch.nn.utils.clip_grad_norm_(params, 1.0) |
| _diag("GRAD-CLIP-DONE", t0) |
| t0 = time.perf_counter() |
| _diag("OPTIMIZER-STEP-BEGIN", detail=f"step={step + 1}") |
| opt.step() |
| scheduler.step() |
| opt.zero_grad(set_to_none=True) |
| _diag("OPTIMIZER-STEP-DONE", t0, detail=f"lr={opt.param_groups[0]['lr']:.2e}") |
| history.append(step_loss) |
| print(f" [{tag}] step {step + 1:>4}/{steps} loss {step_loss:.5f} " |
| f"lr {opt.param_groups[0]['lr']:.2e}", flush=True) |
| del opt, scheduler |
| gc.collect() |
| tail = history[len(history) // 2:] or history |
| return sum(tail) / max(len(tail), 1) |
|
|
|
|
| |
| def _atomic_save_stage2(model, tokenizer, final_dir: Path, meta: Dict, active_ids: List[int], ring=None) -> None: |
| final_dir = Path(final_dir) |
| tmp_dir = final_dir.with_name(final_dir.name + ".tmp") |
| backup_dir = final_dir.with_name(final_dir.name + ".previous") |
| _diag("CHECKPOINT-SAVE-BEGIN", detail=f"tmp={tmp_dir}") |
| started = time.perf_counter() |
| if tmp_dir.exists(): |
| shutil.rmtree(tmp_dir, ignore_errors=True) |
| tmp_dir.mkdir(parents=True, exist_ok=True) |
| model.save_pretrained(str(tmp_dir / "adapter") if meta.get("has_adapter") else str(tmp_dir), |
| safe_serialization=True) |
| if ring is not None and meta.get("has_adapter"): |
| ring.export(tmp_dir) |
| tokenizer.save_pretrained(str(tmp_dir)) |
| with open(tmp_dir / "active_vocab_ids.json", "w", encoding="utf-8") as handle: |
| json.dump(active_ids, handle, separators=(",", ":")) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| with open(tmp_dir / "rss_meta.json", "w", encoding="utf-8") as handle: |
| json.dump(meta, handle, ensure_ascii=False, indent=1) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| if backup_dir.exists(): |
| shutil.rmtree(backup_dir, ignore_errors=True) |
| moved_old = False |
| try: |
| if final_dir.exists(): |
| os.replace(final_dir, backup_dir) |
| moved_old = True |
| os.replace(tmp_dir, final_dir) |
| except Exception: |
| if moved_old and not final_dir.exists() and backup_dir.exists(): |
| os.replace(backup_dir, final_dir) |
| raise |
| else: |
| if backup_dir.exists(): |
| shutil.rmtree(backup_dir, ignore_errors=True) |
| _diag("CHECKPOINT-SAVE-DONE", started, detail=f"dir={final_dir}") |
|
|
|
|
| def run_finetune(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, sft_path: Path, |
| final_dir: Path, adapter_dir: Optional[str] = None) -> Dict: |
| tune = cfg["tune"] |
| gpu = plan["tier"] == 1 |
| ring = getattr(model, "_rss_ring", None) |
| max_len = int(tune.get("sft_max_seq_len") or (1024 if gpu else 384)) |
| epochs = int(tune.get("sft_epochs") or 3) |
| lr = float(tune.get("sft_lr") or (2e-4 if gpu else 1e-4)) |
| stop_weight = float(tune.get("stop_token_weight") or 3.0) |
| rank = int(tune.get("lora_rank", 16)) |
| micro_batch = 1 |
| grad_accum = int(tune.get("gradient_accumulation_steps") or (2 if not gpu else 1)) |
| max_thinking_tokens = tune.get("sft_max_thinking_tokens", 5) |
| if max_thinking_tokens is not None: |
| max_thinking_tokens = max(0, int(max_thinking_tokens)) |
|
|
| _diag("LORA-INIT-BEGIN", detail=f"rank={rank} adapter_dir={adapter_dir or 'new'}") |
| started = time.perf_counter() |
| amp, _scaler = make_amp(plan) |
| model.config.use_cache = False |
| if ring is not None: |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) |
| |
| model = attach_lora(model, plan, rank, False, surgery.get("resized", False), |
| adapter_dir=adapter_dir, save_resized_modules=False) |
| for name, parameter in model.named_parameters(): |
| is_lora = "lora_A" in name or "lora_B" in name |
| parameter.requires_grad_(is_lora) |
| params = [p for p in model.parameters() if p.requires_grad] |
| if not params: |
| raise RuntimeError("LoRA initialised but no lora_A/lora_B parameters are trainable") |
| non_lora_trainable = [name for name, p in model.named_parameters() if p.requires_grad and |
| "lora_A" not in name and "lora_B" not in name] |
| if non_lora_trainable: |
| raise RuntimeError(f"non-LoRA trainable parameters detected: {non_lora_trainable[:8]}") |
| trainable_n = sum(p.numel() for p in params) |
| device = model_device(model) |
| _diag("LORA-INIT-DONE", started, detail=f"trainable={trainable_n:,} device={device}") |
|
|
| _diag("SFT-DATA-INIT-BEGIN", detail=f"path={sft_path}") |
| started = time.perf_counter() |
| examples = build_examples(tokenizer, sft_path, max_len, surgery["ids"]["im_end"], |
| max_thinking_tokens=max_thinking_tokens) |
| if not examples: |
| raise RuntimeError("no SFT examples available") |
| dataset_ids = examples[0].pop("_dataset_token_ids", []) |
| _base, _transformer, head = _base_model_and_head(model) |
| configured_rows = int(getattr(getattr(model, "config", None), "vocab_size", 0) or 0) |
| full_vocab_rows = int(head.weight.shape[0]) if head.weight.numel() else max( |
| configured_rows, int(getattr(tokenizer, "vocab_size", len(tokenizer))), len(tokenizer)) |
| active_ids = build_active_vocab_ids(tokenizer, dataset_ids, vocab_rows=full_vocab_rows) |
| vocab_stats = summarize_active_vocab(tokenizer, dataset_ids, active_ids, full_vocab_rows) |
| steps = max(1, math.ceil((len(examples) * epochs) / max(1, grad_accum))) |
| if tune.get("sft_max_steps"): |
| steps = min(steps, int(tune["sft_max_steps"])) |
| _diag("ACTIVE-VOCAB-SLICE-DONE", started, |
| detail=f"full={full_vocab_rows} active={len(active_ids)} inactive={full_vocab_rows-len(active_ids)} " |
| f"dataset_ids={len(dataset_ids)}") |
| print(f" [vocab-slice] active={len(active_ids):,}/{full_vocab_rows:,}; " |
| f"inactive={full_vocab_rows-len(active_ids):,}; " |
| f"ASCII/code={vocab_stats['english_ascii_code_math_kept']:,}; " |
| f"Devanagari={vocab_stats['devanagari_kept']:,}; " |
| f"special={vocab_stats['special_control_kept']:,}; bytes={vocab_stats['byte_fallback_kept']:,}; " |
| f"dataset-only foreign kept={vocab_stats['dataset_only_foreign_kept']:,}; " |
| f"foreign pruned={vocab_stats['foreign_vocab_pruned']:,}; " |
| "input embedding/checkpoint rows unchanged", flush=True) |
|
|
| _diag("DATALOADER-INIT-BEGIN", detail=f"batch_size={micro_batch} workers=0 pin_memory=False") |
| started = time.perf_counter() |
| batches = _epoch_batch_stream(examples, epochs, surgery["ids"]["pad"], device) |
| |
| _diag("DATALOADER-INIT-DONE", started, detail="num_workers=0 pin_memory=False") |
| loss_core = make_stop_weighted_loss(surgery["ids"]["im_end"], stop_weight, active_ids, ring=ring) |
| def loss_with_batch(model_obj, batch): |
| return loss_core(model_obj, batch) |
| |
| def prepared_batches(): |
| for batch in batches: |
| batch["_loss_fn"] = loss_with_batch |
| yield batch |
|
|
| print(f"[Finetune] Stage 2 rational SFT: {steps} steps (epochs={epochs}, micro-batch={micro_batch}, " |
| f"grad-accum={grad_accum}, effective-batch={micro_batch * grad_accum}, lr={lr}, " |
| f"stop-token weight x{stop_weight})", flush=True) |
| loss = _train_sft_loop(model, params, prepared_batches(), steps, lr, grad_accum, amp) |
| if ring is not None: |
| model.gradient_checkpointing_disable() |
|
|
| final_dir = Path(final_dir) |
| meta = {"stage": 2, "base_model": cfg["base_model_repo"], "tier": plan["tier"], |
| "modes": ["thinking", "non_thinking", "max_thinking"], "created": time.time(), |
| "final_loss": loss, **vocab_stats, |
| "active_vocab_file": "active_vocab_ids.json", "gradient_accumulation_steps": grad_accum, |
| "micro_batch_size": micro_batch, "stop_token_weight": stop_weight} |
|
|
| _diag("CHECKPOINT-PREP-BEGIN") |
| if plan["tier"] == 2 and not plan.get("quant") and ring is None: |
| model = model.merge_and_unload() |
| model.to(torch.bfloat16) |
| meta.update({"has_adapter": False, "adapter_dir": None, "weights": "merged-full"}) |
| else: |
| meta.update({"has_adapter": True, "adapter_dir": "adapter", |
| "weights": "ring-base+adapter" if ring is not None else "adapter-only", |
| "base_is_self": ring is not None}) |
| _atomic_save_stage2(model, tokenizer, final_dir, meta, active_ids, ring=ring) |
| gc.collect() |
| return {"final_dir": str(final_dir), "loss": loss, "meta": meta, |
| "active_vocab_size": len(active_ids), "full_vocab_size": full_vocab_rows} |
|
|
|
|
