dgx-harness-engineering / scripts /101_train_14b_canary_validation.py
koreallmdev's picture
Upload anonymized DGX harness engineering export V1.4 (part 2)
6fda4b7 verified
Raw History Blame Contribute Delete
15.1 kB
#!/usr/bin/env python3
# DGX AI Factory - 14B canary validation training
#
# Goal:
# - Start a small 14B validation LoRA from base 14B model.
# - Do NOT reuse 8B LoRA adapters because 8B adapters are incompatible with 14B.
# - Train only a small canary run.
# - Save under experiments/14b_canary_validation_*.
# - Run smoke test immediately after training.
# - Never overwrite 8B outputs or stable adapters.
import os
import re
import json
import time
import random
import traceback
from pathlib import Path
from typing import Any, Dict, Optional, List
ROOT = Path("/home/harness_user_1/dgx_ai_factory")
EXPERIMENT_ROOT = ROOT / "experiments"
REPORT_DIR = ROOT / "reports"
LOG_DIR = ROOT / "logs"
for d in [EXPERIMENT_ROOT, REPORT_DIR, LOG_DIR]:
d.mkdir(parents=True, exist_ok=True)
BASE_MODEL = os.environ.get("BASE_MODEL_14B", "Qwen/Qwen2.5-14B-Instruct")
MAX_STEPS = int(os.environ.get("MAX_STEPS", "80"))
LR = float(os.environ.get("LR", "5e-6"))
MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "1024"))
SEED = int(os.environ.get("SEED", "42"))
DATASET_CANDIDATES = [
Path(os.environ.get("TRAIN_JSONL", "")) if os.environ.get("TRAIN_JSONL") else None,
ROOT / "data/train_100k_gold/train_canary_1000_gold_clean.jsonl",
ROOT / "data/train_100k/train_canary_1000.jsonl",
ROOT / "dataset_quality_2000_candidate_clean.jsonl",
ROOT / "dataset.jsonl",
]
SYSTEM_PROMPT = (
"๋‹น์‹ ์€ ํ•œ๊ตญ์–ด ์กด๋Œ“๋ง์„ ์‚ฌ์šฉํ•˜๋Š” ์‹ค๋ฌดํ˜• ๊ฐœ๋ฐœ ๋น„์„œ์ž…๋‹ˆ๋‹ค. "
"์ค‘๊ตญ์–ด๋ฅผ ์ ˆ๋Œ€ ์“ฐ์ง€ ๋ง๊ณ , ์งˆ๋ฌธ๊ณผ ๋ฌด๊ด€ํ•œ ๋‚ด์šฉ์„ ๋ฐ˜๋ณตํ•˜์ง€ ๋ง๊ณ , ๋ช…๋ น์–ด์™€ ์„ค๋ช…์„ ๋ถ„๋ฆฌํ•ด์„œ ์ •ํ™•ํ•˜๊ฒŒ ๋‹ต๋ณ€ํ•ฉ๋‹ˆ๋‹ค."
)
SMOKE_TESTS = [
{
"name": "vllm_stop",
"prompt": "vLLM ์„œ๋ฒ„๊ฐ€ ์ผœ์ ธ ์žˆ๋Š”์ง€ ํ™•์ธํ•˜๊ณ  ๊บผ์•ผ ํ•  ๋•Œ ์“ฐ๋Š” ๋ช…๋ น์–ด๋ฅผ ์•Œ๋ ค์ฃผ์„ธ์š”.",
"must": ["ps", "grep", "vllm"],
"any_groups": [["pkill", "kill"], ["8000", "api_server", "vllm.entrypoints"]],
},
{
"name": "systemd_217_user",
"prompt": "systemd ์„œ๋น„์Šค์—์„œ status=217/USER ์˜ค๋ฅ˜๊ฐ€ ๋‚ฌ์Šต๋‹ˆ๋‹ค. ์›์ธ๊ณผ ํ™•์ธ ๋ช…๋ น์–ด๋ฅผ ์•Œ๋ ค์ฃผ์„ธ์š”.",
"must": ["217/USER", "User", "systemctl", "journalctl"],
"any_groups": [["getent passwd", "id "], ["daemon-reload", "restart"]],
},
{
"name": "fastapi_curl",
"prompt": "FastAPI ํ—ฌ์Šค์ฒดํฌ ์—”๋“œํฌ์ธํŠธ๋ฅผ ๋งŒ๋“ค๊ณ  curl๋กœ ํ™•์ธํ•˜๋Š” ์˜ˆ์‹œ๋ฅผ ๋ณด์—ฌ์ฃผ์„ธ์š”.",
"must": ["FastAPI", "curl", "/health"],
"any_groups": [],
},
{
"name": "jsonl_validate",
"prompt": "JSONL ํŒŒ์ผ์ด ๊นจ์กŒ๋Š”์ง€ ํŒŒ์ด์ฌ์œผ๋กœ ๊ฒ€์‚ฌํ•˜๋Š” ์งง์€ ์ฝ”๋“œ๋ฅผ ๋ณด์—ฌ์ฃผ์„ธ์š”.",
"must": ["json", "enumerate", "json.loads"],
"any_groups": [["JSONDecodeError", "Exception"], ["line_no", "line", "์ค„"]],
},
{
"name": "lora_next_step",
"prompt": "LoRA ํŒŒ์ธํŠœ๋‹์ด ๋๋‚œ ๋’ค smoke test์™€ benchmark๋ฅผ ์–ด๋–ค ์ˆœ์„œ๋กœ ํ•ด์•ผ ํ•˜๋Š”์ง€ ์ •๋ฆฌํ•ด์ฃผ์„ธ์š”.",
"must": ["smoke", "benchmark"],
"any_groups": [["stable", "์Šน๊ฒฉ"], ["rejected", "reject", "ํ๊ธฐ", "์ˆ˜๋ฆฌ", "repair"]],
},
]
def has_chinese_ideograph(text: str) -> bool:
return bool(re.search(r"[\u4e00-\u9fff]", text))
def template_leak(text: str) -> bool:
bad = [
"{problem}", "{framework}", "{concept}", "{bug}", "{system}",
"### Instruction", "### Response", "<|im_start|>", "<|im_end|>",
"assistant:", "user:"
]
low = text.lower()
return any(x.lower() in low for x in bad)
def honorific_ok(text: str) -> bool:
return any(m in text for m in ["์Šต๋‹ˆ๋‹ค", "ํ•ฉ๋‹ˆ๋‹ค", "๋ฉ๋‹ˆ๋‹ค", "์ฃผ์„ธ์š”", "ํ•˜์„ธ์š”", "ํ™•์ธํ•˜์‹œ๋ฉด", "์ข‹์Šต๋‹ˆ๋‹ค", "ํ•˜์‹œ๋ฉด"])
def pick_dataset() -> Path:
for p in DATASET_CANDIDATES:
if p and p.exists() and p.stat().st_size > 0:
return p
node-7.example.invalid SystemExit("[FAIL] No usable canary dataset found. Expected data/train_100k_gold/train_canary_1000_gold_clean.jsonl")
def normalize_row(obj: Any) -> Optional[Dict[str, str]]:
if not isinstance(obj, dict):
return None
msgs = obj.get("messages")
if isinstance(msgs, list):
user_parts, assistant_parts = [], []
for m in msgs:
if not isinstance(m, dict):
continue
role = str(m.get("role", "")).lower()
content = str(m.get("content", "") or "").strip()
if not content:
continue
if role == "user":
user_parts.append(content)
elif role == "assistant":
assistant_parts.append(content)
if user_parts and assistant_parts:
return {
"category": str(obj.get("category", "canary")),
"instruction": "\n\n".join(user_parts).strip(),
"input": "",
"output": "\n\n".join(assistant_parts).strip(),
}
instruction = obj.get("instruction") or obj.get("prompt") or obj.get("question") or obj.get("input_text")
output = obj.get("output") or obj.get("response") or obj.get("answer") or obj.get("completion")
if instruction and output:
return {
"category": str(obj.get("category", "canary")),
"instruction": str(instruction).strip(),
"input": str(obj.get("input", "") or "").strip(),
"output": str(output).strip(),
}
return None
def load_rows(path: Path, limit: int = 1000) -> List[Dict[str, str]]:
rows = []
bad_json = 0
skipped = 0
with open(path, "r", encoding="utf-8") as f:
for line_no, line in enumerate(f, 1):
line = line.strip()
if not line:
continue
try:
obj = json.loads(line)
except Exception:
bad_json += 1
continue
row = normalize_row(obj)
if not row:
skipped += 1
continue
if len(row["instruction"]) < 4 or len(row["output"]) < 20:
skipped += 1
continue
if has_chinese_ideograph(row["output"]):
skipped += 1
continue
rows.append(row)
if len(rows) >= limit:
break
print("==== DATASET LOAD ====")
print("dataset:", path)
print("loaded_rows:", len(rows))
print("bad_json:", bad_json)
print("skipped:", skipped)
if len(rows) < 100:
node-7.example.invalid SystemExit("[FAIL] Too few usable canary rows. Need at least 100.")
return rows
def format_chat(tokenizer, row: Dict[str, str]) -> str:
user = row["instruction"]
if row.get("input"):
user += "\n\n์ž…๋ ฅ:\n" + row["input"]
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user},
{"role": "assistant", "content": row["output"]},
]
return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
def build_dataset(tokenizer, rows: List[Dict[str, str]], max_length: int):
from datasets import Dataset
ds = Dataset.from_list(rows)
def tok(row):
text = format_chat(tokenizer, row)
enc = tokenizer(
text,
truncation=True,
max_length=max_length,
padding=False,
)
enc["labels"] = enc["input_ids"].copy()
return enc
return ds.map(tok, remove_columns=ds.column_names)
def generate_one(tokenizer, model, prompt: str, max_new_tokens: int = 320) -> str:
import torch
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": prompt},
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
with torch.no_grad():
out = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=False,
repetition_penalty=1.08,
pad_token_id=tokenizer.eos_token_id,
)
gen = out[0][inputs["input_ids"].shape[-1]:]
return tokenizer.decode(gen, skip_special_tokens=True).strip()
def score_smoke_output(output: str, test: Dict[str, Any]) -> Dict[str, Any]:
score = 100
flags = []
low = output.lower()
if has_chinese_ideograph(output):
score -= 60
flags.append("chinese_cjk_leakage")
if template_leak(output):
score -= 35
flags.append("template_or_role_leakage")
if not honorific_ok(output):
score -= 15
flags.append("weak_korean_honorific")
if len(output.strip()) < 60:
score -= 15
flags.append("too_short")
for item in test["must"]:
if item.lower() not in low:
score -= 18
flags.append("missing:" + item)
for group in test.get("any_groups", []):
if not any(x.lower() in low for x in group):
score -= 10
flags.append("missing_any:" + "|".join(group))
score = max(0, score)
return {
"score": score,
"status": "PASS" if score >= 70 else "FAIL",
"flags": flags,
}
def run_smoke(tokenizer, model, adapter_dir: Path, exp_dir: Path, dataset_path: Path):
print("\n" + "=" * 80)
print("14B CANARY SMOKE TEST")
print("=" * 80)
results = []
pass_count = 0
model.eval()
for test in SMOKE_TESTS:
print("\n" + "-" * 80)
print("TEST:", test["name"])
print("-" * 80)
output = generate_one(tokenizer, model, test["prompt"])
scored = score_smoke_output(output, test)
if scored["status"] == "PASS":
pass_count += 1
print(output)
print("\nSCORE:", scored["score"])
print("STATUS:", scored["status"])
if scored["flags"]:
print("FLAGS:", scored["flags"])
results.append({
"name": test["name"],
"prompt": test["prompt"],
"output": output,
**scored,
})
avg = sum(r["score"] for r in results) / len(results)
decision = "PASS_14B_CANARY_CANDIDATE" if pass_count == len(results) else ("REVIEW_14B_CANARY" if pass_count >= 4 else "REJECT_OR_REPAIR_NEEDED")
report = {
"timestamp": time.strftime("%Y%m%d_%H%M%S"),
"base_model": BASE_MODEL,
"adapter_dir": str(adapter_dir),
"experiment_dir": str(exp_dir),
"dataset_path": str(dataset_path),
"average_score": round(avg, 2),
"smoke_pass": pass_count,
"smoke_total": len(results),
"decision": decision,
"results": results,
}
report_path = exp_dir / "canary_smoke_report.json"
report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
global_report_path = REPORT_DIR / f"14b_canary_smoke_{report['timestamp']}.json"
global_report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
print("\n" + "=" * 80)
print("FINAL 14B CANARY SUMMARY")
print("=" * 80)
print("base_model:", BASE_MODEL)
print("adapter_dir:", adapter_dir)
print("average_score:", round(avg, 2))
print("smoke_pass:", f"{pass_count}/{len(results)}")
print("decision:", decision)
print("report_path:", report_path)
print("global_report_path:", global_report_path)
def main():
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForLanguageModeling
from peft import LoraConfig, get_peft_model, TaskType
random.seed(SEED)
torch.manual_seed(SEED)
dataset_path = pick_dataset()
ts = time.strftime("%Y%m%d_%H%M%S")
exp_dir = EXPERIMENT_ROOT / f"14b_canary_validation_{ts}"
adapter_dir = exp_dir / "adapter"
exp_dir.mkdir(parents=True, exist_ok=True)
config = {
"base_model": BASE_MODEL,
"dataset_path": str(dataset_path),
"max_steps": MAX_STEPS,
"learning_rate": LR,
"max_length": MAX_LENGTH,
"seed": SEED,
"lora_r": 16,
"lora_alpha": 32,
"lora_dropout": 0.05,
"target_modules": ["q_proj", "v_proj"],
"note": "14B canary validation. 8B adapters are not reused.",
}
(exp_dir / "config.json").write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding="utf-8")
print("==== 14B CANARY VALIDATION START ====")
print("base_model:", BASE_MODEL)
print("dataset_path:", dataset_path)
print("experiment_dir:", exp_dir)
print("adapter_dir:", adapter_dir)
print("max_steps:", MAX_STEPS)
print("lr:", LR)
print("max_length:", MAX_LENGTH)
rows = load_rows(dataset_path, limit=1000)
print("\n==== LOAD TOKENIZER ====")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = <REDACTED_SECRET>
tokenized = build_dataset(tokenizer, rows, MAX_LENGTH)
print("\n==== LOAD 14B BASE MODEL ====")
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
torch_dtype=dtype,
device_map="auto",
trust_remote_code=True,
)
model.config.use_cache = False
if hasattr(model, "gradient_checkpointing_enable"):
model.gradient_checkpointing_enable()
print("\n==== APPLY LORA ====")
lora_cfg = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=16,
lora_alpha=32,
lora_dropout=0.05,
target_modules=["q_proj", "v_proj"],
bias="none",
)
model = get_peft_model(model, lora_cfg)
model.print_trainable_parameters()
collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
args = TrainingArguments(
output_dir=str(adapter_dir),
per_device_train_batch_size=1,
gradient_accumulation_steps=8,
learning_rate=LR,
max_steps=MAX_STEPS,
logging_steps=5,
save_steps=MAX_STEPS,
save_total_limit=1,
bf16=torch.cuda.is_available(),
fp16=False,
gradient_checkpointing=True,
report_to=[],
remove_unused_columns=False,
optim="adamw_torch",
)
trainer = Trainer(
model=model,
args=args,
train_dataset=tokenized,
data_collator=collator,
)
print("\n==== TRAIN 14B CANARY ====")
train_result = trainer.train()
metrics = train_result.metrics or {}
(exp_dir / "train_metrics.json").write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")
print("\n==== SAVE 14B CANARY ADAPTER ====")
model.save_pretrained(str(adapter_dir))
tokenizer.save_pretrained(str(adapter_dir))
run_smoke(tokenizer, model, adapter_dir, exp_dir, dataset_path)
if __name__ == "__main__":
try:
main()
except Exception:
print("\n[FATAL ERROR]")
traceback.print_exc()
node-7.example.invalid