| """Stage 1 continual pre-training (PRIVATE to Training_Module). |
| |
| CPU tiers (2 and 3): Layer-by-Layer ELEVATOR unfreezing - only 2 transformer layers are trainable |
| at any moment (promoted to fp32), all the others stay frozen. The window rides |
| from layer 0 up to layer N-1. In tier 3 the window layers are pinned in the ring |
| loader and persisted to disk when the elevator leaves them. |
| GPU tier (1) : rsLoRA on the 4-bit NF4 base model. |
| |
| Data: ./cleaned_data/pretrain_clean.jsonl -> ~75% repo text chunks + ~25% rational/math/identity chunks |
| """ |
| from __future__ import annotations |
|
|
| import contextlib |
| import gc |
| import json |
| import math |
| import random |
| import time |
| from pathlib import Path |
| from typing import Callable, Dict, Iterator, List, Optional |
|
|
| import torch |
|
|
| LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] |
|
|
|
|
| |
| def model_device(model) -> torch.device: |
| try: |
| return next(model.parameters()).device |
| except StopIteration: |
| return torch.device("cpu") |
|
|
|
|
| def make_amp(plan: Dict): |
| """Return (autocast_factory, grad_scaler_or_None) for the plan.""" |
| dtype = plan["dtype"] |
| if plan["tier"] == 1 and torch.cuda.is_available(): |
| dt = getattr(torch, dtype) |
| scaler = None |
| if dt == torch.float16: |
| try: |
| scaler = torch.amp.GradScaler("cuda") |
| except Exception: |
| scaler = torch.cuda.amp.GradScaler() |
| return (lambda: torch.autocast("cuda", dtype=dt)), scaler |
| if dtype == "bfloat16": |
| return (lambda: torch.autocast("cpu", dtype=torch.bfloat16)), None |
| return (lambda: contextlib.nullcontext()), None |
|
|
|
|
| def find_layers(model): |
| base = model.get_base_model() if hasattr(model, "get_base_model") else model |
| return base.model.layers |
|
|
|
|
| def train_loop(model, params: List[torch.nn.Parameter], batches: Iterator[Dict], loss_fn: Callable, |
| steps: int, lr: float, accum: int, amp, scaler, tag: str, log_every: int = 1) -> float: |
| """Generic AdamW loop with warmup + cosine decay, grad clipping and ETA logging.""" |
| 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)) |
|
|
| sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_at) |
| model.train() |
| started = time.time() |
| history: List[float] = [] |
| for step in range(steps): |
| total = 0.0 |
| for _ in range(accum): |
| batch = next(batches) |
| with amp(): |
| loss = loss_fn(model, batch) / accum |
| if scaler is not None: |
| scaler.scale(loss).backward() |
| else: |
| loss.backward() |
| total += float(loss.detach()) |
| if scaler is not None: |
| scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(params, 1.0) |
| if scaler is not None: |
| scaler.step(opt) |
| scaler.update() |
| else: |
| opt.step() |
| sched.step() |
| opt.zero_grad(set_to_none=True) |
| history.append(total) |
| if (step + 1) % log_every == 0 or step == steps - 1: |
| elapsed = time.time() - started |
| eta = elapsed / (step + 1) * (steps - step - 1) |
| print(f" [{tag}] step {step + 1:>4}/{steps} loss {total:.4f} lr {opt.param_groups[0]['lr']:.2e} " |
| f"elapsed {elapsed:5.0f}s eta {eta:5.0f}s") |
| del opt, sched |
| gc.collect() |
| tail = history[len(history) // 2:] or history |
| return sum(tail) / max(len(tail), 1) |
|
|
|
|
| |
| def build_chunks(tokenizer, path: Path, seq_len: int, max_chunks: int, conv_ratio: float, |
| seed: int = 1337, max_chunks_per_doc: int = 8) -> List[List[int]]: |
| """Tokenise pretrain_clean.jsonl into fixed windows; chat.jsonl chunks make up ~conv_ratio.""" |
| eos = tokenizer.convert_tokens_to_ids("<|endoftext|>") |
| if eos is None or eos == tokenizer.unk_token_id: |
| eos = tokenizer.eos_token_id |
| n_chat_target = int(round(max_chunks * conv_ratio)) |
| n_repo_target = max_chunks - n_chat_target |
| repo_chunks: List[List[int]] = [] |
| chat_chunks: List[List[int]] = [] |
| buffers = {"repo": [], "chat": []} |
| with open(path, "r", encoding="utf-8") as handle: |
| for line in handle: |
| rec = json.loads(line) |
| kind = "chat" if rec.get("source") == "chat.jsonl" else "repo" |
| bucket = chat_chunks if kind == "chat" else repo_chunks |
| target = n_chat_target if kind == "chat" else n_repo_target |
| if len(bucket) >= max(target, 1) * (3 if kind == "chat" else 1): |
| if len(repo_chunks) >= n_repo_target and len(chat_chunks) >= n_chat_target: |
| break |
| continue |
| ids = tokenizer(rec["text"], add_special_tokens=False, truncation=False)["input_ids"] |
| ids = ids[: seq_len * max_chunks_per_doc] + [eos] |
| buf = buffers[kind] |
| buf.extend(ids) |
| while len(buf) >= seq_len: |
| bucket.append(buf[:seq_len]) |
| del buf[:seq_len] |
| rng = random.Random(seed) |
| rng.shuffle(repo_chunks) |
| rng.shuffle(chat_chunks) |
| repo_part = repo_chunks[:n_repo_target] |
| if not chat_chunks: |
| chat_part: List[List[int]] = [] |
| elif not repo_part: |
| chat_part = [chat_chunks[i % len(chat_chunks)] for i in range(max_chunks)] |
| else: |
| chat_part = [chat_chunks[i % len(chat_chunks)] for i in range(n_chat_target)] |
| mixed = repo_part + chat_part |
| rng.shuffle(mixed) |
| share = 100.0 * len(chat_part) / max(len(mixed), 1) |
| print(f" [pretrain-data] {len(mixed)} chunks of {seq_len} tokens " |
| f"(repo={len(repo_part)}, chat={len(chat_part)} -> {share:.0f}% conversation share)") |
| return mixed |
|
|
|
|
| def batch_stream(chunks: List[List[int]], batch_size: int, device, seed: int = 7) -> Iterator[Dict]: |
| rng = random.Random(seed) |
| order: List[int] = [] |
| while True: |
| if len(order) < batch_size: |
| fresh = list(range(len(chunks))) |
| rng.shuffle(fresh) |
| order.extend(fresh) |
| take, order = order[:batch_size], order[batch_size:] |
| ids = torch.tensor([chunks[i] for i in take], dtype=torch.long, device=device) |
| yield {"input_ids": ids, "attention_mask": torch.ones_like(ids), "labels": ids.clone()} |
|
|
|
|
| def lm_loss(model, batch): |
| out = model(input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"]) |
| return out.loss |
|
|
|
|
| |
| def attach_lora(model, plan: Dict, rank: int, rslora: bool, resized: bool, adapter_dir: Optional[str] = None, |
| save_resized_modules: bool = True): |
| from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training |
| if plan.get("quant") == "nf4": |
| model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True, |
| gradient_checkpointing_kwargs={"use_reentrant": False}) |
| if adapter_dir: |
| print(f" [lora] continuing adapter from {adapter_dir}") |
| return PeftModel.from_pretrained(model, adapter_dir, is_trainable=True) |
| cfg = LoraConfig(r=rank, lora_alpha=rank if rslora else 2 * rank, lora_dropout=0.05, |
| target_modules=LORA_TARGETS, bias="none", task_type="CAUSAL_LM", use_rslora=rslora, |
| modules_to_save=["embed_tokens", "lm_head"] if (resized and save_resized_modules) else None) |
| model = get_peft_model(model, cfg) |
| trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| print(f" [lora] {'rsLoRA' if rslora else 'LoRA'} r={rank}: {trainable / 1e6:.2f}M trainable params") |
| return model |
|
|
|
|
| def save_stage1_checkpoint(model, tokenizer, out_dir: Path, cfg: Dict, kind: str) -> Path: |
| """Write Stage 1 to a sibling temp directory and publish it without deleting an existing checkpoint.""" |
| out_dir = Path(out_dir) |
| tmp_dir = out_dir.with_name(out_dir.name + ".tmp") |
| if out_dir.exists(): |
| raise FileExistsError(f"refusing to overwrite protected Stage-1 checkpoint: {out_dir}") |
| if tmp_dir.exists(): |
| import shutil |
| shutil.rmtree(tmp_dir, ignore_errors=True) |
| tmp_dir.mkdir(parents=True, exist_ok=True) |
| ring = getattr(model, "_rss_ring", None) |
| if kind == "adapter": |
| model.save_pretrained(str(tmp_dir / "adapter"), safe_serialization=True) |
| elif ring is not None: |
| ring.export(tmp_dir) |
| else: |
| model.to(torch.bfloat16) |
| model.save_pretrained(str(tmp_dir), safe_serialization=True) |
| tokenizer.save_pretrained(str(tmp_dir)) |
| (tmp_dir / "rss_meta.json").write_text(json.dumps({ |
| "stage": 1, "kind": kind, "base_model": cfg["base_model_repo"], |
| "adapter_dir": "adapter" if kind == "adapter" else None, "created": time.time()}, indent=1), encoding="utf-8") |
| out_dir.parent.mkdir(parents=True, exist_ok=True) |
| import os |
| os.replace(tmp_dir, out_dir) |
| print(f" [stage1] checkpoint atomically saved -> {out_dir} ({kind})", flush=True) |
| return out_dir |
|
|
|
|
| |
| def run_pretrain(model, tokenizer, plan: Dict, cfg: Dict, surgery: Dict, pretrain_path: Path, |
| checkpoint_dir: Path) -> Dict: |
| tune = cfg["tune"] |
| gpu = plan["tier"] == 1 |
| seq_len = tune.get("pretrain_seq_len") or (512 if gpu else 192) |
| max_chunks = tune.get("pretrain_max_chunks") or (2000 if gpu else 96) |
| batch_size = tune.get("pretrain_batch_size") or (4 if gpu else 2) |
| lr = tune.get("pretrain_lr") or (2e-4 if gpu else 2e-5) |
| chunks = build_chunks(tokenizer, pretrain_path, seq_len, max_chunks, cfg.get("conv_ratio", 0.25)) |
| if not chunks: |
| raise RuntimeError("no pre-training chunks could be built from pretrain_clean.jsonl") |
| amp, scaler = make_amp(plan) |
| model.config.use_cache = False |
| device = model_device(model) |
| ring = getattr(model, "_rss_ring", None) |
| batches = batch_stream(chunks, batch_size, device) |
|
|
| if gpu: |
| model = attach_lora(model, plan, int(tune.get("lora_rank", 16)), True, surgery.get("resized", False)) |
| steps = int(tune.get("pretrain_steps") or min(300, max(20, len(chunks) // batch_size))) |
| params = [p for p in model.parameters() if p.requires_grad] |
| print(f"[Pretrain] Stage 1 (rsLoRA, 4-bit) for {steps} steps") |
| loss = train_loop(model, params, batches, lm_loss, steps, lr, 1, amp, scaler, "cpt-lora", log_every=5) |
| save_stage1_checkpoint(model, tokenizer, checkpoint_dir, cfg, "adapter") |
| return {"kind": "adapter", "loss": loss, "checkpoint": str(checkpoint_dir)} |
|
|
| |
| layers = find_layers(model) |
| total_layers = len(layers) |
| window = 2 |
| steps_per_window = int(tune.get("pretrain_steps_per_window") or 4) |
| base_dtype = next(model.parameters()).dtype |
| for p in model.parameters(): |
| p.requires_grad_(False) |
| if ring is not None: |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) |
| print(f"[Pretrain] Stage 1 elevator: {total_layers} layers, window={window}, {steps_per_window} steps/window, " |
| f"seq={seq_len} batch={batch_size} lr={lr}") |
| window_losses = [] |
| for start in range(0, total_layers, window): |
| idxs = list(range(start, min(total_layers, start + window))) |
| if ring is not None: |
| ring.pin(idxs) |
| params: List[torch.nn.Parameter] = [] |
| for i in idxs: |
| if base_dtype != torch.float32: |
| layers[i].to(torch.float32) |
| for p in layers[i].parameters(): |
| p.requires_grad_(True) |
| params.append(p) |
| print(f" [elevator] floor {idxs[0]}-{idxs[-1]} unfrozen " |
| f"({sum(p.numel() for p in params) / 1e6:.1f}M params)" + (f" | {ring.describe()}" if ring else "")) |
| loss = train_loop(model, params, batches, lm_loss, steps_per_window, lr, 1, amp, scaler, |
| f"elevator {idxs[0]}-{idxs[-1]}", log_every=max(1, steps_per_window)) |
| window_losses.append(loss) |
| for i in idxs: |
| for p in layers[i].parameters(): |
| p.requires_grad_(False) |
| if base_dtype != torch.float32: |
| layers[i].to(base_dtype) |
| if ring is not None: |
| ring.unpin(idxs, save=True) |
| del params |
| gc.collect() |
| if ring is not None: |
| model.gradient_checkpointing_disable() |
| model.eval() |
| save_stage1_checkpoint(model, tokenizer, checkpoint_dir, cfg, "full") |
| return {"kind": "full", "loss": sum(window_losses) / max(len(window_losses), 1), |
| "checkpoint": str(checkpoint_dir)} |
|
|
