"""Fit the WattGPU power and ITL models and save them for the demo. Two things are produced: 1. `power.joblib` / `itl.joblib` — the pipelines from the paper, refitted on the full Watt Counts subset. 2. `meta.json` — the sets of profiled LLMs and GPUs (which decide the demo's certainty tier), plus the accuracy each tier can be expected to deliver. The tier accuracies come from the paper's own validation protocols, so the number shown next to a prediction is measured under exactly the conditions that prediction is made in: green (seen LLM, seen GPU) -> 5-fold grouped CV yellow (unseen LLM, seen GPU) -> leave-one-LLM-out orange (seen LLM, unseen GPU) -> leave-one-GPU-out Usage: python scripts/train_models.py """ from __future__ import annotations import argparse import ast import json import os import sys import joblib import numpy as np import pandas as pd from sklearn.compose import ColumnTransformer from sklearn.impute import SimpleImputer from sklearn.model_selection import GroupKFold, LeaveOneGroupOut from sklearn.pipeline import Pipeline from sklearn.preprocessing import OrdinalEncoder, StandardScaler from xgboost import XGBRegressor REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, REPO_ROOT) from wattgpu_demo.features import ( # noqa: E402 ID_COLUMNS, ITL_FEATURES, ITL_TARGET, POWER_FEATURES, POWER_TARGET, build_training_frame, clean_frame, ) def find_paper_data(start: str) -> str | None: """Locate the paper's `data/` directory by walking up from `start`. The Space lives in its own git repository nested inside the research repository, and how deeply is not fixed, so the measurement files are found by their contents rather than by a hard-coded number of parent directories. """ current = os.path.abspath(start) while True: candidate = os.path.join(current, "data") if os.path.exists(os.path.join(candidate, "watt_counts_subset.csv")): return candidate parent = os.path.dirname(current) if parent == current: return None current = parent PAPER_DATA_DIR = find_paper_data(REPO_ROOT) # Hyperparameters as reported in the paper. POWER_REGRESSOR = dict(max_depth=6, reg_lambda=150, n_estimators=200) ITL_REGRESSOR = dict(max_depth=5, reg_lambda=100, n_estimators=100) def build_pipeline(X: pd.DataFrame, regressor) -> Pipeline: """Preprocessing + regressor, identical to the notebook's `_build_pipeline`.""" numeric = X.select_dtypes(include=[np.number]).columns.tolist() categorical = X.select_dtypes(include=["object", "category"]).columns.tolist() preprocessor = ColumnTransformer([ ("num", Pipeline([ ("imputer", SimpleImputer(strategy="mean")), ("scaler", StandardScaler()), ]), numeric), ("cat", Pipeline([ ("imputer", SimpleImputer(strategy="most_frequent")), # Unseen categories (a new `model_type` or `memory_type`) encode to # -1 rather than raising, which is what lets the demo predict for # architectures that were never profiled. ("encoder", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1)), ]), categorical), ]) return Pipeline([("preprocessor", preprocessor), ("regressor", regressor)]) def mdape(y_true: np.ndarray, y_pred: np.ndarray) -> float: """Median absolute percentage error, the paper's headline metric.""" return float(np.median(np.abs((y_true - y_pred) / y_true)) * 100) def cross_validate(df: pd.DataFrame, group_col: str | None, target_col: str, regressor_kwargs: dict, log_transform_y: bool) -> np.ndarray: """Out-of-fold predictions under CV, LOGO or LOLO.""" X = df.drop(columns=[c for c in [target_col, *ID_COLUMNS] if c in df.columns]) y = df[target_col] predictions = np.full(len(df), np.nan) if group_col is not None: splits = LeaveOneGroupOut().split(X, y, df[group_col]) else: # Group on the configuration so replicates of one (LLM, GPU) pair never # straddle the train/test boundary. config_id = df[ID_COLUMNS].astype(str).agg("|".join, axis=1) splits = GroupKFold(n_splits=5).split(X, y, groups=config_id) for train_idx, test_idx in splits: pipeline = build_pipeline(X, XGBRegressor(**regressor_kwargs)) y_train = np.log(y.iloc[train_idx]) if log_transform_y else y.iloc[train_idx] pipeline.fit(X.iloc[train_idx], y_train) y_pred = pipeline.predict(X.iloc[test_idx]) predictions[test_idx] = np.exp(y_pred) if log_transform_y else y_pred return predictions def double_holdout_predictions(df: pd.DataFrame, target_col: str, regressor_kwargs: dict, log_transform_y: bool) -> np.ndarray: """Predictions for pairs whose LLM *and* GPU are both held out. The paper validates generalisation one axis at a time (LOGO and LOLO). This is the natural extension: for every measured (LLM, GPU) pair, train on the data with that GPU and that LLM both removed entirely, then predict the pair. It is the only honest way to attach an error to an estimate where neither side was measured -- without it, such an estimate would carry no validated accuracy at all. """ X = df.drop(columns=[c for c in [target_col, *ID_COLUMNS] if c in df.columns]) y = df[target_col] predictions = np.full(len(df), np.nan) models = df["model"].to_numpy() gpus = df["gpu_db_name"].to_numpy() pairs = df[["model", "gpu_db_name"]].drop_duplicates().itertuples(index=False) for n, (model, gpu) in enumerate(pairs, start=1): test = (models == model) & (gpus == gpu) train = (models != model) & (gpus != gpu) if not train.any() or not test.any(): continue pipeline = build_pipeline(X, XGBRegressor(**regressor_kwargs)) y_train = np.log(y[train]) if log_transform_y else y[train] pipeline.fit(X[train], y_train) y_pred = pipeline.predict(X[test]) predictions[test] = np.exp(y_pred) if log_transform_y else y_pred if n % 40 == 0: print(f" double holdout: {n} pairs") return predictions def tier_accuracy(df: pd.DataFrame, target_col: str, regressor_kwargs: dict, log_transform_y: bool, to_watts: bool) -> dict[str, dict[str, float]]: """MdAPE per certainty tier and scenario, each under its own protocol. Reported separately for offline and server operation because, as in the paper's Tables 2 and 3, the two regimes differ substantially -- especially for ITL, where offline throughput depends on batching effects that the features capture only partly. """ scale = df["thermal_design_power_w"].to_numpy() if to_watts else 1.0 y_true = df[target_col].to_numpy() * scale # The paper reports the two server load levels together. regime = np.where(df["scenario"].to_numpy() == "offline", "offline", "server") results: dict[str, dict[str, float]] = {} tiers = (("green", None), ("yellow", "model"), ("orange", "gpu_db_name"), ("red", "both")) for tier, group_col in tiers: if group_col == "both": y_pred = double_holdout_predictions( df, target_col, regressor_kwargs, log_transform_y) * scale else: y_pred = cross_validate( df, group_col, target_col, regressor_kwargs, log_transform_y) * scale valid = ~np.isnan(y_pred) results[tier] = { scenario: round(mdape(y_true[valid & (regime == scenario)], y_pred[valid & (regime == scenario)]), 1) for scenario in ("offline", "server") } print(f" {tier:<7} MdAPE offline {results[tier]['offline']:5.1f}% " f"server {results[tier]['server']:5.1f}%") return results def fit_final(df: pd.DataFrame, target_col: str, regressor_kwargs: dict, log_transform_y: bool) -> Pipeline: """Refit on every row, which is what the demo serves predictions from.""" X = df.drop(columns=[c for c in [target_col, *ID_COLUMNS] if c in df.columns]) y = np.log(df[target_col]) if log_transform_y else df[target_col] pipeline = build_pipeline(X, XGBRegressor(**regressor_kwargs)) pipeline.fit(X, y) return pipeline def _parse_architectures(value) -> list[str]: """`model_features.csv` stores the architecture list as a Python literal.""" if value is None or (isinstance(value, float) and pd.isna(value)): return [] if isinstance(value, list): return [str(v) for v in value] try: parsed = ast.literal_eval(str(value)) except (ValueError, SyntaxError): return [str(value)] return [str(v) for v in parsed] if isinstance(parsed, (list, tuple)) else [str(parsed)] def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--data-dir", default=PAPER_DATA_DIR, required=PAPER_DATA_DIR is None, help="the paper's data/ directory; found automatically when the " "Space repository sits inside the research repository") parser.add_argument("--out-dir", default=os.path.join(REPO_ROOT, "data", "models")) parser.add_argument("--skip-validation", action="store_true", help="fit only; keep the tier accuracies from a previous run") args = parser.parse_args() os.makedirs(args.out_dir, exist_ok=True) print(f"loading measurements from {args.data_dir}") df_all = build_training_frame(args.data_dir) print(f" {len(df_all)} runs, {df_all.model.nunique()} LLMs, {df_all.gpu_db_name.nunique()} GPUs") meta: dict = {} accuracies: dict = {} # --- power draw --------------------------------------------------------- df_power = clean_frame(df_all, [*ID_COLUMNS, POWER_TARGET, *POWER_FEATURES]) tdp = df_all.set_index("gpu_db_name")["thermal_design_power_w"].drop_duplicates() df_power["thermal_design_power_w"] = df_power["gpu_db_name"].map(tdp) print(f"\npower model: {len(df_power)} rows, {len(POWER_FEATURES)} features") if not args.skip_validation: accuracies["power"] = tier_accuracy( df_power.drop(columns=["thermal_design_power_w"]).assign( thermal_design_power_w=df_power["thermal_design_power_w"]), POWER_TARGET, POWER_REGRESSOR, log_transform_y=False, to_watts=True) power_pipeline = fit_final( df_power.drop(columns=["thermal_design_power_w"]), POWER_TARGET, POWER_REGRESSOR, log_transform_y=False) joblib.dump(power_pipeline, os.path.join(args.out_dir, "power.joblib")) # --- inter-token latency ------------------------------------------------ df_itl = clean_frame(df_all, [*ID_COLUMNS, ITL_TARGET, *ITL_FEATURES]) print(f"\nITL model: {len(df_itl)} rows, {len(ITL_FEATURES)} features") if not args.skip_validation: accuracies["itl"] = tier_accuracy( df_itl, ITL_TARGET, ITL_REGRESSOR, log_transform_y=True, to_watts=False) itl_pipeline = fit_final(df_itl, ITL_TARGET, ITL_REGRESSOR, log_transform_y=True) joblib.dump(itl_pipeline, os.path.join(args.out_dir, "itl.joblib")) # --- metadata ----------------------------------------------------------- meta_path = os.path.join(args.out_dir, "meta.json") if args.skip_validation and os.path.exists(meta_path): with open(meta_path) as fh: accuracies = json.load(fh).get("tier_accuracy_mdape", accuracies) # A model or GPU counts as "seen" only if it survived into a training frame. seen_models = sorted(set(df_power["model"]) | set(df_itl["model"])) seen_gpus = sorted(set(df_power["gpu_db_name"]) | set(df_itl["gpu_db_name"])) meta = { "seen_models": seen_models, "seen_gpus": seen_gpus, "power_features": POWER_FEATURES, "itl_features": ITL_FEATURES, "tier_accuracy_mdape": accuracies, # The ITL model is fitted on log(itl); predictions must be exponentiated. "log_transformed_targets": ["itl"], "n_training_runs": {"power": len(df_power), "itl": len(df_itl)}, "gpu_tdp": {k: float(v) for k, v in tdp.items()}, } with open(meta_path, "w") as fh: json.dump(meta, fh, indent=2) # Cache the architecture of every profiled LLM. These features are already # in the paper's data, so a profiled model needs no Hub round-trip -- which # also makes licence-gated models (Llama, Gemma) work without a token. cache = {} for model_id, group in df_all.groupby("model"): if model_id not in seen_models: continue row = group.iloc[0] cache[model_id] = { "model_type": str(row["model_type"]), "num_layers": int(row["num_layers"]), "hidden_size": int(row["hidden_size"]), "num_attention_heads": int(row["num_attention_heads"]), "num_key_value_heads": int(row["num_key_value_heads"]), "total_b_params": float(row["total_b_params"]), "architectures": _parse_architectures(row.get("architectures")), "max_position_embeddings": ( None if pd.isna(row.get("max_position_embeddings")) else int(row["max_position_embeddings"]) ), "torch_dtype": ( None if pd.isna(row.get("torch_dtype")) or row.get("torch_dtype") == "N/A" else str(row["torch_dtype"]) ), } cache_path = os.path.join(args.out_dir, "profiled_llms.json") with open(cache_path, "w") as fh: json.dump(cache, fh, indent=2, sort_keys=True) print(f" cached architecture for {len(cache)} profiled LLMs") print(f"\nsaved models and metadata to {args.out_dir}") print(f" {len(seen_models)} profiled LLMs, {len(seen_gpus)} profiled GPUs") return 0 if __name__ == "__main__": sys.exit(main())