Download scripts/101_train_14b_canary_validation.py from koreallmdev/dgx-harness-engineering: direct link, hf CLI and curl.
- Browser
- Download file 15.1 kB
-
https://huggingface.co/koreallmdev/dgx-harness-engineering/resolve/main/scripts/101_train_14b_canary_validation.py
- Command line
-
hf download hf://koreallmdev/dgx-harness-engineering/scripts/101_train_14b_canary_validation.py
-
curl -L -o 101_train_14b_canary_validation.py https://huggingface.co/koreallmdev/dgx-harness-engineering/resolve/main/scripts/101_train_14b_canary_validation.py
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 | |