| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import os |
| import random |
| import sys |
| import time |
| from dataclasses import dataclass |
| from datetime import timedelta |
| from typing import Any, Dict, List |
|
|
| import torch |
| import torch.distributed as dist |
| from peft import LoraConfig, PeftModel, get_peft_model |
| from torch.optim import AdamW |
| from torch.utils.data import DistributedSampler |
| from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed |
|
|
| CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| PARENT_DIR = os.path.dirname(CURRENT_DIR) |
| if PARENT_DIR not in sys.path: |
| sys.path.insert(0, PARENT_DIR) |
|
|
| from aligned_cell_policy.shared_cell_policy import build_cell_examples_from_row |
| from checkpoint_utils import ensure_final_checkpoint_dir, save_checkpoint_and_update_final |
| from multi_output_cell_policy.prompt_builder import build_multi_output_cell_prompt |
| from multi_output_cell_policy.rewards import score_prediction_text |
| from multi_output_cell_policy.shared_multi_output_policy import ( |
| batched_completion_ce_loss, |
| build_supervised_completion, |
| completion_ce_loss, |
| make_solved_grid_from_row, |
| stage_i_consistent_values, |
| ) |
|
|
|
|
| try: |
| import wandb |
| except Exception: |
| wandb = None |
|
|
|
|
| @dataclass |
| class Args: |
| model_name: str |
| train_jsonl: str |
| eval_jsonl: str |
| output_dir: str |
| cache_dir: str |
| init_adapter_dir: str |
| seed: int |
| gpu_id: int |
| stage_i: int |
| total_empties_hint: int |
| per_device_train_batch_size: int |
| gradient_accumulation_steps: int |
| num_epochs: float |
| learning_rate: float |
| weight_decay: float |
| max_grad_norm: float |
| enable_gradient_checkpointing: bool |
| logging_steps: int |
| save_steps: int |
| eval_steps: int |
| eval_rows: int |
| max_completion_length: int |
| lora_r: int |
| lora_alpha: int |
| lora_dropout: float |
| use_wandb: bool |
| wandb_entity: str |
| wandb_project: str |
| wandb_run_name: str |
| wandb_mode: str |
| debug_print_limit: int |
| limit_train_rows: int |
| eval_exact_set_match_stop: float |
| eval_value_precision_stop: float |
| eval_value_recall_stop: float |
| eval_solve_rate_stop: float |
| min_steps_before_stop: int |
| max_wall_clock_seconds: int |
| max_steps: int |
| multi_value_oversample_factor: int |
| train_target_size_min: int |
| train_target_size_max: int |
| eval_target_size_min: int |
| eval_target_size_max: int |
|
|
|
|
| def configure_hf_cache(cache_dir: str) -> str: |
| cache_dir = os.path.abspath(os.path.expanduser(cache_dir)) |
| hub_dir = os.path.join(cache_dir, "hub") |
| transformers_dir = os.path.join(cache_dir, "transformers") |
| os.makedirs(hub_dir, exist_ok=True) |
| os.makedirs(transformers_dir, exist_ok=True) |
| os.environ["HF_HOME"] = cache_dir |
| os.environ["HF_HUB_CACHE"] = hub_dir |
| os.environ["HUGGINGFACE_HUB_CACHE"] = hub_dir |
| os.environ["TRANSFORMERS_CACHE"] = transformers_dir |
| os.environ.setdefault("HF_HUB_DISABLE_XET", "1") |
| return cache_dir |
|
|
|
|
| def configure_wandb_dirs(output_dir: str) -> None: |
| wandb_dir = os.path.join(output_dir, "wandb_runtime") |
| os.makedirs(wandb_dir, exist_ok=True) |
| os.environ.setdefault("WANDB_DIR", wandb_dir) |
| os.environ.setdefault("WANDB_CACHE_DIR", wandb_dir) |
| os.environ.setdefault("WANDB_CONFIG_DIR", wandb_dir) |
|
|
|
|
| def pick_dtype() -> torch.dtype: |
| if torch.cuda.is_available() and torch.cuda.is_bf16_supported(): |
| return torch.bfloat16 |
| return torch.float16 |
|
|
|
|
| def load_jsonl_rows(path: str, limit_rows: int = 0) -> List[Dict[str, Any]]: |
| rows: List[Dict[str, Any]] = [] |
| with open(path, "r", encoding="utf-8") as f: |
| for line in f: |
| line = line.strip() |
| if not line: |
| continue |
| rows.append(json.loads(line)) |
| if limit_rows > 0 and len(rows) >= limit_rows: |
| break |
| return rows |
|
|
|
|
| def target_size_allowed(target_size: int, min_size: int, max_size: int) -> bool: |
| if int(min_size) > 0 and int(target_size) < int(min_size): |
| return False |
| if int(max_size) > 0 and int(target_size) > int(max_size): |
| return False |
| return True |
|
|
|
|
| def build_training_examples( |
| rows: List[Dict[str, Any]], |
| *, |
| tokenizer: Any, |
| stage_i: int, |
| total_empties_hint: int, |
| progress_every_rows: int = 10, |
| progress_callback: Any = None, |
| ) -> List[Dict[str, Any]]: |
| examples: List[Dict[str, Any]] = [] |
| eos_text = getattr(tokenizer, "eos_token", None) or "" |
| for row_idx, row in enumerate(rows, start=1): |
| solved = make_solved_grid_from_row(row) |
| for ex in build_cell_examples_from_row(row): |
| target_values = stage_i_consistent_values(ex.grid, target_cell=ex.target_cell, stage_i=stage_i) |
| if not target_size_allowed( |
| len(target_values), |
| getattr(tokenizer, "_train_target_size_min", 0), |
| getattr(tokenizer, "_train_target_size_max", 0), |
| ): |
| continue |
| prompt = build_multi_output_cell_prompt( |
| ex.grid, |
| target_cell=ex.target_cell, |
| stage_i=stage_i, |
| tokenizer=tokenizer, |
| turn_idx=ex.turn_idx, |
| total_turns=ex.total_turns, |
| prev_output_flag=None, |
| total_empties_hint=total_empties_hint, |
| ) |
| target_text = build_supervised_completion(ex, stage_i=stage_i) |
| if eos_text: |
| target_text = target_text + eos_text |
| repeat_count = max(1, int(getattr(tokenizer, "_multi_value_oversample_factor", 1))) if len(target_values) > 1 else 1 |
| for _ in range(repeat_count): |
| examples.append( |
| { |
| "prompt_text": prompt, |
| "completion_text": target_text, |
| "target_values": list(target_values), |
| "grid": ex.grid, |
| "solved": solved, |
| "target_cell": ex.target_cell, |
| } |
| ) |
| if progress_callback is not None and ( |
| row_idx == 1 or row_idx == len(rows) or row_idx % max(1, int(progress_every_rows)) == 0 |
| ): |
| progress_callback(row_idx, len(rows), len(examples)) |
| return examples |
|
|
|
|
| def _prepared_data_dir(args: Args) -> str: |
| path = os.path.join(PARENT_DIR, "_prepared_data", "multi_output_cell_policy") |
| os.makedirs(path, exist_ok=True) |
| return path |
|
|
|
|
| def _prepared_sft_cache_path(args: Args) -> str: |
| payload = json.dumps( |
| { |
| "completion_format_version": 2, |
| "train_jsonl": os.path.abspath(args.train_jsonl), |
| "stage_i": int(args.stage_i), |
| "total_empties_hint": int(args.total_empties_hint), |
| "limit_train_rows": int(args.limit_train_rows), |
| "model_name": str(args.model_name), |
| "multi_value_oversample_factor": int(args.multi_value_oversample_factor), |
| "train_target_size_min": int(args.train_target_size_min), |
| "train_target_size_max": int(args.train_target_size_max), |
| }, |
| sort_keys=True, |
| ).encode("utf-8") |
| digest = hashlib.sha1(payload).hexdigest()[:20] |
| return os.path.join(_prepared_data_dir(args), f"sft_stage{int(args.stage_i):02d}_{digest}.jsonl") |
|
|
|
|
| def _to_jsonable(value: Any) -> Any: |
| if isinstance(value, dict): |
| return {k: _to_jsonable(v) for k, v in value.items()} |
| if isinstance(value, (list, tuple)): |
| return [_to_jsonable(v) for v in value] |
| if hasattr(value, "tolist"): |
| return _to_jsonable(value.tolist()) |
| return value |
|
|
|
|
| def _write_jsonl(path: str, rows: List[Dict[str, Any]]) -> None: |
| tmp_path = f"{path}.tmp" |
| with open(tmp_path, "w", encoding="utf-8") as f: |
| for row in rows: |
| f.write(json.dumps(_to_jsonable(row), separators=(",", ":")) + "\n") |
| os.replace(tmp_path, path) |
|
|
|
|
| def _read_jsonl(path: str) -> List[Dict[str, Any]]: |
| rows: List[Dict[str, Any]] = [] |
| with open(path, "r", encoding="utf-8") as f: |
| for line in f: |
| line = line.strip() |
| if line: |
| rows.append(json.loads(line)) |
| return rows |
|
|
|
|
| def _wait_for_cache(path: str, timeout_s: float = 7200.0) -> None: |
| start = time.time() |
| while not os.path.exists(path): |
| if time.time() - start > timeout_s: |
| raise TimeoutError(f"Timed out waiting for prepared cache: {path}") |
| time.sleep(2.0) |
|
|
|
|
| def load_or_build_sft_examples( |
| args: Args, |
| *, |
| rows: List[Dict[str, Any]], |
| tokenizer: Any, |
| rank: int, |
| world_size: int, |
| progress_callback: Any = None, |
| ) -> List[Dict[str, Any]]: |
| cache_path = _prepared_sft_cache_path(args) |
| if os.path.exists(cache_path): |
| return _read_jsonl(cache_path) |
| if rank == 0: |
| print(f"[dataset build][sft stage {args.stage_i}] building prepared cache: {cache_path}", flush=True) |
| examples = build_training_examples( |
| rows, |
| tokenizer=tokenizer, |
| stage_i=args.stage_i, |
| total_empties_hint=args.total_empties_hint, |
| progress_every_rows=10, |
| progress_callback=progress_callback, |
| ) |
| _write_jsonl(cache_path, examples) |
| return examples |
| _wait_for_cache(cache_path) |
| return _read_jsonl(cache_path) |
|
|
|
|
| @torch.no_grad() |
| def run_eval(args: Args, rows: List[Dict[str, Any]], model: torch.nn.Module, tokenizer: Any, device: torch.device): |
| model.eval() |
| total_cells = 0 |
| parse_ok = 0.0 |
| canonical_ok = 0.0 |
| exact_set_match = 0.0 |
| includes_gt = 0.0 |
| precision_sum = 0.0 |
| recall_sum = 0.0 |
| predicted_size_sum = 0.0 |
| good_count_sum = 0.0 |
| bad_count_sum = 0.0 |
| solve_ok = 0 |
| solve_rows = 0 |
| printed = 0 |
| for row in rows: |
| solved = make_solved_grid_from_row(row) |
| row_all_exact = True |
| row_has_eval_cell = False |
| row_debug_lines: List[str] = [] |
| for ex in build_cell_examples_from_row(row): |
| target_values = stage_i_consistent_values(ex.grid, target_cell=ex.target_cell, stage_i=args.stage_i) |
| if not target_size_allowed(len(target_values), int(args.eval_target_size_min), int(args.eval_target_size_max)): |
| continue |
| row_has_eval_cell = True |
| prompt = build_multi_output_cell_prompt( |
| ex.grid, |
| target_cell=ex.target_cell, |
| stage_i=args.stage_i, |
| tokenizer=tokenizer, |
| turn_idx=ex.turn_idx, |
| total_turns=ex.total_turns, |
| prev_output_flag=None, |
| total_empties_hint=args.total_empties_hint, |
| ) |
| enc = tokenizer(prompt, return_tensors="pt", add_special_tokens=False) |
| enc = {k: v.to(device) for k, v in enc.items()} |
| out = model.generate( |
| **enc, |
| max_new_tokens=max(1, int(args.max_completion_length)), |
| do_sample=False, |
| eos_token_id=tokenizer.eos_token_id, |
| pad_token_id=tokenizer.pad_token_id, |
| ) |
| pred_text = tokenizer.decode(out[0][int(enc["input_ids"].shape[1]) :], skip_special_tokens=True).strip() |
| info = score_prediction_text( |
| text=pred_text, |
| grid=ex.grid, |
| solved=solved, |
| target_cell=ex.target_cell, |
| stage_i=args.stage_i, |
| reward_good_value=1.0, |
| penalty_bad_value=1.75, |
| penalty_malformed=4.0, |
| penalty_empty=0.5, |
| penalty_singleton=1.5, |
| ) |
| total_cells += 1 |
| parse_ok += float(info["parse_ok"]) |
| canonical_ok += float(info["strict_canonical"]) |
| exact_set_match += float(info["exact_set_match"]) |
| includes_gt += float(info["includes_ground_truth"]) |
| precision_sum += float(info["value_precision"]) |
| recall_sum += float(info["value_recall"]) |
| predicted_size_sum += float(info["num_predicted_values"]) |
| good_count_sum += float(info["num_i_consistent_values"]) |
| bad_count_sum += float(info["num_non_i_consistent_values"]) |
| if float(info["exact_set_match"]) < 0.5: |
| row_all_exact = False |
| if printed < int(args.debug_print_limit): |
| row_debug_lines.append( |
| f"[baseline sft eval debug] true_values={info['target_values']} " |
| f"predicted_values={info['predicted_values']} output={pred_text!r}" |
| ) |
| if row_has_eval_cell: |
| if printed < int(args.debug_print_limit) and row_debug_lines: |
| print("[baseline sft eval debug] puzzle_outputs_begin", flush=True) |
| for line in row_debug_lines: |
| print(line, flush=True) |
| print("[baseline sft eval debug] puzzle_outputs_end", flush=True) |
| printed += 1 |
| solve_ok += int(row_all_exact) |
| solve_rows += 1 |
| out = { |
| "parse_rate": float(parse_ok / max(1, total_cells)), |
| "strict_canonical_rate": float(canonical_ok / max(1, total_cells)), |
| "exact_set_match_rate": float(exact_set_match / max(1, total_cells)), |
| "includes_ground_truth_rate": float(includes_gt / max(1, total_cells)), |
| "value_precision": float(precision_sum / max(1, total_cells)), |
| "value_recall": float(recall_sum / max(1, total_cells)), |
| "avg_predicted_set_size": float(predicted_size_sum / max(1, total_cells)), |
| "avg_num_i_consistent_values": float(good_count_sum / max(1, total_cells)), |
| "avg_num_non_i_consistent_values": float(bad_count_sum / max(1, total_cells)), |
| "solve_rate": float(solve_ok / max(1, solve_rows)), |
| } |
| print( |
| f"[baseline sft eval] parse={out['parse_rate']:.3f} canonical={out['strict_canonical_rate']:.3f} " |
| f"exact={out['exact_set_match_rate']:.3f} precision={out['value_precision']:.3f} " |
| f"recall={out['value_recall']:.3f} solve={out['solve_rate']:.3f}", |
| flush=True, |
| ) |
| model.train() |
| return out |
|
|
|
|
| def parse_args() -> Args: |
| p = argparse.ArgumentParser() |
| p.add_argument("--model_name", type=str, default="Qwen/Qwen2.5-7B-Instruct") |
| p.add_argument("--train_jsonl", type=str, required=True) |
| p.add_argument("--eval_jsonl", type=str, default="") |
| p.add_argument("--output_dir", type=str, required=True) |
| p.add_argument("--cache_dir", type=str, default="/home/ubuntu/curriculum-CoT/.hf_cache") |
| p.add_argument("--init_adapter_dir", type=str, default="") |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--gpu_id", type=int, default=0) |
| p.add_argument("--stage_i", type=int, default=1) |
| p.add_argument("--total_empties_hint", type=int, default=10) |
| p.add_argument("--per_device_train_batch_size", type=int, default=1) |
| p.add_argument("--gradient_accumulation_steps", type=int, default=8) |
| p.add_argument("--num_epochs", type=float, default=1.0) |
| p.add_argument("--learning_rate", type=float, default=2e-4) |
| p.add_argument("--weight_decay", type=float, default=0.0) |
| p.add_argument( |
| "--max_grad_norm", |
| type=float, |
| default=1.0, |
| help="Clip global grad norm before each optimizer step (0 disables).", |
| ) |
| p.add_argument("--enable_gradient_checkpointing", action="store_true") |
| p.add_argument("--logging_steps", type=int, default=10) |
| p.add_argument("--save_steps", type=int, default=100) |
| p.add_argument("--eval_steps", type=int, default=100) |
| p.add_argument("--eval_rows", type=int, default=20) |
| p.add_argument("--max_completion_length", type=int, default=24) |
| p.add_argument("--lora_r", type=int, default=16) |
| p.add_argument("--lora_alpha", type=int, default=32) |
| p.add_argument("--lora_dropout", type=float, default=0.05) |
| p.add_argument("--use_wandb", action="store_true") |
| p.add_argument("--wandb_entity", type=str, default="") |
| p.add_argument("--wandb_project", type=str, default="sudoku-multi-output-sft") |
| p.add_argument("--wandb_run_name", type=str, default="") |
| p.add_argument("--wandb_mode", type=str, default="online") |
| p.add_argument("--debug_print_limit", type=int, default=3) |
| p.add_argument("--limit_train_rows", type=int, default=0) |
| p.add_argument("--eval_exact_set_match_stop", type=float, default=0.0) |
| p.add_argument("--eval_value_precision_stop", type=float, default=0.0) |
| p.add_argument("--eval_value_recall_stop", type=float, default=0.0) |
| p.add_argument("--eval_solve_rate_stop", type=float, default=0.0) |
| p.add_argument("--min_steps_before_stop", type=int, default=0) |
| p.add_argument("--max_wall_clock_seconds", type=int, default=0) |
| p.add_argument("--max_steps", type=int, default=0) |
| p.add_argument("--multi_value_oversample_factor", type=int, default=1) |
| p.add_argument("--train_target_size_min", type=int, default=0) |
| p.add_argument("--train_target_size_max", type=int, default=0) |
| p.add_argument("--eval_target_size_min", type=int, default=0) |
| p.add_argument("--eval_target_size_max", type=int, default=0) |
| return Args(**vars(p.parse_args())) |
|
|
|
|
| def save_checkpoint(model: torch.nn.Module, tokenizer: Any, output_dir: str, step: int) -> None: |
| save_checkpoint_and_update_final(model, tokenizer, output_dir, f"checkpoint-step-{step:05d}") |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| preset_visible_devices = str(os.environ.get("CUDA_VISIBLE_DEVICES", "")).strip() |
| rank = int(os.environ.get("RANK", "0")) |
| local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| is_distributed = world_size > 1 |
| is_main_process = rank == 0 |
|
|
| if preset_visible_devices: |
| if is_main_process: |
| print(f"Respecting pre-set CUDA_VISIBLE_DEVICES={preset_visible_devices}", flush=True) |
| elif int(args.gpu_id) >= 0: |
| os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" |
| os.environ["CUDA_VISIBLE_DEVICES"] = str(int(args.gpu_id)) |
|
|
| if is_distributed: |
| if torch.cuda.is_available(): |
| torch.cuda.set_device(local_rank) |
| dist.init_process_group(backend="nccl", timeout=timedelta(hours=2)) |
|
|
| set_seed(args.seed + rank) |
| os.makedirs(args.output_dir, exist_ok=True) |
| ensure_final_checkpoint_dir(args.output_dir) |
| cache_dir = configure_hf_cache(args.cache_dir) |
| configure_wandb_dirs(args.output_dir) |
|
|
| wb_run = None |
| if is_main_process and args.use_wandb and wandb is not None: |
| init_kwargs = { |
| "project": args.wandb_project, |
| "name": args.wandb_run_name or None, |
| "mode": args.wandb_mode, |
| } |
| if str(args.wandb_entity).strip(): |
| init_kwargs["entity"] = args.wandb_entity |
| wb_run = wandb.init(**init_kwargs) |
| print(f"W&B run id: {wb_run.id}", flush=True) |
| print(f"W&B run URL: {wb_run.url}", flush=True) |
| wandb.log({"prep/rows_done": 0.0, "prep/examples_built": 0.0, "prep/cache_hit": 0.0}) |
|
|
| rows = load_jsonl_rows(args.train_jsonl, limit_rows=args.limit_train_rows) |
| eval_source = args.eval_jsonl if str(args.eval_jsonl).strip() else args.train_jsonl |
| eval_rows = load_jsonl_rows(eval_source, limit_rows=max(1, int(args.eval_rows))) |
| tokenizer = AutoTokenizer.from_pretrained(args.model_name, cache_dir=cache_dir, use_fast=True) |
| if tokenizer.pad_token_id is None: |
| tokenizer.pad_token = tokenizer.eos_token or "<|endoftext|>" |
| tokenizer._multi_value_oversample_factor = max(1, int(args.multi_value_oversample_factor)) |
| tokenizer._train_target_size_min = max(0, int(args.train_target_size_min)) |
| tokenizer._train_target_size_max = max(0, int(args.train_target_size_max)) |
| if torch.cuda.is_available(): |
| device = torch.device(f"cuda:{local_rank}" if is_distributed else f"cuda:{max(0, int(args.gpu_id))}") |
| else: |
| device = torch.device("cpu") |
|
|
| model = AutoModelForCausalLM.from_pretrained( |
| args.model_name, cache_dir=cache_dir, torch_dtype=pick_dtype(), low_cpu_mem_usage=True |
| ) |
| if str(args.init_adapter_dir).strip(): |
| model = PeftModel.from_pretrained(model, args.init_adapter_dir, is_trainable=True) |
| else: |
| lora = LoraConfig( |
| r=args.lora_r, |
| lora_alpha=args.lora_alpha, |
| lora_dropout=args.lora_dropout, |
| 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) |
| if args.enable_gradient_checkpointing and hasattr(model, "gradient_checkpointing_enable"): |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) |
| if args.enable_gradient_checkpointing and hasattr(model, "enable_input_require_grads"): |
| model.enable_input_require_grads() |
| if hasattr(model, "config"): |
| model.config.use_cache = False |
| model.to(device) |
| model.train() |
| if torch.cuda.is_available(): |
| torch.backends.cudnn.benchmark = True |
|
|
| def on_prep_progress(rows_done: int, total_rows: int, examples_built: int) -> None: |
| if is_main_process: |
| print( |
| f"[dataset build][sft stage {args.stage_i}] rows={rows_done}/{total_rows} examples={examples_built}", |
| flush=True, |
| ) |
| if wb_run is not None: |
| wandb.log({"prep/rows_done": float(rows_done), "prep/examples_built": float(examples_built)}) |
|
|
| train_examples = load_or_build_sft_examples( |
| args, |
| rows=rows, |
| tokenizer=tokenizer, |
| rank=rank, |
| world_size=world_size, |
| progress_callback=on_prep_progress, |
| ) |
| if is_main_process and wb_run is not None: |
| wandb.log( |
| { |
| "prep/cache_hit": float(os.path.exists(_prepared_sft_cache_path(args))), |
| "prep/examples_final": float(len(train_examples)), |
| } |
| ) |
|
|
| optimizer = AdamW((p for p in model.parameters() if p.requires_grad), lr=args.learning_rate, weight_decay=args.weight_decay) |
| denom = max(1, int(args.gradient_accumulation_steps)) * max(1, int(args.per_device_train_batch_size)) |
| total_steps = max(1, math.ceil(len(train_examples) * args.num_epochs / denom)) |
| if int(args.max_steps) > 0: |
| total_steps = min(total_steps, int(args.max_steps)) |
| step = 0 |
| start_time = time.time() |
|
|
| def average_scalar(value: float) -> float: |
| if not is_distributed or not dist.is_initialized(): |
| return float(value) |
| tensor = torch.tensor(float(value), device=device, dtype=torch.float32) |
| dist.all_reduce(tensor, op=dist.ReduceOp.SUM) |
| return float((tensor / float(world_size)).item()) |
|
|
| def all_reduce_gradients() -> None: |
| if not is_distributed or not dist.is_initialized(): |
| return |
| for param in model.parameters(): |
| if param.grad is None: |
| continue |
| dist.all_reduce(param.grad, op=dist.ReduceOp.SUM) |
| param.grad.div_(float(world_size)) |
|
|
| def sync_stop(local_stop: bool) -> bool: |
| if not is_distributed or not dist.is_initialized(): |
| return bool(local_stop) |
| tensor = torch.tensor(1 if local_stop else 0, device=device, dtype=torch.int64) |
| dist.all_reduce(tensor, op=dist.ReduceOp.MAX) |
| return bool(int(tensor.item()) > 0) |
|
|
| for epoch_idx in range(max(1, int(math.ceil(args.num_epochs)))): |
| if is_distributed: |
| sampler = DistributedSampler( |
| train_examples, |
| num_replicas=world_size, |
| rank=rank, |
| shuffle=True, |
| seed=args.seed, |
| drop_last=False, |
| ) |
| sampler.set_epoch(epoch_idx) |
| order = list(iter(sampler)) |
| else: |
| generator = torch.Generator() |
| generator.manual_seed(args.seed + epoch_idx) |
| order = torch.randperm(len(train_examples), generator=generator).tolist() |
| optimizer.zero_grad(set_to_none=True) |
| accum_count = 0 |
| accum_ce_sum = 0.0 |
| microbatch_size = max(1, int(args.per_device_train_batch_size)) |
| for batch_start in range(0, len(order), microbatch_size): |
| batch_indices = order[batch_start : batch_start + microbatch_size] |
| batch_examples = [train_examples[ex_idx] for ex_idx in batch_indices] |
| ce_full = batched_completion_ce_loss( |
| model, |
| tokenizer, |
| [str(ex["prompt_text"]) for ex in batch_examples], |
| [str(ex["completion_text"]) for ex in batch_examples], |
| device, |
| ) |
| loss = ce_full / max(1, int(args.gradient_accumulation_steps)) |
| loss.backward() |
| accum_ce_sum += float(ce_full.detach().item()) |
| accum_count += 1 |
| if accum_count >= int(args.gradient_accumulation_steps): |
| all_reduce_gradients() |
| if float(args.max_grad_norm) > 0.0: |
| torch.nn.utils.clip_grad_norm_(model.parameters(), float(args.max_grad_norm)) |
| optimizer.step() |
| optimizer.zero_grad(set_to_none=True) |
| accum_count = 0 |
| step += 1 |
| mean_ce = accum_ce_sum / max(1, int(args.gradient_accumulation_steps)) |
| accum_ce_sum = 0.0 |
| if step % int(args.logging_steps) == 0: |
| loss_value = average_scalar(mean_ce) |
| if is_main_process: |
| print(f"[baseline sft train step {step:05d}] loss={loss_value:.4f}", flush=True) |
| if wb_run is not None: |
| wandb.log({"train/loss": loss_value, "step": step}) |
| if step % int(args.eval_steps) == 0: |
| if is_distributed and dist.is_initialized(): |
| dist.barrier() |
| should_stop_eval = False |
| if is_main_process: |
| ev = run_eval(args, eval_rows, model, tokenizer, device) |
| if wb_run is not None: |
| wandb.log({f"eval/{k}": float(v) for k, v in ev.items()} | {"step": step}) |
| if ( |
| args.eval_exact_set_match_stop > 0.0 |
| and float(ev["exact_set_match_rate"]) >= args.eval_exact_set_match_stop |
| ): |
| save_checkpoint(model, tokenizer, args.output_dir, step) |
| should_stop_eval = True |
| if ( |
| not should_stop_eval |
| and step >= int(args.min_steps_before_stop) |
| and args.eval_value_precision_stop > 0.0 |
| and args.eval_value_recall_stop > 0.0 |
| and float(ev["value_precision"]) >= args.eval_value_precision_stop |
| and float(ev["value_recall"]) >= args.eval_value_recall_stop |
| ): |
| save_checkpoint(model, tokenizer, args.output_dir, step) |
| should_stop_eval = True |
| if ( |
| not should_stop_eval |
| and args.eval_solve_rate_stop > 0.0 |
| and step >= int(args.min_steps_before_stop) |
| and float(ev["solve_rate"]) >= args.eval_solve_rate_stop |
| ): |
| save_checkpoint(model, tokenizer, args.output_dir, step) |
| should_stop_eval = True |
| should_stop_eval = sync_stop(should_stop_eval) |
| if is_distributed and dist.is_initialized(): |
| dist.barrier() |
| if should_stop_eval: |
| if is_main_process and wb_run is not None: |
| wb_run.finish() |
| if is_distributed and dist.is_initialized(): |
| dist.destroy_process_group() |
| return |
| if step % int(args.save_steps) == 0: |
| if is_distributed and dist.is_initialized(): |
| dist.barrier() |
| if is_main_process: |
| save_checkpoint(model, tokenizer, args.output_dir, step) |
| if is_distributed and dist.is_initialized(): |
| dist.barrier() |
| reached_limit = step >= total_steps |
| exceeded_wall = bool(args.max_wall_clock_seconds) and ( |
| time.time() - start_time >= float(args.max_wall_clock_seconds) |
| ) |
| should_stop = sync_stop(reached_limit or exceeded_wall) |
| if should_stop: |
| break |
| if sync_stop(step >= total_steps): |
| break |
|
|
| if is_distributed and dist.is_initialized(): |
| dist.barrier() |
| if is_main_process: |
| save_checkpoint(model, tokenizer, args.output_dir, step) |
| if wb_run is not None: |
| wb_run.finish() |
| if is_distributed and dist.is_initialized(): |
| dist.barrier() |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
