-
Notifications
You must be signed in to change notification settings - Fork 30
Add graceful Zeus energy profiling metrics across training/search too… #782
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||
|---|---|---|---|---|---|---|---|---|
|
|
@@ -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: | ||||||||
|
||||||||
| if self.zeus_profiler.error: | |
| # Only emit this warning once in distributed setups (rank 0). | |
| if self.zeus_profiler.error and os.environ.get("RANK", "0") == "0": |
Copilot
AI
Mar 29, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
_persist_zeus_metrics hard-codes the expected metrics length (23) and the Zeus indices (21/22). This is brittle because best_val_loss_and_iter.txt’s layout has already grown beyond 23 fields; future insertions will silently corrupt the wrong columns. Consider defining named constants (or deriving indices from the same METRIC_KEYS list used by downstream parsers) and padding based on the max index you need rather than a magic total length.
Copilot
AI
Mar 29, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
There’s a duplicate get_batch('train') call here (one inside the if/else and one immediately after), which does extra work and changes the RNG/data stream (the first batch fetched is discarded). Consider removing the redundant call and only fetching the initial batch once for the selected training_mode.
| self.X, self.Y, current_dataset = self.get_batch('train') |
Copilot
AI
Mar 29, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
self.zeus_profiler.start() is called unconditionally, so in DDP every rank will start a Zeus window even though only rank 0 should profile and report a single set of energy metrics. Consider running Zeus profiling only on self.master_process and leaving metrics as NaN on other ranks to avoid conflicts and overhead.
Copilot
AI
Mar 29, 2026
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
In DDP runs this block will execute on every rank, which can (a) call Zeus begin/end multiple times and skew/duplicate measurements, and (b) race to rewrite best_val_loss_and_iter.txt from multiple processes. Please guard stop()/metric persistence/warn printing with self.master_process (or rank==0) and consider a barrier if other ranks depend on the file.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
With the new
zeus_energy/zeus_powertargets,d_costcan be negative when a candidate reduces energy/power vs baseline. The current efficiency computation later in_evaluatetreats negative cost deltas as flipping the sign (and can even select candidates with worse objective if both objective_improvement and d_cost are negative). Consider explicitly requiringobjective_improvement > 0and handlingd_cost <= 0as a special case (e.g., treat as +inf efficiency when improvement is positive, otherwise skip).