parser/joint_train.py

"""Fine-tune ONE Stoicheia backbone into a full CoNLL-U predictor: lemma (edit-script),
UPOS, factored XPOS/morph, and biaffine HEAD+DEPREL — all heads share the pooled, scalar-mixed
word representation and train jointly (multi-task; morph and syntax reinforce each other).

Reuses the tagger's proven data pipeline (TaggerDataset/LabelVocab, SOTA 94.22 XPOS) and DDP
recipe, adds a Dozat-Manning biaffine head + per-sentence gold alignment + LAS/UAS eval.

  torchrun --nproc_per_node=4 -m parser.joint_train --config configs/joint.json
"""
from __future__ import annotations

import argparse, datetime, json, math, os, sys, time
from pathlib import Path

import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from tagger.backbone import load_backbone_auto
from tagger.conllu import read_conllu
from tagger.dataset import TaggerDataset, batch_chunk, encode_word, pack_dev_items
from tagger.edits import LabelVocab
from tagger.model import TaggerConfig
from tagger.train import param_groups as tagger_param_groups

from parser.biaffine import ParserConfig
from parser.labels import DeprelVocab
from parser.model import build_gold
from parser.joint_model import JointModel


def ddp_setup():
    if "RANK" in os.environ:
        import torch.distributed as dist
        dist.init_process_group("nccl", timeout=datetime.timedelta(minutes=60))
        rank = dist.get_rank(); world = dist.get_world_size()
        torch.cuda.set_device(rank % torch.cuda.device_count())
        return rank, world, True
    return 0, 1, False


def parse_gold(sent_ids, sentences, deprel_vocab, max_w, device):
    """Gold heads/labels/mask (n_sent, max_w) for the sentences in `sent_ids`, in the same
    encodable-token order JointModel._regroup used (build_gold shares that ordering)."""
    n_sent = len(sent_ids)
    heads = torch.full((n_sent, max_w), -100, dtype=torch.long)
    labels = torch.full((n_sent, max_w), -100, dtype=torch.long)
    mask = torch.zeros((n_sent, max_w), dtype=torch.bool)
    for local, si in enumerate(sent_ids):
        n, h, l = build_gold(sentences[si], deprel_vocab)
        n = min(n, max_w)
        if n == 0:
            continue
        heads[local, :n] = torch.tensor(h[:n])
        labels[local, :n] = torch.tensor(l[:n])
        mask[local, :n] = True
    return heads.to(device), labels.to(device), mask.to(device)


