"""Run shared-weight multi-asset close-distribution experiments.""" from __future__ import annotations import argparse import hashlib import json import math import platform import random import sys import time from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any, Iterable import numpy as np import torch from torch import nn from project.data.close_distribution_v2 import ( CANONICAL_HORIZONS, CLASS_COUNT, DISCOVERY_FOLDS, EVALUATION_HORIZONS, SELECTION_ASSETS, ) from project.evaluators.close_distribution_v2 import ( ClassMap, class_map_to_arrays, encode_returns, fit_class_map, metrics, prior_probabilities, score_arrays, ) from project.experiments.forecasting_v4.runner import ( FEATURE_SETS, _context, _feature_indices, _horizon_indices, _load_npz, _masked_nll, _normalization, _optimizer, _optimizer_settings, _ordinal, _parse_horizons, _session_scores, ) from project.models.close_distribution_v3 import ( ASSET_CONDITIONING_MODES, build_shared_model, ) PROVENANCE_PATHS = ( "project/data/close_distribution_v2.py", "project/evaluators/close_distribution_v2.py", "project/experiments/forecasting_v4/runner.py", "project/experiments/forecasting_v5/runner.py", "project/models/close_distribution_v2.py", "project/models/close_distribution_v3.py", "project/studies/forecasting_v5.md", "requirements.txt", ) @dataclass class AssetFoldData: asset: str asset_id: int features: dict[str, np.ndarray] bundle: dict[str, np.ndarray] center: np.ndarray scale: np.ndarray class_map: ClassMap train_rows: np.ndarray train_labels: np.ndarray train_target_mask: np.ndarray diagnostic_rows: np.ndarray diagnostic_labels: np.ndarray diagnostic_target_mask: np.ndarray validation_rows: np.ndarray validation_y: np.ndarray validation_target_mask: np.ndarray validation_session_date: np.ndarray diagnostic_purge_rows: int def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _configure(seed: int, threads: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.set_num_threads(threads) torch.use_deterministic_algorithms(True) def _load_selection_data( study_dir: Path, ) -> tuple[ dict[str, Any], dict[str, tuple[dict[str, np.ndarray], dict[str, np.ndarray]]], ]: manifest_path = study_dir / "data" / "prepared-manifest.json" manifest = json.loads(manifest_path.read_text()) if manifest.get("scope") != "selection": raise ValueError("v5 requires the selection-only prepared manifest") if tuple(manifest.get("selection_assets", ())) != SELECTION_ASSETS: raise ValueError("selection asset cohort changed") data_dir = study_dir / "data" / "runner" loaded = {} for asset in SELECTION_ASSETS: slug = asset.lower() feature_path = data_dir / f"{slug}_features.npz" label_path = data_dir / f"{slug}_labels.npz" for path in (feature_path, label_path): relative = str(path.relative_to(study_dir)) if _sha256(path) != manifest["derived_sha256"][relative]: raise ValueError(f"prepared hash mismatch: {relative}") loaded[asset] = (_load_npz(feature_path), _load_npz(label_path)) return manifest, loaded def _prepare_asset_fold( *, asset: str, asset_id: int, features: dict[str, np.ndarray], bundle: dict[str, np.ndarray], split: dict[str, tuple[str, str]], supervised_indices: tuple[int, ...], selected_indices: tuple[int, ...], max_train_rows_per_asset: int, diagnostic_tail_fraction: float, ) -> AssetFoldData: dates = bundle["session_date"] train = ( (dates >= _ordinal(split["train"][0])) & (dates < _ordinal(split["train"][1])) ) validation = ( (dates >= _ordinal(split["validation"][0])) & (dates < _ordinal(split["validation"][1])) ) train_rows = bundle["row_index"][train] train_y = bundle["y"][train] train_target_mask = bundle["target_mask"][train] if ( max_train_rows_per_asset > 0 and len(train_rows) > max_train_rows_per_asset ): selected = np.linspace( 0, len(train_rows) - 1, max_train_rows_per_asset, dtype=np.int64, ) train_rows = train_rows[selected] train_y = train_y[selected] train_target_mask = train_target_mask[selected] diagnostic_rows = np.empty(0, dtype=train_rows.dtype) diagnostic_y = np.empty((0, train_y.shape[1]), dtype=train_y.dtype) diagnostic_target_mask = np.empty( (0, train_target_mask.shape[1]), dtype=np.bool_, ) diagnostic_purge_rows = 0 normalization_max_row = None if diagnostic_tail_fraction: diagnostic_count = max( 1, round(len(train_rows) * diagnostic_tail_fraction), ) diagnostic_rows = train_rows[-diagnostic_count:] diagnostic_y = train_y[-diagnostic_count:] diagnostic_target_mask = train_target_mask[-diagnostic_count:] train_rows = train_rows[:-diagnostic_count] train_y = train_y[:-diagnostic_count] train_target_mask = train_target_mask[:-diagnostic_count] normalization_max_row = int(diagnostic_rows[0]) unpurged_count = len(train_rows) keep = train_rows + int(EVALUATION_HORIZONS.max()) < int( diagnostic_rows[0] ) train_rows = train_rows[keep] train_y = train_y[keep] train_target_mask = train_target_mask[keep] diagnostic_purge_rows = unpurged_count - len(train_rows) class_map = fit_class_map( train_y, train_target_mask, class_count=CLASS_COUNT, ) center, scale = _normalization( features, train_end=split["train"][1], selected_indices=selected_indices, max_row_exclusive=normalization_max_row, ) all_train_labels = encode_returns( train_y, train_target_mask, class_map.edges, ) all_diagnostic_labels = encode_returns( diagnostic_y, diagnostic_target_mask, class_map.edges, ) return AssetFoldData( asset=asset, asset_id=asset_id, features=features, bundle=bundle, center=center, scale=scale, class_map=class_map, train_rows=train_rows, train_labels=all_train_labels[:, supervised_indices], train_target_mask=train_target_mask[:, supervised_indices], diagnostic_rows=diagnostic_rows, diagnostic_labels=all_diagnostic_labels[:, supervised_indices], diagnostic_target_mask=diagnostic_target_mask[:, supervised_indices], validation_rows=bundle["row_index"][validation], validation_y=bundle["y"][validation], validation_target_mask=bundle["target_mask"][validation], validation_session_date=bundle["session_date"][validation], diagnostic_purge_rows=diagnostic_purge_rows, ) def _balanced_epoch_chunks( row_counts: tuple[int, ...], *, effective_batch_size: int, generator: np.random.Generator, ) -> list[tuple[np.ndarray, ...]]: if not row_counts or any(value <= 0 for value in row_counts): raise ValueError("every asset must have positive training rows") rows_per_asset_update = effective_batch_size // len(row_counts) if rows_per_asset_update <= 0: raise ValueError("effective batch is smaller than the asset cohort") updates = math.ceil(max(row_counts) / rows_per_asset_update) per_asset = [ np.array_split(generator.permutation(count), updates) for count in row_counts ] if any(any(not len(chunk) for chunk in chunks) for chunks in per_asset): raise ValueError("asset-balanced update contains an empty asset slice") return [ tuple(chunks[update] for chunks in per_asset) for update in range(updates) ] def _asset_balanced_loss_part( logits: torch.Tensor, labels: torch.Tensor, mask: torch.Tensor, asset_id: torch.Tensor, denominators: torch.Tensor, ) -> torch.Tensor: if denominators.ndim != 2: raise ValueError("asset/horizon denominators must be two-dimensional") losses = nn.functional.cross_entropy( logits.reshape(-1, logits.shape[-1]), labels.clamp_min(0).reshape(-1), reduction="none", ).reshape(labels.shape) result = losses.new_zeros(()) for index in range(denominators.shape[0]): selected = asset_id == index horizon_sum = (losses[selected] * mask[selected]).sum(dim=0) result = result + torch.sum(horizon_sum / denominators[index]) return result / denominators.numel() def _training_microbatch( data: tuple[AssetFoldData, ...], asset_ids: np.ndarray, positions: np.ndarray, *, context_length: int, selected_indices: tuple[int, ...], ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, ]: tokens = [] padding = [] labels = [] masks = [] ids = [] for asset_data in data: selected = asset_ids == asset_data.asset_id if not np.any(selected): continue local = positions[selected] context, context_padding = _context( asset_data.features["X"], asset_data.train_rows[local], context_length=context_length, selected_indices=selected_indices, center=asset_data.center, scale=asset_data.scale, ) tokens.append(context) padding.append(context_padding) labels.append( torch.from_numpy( asset_data.train_labels[local].astype( np.int64, copy=False, ) ) ) masks.append( torch.from_numpy( asset_data.train_target_mask[local].astype( np.float32, copy=False, ) ) ) ids.append( torch.full( (len(local),), asset_data.asset_id, dtype=torch.long, ) ) return ( torch.cat(tokens), torch.cat(padding), torch.cat(labels), torch.cat(masks), torch.cat(ids), ) def _predict_asset( model: nn.Module, asset_data: AssetFoldData, *, rows: np.ndarray, context_length: int, selected_indices: tuple[int, ...], batch_size: int, device: torch.device, query_horizons: tuple[int, ...] | None, ) -> np.ndarray: horizon_count = ( len(query_horizons) if query_horizons is not None else len(model.output_horizons) ) result = np.empty( (len(rows), horizon_count, CLASS_COUNT), dtype=np.float64, ) query = ( torch.tensor(query_horizons, dtype=torch.float32, device=device) if query_horizons is not None else None ) model.eval() with torch.inference_mode(): for start in range(0, len(rows), batch_size): stop = min(start + batch_size, len(rows)) tokens, padding = _context( asset_data.features["X"], rows[start:stop], context_length=context_length, selected_indices=selected_indices, center=asset_data.center, scale=asset_data.scale, ) asset_id = torch.full( (stop - start,), asset_data.asset_id, dtype=torch.long, device=device, ) logits = model( tokens.to(device), padding.to(device), asset_id, query, ) result[start:stop] = torch.softmax( logits.cpu().to(torch.float64), dim=2, ).numpy() return result def _train_model( model: nn.Module, data: tuple[AssetFoldData, ...], *, context_length: int, selected_indices: tuple[int, ...], seed: int, epochs: int, batch_size: int, gradient_accumulation_steps: int, optimizer_recipe: str, inference_batch_size: int, device: torch.device, ) -> dict[str, Any]: optimizer = _optimizer(model, optimizer_recipe) effective_batch_size = batch_size * gradient_accumulation_steps generator = np.random.default_rng(seed + 10_007) histories: list[float] = [] diagnostic_history: list[float] = [] best_epoch = 0 best_diagnostic_nll = float("inf") best_state: dict[str, torch.Tensor] | None = None update_count = 0 expected_updates_per_epoch = None for epoch in range(epochs): model.train() updates = _balanced_epoch_chunks( tuple(len(item.train_rows) for item in data), effective_batch_size=effective_batch_size, generator=generator, ) if expected_updates_per_epoch is None: expected_updates_per_epoch = len(updates) elif len(updates) != expected_updates_per_epoch: raise AssertionError("updates per epoch changed") epoch_loss_sum = np.zeros( (len(data), data[0].train_labels.shape[1]), dtype=np.float64, ) epoch_target_count = np.zeros_like(epoch_loss_sum, dtype=np.int64) for chunks in updates: update_asset_ids = np.concatenate( [ np.full(len(chunk), index, dtype=np.int64) for index, chunk in enumerate(chunks) ] ) update_positions = np.concatenate(chunks) denominators = np.stack( [ asset_data.train_target_mask[chunk].sum(axis=0) for asset_data, chunk in zip(data, chunks, strict=True) ] ).astype(np.float32) if np.any(denominators == 0): raise ValueError( "an asset-balanced update has no target for a horizon" ) denominator_tensor = torch.from_numpy(denominators).to(device) optimizer.zero_grad(set_to_none=True) for start in range(0, len(update_positions), batch_size): stop = min(start + batch_size, len(update_positions)) ( tokens, padding, labels, mask, asset_id, ) = _training_microbatch( data, update_asset_ids[start:stop], update_positions[start:stop], context_length=context_length, selected_indices=selected_indices, ) tokens = tokens.to(device) padding = padding.to(device) labels = labels.to(device) mask = mask.to(device) asset_id = asset_id.to(device) logits = model(tokens, padding, asset_id) loss = _asset_balanced_loss_part( logits, labels, mask, asset_id, denominator_tensor, ) loss.backward() with torch.no_grad(): item_losses = nn.functional.cross_entropy( logits.reshape(-1, logits.shape[-1]), labels.clamp_min(0).reshape(-1), reduction="none", ).reshape(labels.shape) for index in range(len(data)): selected = asset_id == index epoch_loss_sum[index] += ( item_losses[selected] * mask[selected] ).sum(dim=0).cpu().numpy() epoch_target_count[index] += ( mask[selected] .sum(dim=0) .cpu() .numpy() .astype(np.int64) ) nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() update_count += 1 histories.append( float(np.mean(epoch_loss_sum / epoch_target_count)) ) if all(len(item.diagnostic_rows) for item in data): diagnostic_nll = [] for asset_data in data: probabilities = _predict_asset( model, asset_data, rows=asset_data.diagnostic_rows, context_length=context_length, selected_indices=selected_indices, batch_size=inference_batch_size, device=device, query_horizons=None, ) diagnostic_nll.append( _masked_nll( probabilities, asset_data.diagnostic_labels, asset_data.diagnostic_target_mask, ) ) equal_asset_nll = float(np.mean(diagnostic_nll)) diagnostic_history.append(equal_asset_nll) if equal_asset_nll < best_diagnostic_nll: best_diagnostic_nll = equal_asset_nll best_epoch = epoch + 1 best_state = { name: value.detach().cpu().clone() for name, value in model.state_dict().items() } if best_state is not None: model.load_state_dict(best_state) else: best_epoch = epochs if expected_updates_per_epoch is None: raise AssertionError("training produced no updates") return { "best_epoch": best_epoch, "completed_fixed_update_budget": ( update_count == expected_updates_per_epoch * epochs ), "diagnostic_tail_nll": diagnostic_history, "train_loss": histories, "updates": update_count, "updates_per_epoch": expected_updates_per_epoch, } def _run_manifest(config: dict[str, Any]) -> dict[str, Any]: project_root = Path(__file__).resolve().parents[3] return { "command": sys.argv, "config": config, "created_at_utc": datetime.now(timezone.utc).isoformat(), "environment": { "machine": platform.machine(), "numpy": np.__version__, "platform": platform.platform(), "python": sys.version, "torch": torch.__version__, "torch_mps_available": torch.backends.mps.is_available(), }, "provenance_sha256": { relative: _sha256(project_root / relative) for relative in PROVENANCE_PATHS }, } def run( *, study_dir: Path, output_dir: Path, asset_conditioning: str, supervised_horizons: str | Iterable[int] = "2,4,8,16,32", context_length: int = 128, seed: int = 0, epochs: int = 16, batch_size: int = 256, inference_batch_size: int = 512, target_parameters: int = 25_000, torch_threads: int = 6, max_train_rows_per_asset: int = 100_000, device_name: str = "mps", gradient_accumulation_steps: int = 16, optimizer_recipe: str = "adamw_constant", rotary_base: float = 16.0, feature_set: str = "session_time", diagnostic_tail_fraction: float = 0.1, ) -> dict[str, Any]: started = time.monotonic() horizons = _parse_horizons(supervised_horizons) if asset_conditioning not in ASSET_CONDITIONING_MODES: raise ValueError("unknown asset conditioning mode") if feature_set not in FEATURE_SETS: raise ValueError("unknown feature set") if optimizer_recipe != "adamw_constant": raise ValueError( "v5 currently implements only the frozen AdamW constant recipe" ) if horizons[0] != 2 or horizons[-1] != 32: raise ValueError("supervised horizons must include 2 and 32") if epochs <= 0 or batch_size <= 0 or context_length <= 0: raise ValueError("training dimensions must be positive") if gradient_accumulation_steps <= 0: raise ValueError("gradient accumulation must be positive") if not 0.0 <= diagnostic_tail_fraction < 0.5: raise ValueError("diagnostic fraction must lie in [0, 0.5)") if device_name == "mps" and not torch.backends.mps.is_available(): raise ValueError("MPS was requested but is unavailable") manifest, loaded = _load_selection_data(study_dir) selected_indices = _feature_indices(feature_set) supervised_indices = _horizon_indices(horizons) canonical_indices = _horizon_indices(CANONICAL_HORIZONS) dense_unseen_horizons = tuple( int(value) for value in EVALUATION_HORIZONS if value not in set(CANONICAL_HORIZONS.tolist()) ) dense_unseen_indices = _horizon_indices(dense_unseen_horizons) learning_rate, weight_decay = _optimizer_settings(optimizer_recipe) config = { "asset_conditioning": asset_conditioning, "assets": list(SELECTION_ASSETS), "batch_size": batch_size, "context_length": context_length, "device_name": device_name, "diagnostic_tail_fraction": diagnostic_tail_fraction, "epochs": epochs, "feature_names": list(FEATURE_SETS[feature_set]), "feature_set": feature_set, "gradient_accumulation_steps": gradient_accumulation_steps, "inference_batch_size": inference_batch_size, "learning_rate": learning_rate, "max_train_rows_per_asset": max_train_rows_per_asset, "model_name": "shared_transformer_rotary", "optimizer_recipe": optimizer_recipe, "rotary_base": rotary_base, "seed": seed, "snapshot_sha256": manifest["snapshot_sha256"], "supervised_horizons": list(horizons), "target_parameters": target_parameters, "torch_threads": torch_threads, "weight_decay": weight_decay, } output_dir.mkdir(parents=True, exist_ok=True) results_path = output_dir / "results.json" run_manifest_path = output_dir / "run_manifest.json" if results_path.exists() or run_manifest_path.exists(): raise FileExistsError("experiment output already exists") run_manifest_path.write_text( json.dumps(_run_manifest(config), indent=2, sort_keys=True) + "\n" ) run_manifest_sha256 = _sha256(run_manifest_path) _configure(seed, torch_threads) device = torch.device(device_name) fold_results: dict[str, Any] = {} for fold_name, split in DISCOVERY_FOLDS.items(): fold_started = time.monotonic() data = tuple( _prepare_asset_fold( asset=asset, asset_id=asset_id, features=loaded[asset][0], bundle=loaded[asset][1], split=split, supervised_indices=supervised_indices, selected_indices=selected_indices, max_train_rows_per_asset=max_train_rows_per_asset, diagnostic_tail_fraction=diagnostic_tail_fraction, ) for asset_id, asset in enumerate(SELECTION_ASSETS) ) model = build_shared_model( context_length=context_length, channels=len(selected_indices) + 1, output_horizons=horizons, classes=CLASS_COUNT, target_parameters=target_parameters, rotary_base=rotary_base, asset_count=len(SELECTION_ASSETS), asset_conditioning=asset_conditioning, ).to(device) parameter_count = sum( parameter.numel() for parameter in model.parameters() if parameter.requires_grad ) core_parameter_count = parameter_count - sum( parameter.numel() for name, parameter in model.named_parameters() if name.startswith(("asset_embedding", "asset_affine")) ) training = _train_model( model, data, context_length=context_length, selected_indices=selected_indices, seed=seed, epochs=epochs, batch_size=batch_size, gradient_accumulation_steps=gradient_accumulation_steps, optimizer_recipe=optimizer_recipe, inference_batch_size=inference_batch_size, device=device, ) asset_results = {} for asset_data in data: probabilities = _predict_asset( model, asset_data, rows=asset_data.validation_rows, context_length=context_length, selected_indices=selected_indices, batch_size=inference_batch_size, device=device, query_horizons=tuple( int(value) for value in EVALUATION_HORIZONS ), ) scores = score_arrays( asset_data.validation_y, asset_data.validation_target_mask, probabilities, class_map=asset_data.class_map, ) prior = prior_probabilities( len(asset_data.validation_y), asset_data.class_map, ) prior_scores = score_arrays( asset_data.validation_y, asset_data.validation_target_mask, prior, class_map=asset_data.class_map, ) metric_slices = { "canonical": metrics( asset_data.validation_y, asset_data.validation_target_mask, probabilities, class_map=asset_data.class_map, horizon_indices=canonical_indices, score_values=scores, ), "dense_unseen": metrics( asset_data.validation_y, asset_data.validation_target_mask, probabilities, class_map=asset_data.class_map, horizon_indices=dense_unseen_indices, score_values=scores, ), "prior_canonical": metrics( asset_data.validation_y, asset_data.validation_target_mask, prior, class_map=asset_data.class_map, horizon_indices=canonical_indices, score_values=prior_scores, ), } np.savez_compressed( output_dir / f"{fold_name}_{asset_data.asset.lower()}_session_scores.npz", **_session_scores( asset_data.validation_y, asset_data.validation_target_mask, probabilities, asset_data.validation_session_date, class_map=asset_data.class_map, score_values=scores, ), ) asset_results[asset_data.asset] = { "class_map_rows": int(len(asset_data.train_rows)), "diagnostic_purge_rows": ( asset_data.diagnostic_purge_rows ), "diagnostic_tail_rows": int( len(asset_data.diagnostic_rows) ), "metric_slices": metric_slices, "train_rows": int(len(asset_data.train_rows)), "train_valid_targets_by_supervised_horizon": { str(horizon): int(count) for horizon, count in zip( horizons, asset_data.train_target_mask.sum(axis=0), strict=True, ) }, "validation_rows": int(len(asset_data.validation_rows)), } torch.save( { "asset_calibration": { item.asset: { "center": item.center, "scale": item.scale, **class_map_to_arrays(item.class_map), } for item in data }, "asset_conditioning": asset_conditioning, "assets": list(SELECTION_ASSETS), "context_length": context_length, "feature_names": FEATURE_SETS[feature_set], "rotary_base": rotary_base, "state_dict": model.state_dict(), "supervised_horizons": horizons, }, output_dir / f"{fold_name}_model.pt", ) fold_results[fold_name] = { **training, "assets": asset_results, "checkpoint_selection": ( "lowest_equal_asset_inner_tail_nll" if diagnostic_tail_fraction else "final_epoch" ), "conditioning_parameter_count": ( parameter_count - core_parameter_count ), "core_parameter_count": core_parameter_count, "duration_seconds": time.monotonic() - fold_started, "estimated_madds_per_example": int(model.estimated_madds), "parameter_count": parameter_count, } del model if device.type == "mps": torch.mps.empty_cache() by_asset = {} for asset in SELECTION_ASSETS: canonical = float( np.mean( [ fold_results[fold]["assets"][asset]["metric_slices"][ "canonical" ]["macro_nll"] for fold in DISCOVERY_FOLDS ] ) ) dense = float( np.mean( [ fold_results[fold]["assets"][asset]["metric_slices"][ "dense_unseen" ]["macro_nll"] for fold in DISCOVERY_FOLDS ] ) ) by_asset[asset] = { "mean_canonical_macro_nll": canonical, "mean_dense_unseen_macro_nll": dense, } summary = { **config, "by_asset": by_asset, "duration_seconds": time.monotonic() - started, "equal_asset_mean_canonical_macro_nll": float( np.mean( [ value["mean_canonical_macro_nll"] for value in by_asset.values() ] ) ), "equal_asset_mean_dense_unseen_macro_nll": float( np.mean( [ value["mean_dense_unseen_macro_nll"] for value in by_asset.values() ] ) ), "folds": fold_results, "run_manifest_sha256": run_manifest_sha256, } results_path.write_text( json.dumps(summary, indent=2, sort_keys=True) + "\n" ) return summary def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--study-dir", type=Path, required=True) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument( "--asset-conditioning", choices=ASSET_CONDITIONING_MODES, required=True, ) parser.add_argument( "--supervised-horizons", default="2,4,8,16,32", ) parser.add_argument("--context-length", type=int, default=128) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--epochs", type=int, default=16) parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--inference-batch-size", type=int, default=512) parser.add_argument("--target-parameters", type=int, default=25_000) parser.add_argument("--torch-threads", type=int, default=6) parser.add_argument( "--max-train-rows-per-asset", type=int, default=100_000, ) parser.add_argument( "--device-name", choices=("cpu", "mps"), default="mps", ) parser.add_argument( "--gradient-accumulation-steps", type=int, default=16, ) parser.add_argument( "--optimizer-recipe", choices=("adamw_constant",), default="adamw_constant", ) parser.add_argument("--rotary-base", type=float, default=16.0) parser.add_argument( "--feature-set", choices=tuple(FEATURE_SETS), default="session_time", ) parser.add_argument( "--diagnostic-tail-fraction", type=float, default=0.1, ) args = parser.parse_args() print(json.dumps(run(**vars(args)), sort_keys=True)) if __name__ == "__main__": main()