# QLoRA SFT (4-bit nf4, bitsandbytes + peft + HF Trainer). Loss only on the completion tokens.
import json, sys, time, torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, Trainer, TrainingArguments
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
SYS = "Please reason step by step, and put your final answer within \\boxed{}."
def prompt(q):  # same Qwen3 prompt as gen.py (copied: importing gen.py ran its argparse)
    return ("<|im_start|>system\n" + SYS + "<|im_end|>\n<|im_start|>user\n" + q + "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n")
MODEL = "Qwen/Qwen3-8B"; import os; MAXLEN = int(os.environ.get("E20_MAXLEN", "1536"))
torch.cuda.set_per_process_memory_fraction(0.90)  # stay inside dedicated VRAM, no Windows sysmem fallback
tok = AutoTokenizer.from_pretrained(MODEL)
rows = [json.loads(l) for l in open("data/sft.jsonl", encoding="utf-8")]
data = []
for r in rows:
    p = tok(prompt(r["problem"]), add_special_tokens=False)["input_ids"]
    c = tok(r["completion"] + "<|im_end|>", add_special_tokens=False)["input_ids"]
    ids = (p + c)[:MAXLEN]; lab = ([-100] * len(p) + c)[:MAXLEN]
    if sum(v != -100 for v in lab[1:]) >= 8: data.append({"input_ids": ids, "labels": lab})  # drop rows whose completion was fully truncated
print("examples", len(data), "truncated", sum(len(d["input_ids"]) == MAXLEN for d in data), flush=True)
def collate(b):
    m = max(len(x["input_ids"]) for x in b)
    ids = torch.full((len(b), m), tok.pad_token_id); lab = torch.full((len(b), m), -100); att = torch.zeros((len(b), m), dtype=torch.long)
    for i, x in enumerate(b):
        n = len(x["input_ids"]); ids[i, :n] = torch.tensor(x["input_ids"]); lab[i, :n] = torch.tensor(x["labels"]); att[i, :n] = 1
    return {"input_ids": ids, "labels": lab, "attention_mask": att}
bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
model = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=bnb, dtype=torch.bfloat16, device_map={"": 0}, attn_implementation="sdpa")
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}); model.enable_input_require_grads()  # no fp32 upcast of embeddings/lm_head (peft prepare_model_for_kbit_training upcasts them, +2.3 GB)
model = get_peft_model(model, LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05, task_type="CAUSAL_LM",
                       target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]))
model.print_trainable_parameters()
EPOCHS = float(sys.argv[1]) if len(sys.argv) > 1 else 1.0
args = TrainingArguments(output_dir="out/lora_ckpt", per_device_train_batch_size=1, gradient_accumulation_steps=16,
    num_train_epochs=EPOCHS, learning_rate=1e-4, lr_scheduler_type="cosine", warmup_steps=2, logging_steps=1,
    bf16=True, optim="paged_adamw_8bit", gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False},
    save_strategy="steps", save_steps=20, save_total_limit=2, report_to=[], seed=20)
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
class ChunkedLossTrainer(Trainer):
    # Same loss as standard causal LM SFT, but the 151k-vocab logits are computed 256 tokens at a time
    # (recomputed in backward) so the full logits tensor never sits in VRAM.
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs["labels"]
        inner = model.base_model.model
        h = inner.model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"]).last_hidden_state
        h = h[:, :-1]; y = labels[:, 1:]
        n = (y != -100).sum().clamp(min=1); tot = (h[:, :1].sum() * 0).float()
        def part(hc, yc): return F.cross_entropy(inner.lm_head(hc).float().flatten(0, 1), yc.flatten(), ignore_index=-100, reduction="sum")
        for s in range(0, h.shape[1], 256):
            yc = y[:, s:s+256]
            if (yc != -100).any(): tot = tot + checkpoint(part, h[:, s:s+256], yc, use_reentrant=False)
        anchor = sum(p.reshape(-1)[0] for p in model.parameters() if p.requires_grad) * 0.0  # keeps a grad path even if a micro-batch has no label chunk
        return tot / n + anchor
t = time.time()
tr = ChunkedLossTrainer(model=model, args=args, train_dataset=data, data_collator=collate)
tr.train()
model.save_pretrained("out/lora_adapter")
json.dump(tr.state.log_history, open("logs/train_log.json", "w"), indent=1)
print("TRAINDONE", round(time.time() - t), flush=True)