@torch.no_grad()
def evaluate_dev(model, core, dev_rows, sentences, deprel_vocab, micro, device, T, W,
                 decode="greedy", tokenizer=None):
    """counts: [xpos_exact, script_ok, upos_ok, n_words, uas_ok, las_ok, n_arc]."""
    model.eval()
    cnt = torch.zeros(7, dtype=torch.long, device=device)
    for i in range(0, len(dev_rows), micro):
        batch = batch_chunk(dev_rows[i:i + micro], T, W, tokenizer)
        batch = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()}
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            tag_out, arc, rel, wmask, sent_ids = model(batch)
        m = batch["y_script"] != -100
        if m.any():
            ok = torch.ones_like(m)
            for p, lg in enumerate(tag_out["xpos"]):
                good = lg.argmax(-1) == batch["y_xpos"][:, :, p]
                ok &= good | (batch["y_xpos"][:, :, p] == -100)
            cnt[0] += int((ok & m).sum())
            cnt[1] += int(((tag_out["script"].argmax(-1) == batch["y_script"]) & m).sum())
            cnt[2] += int(((tag_out["upos"].argmax(-1) == batch["y_upos"]) & m).sum())
            cnt[3] += int(m.sum())
        if arc is not None:
            gh, gl, gm = parse_gold(sent_ids, sentences, deprel_vocab, arc.shape[1], device)
            if decode == "mst":
                from parser.mst import mst_heads_labels
                ph, pl = mst_heads_labels(arc, rel, wmask)
            else:
                ph, pl = core.biaffine.decode(arc, rel, wmask)
            gh_c, gl_c, gm_c = gh.cpu(), gl.cpu(), gm.cpu()
            valid = (gh_c != -100) & gm_c
            cnt[4] += int(((ph == gh_c) & valid).sum())
            cnt[5] += int(((ph == gh_c) & (pl == gl_c) & valid).sum())
            cnt[6] += int(valid.sum())
    model.train()
    return cnt


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    a = ap.parse_args()
    cfg = json.loads(Path(a.config).read_text())
    for k in ("out_dir", "kfold_dir"):
        cfg[k] = os.path.expandvars(cfg[k])
    if "ckpt" in cfg:
        cfg["ckpt"] = os.path.expandvars(cfg["ckpt"])

    rank, world, is_ddp = ddp_setup()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(cfg.get("seed", 0) + rank)
    torch.set_float32_matmul_precision("high")

    fold = cfg["fold"]
    T, W, micro = cfg.get("T", 2048), cfg.get("W", 384), cfg["micro_batch"]
    w_parse = cfg.get("w_parse", 1.0)
    out = Path(cfg["out_dir"]); out.mkdir(parents=True, exist_ok=True)

    def log0(*x):
        if rank == 0:
            print(*x, flush=True)

    kdir = Path(cfg["kfold_dir"])
    train_sents = list(read_conllu(kdir / f"train{fold}.conllu"))
    dev_sents = list(read_conllu(kdir / f"dev{fold}.conllu"))
    if cfg.get("max_sents"):
        train_sents = train_sents[:cfg["max_sents"]]
        dev_sents = dev_sents[:max(cfg["max_sents"] // 4, 8)]

    is_enc = lambda f: encode_word(f) is not None
    vocab = LabelVocab.build(train_sents, is_enc)
    deprel_vocab = DeprelVocab.build(train_sents)
    if rank == 0:
        vocab.save(out / "vocab.json")
        deprel_vocab.save(out / "deprel_vocab.json")
    log0(f"fold={fold} train={len(train_sents)} dev={len(dev_sents)} scripts={vocab.n_scripts} "
         f"tags={len(vocab.tags)} deprels={len(deprel_vocab.rels)} world={world}")

    # ---- model (loaded before the datasets since the HF path needs its tokenizer to encode)
    encoder, pcfg_pre, tokenizer = load_backbone_auto(cfg, device)
    is_hf = tokenizer is not None
    hf_max_len = cfg.get("hf_max_len", 512)
    use_cap = cfg.get("use_cap", True)
    if is_hf and use_cap:
        # see tagger/train.py for the same guard: the fine-tune-only cap channel is injected
        # into CharBERT's char-plane embeddings and has no equivalent for subword input
        log0("  [hf backbone] ignoring use_cap=true (no char-plane to inject into)")
        use_cap = False

    train_ds = TaggerDataset(train_sents, vocab, T, W, tokenizer=tokenizer, hf_max_len=hf_max_len)
    dev_ds = TaggerDataset(dev_sents, vocab, T, W, tokenizer=tokenizer, hf_max_len=hf_max_len)
    dev_rows, _ = pack_dev_items(dev_ds.encs, W, tokenizer, T)
    dev_shard = dev_rows[rank::world]

    tcfg = TaggerConfig(pool=cfg.get("pool", "mean"), head_dropout=cfg.get("head_dropout", 0.33),
                        w_xpos=cfg.get("w_xpos", 1.0), w_script=cfg.get("w_script", 1.0),
                        w_upos=cfg.get("w_upos", 0.2), w_flat=cfg.get("w_flat", 1.0),
                        use_cap=use_cap, scalar_mix=cfg.get("scalar_mix", True))
    pcfg = ParserConfig(d_arc=cfg.get("d_arc", 500), d_rel=cfg.get("d_rel", 150),
                        dropout=cfg.get("parse_dropout", 0.33), n_labels=len(deprel_vocab.rels))
    core = JointModel(encoder, vocab, tcfg, pcfg, W=W).to(device)

    model = core
    if is_ddp:
        from torch.nn.parallel import DistributedDataParallel as DDP
        # HF backbones (BERT/RoBERTa-style) instantiate a pooler (pooler.dense) that is never
        # called by HFBackboneWithHidden.forward() (we only consume hidden_states), so it never
        # receives a gradient -- DDP's default strict "all params must be used" check throws
        # without find_unused_parameters=True here. The CharBERT path has no such unused
        # trainable param (frozen LM heads are excluded from requires_grad entirely), so it's
        # left off there to keep that path's DDP bucketing at its original efficiency.
        model = DDP(core, device_ids=[rank % torch.cuda.device_count()],
                   find_unused_parameters=is_hf)
    # core = model.module below stays valid; frozen LM heads keep find_unused_parameters=False safe

    groups = tagger_param_groups(core.tagger, cfg)
    groups.append(dict(params=list(core.biaffine.parameters()), lr=cfg["lr_head"],
                       weight_decay=cfg.get("head_wd", 0.0), base_lr=cfg["lr_head"], is_enc=False))
    groups = [g for g in groups if any(p.requires_grad for p in g["params"])]
    opt = torch.optim.AdamW(groups, betas=(0.9, 0.95), fused=(device.type == "cuda"))

    ref_rows, trunc = pack_dev_items(train_ds.encs, W, tokenizer, T)
    n_rows = len(ref_rows) // world * world
    steps_per_epoch = max(n_rows // world // micro, 1)
    total_steps = steps_per_epoch * cfg["epochs"]
    warmup = max(int(total_steps * cfg.get("warmup_frac", 0.1)), 1)
    log0(f"rows/epoch={n_rows} steps/epoch={steps_per_epoch} total={total_steps} "
         f"trunc={trunc} w_parse={w_parse} eff_batch_rows={micro*world}")

    def lr_scale(step):
        if step < warmup:
            return step / warmup
        f = (step - warmup) / max(total_steps - warmup, 1)
        return 0.5 * (1 + math.cos(math.pi * min(f, 1.0)))

    metrics_f = out / f"metrics_rank{rank}.jsonl"
    t0 = time.time()

    def do_step(batch, lr_s, step):
        for pg in opt.param_groups:
            pg["lr"] = pg["base_lr"] * lr_s
        batch = {k: (v.to(device) if torch.is_tensor(v) else v) for k, v in batch.items()}
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
            tag_out, arc, rel, wmask, sent_ids = model(batch)
            tag_loss, logs = core.tagger.loss(tag_out, batch)
            if arc is not None:
                gh, gl, gm = parse_gold(sent_ids, train_sents, deprel_vocab, arc.shape[1], device)
                p_loss, p_logs = core.biaffine.loss(arc, rel, gh, gl, gm)
                loss = tag_loss + w_parse * p_loss
            else:
                p_logs = {}
                loss = tag_loss
        opt.zero_grad(set_to_none=True)
        loss.backward()
        gnorm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.get("clip", 1.0))
        if torch.isfinite(gnorm):
            opt.step()
        if rank == 0 and step % cfg.get("log_every", 10) == 0:
            rec = dict(step=step, lr=round(opt.param_groups[0]["lr"], 7),
                       loss=round(float(loss.item()), 4), gnorm=round(float(gnorm), 3),
                       **logs, **{f"p_{k}": v for k, v in p_logs.items()},
                       sps=round(step / max(time.time() - t0, 1e-9), 2))
            print("  " + " ".join(f"{k}={v}" for k, v in rec.items()), flush=True)
            with open(metrics_f, "a") as mf:
                mf.write(json.dumps(rec) + "\n")

    best_score, best_epoch, step = -1.0, -1, 0
    for epoch in range(cfg["epochs"]):
        rows, _ = pack_dev_items(train_ds.encs, W, tokenizer, T, order=torch.randperm(
            len(train_ds.encs),
            generator=torch.Generator().manual_seed(cfg.get("seed", 0) * 1000 + epoch)).tolist())
        n = len(rows) // world * world
        shard = rows[rank:n:world]
        nsteps = len(shard) // micro * micro
        for i in range(0, nsteps, micro):
            do_step(batch_chunk(shard[i:i + micro], T, W, tokenizer), lr_scale(step), step)
            step += 1

        cnt = evaluate_dev(model, core, dev_shard, dev_sents, deprel_vocab,
                           cfg.get("eval_micro", micro), device, T, W, tokenizer=tokenizer)
        if is_ddp:
            import torch.distributed as dist
            dist.all_reduce(cnt)
        nw = max(int(cnt[3]), 1); na = max(int(cnt[6]), 1)
        m = dict(epoch=epoch, xpos_exact=round(int(cnt[0]) / nw, 4),
                 lemma_script=round(int(cnt[1]) / nw, 4), upos=round(int(cnt[2]) / nw, 4),
                 uas=round(int(cnt[4]) / na, 4), las=round(int(cnt[5]) / na, 4), n_words=nw)
        score = (m["xpos_exact"] + m["lemma_script"] + m["upos"] + m["las"]) / 4
        m["score"] = round(score, 4)
        log0("  EVAL " + json.dumps(m))
        if rank == 0:
            with open(out / "eval.jsonl", "a") as ef:
                ef.write(json.dumps(m) + "\n")
            if score > best_score:
                tmp = out / "best.pt.tmp"
                torch.save(dict(model=core.state_dict(), tcfg=vars(tcfg), pcfg=vars(pcfg),
                                cfg=cfg, pretrain_cfg=pcfg_pre, epoch=epoch, dev=m, W=W, T=T,
                                deprel_vocab=deprel_vocab.rels), tmp)
                os.replace(tmp, out / "best.pt")
                print(f"  new best.pt (score {score:.4f} las {m['las']:.4f})", flush=True)
        if score > best_score:
            best_score, best_epoch = score, epoch
        if epoch - best_epoch >= cfg.get("patience", 8):
            log0(f"EARLY STOP epoch={epoch} best_epoch={best_epoch} best_score={best_score:.4f}")
            break

    log0(f"DONE best_score={best_score:.4f} best_epoch={best_epoch} ({(time.time()-t0)/60:.1f} min)")
    if is_ddp:
        import torch.distributed as dist
        dist.destroy_process_group()


if __name__ == "__main__":
    main()