From a752260d96211e22823ca3bb1cb7a2f769dd84bf Mon Sep 17 00:00:00 2001 From: Kauna <16511995+klei22@users.noreply.github.com> Date: Sat, 28 Mar 2026 20:19:56 -0700 Subject: [PATCH] Add graceful Zeus energy profiling metrics across training/search tooling --- hyperparam_search.py | 40 ++++++++- optimization_and_search/run_experiments.py | 12 ++- run_exploration_monitor.py | 2 + train.py | 95 ++++++++++++++++++++++ train_args.py | 12 +++ 5 files changed, 157 insertions(+), 4 deletions(-) diff --git a/hyperparam_search.py b/hyperparam_search.py index eeeb718b48..2653811179 100644 --- a/hyperparam_search.py +++ b/hyperparam_search.py @@ -32,7 +32,7 @@ import yaml -TrialMetrics = Tuple[float, float, int, float, float, float, float, float, float] +TrialMetrics = Tuple[float, float, int, float, float, float, float, float, float, float, float] # ───────────────────────── helpers ────────────────────────── @@ -94,6 +94,8 @@ def run_trial_inproc(cfg: Dict[str, Any]) -> TrialMetrics: iter_latency_ms, rankme, areq, + zeus_total_energy_j, + zeus_avg_power_w, ) """ from train import Trainer @@ -119,6 +121,8 @@ def run_trial_inproc(cfg: Dict[str, Any]) -> TrialMetrics: iter_latency_ms = float(getattr(tr, "iter_latency_avg", 0.0)) rankme = float(getattr(tr, "latest_rankme", float("nan"))) areq = float(getattr(tr, "latest_areq", float("nan"))) + zeus_total_energy_j = float(getattr(tr, "zeus_total_energy_j", float("nan"))) + zeus_avg_power_w = float(getattr(tr, "zeus_avg_power_w", float("nan"))) del tr _cleanup_cuda() @@ -132,6 +136,8 @@ def run_trial_inproc(cfg: Dict[str, Any]) -> TrialMetrics: iter_latency_ms, rankme, areq, + zeus_total_energy_j, + zeus_avg_power_w, ) @@ -148,6 +154,8 @@ def _parse_best_metrics_file(metrics_path: Path) -> TrialMetrics: iter_latency_ms = float(line[9]) rankme = float(line[19]) areq = float(line[20]) + zeus_total_energy_j = float(line[21]) if len(line) > 21 else float("nan") + zeus_avg_power_w = float(line[22]) if len(line) > 22 else float("nan") return ( loss, @@ -159,6 +167,8 @@ def _parse_best_metrics_file(metrics_path: Path) -> TrialMetrics: iter_latency_ms, rankme, areq, + zeus_total_energy_j, + zeus_avg_power_w, ) @@ -244,12 +254,12 @@ def main(): ) ap.add_argument( "--efficiency_target", - choices=["params", "vram", "iter", "torch_allocated", "torch_reserved", "process_gpu"], + choices=["params", "vram", "iter", "torch_allocated", "torch_reserved", "process_gpu", "zeus_energy", "zeus_power"], default="params", help=( "Metric to normalize score gain: 'params' (default) for parameter count, " "'vram' (legacy alias for torch_allocated), 'torch_allocated', " - "'torch_reserved', 'process_gpu', or 'iter' for average iteration latency in ms." + "'torch_reserved', 'process_gpu', 'zeus_energy' (total Joules), 'zeus_power' (average Watts), or 'iter' for average iteration latency in ms." ), ) ap.add_argument( @@ -343,6 +353,8 @@ def _extend_layerlists(cfg: Dict[str, Any], dup_idx: int) -> None: base_torch_reserved = last["baseline_metrics"].get("peak_torch_reserved_mb", 0.0) base_process_gpu = last["baseline_metrics"].get("peak_process_gpu_mb", 0.0) base_iter_ms = last["baseline_metrics"].get("iter_latency_avg", 0.0) + base_zeus_energy = last["baseline_metrics"].get("zeus_total_energy_j", float("nan")) + base_zeus_power = last["baseline_metrics"].get("zeus_avg_power_w", float("nan")) cur_iter = last["iter"] + 1 _apply_overrides_to_active_config( baseline_cfg, args.override_cfg, "resumed baseline_cfg" @@ -364,6 +376,8 @@ def _extend_layerlists(cfg: Dict[str, Any], dup_idx: int) -> None: base_iter_ms, base_rankme, base_areq, + base_zeus_energy, + base_zeus_power, ) = run_fn(deepcopy(baseline_cfg)) base_score = 1 / math.exp(base_loss) log["iterations"].append( @@ -380,6 +394,8 @@ def _extend_layerlists(cfg: Dict[str, Any], dup_idx: int) -> None: "best_iter": base_best_iter, "rankme": base_rankme, "areq": base_areq, + "zeus_total_energy_j": base_zeus_energy, + "zeus_avg_power_w": base_zeus_power, }, "baseline_config_after": deepcopy(baseline_cfg), } @@ -430,6 +446,8 @@ def _evaluate( iter_ms, rankme, areq, + zeus_total_energy_j, + zeus_avg_power_w, ) = run_fn(cfg_run) except Exception as exc: print(" ⚠", exc) @@ -448,6 +466,8 @@ def _evaluate( "iter_latency_ms": iter_ms, "rankme": rankme, "areq": areq, + "zeus_total_energy_j": zeus_total_energy_j, + "zeus_avg_power_w": zeus_avg_power_w, } ) scores.append(score) @@ -465,6 +485,8 @@ def _evaluate( avg_iter = sum(s["iter_latency_ms"] for s in seed_runs) / len(seed_runs) avg_rankme = _nanmean([s["rankme"] for s in seed_runs]) avg_areq = _nanmean([s["areq"] for s in seed_runs]) + avg_zeus_energy = _nanmean([s["zeus_total_energy_j"] for s in seed_runs]) + avg_zeus_power = _nanmean([s["zeus_avg_power_w"] for s in seed_runs]) avg_loss = -math.log(avg_score) d_score = avg_score - base_score @@ -483,6 +505,8 @@ def _evaluate( d_torch_reserved = avg_torch_reserved - base_torch_reserved d_process_gpu = avg_process_gpu - base_process_gpu d_iter = avg_iter - base_iter_ms + d_zeus_energy = avg_zeus_energy - base_zeus_energy + d_zeus_power = avg_zeus_power - base_zeus_power if args.efficiency_target == "params": d_cost = d_param @@ -494,6 +518,10 @@ def _evaluate( d_cost = d_process_gpu elif args.efficiency_target == "iter": d_cost = d_iter + elif args.efficiency_target == "zeus_energy": + d_cost = d_zeus_energy + elif args.efficiency_target == "zeus_power": + d_cost = d_zeus_power else: raise ValueError("Unknown efficiency target") @@ -536,6 +564,8 @@ def _evaluate( "peak_torch_reserved_mb": avg_torch_reserved, "peak_process_gpu_mb": avg_process_gpu, "iter_latency_avg": avg_iter, + "zeus_total_energy_j": avg_zeus_energy, + "zeus_avg_power_w": avg_zeus_power, "delta_score": d_score, "delta_rankme": d_rankme, "delta_areq": d_areq, @@ -544,6 +574,8 @@ def _evaluate( "delta_torch_reserved_mb": d_torch_reserved, "delta_process_gpu_mb": d_process_gpu, "delta_iter_latency": d_iter, + "delta_zeus_total_energy_j": d_zeus_energy, + "delta_zeus_avg_power_w": d_zeus_power, "efficiency": eff, "target_metric": args.optimize_target, "target_mode": args.optimize_mode, @@ -725,6 +757,8 @@ def _nlayer_candidate(dup_idx: int, tag: str) -> None: "best_iter": chosen["best_iter"], "rankme": base_rankme, "areq": base_areq, + "zeus_total_energy_j": base_zeus_energy, + "zeus_avg_power_w": base_zeus_power, }, "candidates": candidates, "chosen": chosen, diff --git a/optimization_and_search/run_experiments.py b/optimization_and_search/run_experiments.py index 27202aa76b..fe44cba330 100644 --- a/optimization_and_search/run_experiments.py +++ b/optimization_and_search/run_experiments.py @@ -38,6 +38,8 @@ "ln_f_cosine_95", "rankme", "areq", + "zeus_total_energy_j", + "zeus_avg_power_w", ] @@ -574,9 +576,17 @@ def read_metrics(out_dir: str) -> dict: float, float, float, + float, + float, ] - return {k: typ(v) for k, typ, v in zip(METRIC_KEYS, casts, parts)} + parsed = {} + for idx, (key, typ) in enumerate(zip(METRIC_KEYS, casts)): + if idx >= len(parts): + parsed[key] = float("nan") + continue + parsed[key] = typ(parts[idx]) + return parsed def completed_runs(log_file: Path) -> set[str]: diff --git a/run_exploration_monitor.py b/run_exploration_monitor.py index 441c05a7ab..fa7d0ebe6b 100644 --- a/run_exploration_monitor.py +++ b/run_exploration_monitor.py @@ -202,6 +202,8 @@ def on_mount(self) -> None: "ln_f_cosine_95", "rankme", "areq", + "zeus_total_energy_j", + "zeus_avg_power_w", ] + self.param_keys self.all_columns = base_cols.copy() self.columns = base_cols.copy() diff --git a/train.py b/train.py index 22db1ad8c3..1e48b61434 100644 --- a/train.py +++ b/train.py @@ -11,6 +11,68 @@ import time from collections import deque from datetime import datetime, timedelta +from pathlib import Path + + +def _safe_float(value, default=float("nan")): + try: + return float(value) + except (TypeError, ValueError): + return default + + +class ZeusProfiler: + """Best-effort Zeus energy profiler wrapper. + + If Zeus is unavailable or errors, training proceeds without profiling. + """ + + def __init__(self, enabled: bool, window_name: str = "train_total"): + self.enabled = bool(enabled) + self.window_name = window_name + self.monitor = None + self.active = False + self.total_energy_j = float("nan") + self.avg_power_w = float("nan") + self.error: str | None = None + + if not self.enabled: + return + + try: + from zeus.monitor import ZeusMonitor # type: ignore + self.monitor = ZeusMonitor() + except Exception as exc: + self.error = f"Zeus init failed: {exc}" + self.enabled = False + + def start(self): + if not self.enabled or self.monitor is None or self.active: + return + try: + self.monitor.begin_window(self.window_name) + self.active = True + except Exception as exc: + self.error = f"Zeus begin_window failed: {exc}" + self.enabled = False + self.active = False + + def stop(self): + if not self.enabled or self.monitor is None or not self.active: + return + try: + measurement = self.monitor.end_window(self.window_name) + self.total_energy_j = _safe_float(getattr(measurement, "total_energy", float("nan"))) + start_ts = _safe_float(getattr(measurement, "start_time", float("nan"))) + end_ts = _safe_float(getattr(measurement, "end_time", float("nan"))) + duration_s = end_ts - start_ts + if duration_s > 0 and not math.isnan(self.total_energy_j): + self.avg_power_w = self.total_energy_j / duration_s + self.active = False + except Exception as exc: + self.error = f"Zeus end_window failed: {exc}" + self.enabled = False + self.active = False from rich.console import Group from rich.console import Console @@ -123,6 +185,14 @@ def __init__(self, args, model_group, training_group, logging_group): self.latest_ln_f_cosine_95 = float('nan') self.latest_rankme = float('nan') self.latest_areq = float('nan') + self.zeus_total_energy_j = float('nan') + self.zeus_avg_power_w = float('nan') + self.zeus_profiler = ZeusProfiler( + enabled=getattr(self.args, 'zeus_profile', False), + window_name=getattr(self.args, 'zeus_window_name', 'train_total'), + ) + if self.zeus_profiler.error: + print(f"[WARN] {self.zeus_profiler.error}") # store overall statistics for weights and activations self.latest_overall_weight_stats = { @@ -1818,6 +1888,8 @@ def run_validation_step(self, running_mfu, current_epoch, current_dataset, num_s f"{self.latest_ln_f_cosine_95:.6f}", f"{self.latest_rankme:.6f}", f"{self.latest_areq:.6f}", + f"{self.zeus_total_energy_j:.6f}", + f"{self.zeus_avg_power_w:.6f}", f"{self.latest_overall_weight_stats['stdev']:.6f}", f"{self.latest_overall_weight_stats['kurtosis']:.6f}", f"{self.latest_overall_weight_stats['max']:.6f}", @@ -1914,6 +1986,21 @@ def log_training_step(self, lossf, training_losses, running_mfu, current_epoch, else: self.log_metrics_non_validation(lossf, running_mfu, current_epoch, self.tokens_trained, prior_dataset, better_than_chance) + def _persist_zeus_metrics(self): + """Patch best_val_loss_and_iter.txt with final Zeus metrics if available.""" + metrics_path = os.path.join(self.args.out_dir, 'best_val_loss_and_iter.txt') + if not os.path.exists(metrics_path): + return + try: + parts = [p.strip() for p in Path(metrics_path).read_text().strip().split(',')] + while len(parts) < 23: + parts.append('nan') + parts[21] = f"{self.zeus_total_energy_j:.6f}" + parts[22] = f"{self.zeus_avg_power_w:.6f}" + Path(metrics_path).write_text(", ".join(parts) + "\n") + except Exception as exc: + print(f"[WARN] Failed to persist Zeus metrics: {exc}") + def train(self): if self.args.training_mode == 'multicontext': self.X_dict, self.Y_dict, dataset_list = self.get_batch('train') @@ -1923,6 +2010,7 @@ def train(self): self.X, self.Y, current_dataset = self.get_batch('train') self.X, self.Y, current_dataset = self.get_batch('train') t_start = time.time() + self.zeus_profiler.start() t0 = t_start local_iter_num = 0 running_mfu = -1.0 @@ -2209,6 +2297,13 @@ def train(self): live.start() break + self.zeus_profiler.stop() + self.zeus_total_energy_j = self.zeus_profiler.total_energy_j + self.zeus_avg_power_w = self.zeus_profiler.avg_power_w + self._persist_zeus_metrics() + if self.zeus_profiler.error: + print(f"[WARN] {self.zeus_profiler.error}") + if self.args.plot_statistics: plot_statistics(self.args, self.stats, graph_y_labels) diff --git a/train_args.py b/train_args.py index 9553b7d92b..b403952276 100644 --- a/train_args.py +++ b/train_args.py @@ -1398,6 +1398,18 @@ def parse_args(): logging_group.add_argument('--log_all_metrics', default=False, action=argparse.BooleanOptionalAction, help='Enable logging of all metrics including gns') logging_group.add_argument('--log_rankme', default=True, action=argparse.BooleanOptionalAction, help='Log RankMe representation metric during validation') logging_group.add_argument('--log_areq', default=True, action=argparse.BooleanOptionalAction, help='Log aReQ representation metric during validation') + logging_group.add_argument( + '--zeus_profile', + default=False, + action=argparse.BooleanOptionalAction, + help='Enable optional Zeus energy profiling (gracefully degrades if zeus is not installed).', + ) + logging_group.add_argument( + '--zeus_window_name', + default='train_total', + type=str, + help='Zeus profiling window label used for begin/end window tracking.', + ) # Turn activation/weight statistics off to save CPU RAM and wall time. training_group.add_argument(