| """Explicit single-GPU classifier training with optional development selection.""" |
|
|
| |
|
|
| import argparse |
| import hashlib |
| import json |
| import math |
| import random |
| import time |
| from collections import Counter, defaultdict |
| from pathlib import Path |
|
|
| from .data import assert_disjoint, file_receipts, label_counts, read_records |
| from .evaluate import evaluate_records |
| from .optimization import ( |
| assert_frozen_parameters_preserved, |
| configure_trainable_tail, |
| frozen_parameter_receipt, |
| initial_artifact_receipt, |
| load_trainable_model, |
| optimizer_groups, |
| save_training_checkpoint, |
| validate_method, |
| verify_full_checkpoint, |
| ) |
| from .retention import RetentionCache, masked_teacher_kl, tokenizer_file_receipt |
| from .selection import ( |
| BINARY_FP_SELECTION, |
| binary_checkpoint_key, |
| binary_selection_support, |
| fit_binary_operating_point, |
| validate_selection_options, |
| ) |
| from .task_head import task_head_scope |
| from .training_order import load_training_order |
|
|
|
|
| def microbatch_loss(logits, labels, accumulation_steps): |
| from torch.nn import functional |
|
|
| if accumulation_steps <= 0: |
| raise ValueError("Gradient accumulation must be positive") |
| |
| |
| return ( |
| functional.cross_entropy(logits.float(), labels, reduction="mean") |
| / accumulation_steps |
| ) |
|
|
|
|
| def global_batch_loss(logits, labels, examples_per_step): |
| """Normalize variable microbatches by the same sampled optimizer-step size.""" |
| from torch.nn import functional |
|
|
| if examples_per_step <= 0: |
| raise ValueError("Optimizer-step example count must be positive") |
| return ( |
| functional.cross_entropy(logits.float(), labels, reduction="sum") |
| / examples_per_step |
| ) |
|
|
|
|
| def token_budget_microbatches(indices, lengths, token_budget): |
| """Stable length ordering without dropping or duplicating sampled examples.""" |
| if token_budget <= 0: |
| raise ValueError("Microbatch token budget must be positive") |
| batches, current = [], [] |
| for index in sorted(indices, key=lengths.__getitem__): |
| length = lengths[index] |
| if length > token_budget: |
| raise ValueError("A sampled example exceeds the microbatch token budget") |
| if current and length * (len(current) + 1) > token_budget: |
| batches.append(current) |
| current = [] |
| current.append(index) |
| if current: |
| batches.append(current) |
| return batches |
|
|
|
|
| def selection_score(metrics, selection): |
| if selection == "macro-f1": |
| return metrics["macro_f1"] |
| if selection in { |
| "length-macro-f1", |
| "source-macro-f1", |
| "source-present-macro-f1", |
| }: |
| field = "length_bucket" if selection == "length-macro-f1" else "source" |
| metric = ( |
| "macro_f1_present_labels" |
| if selection == "source-present-macro-f1" |
| else "macro_f1" |
| ) |
| groups = list(metrics["breakdowns"][field].values()) |
| if not groups: |
| raise ValueError("Grouped selection requires measured groups") |
| return sum(group[metric] for group in groups) / len(groups) |
| raise ValueError("Unknown development selection metric") |
|
|
|
|
| def length_sampling_pools(rows): |
| buckets = defaultdict(lambda: defaultdict(list)) |
| for index, row in enumerate(rows): |
| buckets[str(row.get("length_bucket", "unspecified"))][row["label"]].append( |
| index |
| ) |
| return [list(labels.values()) for labels in buckets.values()] |
|
|
|
|
| def sample_length_balanced(pools, rng, balance_labels): |
| bucket = rng.choice(pools) |
| pool = ( |
| rng.choice(bucket) |
| if balance_labels |
| else [index for label in bucket for index in label] |
| ) |
| return rng.choice(pool) |
|
|
|
|
| def source_sampling_pools(rows, balance_lengths): |
| sources = defaultdict(lambda: defaultdict(lambda: defaultdict(list))) |
| for index, row in enumerate(rows): |
| if not row.get("source"): |
| raise ValueError("Source-balanced sampling requires explicit sources") |
| length = ( |
| str(row.get("length_bucket", "unspecified")) if balance_lengths else "all" |
| ) |
| sources[row["source"]][length][row["label"]].append(index) |
| return [ |
| [list(labels.values()) for labels in lengths.values()] |
| for lengths in sources.values() |
| ] |
|
|
|
|
| def source_sampling_weights(rows, specification): |
| """Resolve explicit probabilities in the same order as source pools.""" |
| if specification is None: |
| return None |
| weights = json.loads(specification) |
| sources = list(dict.fromkeys(row.get("source") for row in rows)) |
| if ( |
| not sources |
| or not all(sources) |
| or not isinstance(weights, dict) |
| or set(weights) != set(sources) |
| ): |
| raise ValueError("Source weights must name every eligible source exactly") |
| values = [weights[source] for source in sources] |
| if any( |
| isinstance(value, bool) |
| or not isinstance(value, (int, float)) |
| or not math.isfinite(value) |
| or value <= 0 |
| for value in values |
| ): |
| raise ValueError("Source weights must be finite positive numbers") |
| total = sum(values) |
| if not math.isfinite(total): |
| raise ValueError("Total source weight must be finite") |
| return [value / total for value in values] |
|
|
|
|
| def sample_source_balanced(pools, rng, balance_labels, source_weights=None): |
| source = ( |
| rng.choice(pools) |
| if source_weights is None |
| else rng.choices(pools, weights=source_weights, k=1)[0] |
| ) |
| return sample_length_balanced(source, rng, balance_labels) |
|
|
|
|
| def sampling_exposure(rows, drawn): |
| """Report sampling counts separately from class/length selection metrics.""" |
| sources = defaultdict( |
| lambda: {"rows": 0, "draws": 0, "unique_rows": 0, "label_draws": Counter()} |
| ) |
| for index, row in enumerate(rows): |
| item = sources[row.get("source", "unspecified")] |
| count = drawn[index] |
| item["rows"] += 1 |
| item["draws"] += count |
| item["unique_rows"] += int(count > 0) |
| item["label_draws"][row["label"]] += count |
| return dict(sources) |
|
|
|
|
| def main(): |
| import torch |
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--base", required=True) |
| parser.add_argument("--base-id", required=True) |
| parser.add_argument("--base-revision", required=True) |
| parser.add_argument("--adapter", type=Path) |
| parser.add_argument("--method", choices=["lora", "full"], default="lora") |
| parser.add_argument("--fresh-head", action="store_true") |
| parser.add_argument("--contract", required=True) |
| parser.add_argument("--train", nargs="+", required=True) |
| parser.add_argument("--dev", nargs="+") |
| parser.add_argument("--output", type=Path, required=True) |
| parser.add_argument("--steps", type=int, default=400) |
| parser.add_argument("--batch-size", type=int, default=8) |
| parser.add_argument("--accumulate", type=int, default=2) |
| parser.add_argument("--max-length", type=int, default=2048) |
| parser.add_argument("--learning-rate", type=float, default=2e-5) |
| parser.add_argument("--head-learning-rate", type=float) |
| parser.add_argument("--trainable-last-layers", type=int) |
| parser.add_argument("--retention-targets", type=Path) |
| parser.add_argument("--retention-weight", type=float) |
| parser.add_argument("--retention-temperature", type=float) |
| parser.add_argument("--eval-every", type=int, default=100) |
| parser.add_argument( |
| "--evaluation-dtype", choices=["bfloat16", "float32"], default="bfloat16" |
| ) |
| parser.add_argument( |
| "--selection", |
| choices=[ |
| "macro-f1", |
| "length-macro-f1", |
| "source-macro-f1", |
| "source-present-macro-f1", |
| BINARY_FP_SELECTION, |
| ], |
| default="macro-f1", |
| ) |
| parser.add_argument("--positive-label") |
| parser.add_argument("--selection-false-positive-budget", type=float) |
| parser.add_argument("--balanced-sampling", action="store_true") |
| parser.add_argument("--length-balanced-sampling", action="store_true") |
| parser.add_argument("--source-balanced-sampling", action="store_true") |
| parser.add_argument( |
| "--source-weights", help="JSON object of explicit source weights" |
| ) |
| parser.add_argument("--microbatch-token-budget", type=int) |
| parser.add_argument("--training-order", type=Path) |
| mode = parser.add_mutually_exclusive_group() |
| mode.add_argument("--probe-only", action="store_true") |
| mode.add_argument( |
| "--train-only", |
| action="store_true", |
| help="Skip development data and selection; save only the final checkpoint", |
| ) |
| parser.add_argument("--seed", type=int, default=20260913) |
| args = parser.parse_args() |
| if args.train_only: |
| if args.dev: |
| parser.error("--train-only must not receive development data") |
| if ( |
| args.selection != "macro-f1" |
| or args.positive_label is not None |
| or args.selection_false_positive_budget is not None |
| ): |
| parser.error("--train-only does not select a development operating point") |
| elif not args.dev: |
| parser.error("--dev is required unless --train-only is set") |
| validate_method(args.method, args.adapter, args.fresh_head) |
| if args.trainable_last_layers is not None and args.method != "full": |
| raise ValueError("A trainable encoder tail requires full checkpoint training") |
| if args.retention_targets is None: |
| if args.retention_weight is not None or args.retention_temperature is not None: |
| raise ValueError("Retention settings require explicit TRAIN targets") |
| else: |
| if args.method != "full" or args.fresh_head: |
| raise ValueError("Retention requires the unchanged complete initial head") |
| args.retention_weight = ( |
| 0.5 if args.retention_weight is None else args.retention_weight |
| ) |
| args.retention_temperature = ( |
| 2.0 if args.retention_temperature is None else args.retention_temperature |
| ) |
| if any( |
| not math.isfinite(value) or value <= 0 |
| for value in (args.retention_weight, args.retention_temperature) |
| ): |
| raise ValueError( |
| "Retention weight and temperature must be finite and positive" |
| ) |
| if args.training_order and ( |
| args.balanced_sampling |
| or args.length_balanced_sampling |
| or args.source_balanced_sampling |
| or args.source_weights is not None |
| ): |
| raise ValueError("An explicit training order cannot also configure sampling") |
| if ( |
| min( |
| args.steps, |
| args.batch_size, |
| args.accumulate, |
| args.max_length, |
| args.eval_every, |
| ) |
| <= 0 |
| or args.learning_rate <= 0 |
| or ( |
| args.microbatch_token_budget is not None |
| and args.microbatch_token_budget <= 0 |
| ) |
| ): |
| raise ValueError("Training budgets and learning rate must be positive") |
| if args.output.exists() and any(args.output.iterdir()): |
| raise ValueError("Refusing to overwrite a training run") |
| torch.set_num_threads(8) |
| torch.manual_seed(args.seed) |
| rng = random.Random(args.seed) |
| initial_artifacts = initial_artifact_receipt(args.base, args.adapter) |
| model, tokenizer, label_to_id, id_to_label = load_trainable_model( |
| args.base, args.contract, args.adapter, args.method, args.fresh_head |
| ) |
| if model.config.problem_type != "single_label_classification": |
| raise ValueError("This training loop requires single-label targets") |
| if args.retention_targets is not None: |
| verify_full_checkpoint(model, args.base) |
| trainable_scope = ( |
| configure_trainable_tail(model, args.trainable_last_layers) |
| if args.trainable_last_layers is not None |
| else None |
| ) |
| frozen_parameters = ( |
| frozen_parameter_receipt(model) if trainable_scope is not None else None |
| ) |
| if not args.train_only: |
| validate_selection_options( |
| args.selection, |
| label_to_id, |
| args.positive_label, |
| args.selection_false_positive_budget, |
| ) |
| if args.max_length > model.config.max_position_embeddings: |
| raise ValueError("Training budget exceeds checkpoint position capacity") |
| rows, dev = ( |
| read_records(args.train, label_to_id), |
| [] if args.train_only else read_records(args.dev, label_to_id), |
| ) |
| assert_disjoint(rows, dev) |
| if args.selection == BINARY_FP_SELECTION: |
| binary_selection_support(dev, label_to_id, args.positive_label) |
| train, rejected, pools = [], [], defaultdict(list) |
| for row in rows: |
| encoded = tokenizer(row["text"], truncation=False, padding=False) |
| if len(encoded["input_ids"]) > args.max_length: |
| rejected.append( |
| { |
| "id": row["id"], |
| "reason": "explicit context budget", |
| "tokens": len(encoded["input_ids"]), |
| } |
| ) |
| continue |
| index = len(train) |
| train.append((row, encoded)) |
| pools[row["label"]].append(index) |
| if not train: |
| raise ValueError("No training rows fit the explicit budget") |
| retention = None |
| if args.retention_targets is not None: |
| if any( |
| not isinstance(row.get("retention_replay", False), bool) for row, _ in train |
| ): |
| raise ValueError("retention_replay must be a boolean TRAIN annotation") |
| retention = RetentionCache.load( |
| args.retention_targets, |
| initialization=initial_artifacts, |
| contract_sha256=hashlib.sha256( |
| Path(args.contract).read_bytes() |
| ).hexdigest(), |
| tokenizer_files=tokenizer_file_receipt(args.base), |
| train_rows=[row for row, _ in train], |
| tokenized_rows={row["id"]: encoded for row, encoded in train}, |
| label_to_id=label_to_id, |
| replay_ids=[ |
| row["id"] for row, _ in train if row.get("retention_replay", False) |
| ], |
| ) |
| train_lengths = [len(encoded["input_ids"]) for _, encoded in train] |
| if ( |
| args.microbatch_token_budget is not None |
| and max(train_lengths) > args.microbatch_token_budget |
| ): |
| raise ValueError( |
| "Microbatch token budget cannot fit the longest eligible example" |
| ) |
| training_order = ( |
| load_training_order( |
| args.training_order, |
| [row for row, _ in train], |
| args.train, |
| steps=args.steps, |
| global_batch=args.batch_size * args.accumulate, |
| ) |
| if args.training_order |
| else None |
| ) |
| pool_values = list(pools.values()) |
| length_pools = length_sampling_pools([row for row, _encoded in train]) |
| balance_sources = args.source_balanced_sampling or args.source_weights is not None |
| source_weights = source_sampling_weights( |
| [row for row, _encoded in train], args.source_weights |
| ) |
| source_pools = ( |
| source_sampling_pools( |
| [row for row, _encoded in train], args.length_balanced_sampling |
| ) |
| if balance_sources |
| else None |
| ) |
| model.gradient_checkpointing_enable( |
| gradient_checkpointing_kwargs={"use_reentrant": False} |
| ) |
| model.to("cuda") |
| model.train() |
| parameters = [ |
| parameter for parameter in model.parameters() if parameter.requires_grad |
| ] |
| groups = optimizer_groups(model, args.learning_rate, args.head_learning_rate) |
| optimizer = torch.optim.AdamW(groups, weight_decay=0.01) |
| args.output.mkdir(parents=True, exist_ok=True) |
| metadata = { |
| "base_model": args.base_id, |
| "base_revision": args.base_revision, |
| "method": args.method, |
| "fresh_head": args.fresh_head, |
| "trainable_scope": trainable_scope, |
| "frozen_parameters": frozen_parameters, |
| "retention": ( |
| { |
| "targets": retention.receipt, |
| "weight": args.retention_weight, |
| "temperature": args.retention_temperature, |
| "normalization": "all examples in the optimizer step", |
| } |
| if retention is not None |
| else None |
| ), |
| "initial_artifacts": initial_artifacts, |
| "initial_adapter_sha256": initial_artifacts.get("adapter_model.safetensors"), |
| "initial_adapter_config_sha256": initial_artifacts.get("adapter_config.json"), |
| "contract_sha256": hashlib.sha256(Path(args.contract).read_bytes()).hexdigest(), |
| "classifier_pooling": model.config.classifier_pooling, |
| "task_head": task_head_scope(model), |
| "problem_type": model.config.problem_type, |
| "train_files": file_receipts(args.train), |
| "dev_files": file_receipts(args.dev or []), |
| "train_only": args.train_only, |
| "checkpoint_policy": ( |
| "final-only" if args.train_only else "development-and-final" |
| ), |
| "label2id": label_to_id, |
| "train_rows": len(train), |
| "dev_rows": len(dev), |
| "train_label_counts": label_counts([row for row, _ in train]), |
| "rejected": rejected, |
| "seed": args.seed, |
| "steps": args.steps, |
| "batch_size": args.batch_size, |
| "accumulate": args.accumulate, |
| "examples_per_optimizer_step": args.batch_size * args.accumulate, |
| "microbatch_token_budget": args.microbatch_token_budget, |
| "dropout_trajectory_equivalence": args.microbatch_token_budget is None, |
| "max_length": args.max_length, |
| "learning_rate": args.learning_rate, |
| "optimizer_groups": [ |
| { |
| "name": group["name"], |
| "learning_rate": group["initial_lr"], |
| "parameters": sum(parameter.numel() for parameter in group["params"]), |
| } |
| for group in groups |
| ], |
| "balanced_sampling": args.balanced_sampling, |
| "training_order": training_order.receipt if training_order else None, |
| "length_balanced_sampling": args.length_balanced_sampling, |
| "source_balanced_sampling": balance_sources, |
| "source_sampling_probabilities": ( |
| dict( |
| zip( |
| dict.fromkeys(row["source"] for row, _ in train), |
| source_weights, |
| strict=True, |
| ) |
| ) |
| if source_weights is not None |
| else None |
| ), |
| "implementation_sha256": hashlib.sha256( |
| Path(__file__).read_bytes() |
| ).hexdigest(), |
| "selection": None if args.train_only else args.selection, |
| "evaluation_dtype": None if args.train_only else args.evaluation_dtype, |
| "test_used": False, |
| "precision": "FP32 parameters and explicit FP32 mean CE / BF16 autocast", |
| "attention": "sdpa", |
| "checkpointing": "non-reentrant", |
| "parameters": sum(parameter.numel() for parameter in model.parameters()), |
| "trainable_parameters": sum(parameter.numel() for parameter in parameters), |
| } |
| if args.selection == BINARY_FP_SELECTION: |
| metadata["binary_selection"] = { |
| "positive_label": args.positive_label, |
| "false_positive_budget": args.selection_false_positive_budget, |
| "support": binary_selection_support(dev, label_to_id, args.positive_label), |
| "implementation_sha256": hashlib.sha256( |
| Path(__file__).with_name("selection.py").read_bytes() |
| ).hexdigest(), |
| } |
| (args.output / "run.json").write_text(json.dumps(metadata, indent=2) + "\n") |
| print( |
| json.dumps( |
| { |
| "event": "ready", |
| "train_rows": len(train), |
| "rejected_rows": len(rejected), |
| "trainable_parameters": metadata["trainable_parameters"], |
| } |
| ), |
| flush=True, |
| ) |
| best = -1.0 |
| best_key = None |
| drawn = Counter() |
|
|
| def evaluate_checkpoint(step): |
| nonlocal best, best_key |
| metrics, predictions = evaluate_records( |
| model, |
| tokenizer, |
| dev, |
| label_to_id, |
| id_to_label, |
| args.max_length, |
| dtype=args.evaluation_dtype, |
| ) |
| operating_point = None |
| if args.selection == BINARY_FP_SELECTION: |
| operating_point = fit_binary_operating_point( |
| dev, |
| predictions, |
| label_to_id, |
| args.positive_label, |
| args.selection_false_positive_budget, |
| ) |
| metrics["binary_fp_budget"] = operating_point |
| key = binary_checkpoint_key(operating_point) |
| score = operating_point["recall"] if key is not None else -1.0 |
| (args.output / f"dev-predictions-step-{step}.jsonl").write_text( |
| "".join(json.dumps(row) + "\n" for row in predictions) |
| ) |
| else: |
| score = selection_score(metrics, args.selection) |
| key = (score,) |
| (args.output / f"sampling-step-{step}.json").write_text( |
| json.dumps(sampling_exposure([row for row, _ in train], drawn), indent=2) |
| + "\n" |
| ) |
| (args.output / f"dev-step-{step}.json").write_text( |
| json.dumps(metrics, indent=2) + "\n" |
| ) |
| print( |
| json.dumps( |
| { |
| "event": "dev", |
| "step": step, |
| "score": score, |
| "accuracy": metrics["accuracy"], |
| "macro_f1": metrics["macro_f1"], |
| } |
| ), |
| flush=True, |
| ) |
| if key is not None and (best_key is None or key > best_key): |
| best = score |
| best_key = key |
| if frozen_parameters is not None: |
| assert_frozen_parameters_preserved(frozen_parameters, model) |
| save_training_checkpoint( |
| model, |
| tokenizer, |
| args.output |
| / ("best-adapter" if args.method == "lora" else "best-model"), |
| args.method, |
| args.base_id, |
| args.base_revision, |
| ) |
| (args.output / "selection.json").write_text( |
| json.dumps( |
| { |
| "step": step, |
| "score": score, |
| "metric": args.selection, |
| "test_used": False, |
| **( |
| {"operating_point": operating_point} |
| if operating_point is not None |
| else {} |
| ), |
| }, |
| indent=2, |
| ) |
| + "\n" |
| ) |
|
|
| def draw_microbatch_indices(): |
| return [ |
| ( |
| sample_source_balanced( |
| source_pools, rng, args.balanced_sampling, source_weights |
| ) |
| if balance_sources |
| else ( |
| sample_length_balanced(length_pools, rng, args.balanced_sampling) |
| if args.length_balanced_sampling |
| else ( |
| rng.choice(rng.choice(pool_values)) |
| if args.balanced_sampling |
| else rng.randrange(len(train)) |
| ) |
| ) |
| ) |
| for _ in range(args.batch_size) |
| ] |
|
|
| if not args.probe_only and not args.train_only: |
| evaluate_checkpoint(0) |
| for step in range(1, args.steps + 1): |
| optimizer.zero_grad(set_to_none=True) |
| torch.cuda.reset_peak_memory_stats() |
| torch.cuda.synchronize() |
| start, step_loss, lengths = time.perf_counter(), 0.0, [] |
| step_ce_loss, step_retention_loss, retention_examples = 0.0, 0.0, 0 |
| if training_order: |
| planned = training_order.indices_for_step(step) |
| microbatches = ( |
| token_budget_microbatches( |
| planned, train_lengths, args.microbatch_token_budget |
| ) |
| if args.microbatch_token_budget is not None |
| else [ |
| planned[start : start + args.batch_size] |
| for start in range(0, len(planned), args.batch_size) |
| ] |
| ) |
| elif args.microbatch_token_budget is None: |
| microbatches = (draw_microbatch_indices() for _ in range(args.accumulate)) |
| else: |
| sampled = [ |
| index |
| for _ in range(args.accumulate) |
| for index in draw_microbatch_indices() |
| ] |
| microbatches = token_budget_microbatches( |
| sampled, train_lengths, args.microbatch_token_budget |
| ) |
| microbatch_count, max_padded_tokens = 0, 0 |
| for indices in microbatches: |
| microbatch_count += 1 |
| max_padded_tokens = max( |
| max_padded_tokens, max(train_lengths[i] for i in indices) * len(indices) |
| ) |
| drawn.update(indices) |
| selected = [train[index] for index in indices] |
| lengths.extend(len(encoded["input_ids"]) for _, encoded in selected) |
| batch = tokenizer.pad( |
| [encoded for _, encoded in selected], padding=True, return_tensors="pt" |
| ) |
| batch = {key: value.to("cuda") for key, value in batch.items()} |
| gold = torch.tensor( |
| [label_to_id[row["label"]] for row, _ in selected], device="cuda" |
| ) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| logits = model(**batch).logits |
| loss = ( |
| microbatch_loss(logits, gold, args.accumulate) |
| if args.microbatch_token_budget is None |
| else global_batch_loss( |
| logits, gold, args.batch_size * args.accumulate |
| ) |
| ) |
| step_ce_loss += float(loss.detach()) |
| if retention is not None: |
| teacher_logits, eligible = retention.targets( |
| [row for row, _ in selected], device=logits.device |
| ) |
| preservation = args.retention_weight * masked_teacher_kl( |
| logits, |
| teacher_logits, |
| eligible, |
| examples_per_step=args.batch_size * args.accumulate, |
| temperature=args.retention_temperature, |
| ) |
| loss = loss + preservation |
| step_retention_loss += float(preservation.detach()) |
| retention_examples += int(eligible.sum()) |
| if not bool(torch.isfinite(loss)): |
| raise ValueError(f"Non-finite loss at step {step}") |
| loss.backward() |
| step_loss += float(loss.detach()) |
| del logits, loss, batch |
| gradients = [ |
| parameter.grad for parameter in parameters if parameter.grad is not None |
| ] |
| if not gradients or not all( |
| bool(torch.isfinite(gradient).all()) for gradient in gradients |
| ): |
| raise ValueError(f"Missing or non-finite gradients at step {step}") |
| norm = torch.nn.utils.clip_grad_norm_(parameters, 1.0) |
| warmup = max(1, round(args.steps * 0.06)) |
| rate = ( |
| step / warmup |
| if step <= warmup |
| else 0.5 |
| * (1 + math.cos(math.pi * (step - warmup) / max(1, args.steps - warmup))) |
| ) |
| for group in optimizer.param_groups: |
| group["lr"] = group["initial_lr"] * rate |
| optimizer.step() |
| if training_order: |
| trace = training_order.record_step(step, microbatches) |
| with (args.output / "actual-training-order.jsonl").open("a") as stream: |
| stream.write(json.dumps(trace) + "\n") |
| torch.cuda.synchronize() |
| log = { |
| "step": step, |
| "loss": step_loss, |
| "cross_entropy_loss": step_ce_loss, |
| "retention_loss": step_retention_loss, |
| "retention_examples": retention_examples, |
| "grad_norm": float(norm), |
| "finite": True, |
| "gradient_tensors": len(gradients), |
| "max_tokens": max(lengths), |
| "mean_tokens": sum(lengths) / len(lengths), |
| "microbatches": microbatch_count, |
| "max_padded_tokens": max_padded_tokens, |
| "step_seconds": time.perf_counter() - start, |
| "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30, |
| "peak_reserved_gib": torch.cuda.max_memory_reserved() / 2**30, |
| } |
| with (args.output / "steps.jsonl").open("a") as stream: |
| stream.write(json.dumps(log) + "\n") |
| if step == 1 or step % 10 == 0 or args.probe_only: |
| print(json.dumps(log), flush=True) |
| if ( |
| not args.probe_only |
| and not args.train_only |
| and (step % args.eval_every == 0 or step == args.steps) |
| ): |
| evaluate_checkpoint(step) |
| if training_order: |
| order_receipt = training_order.finish() |
| order_receipt["actual_trace_sha256"] = hashlib.sha256( |
| (args.output / "actual-training-order.jsonl").read_bytes() |
| ).hexdigest() |
| (args.output / "training-order-completed.json").write_text( |
| json.dumps(order_receipt, indent=2) + "\n" |
| ) |
| (args.output / "sampling-final.json").write_text( |
| json.dumps(sampling_exposure([row for row, _ in train], drawn), indent=2) + "\n" |
| ) |
| if frozen_parameters is not None: |
| assert_frozen_parameters_preserved(frozen_parameters, model) |
| if not args.probe_only: |
| save_training_checkpoint( |
| model, |
| tokenizer, |
| args.output / ("last-adapter" if args.method == "lora" else "last-model"), |
| args.method, |
| args.base_id, |
| args.base_revision, |
| ) |
| if best_key is None and not args.train_only: |
| raise ValueError( |
| "No checkpoint met the development FP budget; last checkpoint and " |
| "infeasibility receipts were retained without selecting a best checkpoint" |
| ) |
| print( |
| json.dumps( |
| { |
| "event": "complete", |
| "best_dev_score": ( |
| best if not (args.probe_only or args.train_only) else None |
| ), |
| } |
| ), |
| flush=True, |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
