diff --git a/tools/auto_benchmark/metrics.py b/tools/auto_benchmark/metrics.py new file mode 100644 index 000000000..f15ab0973 --- /dev/null +++ b/tools/auto_benchmark/metrics.py @@ -0,0 +1,714 @@ +#!/usr/bin/env python3 + +import argparse +import csv +import os +import re +import shutil +from collections import defaultdict +from datetime import datetime +from statistics import mean + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +PRIMUS_ROOT = os.path.normpath(os.path.join(SCRIPT_DIR, "..", "..")) +RESULTS_DIR = os.path.join(SCRIPT_DIR, "results") + +WARMUP_SKIP = 5 + +ERROR_STATUS = "Error found - check log" +OK_STATUS = "OK" + +ANSI_ESCAPE_REGEX = re.compile(r"\x1b\[[0-9;]*m") +LOG_EXIT_CODE_REGEX = re.compile(r"primus launcher exited with code (\d+)", re.IGNORECASE) + +LOG_ERROR_PATTERNS = ( + re.compile(r"Traceback \(most recent call last\)", re.IGNORECASE), + re.compile(r"\[ERROR\]", re.IGNORECASE), + re.compile(r"\bRuntimeError:", re.IGNORECASE), + re.compile(r"\bOutOfMemoryError:", re.IGNORECASE), + re.compile(r"\bCUDA out of memory\b", re.IGNORECASE), + re.compile(r"\bHIP out of memory\b", re.IGNORECASE), + re.compile(r"\bNCCL error\b", re.IGNORECASE), + re.compile(r"\bfatal error\b", re.IGNORECASE), +) + +LOG_ERROR_EXCLUSIONS = ( + re.compile(r"error_injection", re.IGNORECASE), + re.compile(r"\[SKIP\].*Import failed", re.IGNORECASE), + re.compile(r"avoid ImportError", re.IGNORECASE), + re.compile(r"TORCH_NCCL_ASYNC_ERROR_HANDLING", re.IGNORECASE), + re.compile(r"destroy_process_group\(\)", re.IGNORECASE), +) + +RUN_LABEL_REGEX = re.compile(r"_run(\d+)\.log$", re.IGNORECASE) +LEGACY_TIMESTAMP_REGEX = re.compile(r"_(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.log$") +LEGACY_SUFFIX_REGEX = re.compile(r"_MI\d+X?_(.+)\.log$", re.IGNORECASE) + +ENV_DEFAULT_REGEX = re.compile(r"\$\{[A-Z0-9_]+:(\d+)\}") +PLAIN_INT_REGEX = re.compile(r"^\d+$") + +MEGATRON_BS_KEYS = ("micro_batch_size",) +MEGATRON_SEQ_KEYS = ("seq_length",) +MEGATRON_GBS_KEYS = ("global_batch_size",) + +TORCHTITAN_BS_KEYS = ("local_batch_size",) +TORCHTITAN_SEQ_KEYS = ("seq_len", "seq_length") +TORCHTITAN_GBS_KEYS = ("global_batch_size",) + +MEGATRON_NOTE_TEXT = """ +NOTE: +- Results are saved to results/metrics_megatron.csv (latest) and a timestamped snapshot. +- "Run" is run1, run2, ... per model and device, ordered oldest to newest. +- "BS", "Seq", and "GBS" come from the benchmark yaml when available. +- Megatron logs often print each iteration twice (Primus log forwarding); metrics + deduplicate by iteration number before averaging. +- Timing/throughput fields use the current value before "/" (e.g. 5896.1/5913.1 -> 5896.1). +- Warm-up: the first five iterations are excluded before averaging. +- "Status" is "Error found - check log" when the log contains errors or metrics could not be computed. +- Numeric fields may contain commas in logs; commas are removed before averaging. +""" + +TORCHTITAN_NOTE_TEXT = """ +NOTE: +- Results are saved to results/metrics_torchtitan.csv (latest) and a timestamped snapshot. +- "Run" is run1, run2, ... per model and device, ordered oldest to newest. +- "BS", "Seq", and "GBS" come from the benchmark yaml when available. +- "Steps" is the number of training steps used after dropping the first five warm-up steps. +- TPS and TFLOPS values may contain commas in logs; commas are removed before averaging. +- "Status" is "Error found - check log" when the log contains errors or metrics could not be computed. +""" + +MEGATRON_NUM = r"[\d,]+(?:\.\d+)?" +MEGATRON_METRIC_VALUE = rf"({MEGATRON_NUM})(?:\s*/\s*{MEGATRON_NUM})?" + +MEGATRON_ITERATION_REGEX = re.compile( + rf"iteration\s+(\d+)/\s*\d+.*?" + rf"elapsed time per iteration \(ms\):\s*{MEGATRON_METRIC_VALUE}.*?" + rf"throughput per GPU \(TFLOP/s/GPU\):\s*{MEGATRON_METRIC_VALUE}.*?" + rf"(?:tokens per GPU \(tokens/s/GPU\):\s*{MEGATRON_METRIC_VALUE}.*?)?" + rf"global batch size:\s*(\d+)", + re.IGNORECASE, +) + +MEGATRON_FILENAME_REGEX = re.compile(r"(?P.+?)_megatron_(?PMI\d+X?)", re.IGNORECASE) + +TORCHTITAN_STEP_REGEX = re.compile( + r"step:\s*(\d+).*?" + r"memory:\s*([\d.]+)GiB.*?" + r"tps:\s*([\d,]+(?:\.\d+)?).*?" + r"tflops:\s*([\d,]+(?:\.\d+)?).*?" + r"mfu:\s*([\d.]+)%" +) + +TORCHTITAN_BS_REGEX = re.compile(r"training\.local_batch_size\s*\.{2,}\s*(\d+)") +TORCHTITAN_SEQ_REGEX = re.compile(r"training\.seq_len\s*\.{2,}\s*(\d+)") + +TORCHTITAN_FILENAME_REGEX = re.compile(r"(?P.+?)_torchtitan_(?PMI\d+X?)", re.IGNORECASE) + +PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE) + + +def _parse_scalar(value): + if value is None: + return None + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return int(value) if float(value).is_integer() else value + if isinstance(value, str): + stripped = value.strip() + if PLAIN_INT_REGEX.match(stripped): + return int(stripped) + m = ENV_DEFAULT_REGEX.search(stripped) + if m: + return int(m.group(1)) + return None + + +def _deep_find(node, keys): + if isinstance(node, dict): + for key in keys: + if key in node: + parsed = _parse_scalar(node[key]) + if parsed is not None: + return parsed + for child in node.values(): + found = _deep_find(child, keys) + if found is not None: + return found + elif isinstance(node, list): + for item in node: + found = _deep_find(item, keys) + if found is not None: + return found + return None + + +def _regex_find(text, keys): + for key in keys: + m = re.search( + rf"^\s*{re.escape(key)}:\s*(.+?)(?:\s+#.*)?$", + text, + re.MULTILINE, + ) + if not m: + continue + parsed = _parse_scalar(m.group(1).strip()) + if parsed is not None: + return parsed + return None + + +def load_yaml_dict(path): + try: + import yaml + except ImportError: + return None + + try: + with open(path, "r", errors="ignore") as f: + data = yaml.safe_load(f) + except Exception: + return None + + return data if isinstance(data, dict) else None + + +def resolve_config_yaml(log_fname, backend, device, model, log_dir): + base = log_fname[:-4] if log_fname.endswith(".log") else log_fname + + for suffix in ("_override", "_edited", ""): + candidate = os.path.join(log_dir, f"{base}{suffix}.yaml") + if os.path.isfile(candidate): + return candidate + + config_dir = os.path.join(PRIMUS_ROOT, "examples", backend, "configs", device) + return os.path.join(config_dir, f"{model}.yaml") + + +def load_training_params(backend, device, model, log_fname, log_dir): + path = resolve_config_yaml(log_fname, backend, device, model, log_dir) + if backend == "megatron": + bs_keys, seq_keys, gbs_keys = MEGATRON_BS_KEYS, MEGATRON_SEQ_KEYS, MEGATRON_GBS_KEYS + else: + bs_keys, seq_keys, gbs_keys = TORCHTITAN_BS_KEYS, TORCHTITAN_SEQ_KEYS, TORCHTITAN_GBS_KEYS + + bs = seq = gbs = None + + data = load_yaml_dict(path) + if data is not None: + bs = _deep_find(data, bs_keys) + seq = _deep_find(data, seq_keys) + gbs = _deep_find(data, gbs_keys) + + if os.path.isfile(path) and (bs is None or seq is None or gbs is None): + with open(path, "r", errors="ignore") as f: + text = f.read() + if bs is None: + bs = _regex_find(text, bs_keys) + if seq is None: + seq = _regex_find(text, seq_keys) + if gbs is None: + gbs = _regex_find(text, gbs_keys) + + return ( + bs if bs is not None else "-", + seq if seq is not None else "-", + gbs if gbs is not None else "-", + ) + + +def log_chronological_key(fname, path): + m = RUN_LABEL_REGEX.search(fname) + if m: + return (0, int(m.group(1)), fname) + + m = LEGACY_TIMESTAMP_REGEX.search(fname) + if m: + return (1, m.group(1), fname) + + m = LEGACY_SUFFIX_REGEX.search(fname) + if m and m.group(1): + return (1, m.group(1), fname) + + return (2, os.path.getmtime(path), fname) + + +def assign_run_labels(entries): + grouped = defaultdict(list) + for entry in entries: + grouped[(entry["model"], entry["device"])].append(entry) + + for group in grouped.values(): + group.sort(key=lambda entry: log_chronological_key(entry["fname"], entry["path"])) + for index, entry in enumerate(group, start=1): + entry["run"] = f"run{index}" + + return entries + + +def run_sort_key(run_label): + m = re.fullmatch(r"run(\d+)", run_label) + if m: + return int(m.group(1)) + return 0 + + +def terminal_width(default=120): + try: + return shutil.get_terminal_size().columns + except OSError: + return default + + +def print_table(headers, rows, max_col_width=36): + if not rows: + return + + widths = [] + for i, header in enumerate(headers): + col_width = max(len(str(header)), *(len(str(row[i])) for row in rows)) + widths.append(min(col_width, max_col_width)) + + def fmt(row): + cells = [] + for i, value in enumerate(row): + text = str(value) + if len(text) > widths[i]: + text = text[: widths[i] - 3] + "..." + cells.append(text.ljust(widths[i])) + return "| " + " | ".join(cells) + " |" + + sep = "+-" + "-+-".join("-" * w for w in widths) + "-+" + print(sep) + print(fmt(headers)) + print(sep) + for row in rows: + print(fmt(row)) + print(sep) + + +def save_csv(headers, rows, backend): + os.makedirs(RESULTS_DIR, exist_ok=True) + timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + latest_path = os.path.join(RESULTS_DIR, f"metrics_{backend}.csv") + snapshot_path = os.path.join(RESULTS_DIR, f"metrics_{backend}_{timestamp}.csv") + + for path in (latest_path, snapshot_path): + with open(path, "w", newline="", encoding="utf-8") as handle: + writer = csv.writer(handle) + writer.writerow(headers) + writer.writerows(rows) + + return latest_path, snapshot_path + + +def strip_ansi(text): + return ANSI_ESCAPE_REGEX.sub("", text) + + +def log_has_error(path): + exit_code = None + saw_error_line = False + + with open(path, "r", errors="ignore") as f: + for line in f: + plain = strip_ansi(line) + + m = LOG_EXIT_CODE_REGEX.search(plain) + if m: + exit_code = int(m.group(1)) + + if any(pattern.search(plain) for pattern in LOG_ERROR_EXCLUSIONS): + continue + + if any(pattern.search(plain) for pattern in LOG_ERROR_PATTERNS): + saw_error_line = True + + if exit_code not in (None, 0): + return True + return saw_error_line + + +def render_results(backend, headers, rows, note_text): + if not rows: + print(f"No {backend} logs found.") + return None + + latest_path, snapshot_path = save_csv(headers, rows, backend) + print_table(headers, rows) + + error_rows = [row for row in rows if len(row) > 2 and row[2] == ERROR_STATUS] + if error_rows: + print( + f"\n{len(error_rows)} run(s) reported errors or incomplete metrics." + " Open the corresponding log file for details." + ) + + table_width = sum(len(str(header)) for header in headers) + 3 * len(headers) + if table_width > terminal_width(): + print("\n(Table may be wider than the terminal. Open the CSV for the full view.)") + + print(note_text) + print("\nMetrics saved to:") + print(f" Latest: {latest_path}") + print(f" Snapshot: {snapshot_path}") + + return latest_path + + +def is_megatron(filename): + name = filename.lower() + return "megatron" in name or "megatorn" in name + + +def megatron_parse_filename(filename): + m = MEGATRON_FILENAME_REGEX.search(filename) + if not m: + return None + + p = PRECISION_REGEX.search(filename) + precision = p.group(1).upper() if p else "-" + + return { + "model": m.group("model"), + "device": m.group("device"), + "precision": precision, + } + + +def megatron_to_float(num_str): + return float(num_str.replace(",", "")) + + +def megatron_parse_log_file(path): + records_by_iter = {} + + with open(path, "r", errors="ignore") as f: + for line in f: + if "iteration" not in line: + continue + + m = MEGATRON_ITERATION_REGEX.search(line) + if not m: + continue + + iter_num = int(m.group(1)) + if iter_num in records_by_iter: + continue + + records_by_iter[iter_num] = { + "iter": iter_num, + "elapsed_ms": megatron_to_float(m.group(2)), + "tflops_gpu": megatron_to_float(m.group(3)), + "tokens_gpu": megatron_to_float(m.group(4)) if m.group(4) else None, + "gbs": int(m.group(5)), + } + + return [records_by_iter[i] for i in sorted(records_by_iter)] + + +def megatron_compute_averages(records): + if len(records) <= WARMUP_SKIP: + return None + + records = records[WARMUP_SKIP:] + token_values = [r["tokens_gpu"] for r in records if r["tokens_gpu"] is not None] + + return { + "count": len(records), + "elapsed_ms": mean(r["elapsed_ms"] for r in records), + "tflops_gpu": mean(r["tflops_gpu"] for r in records), + "tokens_gpu": mean(token_values) if token_values else None, + "gbs": records[0]["gbs"], + } + + +def megatron_collect_rows(): + log_dir = os.path.join(RESULTS_DIR, "logs_megatron") + entries = [] + + if not os.path.isdir(log_dir): + return [] + + for fname in sorted(os.listdir(log_dir)): + if not fname.endswith(".log") or not is_megatron(fname): + continue + + meta = megatron_parse_filename(fname) + if not meta: + continue + + path = os.path.join(log_dir, fname) + has_error = log_has_error(path) + records = megatron_parse_log_file(path) + stats = megatron_compute_averages(records) + + bs, seq, gbs = load_training_params("megatron", meta["device"], meta["model"], fname, log_dir) + + if has_error or not stats: + if gbs == "-" and records: + gbs = records[0]["gbs"] + values = [ + ERROR_STATUS, + "megatron", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + "-", + "-", + "-", + "-", + ] + else: + if gbs == "-": + gbs = stats["gbs"] + + tokens_gpu = f"{stats['tokens_gpu']:.2f}" if stats["tokens_gpu"] is not None else "-" + values = [ + OK_STATUS, + "megatron", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + stats["count"], + f"{stats['elapsed_ms']:.2f}", + f"{stats['tflops_gpu']:.2f}", + tokens_gpu, + ] + + entries.append( + { + "model": meta["model"], + "device": meta["device"], + "fname": fname, + "path": path, + "values": values, + } + ) + + assign_run_labels(entries) + + rows = [[entry["model"], entry["run"], *entry["values"]] for entry in entries] + rows.sort(key=lambda row: (row[0], run_sort_key(row[1]), row[7])) + return rows + + +def megatron_main(): + headers = [ + "Model", + "Run", + "Status", + "Backend", + "Device", + "BS", + "Seq", + "GBS", + "Precision", + "Iterations", + "Iter Time (ms)", + "TFLOPS/GPU", + "Tokens/GPU", + ] + render_results("megatron", headers, megatron_collect_rows(), MEGATRON_NOTE_TEXT) + + +def is_torchtitan(filename): + return "torchtitan" in filename.lower() + + +def torchtitan_parse_filename(filename): + m = TORCHTITAN_FILENAME_REGEX.search(filename) + if not m: + return None + + p = PRECISION_REGEX.search(filename) + precision = p.group(1).upper() if p else "-" + + return { + "model": m.group("model"), + "device": m.group("device"), + "precision": precision, + } + + +def torchtitan_parse_log_file(path): + steps_by_num = {} + + with open(path, "r", errors="ignore") as f: + for line in f: + m = TORCHTITAN_STEP_REGEX.search(line) + if not m: + continue + + step_num = int(m.group(1)) + if step_num in steps_by_num: + continue + + steps_by_num[step_num] = { + "step": step_num, + "memory": float(m.group(2)), + "tps": float(m.group(3).replace(",", "")), + "tflops": float(m.group(4).replace(",", "")), + "mfu": float(m.group(5)), + } + + return [steps_by_num[i] for i in sorted(steps_by_num)] + + +def torchtitan_parse_log_training_fallback(path): + bs = None + seq = None + + with open(path, "r", errors="ignore") as f: + for line in f: + if bs is None: + m = TORCHTITAN_BS_REGEX.search(line) + if m: + bs = int(m.group(1)) + if seq is None: + m = TORCHTITAN_SEQ_REGEX.search(line) + if m: + seq = int(m.group(1)) + if bs is not None and seq is not None: + break + + return bs, seq + + +def torchtitan_compute_averages(steps): + if len(steps) <= WARMUP_SKIP: + return None + + steps = steps[WARMUP_SKIP:] + + return { + "count": len(steps), + "memory": mean(s["memory"] for s in steps), + "tps": mean(s["tps"] for s in steps), + "tflops": mean(s["tflops"] for s in steps), + "mfu": mean(s["mfu"] for s in steps), + } + + +def torchtitan_collect_rows(): + log_dir = os.path.join(RESULTS_DIR, "logs_torchtitan") + entries = [] + + if not os.path.isdir(log_dir): + return [] + + for fname in sorted(os.listdir(log_dir)): + if not fname.endswith(".log") or not is_torchtitan(fname): + continue + + meta = torchtitan_parse_filename(fname) + if not meta: + continue + + path = os.path.join(log_dir, fname) + has_error = log_has_error(path) + steps = torchtitan_parse_log_file(path) + stats = torchtitan_compute_averages(steps) + + bs, seq, gbs = load_training_params("torchtitan", meta["device"], meta["model"], fname, log_dir) + if bs == "-": + log_bs, _ = torchtitan_parse_log_training_fallback(path) + if log_bs is not None: + bs = log_bs + if seq == "-": + _, log_seq = torchtitan_parse_log_training_fallback(path) + if log_seq is not None: + seq = log_seq + + if has_error or not stats: + values = [ + ERROR_STATUS, + "torchtitan", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + "-", + "-", + "-", + "-", + "-", + ] + else: + values = [ + OK_STATUS, + "torchtitan", + meta["device"], + bs, + seq, + gbs, + meta["precision"], + stats["count"], + f"{stats['memory']:.2f}", + f"{stats['tps']:.2f}", + f"{stats['tflops']:.2f}", + f"{stats['mfu']:.2f}", + ] + + entries.append( + { + "model": meta["model"], + "device": meta["device"], + "fname": fname, + "path": path, + "values": values, + } + ) + + assign_run_labels(entries) + + rows = [[entry["model"], entry["run"], *entry["values"]] for entry in entries] + rows.sort(key=lambda row: (row[0], run_sort_key(row[1]), row[7])) + return rows + + +def torchtitan_main(): + headers = [ + "Model", + "Run", + "Status", + "Backend", + "Device", + "BS", + "Seq", + "GBS", + "Precision", + "Steps", + "Mem(GiB)", + "TPS", + "TFLOPS", + "MFU(%)", + ] + render_results("torchtitan", headers, torchtitan_collect_rows(), TORCHTITAN_NOTE_TEXT) + + +BACKENDS = { + "megatron": megatron_main, + "torchtitan": torchtitan_main, +} + + +def main(): + parser = argparse.ArgumentParser(description="Generate benchmark metrics tables and CSVs.") + parser.add_argument( + "backend", + choices=sorted(BACKENDS), + help="Training backend to process logs for", + ) + args = parser.parse_args() + BACKENDS[args.backend]() + + +if __name__ == "__main__": + main() diff --git a/tools/auto_benchmark/metrics_megatron.py b/tools/auto_benchmark/metrics_megatron.py deleted file mode 100644 index ff70c9c79..000000000 --- a/tools/auto_benchmark/metrics_megatron.py +++ /dev/null @@ -1,212 +0,0 @@ -#!/usr/bin/env python3 - -import os -import re -from statistics import mean - -LOG_DIR = "/workspace/Primus/tools/auto_benchmark/results/logs_megatron" - -# ============================================================ -# NOTE ON ITERATIONS AND AVERAGING -# ============================================================ - -NOTE_TEXT = """ -NOTE: -- "Iterations" represents the number of training iterations USED to compute - the averages shown above (after warm-up removal). - -- All iteration records are extracted from the log file, sorted by iteration - number, and the FIRST TWO iterations are discarded to remove warm-up effects. - -- Iteration numbers do NOT need to be sequential. Valid examples: - * iteration 1, iteration 2, iteration 3 - * iteration 1, iteration 5, iteration 10 - * iteration 1, iteration 10, iteration 20 - -- Numeric fields (elapsed time, TFLOPS/GPU, tokens/GPU) may appear as: - 1234 - 1,234 - 1,234.56 - Commas are removed before averaging. -""" - -# ============================================================ -# Regex patterns (comma-safe) -# ============================================================ - -NUM = r"[\d,]+(?:\.\d+)?" - -ITERATION_REGEX = re.compile( - rf"iteration\s+(\d+)/\s*\d+.*?" - rf"elapsed time per iteration \(ms\):\s*({NUM}).*?" - rf"throughput per GPU \(TFLOP/s/GPU\):\s*({NUM}).*?" - rf"tokens per GPU \(tokens/s/GPU\):\s*({NUM}).*?" - rf"global batch size:\s*(\d+)", - re.IGNORECASE, -) - -# Filename metadata -FILENAME_REGEX = re.compile(r"(?P.+?)_megatron_(?PMI\d+X?)", re.IGNORECASE) - -PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE) - -# ============================================================ -# Helpers -# ============================================================ - - -def is_megatron(filename: str) -> bool: - name = filename.lower() - return "megatron" in name or "megatorn" in name - - -def parse_filename(filename: str): - m = FILENAME_REGEX.search(filename) - if not m: - return None - - p = PRECISION_REGEX.search(filename) - precision = p.group(1).upper() if p else "-" - - return { - "model": m.group("model"), - "device": m.group("device"), - "precision": precision, - } - - -def to_float(num_str: str) -> float: - """Convert numbers like '1,234.56' safely to float.""" - return float(num_str.replace(",", "")) - - -# ============================================================ -# Log parsing (FULL FILE SCAN) -# ============================================================ - - -def parse_log_file(path): - records = [] - - with open(path, "r", errors="ignore") as f: - for line in f: - m = ITERATION_REGEX.search(line) - if m: - records.append( - { - "iter": int(m.group(1)), - "elapsed_ms": to_float(m.group(2)), - "tflops_gpu": to_float(m.group(3)), - "tokens_gpu": to_float(m.group(4)), - "gbs": int(m.group(5)), - } - ) - - return records - - -def compute_averages(records): - if len(records) <= 2: - return None - - # Sort by iteration number (non-sequential safe) - records = sorted(records, key=lambda x: x["iter"]) - - # Drop first two warm-up iterations - records = records[2:] - - return { - "count": len(records), - "elapsed_ms": mean(r["elapsed_ms"] for r in records), - "tflops_gpu": mean(r["tflops_gpu"] for r in records), - "tokens_gpu": mean(r["tokens_gpu"] for r in records), - "gbs": records[0]["gbs"], - } - - -# ============================================================ -# Table printer -# ============================================================ - - -def print_table(headers, rows): - widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))] - - def fmt(row): - return "| " + " | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))) + " |" - - sep = "+-" + "-+-".join("-" * w for w in widths) + "-+" - - print(sep) - print(fmt(headers)) - print(sep) - for r in rows: - print(fmt(r)) - print(sep) - - -def print_note(): - print(NOTE_TEXT) - - -# ============================================================ -# Main -# ============================================================ - - -def main(): - rows = [] - - for fname in sorted(os.listdir(LOG_DIR)): - if not fname.endswith(".log"): - continue - - if not is_megatron(fname): - continue - - meta = parse_filename(fname) - if not meta: - continue - - path = os.path.join(LOG_DIR, fname) - - records = parse_log_file(path) - stats = compute_averages(records) - if not stats: - continue - - rows.append( - [ - meta["model"], - "megatron", - meta["device"], - meta["precision"], - stats["count"], - f"{stats['elapsed_ms']:.2f}", - f"{stats['tflops_gpu']:.2f}", - f"{stats['tokens_gpu']:.2f}", - stats["gbs"], - ] - ) - - headers = [ - "Model", - "Backend", - "Device", - "Precision", - "Iterations", - "Iter Time (ms)", - "TFLOPS/GPU", - "Tokens/GPU", - "GBS", - ] - - if rows: - print_table(headers, rows) - print_note() - else: - print("No Megatron logs found.") - - -if __name__ == "__main__": - main() diff --git a/tools/auto_benchmark/metrics_torchtitan.py b/tools/auto_benchmark/metrics_torchtitan.py deleted file mode 100644 index 88770f72f..000000000 --- a/tools/auto_benchmark/metrics_torchtitan.py +++ /dev/null @@ -1,224 +0,0 @@ -#!/usr/bin/env python3 - -import os -import re -from statistics import mean - -LOG_DIR = "/workspace/Primus/tools/auto_benchmark/results/logs_torchtitan" - -# ============================================================ -# Regex patterns -# ============================================================ - -STEP_REGEX = re.compile( - r"step:\s*(\d+).*?" - r"memory:\s*([\d.]+)GiB.*?" - r"tps:\s*([\d,]+(?:\.\d+)?).*?" - r"tflops:\s*([\d,]+(?:\.\d+)?).*?" - r"mfu:\s*([\d.]+)%" -) - -# Dot-aligned config values -BS_REGEX = re.compile(r"training\.local_batch_size\s*\.{2,}\s*(\d+)") - -SEQ_REGEX = re.compile(r"training\.seq_len\s*\.{2,}\s*(\d+)") - -# Filename metadata -FILENAME_REGEX = re.compile(r"(?P.+?)_torchtitan_(?PMI\d+X?)", re.IGNORECASE) - -PRECISION_REGEX = re.compile(r"(BF16|FP8)", re.IGNORECASE) - -# ============================================================ -# Helpers -# ============================================================ - - -def is_torchtitan(filename): - return "torchtitan" in filename.lower() - - -def parse_filename(filename): - m = FILENAME_REGEX.search(filename) - if not m: - return None - - p = PRECISION_REGEX.search(filename) - precision = p.group(1).upper() if p else "-" - - return { - "model": m.group("model"), - "device": m.group("device"), - "precision": precision, - } - - -# ============================================================ -# Log parsing (FULL FILE SCAN) -# ============================================================ - - -def parse_log_file(path): - """ - Fully scans the log file to extract: - - BS from training.local_batch_size - - SEQ from training.seq_len - - Per-step performance metrics - """ - bs = None - seq = None - steps = [] - - with open(path, "r", errors="ignore") as f: - for line in f: - # Batch size - if bs is None: - m = BS_REGEX.search(line) - if m: - bs = int(m.group(1)) - - # Sequence length - if seq is None: - m = SEQ_REGEX.search(line) - if m: - seq = int(m.group(1)) - - # Step metrics - m = STEP_REGEX.search(line) - if m: - steps.append( - { - "step": int(m.group(1)), - "memory": float(m.group(2)), - "tps": float(m.group(3).replace(",", "")), - "tflops": float(m.group(4).replace(",", "")), - "mfu": float(m.group(5)), - } - ) - - return bs if bs is not None else "-", seq if seq is not None else "-", steps - - -def compute_averages(steps): - if len(steps) <= 2: - return None - - steps = sorted(steps, key=lambda x: x["step"]) - steps = steps[2:] # drop first two warm-up steps - - return { - "count": len(steps), - "memory": mean(s["memory"] for s in steps), - "tps": mean(s["tps"] for s in steps), - "tflops": mean(s["tflops"] for s in steps), - "mfu": mean(s["mfu"] for s in steps), - } - - -# ============================================================ -# Table printer -# ============================================================ - - -def print_table(headers, rows): - widths = [max(len(str(r[i])) for r in ([headers] + rows)) for i in range(len(headers))] - - def fmt(row): - return "| " + " | ".join(str(row[i]).ljust(widths[i]) for i in range(len(row))) + " |" - - sep = "+-" + "-+-".join("-" * w for w in widths) + "-+" - - print(sep) - print(fmt(headers)) - print(sep) - for r in rows: - print(fmt(r)) - print(sep) - - -# ============================================================ -# Note printer -# ============================================================ - - -def print_note(): - print( - """ -NOTE: -- "Steps" represents the number of training steps USED to compute averages. -- Steps are sorted by step number and the FIRST TWO steps are dropped - to remove warm-up effects. -- Step numbers do NOT need to be sequential. Valid examples: - * step 1, step 2, step 3 - * step 1, step 5, step 10 - * step 1, step 10, step 20 -- TPS and TFLOPS values may contain commas (e.g. 1,202.17); - commas are removed before averaging. -""" - ) - - -# ============================================================ -# Main -# ============================================================ - - -def main(): - rows = [] - - for fname in sorted(os.listdir(LOG_DIR)): - if not fname.endswith(".log"): - continue - - if not is_torchtitan(fname): - continue - - meta = parse_filename(fname) - if not meta: - continue - - path = os.path.join(LOG_DIR, fname) - - bs, seq, steps = parse_log_file(path) - stats = compute_averages(steps) - if not stats: - continue - - rows.append( - [ - meta["model"], - "torchtitan", - meta["device"], - bs, - seq, - meta["precision"], - stats["count"], - f"{stats['memory']:.2f}", - f"{stats['tps']:.2f}", - f"{stats['tflops']:.2f}", - f"{stats['mfu']:.2f}", - ] - ) - - headers = [ - "Model", - "Backend", - "Device", - "BS", - "Seq", - "Precision", - "Steps", - "Mem(GiB)", - "TPS", - "TFLOPS", - "MFU(%)", - ] - - if rows: - print_table(headers, rows) - print_note() - else: - print("No TorchTitan logs found.") - - -if __name__ == "__main__": - main() diff --git a/tools/auto_benchmark/run_primus_autobenchmark.sh b/tools/auto_benchmark/run_primus_autobenchmark.sh index 53201b858..4af17ecfe 100644 --- a/tools/auto_benchmark/run_primus_autobenchmark.sh +++ b/tools/auto_benchmark/run_primus_autobenchmark.sh @@ -1,6 +1,11 @@ #!/usr/bin/env bash # Removed: set -e (allow script to continue on errors in benchmarks) +if [[ -z "${BASH_VERSION:-}" ]]; then + echo "This script must be run with bash (not sh)." >&2 + exit 1 +fi + # Set up trap to debug unexpected exits trap 'echo "[DEBUG] Script exiting at line $LINENO with exit code $?"' EXIT @@ -29,6 +34,7 @@ PRIMUS_ROOT="/workspace/Primus" MEGATRON_BASE_DIR="${PRIMUS_ROOT}/examples/megatron/configs" TORCHTITAN_BASE_DIR="${PRIMUS_ROOT}/examples/torchtitan/configs" RUN_SCRIPT="${PRIMUS_ROOT}/examples/run_pretrain.sh" +VALID_DEVICES=(MI300X MI325X MI355X) # Check if run_pretrain.sh exists, otherwise try run_pretrain_1.sh if [[ ! -f "$RUN_SCRIPT" && -f "${PRIMUS_ROOT}/examples/run_pretrain_1.sh" ]]; then @@ -36,6 +42,213 @@ if [[ ! -f "$RUN_SCRIPT" && -f "${PRIMUS_ROOT}/examples/run_pretrain_1.sh" ]]; t echo "[DEBUG] Using run_pretrain_1.sh instead" fi +# ------------------------------------------ +# Helpers +# ------------------------------------------ +install_vim_editor() { + echo -e "${YELLOW}⚠ No editor found. Installing vim...${RESET}" + + if [[ $EUID -eq 0 ]]; then + apt-get update && apt-get install -y vim + elif command -v sudo &>/dev/null; then + sudo apt-get update && sudo apt-get install -y vim + else + echo -e "${RED}✗ Cannot install vim: not root and sudo is unavailable.${RESET}" + echo -e " ${DOT} Install vim manually, then re-run config editing:" + echo -e " ${CYAN}apt-get update && apt-get install -y vim${RESET}" + return 1 + fi +} + +open_config_editor() { + local config_file="$1" + local candidate editor_bin editor_args editor_label + + for candidate in \ + "${EDITOR:-}" \ + "${VISUAL:-}" \ + "vim" \ + "vi" \ + "nano" \ + "emacs -nw" \ + "code --wait" \ + "cursor --wait"; do + if [[ -z "$candidate" ]]; then + continue + fi + + editor_bin="${candidate%% *}" + if ! command -v "$editor_bin" &>/dev/null; then + continue + fi + + editor_args="${candidate#"$editor_bin"}" + editor_label="$candidate" + echo -e " ${DOT} Using editor: ${CYAN}$editor_label${RESET}" + # shellcheck disable=SC2086 + "$editor_bin" $editor_args "$config_file" + return 0 + done + + if install_vim_editor && command -v vim &>/dev/null; then + echo -e " ${DOT} Using editor: ${CYAN}vim${RESET}" + vim "$config_file" + return 0 + fi + + echo -e "${RED}✗ Failed to open an editor for:${RESET} ${CYAN}$config_file${RESET}" + return 1 +} + +next_run_number() { + local model_name="$1" + local prefix="${model_name}_${BACKEND}_${DEVICE}" + local max_run=0 + local f bn run_n legacy_count=0 + + shopt -s nullglob + for f in "$LOG_DIR"/"${prefix}"_run*.log; do + bn=$(basename "$f" .log) + if [[ "$bn" =~ _run([0-9]+)$ ]]; then + run_n="${BASH_REMATCH[1]}" + if (( run_n > max_run )); then + max_run=$run_n + fi + fi + done + + for f in "$LOG_DIR"/"${prefix}"_*.log; do + bn=$(basename "$f" .log) + if [[ "$bn" =~ _run[0-9]+$ ]]; then + continue + fi + legacy_count=$((legacy_count + 1)) + done + shopt -u nullglob + + echo $((max_run + legacy_count + 1)) +} + +prepare_benchmark_artifacts() { + local cfg_file="$1" + local model_name run_num artifact_prefix + + PREP_CFG_FILE="$cfg_file" + PREP_MODEL_NAME=$(basename "$cfg_file" .yaml) + model_name="$PREP_MODEL_NAME" + run_num=$(next_run_number "$model_name") + PREP_RUN_LABEL="run${run_num}" + artifact_prefix="${model_name}_${BACKEND}_${DEVICE}_${PREP_RUN_LABEL}" + + PREP_LOG_FILE="$LOG_DIR/${artifact_prefix}.log" + + if [[ -n "${EDITED_CONFIGS[$cfg_file]:-}" ]]; then + PREP_WORKING_CONFIG="${EDITED_CONFIGS[$cfg_file]}" + else + PREP_WORKING_CONFIG="$cfg_file" + fi + + if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then + PREP_WORKING_CONFIG="$LOG_DIR/${artifact_prefix}_override.yaml" + cp "${EDITED_CONFIGS[$cfg_file]:-$cfg_file}" "$PREP_WORKING_CONFIG" + for KEY in "${!PARAM_OVERRIDES[@]}"; do + sed -i "s|^\([[:space:]]*${KEY}:[[:space:]]*\).*|\1${PARAM_OVERRIDES[$KEY]}|g" "$PREP_WORKING_CONFIG" + done + elif [[ -n "${EDITED_CONFIGS[$cfg_file]:-}" ]]; then + PREP_WORKING_CONFIG="$LOG_DIR/${artifact_prefix}_edited.yaml" + cp "${EDITED_CONFIGS[$cfg_file]}" "$PREP_WORKING_CONFIG" + fi +} + +execute_benchmark_run() { + local cfg_file="$1" + local working_config="$2" + local log_file="$3" + local model_name="$4" + local current="$5" + local total="$6" + + local original_config_backup="" + local run_exit_code=0 + + echo -e "${STAR} ${BOLD}Starting Benchmark ${current}/${total}...${RESET}" + echo -e " ${DOT} Model: ${CYAN}$model_name${RESET}" + echo -e " ${DOT} Backend: ${CYAN}$BACKEND${RESET}" + echo -e " ${DOT} Device: ${CYAN}$DEVICE${RESET}" + echo -e " ${DOT} Config: ${YELLOW}$working_config${RESET}" + echo -e " ${DOT} Log: ${YELLOW}$log_file${RESET}\n" + + if [[ "$working_config" != "$cfg_file" ]]; then + original_config_backup="${cfg_file}.backup_$$" + cp "$cfg_file" "$original_config_backup" + cp "$working_config" "$cfg_file" + echo -e " ${CHECK} Copied edited/overridden config to: ${CYAN}$cfg_file${RESET}" + fi + + EXP="${BACKEND_BASE_DIR}/${DEVICE}/$(basename "$cfg_file")" + export EXP + echo -e " ${CHECK} EXP set to: ${CYAN}$EXP${RESET}\n" + + echo -e " ${DOT} Changing to Primus root directory: ${CYAN}$PRIMUS_ROOT${RESET}" + cd "$PRIMUS_ROOT" || return 1 + + set +e + bash "$RUN_SCRIPT" 2>&1 | tee "$log_file" || true + run_exit_code=$? + set +e + + cd "$SCRIPT_DIR" || return 1 + + if [[ -n "$original_config_backup" && -f "$original_config_backup" ]]; then + mv "$original_config_backup" "$cfg_file" + echo -e " ${CHECK} Restored original config file" + fi + + echo + echo -e "${GREEN}==========================================${RESET}" + if [[ $run_exit_code -eq 0 ]]; then + echo -e " ${BOLD}${GREEN}✓ Benchmark ${current}/${total} Completed Successfully!${RESET}" + else + echo -e " ${BOLD}${YELLOW}⚠ Benchmark ${current}/${total} Completed with Exit Code: $run_exit_code${RESET}" + fi + echo -e " Log saved at:" + echo -e " ${CYAN}$log_file${RESET}" + if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then + echo -e " Override config saved at:" + echo -e " ${CYAN}$working_config${RESET}" + fi + echo -e "${GREEN}==========================================${RESET}" + echo + + return "$run_exit_code" +} + +generate_metrics_table() { + local metrics_script="metrics.py" + + echo + echo -e "${STAR} ${BOLD}Generating Metrics Table...${RESET}\n" + + if [[ -f "$SCRIPT_DIR/$metrics_script" && ( "$BACKEND" == "megatron" || "$BACKEND" == "torchtitan" ) ]]; then + echo -e " ${CHECK} Running: ${CYAN}python $metrics_script $BACKEND${RESET}\n" + metrics_output=$(cd "$SCRIPT_DIR" && python "$metrics_script" "$BACKEND") + metrics_status=$? + printf '%s\n' "$metrics_output" + echo + if [[ $metrics_status -eq 0 ]]; then + csv_path=$(printf '%s\n' "$metrics_output" | awk '/^ Latest:/{print $2}') + echo -e " ${CHECK} ${GREEN}Metrics table generated successfully${RESET}" + if [[ -n "$csv_path" ]]; then + echo -e " ${DOT} CSV: ${CYAN}$csv_path${RESET}" + fi + else + echo -e " ${RED}✗ Metrics generation failed${RESET}" + fi + else + echo -e " ${RED}✗ Metrics script not found: ${metrics_script:-unknown}${RESET}" + fi +} + # ------------------------------------------ # Banner # ------------------------------------------ @@ -85,14 +298,49 @@ sleep 0.2 # ------------------------------------------ echo -e "${STAR} ${BOLD}Detecting Device...${RESET}" -DEVICE=$(/opt/rocm/bin/rocminfo | grep "AMD Instinct" | head -n1 | awk '{print $5}') -echo -e " ${DOT} Device found: ${CYAN}$DEVICE${RESET}" +ROCMINFO="" +for candidate in \ + "$(command -v rocminfo 2>/dev/null)" \ + "/opt/rocm/bin/rocminfo" \ + "${ROCM_PATH:+$ROCM_PATH/bin/rocminfo}" \ + /opt/rocm-*/bin/rocminfo; do + if [[ -n "$candidate" && -x "$candidate" ]]; then + ROCMINFO="$candidate" + break + fi +done -if [[ -z "$DEVICE" || "$DEVICE" != "MI300X" && "$DEVICE" != "MI355X" ]]; then - ARCH=$(/opt/rocm/bin/rocminfo | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') +is_valid_device() { + local candidate="$1" + for dev in "${VALID_DEVICES[@]}"; do + if [[ "$candidate" == "$dev" ]]; then + return 0 + fi + done + return 1 +} + +if [[ -z "$ROCMINFO" ]]; then + echo -e " ${YELLOW}⚠ rocminfo not found (checked PATH, /opt/rocm/bin, ROCM_PATH)${RESET}" + DEVICE="" +else + echo -e " ${DOT} Using rocminfo: ${CYAN}$ROCMINFO${RESET}" + DEVICE=$("$ROCMINFO" 2>/dev/null | grep -oE 'MI3[0-9]{2}X' | head -n1) + if [[ -z "$DEVICE" ]]; then + DEVICE=$("$ROCMINFO" 2>/dev/null | grep "AMD Instinct" | head -n1 | awk '{print $5}') + fi + echo -e " ${DOT} Device found: ${CYAN}$DEVICE${RESET}" +fi + +if ! is_valid_device "$DEVICE"; then + if [[ -n "$ROCMINFO" ]]; then + ARCH=$("$ROCMINFO" 2>/dev/null | grep -o 'gfx942\|gfx950' | head -n 1 | tr -d '[:space:]') + else + ARCH="" + fi case "$ARCH" in - "gfx942") DEVICE="MI300X" ;; "gfx950") DEVICE="MI355X" ;; + # gfx942 is shared by MI300X and MI325X; require manual selection if marketing name is missing *) DEVICE="" ;; esac fi @@ -101,7 +349,8 @@ if [[ -z "$DEVICE" ]]; then echo -e "${RED}✗ Could not detect device automatically${RESET}" echo -e "${STAR} ${BOLD}Please select Device manually:${RESET}" echo -e " ${DOT} 1) MI300X" - echo -e " ${DOT} 2) MI355X" + echo -e " ${DOT} 2) MI325X" + echo -e " ${DOT} 3) MI355X" echo -en " ${ARROW} Enter number or name: " read -r DEV_IN @@ -110,7 +359,10 @@ if [[ -z "$DEVICE" ]]; then 1|MI300X|mi300x|Mi300x) DEVICE="MI300X" ;; - 2|MI355X|mi355x|Mi355x) + 2|MI325X|mi325x|Mi325x) + DEVICE="MI325X" + ;; + 3|MI355X|mi355x|Mi355x) DEVICE="MI355X" ;; *) @@ -175,10 +427,10 @@ SELECTED_CONFIGS=() if [[ "$CFG_NUM" == "all" ]]; then # Select all configs SELECTED_CONFIGS=("${CONFIG_LIST[@]}") -elif [[ "$CFG_NUM" =~ ^([0-9]+)-([0-9]+)$ ]]; then +elif [[ "$CFG_NUM" =~ ^[0-9]+-[0-9]+$ ]]; then # Handle range input like 4-8 - START="${BASH_REMATCH[1]}" - END="${BASH_REMATCH[2]}" + START="${CFG_NUM%%-*}" + END="${CFG_NUM##*-}" if [[ $START -lt 1 || $END -gt ${#CONFIG_LIST[@]} || $START -gt $END ]]; then echo -e "${RED}✗ Invalid range: $START-$END${RESET}" @@ -190,7 +442,9 @@ elif [[ "$CFG_NUM" =~ ^([0-9]+)-([0-9]+)$ ]]; then done else # Handle comma-separated input + _saved_ifs=$IFS IFS=',' read -ra CFG_NUMS <<< "$CFG_NUM" + IFS=$_saved_ifs for num in "${CFG_NUMS[@]}"; do # Trim whitespace @@ -205,7 +459,8 @@ else done fi -echo -e " ${CHECK} Selected ${GREEN}${#SELECTED_CONFIGS[@]}${RESET} config(s):" +SELECTED_CONFIG_COUNT=${#SELECTED_CONFIGS[@]} +echo -e " ${CHECK} Selected ${GREEN}${SELECTED_CONFIG_COUNT}${RESET} configs:" for cfg in "${SELECTED_CONFIGS[@]}"; do echo -e " ${DOT} $(basename "$cfg")" done @@ -255,7 +510,9 @@ if [[ ${#SELECTED_CONFIGS[@]} -gt 1 ]]; then if [[ "$EDIT_SELECTION" == "all" ]]; then MODELS_TO_EDIT=("${!SELECTED_CONFIGS[@]}") else + _saved_ifs=$IFS IFS=',' read -ra EDIT_NUMS <<< "$EDIT_SELECTION" + IFS=$_saved_ifs MODELS_TO_EDIT=() for num in "${EDIT_NUMS[@]}"; do num=$(echo "$num" | xargs) @@ -277,18 +534,7 @@ if [[ ${#SELECTED_CONFIGS[@]} -gt 1 ]]; then TEMP_EDIT_CONFIG="/tmp/primus_edit_${model_name}_$$.yaml" cp "$cfg" "$TEMP_EDIT_CONFIG" - # Try to find an editor - if command -v nano &> /dev/null; then - nano "$TEMP_EDIT_CONFIG" - elif command -v vim &> /dev/null; then - vim "$TEMP_EDIT_CONFIG" - elif command -v vi &> /dev/null; then - vi "$TEMP_EDIT_CONFIG" - elif command -v code &> /dev/null; then - code --wait "$TEMP_EDIT_CONFIG" - else - ${EDITOR:-vi} "$TEMP_EDIT_CONFIG" - fi + open_config_editor "$TEMP_EDIT_CONFIG" # Store the edited config EDITED_CONFIGS["$cfg"]="$TEMP_EDIT_CONFIG" @@ -311,18 +557,7 @@ elif [[ ${#SELECTED_CONFIGS[@]} -eq 1 ]]; then TEMP_EDIT_CONFIG="/tmp/primus_edit_${model_name}_$$.yaml" cp "$cfg" "$TEMP_EDIT_CONFIG" - # Try to find an editor - if command -v nano &> /dev/null; then - nano "$TEMP_EDIT_CONFIG" - elif command -v vim &> /dev/null; then - vim "$TEMP_EDIT_CONFIG" - elif command -v vi &> /dev/null; then - vi "$TEMP_EDIT_CONFIG" - elif command -v code &> /dev/null; then - code --wait "$TEMP_EDIT_CONFIG" - else - ${EDITOR:-vi} "$TEMP_EDIT_CONFIG" - fi + open_config_editor "$TEMP_EDIT_CONFIG" # Store the edited config EDITED_CONFIGS["$cfg"]="$TEMP_EDIT_CONFIG" @@ -360,7 +595,8 @@ if [[ "$OVERRIDE_PARAMS" == "y" || "$OVERRIDE_PARAMS" == "Y" ]]; then done if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then - echo -e "\n ${CHECK} ${GREEN}${#PARAM_OVERRIDES[@]}${RESET} parameter(s) will be overridden\n" + PARAM_OVERRIDE_COUNT=${#PARAM_OVERRIDES[@]} + echo -e "\n ${CHECK} ${GREEN}${PARAM_OVERRIDE_COUNT}${RESET} parameters will be overridden\n" fi fi @@ -398,7 +634,8 @@ if [[ "$ADD_ENV_VARS" == "y" || "$ADD_ENV_VARS" == "Y" ]]; then done if [[ ${#DEVICE_ENV_VARS[@]} -gt 0 ]]; then - echo -e "\n ${CHECK} ${GREEN}${#DEVICE_ENV_VARS[@]}${RESET} environment variable(s) will be set\n" + DEVICE_ENV_VAR_COUNT=${#DEVICE_ENV_VARS[@]} + echo -e "\n ${CHECK} ${GREEN}${DEVICE_ENV_VAR_COUNT}${RESET} environment variables will be set\n" fi fi @@ -453,105 +690,30 @@ for CFG_FILE in "${SELECTED_CONFIGS[@]}"; do echo -e "${MAGENTA}${BOLD}║ CONFIG FILE: $(basename "$CFG_FILE")${RESET}" echo -e "${MAGENTA}${BOLD}╚════════════════════════════════════════════════════════════╝${RESET}\n" - # Extract full filename without extension to preserve all details - CONFIG_FILENAME=$(basename "$CFG_FILE" .yaml) - MODEL_NAME="$CONFIG_FILENAME" - TIMESTAMP=$(date +"%Y-%m-%d_%H-%M-%S") + prepare_benchmark_artifacts "$CFG_FILE" - # Use the full config filename in the log name to preserve all details - LOG_FILE="$LOG_DIR/${CONFIG_FILENAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}.log" - - # Use edited config if available, otherwise use original - if [[ -n "${EDITED_CONFIGS[$CFG_FILE]}" ]]; then - WORKING_CONFIG="${EDITED_CONFIGS[$CFG_FILE]}" - echo -e "${INFO} ${BOLD}Using edited config for ${CYAN}$MODEL_NAME${RESET}" - else - WORKING_CONFIG="$CFG_FILE" + if [[ -n "${EDITED_CONFIGS[$CFG_FILE]:-}" && ${#PARAM_OVERRIDES[@]} -eq 0 ]]; then + echo -e "${INFO} ${BOLD}Using edited config for ${CYAN}$PREP_MODEL_NAME${RESET}" fi - - # Apply parameter overrides if any + echo -e " ${DOT} Run label: ${CYAN}$PREP_RUN_LABEL${RESET}" if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then - OVERRIDE_CONFIG="$LOG_DIR/${MODEL_NAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}_override.yaml" - cp "$WORKING_CONFIG" "$OVERRIDE_CONFIG" - echo -e "${STAR} ${BOLD}Applying parameter overrides...${RESET}" for KEY in "${!PARAM_OVERRIDES[@]}"; do - VALUE="${PARAM_OVERRIDES[$KEY]}" - # Use sed to replace the parameter value in YAML (handles both 'key: value' and 'key:value') - sed -i "s|^\([[:space:]]*${KEY}:[[:space:]]*\).*|\1${VALUE}|g" "$OVERRIDE_CONFIG" - echo -e " ${DOT} ${CYAN}$KEY${RESET}: $VALUE" + echo -e " ${DOT} ${CYAN}$KEY${RESET}: ${PARAM_OVERRIDES[$KEY]}" done - WORKING_CONFIG="$OVERRIDE_CONFIG" - echo -e " ${CHECK} Override config saved: ${YELLOW}$OVERRIDE_CONFIG${RESET}\n" - elif [[ -n "${EDITED_CONFIGS[$CFG_FILE]}" ]]; then - # Save edited config to logs directory - SAVED_CONFIG="$LOG_DIR/${MODEL_NAME}_${BACKEND}_${DEVICE}_${TIMESTAMP}_edited.yaml" - cp "$WORKING_CONFIG" "$SAVED_CONFIG" - WORKING_CONFIG="$SAVED_CONFIG" + echo -e " ${CHECK} Override config saved: ${YELLOW}$PREP_WORKING_CONFIG${RESET}\n" fi - echo -e "${STAR} ${BOLD}Starting Benchmark ${CURRENT}/${TOTAL_CONFIGS}...${RESET}" - echo -e " ${DOT} Model: ${CYAN}$MODEL_NAME${RESET}" - echo -e " ${DOT} Backend: ${CYAN}$BACKEND${RESET}" - echo -e " ${DOT} Device: ${CYAN}$DEVICE${RESET}" - echo -e " ${DOT} Config: ${YELLOW}$WORKING_CONFIG${RESET}" - echo -e " ${DOT} Log: ${YELLOW}$LOG_FILE${RESET}\n" - - # Set EXP to the working config (edited or overridden version if available) - # For edited/overridden configs, we need to copy them to the expected location - if [[ "$WORKING_CONFIG" != "$CFG_FILE" ]]; then - # Config was edited or overridden, copy it to the original location temporarily - ORIGINAL_CONFIG_BACKUP="${CFG_FILE}.backup_$$" - cp "$CFG_FILE" "$ORIGINAL_CONFIG_BACKUP" - cp "$WORKING_CONFIG" "$CFG_FILE" - echo -e " ${CHECK} Copied edited/overridden config to: ${CYAN}$CFG_FILE${RESET}" - fi - - # Set EXP to the device-specific config path (now contains edited content if applicable) - EXP_CONFIG_PATH="${BACKEND_BASE_DIR}/${DEVICE}/$(basename "$CFG_FILE")" - export EXP="$EXP_CONFIG_PATH" - echo -e " ${CHECK} EXP set to: ${CYAN}$EXP${RESET}\n" - - # Change to Primus root directory before running the script - echo -e " ${DOT} Changing to Primus root directory: ${CYAN}$PRIMUS_ROOT${RESET}" - cd "$PRIMUS_ROOT" - - # Run the script and capture exit code, but don't stop on failure - # Use 'set +e' locally to ensure we continue even on errors - set +e - bash $RUN_SCRIPT 2>&1 | tee "$LOG_FILE" || true - RUN_EXIT_CODE=$? - set -e - - # Return to script directory - cd "$SCRIPT_DIR" - - # Restore original config if it was temporarily replaced - if [[ -n "$ORIGINAL_CONFIG_BACKUP" && -f "$ORIGINAL_CONFIG_BACKUP" ]]; then - mv "$ORIGINAL_CONFIG_BACKUP" "$CFG_FILE" - echo -e " ${CHECK} Restored original config file" - unset ORIGINAL_CONFIG_BACKUP - fi - - echo - echo -e "${GREEN}==========================================${RESET}" - if [[ $RUN_EXIT_CODE -eq 0 ]]; then - echo -e " ${BOLD}${GREEN}✓ Benchmark ${CURRENT}/${TOTAL_CONFIGS} Completed Successfully!${RESET}" - else - echo -e " ${BOLD}${YELLOW}⚠ Benchmark ${CURRENT}/${TOTAL_CONFIGS} Completed with Exit Code: $RUN_EXIT_CODE${RESET}" - fi - echo -e " Log saved at:" - echo -e " ${CYAN}$LOG_FILE${RESET}" - if [[ ${#PARAM_OVERRIDES[@]} -gt 0 ]]; then - echo -e " Override config saved at:" - echo -e " ${CYAN}$OVERRIDE_CONFIG${RESET}" - fi - echo -e "${GREEN}==========================================${RESET}" - echo + execute_benchmark_run \ + "$PREP_CFG_FILE" \ + "$PREP_WORKING_CONFIG" \ + "$PREP_LOG_FILE" \ + "$PREP_MODEL_NAME" \ + "$CURRENT" \ + "$TOTAL_CONFIGS" || true CURRENT=$((CURRENT + 1)) - # Add a short delay between runs if [[ $CURRENT -le $TOTAL_CONFIGS ]]; then echo -e "${YELLOW}Preparing next benchmark...${RESET}\n" echo -e "${INFO} ${BOLD}Next: Config ${CURRENT}/${TOTAL_CONFIGS}${RESET}\n" @@ -561,26 +723,10 @@ done echo echo -e "${MAGENTA}${BOLD}=========================================${RESET}" -echo -e "${MAGENTA}${BOLD} All ${TOTAL_CONFIGS} Benchmark(s) Completed!${RESET}" +echo -e "${MAGENTA}${BOLD} All ${TOTAL_CONFIGS} benchmarks completed!${RESET}" echo -e "${MAGENTA}${BOLD}=========================================${RESET}" # ------------------------------------------ # 7. GENERATE METRICS TABLE # ------------------------------------------ -echo -echo -e "${STAR} ${BOLD}Generating Metrics Table...${RESET}\n" - -if [[ "$BACKEND" == "megatron" ]]; then - METRICS_SCRIPT="metrics_megatron.py" -elif [[ "$BACKEND" == "torchtitan" ]]; then - METRICS_SCRIPT="metrics_torchtitan.py" -fi - -if [[ -f "$METRICS_SCRIPT" ]]; then - echo -e " ${CHECK} Running: ${CYAN}python $METRICS_SCRIPT${RESET}\n" - python "$METRICS_SCRIPT" - echo - echo -e " ${CHECK} ${GREEN}Metrics table generated successfully${RESET}" -else - echo -e " ${RED}✗ Metrics script not found: $METRICS_SCRIPT${RESET}" -fi +generate_metrics_table