"""Phase 3 — continued pretraining of Qwen3-1.7B-Base on the packed RU corpus.

Fits a 16GB RX 7900 GRE (gfx1100). Full-weight CPT of 1.7B with plain Adam does
NOT fit; we use adafactor (no momentum states) + gradient checkpointing (path A,
plan §3.1). If it still OOMs, set PATH="dora" for high-rank DoRA on all linears.

NO bitsandbytes (ROCm flaky) → no 8-bit Adam, hence adafactor.

Critical hyperparams (plan §3.2): lr=1e-5 (10x lower than LoRA — nudge, don't
forget), 1 epoch (more = catastrophic forgetting). Checkpoints every 500 steps;
resume with --resume.

Usage:
    HSA_OVERRIDE_GFX_VERSION=11.0.0 python train_cpt.py 2>&1 | tee cpt_run.log
    HSA_OVERRIDE_GFX_VERSION=11.0.0 python train_cpt.py --resume
"""
import os
import sys

os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")
os.environ.setdefault("TMPDIR", "/mnt/D/tmp")
os.environ.setdefault("HSA_OVERRIDE_GFX_VERSION", "11.0.0")
# lever 2: kill allocator fragmentation so the caching allocator doesn't hoard the card
os.environ.setdefault("PYTORCH_HIP_ALLOC_CONF", "expandable_segments:True")

import torch
from datasets import load_from_disk
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments,
    DataCollatorForLanguageModeling,
)

from corpus_common import PACKED

MODEL = "Qwen/Qwen3-1.7B-Base"          # BASE, not Instruct (plan hard rule 4)
OUTPUT = "./Qwen3-1.7B-ru-cpt"
PATH = os.environ.get("CPT_PATH", "full")   # "full" (adafactor) | "dora"


def main():
    resume = "--resume" in sys.argv

    tok = AutoTokenizer.from_pretrained(MODEL)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token

    ds = load_from_disk(str(PACKED))
    print(f"[cpt] {len(ds)} blocks x 2048 = {len(ds)*2048/1e6:.0f}M tokens")

    # lever 3: hard cap VRAM so an OOM raises (recoverable) instead of starving the compositor
    if torch.cuda.is_available():
        torch.cuda.set_per_process_memory_fraction(0.85)

    model = AutoModelForCausalLM.from_pretrained(
        MODEL, torch_dtype=torch.bfloat16, attn_implementation="sdpa")  # lever 1: flash/mem-efficient
    model.config.use_cache = False
    model.gradient_checkpointing_enable()

    optim = "adafactor"
    if PATH == "dora":                  # fallback if full-weight OOMs (plan §3.1B)
        from peft import LoraConfig, get_peft_model
        cfg = LoraConfig(
            r=64, lora_alpha=128, lora_dropout=0.05, use_dora=True,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                            "gate_proj", "up_proj", "down_proj"],
            task_type="CAUSAL_LM",
        )
        model = get_peft_model(model, cfg)
        model.print_trainable_parameters()
        optim = "adamw_torch_fused"     # LoRA params are tiny, Adam fits

    args = TrainingArguments(
        output_dir=OUTPUT,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=16,     # eff batch 16 x 2048 tok
        num_train_epochs=1,                 # ONE pass (hard rule 3)
        learning_rate=1e-5,                 # LOW (hard rule 3)
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        optim=optim,
        bf16=True,
        gradient_checkpointing=True,
        logging_steps=20,
        save_steps=200,                     # ~1.1h between ckpts → clean daily stop/resume, low loss on interrupt
        save_total_limit=3,
        report_to="none",
        dataloader_num_workers=0,   # py3.14 mp leaks semaphores + crashes; data is pre-tokenized so workers buy nothing
    )

    collator = DataCollatorForLanguageModeling(tok, mlm=False)
    trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=collator)
    trainer.train(resume_from_checkpoint=resume)

    if PATH == "dora":
        model = model.merge_and_unload()
    model.save_pretrained(OUTPUT)
    tok.save_pretrained(OUTPUT)
    print(f"[cpt] DONE -> {OUTPUT}")


if __name__ == "__main__":
    main()