792 lines
31 KiB
Python
792 lines
31 KiB
Python
"""Modal PKPO pipeline: baseline eval, tiny warmup, PKPO LoRA, upload."""
|
|
from __future__ import annotations
|
|
|
|
import gc
|
|
import inspect
|
|
import json
|
|
import os
|
|
import shutil
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import modal
|
|
|
|
from modal_common import BASE_MODEL, CACHE_DIR, HF_REPO, WORK_DIR, hf_cache_vol, image, work_vol
|
|
|
|
app = modal.App("pkpo-train")
|
|
|
|
|
|
def _copy_source_bundle(dst: Path):
|
|
import agent_core
|
|
import eval_lib
|
|
import pkpo
|
|
import shipped_tool
|
|
import train_pkpo_modal
|
|
|
|
for mod, name in [
|
|
(agent_core, "agent_core.py"),
|
|
(eval_lib, "eval_lib.py"),
|
|
(pkpo, "pkpo.py"),
|
|
(shipped_tool, "shipped_tool.py"),
|
|
(train_pkpo_modal, "train_pkpo_modal.py"),
|
|
]:
|
|
Path(dst / name).write_text(Path(inspect.getfile(mod)).read_text(encoding="utf-8"), encoding="utf-8")
|
|
|
|
|
|
def _load_tokenizer(model_ref: str):
|
|
from transformers import AutoTokenizer
|
|
from agent_core import CHAT_TEMPLATE
|
|
|
|
tok = AutoTokenizer.from_pretrained(model_ref, trust_remote_code=True)
|
|
tok.chat_template = CHAT_TEMPLATE
|
|
if tok.pad_token is None:
|
|
tok.pad_token = tok.eos_token
|
|
return tok
|
|
|
|
|
|
def _write_generation_config(model, tok, dst: Path):
|
|
from transformers import GenerationConfig
|
|
from agent_core import GEN_TEMPERATURE, GEN_TOP_P, MAX_TURN_TOKENS, STOP_STRING
|
|
|
|
gen = GenerationConfig.from_model_config(model.config)
|
|
gen.do_sample = True
|
|
gen.temperature = GEN_TEMPERATURE
|
|
gen.top_p = GEN_TOP_P
|
|
gen.max_new_tokens = MAX_TURN_TOKENS
|
|
gen.pad_token_id = tok.pad_token_id
|
|
gen.eos_token_id = tok.eos_token_id
|
|
gen.stop_strings = [STOP_STRING]
|
|
gen.save_pretrained(dst)
|
|
|
|
|
|
def _read_json(path: Path, default):
|
|
if path.exists():
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
return default
|
|
|
|
|
|
def _model_card(baseline, final, schedule, notes: str) -> str:
|
|
b = baseline.get("pass_at_1", 0.0) if baseline else 0.0
|
|
f = final.get("pass_at_1", 0.0) if final else 0.0
|
|
return f"""---
|
|
license: other
|
|
base_model: Qwen/Qwen3-8B-Base
|
|
tags:
|
|
- code
|
|
- reinforcement-learning
|
|
- pkpo
|
|
- livecodebench
|
|
---
|
|
|
|
# Qwen3-8B Code PKPO
|
|
|
|
This repo contains a Qwen/Qwen3-8B-Base derivative trained for a small agentic
|
|
coding experiment using the shared tool path in `agent_core.py` and
|
|
`shipped_tool.py`.
|
|
|
|
## Method
|
|
|
|
- Base: `Qwen/Qwen3-8B-Base`.
|
|
- Prompt/template: custom `<think>...</think><answer>...</answer>` template saved
|
|
in the tokenizer. The generation prompt ends with `Assistant: <think>`.
|
|
- Tool protocol: no system role; instructions are merged into the first user
|
|
message; strict user/assistant alternation; plain-text `Tool type` and
|
|
`Tool query` calls.
|
|
- Training data: `deepmind/code_contests` train split only, filtered to old
|
|
stdin/stdout problems. The LiveCodeBench eval subset is not used for training.
|
|
- Reward: binary hidden-test pass/fail.
|
|
- PKPO: `sloo_minus_one` from the paper for `k >= 2`; centered `k=1` rewards for
|
|
the first and final stages. No GRPO-style reward normalization is applied.
|
|
- Schedule actually run: `{schedule}`.
|
|
|
|
The run was intentionally small to fit the free-credit budget and deadline.
|
|
Results should be treated as a reproducible experiment, not a leaderboard model.
|
|
|
|
## Results
|
|
|
|
Evaluation uses `livecodebench/code_generation_lite` `v6`, a fixed subset saved
|
|
at `eval/eval_subset.json`, temperature 1.0, and the same one-turn tool path used
|
|
for training.
|
|
|
|
| model | pass@1 estimate |
|
|
|---|---:|
|
|
| base before training | {b:.4f} |
|
|
| final merged model | {f:.4f} |
|
|
|
|
Raw files:
|
|
|
|
- `eval/baseline_results.json`
|
|
- `eval/final_results.json`
|
|
- `eval/eval_subset.json`
|
|
|
|
## Usage
|
|
|
|
Serve with vLLM:
|
|
|
|
```bash
|
|
vllm serve bk1dr/qwen3-8b-code-pkpo --trust-remote-code --max-model-len 8192
|
|
```
|
|
|
|
Run the shipped tool:
|
|
|
|
```bash
|
|
python shipped_tool.py --base-url http://127.0.0.1:8000/v1 --model bk1dr/qwen3-8b-code-pkpo --max-turns 1 --cp < problem.txt
|
|
```
|
|
|
|
## Run Notes
|
|
|
|
{notes}
|
|
"""
|
|
|
|
|
|
@app.function(
|
|
image=image,
|
|
volumes={CACHE_DIR: hf_cache_vol, WORK_DIR: work_vol},
|
|
secrets=[modal.Secret.from_name("hf-secret")],
|
|
timeout=1800,
|
|
cpu=4,
|
|
memory=16384,
|
|
)
|
|
def push_scaffold():
|
|
from huggingface_hub import HfApi, snapshot_download
|
|
from transformers import AutoConfig
|
|
from agent_core import CHAT_TEMPLATE
|
|
|
|
base_path = snapshot_download(BASE_MODEL, cache_dir=CACHE_DIR)
|
|
out = Path(WORK_DIR) / "scaffold_repo"
|
|
if out.exists():
|
|
shutil.rmtree(out)
|
|
out.mkdir(parents=True)
|
|
tok = _load_tokenizer(base_path)
|
|
tok.save_pretrained(out)
|
|
cfg = AutoConfig.from_pretrained(base_path, trust_remote_code=True)
|
|
cfg.save_pretrained(out)
|
|
(out / "chat_template.jinja").write_text(CHAT_TEMPLATE, encoding="utf-8")
|
|
_copy_source_bundle(out)
|
|
(out / "README.md").write_text(_model_card({}, {}, "pending", "Scaffold pushed before GPU training."), encoding="utf-8")
|
|
HfApi().upload_folder(repo_id=HF_REPO, repo_type="model", folder_path=str(out), commit_message="Initial scaffold")
|
|
work_vol.commit()
|
|
return {"uploaded": str(out)}
|
|
|
|
|
|
@app.function(
|
|
image=image,
|
|
gpu="H100",
|
|
volumes={CACHE_DIR: hf_cache_vol},
|
|
timeout=15 * 60,
|
|
cpu=4,
|
|
memory=98304,
|
|
)
|
|
def smoke_batched_generation(max_new_tokens: int = 96):
|
|
"""Low-cost preflight for the exact batched stop-string generation API."""
|
|
import torch
|
|
from huggingface_hub import snapshot_download
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from agent_core import CHAT_TEMPLATE, STOP_STRING, coding_task_message, render_messages
|
|
|
|
base_path = snapshot_download(BASE_MODEL, cache_dir=CACHE_DIR)
|
|
tok = AutoTokenizer.from_pretrained(base_path, trust_remote_code=True)
|
|
tok.chat_template = CHAT_TEMPLATE
|
|
if tok.pad_token is None:
|
|
tok.pad_token = tok.eos_token
|
|
tok.padding_side = "left"
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
base_path,
|
|
torch_dtype=torch.bfloat16,
|
|
device_map={"": 0},
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
prompt = render_messages([{
|
|
"role": "user",
|
|
"content": coding_task_message(
|
|
"Read an integer N and print N plus one.", max_turns=1
|
|
),
|
|
}], add_generation_prompt=True)
|
|
inputs = tok([prompt, prompt], return_tensors="pt", padding=True).to(model.device)
|
|
prompt_width = inputs.input_ids.shape[1]
|
|
with torch.no_grad():
|
|
out = model.generate(
|
|
**inputs,
|
|
do_sample=True,
|
|
temperature=1.0,
|
|
top_p=0.95,
|
|
max_new_tokens=max_new_tokens,
|
|
pad_token_id=tok.pad_token_id,
|
|
eos_token_id=tok.eos_token_id,
|
|
stop_strings=[STOP_STRING],
|
|
tokenizer=tok,
|
|
)
|
|
completions = [tok.decode(row[prompt_width:], skip_special_tokens=False) for row in out]
|
|
result = {
|
|
"batch_size": len(completions),
|
|
"completion_lengths": [len(c) for c in completions],
|
|
"stopped": [STOP_STRING in c for c in completions],
|
|
"samples": [c[:500] for c in completions],
|
|
}
|
|
print(json.dumps(result), flush=True)
|
|
return result
|
|
|
|
|
|
@app.function(
|
|
image=image,
|
|
gpu="H100",
|
|
volumes={CACHE_DIR: hf_cache_vol, WORK_DIR: work_vol},
|
|
secrets=[modal.Secret.from_name("hf-secret")],
|
|
timeout=125 * 60,
|
|
cpu=10,
|
|
memory=98304,
|
|
)
|
|
def train_eval_upload(
|
|
eval_limit: int = 6,
|
|
eval_samples: int = 6,
|
|
train_groups_per_stage: int = 4,
|
|
rollouts: int = 12,
|
|
max_new_tokens: int = 1024,
|
|
sft_examples: int = 96,
|
|
finish_by_epoch: float = 0.0,
|
|
):
|
|
"""Train a checkpointed one-turn agentic coding LoRA with PKPO.
|
|
|
|
A one-turn episode is intentional here: it uses the same ``Episode`` parser
|
|
and final-action contract as the shipped tool while making the short budget
|
|
practical. Every generated completion is durably written to a private text
|
|
trace and JSONL before it is judged.
|
|
"""
|
|
import hashlib
|
|
import numpy as np
|
|
import torch
|
|
from huggingface_hub import HfApi, snapshot_download
|
|
from peft import LoraConfig, get_peft_model
|
|
from torch.nn.utils import clip_grad_norm_
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
from agent_core import (CHAT_TEMPLATE, STOP_STRING, THINK_PREFIX, Episode,
|
|
extract_code, parse_action, render_messages)
|
|
from eval_lib import (compact_json_dump, judge_final_answer, load_codecontest_train,
|
|
load_lcb_v6_subset, python3_verified_solutions,
|
|
short_completion_for_sft)
|
|
from pkpo import transform_rewards
|
|
|
|
torch.backends.cuda.matmul.allow_tf32 = True
|
|
started = time.time()
|
|
# Wall-clock guard: the final eval + merge + save + upload need a fixed
|
|
# reserve that SFT/RL must never eat into. finish_by_epoch is when the
|
|
# merged upload must be DONE (epoch seconds).
|
|
finish_by = finish_by_epoch if finish_by_epoch > started else started + 100 * 60
|
|
FINAL_RESERVE_S = 30 * 60
|
|
rl_stop_at = finish_by - FINAL_RESERVE_S
|
|
print("time_guard config", {"now": int(started), "finish_by": int(finish_by),
|
|
"rl_stop_at": int(rl_stop_at)}, flush=True)
|
|
run_id = time.strftime("pkpo_%Y%m%dT%H%M%SZ", time.gmtime())
|
|
run_dir = Path(WORK_DIR) / "results" / run_id
|
|
checkpoint_root = Path(WORK_DIR) / "checkpoints" / run_id
|
|
run_dir.mkdir(parents=True, exist_ok=False)
|
|
checkpoint_root.mkdir(parents=True, exist_ok=True)
|
|
|
|
def save_json(name: str, value):
|
|
compact_json_dump(run_dir / name, value)
|
|
|
|
def safe_verdict(verdict: dict) -> dict:
|
|
"""Keep useful diagnostics without saving hidden test values."""
|
|
out = {
|
|
"passed": bool(verdict.get("passed", False)),
|
|
"error": verdict.get("error", "ok"),
|
|
"passed_tests": int(verdict.get("passed_tests", 0)),
|
|
"total_tests": int(verdict.get("total_tests", 0)),
|
|
}
|
|
for key in ("seconds", "exit_code"):
|
|
if key in verdict:
|
|
out[key] = verdict[key]
|
|
if verdict.get("stderr"):
|
|
out["stderr_tail"] = str(verdict["stderr"])[-500:]
|
|
return out
|
|
|
|
class RolloutTrace:
|
|
"""Private structured and plain-text traces, flushed per sampled output."""
|
|
def __init__(self, root: Path):
|
|
self.jsonl = (root / "rollouts.jsonl").open("a", encoding="utf-8")
|
|
self.text = (root / "raw_rollouts.txt").open("a", encoding="utf-8")
|
|
|
|
def _write(self, record: dict):
|
|
self.jsonl.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
|
|
self.jsonl.flush()
|
|
|
|
def rollout(self, record: dict):
|
|
self._write(record)
|
|
header = {
|
|
key: record.get(key)
|
|
for key in ("run_id", "phase", "stage", "k", "group", "candidate_attempt",
|
|
"sample", "problem_id", "tool_type", "stop_seen")
|
|
}
|
|
self.text.write("\n===== ROLLOUT =====\n")
|
|
self.text.write(json.dumps(header, sort_keys=True) + "\n")
|
|
self.text.write("--- completion_raw ---\n")
|
|
self.text.write(record.get("completion_raw", "") + "\n")
|
|
self.text.write("--- final_answer ---\n")
|
|
self.text.write(record.get("final_answer", "") + "\n")
|
|
self.text.write("--- verdict ---\n")
|
|
self.text.write(json.dumps(record.get("verdict", {}), ensure_ascii=False, sort_keys=True) + "\n")
|
|
self.text.write("===== END ROLLOUT =====\n")
|
|
self.text.flush()
|
|
|
|
def summary(self, record: dict):
|
|
self._write(record)
|
|
self.text.write("\n===== GROUP SUMMARY =====\n")
|
|
self.text.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
|
|
self.text.write("===== END GROUP SUMMARY =====\n")
|
|
self.text.flush()
|
|
|
|
def close(self):
|
|
self.jsonl.close()
|
|
self.text.close()
|
|
|
|
trace = RolloutTrace(run_dir)
|
|
|
|
base_path = snapshot_download(BASE_MODEL, cache_dir=CACHE_DIR)
|
|
tok = AutoTokenizer.from_pretrained(base_path, trust_remote_code=True)
|
|
tok.chat_template = CHAT_TEMPLATE
|
|
if tok.pad_token is None:
|
|
tok.pad_token = tok.eos_token
|
|
tok.padding_side = "left"
|
|
|
|
def load_model():
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
base_path,
|
|
torch_dtype=torch.bfloat16,
|
|
device_map={"": 0},
|
|
trust_remote_code=True,
|
|
low_cpu_mem_usage=True,
|
|
)
|
|
model.config.use_cache = False
|
|
model.gradient_checkpointing_enable()
|
|
lora = LoraConfig(
|
|
r=32,
|
|
lora_alpha=64,
|
|
lora_dropout=0.02,
|
|
bias="none",
|
|
task_type="CAUSAL_LM",
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
|
)
|
|
model = get_peft_model(model, lora)
|
|
model.enable_input_require_grads()
|
|
model.print_trainable_parameters()
|
|
return model
|
|
|
|
model = load_model()
|
|
|
|
def generate_many(first_user_msg: str, count: int) -> list[str]:
|
|
"""Sample a group together; each row has an independent stop string."""
|
|
if count <= 0:
|
|
return []
|
|
prompt = render_messages([{"role": "user", "content": first_user_msg}], add_generation_prompt=True)
|
|
inputs = tok([prompt] * count, return_tensors="pt", padding=True).to(model.device)
|
|
input_width = inputs.input_ids.shape[1]
|
|
was_training = model.training
|
|
old_cache = getattr(model.config, "use_cache", False)
|
|
model.eval()
|
|
model.config.use_cache = True
|
|
try:
|
|
with torch.no_grad():
|
|
out = model.generate(
|
|
**inputs,
|
|
do_sample=True,
|
|
temperature=1.0,
|
|
top_p=0.95,
|
|
max_new_tokens=max_new_tokens,
|
|
pad_token_id=tok.pad_token_id,
|
|
eos_token_id=tok.eos_token_id,
|
|
use_cache=True,
|
|
stop_strings=[STOP_STRING],
|
|
tokenizer=tok,
|
|
)
|
|
finally:
|
|
model.config.use_cache = old_cache
|
|
if was_training:
|
|
model.train()
|
|
return [tok.decode(row[input_width:], skip_special_tokens=False) for row in out]
|
|
|
|
eval_subset = load_lcb_v6_subset(limit=eval_limit, seed=7341)
|
|
save_json("eval_subset.json", [p.to_public_dict() for p in eval_subset])
|
|
|
|
def sample_episodes(problem, count: int, phase: str, stage=None, k=None, group=None,
|
|
candidate_attempt=None, max_tests: int = 18):
|
|
"""Drive the exact one-turn ``Episode`` state machine for every sample."""
|
|
first_user_msg = problem.first_user_message(max_turns=1)
|
|
prompt_hash = hashlib.sha256(
|
|
render_messages([{"role": "user", "content": first_user_msg}], add_generation_prompt=True).encode()
|
|
).hexdigest()
|
|
completions = generate_many(first_user_msg, count)
|
|
samples = []
|
|
for sample_idx, completion in enumerate(completions):
|
|
episode = Episode(first_user_msg, max_turns=1)
|
|
kind, _ = episode.step(completion)
|
|
tool_type, _, parse_error = parse_action(completion)
|
|
stop_seen = STOP_STRING in completion
|
|
valid_final = bool(stop_seen and kind == "final" and episode.final_answer is not None)
|
|
if valid_final:
|
|
final_answer = episode.final_answer
|
|
extracted = extract_code(final_answer)
|
|
verdict = judge_final_answer(
|
|
final_answer, problem.hidden_tests, timeout_s=3, memory_mb=768, max_tests=max_tests
|
|
)
|
|
else:
|
|
final_answer = episode.final_answer or ""
|
|
extracted = ""
|
|
reason = parse_error or episode.end_reason or "missing_final_action"
|
|
verdict = {
|
|
"passed": False,
|
|
"passed_tests": 0,
|
|
"total_tests": min(max_tests, len(problem.hidden_tests)),
|
|
"error": "format_error: " + str(reason),
|
|
}
|
|
safe = safe_verdict(verdict)
|
|
record = {
|
|
"schema_version": 1,
|
|
"record_type": "rollout",
|
|
"run_id": run_id,
|
|
"timestamp_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
|
"phase": phase,
|
|
"stage": stage,
|
|
"k": k,
|
|
"group": group,
|
|
"candidate_attempt": candidate_attempt,
|
|
"sample": sample_idx,
|
|
"problem_id": problem.problem_id,
|
|
"source": problem.source,
|
|
"sampling": {"temperature": 1.0, "top_p": 0.95, "max_new_tokens": max_new_tokens},
|
|
"prompt_sha256": prompt_hash,
|
|
"prompt_chars": len(first_user_msg),
|
|
"completion_raw": completion,
|
|
"assistant_content": THINK_PREFIX + completion,
|
|
"stop_seen": stop_seen,
|
|
"tool_type": tool_type,
|
|
"parse_error": parse_error,
|
|
"final_answer": final_answer,
|
|
"extracted_code": extracted,
|
|
"verdict": safe,
|
|
}
|
|
trace.rollout(record)
|
|
samples.append({
|
|
"completion": completion,
|
|
"valid_final": valid_final,
|
|
"verdict": verdict,
|
|
"safe_verdict": safe,
|
|
"code": extracted,
|
|
"tool_type": tool_type,
|
|
})
|
|
print(phase, problem.problem_id, sample_idx + 1,
|
|
"pass" if verdict.get("passed") else verdict.get("error", "fail"), flush=True)
|
|
return samples
|
|
|
|
def evaluate(tag: str) -> dict:
|
|
details = []
|
|
pass_count = 0
|
|
total = 0
|
|
for p in eval_subset:
|
|
sample_results = []
|
|
samples = sample_episodes(p, eval_samples, phase=tag, max_tests=80)
|
|
for sample in samples:
|
|
verdict = sample["verdict"]
|
|
ok = bool(verdict["passed"])
|
|
pass_count += int(ok)
|
|
total += 1
|
|
sample_results.append({
|
|
"passed": ok,
|
|
"end": sample["safe_verdict"].get("error", "ok"),
|
|
"completion_chars": len(sample["completion"]),
|
|
"code_chars": len(sample["code"]),
|
|
"stop_seen": STOP_STRING in sample["completion"],
|
|
"tool_type": sample["tool_type"],
|
|
})
|
|
details.append({"problem_id": p.problem_id, "difficulty": p.difficulty, "samples": sample_results})
|
|
work_vol.commit()
|
|
res = {
|
|
"tag": tag,
|
|
"pass_at_1": pass_count / max(total, 1),
|
|
"passed_samples": pass_count,
|
|
"total_samples": total,
|
|
"eval_limit": eval_limit,
|
|
"samples_per_problem": eval_samples,
|
|
"temperature": 1.0,
|
|
"top_p": 0.95,
|
|
"max_new_tokens": max_new_tokens,
|
|
"tool_turn_limit": 1,
|
|
"details": details,
|
|
}
|
|
save_json(f"{tag}_results.json", res)
|
|
work_vol.commit()
|
|
print(tag, "pass@1", res["pass_at_1"], flush=True)
|
|
return res
|
|
|
|
def checkpoint(tag: str, state: dict):
|
|
"""Persist the small LoRA adapter after every recoverable unit of work."""
|
|
path = checkpoint_root / tag
|
|
if path.exists():
|
|
shutil.rmtree(path)
|
|
path.mkdir(parents=True)
|
|
model.save_pretrained(path)
|
|
tok.save_pretrained(path)
|
|
compact_json_dump(path / "state.json", state)
|
|
compact_json_dump(run_dir / "latest_state.json", {"checkpoint": str(path), **state})
|
|
work_vol.commit()
|
|
return str(path)
|
|
|
|
def train_completion(first_user_msg: str, completion: str, weight: float):
|
|
"""Backprop only a complete assistant final turn; never truncate silently."""
|
|
if abs(weight) < 1e-9:
|
|
return None, 0.0, "zero_weight"
|
|
stop_idx = completion.find(STOP_STRING)
|
|
if stop_idx == -1:
|
|
return None, 0.0, "missing_stop"
|
|
prompt = render_messages([{"role": "user", "content": first_user_msg}], add_generation_prompt=True)
|
|
prompt_ids = tok(prompt, return_tensors="pt", add_special_tokens=False).input_ids[0]
|
|
# Batched generate right-pads finished rows with eos/pad until the longest
|
|
# row stops; training on that tail drowns the real signal. The target is
|
|
# exactly the turn content through </answer> plus ONE terminating EOS.
|
|
target = completion[: stop_idx + len(STOP_STRING)]
|
|
if tok.eos_token:
|
|
target += tok.eos_token
|
|
full = tok(prompt + target, return_tensors="pt", add_special_tokens=False)
|
|
if full.input_ids.shape[1] > 8192:
|
|
return None, 0.0, "overlength"
|
|
input_ids = full.input_ids.to(model.device)
|
|
labels = input_ids.clone()
|
|
labels[:, : min(prompt_ids.numel(), labels.shape[1])] = -100
|
|
if (labels != -100).sum() == 0:
|
|
return None, 0.0, "empty_target"
|
|
out = model(input_ids=input_ids, labels=labels)
|
|
nll = float(out.loss.detach().cpu())
|
|
weighted_loss = out.loss * float(weight)
|
|
weighted_loss.backward()
|
|
return nll, float(weighted_loss.detach().cpu()), None
|
|
|
|
baseline = evaluate("baseline")
|
|
opt = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=4e-5, betas=(0.9, 0.95))
|
|
|
|
print("verified Python-3 format warmup", flush=True)
|
|
model.train()
|
|
sft_log = []
|
|
sft_pairs = python3_verified_solutions(limit=sft_examples)
|
|
sft_time_floor = rl_stop_at - 25 * 60 # keep at least ~25 min of RL window
|
|
for sft_idx, (prob, sol) in enumerate(sft_pairs, start=1):
|
|
if time.time() > sft_time_floor:
|
|
print("time_guard: stopping SFT early at", sft_idx - 1, "examples", flush=True)
|
|
break
|
|
opt.zero_grad(set_to_none=True)
|
|
nll, loss, skipped = train_completion(
|
|
prob.first_user_message(max_turns=1), short_completion_for_sft(sol), 1.0
|
|
)
|
|
entry = {"index": sft_idx, "problem_id": prob.problem_id, "nll": nll, "loss": loss,
|
|
"skipped": skipped}
|
|
if skipped is None:
|
|
entry["grad_norm"] = float(clip_grad_norm_(model.parameters(), 1.0))
|
|
opt.step()
|
|
sft_log.append(entry)
|
|
print("sft", prob.problem_id, entry, flush=True)
|
|
if sft_idx % 24 == 0:
|
|
save_json("sft_log.json", sft_log)
|
|
checkpoint(f"sft_{sft_idx:04d}", {"phase": "sft", "completed_examples": sft_idx})
|
|
save_json("sft_log.json", sft_log)
|
|
sft_checkpoint = checkpoint("after_sft", {
|
|
"phase": "sft_complete",
|
|
"completed_examples": len(sft_log),
|
|
"verified_examples": len(sft_pairs),
|
|
})
|
|
|
|
print("loading disjoint RL curriculum", flush=True)
|
|
schedule = [1, 8, 1]
|
|
sft_ids = {p.problem_id for p, _ in sft_pairs}
|
|
train_problems = load_codecontest_train(
|
|
limit=max(72, train_groups_per_stage * len(schedule) * 8),
|
|
seed=20260709,
|
|
exclude_problem_ids=sft_ids,
|
|
)
|
|
print("loaded RL candidates", len(train_problems), flush=True)
|
|
stage_log = []
|
|
candidate_cursor = 0
|
|
# Split the remaining RL window across stages so the mandatory final k=1
|
|
# stage always gets its share even when earlier stages run long.
|
|
rl_start = time.time()
|
|
rl_window = max(rl_stop_at - rl_start, 0.0)
|
|
stage_shares = [0.30, 0.40, 0.30]
|
|
assert len(stage_shares) == len(schedule)
|
|
stage_deadline = [rl_start + rl_window * sum(stage_shares[: i + 1]) for i in range(len(schedule))]
|
|
print("rl window minutes", round(rl_window / 60, 1), flush=True)
|
|
for stage_idx, k in enumerate(schedule):
|
|
opt.param_groups[0]["lr"] = 2e-5 if k == 1 else 1e-5
|
|
for group_idx in range(train_groups_per_stage):
|
|
if time.time() > stage_deadline[stage_idx]:
|
|
print("time_guard: closing stage", stage_idx, "k", k, "after", group_idx, "groups", flush=True)
|
|
stage_log.append({"stage": stage_idx, "k": k, "group": group_idx,
|
|
"skipped": "stage_time_guard"})
|
|
save_json("training_stage_log.json", stage_log)
|
|
break
|
|
selected = None
|
|
attempts = []
|
|
for attempt_idx in range(1, 6):
|
|
if candidate_cursor >= len(train_problems) or time.time() > stage_deadline[stage_idx]:
|
|
break
|
|
p = train_problems[candidate_cursor]
|
|
candidate_cursor += 1
|
|
samples = sample_episodes(
|
|
p, rollouts, phase="rl_candidate", stage=stage_idx, k=k, group=group_idx,
|
|
candidate_attempt=attempt_idx,
|
|
)
|
|
rewards = [1.0 if sample["verdict"].get("passed") else 0.0 for sample in samples]
|
|
advantages = transform_rewards(np.array(rewards, dtype=np.float64), k)
|
|
eligible = bool(np.any(np.abs(advantages) > 1e-10))
|
|
attempt = {
|
|
"problem_id": p.problem_id,
|
|
"passes": int(sum(rewards)),
|
|
"rewards": rewards,
|
|
"eligible": eligible,
|
|
"advantages": [float(x) for x in advantages],
|
|
}
|
|
attempts.append(attempt)
|
|
trace.summary({
|
|
"schema_version": 1,
|
|
"record_type": "candidate_summary",
|
|
"run_id": run_id,
|
|
"stage": stage_idx,
|
|
"k": k,
|
|
"group": group_idx,
|
|
"candidate_attempt": attempt_idx,
|
|
**attempt,
|
|
})
|
|
if eligible:
|
|
selected = (p, samples, rewards, advantages)
|
|
break
|
|
|
|
if selected is None:
|
|
entry = {
|
|
"stage": stage_idx,
|
|
"k": k,
|
|
"group": group_idx,
|
|
"skipped": "no_pkpo_eligible_group",
|
|
"candidate_attempts": attempts,
|
|
}
|
|
stage_log.append(entry)
|
|
save_json("training_stage_log.json", stage_log)
|
|
checkpoint(f"stage{stage_idx}_group{group_idx}_skipped", {
|
|
"phase": "rl", "stage": stage_idx, "k": k, "group": group_idx,
|
|
"status": "skipped", "candidate_cursor": candidate_cursor,
|
|
})
|
|
print("stage skipped", entry, flush=True)
|
|
continue
|
|
|
|
p, samples, rewards, adv = selected
|
|
opt.zero_grad(set_to_none=True)
|
|
losses = []
|
|
nlls = []
|
|
skipped_samples = []
|
|
first_user_msg = p.first_user_message(max_turns=1)
|
|
for sample_idx, (sample, a) in enumerate(zip(samples, adv)):
|
|
if not sample["valid_final"]:
|
|
skipped_samples.append({"sample": sample_idx, "reason": "invalid_final"})
|
|
losses.append(0.0)
|
|
nlls.append(None)
|
|
continue
|
|
nll, loss, skipped = train_completion(first_user_msg, sample["completion"], float(a) / rollouts)
|
|
nlls.append(nll)
|
|
losses.append(loss)
|
|
if skipped not in (None, "zero_weight"):
|
|
skipped_samples.append({"sample": sample_idx, "reason": skipped})
|
|
if any(nll is not None for nll in nlls):
|
|
grad_norm = float(clip_grad_norm_(model.parameters(), 1.0))
|
|
opt.step()
|
|
else:
|
|
grad_norm = 0.0
|
|
entry = {
|
|
"stage": stage_idx,
|
|
"k": k,
|
|
"group": group_idx,
|
|
"problem_id": p.problem_id,
|
|
"rewards": rewards,
|
|
"advantages": [float(x) for x in adv],
|
|
"losses": losses,
|
|
"nlls": nlls,
|
|
"skipped_samples": skipped_samples,
|
|
"candidate_attempts": attempts,
|
|
"grad_norm": grad_norm,
|
|
}
|
|
stage_log.append(entry)
|
|
save_json("training_stage_log.json", stage_log)
|
|
ckpt = checkpoint(f"stage{stage_idx}_group{group_idx}", {
|
|
"phase": "rl", "stage": stage_idx, "k": k, "group": group_idx,
|
|
"problem_id": p.problem_id, "candidate_cursor": candidate_cursor,
|
|
"after_sft_checkpoint": sft_checkpoint,
|
|
})
|
|
entry["checkpoint"] = ckpt
|
|
save_json("training_stage_log.json", stage_log)
|
|
trace.summary({
|
|
"schema_version": 1,
|
|
"record_type": "training_group_summary",
|
|
"run_id": run_id,
|
|
**entry,
|
|
})
|
|
print("stage", entry, flush=True)
|
|
|
|
final = evaluate("final")
|
|
|
|
final_dir = Path(WORK_DIR) / "final_model"
|
|
if final_dir.exists():
|
|
shutil.rmtree(final_dir)
|
|
final_dir.mkdir(parents=True)
|
|
model.eval()
|
|
merged = model.merge_and_unload()
|
|
merged.config.use_cache = True
|
|
merged.save_pretrained(final_dir, safe_serialization=True, max_shard_size="4GB")
|
|
tok.save_pretrained(final_dir)
|
|
_write_generation_config(merged, tok, final_dir)
|
|
_copy_source_bundle(final_dir)
|
|
eval_dir = final_dir / "eval"
|
|
eval_dir.mkdir()
|
|
for name in ["baseline_results.json", "final_results.json", "eval_subset.json", "training_stage_log.json", "sft_log.json"]:
|
|
src = run_dir / name
|
|
if src.exists():
|
|
shutil.copyfile(src, eval_dir / name)
|
|
notes = (
|
|
f"Modal H100 pipeline elapsed {(time.time() - started) / 60:.1f} minutes. "
|
|
f"Run id {run_id}; eval subset size {eval_limit}, samples/problem {eval_samples}, "
|
|
f"rollouts/group {rollouts}; verified Python-3 SFT examples {len(sft_pairs)}. "
|
|
"Raw rollout traces remain in the private Modal work volume and are not published."
|
|
)
|
|
(final_dir / "README.md").write_text(_model_card(baseline, final, schedule, notes), encoding="utf-8")
|
|
HfApi().upload_folder(repo_id=HF_REPO, repo_type="model", folder_path=str(final_dir), commit_message="Upload merged PKPO run")
|
|
trace.close()
|
|
work_vol.commit()
|
|
return {
|
|
"run_id": run_id,
|
|
"baseline": baseline["pass_at_1"],
|
|
"final": final["pass_at_1"],
|
|
"model_dir": str(final_dir),
|
|
"trace_text": str(run_dir / "raw_rollouts.txt"),
|
|
}
|
|
|
|
|
|
@app.local_entrypoint()
|
|
def main(
|
|
eval_limit: int = 6,
|
|
eval_samples: int = 6,
|
|
train_groups_per_stage: int = 4,
|
|
rollouts: int = 12,
|
|
max_new_tokens: int = 1024,
|
|
sft_examples: int = 96,
|
|
finish_by_epoch: float = 0.0,
|
|
):
|
|
print("Spawning scaffold push (runs in parallel on CPU)...")
|
|
scaffold_call = push_scaffold.spawn()
|
|
print("Running GPU pipeline...")
|
|
print(train_eval_upload.remote(
|
|
eval_limit,
|
|
eval_samples,
|
|
train_groups_per_stage,
|
|
rollouts,
|
|
max_new_tokens,
|
|
sft_examples,
|
|
finish_by_epoch,
|
|
))
|
|
print("scaffold:", scaffold_call.get())
|