scripts/pretrain.py

"""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()