| """Real miniature SFT/LoRA and DPO optimization; no general capability claim.""" |
| from pathlib import Path |
| import json |
| import sys |
| import copy |
| import time |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import torch |
| from safetensors.torch import load_file, save_file |
| from nexora.model import NexoraLM, ModelConfig |
| from nexora.tokenizer import ByteTokenizer |
| from nexora.adapters import inject_lora, merge_lora |
| from nexora.posttraining import masked_sft_loss, dpo_loss |
|
|
|
|
| def main(): |
| torch.set_num_threads(4) |
| torch.manual_seed(42) |
| root = Path("artifacts/tiny") |
| cfg = ModelConfig(**json.loads((root / "config.json").read_text())) |
| model = NexoraLM(cfg) |
| model.load_state_dict(load_file(str(root / "model.safetensors"))) |
| reference = copy.deepcopy(model).eval().requires_grad_(False) |
| replaced = inject_lora(model) |
| opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=.005) |
| tok = ByteTokenizer() |
| rows = [ |
| ("User: What confirms a code change?\nAssistant: ", "Run tests and inspect their actual results."), |
| ("User: A test failed. Is the repair verified?\nAssistant: ", "No. Inspect the failure, repair it, and run the test again."), |
| ("User: What belongs in a tool receipt?\nAssistant: ", "The operation, observed output, exit status, and elapsed time."), |
| ("User: How should memory represent a guess?\nAssistant: ", "Record its source and uncertainty instead of storing it as a fact."), |
| ] |
| def tensors(prompt, answer): |
| prefix = [tok.bos_id, *tok.encode(prompt)] |
| ids = prefix + tok.encode(answer) + [tok.eos_id] |
| x, y = torch.tensor([ids[:-1]]), torch.tensor([ids[1:]]) |
| mask = torch.arange(y.shape[1])[None] >= len(prefix)-1 |
| return x, y, mask |
| def logprob(m, prompt, answer): |
| x, y, mask = tensors(prompt, answer) |
| logits, _ = m(x) |
| return (logits.log_softmax(-1).gather(-1, y[..., None]).squeeze(-1)*mask).sum(-1) |
| history = [] |
| start = time.perf_counter() |
| for step in range(40): |
| prompt, answer = rows[step % len(rows)] |
| x, y, mask = tensors(prompt, answer) |
| logits, _ = model(x) |
| loss = masked_sft_loss(logits, y, mask) |
| opt.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1) |
| opt.step() |
| if step % 10 == 0 or step == 39: |
| history.append({"stage": "SFT", "step": step+1, "loss": loss.item()}) |
| prompt = "User: A test failed. Is the repair verified?\nAssistant: " |
| chosen, rejected = "No. Inspect the failure and test again.", "Yes. Everything passed successfully." |
| with torch.no_grad(): |
| rc, rr = logprob(reference, prompt, chosen), logprob(reference, prompt, rejected) |
| for step in range(10): |
| loss = dpo_loss(logprob(model, prompt, chosen), logprob(model, prompt, rejected), rc, rr) |
| opt.zero_grad() |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1) |
| opt.step() |
| history.append({"stage": "DPO", "step": step+1, "loss": loss.item()}) |
| out = Path("artifacts/posttraining-experiment") |
| out.mkdir(exist_ok=True) |
| save_file({k: v.detach().contiguous() for k, v in model.state_dict().items() if k.endswith((".a", ".b"))}, str(out / "adapter.safetensors")) |
| probe, _, _ = tensors(*rows[0]) |
| with torch.no_grad(): |
| before, _ = model(probe) |
| merged = merge_lora(copy.deepcopy(model)) |
| after, _ = merged(probe) |
| torch.testing.assert_close(before, after, atol=2e-5, rtol=2e-5) |
| report = {"status": "VALIDATED_TOY_OPTIMIZATION_ONLY", "base": "artifacts/tiny", "rank": 4, "alpha": 8, "targets": replaced, |
| "seconds": time.perf_counter()-start, "history": history, "merge_max_abs_error": (before-after).abs().max().item(), |
| "limitations": "Four synthetic SFT examples and one preference pair; no evidence of reasoning improvement or generalization. Adapter not enabled by default."} |
| (out / "config.json").write_text(json.dumps(report, indent=2)) |
| Path("reports/posttraining.json").write_text(json.dumps(report, indent=2)) |
| print(json.dumps(report)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
