"""Stage 4: pretrain the generator on the (direction, target_text) firehose. Conditioning vector for each record is row `vec_idx` of the memmap vec bank (probe directions in READ_LAYER residual space); injected at INJECT_LAYER at the marker. Teacher-force the target. torchrun --standalone --nproc_per_node=8 scripts/pretrain.py --data-dir data/pretrain --epochs 1 """ import argparse import json import math import os import time import numpy as np import torch import torch.distributed as dist from peft import LoraConfig, PeftModel, get_peft_model from torch.nn.parallel import DistributedDataParallel as DDP from transformers import AutoModelForCausalLM, AutoTokenizer import wandb from mxf.config import D_MODEL, INJECT_LAYER, MODEL, STEER_COEFF, TrainConfig from mxf.inject import get_layer, hooked, make_inject_hook, make_packed_inject_hook from mxf.mfu import mfu from mxf.prompts import build_sft_ids def pack_examples(toks, pack_len, seed=0): """Greedy-pack (ids, labels, marker_pos, vec_idx) tuples end-to-end into blocks of <= pack_len tokens (example order shuffled once; an example that would overflow starts the next block, so examples are never split). Each block: ids/labels concat + per-example seg_lens, absolute marker positions, vec idxs.""" order = np.random.default_rng(seed).permutation(len(toks)) blocks, cur = [], None for i in order: ids, labs, pos, vidx = toks[i] if len(ids) > pack_len: continue if cur is None or len(cur["ids"]) + len(ids) > pack_len: if cur is not None: blocks.append(cur) cur = {"ids": [], "labels": [], "seg_lens": [], "markers": [], "vec_idxs": []} cur["markers"].append(len(cur["ids"]) + pos[0]) cur["ids"] += ids cur["labels"] += labs cur["seg_lens"].append(len(ids)) cur["vec_idxs"].append(vidx) if cur is not None and cur["seg_lens"]: blocks.append(cur) return blocks def pack_batch(bblocks, pack_len, pad_id): """CPU tensors for a batch of packed blocks. seg = example index per token; tail pads get a unique seg id each (self-attention only, never attended by real tokens, labels -100). position_ids restart at 0 for every example. Returns per-marker (rows, cols) for injection.""" B = len(bblocks) input_ids = torch.full((B, pack_len), pad_id, dtype=torch.long) labels = torch.full((B, pack_len), -100, dtype=torch.long) pos_ids = torch.zeros((B, pack_len), dtype=torch.long) seg = torch.arange(pack_len, dtype=torch.long).repeat(B, 1) + 1_000_000 # pads: isolated rows, cols, n_real = [], [], 0 for b, blk in enumerate(bblocks): n = len(blk["ids"]) n_real += n input_ids[b, :n] = torch.tensor(blk["ids"]) labels[b, :n] = torch.tensor(blk["labels"]) s = 0 for j, sl in enumerate(blk["seg_lens"]): seg[b, s : s + sl] = j pos_ids[b, s : s + sl] = torch.arange(sl) s += sl rows += [b] * len(blk["markers"]) cols += blk["markers"] return input_ids, labels, pos_ids, seg, torch.tensor(rows), torch.tensor(cols), n_real def packed_attn_mask(seg, causal, dtype): """Additive [B,1,L,L] mask: 0 where (same example ∧ causal), finfo.min elsewhere. transformers returns already-4D masks as-is, so this reaches sdpa untouched.""" allowed = (seg[:, None, :, None] == seg[:, None, None, :]) & causal zero = torch.zeros((), dtype=dtype, device=seg.device) neg = torch.full((), torch.finfo(dtype).min, dtype=dtype, device=seg.device) return torch.where(allowed, zero, neg) def main(): cfg = TrainConfig() ap = argparse.ArgumentParser() ap.add_argument("--data-dir", default="data/pretrain") ap.add_argument("--init-adapter", default=cfg.init_adapter) ap.add_argument("--save-dir", default=cfg.save_dir) ap.add_argument("--lr", type=float, default=cfg.lr) ap.add_argument("--batch-size", type=int, default=cfg.batch_size) ap.add_argument("--epochs", type=int, default=cfg.epochs) ap.add_argument("--max-seq", type=int, default=cfg.max_seq) ap.add_argument("--pack-len", type=int, default=0, help="0 = per-example padded batches + compile (validated 57%% MFU, the default). " ">0 packs into fixed blocks but REGRESSES on Blackwell (no flash-attn → dense " "attn mask wastes off-block compute); only use with a block-sparse attn backend.") ap.add_argument("--pack-blocks", type=int, default=8, help="packed blocks per device micro-batch (tokens/step = pack-blocks * pack-len)") ap.add_argument("--run-name", default=cfg.run_name) ap.add_argument("--compile", action="store_true", help="torch.compile the policy (test injection still fires)") ap.add_argument("--no-wandb", action="store_true") a = ap.parse_args() world = int(os.environ.get("WORLD_SIZE", 1)); rank = int(os.environ.get("RANK", 0)) local = int(os.environ.get("LOCAL_RANK", 0)); is_main = rank == 0 if world > 1: dist.init_process_group("nccl"); torch.cuda.set_device(local) device = f"cuda:{local}" tok = AutoTokenizer.from_pretrained(MODEL) if tok.pad_token is None: tok.pad_token = tok.eos_token records = [json.loads(l) for l in open(f"{a.data_dir}/records.jsonl")] n_vecs = max(r["vec_idx"] for r in records) + 1 vecs = np.memmap(f"{a.data_dir}/vecs.f32", dtype=np.float32, mode="r", shape=(n_vecs, D_MODEL)) records = records[rank::world][: len(records) // world] # equal shards: unequal lengths deadlock DDP on the last batch if is_main: print(f"{len(records)*world} records, {n_vecs} vectors, world={world}", flush=True) model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, attn_implementation="sdpa", # flash-attn has no sm_103 build device_map={"": device}) model.enable_input_require_grads() if a.init_adapter: model = PeftModel.from_pretrained(model, a.init_adapter, is_trainable=True) else: model = get_peft_model(model, LoraConfig( r=cfg.lora_r, lora_alpha=cfg.lora_alpha, lora_dropout=0.0, use_rslora=True, target_modules="all-linear", bias="none", task_type="CAUSAL_LM")) model.train() n_params = sum(p.numel() for p in model.parameters()) # ~8.19B; LoRA adds <0.2%, fine for MFU if a.compile: model.forward = torch.compile(model.forward) # hook still fires (graph-breaks at layer-1) ddp = DDP(model, device_ids=[local]) if world > 1 else model opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=a.lr, weight_decay=0.0) submodule = get_layer(model, INJECT_LAYER) # pre-tokenize once. Packed path: shuffle-once greedy packing into fixed pack-len blocks (zero # intra-block padding, one static shape for compile). Legacy path: length-bucketed padded batches. toks_cache = [] for r in records: ids, labs, pos = build_sft_ids(tok, r["target_text"]) toks_cache.append((ids[: a.max_seq], labs[: a.max_seq], pos, r["vec_idx"])) if a.pack_len: blocks = pack_examples(toks_cache, a.pack_len, seed=0) if world > 1: # equalize block count across ranks (packing yields ±1 per rank → DDP deadlock) t = torch.tensor([len(blocks)], device=device) dist.all_reduce(t, op=dist.ReduceOp.MIN) blocks = blocks[: int(t.item())] bper = len(blocks) // a.pack_blocks # drop remainder batch: keeps a single static shape steps_total = bper * a.epochs causal = torch.tril(torch.ones(a.pack_len, a.pack_len, dtype=torch.bool, device=device)) if is_main: fill = sum(len(b["ids"]) for b in blocks) / (len(blocks) * a.pack_len) print(f"packed: {len(blocks)} blocks of {a.pack_len} (fill {fill:.1%}), " f"{bper} steps/epoch x {a.pack_blocks} blocks", flush=True) else: toks_cache.sort(key=lambda t: len(t[0])) steps_total = math.ceil(len(toks_cache) / a.batch_size) * a.epochs sched = torch.optim.lr_scheduler.OneCycleLR(opt, a.lr, total_steps=steps_total, pct_start=cfg.warmup_frac, anneal_strategy="linear") if is_main and not a.no_wandb: wandb.init(project="maxact-fast", name=a.run_name, config=vars(a)) os.makedirs(a.save_dir, exist_ok=True) step = 0 for ep in range(a.epochs): if a.pack_len: order = np.random.default_rng(ep).permutation(len(blocks)) batches = [[blocks[i] for i in order[s : s + a.pack_blocks]] for s in range(0, bper * a.pack_blocks, a.pack_blocks)] else: batches = [toks_cache[s : s + a.batch_size] for s in range(0, len(toks_cache), a.batch_size)] np.random.default_rng(ep).shuffle(batches) for batch in batches: t0 = time.time() if a.pack_len: input_ids, labels, pos_ids, seg, rows, cols, n_real = pack_batch( batch, a.pack_len, tok.pad_token_id) mask4 = packed_attn_mask(seg.to(device), causal, torch.bfloat16) vmat = torch.from_numpy(np.asarray(vecs[[v for blk in batch for v in blk["vec_idxs"]]])) hook = make_packed_inject_hook(vmat, rows, cols, STEER_COEFF, device, torch.bfloat16) with hooked(submodule, hook): out = ddp(input_ids=input_ids.to(device), attention_mask=mask4, position_ids=pos_ids.to(device), labels=labels.to(device), use_cache=False) else: L = max(len(t[0]) for t in batch) L = min(((L + 63) // 64) * 64, a.max_seq) # round to mult-of-64 → ≤3 static shapes for compile input_ids = torch.full((len(batch), L), tok.pad_token_id, dtype=torch.long) labels = torch.full((len(batch), L), -100, dtype=torch.long) attn = torch.zeros((len(batch), L), dtype=torch.bool) pos = batch[0][2] for i, (ii, ll, _, _) in enumerate(batch): input_ids[i, : len(ii)] = torch.tensor(ii) labels[i, : len(ll)] = torch.tensor(ll) attn[i, : len(ii)] = True n_real = int(attn.sum()) vlist = [torch.from_numpy(np.asarray(vecs[t[3]])).unsqueeze(0) for t in batch] hook = make_inject_hook(vlist, [pos] * len(batch), STEER_COEFF, device, torch.bfloat16) with hooked(submodule, hook): out = ddp(input_ids=input_ids.to(device), attention_mask=attn.to(device), labels=labels.to(device)) out.loss.backward() torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0) opt.step(); sched.step(); opt.zero_grad() if is_main and step % 20 == 0: torch.cuda.synchronize() tfl, m = mfu(n_real, time.time() - t0, n_params, fwd_bwd=True) print(f"ep{ep} step {step}/{steps_total} loss {out.loss.item():.4f} | " f"{tfl:.0f} TFLOP/s MFU {m:.0%}", flush=True) if not a.no_wandb: wandb.log({"loss": out.loss.item(), "lr": sched.get_last_lr()[0], "mfu": m, "tflops": tfl}, step=step) if is_main and step % 2000 == 0 and step: model.save_pretrained(f"{a.save_dir}/step_{step}") step += 1 if is_main: model.save_pretrained(f"{a.save_dir}/final") print("PRETRAIN_DONE", flush=True) if world > 1: dist.destroy_process_group() if __name__ == "__main__": main()