| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import os |
| import sys |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
|
|
|
|
| SCRIPT_DIR = Path(__file__).resolve().parent |
| V3P5_SCRIPT_DIR = SCRIPT_DIR.parents[1] / "v3p5_static" / "scripts" |
| if str(SCRIPT_DIR) not in sys.path: |
| sys.path.insert(0, str(SCRIPT_DIR)) |
| if str(V3P5_SCRIPT_DIR) not in sys.path: |
| sys.path.insert(0, str(V3P5_SCRIPT_DIR)) |
|
|
| from loss_v4p4 import sctm_v4p4_loss, sctm_v4p4_next_visit_metrics |
| from model_v4p4 import SCTMv4p4, SCTMv4p4Config |
| from smoke_v3p4_architecture import build_field_value_mask, load_json, load_service_prior |
| from config_v4p4_train import ( |
| CURRICULUM_PHASES, |
| CURRICULUM_PHASE_TO_LEGACY_STAGE, |
| LEGACY_STAGE_TO_CURRICULUM_PHASE, |
| LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523, |
| ROLLOUT_SCHEDULE_DEFAULTS_V4P4_WORLD_MODEL_20260523, |
| ROLLOUT_WEIGHTS_V4P4_WORLD_MODEL_20260523, |
| curriculum_phase_summary as config_curriculum_phase_summary, |
| ) |
|
|
|
|
| DEFAULT_TENSOR_DIR = Path("/Users/yang/Downloads/精卫_SCTM_数据交付包/60_pretraining/sctm_v3p1/tensor_shards_v3p5_static_collapsed_seq128_stride64") |
| DEFAULT_STAGE0_DIR = Path("/Users/yang/Downloads/精卫_SCTM_数据交付包/60_pretraining/sctm_v3p1/stage0") |
| DEFAULT_OUT_DIR = Path("/Users/yang/Downloads/精卫_SCTM_数据交付包/60_pretraining/sctm_v3p1/v4p4_world_model_train") |
|
|
| TENSOR_KEYS = [ |
| "cat_value_ids", |
| "missing_ids", |
| "numeric_values", |
| "numeric_mask", |
| "ordinal_cbe", |
| "ordinal_mask", |
| "drug_name_ids", |
| "drug_class_ids", |
| "drug_mask", |
| "service_state", |
| "terminal_label", |
| "event_labels", |
| "delta_t_next_log", |
| "time_since_start_days", |
| "visit_year", |
| "visit_indices", |
| "valid_mask", |
| "static_value_ids", |
| ] |
| INTEGER_TENSOR_KEYS = { |
| "cat_value_ids", |
| "missing_ids", |
| "drug_name_ids", |
| "drug_class_ids", |
| "service_state", |
| "terminal_label", |
| "visit_year", |
| "visit_indices", |
| "static_value_ids", |
| } |
| FLOAT_TENSOR_KEYS = {"numeric_values", "ordinal_cbe", "event_labels", "delta_t_next_log", "time_since_start_days"} |
| BOOL_TENSOR_KEYS = {"numeric_mask", "ordinal_mask", "drug_mask", "valid_mask"} |
|
|
|
|
| def resolve_split_path(meta: dict[str, Any], tensor_dir: Path, split: str) -> Path: |
| seq_len = int(meta["max_seq_len"]) |
| local = tensor_dir / f"v3p1_{split}_seq{seq_len}_tensor_shard.npz" |
| if local.exists(): |
| return local |
| original = Path(meta["split_paths"][split]) |
| if original.exists(): |
| return original |
| raise FileNotFoundError(f"Could not resolve {split} tensor shard in {tensor_dir} or {original}") |
|
|
|
|
| def load_npz_to_memory(path: Path) -> dict[str, np.ndarray]: |
| with np.load(path, allow_pickle=True) as z: |
| keep = set(TENSOR_KEYS) | {"patient_ids", "window_starts"} |
| return {k: z[k] for k in z.files if k in keep} |
|
|
|
|
| def inspect_npz_schema(path: Path) -> dict[str, Any]: |
| with np.load(path, allow_pickle=True) as z: |
| return {k: {"shape": list(z[k].shape), "dtype": str(z[k].dtype)} for k in z.files if k in set(TENSOR_KEYS)} |
|
|
|
|
| def tensor_alignment_audit(meta: dict[str, Any], paths: dict[str, Path]) -> dict[str, Any]: |
| expected = { |
| "cat_value_ids": [int(meta["max_seq_len"]), len(meta["cat_cols"])], |
| "missing_ids": [int(meta["max_seq_len"]), len(meta["cat_cols"])], |
| "numeric_values": [int(meta["max_seq_len"]), len(meta["num_cols"])], |
| "numeric_mask": [int(meta["max_seq_len"]), len(meta["num_cols"])], |
| "ordinal_cbe": [int(meta["max_seq_len"]), len(meta["ord_cols"]), int(meta["cbe_dim"])], |
| "ordinal_mask": [int(meta["max_seq_len"]), len(meta["ord_cols"])], |
| "drug_name_ids": [int(meta["max_seq_len"]), int(meta["max_drug_atoms"])], |
| "drug_class_ids": [int(meta["max_seq_len"]), int(meta["max_drug_atoms"])], |
| "drug_mask": [int(meta["max_seq_len"]), int(meta["max_drug_atoms"])], |
| "service_state": [int(meta["max_seq_len"])], |
| "terminal_label": [int(meta["max_seq_len"])], |
| "event_labels": [int(meta["max_seq_len"]), len(meta["event_cols"])], |
| "delta_t_next_log": [int(meta["max_seq_len"])], |
| "time_since_start_days": [int(meta["max_seq_len"])], |
| "visit_year": [int(meta["max_seq_len"])], |
| "visit_indices": [int(meta["max_seq_len"])], |
| "valid_mask": [int(meta["max_seq_len"])], |
| "static_value_ids": [len(meta.get("static_cols", []))], |
| } |
| split_audits: dict[str, Any] = {} |
| failures: list[str] = [] |
| for split, path in paths.items(): |
| schema = inspect_npz_schema(path) |
| missing = [key for key in TENSOR_KEYS if key not in schema] |
| shape_errors: dict[str, Any] = {} |
| for key, tail in expected.items(): |
| if key not in schema: |
| continue |
| actual = schema[key]["shape"] |
| if actual[1:] != tail: |
| shape_errors[key] = {"actual": actual, "expected_tail": ["n_windows", *tail]} |
| if missing or shape_errors: |
| failures.append(split) |
| split_audits[split] = { |
| "path": str(path), |
| "n_windows": int(schema["valid_mask"]["shape"][0]) if "valid_mask" in schema else None, |
| "missing_required_keys": missing, |
| "shape_errors": shape_errors, |
| } |
| return { |
| "status": "passed" if not failures else "failed", |
| "failed_splits": failures, |
| "meta_dims": { |
| "max_seq_len": int(meta["max_seq_len"]), |
| "n_cat_fields": len(meta["cat_cols"]), |
| "n_numeric_fields": len(meta["num_cols"]), |
| "n_ordinal_fields": len(meta["ord_cols"]), |
| "n_events": len(meta["event_cols"]), |
| "n_static_fields": len(meta.get("static_cols", [])), |
| "cat_vocab_size": int(meta["cat_vocab_size"]), |
| "drug_vocab_size": int(meta["drug_vocab_size"]), |
| "static_vocab_size": int(meta.get("static_vocab_size", 0)), |
| }, |
| "splits": split_audits, |
| } |
|
|
|
|
| def train_value_alignment_audit(arrays: dict[str, np.ndarray], meta: dict[str, Any], config: SCTMv4p4Config) -> dict[str, Any]: |
| valid = arrays["valid_mask"].astype(bool) |
| service = arrays["service_state"] |
| terminal = arrays["terminal_label"] |
| event_labels = arrays["event_labels"] |
| delta_days = np.expm1(arrays["delta_t_next_log"]) |
| valid_next = valid[:, :-1] & valid[:, 1:] |
| time_gap = arrays["time_since_start_days"][:, 1:] - arrays["time_since_start_days"][:, :-1] |
| dt_diff = np.abs(time_gap[valid_next] - delta_days[:, :-1][valid_next]) if valid_next.any() else np.array([0.0]) |
| service_ok = bool((service[valid] < config.n_service_states).all()) if valid.any() else True |
| terminal_ok = bool(((terminal[valid] >= 0) & (terminal[valid] <= 4)).all()) if valid.any() else True |
| event_binary = bool(((event_labels == 0.0) | (event_labels == 1.0)).all()) |
| cat_ok = bool(arrays["cat_value_ids"].max() < config.cat_vocab_size) |
| missing_ok = bool(arrays["missing_ids"].max() < config.n_missing) |
| static_ok = bool((not config.n_static_fields) or arrays["static_value_ids"].max() < config.static_vocab_size) |
| delta_ok = bool(float(dt_diff.max()) <= 1.0e-2) |
| status = "passed" if all([service_ok, terminal_ok, event_binary, cat_ok, missing_ok, static_ok, delta_ok]) else "failed" |
| return { |
| "status": status, |
| "valid_slots": int(valid.sum()), |
| "valid_next_positions": int(valid_next.sum()), |
| "service_state_min": int(service[valid].min()) if valid.any() else None, |
| "service_state_max": int(service[valid].max()) if valid.any() else None, |
| "terminal_label_min": int(terminal[valid].min()) if valid.any() else None, |
| "terminal_label_max": int(terminal[valid].max()) if valid.any() else None, |
| "event_labels_last_dim": int(event_labels.shape[-1]), |
| "event_cols": list(meta["event_cols"]), |
| "event_positive_by_col": event_labels[valid].sum(axis=0).astype(int).tolist() if valid.any() else [], |
| "cat_value_min": int(arrays["cat_value_ids"].min()), |
| "cat_value_max": int(arrays["cat_value_ids"].max()), |
| "cat_vocab_size": int(config.cat_vocab_size), |
| "missing_id_min": int(arrays["missing_ids"].min()), |
| "missing_id_max": int(arrays["missing_ids"].max()), |
| "raw_missing_map": dict(meta["missing_map"]), |
| "static_value_min": int(arrays["static_value_ids"].min()) if config.n_static_fields else None, |
| "static_value_max": int(arrays["static_value_ids"].max()) if config.n_static_fields else None, |
| "static_vocab_size": int(config.static_vocab_size), |
| "delta_days_min_valid_next": float(delta_days[:, :-1][valid_next].min()) if valid_next.any() else None, |
| "delta_days_max_valid_next": float(delta_days[:, :-1][valid_next].max()) if valid_next.any() else None, |
| "delta_matches_time_gap_abs_diff_max": float(dt_diff.max()), |
| "delta_matches_time_gap_abs_diff_mean": float(dt_diff.mean()), |
| "service_state_within_config": service_ok, |
| "terminal_label_within_expected_0_to_4": terminal_ok, |
| "event_labels_binary": event_binary, |
| "cat_value_within_vocab": cat_ok, |
| "missing_id_within_v4_vocab": missing_ok, |
| "static_value_within_vocab": static_ok, |
| "delta_matches_time_gap_within_tolerance": delta_ok, |
| } |
|
|
|
|
| def batch_from_indices(arrays: dict[str, np.ndarray], idx: np.ndarray, device: torch.device) -> dict[str, torch.Tensor]: |
| out: dict[str, torch.Tensor] = {} |
| for key in TENSOR_KEYS: |
| if key not in arrays: |
| continue |
| arr = arrays[key][idx] |
| if key in INTEGER_TENSOR_KEYS: |
| out[key] = torch.as_tensor(arr, dtype=torch.long, device=device) |
| elif key in BOOL_TENSOR_KEYS: |
| out[key] = torch.as_tensor(arr, dtype=torch.bool, device=device) |
| elif key in FLOAT_TENSOR_KEYS: |
| out[key] = torch.as_tensor(arr, dtype=torch.float32, device=device) |
| else: |
| raise KeyError(f"untyped tensor key in batch_from_indices: {key}") |
| return out |
|
|
|
|
| class TrainIndexSampler: |
| def __init__(self, rng: np.random.Generator, n: int, sample_with_replacement: bool) -> None: |
| self.rng = rng |
| self.n = int(n) |
| self.sample_with_replacement = bool(sample_with_replacement) |
| self.perm = self.rng.permutation(self.n) if self.n > 0 and not self.sample_with_replacement else np.empty(0, dtype=np.int64) |
| self.ptr = 0 |
|
|
| def next(self, batch_size: int) -> np.ndarray: |
| size = min(int(batch_size), self.n) |
| if size <= 0: |
| return np.empty(0, dtype=np.int64) |
| if self.sample_with_replacement: |
| return self.rng.integers(0, self.n, size=size) |
| if self.ptr + size <= self.n: |
| idx = self.perm[self.ptr : self.ptr + size] |
| self.ptr += size |
| return idx |
| tail = self.perm[self.ptr :] |
| self.perm = self.rng.permutation(self.n) |
| take = size - int(tail.shape[0]) |
| head = self.perm[:take] |
| self.ptr = take |
| return np.concatenate([tail, head]) |
|
|
|
|
| def autocast_context(device: torch.device, precision: str): |
| enabled = device.type == "cuda" and precision in {"bf16", "fp16"} |
| dtype = torch.bfloat16 if precision == "bf16" else torch.float16 |
| return torch.autocast(device_type=device.type, dtype=dtype, enabled=enabled) |
|
|
|
|
| def cosine_lr(step: int, *, max_steps: int, warmup_steps: int, lr: float, min_lr: float) -> float: |
| if step <= warmup_steps: |
| return lr * step / max(1, warmup_steps) |
| progress = (step - warmup_steps) / max(1, max_steps - warmup_steps) |
| coeff = 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) |
| return min_lr + coeff * (lr - min_lr) |
|
|
|
|
| def build_conditional_specs(meta: dict[str, Any], vocab: dict[str, int]) -> tuple[dict[str, Any], ...]: |
| cat_cols = list(meta["cat_cols"]) |
|
|
| def field_idx(name: str) -> int: |
| return cat_cols.index(name) if name in cat_cols else -1 |
|
|
| specs = [] |
| tables = [ |
| ("是否转诊", ["是"], ["转诊类型", "转诊原因", "转诊至机构"]), |
| ("治疗方式", ["住院"], ["本次入院形式", "本次住院名称", "本次入院日期", "末次出院日期"]), |
| ] |
| for parent, active_values, children in tables: |
| pidx = field_idx(parent) |
| child_indices = [field_idx(c) for c in children if field_idx(c) >= 0] |
| active_ids = [int(vocab[f"{parent}={v}"]) for v in active_values if f"{parent}={v}" in vocab] |
| if pidx >= 0 and child_indices and active_ids: |
| specs.append({"parent": parent, "parent_idx": pidx, "active_value_ids": active_ids, "child_indices": child_indices}) |
| return tuple(specs) |
|
|
|
|
| def build_color_map(meta: dict[str, Any], vocab: dict[str, int]) -> tuple[int, tuple[tuple[int, int], ...]]: |
| cat_cols = list(meta["cat_cols"]) |
| color_field = "新规评估信号（颜色）" |
| if color_field not in cat_cols: |
| return -1, tuple() |
| order = {"绿色": 0, "绿": 0, "蓝色": 1, "蓝": 1, "黄色": 2, "黄": 2, "橙色": 3, "橙": 3, "红色": 4, "红": 4} |
| pairs: list[tuple[int, int]] = [] |
| prefix = f"{color_field}=" |
| for token, value_id in vocab.items(): |
| if not token.startswith(prefix): |
| continue |
| value = token.removeprefix(prefix) |
| if value in order: |
| pairs.append((int(value_id), int(order[value]))) |
| return cat_cols.index(color_field), tuple(sorted(pairs)) |
|
|
|
|
| def metadata_config(meta: dict[str, Any], vocab: dict[str, int], args: argparse.Namespace) -> SCTMv4p4Config: |
| color_idx, color_map = build_color_map(meta, vocab) |
| return SCTMv4p4Config( |
| cat_vocab_size=int(meta["cat_vocab_size"]), |
| drug_vocab_size=int(meta["drug_vocab_size"]), |
| n_cat_fields=len(meta["cat_cols"]), |
| n_numeric_fields=len(meta["num_cols"]), |
| n_ordinal_fields=len(meta["ord_cols"]), |
| cbe_dim=int(meta["cbe_dim"]), |
| max_drug_atoms=int(meta["max_drug_atoms"]), |
| n_events=len(meta["event_cols"]), |
| static_vocab_size=int(meta.get("static_vocab_size", 0)), |
| n_static_fields=len(meta.get("static_cols", [])), |
| d_model=args.d_model, |
| n_heads=args.n_heads, |
| n_layers=args.n_layers, |
| dropout=args.dropout, |
| n_missing=args.n_missing, |
| year_min=int(meta.get("year_min", 2010)), |
| year_max=int(meta.get("year_max", 2026)), |
| conditional_specs=build_conditional_specs(meta, vocab), |
| color_cat_field_index=color_idx, |
| color_value_to_class=color_map, |
| ) |
|
|
|
|
| def try_train_microbatch( |
| model: SCTMv4p4, |
| config: SCTMv4p4Config, |
| arrays: dict[str, np.ndarray], |
| batch_size: int, |
| device: torch.device, |
| args: argparse.Namespace, |
| ) -> bool: |
| idx = np.arange(min(batch_size, arrays["valid_mask"].shape[0])) |
| try: |
| model.train() |
| batch = batch_from_indices(arrays, idx, device) |
| rollout_steps, teacher_forcing_rate, gumbel_temperature = rollout_schedule(args, 1) |
| with autocast_context(device, args.precision): |
| out = model( |
| batch, |
| rollout_steps=rollout_steps, |
| teacher_forcing_rate=teacher_forcing_rate, |
| gumbel_temperature=gumbel_temperature, |
| compute_pwe_diagnostics=False, |
| ) |
| loss, _ = sctm_v4p4_loss( |
| out, |
| batch, |
| config, |
| field_loss_weight=args.field_loss_weight, |
| numeric_loss_weight=args.numeric_loss_weight, |
| ordinal_loss_weight=args.ordinal_loss_weight, |
| missingness_loss_weight=args.missingness_loss_weight, |
| pre_event_risk_weight=args.pre_event_risk_weight, |
| event_risk_weight=args.event_risk_weight, |
| post_event_risk_weight=args.post_event_risk_weight, |
| event_generation_weight=args.event_generation_weight, |
| ontology_event_weight=args.ontology_event_weight, |
| history_aux_weight=args.history_aux_weight, |
| contact_aux_weight=args.contact_aux_weight, |
| color_loss_weight=args.color_loss_weight, |
| pwe_pre_nll_weight=args.pwe_pre_nll_weight, |
| pwe_contact_nll_weight=args.pwe_contact_nll_weight, |
| pwe_post_nll_weight=args.pwe_post_nll_weight, |
| pwe_pre_distill_weight=args.pwe_pre_distill_weight, |
| pwe_contact_distill_weight=args.pwe_contact_distill_weight, |
| rollout_weights=(args.rollout_weight_1, args.rollout_weight_2, args.rollout_weight_3), |
| return_metrics=False, |
| ) |
| loss.backward() |
| model.zero_grad(set_to_none=True) |
| if device.type == "cuda": |
| torch.cuda.empty_cache() |
| return True |
| except RuntimeError as exc: |
| if "out of memory" in str(exc).lower(): |
| model.zero_grad(set_to_none=True) |
| if device.type == "cuda": |
| torch.cuda.empty_cache() |
| return False |
| raise |
|
|
|
|
| def select_batch_size( |
| model: SCTMv4p4, |
| config: SCTMv4p4Config, |
| arrays: dict[str, np.ndarray], |
| args: argparse.Namespace, |
| device: torch.device, |
| ) -> int: |
| if not args.auto_batch: |
| return args.batch_size |
| candidates = [] |
| value = args.batch_size |
| while value >= args.min_batch_size: |
| candidates.append(value) |
| value = int(value * 0.75) |
| candidates.append(args.min_batch_size) |
| for batch_size in sorted(set(candidates), reverse=True): |
| if try_train_microbatch(model, config, arrays, batch_size, device, args): |
| return batch_size |
| raise RuntimeError(f"No valid batch size found down to {args.min_batch_size}") |
|
|
|
|
| def adapt_v3p5_tensor_for_v4p4(key: str, src: torch.Tensor, dst: torch.Tensor, config: SCTMv4p4Config) -> torch.Tensor | None: |
| """Copy v3p5 weights into v4p4 tensors whose leading semantics expanded. |
| |
| Generic prefix-copying is intentionally avoided for flattened heads because |
| rows encode `(field, class)` pairs. Only known v3p5->v4p4 expansions are |
| adapted here. |
| """ |
|
|
| if key == "tab_encoder.axis_slot_emb" and src.ndim == 4 and dst.ndim == 4 and src.shape[-1] == dst.shape[-1]: |
| adapted = torch.zeros_like(dst) |
| n_slots = min(src.shape[2], dst.shape[2]) |
| adapted[:, :, :n_slots, :] = src[:, :, :n_slots, :].to(dtype=dst.dtype, device=dst.device) |
| return adapted |
| if key == "tab_encoder.missing_emb.weight" and src.ndim == 2 and dst.ndim == 2 and src.shape[1] == dst.shape[1]: |
| adapted = torch.zeros_like(dst) |
| n_rows = min(src.shape[0], dst.shape[0]) |
| adapted[:n_rows, :] = src[:n_rows, :].to(dtype=dst.dtype, device=dst.device) |
| return adapted |
| if key == "missingness_head.weight" and src.ndim == 2 and dst.ndim == 2 and src.shape[1] == dst.shape[1]: |
| n_fields = int(config.n_cat_fields) |
| if src.shape[0] % n_fields == 0 and dst.shape[0] == n_fields * int(config.n_missing): |
| src_missing = src.shape[0] // n_fields |
| n_missing = min(src_missing, int(config.n_missing)) |
| adapted = torch.zeros_like(dst) |
| src_view = src.reshape(n_fields, src_missing, src.shape[1]) |
| dst_view = adapted.reshape(n_fields, int(config.n_missing), dst.shape[1]) |
| dst_view[:, :n_missing, :] = src_view[:, :n_missing, :].to(dtype=dst.dtype, device=dst.device) |
| return adapted |
| if key == "missingness_head.bias" and src.ndim == 1 and dst.ndim == 1: |
| n_fields = int(config.n_cat_fields) |
| if src.shape[0] % n_fields == 0 and dst.shape[0] == n_fields * int(config.n_missing): |
| src_missing = src.shape[0] // n_fields |
| n_missing = min(src_missing, int(config.n_missing)) |
| adapted = torch.zeros_like(dst) |
| src_view = src.reshape(n_fields, src_missing) |
| dst_view = adapted.reshape(n_fields, int(config.n_missing)) |
| dst_view[:, :n_missing] = src_view[:, :n_missing].to(dtype=dst.dtype, device=dst.device) |
| return adapted |
| return None |
|
|
|
|
| def load_compatible_state_dict(model: torch.nn.Module, checkpoint: Path, device: torch.device) -> dict[str, Any]: |
| payload = torch.load(checkpoint, map_location="cpu", weights_only=False) |
| payload.pop("optimizer_state_dict", None) |
| state = payload["model_state_dict"] |
| if any(k.startswith("_orig_mod.") for k in state): |
| state = {k.removeprefix("_orig_mod."): v for k, v in state.items()} |
| own = model.state_dict() |
| compatible: dict[str, torch.Tensor] = {} |
| skipped: dict[str, list[int]] = {} |
| adapted: dict[str, dict[str, list[int]]] = {} |
| base_model = model._orig_mod if hasattr(model, "_orig_mod") else model |
| config = base_model.config |
| for key, value in state.items(): |
| if key not in own: |
| continue |
| target = own[key] |
| if tuple(value.shape) == tuple(target.shape): |
| compatible[key] = value |
| continue |
| adapted_value = adapt_v3p5_tensor_for_v4p4(key, value, target, config) |
| if adapted_value is None: |
| skipped[key] = list(value.shape) |
| continue |
| compatible[key] = adapted_value |
| adapted[key] = {"source_shape": list(value.shape), "target_shape": list(target.shape)} |
| for suffix in ("weight", "bias"): |
| src_key = f"ontology_event_head.{suffix}" |
| dst_key = f"event_generation_head.{suffix}" |
| if dst_key not in compatible and src_key in state and dst_key in own and tuple(state[src_key].shape) == tuple(own[dst_key].shape): |
| compatible[dst_key] = state[src_key].to(dtype=own[dst_key].dtype, device=own[dst_key].device) |
| adapted[dst_key] = {"source_key": src_key, "source_shape": list(state[src_key].shape), "target_shape": list(own[dst_key].shape)} |
| incompatible = model.load_state_dict(compatible, strict=False) |
| return { |
| "checkpoint": str(checkpoint), |
| "source_step": int(payload.get("step", -1)), |
| "loaded_keys": len(compatible), |
| "shape_adapted": adapted, |
| "missing_keys": list(incompatible.missing_keys), |
| "unexpected_keys": list(incompatible.unexpected_keys), |
| "shape_skipped": skipped, |
| } |
|
|
|
|
| def unwrap_compiled_model(model: torch.nn.Module) -> torch.nn.Module: |
| return model._orig_mod if hasattr(model, "_orig_mod") else model |
|
|
|
|
| def strip_compile_prefix(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: |
| if any(key.startswith("_orig_mod.") for key in state): |
| return {key.removeprefix("_orig_mod."): value for key, value in state.items()} |
| return state |
|
|
|
|
| V4P4_ALLOWED_RESUME_MISSING_PREFIXES = ( |
| "deployable_risk_post_residual_head.", |
| ) |
|
|
|
|
| def load_resume_state_compatible(model: torch.nn.Module, state: dict[str, torch.Tensor]) -> dict[str, Any]: |
| incompatible = unwrap_compiled_model(model).load_state_dict(state, strict=False) |
| bad_missing = [key for key in incompatible.missing_keys if not key.startswith(V4P4_ALLOWED_RESUME_MISSING_PREFIXES)] |
| if bad_missing or incompatible.unexpected_keys: |
| raise RuntimeError( |
| "resume checkpoint is not compatible with current v4p4 model: " |
| f"bad_missing={bad_missing}, unexpected={list(incompatible.unexpected_keys)}" |
| ) |
| return { |
| "missing_keys": list(incompatible.missing_keys), |
| "unexpected_keys": list(incompatible.unexpected_keys), |
| "allowed_missing_prefixes": list(V4P4_ALLOWED_RESUME_MISSING_PREFIXES), |
| } |
|
|
|
|
| BACKBONE_PREFIXES = ("tab_encoder", "traj_decoder", "static_encoder", "static_injection_scale") |
|
|
|
|
| def normalize_curriculum_args(args: argparse.Namespace) -> argparse.Namespace: |
| """Keep legacy --stage reproducible while exposing the unified-model curriculum.""" |
|
|
| requested = getattr(args, "curriculum_phase", "") or "" |
| if requested: |
| if requested not in CURRICULUM_PHASE_TO_LEGACY_STAGE: |
| raise ValueError(f"unknown curriculum phase: {requested}") |
| legacy = CURRICULUM_PHASE_TO_LEGACY_STAGE[requested] |
| if getattr(args, "stage", legacy) != legacy: |
| raise ValueError(f"--stage {args.stage!r} conflicts with --curriculum-phase {requested!r}") |
| args.stage = legacy |
| args.curriculum_phase = requested |
| else: |
| args.curriculum_phase = LEGACY_STAGE_TO_CURRICULUM_PHASE.get(args.stage, args.stage) |
| return args |
|
|
|
|
| def curriculum_phase_summary(args: argparse.Namespace) -> dict[str, Any]: |
| phase = getattr(args, "curriculum_phase", LEGACY_STAGE_TO_CURRICULUM_PHASE.get(args.stage, args.stage)) |
| return config_curriculum_phase_summary(phase, legacy_stage=args.stage) |
|
|
|
|
| def apply_stage_trainability(model: torch.nn.Module, args: argparse.Namespace) -> dict[str, Any]: |
| for param in model.parameters(): |
| param.requires_grad = True |
| frozen: list[str] = [] |
| phase = getattr(args, "curriculum_phase", LEGACY_STAGE_TO_CURRICULUM_PHASE.get(args.stage, args.stage)) |
| if phase == "adapter_warmup" and args.freeze_backbone_stage1: |
| for name, param in model.named_parameters(): |
| if name == "static_injection_scale" or name.startswith(BACKBONE_PREFIXES[:3]): |
| param.requires_grad = False |
| frozen.append(name) |
| trainable_params = int(sum(p.numel() for p in model.parameters() if p.requires_grad)) |
| frozen_params = int(sum(p.numel() for p in model.parameters() if not p.requires_grad)) |
| if trainable_params == 0: |
| raise RuntimeError("stage trainability configuration froze all parameters") |
| return { |
| "stage": args.stage, |
| "curriculum_phase": phase, |
| "freeze_backbone_stage1": bool(args.freeze_backbone_stage1), |
| "trainable_params": trainable_params, |
| "frozen_params": frozen_params, |
| "frozen_param_names_sample": frozen[:20], |
| "n_frozen_param_names": len(frozen), |
| } |
|
|
|
|
| def curriculum_schedule(args: argparse.Namespace, step: int) -> tuple[int, float, float]: |
| phase = getattr(args, "curriculum_phase", LEGACY_STAGE_TO_CURRICULUM_PHASE.get(args.stage, args.stage)) |
| if phase == "joint_one_step": |
| return 1, 1.0, args.gumbel_temperature_start |
| return 1, 1.0, args.gumbel_temperature_start |
|
|
|
|
| def rollout_schedule(args: argparse.Namespace, step: int) -> tuple[int, float, float]: |
| return curriculum_schedule(args, step) |
|
|
|
|
| @torch.no_grad() |
| def evaluate( |
| model: SCTMv4p4, |
| config: SCTMv4p4Config, |
| arrays: dict[str, np.ndarray], |
| batch_size: int, |
| max_batches: int, |
| device: torch.device, |
| precision: str, |
| loss_kwargs: dict[str, Any] | None = None, |
| rollout_steps: int = 1, |
| teacher_forcing_rate: float = 1.0, |
| gumbel_temperature: float = 1.0, |
| ) -> dict[str, float]: |
| model.eval() |
| weighted: dict[str, float] = {} |
| weights: dict[str, float] = {} |
| batches = 0 |
| n = arrays["valid_mask"].shape[0] |
| for start in range(0, n, batch_size): |
| if max_batches > 0 and batches >= max_batches: |
| break |
| idx = np.arange(start, min(start + batch_size, n)) |
| batch = batch_from_indices(arrays, idx, device) |
| with autocast_context(device, precision): |
| out = model( |
| batch, |
| rollout_steps=rollout_steps, |
| teacher_forcing_rate=teacher_forcing_rate, |
| gumbel_temperature=gumbel_temperature, |
| compute_pwe_diagnostics=False, |
| ) |
| _, loss_metrics = sctm_v4p4_loss(out, batch, config, **(loss_kwargs or {}), return_metrics=True) |
| metric = sctm_v4p4_next_visit_metrics(out, batch, config) |
| w = float(loss_metrics.get("valid_next_positions", 1.0)) |
| for row in (loss_metrics, {f"metric_{k}": v for k, v in metric.items()}): |
| for key, value in row.items(): |
| if isinstance(value, (int, float)) and np.isfinite(value): |
| if key == "valid_next_positions": |
| continue |
| weighted[key] = weighted.get(key, 0.0) + float(value) * w |
| weights[key] = weights.get(key, 0.0) + w |
| weighted["valid_next_positions"] = weighted.get("valid_next_positions", 0.0) + w |
| batches += 1 |
| out = {key: weighted[key] / max(1.0e-12, weights[key]) for key in weights} |
| out["valid_next_positions"] = weighted.get("valid_next_positions", 0.0) |
| out["eval_batches"] = float(batches) |
| return out |
|
|
|
|
| def json_ready(value: Any) -> Any: |
| if isinstance(value, Path): |
| return str(value) |
| if isinstance(value, dict): |
| return {str(k): json_ready(v) for k, v in value.items()} |
| if isinstance(value, (list, tuple)): |
| return [json_ready(v) for v in value] |
| if isinstance(value, np.generic): |
| return value.item() |
| return value |
|
|
|
|
| def save_checkpoint( |
| path: Path, |
| *, |
| model: SCTMv4p4, |
| optimizer: torch.optim.Optimizer, |
| step: int, |
| config: SCTMv4p4Config, |
| tensor_metadata: dict[str, Any], |
| args: argparse.Namespace, |
| train_log_tail: dict[str, Any], |
| ) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| torch.save( |
| { |
| "step": step, |
| "model_state_dict": unwrap_compiled_model(model).state_dict(), |
| "optimizer_state_dict": optimizer.state_dict(), |
| "config": config.__dict__, |
| "tensor_metadata": tensor_metadata, |
| "args": json_ready(vars(args)), |
| "train_log_tail": train_log_tail, |
| }, |
| path, |
| ) |
|
|
|
|
| def maybe_upload(args: argparse.Namespace, folder: Path, stage: str) -> None: |
| if not args.hub_repo_id: |
| return |
| token = os.environ.get("HF_TOKEN") |
| if not token: |
| print(json.dumps({"stage": "hub_upload_skipped", "reason": "HF_TOKEN missing", "upload_stage": stage}), flush=True) |
| return |
| from huggingface_hub import HfApi |
|
|
| api = HfApi(token=token) |
| api.upload_folder( |
| folder_path=str(folder), |
| path_in_repo=args.hub_path, |
| repo_id=args.hub_repo_id, |
| repo_type=args.hub_repo_type, |
| token=token, |
| commit_message=f"Upload SCTM-v4p4 {args.run_label} at {stage}", |
| ) |
|
|
|
|
| def run(args: argparse.Namespace) -> dict[str, Any]: |
| args = normalize_curriculum_args(args) |
| args.out_dir.mkdir(parents=True, exist_ok=True) |
| torch.manual_seed(args.seed) |
| np.random.seed(args.seed) |
| if args.matmul_precision: |
| torch.set_float32_matmul_precision(args.matmul_precision) |
| device = torch.device(args.device) |
| meta = load_json(args.tensor_dir / "tensor_metadata.json") |
| vocab = load_json(args.tensor_dir / "cat_value_vocab.json") |
| paths = {split: resolve_split_path(meta, args.tensor_dir, split) for split in ["train", "val", "test"]} |
| shape_audit = tensor_alignment_audit(meta, paths) |
| print(json.dumps({"stage": "tensor_alignment_audit", **shape_audit}, ensure_ascii=False), flush=True) |
| if shape_audit["status"] != "passed": |
| raise RuntimeError(f"tensor schema does not align with v4p4 config contract: {shape_audit}") |
| print(json.dumps({"stage": "load_npz_to_memory", "paths": {k: str(v) for k, v in paths.items()}}, ensure_ascii=False), flush=True) |
| train_arrays = load_npz_to_memory(paths["train"]) |
| val_arrays = load_npz_to_memory(paths["val"]) |
| missing = [key for key in TENSOR_KEYS if key not in train_arrays] |
| if missing: |
| raise KeyError(f"train tensor shard missing required v4p4 keys: {missing}") |
| field_value_mask, field_mask_audit = build_field_value_mask(meta, vocab, {k: train_arrays[k] for k in ["cat_value_ids", "missing_ids"]}) |
| service_prior, prior_audit = load_service_prior(args.stage0_dir, int(meta.get("n_service_states", 8)), args.service_prior_file) |
| config = metadata_config(meta, vocab, args) |
| value_audit = train_value_alignment_audit(train_arrays, meta, config) |
| print(json.dumps({"stage": "train_value_alignment_audit", **value_audit}, ensure_ascii=False), flush=True) |
| if value_audit["status"] != "passed": |
| raise RuntimeError(f"tensor values do not align with v4p4 config contract: {value_audit}") |
| model = SCTMv4p4(config, field_value_mask=field_value_mask, service_prior_bias=service_prior).to(device) |
| init_audit: dict[str, Any] = {} |
| if args.init_checkpoint: |
| init_audit = load_compatible_state_dict(model, args.init_checkpoint, device) |
| print(json.dumps({"stage": "init_checkpoint_loaded", **init_audit}, ensure_ascii=False), flush=True) |
| trainability_audit = apply_stage_trainability(model, args) |
| param_count = int(sum(p.numel() for p in model.parameters())) |
| selected_batch_size = select_batch_size(model, config, train_arrays, args, device) |
| compile_enabled = bool(args.compile and hasattr(torch, "compile")) |
| if compile_enabled: |
| model = torch.compile(model) |
| print( |
| json.dumps( |
| { |
| "stage": "model_ready", |
| "param_count": param_count, |
| "selected_batch_size": selected_batch_size, |
| "requested_batch_size": args.batch_size, |
| "device": str(device), |
| "run_stage": args.stage, |
| "cuda_device": torch.cuda.get_device_name(0) if device.type == "cuda" else None, |
| "compile_enabled": compile_enabled, |
| "sample_with_replacement": bool(args.sample_with_replacement), |
| }, |
| ensure_ascii=False, |
| ), |
| flush=True, |
| ) |
| use_fused = bool(args.fused_adamw and device.type == "cuda") |
| optimizer = torch.optim.AdamW( |
| [p for p in model.parameters() if p.requires_grad], |
| lr=args.learning_rate, |
| betas=(args.beta1, args.beta2), |
| weight_decay=args.weight_decay, |
| fused=use_fused, |
| ) |
| resume_step = 0 |
| resume_audit: dict[str, Any] = {} |
| if args.resume_checkpoint: |
| payload = torch.load(args.resume_checkpoint, map_location=device, weights_only=False) |
| state = strip_compile_prefix(payload["model_state_dict"]) |
| resume_state_audit = load_resume_state_compatible(model, state) |
| resume_audit = { |
| "checkpoint": str(args.resume_checkpoint), |
| "source_step": int(payload.get("step", 0)), |
| "model_state_loaded": True, |
| "model_state_audit": resume_state_audit, |
| "optimizer_state_loaded": False, |
| "optimizer_state_skip_reason": "", |
| } |
| if "optimizer_state_dict" in payload: |
| try: |
| optimizer.load_state_dict(payload["optimizer_state_dict"]) |
| resume_audit["optimizer_state_loaded"] = True |
| except ValueError as exc: |
| resume_audit["optimizer_state_skip_reason"] = str(exc) |
| resume_step = int(payload.get("step", 0)) |
| print(json.dumps({"stage": "resume_checkpoint_loaded", **resume_audit}, ensure_ascii=False), flush=True) |
|
|
| run_meta = { |
| "run_label": args.run_label, |
| "model_family": "sctm_v4p4_world_model", |
| "stage": args.stage, |
| "curriculum_phase": args.curriculum_phase, |
| "curriculum_phase_summary": curriculum_phase_summary(args), |
| "resolved_train_config": { |
| "preset": "v4p4_world_model_20260523", |
| "loss_weights": {key: getattr(args, key) for key in LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523}, |
| "rollout_weights": { |
| "rollout_weight_1": args.rollout_weight_1, |
| "rollout_weight_2": args.rollout_weight_2, |
| "rollout_weight_3": args.rollout_weight_3, |
| }, |
| "rollout_schedule": { |
| "rollout_steps": args.rollout_steps, |
| "teacher_forcing_rate": args.teacher_forcing_rate, |
| "teacher_forcing_start": args.teacher_forcing_start, |
| "teacher_forcing_end": args.teacher_forcing_end, |
| "gumbel_temperature_start": args.gumbel_temperature_start, |
| "gumbel_temperature_end": args.gumbel_temperature_end, |
| }, |
| }, |
| "param_count": param_count, |
| "config": json_ready(config.__dict__), |
| "tensor_dir": str(args.tensor_dir), |
| "tensor_paths": {k: str(v) for k, v in paths.items()}, |
| "tensor_alignment_audit": shape_audit, |
| "train_value_alignment_audit": value_audit, |
| "field_value_mask_audit": field_mask_audit, |
| "service_prior_audit": prior_audit, |
| "init_audit": init_audit, |
| "trainability_audit": trainability_audit, |
| "selected_batch_size": selected_batch_size, |
| "requested_batch_size": args.batch_size, |
| "device": str(device), |
| "precision": args.precision, |
| "compile_enabled": compile_enabled, |
| "sample_with_replacement": bool(args.sample_with_replacement), |
| "finite_check_every": int(args.finite_check_every), |
| "resume_checkpoint": str(args.resume_checkpoint) if args.resume_checkpoint else "", |
| "resume_step": resume_step, |
| "resume_audit": resume_audit, |
| } |
| (args.out_dir / "run_meta.json").write_text(json.dumps(run_meta, ensure_ascii=False, indent=2), encoding="utf-8") |
|
|
| rng = np.random.default_rng(args.seed) |
| n_train = train_arrays["valid_mask"].shape[0] |
| train_sampler = TrainIndexSampler(rng, n_train, args.sample_with_replacement) |
| start_time = time.time() |
| log_rows: list[dict[str, Any]] = [] |
| latest_eval: dict[str, float] = {} |
| best_metric_value: float | None = None |
| best_step = resume_step |
| evals_since_best = 0 |
| early_stopped = False |
| stop_reason = "max_steps" |
| loss_kwargs = { |
| key: getattr(args, key) |
| for key in LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523 |
| } | { |
| "rollout_weights": (args.rollout_weight_1, args.rollout_weight_2, args.rollout_weight_3), |
| } |
|
|
| for step in range(resume_step + 1, args.max_steps + 1): |
| need_step_metrics = ( |
| step == 1 |
| or step % args.log_every == 0 |
| or step % args.eval_every == 0 |
| or step % args.save_every == 0 |
| or step % args.hub_upload_every == 0 |
| or step == args.max_steps |
| ) |
| need_finite_check = args.finite_check_every > 0 and ( |
| step == 1 |
| or step % args.finite_check_every == 0 |
| or step == args.max_steps |
| ) |
| lr = cosine_lr(step, max_steps=args.max_steps, warmup_steps=args.warmup_steps, lr=args.learning_rate, min_lr=args.min_learning_rate) |
| for group in optimizer.param_groups: |
| group["lr"] = lr |
| step_start = time.time() |
| if device.type == "cuda" and need_step_metrics: |
| torch.cuda.reset_peak_memory_stats(device) |
| model.train() |
| optimizer.zero_grad(set_to_none=True) |
| rollout_steps, teacher_forcing_rate, gumbel_temperature = rollout_schedule(args, step) |
| accum_losses: list[dict[str, float]] = [] |
| for _ in range(args.grad_accum_steps): |
| idx = train_sampler.next(selected_batch_size) |
| batch = batch_from_indices(train_arrays, idx, device) |
| with autocast_context(device, args.precision): |
| out = model( |
| batch, |
| rollout_steps=rollout_steps, |
| teacher_forcing_rate=teacher_forcing_rate, |
| gumbel_temperature=gumbel_temperature, |
| compute_pwe_diagnostics=need_step_metrics, |
| ) |
| loss, loss_metrics = sctm_v4p4_loss( |
| out, |
| batch, |
| config, |
| **loss_kwargs, |
| return_metrics=need_step_metrics, |
| ) |
| loss = loss / args.grad_accum_steps |
| if need_finite_check and not torch.isfinite(loss): |
| raise RuntimeError(f"non-finite loss at step {step}: {float(loss.detach().cpu())}") |
| loss.backward() |
| if need_step_metrics: |
| accum_losses.append(loss_metrics) |
| grad_norm_tensor = torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) |
| optimizer.step() |
| if need_step_metrics and device.type == "cuda": |
| torch.cuda.synchronize(device) |
| step_sec = max(1.0e-12, time.time() - step_start) |
| effective_windows = int(selected_batch_size * args.grad_accum_steps) |
| seq_len = int(train_arrays["valid_mask"].shape[1]) |
| base_model = model._orig_mod if hasattr(model, "_orig_mod") else model |
| row: dict[str, Any] = { |
| "step": step, |
| "curriculum_phase": args.curriculum_phase, |
| "legacy_stage": args.stage, |
| "lr": lr, |
| "grad_norm": float(grad_norm_tensor.detach().cpu()) if need_step_metrics else float("nan"), |
| "elapsed_sec": round(time.time() - start_time, 3), |
| "param_count": param_count, |
| "step_sec": round(step_sec, 4), |
| "selected_batch_size": selected_batch_size, |
| "effective_windows": effective_windows, |
| "visit_slots_per_sec": float(effective_windows * seq_len / step_sec), |
| "rollout_steps": rollout_steps, |
| "teacher_forcing_rate": teacher_forcing_rate, |
| "gumbel_temperature": gumbel_temperature, |
| "contact_injection_scale": ( |
| float(base_model.contact_injector.scale.detach().float().cpu().item()) |
| if need_step_metrics and hasattr(base_model, "contact_injector") |
| else float("nan") |
| ), |
| } |
| if device.type == "cuda" and need_step_metrics: |
| row["cuda_max_memory_allocated_gb"] = float(torch.cuda.max_memory_allocated(device) / 1024**3) |
| row["cuda_memory_reserved_gb"] = float(torch.cuda.memory_reserved(device) / 1024**3) |
| for key in sorted(set().union(*(x.keys() for x in accum_losses))): |
| vals = [x[key] for x in accum_losses if key in x and isinstance(x[key], (int, float))] |
| if vals: |
| row[f"train_{key}"] = float(np.mean(vals)) |
| if step == 1 or step % args.eval_every == 0 or step == args.max_steps: |
| val_metrics = evaluate( |
| model, |
| config, |
| val_arrays, |
| args.eval_batch_size, |
| args.eval_batches, |
| device, |
| args.precision, |
| loss_kwargs=loss_kwargs, |
| rollout_steps=rollout_steps, |
| teacher_forcing_rate=teacher_forcing_rate, |
| gumbel_temperature=gumbel_temperature, |
| ) |
| latest_eval = {f"val_{k}": v for k, v in val_metrics.items()} |
| row.update(latest_eval) |
| metric_value = row.get(args.early_stopping_metric) |
| if args.early_stopping_patience > 0 and isinstance(metric_value, (int, float)) and math.isfinite(float(metric_value)): |
| metric_value = float(metric_value) |
| improved = best_metric_value is None or ( |
| metric_value < best_metric_value - args.early_stopping_min_delta |
| if args.early_stopping_mode == "min" |
| else metric_value > best_metric_value + args.early_stopping_min_delta |
| ) |
| if improved: |
| best_metric_value = metric_value |
| best_step = step |
| evals_since_best = 0 |
| row["early_stopping_best_step"] = best_step |
| row["early_stopping_best_metric"] = best_metric_value |
| save_checkpoint(args.out_dir / "checkpoints" / "best.pt", model=model, optimizer=optimizer, step=step, config=config, tensor_metadata=meta, args=args, train_log_tail=row) |
| else: |
| evals_since_best += 1 |
| row["early_stopping_evals_since_best"] = evals_since_best |
| row["early_stopping_best_step"] = best_step |
| row["early_stopping_best_metric"] = best_metric_value |
| log_rows.append(row) |
| if step == 1 or step % args.log_every == 0 or step % args.eval_every == 0 or step == args.max_steps: |
| print(json.dumps(row, ensure_ascii=False), flush=True) |
| should_stop = ( |
| args.early_stopping_patience > 0 |
| and step >= args.early_stopping_start_step |
| and evals_since_best >= args.early_stopping_patience |
| ) |
| should_save = step % args.save_every == 0 or step == args.max_steps or should_stop |
| should_upload = step % args.hub_upload_every == 0 or step == args.max_steps or should_stop |
| if should_save or should_upload: |
| pd.DataFrame(log_rows).to_csv(args.out_dir / "train_log.csv", index=False) |
| save_checkpoint(args.out_dir / "checkpoints" / "latest.pt", model=model, optimizer=optimizer, step=step, config=config, tensor_metadata=meta, args=args, train_log_tail=row) |
| if args.keep_step_checkpoints and should_save: |
| save_checkpoint(args.out_dir / "checkpoints" / f"step_{step:06d}.pt", model=model, optimizer=optimizer, step=step, config=config, tensor_metadata=meta, args=args, train_log_tail=row) |
| if should_upload: |
| maybe_upload(args, args.out_dir, f"step_{step}") |
| if should_stop: |
| early_stopped = True |
| stop_reason = f"early_stopping_patience_{args.early_stopping_patience}" |
| print( |
| json.dumps( |
| { |
| "stage": "early_stopping_triggered", |
| "step": step, |
| "metric": args.early_stopping_metric, |
| "best_step": best_step, |
| "best_metric": best_metric_value, |
| "evals_since_best": evals_since_best, |
| }, |
| ensure_ascii=False, |
| ), |
| flush=True, |
| ) |
| break |
|
|
| summary = { |
| "status": "completed", |
| "run_label": args.run_label, |
| "stop_reason": stop_reason, |
| "early_stopped": early_stopped, |
| "best_step": best_step, |
| "best_metric": best_metric_value, |
| "early_stopping_metric": args.early_stopping_metric if args.early_stopping_patience > 0 else "", |
| "selected_batch_size": selected_batch_size, |
| "max_steps": args.max_steps, |
| "actual_steps": int(log_rows[-1]["step"]) if log_rows else resume_step, |
| "latest_eval": latest_eval, |
| "elapsed_sec": time.time() - start_time, |
| "out_dir": str(args.out_dir), |
| "checkpoint": str(args.out_dir / "checkpoints" / "latest.pt"), |
| } |
| (args.out_dir / "train_result.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") |
| pd.DataFrame(log_rows).to_csv(args.out_dir / "train_log.csv", index=False) |
| maybe_upload(args, args.out_dir, "final") |
| return summary |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Train the SCTM-v4.4 closed-loop world-model heads on v3p5-compatible tensors.") |
| parser.add_argument("--tensor-dir", type=Path, default=DEFAULT_TENSOR_DIR) |
| parser.add_argument("--stage0-dir", type=Path, default=DEFAULT_STAGE0_DIR) |
| parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR) |
| parser.add_argument("--service-prior-file", type=str, default="service_state_transitions_train.json") |
| parser.add_argument("--run-label", type=str, default="sctm_v4p4_world_model") |
| parser.add_argument("--stage", choices=["stage1", "stage2a"], default="stage2a") |
| parser.add_argument( |
| "--curriculum-phase", |
| choices=CURRICULUM_PHASES, |
| default="", |
| help="Unified-model curriculum alias for --stage; kept separate so old run scripts remain reproducible.", |
| ) |
| parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu") |
| parser.add_argument("--precision", choices=["fp32", "bf16", "fp16"], default="bf16") |
| parser.add_argument("--seed", type=int, default=20260521) |
| parser.add_argument("--d-model", type=int, default=896) |
| parser.add_argument("--n-heads", type=int, default=14) |
| parser.add_argument("--n-layers", type=int, default=10) |
| parser.add_argument("--dropout", type=float, default=0.05) |
| parser.add_argument("--n-missing", type=int, default=5) |
| parser.add_argument("--batch-size", type=int, default=128) |
| parser.add_argument("--min-batch-size", type=int, default=32) |
| parser.add_argument("--auto-batch", action="store_true") |
| parser.add_argument("--eval-batch-size", type=int, default=128) |
| parser.add_argument("--grad-accum-steps", type=int, default=2) |
| parser.add_argument("--max-steps", type=int, default=1000) |
| parser.add_argument("--warmup-steps", type=int, default=100) |
| parser.add_argument("--learning-rate", type=float, default=1.5e-4) |
| parser.add_argument("--min-learning-rate", type=float, default=1.0e-5) |
| parser.add_argument("--weight-decay", type=float, default=0.05) |
| parser.add_argument("--beta1", type=float, default=0.9) |
| parser.add_argument("--beta2", type=float, default=0.95) |
| parser.add_argument("--grad-clip", type=float, default=1.0) |
| parser.add_argument("--log-every", type=int, default=25) |
| parser.add_argument("--eval-every", type=int, default=500) |
| parser.add_argument("--eval-batches", type=int, default=5) |
| parser.add_argument("--save-every", type=int, default=1000) |
| parser.add_argument("--hub-upload-every", type=int, default=1000) |
| parser.add_argument("--field-loss-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["field_loss_weight"]) |
| parser.add_argument("--numeric-loss-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["numeric_loss_weight"]) |
| parser.add_argument("--ordinal-loss-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["ordinal_loss_weight"]) |
| parser.add_argument("--missingness-loss-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["missingness_loss_weight"]) |
| parser.add_argument("--pre-event-risk-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["pre_event_risk_weight"]) |
| parser.add_argument("--event-risk-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["event_risk_weight"]) |
| parser.add_argument("--post-event-risk-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["post_event_risk_weight"]) |
| parser.add_argument("--event-generation-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["event_generation_weight"]) |
| parser.add_argument("--ontology-event-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["ontology_event_weight"]) |
| parser.add_argument("--history-aux-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["history_aux_weight"]) |
| parser.add_argument("--contact-aux-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["contact_aux_weight"]) |
| parser.add_argument("--color-loss-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["color_loss_weight"]) |
| parser.add_argument("--pwe-pre-nll-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["pwe_pre_nll_weight"]) |
| parser.add_argument("--pwe-contact-nll-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["pwe_contact_nll_weight"]) |
| parser.add_argument("--pwe-post-nll-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["pwe_post_nll_weight"]) |
| parser.add_argument("--pwe-pre-distill-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["pwe_pre_distill_weight"]) |
| parser.add_argument("--pwe-contact-distill-weight", type=float, default=LOSS_WEIGHTS_V4P4_WORLD_MODEL_20260523["pwe_contact_distill_weight"]) |
| parser.add_argument("--rollout-steps", type=int, default=ROLLOUT_SCHEDULE_DEFAULTS_V4P4_WORLD_MODEL_20260523["rollout_steps"]) |
| parser.add_argument("--rollout-weight-1", type=float, default=ROLLOUT_WEIGHTS_V4P4_WORLD_MODEL_20260523["rollout_weight_1"]) |
| parser.add_argument("--rollout-weight-2", type=float, default=ROLLOUT_WEIGHTS_V4P4_WORLD_MODEL_20260523["rollout_weight_2"]) |
| parser.add_argument("--rollout-weight-3", type=float, default=ROLLOUT_WEIGHTS_V4P4_WORLD_MODEL_20260523["rollout_weight_3"]) |
| parser.add_argument("--teacher-forcing-rate", type=float, default=ROLLOUT_SCHEDULE_DEFAULTS_V4P4_WORLD_MODEL_20260523["teacher_forcing_rate"]) |
| parser.add_argument("--teacher-forcing-start", type=float, default=ROLLOUT_SCHEDULE_DEFAULTS_V4P4_WORLD_MODEL_20260523["teacher_forcing_start"]) |
| parser.add_argument("--teacher-forcing-end", type=float, default=ROLLOUT_SCHEDULE_DEFAULTS_V4P4_WORLD_MODEL_20260523["teacher_forcing_end"]) |
| parser.add_argument("--gumbel-temperature-start", type=float, default=ROLLOUT_SCHEDULE_DEFAULTS_V4P4_WORLD_MODEL_20260523["gumbel_temperature_start"]) |
| parser.add_argument("--gumbel-temperature-end", type=float, default=ROLLOUT_SCHEDULE_DEFAULTS_V4P4_WORLD_MODEL_20260523["gumbel_temperature_end"]) |
| parser.add_argument("--compile", action="store_true") |
| parser.add_argument("--sample-with-replacement", action="store_true") |
| parser.add_argument("--finite-check-every", type=int, default=25) |
| parser.add_argument("--no-freeze-backbone-stage1", dest="freeze_backbone_stage1", action="store_false") |
| parser.set_defaults(freeze_backbone_stage1=True) |
| parser.add_argument("--fused-adamw", action="store_true") |
| parser.add_argument("--matmul-precision", type=str, default="high") |
| parser.add_argument("--init-checkpoint", type=Path, default=None) |
| parser.add_argument("--resume-checkpoint", type=Path, default=None) |
| parser.add_argument("--keep-step-checkpoints", action="store_true") |
| parser.add_argument("--early-stopping-patience", type=int, default=0) |
| parser.add_argument("--early-stopping-metric", type=str, default="val_loss_total") |
| parser.add_argument("--early-stopping-mode", choices=["min", "max"], default="min") |
| parser.add_argument("--early-stopping-min-delta", type=float, default=0.0) |
| parser.add_argument("--early-stopping-start-step", type=int, default=0) |
| parser.add_argument("--hub-repo-id", type=str, default="") |
| parser.add_argument("--hub-repo-type", type=str, default="model") |
| parser.add_argument("--hub-path", type=str, default="sctm_v4p4_world_model") |
| parsed = parser.parse_args() |
| if parsed.curriculum_phase and "--stage" not in sys.argv[1:]: |
| parsed.stage = CURRICULUM_PHASE_TO_LEGACY_STAGE[parsed.curriculum_phase] |
| return normalize_curriculum_args(parsed) |
|
|
|
|
| if __name__ == "__main__": |
| result = run(parse_args()) |
| print(json.dumps(result, ensure_ascii=False), flush=True) |
|
|
