From c21d05b3b51db333a48ead9844e04c503a9c2d61 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Jul 2026 01:37:15 +0000 Subject: [PATCH 1/4] Add CPU-only training simulator --- .../local_benchmark/training_sim/README.md | 185 +++ .../local_benchmark/training_sim/__init__.py | 1 + .../training_sim/benchmark_behavior.py | 1026 +++++++++++++++++ .../training_sim/calibration_evaluator.py | 256 ++++ .../training_sim/collect_calibration.py | 224 ++++ .../training_sim/config_fingerprint.py | 213 ++++ .../training_sim/k8s/README.md | 13 + .../training_sim/memory_ledger.py | 246 ++++ .../training_sim/model_metadata.py | 255 ++++ .../local_benchmark/training_sim/predict.py | 124 ++ .../training_sim/scenario_planner.py | 1024 ++++++++++++++++ .../local_benchmark/training_sim/schemas.py | 322 ++++++ .../training_sim/shape_engine.py | 109 ++ .../training_sim/tradeoff_ranker.py | 202 ++++ .../training_sim/validate_benchmarks.py | 457 ++++++++ tests/experiments/test_training_sim.py | 624 ++++++++++ 16 files changed, 5281 insertions(+) create mode 100644 experiments/local_benchmark/training_sim/README.md create mode 100644 experiments/local_benchmark/training_sim/__init__.py create mode 100644 experiments/local_benchmark/training_sim/benchmark_behavior.py create mode 100644 experiments/local_benchmark/training_sim/calibration_evaluator.py create mode 100644 experiments/local_benchmark/training_sim/collect_calibration.py create mode 100644 experiments/local_benchmark/training_sim/config_fingerprint.py create mode 100644 experiments/local_benchmark/training_sim/k8s/README.md create mode 100644 experiments/local_benchmark/training_sim/memory_ledger.py create mode 100644 experiments/local_benchmark/training_sim/model_metadata.py create mode 100644 experiments/local_benchmark/training_sim/predict.py create mode 100644 experiments/local_benchmark/training_sim/scenario_planner.py create mode 100644 experiments/local_benchmark/training_sim/schemas.py create mode 100644 experiments/local_benchmark/training_sim/shape_engine.py create mode 100644 experiments/local_benchmark/training_sim/tradeoff_ranker.py create mode 100644 experiments/local_benchmark/training_sim/validate_benchmarks.py create mode 100644 tests/experiments/test_training_sim.py diff --git a/experiments/local_benchmark/training_sim/README.md b/experiments/local_benchmark/training_sim/README.md new file mode 100644 index 00000000..ca8d3139 --- /dev/null +++ b/experiments/local_benchmark/training_sim/README.md @@ -0,0 +1,185 @@ +# XoRL Training-Engine Simulator + +This is a CPU-only first slice of the local training-engine simulator. It resolves the launch topology from a +XoRL YAML config, computes deterministic balanced-routing token shapes, estimates a sharded persistent model-state +memory floor, and parses structured trainer logs into calibration summaries. + +It does not yet model activation, attention workspace, MoE kernel workspace, FSDP transient, or allocator slack +memory. Those are left as explicit `unsupported_buckets` in the JSON report until calibrated formulas are added. + +## Predict From A Config + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.predict \ + --config path/to/xorl_config.yaml \ + --world-size 16 \ + --local-world-size 8 \ + --balanced-routing +``` + +## Add Log Calibration + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.predict \ + --config path/to/xorl_config.yaml \ + --world-size 16 \ + --local-world-size 8 \ + --balanced-routing \ + --num-experts 128 \ + --top-k 8 \ + --logs /shared/path/to/trainer-head/logs/run.log \ + --warmup-steps 3 \ + --output experiments/local_benchmark/training_sim/calibration/report.json +``` + +Pass `--benchmark-dir` to include empirical behavior calibrated from a recipe README and result JSON: + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.predict \ + --config path/to/benchmark/configs/xorl_cli.yaml \ + --balanced-routing \ + --benchmark-dir path/to/benchmark +``` + +Empirical matches are config-specific. The simulator checks both topology/workload shape and runtime knobs such as +`deepep_async_combine`, `deepep_num_sms`, `deepep_buffer_size_gb`, `enable_compile`, +`gradient_checkpointing_method`, activation offload, and prefetch count before treating a benchmark row as an exact +calibration point. + +`--benchmark-dir` can also point at a results root containing resolved run directories. Any subdirectory with +`xorl_cli.yaml` is treated as a candidate calibration source. If a matching `node-0.log` is available directly beside +the config or through `startup_metrics.json`'s `startup/master_addr`, the loader parses measured `[STEP ...]` rows +with two warmup steps excluded. OOM logs become calibrated failure boundaries, and runs that report throughput before +crashing are kept as partial-failure calibration points rather than clean promotion candidates. + +## Rank Benchmark Tradeoffs + +Use the tradeoff ranker to compare autotune rows in a benchmark folder. It keeps the fastest raw +candidate separate from the fastest correctness-promotable candidate, so a raw-speed win is not promoted unless it +has a matching `k3_pass` gate. + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.tradeoff_ranker \ + path/to/benchmark +``` + +The report keeps a faster raw candidate separate from a slower promotable candidate when only the latter has a +matching correctness gate. + +## Plan What-If Scenarios + +Use the scenario planner when the question is not just "what already won?" but "what should we try next?". It +mutates a base config across micro-batch and parallelism choices, computes a topology and sharded model-state memory +floor for each candidate, then ranks exact calibrated matches ahead of lower-confidence extrapolations. Extrapolated +candidates are never marked promotable; they need a fresh K3 gate. + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.scenario_planner \ + --config path/to/benchmark/configs/xorl_cli.yaml \ + --benchmark-dir path/to/benchmark \ + --micro-batch-sizes 5 \ + --expert-parallel-sizes 32,64 +``` + +The planner compares concrete parallelism tradeoffs while preserving correctness, runtime-compatibility, and +memory-feasibility caveats. + +For wider topology searches, add `--topology-sweep auto`. Auto mode derives legal candidate values for EP, TP, PP, +Ulysses, and Ring from the resolved world size and model metadata, while explicit comma lists still override any +individual dimension. Exact empirical matches are conservative: an observation only matches TP/PP/Ulysses/Ring values +known from that artifact, and legacy artifacts with missing topology dimensions only exact-match the default value of +1. Non-default TP/PP/CP candidates therefore remain extrapolated unless there is a measured row for that exact +topology. + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.scenario_planner \ + --config path/to/benchmark/configs/xorl_cli.yaml \ + --benchmark-dir path/to/benchmark \ + --gradient-accumulation-steps 1,2,4,8 \ + --topology-sweep auto +``` + +The planner also understands markdown result tables with `tok/s tot`, `tok/step`, and `peak GB` columns, such as the +Qwen3-235B 2k-context sweep. Observed peak memory overrides the analytic floor for feasibility checks, and OOM rows +are kept as calibrated failures when their topology and pack length match a scenario. When two or more +global-batch/GA points are calibrated for the same topology, it fits a simple +`step_time = fixed_overhead + token_slope * tokens` curve for larger GA what-ifs: + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.scenario_planner \ + --config path/to/benchmark/configs/xorl_cli.yaml \ + --benchmark-dir path/to/benchmark \ + --micro-batch-sizes 1 \ + --gradient-accumulation-steps 1,2,4,8 \ + --expert-parallel-sizes 8 +``` + +For that q235 scenario, GA2 is a calibrated measured point and GA4/GA8 are step-time-fit extrapolations. They are +useful next-run candidates, not correctness-promotable results. + +Planner candidates include `calibration_scope` and `risk_flags` fields. `exact_calibrated` means an empirical row +matched the full scenario topology. `inside_measured_envelope` and `outside_measured_envelope` describe whether an +extrapolated candidate stays inside the measured micro-batch/global-batch/parallelism range for that sequence length. +Risk flags call out cases like `requires_remeasurement`, `matched_allocator_pressure_slowdown`, an +`allocator_pressure_boundary:*`, an `observed_oom_boundary:*`, `correctness_runtime_failure_after_steps`, or +`runtime_mismatch:*` when extrapolation had to fall back to a row with different runtime knobs. Treat these flags as +launch-planning constraints: they do not erase the raw score, but they mean the row needs a fresh measurement or debug +pass before it can be used as an optimum. + +Scenario reports keep `best_raw` as the fastest feasible throughput hypothesis, then add `best_risk_adjusted` and +`best_next_measurement`. The risk-adjusted score penalizes extrapolation, memory pressure, missing correctness gates, +allocator-pressure slowdowns, and observed-OOM boundaries. This makes the planner useful as an optimizer loop: launch +the best next measurement when it is a hypothesis, but prefer the risk-adjusted or promotable row when choosing what is +already defensible. + +For exact calibrated rows, an observed peak below the device limit remains feasible even when the configured safety +factor would reserve slightly more than the device capacity. Those rows are marked `feasible_calibrated_peak_high_pressure`. +The safety margin still gates extrapolated peaks and analytic floors. + +## Validate Prediction Fidelity + +Use the calibration evaluator before trusting a scenario sweep as an optimizer. It runs leave-one-out validation over +measured benchmark rows, rebuilds the held-out topology, predicts it from the remaining calibration points, and reports +actual-vs-predicted error. + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.calibration_evaluator \ + --config path/to/benchmark/configs/xorl_cli.yaml \ + --benchmark-dir path/to/benchmark +``` + +Treat a large holdout error as a sign that the relevant lever needs a new calibration point or a more specific +simulator feature before promotion decisions rely on it. + +## Parse Logs Only + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.collect_calibration \ + /shared/path/to/trainer-head/logs/run.log \ + --warmup-steps 3 \ + --world-size 16 +``` + +The log parser recognizes `[STEP ...]`, `[STEP_PHASES ...]`, and `[STEP_MEMORY ...]` lines emitted by +`src/xorl/trainers/trainer.py`. + +## Validate Checked-In Benchmarks + +```bash +cd "$(git rev-parse --show-toplevel)" +python -m experiments.local_benchmark.training_sim.validate_benchmarks \ + --benchmarks-root path/to/benchmarks \ + --model benchmark_name +``` + +The validator checks benchmark YAML, README target metrics, synthetic-routing render scripts, stored throughput +summaries, and static-K3 gate status when those artifacts are present. diff --git a/experiments/local_benchmark/training_sim/__init__.py b/experiments/local_benchmark/training_sim/__init__.py new file mode 100644 index 00000000..fd63eca3 --- /dev/null +++ b/experiments/local_benchmark/training_sim/__init__.py @@ -0,0 +1 @@ +"""Local training-engine simulator helpers for benchmark planning.""" diff --git a/experiments/local_benchmark/training_sim/benchmark_behavior.py b/experiments/local_benchmark/training_sim/benchmark_behavior.py new file mode 100644 index 00000000..d9fd6462 --- /dev/null +++ b/experiments/local_benchmark/training_sim/benchmark_behavior.py @@ -0,0 +1,1026 @@ +"""Empirical benchmark behavior calibration for checked-in benchmark recipes.""" + +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path +from typing import Any + + +try: + from .collect_calibration import parse_log_path, summarize_observed_run + from .config_fingerprint import load_training_config, resolve_topology + from .schemas import BenchmarkBehaviorPoint, BenchmarkBehaviorPrediction, ShapeLedger, Topology, to_jsonable +except ImportError: # pragma: no cover - exercised by direct script execution + from collect_calibration import parse_log_path, summarize_observed_run + from config_fingerprint import load_training_config, resolve_topology + from schemas import BenchmarkBehaviorPoint, BenchmarkBehaviorPrediction, ShapeLedger, Topology, to_jsonable + + +H100_BF16_PROMISED_TFLOPS_PER_GPU = 989.0 + + +def _gpu_count_from_text(text: str) -> int | None: + match = re.search(r"(?P\d+)\s+nodes?\s+x\s+(?P\d+)\s+H100", text, re.IGNORECASE) + if match: + return int(match.group("nodes")) * int(match.group("gpus")) + match = re.search(r"(?P\d+)\s*[x×]\s*H100", text, re.IGNORECASE) + if match: + return int(match.group("gpus")) + match = re.search(r"(?P\d+)\s+GPUs?", text, re.IGNORECASE) + return int(match.group("gpus")) if match else None + + +def human_number(value: str) -> float: + cleaned = value.strip().replace(",", "").lstrip("~") + multiplier = 1.0 + if cleaned.endswith(("K", "k")): + cleaned = cleaned[:-1] + multiplier = 1_000.0 + elif cleaned.endswith(("M", "m")): + cleaned = cleaned[:-1] + multiplier = 1_000_000.0 + return float(cleaned) * multiplier + + +def _first_non_none(*values: Any) -> Any: + for value in values: + if value is not None: + return value + return None + + +def _readme_point(readme_text: str, *, source: str) -> BenchmarkBehaviorPoint | None: + tps_match = re.search(r"\|\s*tokens/sec\s*\|\s*(?P~?[0-9.]+[KkMm]?)\s*\|", readme_text) + step_match = re.search(r"\|\s*step time\s*\|\s*(?P~?[0-9.]+)s\s*\|", readme_text) + mfu_match = re.search(r"\|\s*MFU\s*\|\s*(?P~?[0-9.]+)%", readme_text) + memory_match = re.search(r"\|\s*allocated memory\s*\|\s*(?P~?[0-9.]+)GB\s*\|", readme_text) + retries_match = re.search(r"\|\s*allocator retries\s*\|\s*(?P\d+)\s*\|", readme_text) + mbs_match = re.search(r"micro_batch_size:\s*(?P\d+)", readme_text) + global_batch_match = re.search(r"global_batch_size:\s*(?P\d+)", readme_text) + if not tps_match: + return None + return BenchmarkBehaviorPoint( + label="readme_reference_mbs8", + source=source, + micro_batch_size=int(mbs_match.group("value")) if mbs_match else None, + global_batch_size=int(global_batch_match.group("value")) if global_batch_match else None, + tokens_per_sec=human_number(tps_match.group("value")), + step_time_sec=float(step_match.group("value").lstrip("~")) if step_match else None, + mfu_percent=float(mfu_match.group("value").lstrip("~")) if mfu_match else None, + tflops_per_gpu=None, + peak_mem_gb=float(memory_match.group("value").lstrip("~")) if memory_match else None, + allocator_retries=int(retries_match.group("value")) if retries_match else None, + gpu_count=_gpu_count_from_text(readme_text), + sample_packing_sequence_len=_seq_len_from_readme(readme_text), + tensor_parallel_size=_readme_parallel_int( + readme_text, + "tensor_parallel_size", + (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), + ), + pipeline_parallel_size=_readme_parallel_int( + readme_text, + "pipeline_parallel_size", + (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), + ), + ulysses_parallel_size=_readme_parallel_int( + readme_text, + "ulysses_parallel_size", + (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), + ), + ringattn_parallel_size=_readme_parallel_int( + readme_text, + "ringattn_parallel_size", + (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), + ), + expert_parallel_size=_readme_parallel_int( + readme_text, + "expert_parallel_size", + (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), + ), + ep_fsdp_size=_readme_parallel_int( + readme_text, + "ep_fsdp_size", + (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), + ), + deepep_async_combine=_readme_bool_from_text(readme_text, "deepep_async_combine"), + deepep_num_sms=_readme_parallel_int( + readme_text, + "deepep_num_sms", + (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), + ), + deepep_buffer_size_gb=_readme_float_from_text(readme_text, "deepep_buffer_size_gb"), + enable_compile=_readme_bool_from_text(readme_text, "enable_compile"), + gradient_checkpointing_method=_checkpointing_method_from_text(readme_text), + status="reference_speed", + correctness_status="raw_speed_not_promoted_without_matching_k3_pass", + notes=["current-main logical FLOPs accounting", "balanced synthetic routing", "deepep_async_combine true"], + ) + + +def _readme_adjacent_mbs10_point( + readme_text: str, *, source: str, seq_len: int | None +) -> BenchmarkBehaviorPoint | None: + match = re.search(r"`mbs=10`[^~]+~(?P[0-9.]+)K tok/s", readme_text) + if not match: + return None + tokens_per_sec = human_number(match.group("value") + "K") + global_batch_size = 320 + step_time_sec = (global_batch_size * seq_len / tokens_per_sec) if seq_len else None + return BenchmarkBehaviorPoint( + label="readme_adjacent_mbs10_allocator_pressure", + source=source, + micro_batch_size=10, + global_batch_size=global_batch_size, + tokens_per_sec=tokens_per_sec, + step_time_sec=step_time_sec, + mfu_percent=None, + tflops_per_gpu=None, + peak_mem_gb=None, + allocator_retries=None, + gpu_count=_gpu_count_from_text(readme_text), + sample_packing_sequence_len=seq_len, + tensor_parallel_size=_readme_parallel_int( + readme_text, + "tensor_parallel_size", + (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), + ), + pipeline_parallel_size=_readme_parallel_int( + readme_text, + "pipeline_parallel_size", + (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), + ), + ulysses_parallel_size=_readme_parallel_int( + readme_text, + "ulysses_parallel_size", + (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), + ), + ringattn_parallel_size=_readme_parallel_int( + readme_text, + "ringattn_parallel_size", + (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), + ), + expert_parallel_size=_readme_parallel_int( + readme_text, + "expert_parallel_size", + (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), + ), + ep_fsdp_size=_readme_parallel_int( + readme_text, + "ep_fsdp_size", + (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), + ), + deepep_async_combine=_readme_bool_from_text(readme_text, "deepep_async_combine"), + deepep_num_sms=_readme_parallel_int( + readme_text, + "deepep_num_sms", + (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), + ), + deepep_buffer_size_gb=_readme_float_from_text(readme_text, "deepep_buffer_size_gb"), + enable_compile=_readme_bool_from_text(readme_text, "enable_compile"), + gradient_checkpointing_method=_checkpointing_method_from_text(readme_text), + status="allocator_pressure_slowdown", + correctness_status="not_promoted", + notes=["fit but slowed with allocator retries"], + ) + + +def _result_throughput_point( + result_path: Path, + result: dict[str, Any], + *, + topology_defaults: dict[str, int | float | bool | str], +) -> BenchmarkBehaviorPoint: + throughput = result["throughput"] + candidate = ( + throughput.get("candidate") + or result.get("candidate") + or (result.get("replay_candidate", {}) if isinstance(result.get("replay_candidate"), dict) else {}).get( + "candidate" + ) + or "throughput" + ) + return BenchmarkBehaviorPoint( + label=f"{result_path.stem}:{candidate}", + source=str(result_path), + micro_batch_size=throughput.get("micro_batch_size"), + global_batch_size=throughput.get("global_batch_size"), + tokens_per_sec=throughput.get("tokens_per_sec"), + step_time_sec=throughput.get("step_time_sec"), + mfu_percent=throughput.get("mfu_percent"), + tflops_per_gpu=throughput.get("mean_tflops_per_gpu"), + peak_mem_gb=throughput.get("gpu_alloc_gb"), + allocator_retries=None, + measured_steps=throughput.get("measured_steps"), + warmup_steps=throughput.get("warmup_steps"), + gpu_count=throughput.get("gpus"), + sample_packing_sequence_len=throughput.get("sample_packing_sequence_len"), + tensor_parallel_size=_first_non_none( + throughput.get("tensor_parallel_size"), topology_defaults.get("tensor_parallel_size") + ), + pipeline_parallel_size=_first_non_none( + throughput.get("pipeline_parallel_size"), topology_defaults.get("pipeline_parallel_size") + ), + ulysses_parallel_size=_first_non_none( + throughput.get("ulysses_parallel_size"), topology_defaults.get("ulysses_parallel_size") + ), + ringattn_parallel_size=_first_non_none( + throughput.get("ringattn_parallel_size"), topology_defaults.get("ringattn_parallel_size") + ), + expert_parallel_size=_first_non_none( + throughput.get("expert_parallel_size"), topology_defaults.get("expert_parallel_size") + ), + ep_fsdp_size=_first_non_none( + throughput.get("ep_fsdp"), throughput.get("ep_fsdp_size"), topology_defaults.get("ep_fsdp_size") + ), + deepep_async_combine=_first_non_none( + throughput.get("deepep_async_combine"), topology_defaults.get("deepep_async_combine") + ), + deepep_num_sms=_first_non_none(throughput.get("deepep_num_sms"), topology_defaults.get("deepep_num_sms")), + deepep_buffer_size_gb=_first_non_none( + throughput.get("deepep_buffer_size_gb"), topology_defaults.get("deepep_buffer_size_gb") + ), + enable_compile=_first_non_none(throughput.get("enable_compile"), topology_defaults.get("enable_compile")), + gradient_checkpointing_method=_first_non_none( + throughput.get("gradient_checkpointing_method"), topology_defaults.get("gradient_checkpointing_method") + ), + enable_activation_offload=_first_non_none( + throughput.get("enable_activation_offload"), topology_defaults.get("enable_activation_offload") + ), + activation_offload_prefetch_count=_first_non_none( + throughput.get("activation_offload_prefetch_count"), + topology_defaults.get("activation_offload_prefetch_count"), + ), + status="historical_throughput_artifact", + correctness_status=None, + notes=[f"commit={throughput.get('commit')}"] if throughput.get("commit") else [], + ) + + +def _with_k3_status(point: BenchmarkBehaviorPoint, result: dict[str, Any]) -> BenchmarkBehaviorPoint: + k3_gate = result.get("k3_gate", {}) + if not k3_gate or k3_gate.get("candidate") not in (None, point.label.split(":", 1)[-1]): + return point + notes = list(point.notes) + if k3_gate.get("primary_failure"): + notes.append(f"k3_primary_failure={k3_gate['primary_failure']}") + return BenchmarkBehaviorPoint( + label=point.label, + source=point.source, + micro_batch_size=point.micro_batch_size, + global_batch_size=point.global_batch_size, + tokens_per_sec=point.tokens_per_sec, + step_time_sec=point.step_time_sec, + mfu_percent=point.mfu_percent, + tflops_per_gpu=point.tflops_per_gpu, + peak_mem_gb=point.peak_mem_gb, + allocator_retries=point.allocator_retries, + measured_steps=point.measured_steps, + warmup_steps=point.warmup_steps, + gpu_count=point.gpu_count, + sample_packing_sequence_len=point.sample_packing_sequence_len, + tensor_parallel_size=point.tensor_parallel_size, + pipeline_parallel_size=point.pipeline_parallel_size, + ulysses_parallel_size=point.ulysses_parallel_size, + ringattn_parallel_size=point.ringattn_parallel_size, + expert_parallel_size=point.expert_parallel_size, + ep_fsdp_size=point.ep_fsdp_size, + deepep_async_combine=point.deepep_async_combine, + deepep_num_sms=point.deepep_num_sms, + deepep_buffer_size_gb=point.deepep_buffer_size_gb, + enable_compile=point.enable_compile, + gradient_checkpointing_method=point.gradient_checkpointing_method, + enable_activation_offload=point.enable_activation_offload, + activation_offload_prefetch_count=point.activation_offload_prefetch_count, + status=point.status, + correctness_status=f"k3_{k3_gate.get('status')}", + notes=notes, + ) + + +def _seq_len_from_readme(readme_text: str) -> int | None: + match = re.search(r"sample_packing_sequence_len:\s*(?P\d+)", readme_text) + return int(match.group("seq")) if match else None + + +def _config_int_from_text(text: str, key: str) -> int | None: + match = re.search(rf"{re.escape(key)}:\s*(?P\d+)", text) + return int(match.group("value")) if match else None + + +def _readme_float_from_text(text: str, key: str) -> float | None: + match = re.search(rf"{re.escape(key)}:\s*(?P\d+(?:\.\d+)?)", text) + return float(match.group("value")) if match else None + + +def _readme_bool_from_text(text: str, key: str) -> bool | None: + match = re.search(rf"{re.escape(key)}:\s*(?Ptrue|false)", text, re.IGNORECASE) + if not match: + return None + return match.group("value").lower() == "true" + + +def _checkpointing_method_from_text(text: str) -> str | None: + lowered = text.lower() + if "recompute_before_dispatch" in lowered or "before_dispatch" in lowered: + return "recompute_before_dispatch" + if "recompute_full_layer" in lowered or "full-layer recompute" in lowered or "fullrecompute" in lowered: + return "recompute_full_layer" + if "no_recompute" in lowered or "no recompute" in lowered: + return "no_recompute" + return None + + +def _trial_checkpointing_method(trial: str) -> str | None: + return _checkpointing_method_from_text(trial) + + +def _trial_activation_offload(trial: str) -> bool | None: + if "noactivationoffload" in trial: + return False + if "activationoffload" in trial: + return True + return None + + +def _trial_prefetch_count(trial: str) -> int | None: + match = re.search(r"prefetch(?P\d+)", trial) + return int(match.group("value")) if match else None + + +def _trial_compile_enabled(trial: str) -> bool | None: + if "nocompile" in trial: + return False + if "compile" in trial: + return True + return None + + +def _trial_deepep_async_combine(trial: str) -> bool | None: + if "noasync" in trial: + return False + if "async" in trial: + return True + return None + + +def _trial_sms_count(trial: str) -> int | None: + match = re.search(r"sms(?P\d+)", trial) + return int(match.group("value")) if match else None + + +def _trial_buffer_size_gb(trial: str) -> float | None: + match = re.search(r"buf(?P\d+)", trial) + if not match: + return None + raw = match.group("value") + if len(raw) == 1: + return float(raw) + return float(f"{raw[:-1]}.{raw[-1]}") + + +def _last_regex_int(line: str, patterns: tuple[str, ...]) -> int | None: + value = None + for pattern in patterns: + for match in re.finditer(pattern, line, re.IGNORECASE): + groupdict = match.groupdict() + for key in ("value", "tp", "pp", "u", "ring"): + if groupdict.get(key) is not None: + value = int(groupdict[key]) + break + return value + + +def _readme_parallel_int(readme_text: str, config_key: str, patterns: tuple[str, ...]) -> int | None: + if value := _config_int_from_text(readme_text, config_key): + return value + for line in readme_text.splitlines(): + if value := _last_regex_int(line, patterns): + return value + return None + + +def _readme_topology_defaults(readme_text: str) -> dict[str, int | float | bool | str]: + defaults: dict[str, int | float | bool | str] = {} + field_patterns = { + "tensor_parallel_size": (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), + "pipeline_parallel_size": (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), + "ulysses_parallel_size": (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), + "ringattn_parallel_size": (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), + "expert_parallel_size": (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), + "ep_fsdp_size": (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), + } + for field, patterns in field_patterns.items(): + if value := _readme_parallel_int(readme_text, field, patterns): + defaults[field] = value + if (value := _readme_bool_from_text(readme_text, "deepep_async_combine")) is not None: + defaults["deepep_async_combine"] = value + if value := _readme_parallel_int( + readme_text, + "deepep_num_sms", + (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), + ): + defaults["deepep_num_sms"] = value + if (value := _readme_float_from_text(readme_text, "deepep_buffer_size_gb")) is not None: + defaults["deepep_buffer_size_gb"] = value + if (value := _readme_bool_from_text(readme_text, "enable_compile")) is not None: + defaults["enable_compile"] = value + if value := _checkpointing_method_from_text(readme_text): + defaults["gradient_checkpointing_method"] = value + return defaults + + +def _first_markdown_number(value: str) -> float | None: + match = re.search(r"~?\s*(?P[0-9][0-9,.]*)(?P[KkMm]?)", value.replace("*", "")) + if not match: + return None + return human_number(match.group("value") + match.group("suffix")) + + +def _markdown_value(values: dict[str, str], key_substring: str) -> str: + for key, value in values.items(): + if key_substring in key: + return value + return "" + + +def _markdown_peak_gb(value: str) -> float | None: + if "oom" in value.lower(): + return None + return _first_markdown_number(value) + + +def _q235_markdown_points(readme_text: str, *, source: str) -> list[BenchmarkBehaviorPoint]: + if "Qwen3-235B" not in readme_text or "tok/s tot" not in readme_text: + return [] + + points: list[BenchmarkBehaviorPoint] = [] + current_header: list[str] | None = None + current_gpu_count: int | None = None + current_ep_size: int | None = None + current_ep_fsdp_size: int | None = None + current_tensor_parallel_size = 1 + current_pipeline_parallel_size = 1 + current_ulysses_parallel_size = 1 + current_ringattn_parallel_size = 1 + for line in readme_text.splitlines(): + if gpu_count := _gpu_count_from_text(line): + current_gpu_count = gpu_count + ep_matches = list(re.finditer(r"\bEP(?P\d+)\b", line)) + if ep_matches: + current_ep_size = int(ep_matches[-1].group("ep")) + efsdp_matches = list(re.finditer(r"(?:ep_fsdp|eFSDP)(?:[= ]|)(?P\d+)", line)) + if efsdp_matches: + current_ep_fsdp_size = int(efsdp_matches[-1].group("efsdp")) + if tp := _last_regex_int( + line, + ( + r"\bTP(?P\d+)\b", + r"\btensor_parallel_size[:= ]+(?P\d+)\b", + r"\btp[=_](?P\d+)\b", + ), + ): + current_tensor_parallel_size = tp + if pp := _last_regex_int( + line, + ( + r"\bPP(?P\d+)\b", + r"\bpipeline_parallel_size[:= ]+(?P\d+)\b", + r"\bpp[=_](?P\d+)\b", + ), + ): + current_pipeline_parallel_size = pp + if ulysses := _last_regex_int( + line, + ( + r"\bU(?P\d+)\b", + r"\bul[y]?sses_parallel_size[:= ]+(?P\d+)\b", + r"\bu[=_]?(?P\d+)\b", + ), + ): + current_ulysses_parallel_size = ulysses + if ringattn := _last_regex_int( + line, + ( + r"\bR(?P\d+)\b", + r"\bringattn_parallel_size[:= ]+(?P\d+)\b", + r"\bring[=_]?(?P\d+)\b", + ), + ): + current_ringattn_parallel_size = ringattn + if not line.startswith("|"): + continue + cells = [cell.strip() for cell in line.strip().strip("|").split("|")] + lowered = [cell.lower() for cell in cells] + if "run" in lowered and "tok/s tot" in lowered: + current_header = lowered + continue + if current_header is None or set(cells) == {"---"} or not cells: + continue + values = dict(zip(current_header, cells, strict=False)) + run = values.get("run", "").replace("*", "").strip("` ") + if not run or run.lower() in {"run", "-----"}: + continue + status_text = _markdown_value(values, "status") + is_failure = "oom" in status_text.lower() or "fail" in status_text.lower() + tokens_per_sec = _first_markdown_number(_markdown_value(values, "tok/s tot")) + tok_step = _first_markdown_number(_markdown_value(values, "tok/step")) + pack = _first_markdown_number(_markdown_value(values, "pack")) + if tok_step is None or pack in (None, 0): + continue + if tokens_per_sec is None and not is_failure: + continue + global_batch_size = int(round(tok_step / pack)) + step_time_sec = _first_markdown_number(_markdown_value(values, "step s")) + mfu_percent = _first_markdown_number(_markdown_value(values, "mfu")) + peak_mem_gb = _markdown_peak_gb(_markdown_value(values, "peak gb")) + points.append( + BenchmarkBehaviorPoint( + label=f"q235_markdown:{run}", + source=source, + micro_batch_size=int(_first_markdown_number(values.get("mbs", "")) or 1), + global_batch_size=global_batch_size, + tokens_per_sec=tokens_per_sec, + step_time_sec=step_time_sec, + mfu_percent=mfu_percent, + peak_mem_gb=peak_mem_gb, + allocator_retries=None, + gpu_count=current_gpu_count, + sample_packing_sequence_len=int(pack), + tensor_parallel_size=current_tensor_parallel_size, + pipeline_parallel_size=current_pipeline_parallel_size, + ulysses_parallel_size=current_ulysses_parallel_size, + ringattn_parallel_size=current_ringattn_parallel_size, + expert_parallel_size=current_ep_size, + ep_fsdp_size=current_ep_fsdp_size, + status="historical_q235_markdown_oom" if is_failure else "historical_q235_markdown", + correctness_status="oom" if is_failure else "not_promoted", + notes=[status_text] if status_text else [], + ) + ) + return points + + +def _best_by_mfu_point( + result_path: Path, + result: dict[str, Any], + row: dict[str, Any], + *, + topology_defaults: dict[str, int | float | bool | str], +) -> BenchmarkBehaviorPoint: + trial = str(row["trial"]) + caveat = row.get("caveat") + k3_gate = row.get("k3_gate") + notes = [] + if caveat: + notes.append(str(caveat)) + if k3_gate: + notes.append(str(k3_gate)) + correctness_status = None + if k3_gate and str(k3_gate).startswith("pass"): + correctness_status = "k3_pass" + elif _first_non_none(row.get("deepep_async_combine"), _trial_deepep_async_combine(trial)): + correctness_status = "raw_speed_not_promoted_without_matching_k3_pass" + return BenchmarkBehaviorPoint( + label=f"best_by_mfu:{trial}", + source=str(result_path), + micro_batch_size=row.get("micro_batch_size"), + global_batch_size=row.get("global_batch_size"), + tokens_per_sec=row.get("tokens_per_sec"), + step_time_sec=row.get("step_time_sec"), + mfu_percent=row.get("mfu_percent"), + tflops_per_gpu=row.get("mean_tflops_per_gpu"), + peak_mem_gb=None, + allocator_retries=None, + measured_steps=row.get("measured_steps"), + warmup_steps=row.get("warmup_steps"), + gpu_count=row.get("gpus") or _gpu_count_from_text(str(result.get("workload", ""))), + sample_packing_sequence_len=row.get("sample_packing_sequence_len"), + tensor_parallel_size=_first_non_none( + row.get("tensor_parallel_size"), topology_defaults.get("tensor_parallel_size") + ), + pipeline_parallel_size=_first_non_none( + row.get("pipeline_parallel_size"), topology_defaults.get("pipeline_parallel_size") + ), + ulysses_parallel_size=_first_non_none( + row.get("ulysses_parallel_size"), topology_defaults.get("ulysses_parallel_size") + ), + ringattn_parallel_size=_first_non_none( + row.get("ringattn_parallel_size"), topology_defaults.get("ringattn_parallel_size") + ), + expert_parallel_size=_first_non_none( + row.get("expert_parallel_size"), topology_defaults.get("expert_parallel_size") + ), + ep_fsdp_size=_first_non_none(row.get("ep_fsdp"), topology_defaults.get("ep_fsdp_size")), + deepep_async_combine=_first_non_none( + row.get("deepep_async_combine"), + _trial_deepep_async_combine(trial), + topology_defaults.get("deepep_async_combine"), + ), + deepep_num_sms=_first_non_none( + row.get("deepep_num_sms"), _trial_sms_count(trial), topology_defaults.get("deepep_num_sms") + ), + deepep_buffer_size_gb=_first_non_none( + row.get("deepep_buffer_size_gb"), + _trial_buffer_size_gb(trial), + topology_defaults.get("deepep_buffer_size_gb"), + ), + enable_compile=_first_non_none( + row.get("enable_compile"), _trial_compile_enabled(trial), topology_defaults.get("enable_compile") + ), + gradient_checkpointing_method=_first_non_none( + row.get("gradient_checkpointing_method"), + _trial_checkpointing_method(trial), + topology_defaults.get("gradient_checkpointing_method"), + ), + enable_activation_offload=_first_non_none( + row.get("enable_activation_offload"), + _trial_activation_offload(trial), + topology_defaults.get("enable_activation_offload"), + ), + activation_offload_prefetch_count=_first_non_none( + row.get("activation_offload_prefetch_count"), + _trial_prefetch_count(trial), + topology_defaults.get("activation_offload_prefetch_count"), + ), + status="autotune_result", + correctness_status=correctness_status, + notes=notes, + ) + + +def _load_startup_metrics(run_dir: Path) -> dict[str, Any]: + startup_path = run_dir / "startup_metrics.json" + if not startup_path.is_file(): + return {} + return json.loads(startup_path.read_text(encoding="utf-8")) + + +def _startup_master_log_path(benchmark_path: Path, startup_metrics: dict[str, Any]) -> Path | None: + metrics = startup_metrics.get("metrics", {}) + master_addr = metrics.get("startup/master_addr") + if not isinstance(master_addr, str) or not master_addr: + return None + run_name = master_addr.removesuffix("-master") + return benchmark_path / run_name / "node-0.log" + + +def _resolved_run_log_path(benchmark_path: Path, run_dir: Path, startup_metrics: dict[str, Any]) -> Path | None: + candidates = [ + run_dir / "node-0.log", + _startup_master_log_path(benchmark_path, startup_metrics), + ] + for candidate in candidates: + if candidate is not None and candidate.is_file(): + return candidate + return None + + +def _log_failure_status(text: str) -> str | None: + lowered = text.lower() + if "outofmemoryerror" in lowered or "cuda out of memory" in lowered: + return "oom" + if "childfailederror" in lowered or "traceback" in lowered: + return "runtime_failure_after_steps" + return None + + +def _oom_peak_mem_gb(text: str) -> float | None: + values = [ + float(match.group("value")) + for match in re.finditer(r"process has (?P\d+(?:\.\d+)?)\s+GiB memory in use", text) + ] + return max(values) if values else None + + +def _round_or_none(value: Any, ndigits: int) -> float | None: + return round(float(value), ndigits) if value is not None else None + + +def _resolved_run_behavior_point(benchmark_path: Path, config_path: Path) -> BenchmarkBehaviorPoint | None: + run_dir = config_path.parent + raw_config = load_training_config(config_path) + try: + topology = resolve_topology(raw_config) + except ValueError: + return None + + startup_metrics = _load_startup_metrics(run_dir) + log_path = _resolved_run_log_path(benchmark_path, run_dir, startup_metrics) + log_text = log_path.read_text(encoding="utf-8", errors="replace") if log_path is not None else "" + failure_status = _log_failure_status(log_text) + observed_summary: dict[str, Any] = {} + if log_path is not None: + observed = parse_log_path(log_path) + warmup_steps = 2 if len(observed.steps) > 2 else 0 + observed_summary = summarize_observed_run(observed, warmup_steps=warmup_steps, world_size=topology.world_size) + + tokens_per_sec = _round_or_none(observed_summary.get("tokens_per_sec_mean"), 3) + peak_mem_gb = _round_or_none(observed_summary.get("peak_mem_gb_max"), 3) + if peak_mem_gb is None and failure_status == "oom": + peak_mem_gb = _round_or_none(_oom_peak_mem_gb(log_text), 3) + if tokens_per_sec is None and failure_status is None: + return None + + if failure_status == "oom" and tokens_per_sec is None: + status = "observed_log_oom" + correctness_status = "oom" + elif failure_status is not None: + status = "observed_log_partial_failure" + correctness_status = failure_status + else: + status = "observed_log_summary" + correctness_status = "not_promoted" + + metrics = startup_metrics.get("metrics", {}) + notes = [ + f"warmup_excluded={observed_summary.get('warmup_excluded', 0)}", + f"parsed_steps={observed_summary.get('parsed_step_count', 0)}", + ] + if startup_metrics.get("repo_commit"): + notes.append(f"commit={startup_metrics['repo_commit']}") + if isinstance(metrics.get("startup/master_addr"), str): + notes.append(f"master_addr={metrics['startup/master_addr']}") + if failure_status is not None: + notes.append(f"log_failure_status={failure_status}") + + return BenchmarkBehaviorPoint( + label=f"resolved_run:{config_path.parent.relative_to(benchmark_path)}", + source=str(log_path or config_path), + micro_batch_size=topology.micro_batch_size, + global_batch_size=topology.global_batch_size, + tokens_per_sec=tokens_per_sec, + step_time_sec=_round_or_none(observed_summary.get("step_time_s_mean"), 6), + mfu_percent=_round_or_none((observed_summary.get("mfu_mean") or 0.0) * 100.0, 3) + if observed_summary.get("mfu_mean") is not None + else None, + tflops_per_gpu=_round_or_none(observed_summary.get("tflops_per_gpu_mean"), 3), + peak_mem_gb=peak_mem_gb, + allocator_retries=None, + measured_steps=observed_summary.get("measured_steps"), + warmup_steps=observed_summary.get("warmup_excluded"), + gpu_count=topology.world_size, + sample_packing_sequence_len=topology.sample_packing_sequence_len, + tensor_parallel_size=topology.tensor_parallel_size, + pipeline_parallel_size=topology.pipeline_parallel_size, + ulysses_parallel_size=topology.ulysses_parallel_size, + ringattn_parallel_size=topology.ringattn_parallel_size, + expert_parallel_size=topology.expert_parallel_size, + ep_fsdp_size=topology.ep_fsdp_size, + deepep_async_combine=_config_bool(raw_config, "model", "deepep_async_combine", False), + deepep_num_sms=_config_int(raw_config, "model", "deepep_num_sms"), + deepep_buffer_size_gb=_config_float(raw_config, "model", "deepep_buffer_size_gb"), + enable_compile=_config_bool(raw_config, "train", "enable_compile", False), + gradient_checkpointing_method=_config_str(raw_config, "train", "gradient_checkpointing_method"), + enable_activation_offload=_config_bool(raw_config, "train", "enable_activation_offload", False), + activation_offload_prefetch_count=_config_int(raw_config, "train", "activation_offload_prefetch_count"), + status=status, + correctness_status=correctness_status, + notes=notes, + ) + + +def _resolved_run_points(benchmark_path: Path) -> list[BenchmarkBehaviorPoint]: + points: list[BenchmarkBehaviorPoint] = [] + for config_path in sorted(benchmark_path.rglob("xorl_cli.yaml")): + if not config_path.is_file(): + continue + point = _resolved_run_behavior_point(benchmark_path, config_path) + if point is not None: + points.append(point) + return points + + +def load_benchmark_behavior_points(benchmark_dir: str | Path) -> list[BenchmarkBehaviorPoint]: + benchmark_path = Path(benchmark_dir) + points: list[BenchmarkBehaviorPoint] = [] + topology_defaults: dict[str, int | float | bool | str] = {} + + for readme_path in (benchmark_path / "README.md", benchmark_path / "RESULTS.md"): + if not readme_path.is_file(): + continue + readme_text = readme_path.read_text(encoding="utf-8") + topology_defaults.update(_readme_topology_defaults(readme_text)) + seq_len = _seq_len_from_readme(readme_text) + readme_reference = _readme_point(readme_text, source=str(readme_path)) + if readme_reference is not None: + points.append(readme_reference) + adjacent_mbs10 = _readme_adjacent_mbs10_point(readme_text, source=str(readme_path), seq_len=seq_len) + if adjacent_mbs10 is not None: + points.append(adjacent_mbs10) + points.extend(_q235_markdown_points(readme_text, source=str(readme_path))) + + for result_path in sorted((benchmark_path / "results").glob("*.json")): + result = json.loads(result_path.read_text(encoding="utf-8")) + for row in result.get("best_by_mfu", []): + if isinstance(row, dict) and row.get("trial"): + points.append(_best_by_mfu_point(result_path, result, row, topology_defaults=topology_defaults)) + throughput = result.get("throughput") + if isinstance(throughput, dict): + points.append( + _with_k3_status( + _result_throughput_point(result_path, result, topology_defaults=topology_defaults), result + ) + ) + + points.extend(_resolved_run_points(benchmark_path)) + return points + + +def behavior_point_matches_topology(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + if point.micro_batch_size != topology.micro_batch_size or point.global_batch_size != topology.global_batch_size: + return False + if not _point_parallel_size_matches(point.tensor_parallel_size, topology.tensor_parallel_size): + return False + if not _point_parallel_size_matches(point.pipeline_parallel_size, topology.pipeline_parallel_size): + return False + if not _point_parallel_size_matches(point.ulysses_parallel_size, topology.ulysses_parallel_size): + return False + if not _point_parallel_size_matches(point.ringattn_parallel_size, topology.ringattn_parallel_size): + return False + if point.expert_parallel_size is None: + if topology.expert_parallel_size != 1: + return False + elif point.expert_parallel_size != topology.expert_parallel_size: + return False + if point.ep_fsdp_size is not None and point.ep_fsdp_size != topology.ep_fsdp_size: + return False + if ( + point.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len is not None + and point.sample_packing_sequence_len != topology.sample_packing_sequence_len + ): + return False + return True + + +def _section(raw_config: dict[str, Any], name: str) -> dict[str, Any]: + value = raw_config.get(name, {}) + return value if isinstance(value, dict) else {} + + +def _config_bool(raw_config: dict[str, Any], section_name: str, key: str, default: bool | None = None) -> bool | None: + section = _section(raw_config, section_name) + value = section.get(key, default) + if value is None: + return None + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + + +def _config_int(raw_config: dict[str, Any], section_name: str, key: str) -> int | None: + section = _section(raw_config, section_name) + value = section.get(key) + return int(value) if value is not None else None + + +def _config_float(raw_config: dict[str, Any], section_name: str, key: str) -> float | None: + section = _section(raw_config, section_name) + value = section.get(key) + return float(value) if value is not None else None + + +def _config_str(raw_config: dict[str, Any], section_name: str, key: str) -> str | None: + section = _section(raw_config, section_name) + value = section.get(key) + return str(value) if value is not None else None + + +def behavior_point_workload_mismatches(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> list[str]: + checks: tuple[tuple[str, Any, Any], ...] = ( + ( + "deepep_async_combine", + point.deepep_async_combine, + _config_bool(raw_config, "model", "deepep_async_combine", False), + ), + ("deepep_num_sms", point.deepep_num_sms, _config_int(raw_config, "model", "deepep_num_sms")), + ( + "deepep_buffer_size_gb", + point.deepep_buffer_size_gb, + _config_float(raw_config, "model", "deepep_buffer_size_gb"), + ), + ("enable_compile", point.enable_compile, _config_bool(raw_config, "train", "enable_compile", False)), + ( + "gradient_checkpointing_method", + point.gradient_checkpointing_method, + _config_str(raw_config, "train", "gradient_checkpointing_method"), + ), + ( + "enable_activation_offload", + point.enable_activation_offload, + _config_bool(raw_config, "train", "enable_activation_offload", False), + ), + ( + "activation_offload_prefetch_count", + point.activation_offload_prefetch_count, + _config_int(raw_config, "train", "activation_offload_prefetch_count"), + ), + ) + mismatches: list[str] = [] + for field_name, point_value, config_value in checks: + if point_value is None: + continue + if isinstance(point_value, float): + if config_value is None or abs(float(point_value) - float(config_value)) > 1e-9: + mismatches.append(field_name) + elif point_value != config_value: + mismatches.append(field_name) + return mismatches + + +def behavior_point_matches_workload(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> bool: + return not behavior_point_workload_mismatches(point, raw_config) + + +def _point_parallel_size_matches(point_value: int | None, topology_value: int) -> bool: + if point_value is None: + return topology_value == 1 + return point_value == topology_value + + +def predict_benchmark_behavior( + points: list[BenchmarkBehaviorPoint], + topology: Topology, + shape: ShapeLedger, + raw_config: dict[str, Any] | None = None, +) -> BenchmarkBehaviorPrediction: + matches = [ + point + for point in points + if behavior_point_matches_topology(point, topology) + and (raw_config is None or behavior_point_matches_workload(point, raw_config)) + ] + warnings: list[str] = [] + if not matches: + known = ", ".join( + f"{point.label}(mbs={point.micro_batch_size},gb={point.global_batch_size})" for point in points + ) + return BenchmarkBehaviorPrediction( + status="no_calibrated_match", + matched_label=None, + source=None, + tokens_per_sec=None, + tokens_per_sec_per_gpu=None, + step_time_sec=None, + mfu_percent=None, + tflops_per_gpu=None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=None, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + warnings=[ + f"no empirical behavior point for mbs={topology.micro_batch_size}, gb={topology.global_batch_size}; known: {known}" + ], + ) + + point = matches[0] + tokens_per_sec_per_gpu = None + if point.tokens_per_sec is not None and topology.world_size: + tokens_per_sec_per_gpu = point.tokens_per_sec / topology.world_size + step_time_sec = point.step_time_sec + if step_time_sec is None and shape.global_tokens_per_train_step and point.tokens_per_sec: + step_time_sec = shape.global_tokens_per_train_step / point.tokens_per_sec + tflops_per_gpu = None + if point.tflops_per_gpu is not None: + tflops_per_gpu = point.tflops_per_gpu + elif point.mfu_percent is not None: + tflops_per_gpu = H100_BF16_PROMISED_TFLOPS_PER_GPU * point.mfu_percent / 100.0 + + if point.status == "allocator_pressure_slowdown": + warnings.append("matched behavior point is an allocator-pressure slowdown, not a promotable speed target") + if point.correctness_status and point.correctness_status != "k3_pass": + warnings.append(f"correctness status is {point.correctness_status}") + + prediction_status = "calibrated_failure" if point.correctness_status == "oom" else "calibrated" + + return BenchmarkBehaviorPrediction( + status=prediction_status, + matched_label=point.label, + source=point.source, + tokens_per_sec=point.tokens_per_sec, + tokens_per_sec_per_gpu=tokens_per_sec_per_gpu, + step_time_sec=step_time_sec, + mfu_percent=point.mfu_percent, + tflops_per_gpu=tflops_per_gpu, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=point.peak_mem_gb, + allocator_retries=point.allocator_retries, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + correctness_status=point.correctness_status, + warnings=warnings, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("benchmark_dir", type=Path) + args = parser.parse_args() + points = load_benchmark_behavior_points(args.benchmark_dir) + print(json.dumps(to_jsonable({"points": points}), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/experiments/local_benchmark/training_sim/calibration_evaluator.py b/experiments/local_benchmark/training_sim/calibration_evaluator.py new file mode 100644 index 00000000..cd82fd24 --- /dev/null +++ b/experiments/local_benchmark/training_sim/calibration_evaluator.py @@ -0,0 +1,256 @@ +"""Evaluate scenario-prediction fidelity with leave-one-out benchmark holdouts.""" + +from __future__ import annotations + +import argparse +import json +import statistics +from pathlib import Path +from typing import Any + + +try: + from .benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from .config_fingerprint import load_training_config, resolve_topology + from .memory_ledger import build_memory_ledger + from .model_metadata import resolve_model_metadata + from .scenario_planner import _extrapolate_behavior, _mutated_config, _topology_label + from .schemas import ( + BenchmarkBehaviorPoint, + CalibrationHoldout, + CalibrationReport, + Topology, + to_jsonable, + ) + from .shape_engine import build_shape_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from config_fingerprint import load_training_config, resolve_topology + from memory_ledger import build_memory_ledger + from model_metadata import resolve_model_metadata + from scenario_planner import _extrapolate_behavior, _mutated_config, _topology_label + from schemas import BenchmarkBehaviorPoint, CalibrationHoldout, CalibrationReport, Topology, to_jsonable + from shape_engine import build_shape_ledger + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + if isinstance(value, dict): + return value + raw[name] = {} + return raw[name] + + +def _point_parallel_size(value: int | None, fallback: int) -> int: + return value if value is not None else fallback + + +def _set_if_known(section: dict[str, Any], key: str, value: Any) -> None: + if value is not None: + section[key] = value + + +def _apply_point_runtime_signature(raw_config: dict[str, Any], point: BenchmarkBehaviorPoint) -> None: + model = _section(raw_config, "model") + train = _section(raw_config, "train") + _set_if_known(model, "deepep_async_combine", point.deepep_async_combine) + _set_if_known(model, "deepep_num_sms", point.deepep_num_sms) + _set_if_known(model, "deepep_buffer_size_gb", point.deepep_buffer_size_gb) + _set_if_known(train, "enable_compile", point.enable_compile) + _set_if_known(train, "gradient_checkpointing_method", point.gradient_checkpointing_method) + _set_if_known(train, "enable_activation_offload", point.enable_activation_offload) + _set_if_known(train, "activation_offload_prefetch_count", point.activation_offload_prefetch_count) + + +def _topology_for_point( + base_config: dict[str, Any], + base_topology: Topology, + point: BenchmarkBehaviorPoint, + *, + world_size: int | None, + local_world_size: int | None, +) -> tuple[dict[str, Any] | None, Topology | None, str | None]: + if point.micro_batch_size is None or point.global_batch_size is None: + return None, None, "missing micro_batch_size/global_batch_size" + if point.tokens_per_sec is None: + return None, None, "missing tokens_per_sec" + + resolved_world_size = point.gpu_count or world_size or base_topology.world_size + resolved_local_world_size = local_world_size or base_topology.local_world_size + tensor_parallel = _point_parallel_size(point.tensor_parallel_size, base_topology.tensor_parallel_size) + pipeline_parallel = _point_parallel_size(point.pipeline_parallel_size, base_topology.pipeline_parallel_size) + ulysses_parallel = _point_parallel_size(point.ulysses_parallel_size, base_topology.ulysses_parallel_size) + ringattn_parallel = _point_parallel_size(point.ringattn_parallel_size, base_topology.ringattn_parallel_size) + expert_parallel = _point_parallel_size(point.expert_parallel_size, base_topology.expert_parallel_size) + non_dp = tensor_parallel * pipeline_parallel * ulysses_parallel * ringattn_parallel + if non_dp <= 0 or resolved_world_size % non_dp: + return None, None, "world_size is not divisible by heldout non-DP topology" + data_parallel_size = resolved_world_size // non_dp + denominator = point.micro_batch_size * data_parallel_size + if denominator <= 0 or point.global_batch_size % denominator: + return None, None, "global_batch_size is not divisible by micro_batch_size * data_parallel_size" + gradient_accumulation_steps = point.global_batch_size // denominator + + raw_config = _mutated_config( + base_config, + world_size=resolved_world_size, + micro_batch_size=point.micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + expert_parallel_size=expert_parallel, + tensor_parallel_size=tensor_parallel, + pipeline_parallel_size=pipeline_parallel, + ulysses_parallel_size=ulysses_parallel, + ringattn_parallel_size=ringattn_parallel, + ) + if point.sample_packing_sequence_len is not None: + _section(raw_config, "data")["sample_packing_sequence_len"] = point.sample_packing_sequence_len + _apply_point_runtime_signature(raw_config, point) + try: + topology = resolve_topology( + raw_config, + world_size=resolved_world_size, + local_world_size=resolved_local_world_size, + ) + except ValueError as exc: + return None, None, str(exc) + if point.ep_fsdp_size is not None and topology.ep_fsdp_size != point.ep_fsdp_size: + return None, None, "heldout ep_fsdp_size does not match resolved topology" + return raw_config, topology, None + + +def _without_point( + behavior_points: list[BenchmarkBehaviorPoint], + heldout: BenchmarkBehaviorPoint, +) -> list[BenchmarkBehaviorPoint]: + return [point for point in behavior_points if not (point.label == heldout.label and point.source == heldout.source)] + + +def evaluate_calibration( + base_config_path: str | Path, + *, + benchmark_dir: str | Path, + world_size: int | None = None, + local_world_size: int | None = None, + device_memory_limit_gb: float = 80.0, + memory_safety_factor: float = 1.15, +) -> CalibrationReport: + base_path = Path(base_config_path) + benchmark_path = Path(benchmark_dir) + base_config = load_training_config(base_path) + base_topology = resolve_topology(base_config, world_size=world_size, local_world_size=local_world_size) + metadata = resolve_model_metadata(base_config) + behavior_points = load_benchmark_behavior_points(benchmark_path) + measured_points = [point for point in behavior_points if point.tokens_per_sec is not None] + + holdouts: list[CalibrationHoldout] = [] + warnings: list[str] = [] + skipped_count = 0 + for heldout in measured_points: + raw_config, topology, skip_reason = _topology_for_point( + base_config, + base_topology, + heldout, + world_size=world_size, + local_world_size=local_world_size, + ) + if raw_config is None or topology is None: + skipped_count += 1 + warnings.append(f"skipped {heldout.label}: {skip_reason}") + continue + + training_points = _without_point(behavior_points, heldout) + shape = build_shape_ledger(topology, balanced_routing=True) + exact_prediction = predict_benchmark_behavior(training_points, topology, shape, raw_config) + if exact_prediction.status == "calibrated": + prediction = exact_prediction + else: + memory = build_memory_ledger(raw_config, topology=topology, model_metadata=metadata) + prediction, _ = _extrapolate_behavior( + training_points, + topology, + shape, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + ) + + predicted = prediction.tokens_per_sec + absolute_error = None + absolute_percentage_error = None + if predicted is not None: + absolute_error = abs(predicted - heldout.tokens_per_sec) + absolute_percentage_error = 100.0 * absolute_error / heldout.tokens_per_sec + holdouts.append( + CalibrationHoldout( + label=heldout.label, + source=heldout.source, + topology_label=_topology_label(topology), + actual_tokens_per_sec=heldout.tokens_per_sec, + predicted_tokens_per_sec=predicted, + prediction_status=prediction.status, + matched_label=prediction.matched_label, + absolute_error_tokens_per_sec=round(absolute_error, 3) if absolute_error is not None else None, + absolute_percentage_error=round(absolute_percentage_error, 3) + if absolute_percentage_error is not None + else None, + calibrated_from_count=len(training_points), + warnings=prediction.warnings, + ) + ) + + errors = [ + holdout.absolute_percentage_error for holdout in holdouts if holdout.absolute_percentage_error is not None + ] + status_counts: dict[str, int] = {} + for holdout in holdouts: + status_counts[holdout.prediction_status] = status_counts.get(holdout.prediction_status, 0) + 1 + status = "ok" if errors else "insufficient_data" + if holdouts and not errors: + warnings.append("all holdouts were unscored") + + return CalibrationReport( + base_config_path=str(base_path), + benchmark_dir=str(benchmark_path), + status=status, + measured_point_count=len(measured_points), + evaluated_count=len(holdouts), + skipped_count=skipped_count, + mean_absolute_percentage_error=round(statistics.fmean(errors), 3) if errors else None, + median_absolute_percentage_error=round(statistics.median(errors), 3) if errors else None, + max_absolute_percentage_error=round(max(errors), 3) if errors else None, + prediction_status_counts=dict(sorted(status_counts.items())), + holdouts=holdouts, + warnings=warnings, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--benchmark-dir", type=Path, required=True) + parser.add_argument("--world-size", type=int, default=None) + parser.add_argument("--local-world-size", type=int, default=None) + parser.add_argument("--device-memory-limit-gb", type=float, default=80.0) + parser.add_argument("--memory-safety-factor", type=float, default=1.15) + parser.add_argument("--output", type=Path, default=None) + args = parser.parse_args() + + report = evaluate_calibration( + args.config, + benchmark_dir=args.benchmark_dir, + world_size=args.world_size, + local_world_size=args.local_world_size, + device_memory_limit_gb=args.device_memory_limit_gb, + memory_safety_factor=args.memory_safety_factor, + ) + rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/experiments/local_benchmark/training_sim/collect_calibration.py b/experiments/local_benchmark/training_sim/collect_calibration.py new file mode 100644 index 00000000..d203be37 --- /dev/null +++ b/experiments/local_benchmark/training_sim/collect_calibration.py @@ -0,0 +1,224 @@ +"""Parse XoRL trainer structured logs into calibration observations.""" + +from __future__ import annotations + +import argparse +import json +import re +import statistics +from pathlib import Path +from typing import Any, Iterable + + +try: + from .schemas import MemoryPhaseObservation, ObservedRun, PhaseObservation, StepObservation, to_jsonable +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import MemoryPhaseObservation, ObservedRun, PhaseObservation, StepObservation, to_jsonable + + +STEP_RE = re.compile(r"\[STEP\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") +PHASE_RE = re.compile(r"\[(?PSTEP_PHASES(?:_PARTIAL)?)\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") +MEMORY_RE = re.compile(r"\[(?PSTEP_MEMORY(?:_PARTIAL)?)\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") +KV_RE = re.compile(r"(?P[A-Za-z0-9_+./-]+)=(?P\S+)") + + +def _float_or_none(value: str | None) -> float | None: + if value is None: + return None + cleaned = value.strip().rstrip(",") + for suffix in ("GB", "gb", "s"): + if cleaned.endswith(suffix): + cleaned = cleaned[: -len(suffix)] + break + try: + return float(cleaned) + except ValueError: + return None + + +def _parse_metric_body(body: str) -> dict[str, float]: + metrics: dict[str, float] = {} + for match in KV_RE.finditer(body): + numeric = _float_or_none(match.group("value")) + if numeric is not None: + metrics[match.group("key")] = numeric + return metrics + + +def _step_from_match(match: re.Match[str], source: str) -> StepObservation: + metrics = _parse_metric_body(match.group("body")) + phase_memory: dict[str, float] = {} + for key in ("fwd", "bwd", "optim", "fwd+bwd", "offload"): + if key in metrics: + phase_memory[key] = metrics[key] + + known_keys = { + "loss", + "grad_norm", + "lr", + "tflops", + "mfu", + "tokens_per_sec", + "time", + "peak_mem", + "fwd", + "bwd", + "optim", + "fwd+bwd", + "offload", + } + extra = {key: value for key, value in metrics.items() if key not in known_keys} + return StepObservation( + source=source, + step=int(match.group("step")), + max_steps=match.group("max"), + loss=metrics.get("loss"), + grad_norm=metrics.get("grad_norm"), + lr=metrics.get("lr"), + tflops_per_gpu=metrics.get("tflops"), + mfu=metrics.get("mfu"), + tokens_per_sec=metrics.get("tokens_per_sec"), + step_time_s=metrics.get("time"), + peak_mem_gb=metrics.get("peak_mem"), + phase_memory_gb=phase_memory, + extra=extra, + ) + + +def parse_log_text(text: str, *, source: str = "") -> ObservedRun: + steps: list[StepObservation] = [] + phases: list[PhaseObservation] = [] + memory_phases: list[MemoryPhaseObservation] = [] + + for line in text.splitlines(): + if phase_match := PHASE_RE.search(line): + phases.append( + PhaseObservation( + source=source, + prefix=phase_match.group("prefix"), + step=int(phase_match.group("step")), + max_steps=phase_match.group("max"), + metrics=_parse_metric_body(phase_match.group("body")), + ) + ) + continue + + if memory_match := MEMORY_RE.search(line): + memory_phases.append( + MemoryPhaseObservation( + source=source, + prefix=memory_match.group("prefix"), + step=int(memory_match.group("step")), + max_steps=memory_match.group("max"), + metrics=_parse_metric_body(memory_match.group("body")), + ) + ) + continue + + if step_match := STEP_RE.search(line): + steps.append(_step_from_match(step_match, source)) + + return ObservedRun(sources=[source], steps=steps, phases=phases, memory_phases=memory_phases) + + +def parse_log_path(path: str | Path) -> ObservedRun: + log_path = Path(path) + text = log_path.read_text(encoding="utf-8", errors="replace") + return parse_log_text(text, source=str(log_path)) + + +def merge_observed_runs(runs: Iterable[ObservedRun]) -> ObservedRun: + sources: list[str] = [] + steps: list[StepObservation] = [] + phases: list[PhaseObservation] = [] + memory_phases: list[MemoryPhaseObservation] = [] + for run in runs: + sources.extend(run.sources) + steps.extend(run.steps) + phases.extend(run.phases) + memory_phases.extend(run.memory_phases) + return ObservedRun(sources=sources, steps=steps, phases=phases, memory_phases=memory_phases) + + +def _mean(values: list[float]) -> float | None: + return statistics.fmean(values) if values else None + + +def _median(values: list[float]) -> float | None: + return statistics.median(values) if values else None + + +def summarize_observed_run( + run: ObservedRun, + *, + warmup_steps: int = 0, + world_size: int | None = None, +) -> dict[str, Any]: + ordered_steps = sorted(run.steps, key=lambda row: (row.source, row.step)) + measured = ordered_steps[warmup_steps:] + tps = [row.tokens_per_sec for row in measured if row.tokens_per_sec is not None] + tflops = [row.tflops_per_gpu for row in measured if row.tflops_per_gpu is not None] + mfu = [row.mfu for row in measured if row.mfu is not None] + step_time = [row.step_time_s for row in measured if row.step_time_s is not None] + peaks = [row.peak_mem_gb for row in measured if row.peak_mem_gb is not None] + + summary: dict[str, Any] = { + "sources": run.sources, + "parsed_step_count": len(run.steps), + "parsed_phase_count": len(run.phases), + "parsed_memory_phase_count": len(run.memory_phases), + "warmup_excluded": warmup_steps, + "measured_steps": len(measured), + "tokens_per_sec_mean": _mean(tps), + "tokens_per_sec_median": _median(tps), + "tflops_per_gpu_mean": _mean(tflops), + "mfu_mean": _mean(mfu), + "step_time_s_mean": _mean(step_time), + "peak_mem_gb_max": max(peaks) if peaks else None, + } + if world_size and summary["tokens_per_sec_mean"] is not None: + summary["tokens_per_sec_per_gpu_mean"] = summary["tokens_per_sec_mean"] / world_size + if measured: + summary["first_measured_step"] = measured[0].step + summary["last_measured_step"] = measured[-1].step + summary["loss_last"] = measured[-1].loss + return summary + + +def _expand_paths(paths: list[Path]) -> list[Path]: + expanded: list[Path] = [] + for path in paths: + if path.is_dir(): + expanded.extend(sorted(child for child in path.rglob("*") if child.is_file())) + else: + expanded.append(path) + return expanded + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", nargs="+", type=Path, help="Log files or directories to parse") + parser.add_argument("--warmup-steps", type=int, default=0, help="Drop this many parsed [STEP] rows from summary") + parser.add_argument("--world-size", type=int, default=None, help="Optional GPU count for per-GPU throughput") + parser.add_argument("--output", type=Path, default=None, help="Write JSON output to this path") + parser.add_argument("--include-rows", action="store_true", help="Include parsed row details, not just the summary") + args = parser.parse_args() + + runs = [parse_log_path(path) for path in _expand_paths(args.paths)] + observed = merge_observed_runs(runs) + payload: dict[str, Any] = { + "summary": summarize_observed_run(observed, warmup_steps=args.warmup_steps, world_size=args.world_size) + } + if args.include_rows: + payload["observed"] = to_jsonable(observed) + + rendered = json.dumps(to_jsonable(payload), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/experiments/local_benchmark/training_sim/config_fingerprint.py b/experiments/local_benchmark/training_sim/config_fingerprint.py new file mode 100644 index 00000000..428840dd --- /dev/null +++ b/experiments/local_benchmark/training_sim/config_fingerprint.py @@ -0,0 +1,213 @@ +"""Resolve the subset of XoRL config state needed by the simulator.""" + +from __future__ import annotations + +import hashlib +import os +import subprocess +from pathlib import Path +from typing import Any + +import yaml + + +try: + from .model_metadata import resolve_model_metadata + from .schemas import RunFingerprint, Topology +except ImportError: # pragma: no cover - exercised by direct script execution + from model_metadata import resolve_model_metadata + from schemas import RunFingerprint, Topology + + +REPO_ROOT = Path(__file__).resolve().parents[3] + + +def load_training_config(path: str | Path) -> dict[str, Any]: + with Path(path).open("r", encoding="utf-8") as handle: + return yaml.safe_load(handle) or {} + + +def config_sha256(path: str | Path) -> str: + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +def repo_commit(repo_root: str | Path = REPO_ROOT) -> str | None: + result = subprocess.run( + ["git", "-C", str(repo_root), "rev-parse", "--short=12", "HEAD"], + check=False, + capture_output=True, + text=True, + ) + if result.returncode != 0: + return None + return result.stdout.strip() or None + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + return value if isinstance(value, dict) else {} + + +def _int_value(section: dict[str, Any], key: str, default: int | None = None) -> int | None: + value = section.get(key, default) + if value is None: + return None + return int(value) + + +def _find_int(sections: list[dict[str, Any]], keys: tuple[str, ...]) -> int | None: + for section in sections: + for key in keys: + if key in section and section[key] is not None: + return int(section[key]) + return None + + +def _infer_world_size( + train: dict[str, Any], + *, + non_dp_size: int, + world_size: int | None, +) -> int: + if world_size is not None: + return int(world_size) + if "WORLD_SIZE" in os.environ: + return int(os.environ["WORLD_SIZE"]) + + replicate_size = _int_value(train, "data_parallel_replicate_size", -1) or -1 + shard_size = _int_value(train, "data_parallel_shard_size", -1) or -1 + if replicate_size > 0 and shard_size > 0: + return replicate_size * shard_size * non_dp_size + return 1 + + +def _infer_local_world_size(*, local_world_size: int | None, world_size: int) -> int: + if local_world_size is not None: + return int(local_world_size) + if "LOCAL_WORLD_SIZE" in os.environ: + return int(os.environ["LOCAL_WORLD_SIZE"]) + if world_size > 8 and world_size % 8 == 0: + return 8 + return world_size + + +def _resolve_dp_split(data_parallel_size: int, replicate_size: int, shard_size: int) -> tuple[int, int]: + if replicate_size > 0 and shard_size > 0: + if data_parallel_size != replicate_size * shard_size: + raise ValueError( + f"data_parallel_size ({data_parallel_size}) should equal " + f"data_parallel_replicate_size ({replicate_size}) * data_parallel_shard_size ({shard_size})." + ) + return replicate_size, shard_size + + if replicate_size > 0: + if data_parallel_size % replicate_size != 0: + raise ValueError("data_parallel_size should be a multiple of data_parallel_replicate_size.") + return replicate_size, data_parallel_size // replicate_size + + if shard_size > 0: + if data_parallel_size % shard_size != 0: + raise ValueError("data_parallel_size should be a multiple of data_parallel_shard_size.") + return data_parallel_size // shard_size, shard_size + + return 1, data_parallel_size + + +def resolve_topology( + raw_config: dict[str, Any], + *, + world_size: int | None = None, + local_world_size: int | None = None, + num_experts: int | None = None, + top_k: int | None = None, +) -> Topology: + train = _section(raw_config, "train") + data = _section(raw_config, "data") + + ulysses = _int_value(train, "ulysses_parallel_size", 1) or 1 + ringattn = _int_value(train, "ringattn_parallel_size", 1) or 1 + tensor_parallel = _int_value(train, "tensor_parallel_size", 1) or 1 + pipeline_parallel = _int_value(train, "pipeline_parallel_size", 1) or 1 + expert_parallel = _int_value(train, "expert_parallel_size", 1) or 1 + non_dp_size = ulysses * ringattn * tensor_parallel * pipeline_parallel + resolved_world_size = _infer_world_size(train, non_dp_size=non_dp_size, world_size=world_size) + resolved_local_world_size = _infer_local_world_size( + local_world_size=local_world_size, world_size=resolved_world_size + ) + if resolved_world_size <= 0 or resolved_local_world_size <= 0: + raise ValueError("world_size and local_world_size must be positive") + if resolved_world_size % resolved_local_world_size != 0: + raise ValueError("world_size must be divisible by local_world_size") + if resolved_world_size % non_dp_size != 0: + raise ValueError( + f"world_size ({resolved_world_size}) must be divisible by ulysses ({ulysses}) * ringattn " + f"({ringattn}) * tensor_parallel ({tensor_parallel}) * pipeline_parallel ({pipeline_parallel})." + ) + + data_parallel_size = resolved_world_size // non_dp_size + ranks_per_pipeline_stage = resolved_world_size // pipeline_parallel + ep_fsdp_size = ( + ranks_per_pipeline_stage // expert_parallel if ranks_per_pipeline_stage % expert_parallel == 0 else None + ) + replicate_size = _int_value(train, "data_parallel_replicate_size", -1) or -1 + shard_size = _int_value(train, "data_parallel_shard_size", -1) or -1 + replicate_size, shard_size = _resolve_dp_split(data_parallel_size, replicate_size, shard_size) + + micro_batch_size = _int_value(train, "micro_batch_size", 1) or 1 + gradient_accumulation_steps = _int_value(train, "gradient_accumulation_steps", 1) or 1 + global_batch_size = micro_batch_size * gradient_accumulation_steps * data_parallel_size + sample_packing_sequence_len = _int_value(data, "sample_packing_sequence_len", 32000) + + model_metadata = resolve_model_metadata(raw_config, num_experts=num_experts, top_k=top_k) + + return Topology( + world_size=resolved_world_size, + local_world_size=resolved_local_world_size, + node_count=resolved_world_size // resolved_local_world_size, + data_parallel_size=data_parallel_size, + data_parallel_replicate_size=replicate_size, + data_parallel_shard_size=shard_size, + tensor_parallel_size=tensor_parallel, + pipeline_parallel_size=pipeline_parallel, + expert_parallel_size=expert_parallel, + ep_fsdp_size=ep_fsdp_size, + ulysses_parallel_size=ulysses, + ringattn_parallel_size=ringattn, + micro_batch_size=micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + global_batch_size=global_batch_size, + sample_packing_sequence_len=sample_packing_sequence_len, + num_experts=model_metadata.num_experts, + top_k=model_metadata.top_k, + ) + + +def build_fingerprint( + config_path: str | Path, + *, + world_size: int | None = None, + local_world_size: int | None = None, + balanced_routing: bool = False, + num_experts: int | None = None, + top_k: int | None = None, + repo_root: str | Path = REPO_ROOT, +) -> RunFingerprint: + path = Path(config_path) + raw_config = load_training_config(path) + model_metadata = resolve_model_metadata(raw_config, num_experts=num_experts, top_k=top_k) + topology = resolve_topology( + raw_config, + world_size=world_size, + local_world_size=local_world_size, + num_experts=model_metadata.num_experts, + top_k=model_metadata.top_k, + ) + return RunFingerprint( + config_path=str(path), + config_sha256=config_sha256(path), + config_name=path.name, + repo_commit=repo_commit(repo_root), + balanced_routing=balanced_routing, + topology=topology, + model_metadata=model_metadata, + ) diff --git a/experiments/local_benchmark/training_sim/k8s/README.md b/experiments/local_benchmark/training_sim/k8s/README.md new file mode 100644 index 00000000..c53e3687 --- /dev/null +++ b/experiments/local_benchmark/training_sim/k8s/README.md @@ -0,0 +1,13 @@ +# Simulator Calibration K8s Notes + +Calibration jobs should be normal XoRL training benchmark jobs with these trainer flags enabled: + +```yaml +train: + enable_step_phase_timing: true + enable_per_component_timing: true + enable_step_memory_profiling: true +``` + +Any pod requesting GPUs on the research-common-h100 cluster must set `team: turbo` on the pod template labels. +Keep the run short, preserve the trainer-head log, and feed that log to `collect_calibration.py`. diff --git a/experiments/local_benchmark/training_sim/memory_ledger.py b/experiments/local_benchmark/training_sim/memory_ledger.py new file mode 100644 index 00000000..e24b6cd4 --- /dev/null +++ b/experiments/local_benchmark/training_sim/memory_ledger.py @@ -0,0 +1,246 @@ +"""Initial memory ledger built from config constants and observed structured logs.""" + +from __future__ import annotations + +from typing import Any + + +try: + from .schemas import MemoryBucket, MemoryLedger, ModelMetadata, ObservedRun, Topology +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import MemoryBucket, MemoryLedger, ModelMetadata, ObservedRun, Topology + + +BYTES_PER_GIB = 1024**3 + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + return value if isinstance(value, dict) else {} + + +def _float_field(section: dict[str, Any], key: str) -> float | None: + value = section.get(key) + if value is None: + return None + return float(value) + + +def _dtype_bytes(dtype: Any, *, default: int) -> int: + if dtype is None: + return default + normalized = str(dtype).lower() + if normalized in {"bf16", "bfloat16", "fp16", "float16", "half"}: + return 2 + if normalized in {"fp32", "float32", "float"}: + return 4 + if normalized in {"fp8", "float8", "e4m3", "e5m2"}: + return 1 + return default + + +def _gb(byte_count: float) -> float: + return byte_count / BYTES_PER_GIB + + +def _round_gb(value: float | None) -> float | None: + return round(value, 3) if value is not None else None + + +def _estimate_param_counts(metadata: ModelMetadata) -> tuple[float, float, float] | None: + hidden = metadata.hidden_size + layers = metadata.num_hidden_layers + vocab = metadata.vocab_size + if hidden is None or layers is None or vocab is None: + return None + + head_dim = metadata.head_dim + if head_dim is None and metadata.num_attention_heads: + head_dim = hidden // metadata.num_attention_heads + attention_heads = metadata.num_attention_heads or 1 + key_value_heads = metadata.num_key_value_heads or attention_heads + if head_dim is None: + return None + + q_proj = hidden * attention_heads * head_dim + k_proj = hidden * key_value_heads * head_dim + v_proj = hidden * key_value_heads * head_dim + o_proj = attention_heads * head_dim * hidden + attention_params = layers * (q_proj + k_proj + v_proj + o_proj) + + dense_mlp_params = 0 + has_routed_experts = metadata.num_experts is not None and metadata.moe_intermediate_size is not None + if metadata.intermediate_size is not None and not has_routed_experts: + dense_mlp_params = layers * 3 * hidden * metadata.intermediate_size + + shared_expert_params = 0 + if metadata.shared_expert_intermediate_size is not None: + shared_expert_params = layers * 3 * hidden * metadata.shared_expert_intermediate_size + + expert_params = 0 + if has_routed_experts and metadata.num_experts is not None and metadata.moe_intermediate_size is not None: + expert_params = layers * metadata.num_experts * 3 * hidden * metadata.moe_intermediate_size + + embedding_params = vocab * hidden + lm_head_params = 0 if metadata.tie_word_embeddings else vocab * hidden + norm_params = (2 * layers + 1) * hidden + non_expert_params = attention_params + dense_mlp_params + shared_expert_params + embedding_params + lm_head_params + non_expert_params += norm_params + return float(non_expert_params + expert_params), float(non_expert_params), float(expert_params) + + +def _model_state_buckets( + raw_config: dict[str, Any], + topology: Topology | None, + metadata: ModelMetadata | None, + deepep_buffer_size_gb: float | None, +) -> tuple[float | None, float | None, float | None, float | None, float | None, list[MemoryBucket], list[str]]: + if topology is None or metadata is None: + return None, None, None, None, None, [], ["parameter_and_optimizer_bytes"] + + counts = _estimate_param_counts(metadata) + if counts is None: + return None, None, None, None, None, [], ["parameter_and_optimizer_bytes"] + + total_params, non_expert_params, expert_params = counts + train = _section(raw_config, "train") + expert_shard_size = topology.expert_parallel_size * (topology.ep_fsdp_size or 1) + local_non_expert_params = non_expert_params / max(topology.data_parallel_shard_size, 1) + local_expert_params = expert_params / max(expert_shard_size, 1) + local_params = local_non_expert_params + local_expert_params + + weight_bytes = _dtype_bytes(train.get("param_dtype"), default=2 if train.get("enable_mixed_precision") else 4) + optimizer = str(train.get("optimizer", "")).lower() + optimizer_dtype_bytes = _dtype_bytes(train.get("optimizer_dtype"), default=4) + gradient_dtype = train.get("gradient_dtype") or train.get("fsdp_reduce_dtype") or train.get("optimizer_dtype") + gradient_bytes = _dtype_bytes(gradient_dtype, default=4) + + sharded_param_gb = _gb(local_params * weight_bytes) + master_param_gb = 0.0 + if optimizer == "adamw": + master_param_gb = _gb(local_params * optimizer_dtype_bytes) + persistent_model_state_gb = sharded_param_gb + master_param_gb + + gradient_state_gb = _gb(local_params * gradient_bytes) + optimizer_state_gb = 0.0 + if optimizer == "adamw": + optimizer_state_gb = _gb(local_params * 2 * optimizer_dtype_bytes) + elif optimizer == "muon" and float(train.get("muon_momentum", 0.0) or 0.0) > 0: + optimizer_state_gb = _gb(local_params * optimizer_dtype_bytes) + + buckets = [ + MemoryBucket( + name="sharded_trainable_params", + gb=_round_gb(sharded_param_gb) or 0.0, + source="analytic_model_metadata", + notes=[ + f"weight_bytes={weight_bytes}", + f"local_non_expert_params={local_non_expert_params:.0f}", + f"local_expert_params={local_expert_params:.0f}", + ], + ), + MemoryBucket( + name="gradient_storage", + gb=_round_gb(gradient_state_gb) or 0.0, + source="analytic_dtype_policy", + notes=[f"gradient_bytes={gradient_bytes}"], + ), + ] + if master_param_gb: + buckets.append( + MemoryBucket( + name="optimizer_master_params", + gb=_round_gb(master_param_gb) or 0.0, + source="analytic_optimizer_policy", + notes=[f"optimizer={optimizer}", f"optimizer_dtype_bytes={optimizer_dtype_bytes}"], + ) + ) + if optimizer_state_gb: + buckets.append( + MemoryBucket( + name=f"{optimizer}_optimizer_state", + gb=_round_gb(optimizer_state_gb) or 0.0, + source="analytic_optimizer_policy", + notes=[f"optimizer_dtype_bytes={optimizer_dtype_bytes}"], + ) + ) + if deepep_buffer_size_gb: + buckets.append( + MemoryBucket( + name="deepep_static_buffer", + gb=deepep_buffer_size_gb, + source="config", + ) + ) + + unsupported = [ + "activation_recompute_schedule", + "attention_workspace", + "moe_kernel_workspace", + "fsdp_unshard_and_reduce_scatter_transients", + "allocator_reserved_slack", + ] + return ( + total_params / 1_000_000_000, + local_params / 1_000_000_000, + persistent_model_state_gb, + gradient_state_gb, + optimizer_state_gb, + sorted(buckets, key=lambda bucket: bucket.gb, reverse=True), + unsupported, + ) + + +def build_memory_ledger( + raw_config: dict[str, Any], + observed: ObservedRun | None = None, + *, + topology: Topology | None = None, + model_metadata: ModelMetadata | None = None, +) -> MemoryLedger: + model = _section(raw_config, "model") + train = _section(raw_config, "train") + observed_peak = None + observed_phase_peak: dict[str, float] = {} + + if observed is not None: + peaks = [row.peak_mem_gb for row in observed.steps if row.peak_mem_gb is not None] + observed_peak = max(peaks) if peaks else None + for row in observed.steps: + for phase, value in row.phase_memory_gb.items(): + observed_phase_peak[phase] = max(value, observed_phase_peak.get(phase, value)) + for memory_row in observed.memory_phases: + for key, value in memory_row.metrics.items(): + observed_phase_peak[key] = max(value, observed_phase_peak.get(key, value)) + + deepep_buffer_size_gb = _float_field(model, "deepep_buffer_size_gb") + if deepep_buffer_size_gb is None: + deepep_buffer_size_gb = _float_field(train, "deepep_buffer_size_gb") + + ( + estimated_total_params_b, + estimated_local_params_b, + persistent_model_state_gb, + gradient_state_gb, + optimizer_state_gb, + top_memory_buckets, + unsupported_buckets, + ) = _model_state_buckets(raw_config, topology, model_metadata, deepep_buffer_size_gb) + analytic_peak_floor_gb = None + if persistent_model_state_gb is not None and gradient_state_gb is not None and optimizer_state_gb is not None: + analytic_peak_floor_gb = persistent_model_state_gb + gradient_state_gb + optimizer_state_gb + analytic_peak_floor_gb += deepep_buffer_size_gb or 0.0 + + return MemoryLedger( + deepep_buffer_size_gb=deepep_buffer_size_gb, + observed_peak_mem_gb_max=observed_peak, + observed_phase_peak_gb=observed_phase_peak, + estimated_total_params_b=_round_gb(estimated_total_params_b), + estimated_local_params_b=_round_gb(estimated_local_params_b), + persistent_model_state_gb=_round_gb(persistent_model_state_gb), + gradient_state_gb=_round_gb(gradient_state_gb), + optimizer_state_gb=_round_gb(optimizer_state_gb), + analytic_peak_floor_gb=_round_gb(analytic_peak_floor_gb), + top_memory_buckets=top_memory_buckets, + unsupported_buckets=unsupported_buckets, + ) diff --git a/experiments/local_benchmark/training_sim/model_metadata.py b/experiments/local_benchmark/training_sim/model_metadata.py new file mode 100644 index 00000000..e2267d2e --- /dev/null +++ b/experiments/local_benchmark/training_sim/model_metadata.py @@ -0,0 +1,255 @@ +"""Resolve lightweight model metadata needed by the simulator.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + + +try: + from .schemas import ModelMetadata +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import ModelMetadata + + +KNOWN_MODEL_METADATA: dict[str, dict[str, int]] = { + "Qwen/Qwen3-235B-A22B": { + "num_experts": 128, + "top_k": 8, + "num_hidden_layers": 94, + "hidden_size": 4096, + "intermediate_size": 12288, + "moe_intermediate_size": 1536, + "num_attention_heads": 64, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 151936, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3.6-35B-A3B": { + "num_experts": 256, + "top_k": 8, + "num_hidden_layers": 40, + "hidden_size": 2048, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "vocab_size": 248320, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3.6-35B-A3B-FP8": { + "num_experts": 256, + "top_k": 8, + "num_hidden_layers": 40, + "hidden_size": 2048, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "vocab_size": 248320, + "tie_word_embeddings": False, + }, +} + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + return value if isinstance(value, dict) else {} + + +def model_ref_from_config(raw_config: dict[str, Any]) -> str | None: + model = _section(raw_config, "model") + candidates = ( + model.get("config_path"), + model.get("model_path"), + model.get("model_name"), + raw_config.get("config_path"), + raw_config.get("model_path"), + raw_config.get("model_name"), + ) + for value in candidates: + if value: + return str(value) + return None + + +def default_hf_cache_roots() -> list[Path]: + roots: list[Path] = [] + hub_cache = os.environ.get("HUGGINGFACE_HUB_CACHE") + if hub_cache: + roots.append(Path(hub_cache)) + hf_home = os.environ.get("HF_HOME") + if hf_home: + roots.append(Path(hf_home) / "hub") + roots.extend( + [ + Path("/shared/huggingface/hub"), + Path.home() / ".cache" / "huggingface" / "hub", + ] + ) + deduped: list[Path] = [] + seen: set[Path] = set() + for root in roots: + expanded = root.expanduser() + if expanded not in seen: + seen.add(expanded) + deduped.append(expanded) + return deduped + + +def _candidate_config_paths(model_ref: str, hf_cache_roots: list[Path]) -> list[Path]: + ref_path = Path(model_ref).expanduser() + candidates: list[Path] = [] + if ref_path.is_file(): + candidates.append(ref_path) + elif ref_path.is_dir(): + candidates.append(ref_path / "config.json") + + if "/" in model_ref and not ref_path.exists(): + cache_name = "models--" + model_ref.replace("/", "--") + for root in hf_cache_roots: + snapshots_dir = root / cache_name / "snapshots" + if snapshots_dir.is_dir(): + candidates.extend( + sorted(snapshots_dir.glob("*/config.json"), key=lambda path: path.stat().st_mtime, reverse=True) + ) + return [path for path in candidates if path.is_file()] + + +def _find_int(sections: list[dict[str, Any]], keys: tuple[str, ...]) -> int | None: + for section in sections: + for key in keys: + if key in section and section[key] is not None: + return int(section[key]) + return None + + +def _find_bool(sections: list[dict[str, Any]], keys: tuple[str, ...]) -> bool | None: + for section in sections: + for key in keys: + if key in section and section[key] is not None: + return bool(section[key]) + return None + + +def _read_metadata_file(config_path: Path, model_ref: str | None) -> ModelMetadata: + data = json.loads(config_path.read_text(encoding="utf-8")) + text_config = data.get("text_config") if isinstance(data.get("text_config"), dict) else {} + sections = [text_config, data] + return ModelMetadata( + model_path=model_ref, + config_path=str(config_path), + source="hf_config", + num_experts=_find_int(sections, ("num_experts", "n_routed_experts")), + top_k=_find_int(sections, ("num_experts_per_tok", "moe_top_k", "top_k")), + num_hidden_layers=_find_int(sections, ("num_hidden_layers",)), + hidden_size=_find_int(sections, ("hidden_size",)), + intermediate_size=_find_int(sections, ("intermediate_size",)), + moe_intermediate_size=_find_int(sections, ("moe_intermediate_size",)), + shared_expert_intermediate_size=_find_int(sections, ("shared_expert_intermediate_size",)), + num_attention_heads=_find_int(sections, ("num_attention_heads",)), + num_key_value_heads=_find_int(sections, ("num_key_value_heads",)), + head_dim=_find_int(sections, ("head_dim",)), + vocab_size=_find_int(sections, ("vocab_size",)), + tie_word_embeddings=_find_bool(sections, ("tie_word_embeddings",)), + ) + + +def _known_metadata(model_ref: str) -> ModelMetadata | None: + values = KNOWN_MODEL_METADATA.get(model_ref) + if values is None: + return None + return ModelMetadata( + model_path=model_ref, + config_path=None, + source="known_model", + num_experts=values.get("num_experts"), + top_k=values.get("top_k"), + num_hidden_layers=values.get("num_hidden_layers"), + hidden_size=values.get("hidden_size"), + intermediate_size=values.get("intermediate_size"), + moe_intermediate_size=values.get("moe_intermediate_size"), + shared_expert_intermediate_size=values.get("shared_expert_intermediate_size"), + num_attention_heads=values.get("num_attention_heads"), + num_key_value_heads=values.get("num_key_value_heads"), + head_dim=values.get("head_dim"), + vocab_size=values.get("vocab_size"), + tie_word_embeddings=values.get("tie_word_embeddings"), + ) + + +def resolve_model_metadata( + raw_config: dict[str, Any], + *, + num_experts: int | None = None, + top_k: int | None = None, + hf_cache_roots: list[Path] | None = None, +) -> ModelMetadata: + model = _section(raw_config, "model") + model_config = _section(raw_config, "model_config") + nested_model_config = model.get("config", {}) if isinstance(model.get("config"), dict) else {} + config_sections = [model, nested_model_config, model_config] + model_ref = model_ref_from_config(raw_config) + + config_metadata = ModelMetadata( + model_path=model_ref, + config_path=None, + source="config", + num_experts=_find_int(config_sections, ("num_experts", "n_routed_experts")), + top_k=_find_int(config_sections, ("num_experts_per_tok", "moe_top_k", "top_k")), + num_hidden_layers=_find_int(config_sections, ("num_hidden_layers",)), + hidden_size=_find_int(config_sections, ("hidden_size",)), + intermediate_size=_find_int(config_sections, ("intermediate_size",)), + moe_intermediate_size=_find_int(config_sections, ("moe_intermediate_size",)), + shared_expert_intermediate_size=_find_int(config_sections, ("shared_expert_intermediate_size",)), + num_attention_heads=_find_int(config_sections, ("num_attention_heads",)), + num_key_value_heads=_find_int(config_sections, ("num_key_value_heads",)), + head_dim=_find_int(config_sections, ("head_dim",)), + vocab_size=_find_int(config_sections, ("vocab_size",)), + tie_word_embeddings=_find_bool(config_sections, ("tie_word_embeddings",)), + ) + + file_metadata = None + if model_ref: + roots = hf_cache_roots if hf_cache_roots is not None else default_hf_cache_roots() + candidate_paths = _candidate_config_paths(model_ref, roots) + if candidate_paths: + file_metadata = _read_metadata_file(candidate_paths[0], model_ref) + + known_metadata = _known_metadata(model_ref) if model_ref else None + source_metadata = file_metadata or known_metadata or config_metadata + + resolved_num_experts = num_experts if num_experts is not None else config_metadata.num_experts + if resolved_num_experts is None: + resolved_num_experts = source_metadata.num_experts + + resolved_top_k = top_k if top_k is not None else config_metadata.top_k + if resolved_top_k is None: + resolved_top_k = source_metadata.top_k + + source = source_metadata.source + if num_experts is not None or top_k is not None: + source = f"{source}+explicit_override" + + return ModelMetadata( + model_path=model_ref, + config_path=source_metadata.config_path, + source=source, + num_experts=resolved_num_experts, + top_k=resolved_top_k, + num_hidden_layers=source_metadata.num_hidden_layers, + hidden_size=source_metadata.hidden_size, + intermediate_size=source_metadata.intermediate_size, + moe_intermediate_size=source_metadata.moe_intermediate_size, + shared_expert_intermediate_size=source_metadata.shared_expert_intermediate_size, + num_attention_heads=source_metadata.num_attention_heads, + num_key_value_heads=source_metadata.num_key_value_heads, + head_dim=source_metadata.head_dim, + vocab_size=source_metadata.vocab_size, + tie_word_embeddings=source_metadata.tie_word_embeddings, + ) diff --git a/experiments/local_benchmark/training_sim/predict.py b/experiments/local_benchmark/training_sim/predict.py new file mode 100644 index 00000000..d967856f --- /dev/null +++ b/experiments/local_benchmark/training_sim/predict.py @@ -0,0 +1,124 @@ +"""Emit a static simulator report for one XoRL training config.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +try: + from .benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from .collect_calibration import merge_observed_runs, parse_log_path, summarize_observed_run + from .config_fingerprint import build_fingerprint, load_training_config + from .memory_ledger import build_memory_ledger + from .schemas import PredictionReport, to_jsonable + from .shape_engine import build_shape_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from collect_calibration import merge_observed_runs, parse_log_path, summarize_observed_run + from config_fingerprint import build_fingerprint, load_training_config + from memory_ledger import build_memory_ledger + from schemas import PredictionReport, to_jsonable + from shape_engine import build_shape_ledger + + +def build_report( + config_path: str | Path, + *, + world_size: int | None, + local_world_size: int | None, + balanced_routing: bool, + num_experts: int | None, + top_k: int | None, + log_paths: list[Path] | None = None, + warmup_steps: int = 0, + benchmark_dir: Path | None = None, +) -> PredictionReport: + fingerprint = build_fingerprint( + config_path, + world_size=world_size, + local_world_size=local_world_size, + balanced_routing=balanced_routing, + num_experts=num_experts, + top_k=top_k, + ) + raw_config = load_training_config(config_path) + shape = build_shape_ledger(fingerprint.topology, balanced_routing=balanced_routing) + observed = None + observed_summary: dict[str, Any] | None = None + calibration_sources: list[str] = [] + if log_paths: + observed = merge_observed_runs(parse_log_path(path) for path in log_paths) + calibration_sources = observed.sources + observed_summary = summarize_observed_run( + observed, + warmup_steps=warmup_steps, + world_size=fingerprint.topology.world_size, + ) + + memory = build_memory_ledger( + raw_config, + observed, + topology=fingerprint.topology, + model_metadata=fingerprint.model_metadata, + ) + benchmark_behavior = None + if benchmark_dir is not None: + behavior_points = load_benchmark_behavior_points(benchmark_dir) + benchmark_behavior = predict_benchmark_behavior(behavior_points, fingerprint.topology, shape, raw_config) + warnings = list(shape.warnings) + if memory.observed_peak_mem_gb_max is None: + warnings.append("no observed memory calibration was supplied") + if benchmark_behavior is not None: + warnings.extend(benchmark_behavior.warnings) + + return PredictionReport( + fingerprint=fingerprint, + shape=shape, + memory=memory, + benchmark_behavior=benchmark_behavior, + observed_summary=observed_summary, + calibration_sources=calibration_sources, + warnings=warnings, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True, help="XoRL YAML config") + parser.add_argument("--world-size", type=int, default=None, help="Override WORLD_SIZE for config resolution") + parser.add_argument("--local-world-size", type=int, default=None, help="Override LOCAL_WORLD_SIZE") + parser.add_argument("--balanced-routing", action="store_true", help="Assume deterministic balanced MoE routing") + parser.add_argument("--num-experts", type=int, default=None, help="Override model num_experts when config omits it") + parser.add_argument("--top-k", type=int, default=None, help="Override model top-k routing when config omits it") + parser.add_argument("--logs", nargs="*", type=Path, default=None, help="Optional trainer logs for calibration") + parser.add_argument( + "--warmup-steps", type=int, default=0, help="Drop this many parsed [STEP] rows from log summary" + ) + parser.add_argument("--benchmark-dir", type=Path, default=None, help="Optional benchmark recipe directory") + parser.add_argument("--output", type=Path, default=None, help="Write JSON report to this path") + args = parser.parse_args() + + report = build_report( + args.config, + world_size=args.world_size, + local_world_size=args.local_world_size, + balanced_routing=args.balanced_routing, + num_experts=args.num_experts, + top_k=args.top_k, + log_paths=args.logs, + warmup_steps=args.warmup_steps, + benchmark_dir=args.benchmark_dir, + ) + rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/experiments/local_benchmark/training_sim/scenario_planner.py b/experiments/local_benchmark/training_sim/scenario_planner.py new file mode 100644 index 00000000..a7859492 --- /dev/null +++ b/experiments/local_benchmark/training_sim/scenario_planner.py @@ -0,0 +1,1024 @@ +"""Plan and score topology scenarios from a base XoRL training config.""" + +from __future__ import annotations + +import argparse +import copy +import json +import math +from pathlib import Path +from typing import Any + + +try: + from .benchmark_behavior import ( + H100_BF16_PROMISED_TFLOPS_PER_GPU, + behavior_point_matches_topology, + behavior_point_matches_workload, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from .config_fingerprint import load_training_config, resolve_topology + from .memory_ledger import build_memory_ledger + from .model_metadata import resolve_model_metadata + from .schemas import ( + BenchmarkBehaviorPoint, + BenchmarkBehaviorPrediction, + ModelMetadata, + ScenarioCandidate, + ScenarioReport, + Topology, + to_jsonable, + ) + from .shape_engine import ShapeLedger, build_shape_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from benchmark_behavior import ( + H100_BF16_PROMISED_TFLOPS_PER_GPU, + behavior_point_matches_topology, + behavior_point_matches_workload, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from config_fingerprint import load_training_config, resolve_topology + from memory_ledger import build_memory_ledger + from model_metadata import resolve_model_metadata + from schemas import ( + BenchmarkBehaviorPoint, + BenchmarkBehaviorPrediction, + ModelMetadata, + ScenarioCandidate, + ScenarioReport, + Topology, + to_jsonable, + ) + from shape_engine import ShapeLedger, build_shape_ledger + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + if isinstance(value, dict): + return value + raw[name] = {} + return raw[name] + + +def _parse_int_list(raw: str | None) -> list[int] | None: + if raw is None or raw == "auto": + return None + values = sorted({int(part.strip()) for part in raw.split(",") if part.strip()}) + if not values: + raise ValueError("expected at least one integer") + return values + + +def _divisors(value: int) -> list[int]: + return [candidate for candidate in range(1, value + 1) if value % candidate == 0] + + +def _power_of_two_divisors(value: int, *, max_value: int | None = None) -> list[int]: + limit = value if max_value is None else min(value, max_value) + return [candidate for candidate in _divisors(value) if candidate <= limit and candidate & (candidate - 1) == 0] + + +def _dedupe_sorted(values: list[int] | set[int]) -> list[int]: + return sorted(value for value in set(values) if value > 0) + + +def _default_micro_batch_sizes( + base_topology: Topology, + behavior_points: list[BenchmarkBehaviorPoint], +) -> list[int]: + values = {base_topology.micro_batch_size} + values.update(point.micro_batch_size for point in behavior_points if point.micro_batch_size is not None) + return sorted(values) + + +def _default_ep_sizes(base_topology: Topology) -> list[int]: + if base_topology.num_experts is None: + return [base_topology.expert_parallel_size] + ranks_per_pipeline_stage = base_topology.world_size // base_topology.pipeline_parallel_size + values = { + value for value in _divisors(ranks_per_pipeline_stage) if value > 0 and base_topology.num_experts % value == 0 + } + if base_topology.expert_parallel_size in values: + return [base_topology.expert_parallel_size] + return sorted(values) or [base_topology.expert_parallel_size] + + +def _auto_ep_sizes(base_topology: Topology) -> list[int]: + if base_topology.num_experts is None: + return [base_topology.expert_parallel_size] + values = { + value + for value in _divisors(base_topology.world_size) + if base_topology.num_experts % value == 0 and value <= base_topology.world_size + } + values.add(base_topology.expert_parallel_size) + return _dedupe_sorted(values) + + +def _auto_tensor_parallel_sizes(base_topology: Topology, metadata: ModelMetadata) -> list[int]: + values = set(_power_of_two_divisors(base_topology.world_size, max_value=base_topology.local_world_size)) + values.add(base_topology.tensor_parallel_size) + if metadata.hidden_size is not None: + values = {value for value in values if metadata.hidden_size % value == 0} + if metadata.num_attention_heads is not None: + values = {value for value in values if metadata.num_attention_heads % value == 0} + return _dedupe_sorted(values) or [base_topology.tensor_parallel_size] + + +def _auto_pipeline_parallel_sizes(base_topology: Topology, metadata: ModelMetadata) -> list[int]: + values = set(_power_of_two_divisors(base_topology.world_size, max_value=4)) + values.add(base_topology.pipeline_parallel_size) + if metadata.num_hidden_layers is not None: + values = {value for value in values if metadata.num_hidden_layers % value == 0} + values.add(base_topology.pipeline_parallel_size) + return _dedupe_sorted(values) or [base_topology.pipeline_parallel_size] + + +def _auto_ulysses_parallel_sizes(base_topology: Topology) -> list[int]: + values = {base_topology.ulysses_parallel_size, 1} + seq_len = base_topology.sample_packing_sequence_len or 0 + if seq_len >= 16_384: + values.update(_power_of_two_divisors(base_topology.world_size, max_value=64)) + return _dedupe_sorted(values) + + +def _auto_ringattn_parallel_sizes(base_topology: Topology) -> list[int]: + values = {base_topology.ringattn_parallel_size, 1} + seq_len = base_topology.sample_packing_sequence_len or 0 + if seq_len >= 64_000: + values.update(_power_of_two_divisors(base_topology.world_size, max_value=4)) + return _dedupe_sorted(values) + + +def _known_or_default_parallel_size(point_value: int | None) -> int: + return point_value if point_value is not None else 1 + + +def _point_matches_topology_parallel_dims(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + return ( + _known_or_default_parallel_size(point.tensor_parallel_size) == topology.tensor_parallel_size + and _known_or_default_parallel_size(point.pipeline_parallel_size) == topology.pipeline_parallel_size + and _known_or_default_parallel_size(point.ulysses_parallel_size) == topology.ulysses_parallel_size + and _known_or_default_parallel_size(point.ringattn_parallel_size) == topology.ringattn_parallel_size + ) + + +def _point_matches_parallel_dims_for_risk(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + if point.expert_parallel_size is not None and point.expert_parallel_size != topology.expert_parallel_size: + return False + if point.ep_fsdp_size is not None and point.ep_fsdp_size != topology.ep_fsdp_size: + return False + return _point_matches_topology_parallel_dims(point, topology) + + +def _calibration_scope( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + *, + prediction_confidence: str, +) -> str: + if prediction_confidence == "calibrated": + return "exact_calibrated" + + throughput_points = [point for point in behavior_points if point.tokens_per_sec is not None] + if not throughput_points: + return "no_calibration" + + same_sequence_points = [ + point + for point in throughput_points + if point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) + ] + if not same_sequence_points: + return "outside_sequence_calibration_envelope" + + dimensions: tuple[tuple[str, int], ...] = ( + ("micro_batch_size", topology.micro_batch_size), + ("global_batch_size", topology.global_batch_size), + ("expert_parallel_size", topology.expert_parallel_size), + ("ep_fsdp_size", topology.ep_fsdp_size or 0), + ("tensor_parallel_size", topology.tensor_parallel_size), + ("pipeline_parallel_size", topology.pipeline_parallel_size), + ("ulysses_parallel_size", topology.ulysses_parallel_size), + ("ringattn_parallel_size", topology.ringattn_parallel_size), + ) + for field_name, topology_value in dimensions: + observed_values = [ + _known_or_default_parallel_size(getattr(point, field_name)) + if field_name.endswith("_parallel_size") + else getattr(point, field_name) + for point in same_sequence_points + if getattr(point, field_name) is not None + ] + if not observed_values: + continue + if topology_value < min(observed_values) or topology_value > max(observed_values): + return "outside_measured_envelope" + return "inside_measured_envelope" + + +def _candidate_risk_flags( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + behavior: BenchmarkBehaviorPrediction, + *, + raw_config: dict[str, Any] | None, + calibration_scope: str, + prediction_confidence: str, +) -> list[str]: + flags: list[str] = [] + if prediction_confidence != "calibrated": + flags.append("requires_remeasurement") + if calibration_scope.startswith("outside"): + flags.append(calibration_scope) + if behavior.correctness_status and behavior.correctness_status != "k3_pass": + flags.append(f"correctness_{behavior.correctness_status}") + + matched_labels = {part.strip() for part in (behavior.matched_label or "").split(",") if part.strip()} + for point in behavior_points: + if raw_config is not None and point.label in matched_labels: + for mismatch in behavior_point_workload_mismatches(point, raw_config): + flags.append(f"runtime_mismatch:{mismatch}") + same_sequence = point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) + if not same_sequence and point.sample_packing_sequence_len is not None: + same_sequence = ( + topology.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len >= point.sample_packing_sequence_len + ) + if not same_sequence or not _point_matches_parallel_dims_for_risk(point, topology): + continue + + if point.status == "allocator_pressure_slowdown": + at_or_beyond_mbs = ( + point.micro_batch_size is not None and topology.micro_batch_size >= point.micro_batch_size + ) + at_or_beyond_global_batch = ( + point.global_batch_size is not None and topology.global_batch_size >= point.global_batch_size + ) + if point.label in matched_labels: + flags.append("matched_allocator_pressure_slowdown") + elif at_or_beyond_mbs or at_or_beyond_global_batch: + flags.append(f"allocator_pressure_boundary:{point.label}") + + if point.correctness_status == "oom": + at_or_beyond_sequence = ( + point.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len >= point.sample_packing_sequence_len + ) + at_or_beyond_mbs = ( + point.micro_batch_size is not None and topology.micro_batch_size >= point.micro_batch_size + ) + at_or_beyond_global_batch = ( + point.global_batch_size is not None and topology.global_batch_size >= point.global_batch_size + ) + if point.label in matched_labels or ( + at_or_beyond_sequence and (at_or_beyond_mbs or at_or_beyond_global_batch) + ): + flags.append(f"observed_oom_boundary:{point.label}") + + return sorted(set(flags)) + + +def _risk_adjusted_score( + score_tokens_per_sec: float | None, + *, + calibration_scope: str, + risk_flags: list[str], + feasibility_status: str, +) -> float | None: + if score_tokens_per_sec is None: + return None + + multiplier = 1.0 + if calibration_scope == "inside_measured_envelope": + multiplier *= 0.85 + elif calibration_scope == "outside_measured_envelope": + multiplier *= 0.65 + elif calibration_scope == "outside_sequence_calibration_envelope": + multiplier *= 0.35 + elif calibration_scope == "no_calibration": + multiplier *= 0.20 + + if "matched_allocator_pressure_slowdown" in risk_flags: + multiplier *= 0.35 + elif any(flag.startswith("allocator_pressure_boundary:") for flag in risk_flags): + multiplier *= 0.50 + if any(flag.startswith("observed_oom_boundary:") for flag in risk_flags): + multiplier *= 0.25 + + for flag in risk_flags: + if flag == "correctness_k3_fail": + multiplier *= 0.50 + elif flag in {"correctness_not_promoted", "correctness_raw_speed_not_promoted_without_matching_k3_pass"}: + multiplier *= 0.95 + elif flag == "correctness_not_promoted_extrapolated": + multiplier *= 0.90 + elif flag == "correctness_runtime_failure_after_steps": + multiplier *= 0.45 + elif flag == "correctness_missing_calibration": + multiplier *= 0.50 + elif flag.startswith("correctness_"): + multiplier *= 0.75 + + if feasibility_status.endswith("_high_pressure"): + multiplier *= 0.85 + elif feasibility_status.endswith("_moderate_pressure"): + multiplier *= 0.95 + + return round(score_tokens_per_sec * multiplier, 3) + + +def _recommendation( + *, + feasible: bool, + promotable: bool, + feasibility_status: str, + risk_flags: list[str], +) -> str: + if feasibility_status == "observed_oom": + return "avoid_observed_oom" + if not feasible: + return "do_not_launch_unscored" + if promotable: + return "promote_candidate" + if "matched_allocator_pressure_slowdown" in risk_flags or any( + flag.startswith("allocator_pressure_boundary:") for flag in risk_flags + ): + return "measure_allocator_boundary" + if any(flag.startswith("observed_oom_boundary:") for flag in risk_flags): + return "remeasure_after_memory_fix" + if "requires_remeasurement" in risk_flags: + return "remeasure_before_ranking" + if "correctness_runtime_failure_after_steps" in risk_flags: + return "debug_runtime_failure" + if any(flag.startswith("correctness_") for flag in risk_flags): + return "correctness_gate_required" + return "review_candidate" + + +def _mutated_config( + base_config: dict[str, Any], + *, + world_size: int, + micro_batch_size: int, + gradient_accumulation_steps: int, + expert_parallel_size: int, + tensor_parallel_size: int, + pipeline_parallel_size: int, + ulysses_parallel_size: int, + ringattn_parallel_size: int, +) -> dict[str, Any]: + raw_config = copy.deepcopy(base_config) + train = _section(raw_config, "train") + train["micro_batch_size"] = micro_batch_size + train["gradient_accumulation_steps"] = gradient_accumulation_steps + train["expert_parallel_size"] = expert_parallel_size + train["tensor_parallel_size"] = tensor_parallel_size + train["pipeline_parallel_size"] = pipeline_parallel_size + train["ulysses_parallel_size"] = ulysses_parallel_size + train["ringattn_parallel_size"] = ringattn_parallel_size + + non_dp_size = tensor_parallel_size * pipeline_parallel_size * ulysses_parallel_size * ringattn_parallel_size + if non_dp_size <= 0 or world_size % non_dp_size != 0: + raise ValueError("world_size is not divisible by non-DP parallelism product") + data_parallel_size = world_size // non_dp_size + train["data_parallel_replicate_size"] = 1 + train["data_parallel_shard_size"] = data_parallel_size + if pipeline_parallel_size > 1: + train["gradient_accumulation_steps"] = max( + int(train.get("gradient_accumulation_steps", 1) or 1), pipeline_parallel_size + ) + return raw_config + + +def _topology_label(topology: Topology) -> str: + return ( + f"mbs{topology.micro_batch_size}-gb{topology.global_batch_size}-" + f"ep{topology.expert_parallel_size}-efsdp{topology.ep_fsdp_size}-" + f"tp{topology.tensor_parallel_size}-pp{topology.pipeline_parallel_size}-" + f"u{topology.ulysses_parallel_size}-r{topology.ringattn_parallel_size}" + ) + + +def _reference_tokens_per_gpu(point: BenchmarkBehaviorPoint, topology: Topology) -> float | None: + if point.tokens_per_sec is None: + return None + gpu_count = point.gpu_count or topology.world_size + if gpu_count <= 0: + return None + return point.tokens_per_sec / gpu_count + + +def _select_reference_point( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + raw_config: dict[str, Any] | None = None, +) -> BenchmarkBehaviorPoint | None: + usable = [ + point + for point in behavior_points + if point.tokens_per_sec is not None + and point.micro_batch_size is not None + and point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) + and _reference_tokens_per_gpu(point, topology) is not None + ] + if not usable: + return None + + workload_compatible = ( + [point for point in usable if behavior_point_matches_workload(point, raw_config)] + if raw_config is not None + else usable + ) + same_ep = [ + point for point in workload_compatible if point.expert_parallel_size in (None, topology.expert_parallel_size) + ] + candidates = same_ep or workload_compatible or usable + + def key(point: BenchmarkBehaviorPoint) -> tuple[float, float, float]: + mismatch_count = len(behavior_point_workload_mismatches(point, raw_config)) if raw_config is not None else 0 + mbs_distance = abs((point.micro_batch_size or 1) - topology.micro_batch_size) + per_gpu = _reference_tokens_per_gpu(point, topology) or 0.0 + return (-mismatch_count, -mbs_distance, per_gpu) + + return max(candidates, key=key) + + +def _parallelism_factor(reference: BenchmarkBehaviorPoint, topology: Topology) -> tuple[float, list[str]]: + notes: list[str] = [] + factor = 1.0 + if reference.micro_batch_size: + mbs_ratio = topology.micro_batch_size / reference.micro_batch_size + factor *= min(1.15, max(0.55, mbs_ratio**0.20)) + if reference.expert_parallel_size and reference.expert_parallel_size != topology.expert_parallel_size: + ep_ratio = topology.expert_parallel_size / reference.expert_parallel_size + factor *= max(0.70, 1.0 - 0.04 * abs(math.log2(ep_ratio))) + notes.append(f"EP extrapolated from {reference.expert_parallel_size} to {topology.expert_parallel_size}") + reference_tp = _known_or_default_parallel_size(reference.tensor_parallel_size) + if reference_tp != topology.tensor_parallel_size: + tp_ratio = topology.tensor_parallel_size / reference_tp + factor *= 0.90 ** abs(math.log2(tp_ratio)) + notes.append("TP extrapolation uses conservative communication penalty") + reference_pp = _known_or_default_parallel_size(reference.pipeline_parallel_size) + if reference_pp != topology.pipeline_parallel_size: + pp_delta = abs(topology.pipeline_parallel_size - reference_pp) + factor *= 0.88**pp_delta + notes.append("PP extrapolation uses conservative bubble penalty") + reference_cp = _known_or_default_parallel_size(reference.ulysses_parallel_size) * _known_or_default_parallel_size( + reference.ringattn_parallel_size + ) + if reference_cp != topology.sequence_parallel_size: + cp_ratio = topology.sequence_parallel_size / reference_cp + if topology.sample_packing_sequence_len and topology.sample_packing_sequence_len >= 32768: + factor *= min(1.10, 1.0 + 0.04 * abs(math.log2(cp_ratio))) + else: + factor *= 0.94 ** abs(math.log2(cp_ratio)) + notes.append("SP/CP extrapolation penalized for short-context workload") + return factor, notes + + +def _step_time_fit_prediction( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + shape: ShapeLedger, + raw_config: dict[str, Any] | None = None, +) -> BenchmarkBehaviorPrediction | None: + if topology.sample_packing_sequence_len is None or shape.global_tokens_per_train_step is None: + return None + compatible = [ + point + for point in behavior_points + if point.tokens_per_sec is not None + and point.global_batch_size is not None + and point.micro_batch_size == topology.micro_batch_size + and point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) + and point.expert_parallel_size in (None, topology.expert_parallel_size) + and point.ep_fsdp_size in (None, topology.ep_fsdp_size) + and _point_matches_topology_parallel_dims(point, topology) + and (raw_config is None or behavior_point_matches_workload(point, raw_config)) + ] + best_by_global_batch: dict[int, BenchmarkBehaviorPoint] = {} + for point in compatible: + current = best_by_global_batch.get(point.global_batch_size) + if current is None or (point.tokens_per_sec or 0.0) > (current.tokens_per_sec or 0.0): + best_by_global_batch[point.global_batch_size] = point + fit_points = sorted(best_by_global_batch.values(), key=lambda point: point.global_batch_size or 0) + if len(fit_points) < 2: + return None + + x_values: list[float] = [] + y_values: list[float] = [] + for point in fit_points: + tokens = point.global_batch_size * topology.sample_packing_sequence_len + step_time = point.step_time_sec + if step_time is None and point.tokens_per_sec: + step_time = tokens / point.tokens_per_sec + if step_time is None: + continue + x_values.append(float(tokens)) + y_values.append(float(step_time)) + if len(x_values) < 2 or len(set(x_values)) < 2: + return None + + x_mean = sum(x_values) / len(x_values) + y_mean = sum(y_values) / len(y_values) + denominator = sum((x_value - x_mean) ** 2 for x_value in x_values) + if denominator == 0: + return None + slope = sum((x_value - x_mean) * (y_value - y_mean) for x_value, y_value in zip(x_values, y_values, strict=False)) + slope /= denominator + intercept = y_mean - slope * x_mean + predicted_step = intercept + slope * shape.global_tokens_per_train_step + if predicted_step <= 0: + return None + tokens_per_sec = shape.global_tokens_per_train_step / predicted_step + tokens_per_sec_per_gpu = tokens_per_sec / topology.world_size + labels = ", ".join(point.label for point in fit_points) + peak_mem_gb = max((point.peak_mem_gb for point in fit_points if point.peak_mem_gb is not None), default=None) + return BenchmarkBehaviorPrediction( + status="extrapolated_step_time_fit", + matched_label=labels, + source="step_time_fit", + tokens_per_sec=round(tokens_per_sec, 3), + tokens_per_sec_per_gpu=round(tokens_per_sec_per_gpu, 3), + step_time_sec=round(predicted_step, 6), + mfu_percent=None, + tflops_per_gpu=None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=peak_mem_gb, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + correctness_status="not_promoted_extrapolated", + warnings=[ + f"extrapolated step time from calibrated global batches: {labels}", + f"fit_intercept_sec={intercept:.6f}", + f"fit_sec_per_token={slope:.12f}", + "correctness must be re-gated before promotion", + ], + ) + + +def _memory_factor( + memory_estimate_gb: float | None, + *, + memory_basis: str, + device_memory_limit_gb: float, + memory_safety_factor: float, +) -> tuple[float, float | None, str]: + if memory_estimate_gb is None: + return 0.0, None, "unknown_memory_estimate" + reserved_memory = memory_estimate_gb * memory_safety_factor + headroom = device_memory_limit_gb - reserved_memory + status_basis = "floor" if memory_basis == "analytic_floor" else memory_basis + if headroom < 0: + if memory_basis == "calibrated_peak" and memory_estimate_gb <= device_memory_limit_gb: + return 0.75, headroom, f"feasible_{status_basis}_high_pressure" + if memory_basis == "analytic_floor": + return 0.0, headroom, "memory_floor_exceeds_limit" + return 0.0, headroom, f"{status_basis}_exceeds_limit" + utilization = reserved_memory / device_memory_limit_gb if device_memory_limit_gb else 1.0 + if utilization >= 0.90: + return 0.75, headroom, f"feasible_{status_basis}_high_pressure" + if utilization >= 0.80: + return 0.90, headroom, f"feasible_{status_basis}_moderate_pressure" + return 1.0, headroom, f"feasible_{status_basis}" + + +def _extrapolate_behavior( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + shape: ShapeLedger, + *, + raw_config: dict[str, Any] | None = None, + device_memory_limit_gb: float, + memory_safety_factor: float, + analytic_peak_floor_gb: float | None, +) -> tuple[BenchmarkBehaviorPrediction, list[str]]: + step_fit = _step_time_fit_prediction(behavior_points, topology, shape, raw_config=raw_config) + if step_fit is not None: + return step_fit, ["step_time_fit_extrapolation"] + + reference = _select_reference_point(behavior_points, topology, raw_config=raw_config) + if reference is None: + return ( + BenchmarkBehaviorPrediction( + status="unscored", + matched_label=None, + source=None, + tokens_per_sec=None, + tokens_per_sec_per_gpu=None, + step_time_sec=None, + mfu_percent=None, + tflops_per_gpu=None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=None, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + correctness_status="missing_calibration", + warnings=["no benchmark behavior point is available for extrapolation"], + ), + [], + ) + + ref_per_gpu = _reference_tokens_per_gpu(reference, topology) or 0.0 + parallel_factor, notes = _parallelism_factor(reference, topology) + memory_factor, _, memory_status = _memory_factor( + analytic_peak_floor_gb, + memory_basis="analytic_floor", + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + tokens_per_sec_per_gpu = ref_per_gpu * parallel_factor * memory_factor + tokens_per_sec = tokens_per_sec_per_gpu * topology.world_size + step_time_sec = None + if shape.global_tokens_per_train_step and tokens_per_sec: + step_time_sec = shape.global_tokens_per_train_step / tokens_per_sec + tflops_per_gpu = reference.tflops_per_gpu + if tflops_per_gpu is None and reference.mfu_percent is not None: + tflops_per_gpu = H100_BF16_PROMISED_TFLOPS_PER_GPU * reference.mfu_percent / 100.0 + if tflops_per_gpu is not None and ref_per_gpu: + tflops_per_gpu *= tokens_per_sec_per_gpu / ref_per_gpu + + warnings = [ + f"extrapolated from {reference.label}; correctness must be re-gated before promotion", + f"memory feasibility status is {memory_status}", + ] + if raw_config is not None: + mismatches = behavior_point_workload_mismatches(reference, raw_config) + if mismatches: + warnings.append(f"reference runtime knobs differ: {', '.join(mismatches)}") + warnings.extend(notes) + return ( + BenchmarkBehaviorPrediction( + status="extrapolated", + matched_label=reference.label, + source=reference.source, + tokens_per_sec=round(tokens_per_sec, 3), + tokens_per_sec_per_gpu=round(tokens_per_sec_per_gpu, 3), + step_time_sec=round(step_time_sec, 6) if step_time_sec is not None else None, + mfu_percent=None, + tflops_per_gpu=round(tflops_per_gpu, 3) if tflops_per_gpu is not None else None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=reference.peak_mem_gb, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + correctness_status="not_promoted_extrapolated", + warnings=warnings, + ), + notes, + ) + + +def _candidate_from_prediction( + *, + label: str, + config_path: str | None, + topology: Topology, + shape: ShapeLedger, + behavior: BenchmarkBehaviorPrediction, + prediction_confidence: str, + promotable: bool, + behavior_points: list[BenchmarkBehaviorPoint], + raw_config: dict[str, Any] | None, + device_memory_limit_gb: float, + memory_safety_factor: float, + analytic_peak_floor_gb: float | None, + notes: list[str], +) -> ScenarioCandidate: + estimated_peak_mem_gb = analytic_peak_floor_gb + memory_basis = "analytic_floor" + if behavior.peak_mem_gb is not None: + if analytic_peak_floor_gb is None or behavior.peak_mem_gb >= analytic_peak_floor_gb: + estimated_peak_mem_gb = behavior.peak_mem_gb + memory_basis = "calibrated_peak" if prediction_confidence == "calibrated" else "extrapolated_peak" + + _, headroom, feasibility_status = _memory_factor( + estimated_peak_mem_gb, + memory_basis=memory_basis, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + if behavior.status == "calibrated_failure" or behavior.correctness_status == "oom": + feasibility_status = "observed_oom" + if behavior.tokens_per_sec is None and feasibility_status.startswith("feasible"): + feasibility_status = "unscored" + feasible = feasibility_status.startswith("feasible") and behavior.tokens_per_sec is not None + score_tokens_per_sec = behavior.tokens_per_sec if feasible else None + score_tokens_per_gpu_per_sec = behavior.tokens_per_sec_per_gpu if feasible else None + max_ep_slots = max(shape.ep_rank_slots_per_microbatch) if shape.ep_rank_slots_per_microbatch else None + calibration_scope = _calibration_scope( + behavior_points, + topology, + prediction_confidence=prediction_confidence, + ) + risk_flags = _candidate_risk_flags( + behavior_points, + topology, + behavior, + raw_config=raw_config, + calibration_scope=calibration_scope, + prediction_confidence=prediction_confidence, + ) + score_risk_adjusted_tokens_per_sec = _risk_adjusted_score( + score_tokens_per_sec, + calibration_scope=calibration_scope, + risk_flags=risk_flags, + feasibility_status=feasibility_status, + ) + recommendation = _recommendation( + feasible=feasible, + promotable=promotable and feasible, + feasibility_status=feasibility_status, + risk_flags=risk_flags, + ) + return ScenarioCandidate( + label=label, + config_path=config_path, + topology=topology, + behavior=behavior, + prediction_confidence=prediction_confidence, + promotable=promotable and feasible, + feasibility_status=feasibility_status, + score_tokens_per_sec=score_tokens_per_sec, + score_tokens_per_gpu_per_sec=score_tokens_per_gpu_per_sec, + score_risk_adjusted_tokens_per_sec=score_risk_adjusted_tokens_per_sec, + analytic_peak_floor_gb=analytic_peak_floor_gb, + estimated_peak_mem_gb=estimated_peak_mem_gb, + memory_basis=memory_basis, + memory_headroom_gb=round(headroom, 3) if headroom is not None else None, + max_ep_rank_slots_per_microbatch=max_ep_slots, + calibration_scope=calibration_scope, + recommendation=recommendation, + risk_flags=risk_flags, + notes=notes, + ) + + +def _candidate_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: + return ( + candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), + candidate.score_tokens_per_gpu_per_sec if candidate.score_tokens_per_gpu_per_sec is not None else float("-inf"), + ) + + +def _risk_adjusted_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: + return ( + candidate.score_risk_adjusted_tokens_per_sec + if candidate.score_risk_adjusted_tokens_per_sec is not None + else float("-inf"), + candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), + ) + + +def plan_scenario( + base_config_path: str | Path, + *, + benchmark_dir: str | Path | None = None, + world_size: int | None = None, + local_world_size: int | None = None, + micro_batch_sizes: list[int] | None = None, + gradient_accumulation_steps: list[int] | None = None, + expert_parallel_sizes: list[int] | None = None, + tensor_parallel_sizes: list[int] | None = None, + pipeline_parallel_sizes: list[int] | None = None, + ulysses_parallel_sizes: list[int] | None = None, + ringattn_parallel_sizes: list[int] | None = None, + topology_sweep: str = "base", + device_memory_limit_gb: float = 80.0, + memory_safety_factor: float = 1.15, +) -> ScenarioReport: + if topology_sweep not in {"base", "auto"}: + raise ValueError("topology_sweep must be 'base' or 'auto'") + base_path = Path(base_config_path) + base_config = load_training_config(base_path) + base_topology = resolve_topology(base_config, world_size=world_size, local_world_size=local_world_size) + resolved_world_size = world_size or base_topology.world_size + resolved_local_world_size = local_world_size or base_topology.local_world_size + behavior_points = load_benchmark_behavior_points(benchmark_dir) if benchmark_dir is not None else [] + metadata = resolve_model_metadata(base_config) + + micro_batch_values = micro_batch_sizes or _default_micro_batch_sizes(base_topology, behavior_points) + gradient_accumulation_values = gradient_accumulation_steps or [base_topology.gradient_accumulation_steps] + if topology_sweep == "auto": + ep_values = expert_parallel_sizes or _auto_ep_sizes(base_topology) + tp_values = tensor_parallel_sizes or _auto_tensor_parallel_sizes(base_topology, metadata) + pp_values = pipeline_parallel_sizes or _auto_pipeline_parallel_sizes(base_topology, metadata) + ulysses_values = ulysses_parallel_sizes or _auto_ulysses_parallel_sizes(base_topology) + ring_values = ringattn_parallel_sizes or _auto_ringattn_parallel_sizes(base_topology) + else: + ep_values = expert_parallel_sizes or [base_topology.expert_parallel_size] + tp_values = tensor_parallel_sizes or [base_topology.tensor_parallel_size] + pp_values = pipeline_parallel_sizes or [base_topology.pipeline_parallel_size] + ulysses_values = ulysses_parallel_sizes or [base_topology.ulysses_parallel_size] + ring_values = ringattn_parallel_sizes or [base_topology.ringattn_parallel_size] + + candidates: list[ScenarioCandidate] = [] + warnings: list[str] = [] + seen: set[tuple[str, str]] = set() + for pp in pp_values: + for tp in tp_values: + for ulysses in ulysses_values: + for ringattn in ring_values: + for ep in ep_values: + for micro_batch_size in micro_batch_values: + for gradient_accumulation_step in gradient_accumulation_values: + try: + raw_config = _mutated_config( + base_config, + world_size=resolved_world_size, + micro_batch_size=micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_step, + expert_parallel_size=ep, + tensor_parallel_size=tp, + pipeline_parallel_size=pp, + ulysses_parallel_size=ulysses, + ringattn_parallel_size=ringattn, + ) + topology = resolve_topology( + raw_config, + world_size=resolved_world_size, + local_world_size=resolved_local_world_size, + ) + except ValueError as exc: + warnings.append( + f"skipped mbs={micro_batch_size}, ga={gradient_accumulation_step}, " + f"ep={ep}, tp={tp}, pp={pp}, u={ulysses}, r={ringattn}: {exc}" + ) + continue + if topology.ep_fsdp_size is None: + warnings.append(f"skipped {_topology_label(topology)}: ep_fsdp is not integral") + continue + if ( + topology.num_experts is not None + and topology.num_experts % topology.expert_parallel_size + ): + warnings.append( + f"skipped {_topology_label(topology)}: EP does not divide num_experts" + ) + continue + + shape = build_shape_ledger(topology, balanced_routing=True) + memory = build_memory_ledger( + raw_config, + topology=topology, + model_metadata=metadata, + ) + exact_points = [ + point + for point in behavior_points + if behavior_point_matches_topology(point, topology) + and behavior_point_matches_workload(point, raw_config) + ] + if exact_points: + for point in exact_points: + behavior = predict_benchmark_behavior([point], topology, shape, raw_config) + label = f"{_topology_label(topology)}:{point.label}" + key = (label, point.source) + if key in seen: + continue + seen.add(key) + candidates.append( + _candidate_from_prediction( + label=label, + config_path=str(base_path), + topology=topology, + shape=shape, + behavior=behavior, + prediction_confidence="calibrated", + promotable=point.correctness_status == "k3_pass", + behavior_points=behavior_points, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + notes=list(point.notes), + ) + ) + continue + + behavior, extrapolation_notes = _extrapolate_behavior( + behavior_points, + topology, + shape, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + ) + label = f"{_topology_label(topology)}:extrapolated" + key = (label, behavior.source or "") + if key in seen: + continue + seen.add(key) + candidates.append( + _candidate_from_prediction( + label=label, + config_path=None, + topology=topology, + shape=shape, + behavior=behavior, + prediction_confidence=behavior.status, + promotable=False, + behavior_points=behavior_points, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + notes=extrapolation_notes, + ) + ) + + candidates = sorted(candidates, key=_candidate_sort_key, reverse=True) + feasible = [candidate for candidate in candidates if candidate.score_tokens_per_sec is not None] + best_raw = feasible[0] if feasible else None + risk_adjusted = [candidate for candidate in feasible if candidate.score_risk_adjusted_tokens_per_sec is not None] + best_risk_adjusted = max(risk_adjusted, key=_risk_adjusted_sort_key) if risk_adjusted else None + next_measurement = [candidate for candidate in risk_adjusted if "requires_remeasurement" in candidate.risk_flags] + best_next_measurement = max(next_measurement, key=_risk_adjusted_sort_key) if next_measurement else None + promotable = [candidate for candidate in feasible if candidate.promotable] + best_promotable = promotable[0] if promotable else None + if best_raw is not None and not best_raw.promotable: + warnings.append(f"best raw scenario {best_raw.label} is not correctness-promotable") + if best_raw is not None and best_risk_adjusted is not None and best_raw.label != best_risk_adjusted.label: + warnings.append( + f"best raw scenario {best_raw.label} differs from risk-adjusted choice {best_risk_adjusted.label}" + ) + if best_promotable is None: + warnings.append("no correctness-promotable scenario found") + + return ScenarioReport( + base_config_path=str(base_path), + benchmark_dir=str(benchmark_dir) if benchmark_dir is not None else None, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + topology_sweep=topology_sweep, + candidate_count=len(candidates), + feasible_count=len(feasible), + best_raw=best_raw, + best_risk_adjusted=best_risk_adjusted, + best_next_measurement=best_next_measurement, + best_promotable=best_promotable, + candidates=candidates, + warnings=warnings, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", type=Path, required=True) + parser.add_argument("--benchmark-dir", type=Path, default=None) + parser.add_argument("--world-size", type=int, default=None) + parser.add_argument("--local-world-size", type=int, default=None) + parser.add_argument("--micro-batch-sizes", default=None, help="Comma list, or auto when omitted") + parser.add_argument( + "--gradient-accumulation-steps", default=None, help="Comma list, or base config GA when omitted" + ) + parser.add_argument("--expert-parallel-sizes", default=None, help="Comma list, or base config EP when omitted") + parser.add_argument("--tensor-parallel-sizes", default=None, help="Comma list, or base config TP when omitted") + parser.add_argument("--pipeline-parallel-sizes", default=None, help="Comma list, or base config PP when omitted") + parser.add_argument( + "--ulysses-parallel-sizes", default=None, help="Comma list, or base config Ulysses when omitted" + ) + parser.add_argument("--ringattn-parallel-sizes", default=None, help="Comma list, or base config Ring when omitted") + parser.add_argument( + "--topology-sweep", + choices=("base", "auto"), + default="base", + help="Use base topology dimensions, or derive an automatic TP/PP/CP/EP sweep", + ) + parser.add_argument("--device-memory-limit-gb", type=float, default=80.0) + parser.add_argument("--memory-safety-factor", type=float, default=1.15) + parser.add_argument("--output", type=Path, default=None) + args = parser.parse_args() + + report = plan_scenario( + args.config, + benchmark_dir=args.benchmark_dir, + world_size=args.world_size, + local_world_size=args.local_world_size, + micro_batch_sizes=_parse_int_list(args.micro_batch_sizes), + gradient_accumulation_steps=_parse_int_list(args.gradient_accumulation_steps), + expert_parallel_sizes=_parse_int_list(args.expert_parallel_sizes), + tensor_parallel_sizes=_parse_int_list(args.tensor_parallel_sizes), + pipeline_parallel_sizes=_parse_int_list(args.pipeline_parallel_sizes), + ulysses_parallel_sizes=_parse_int_list(args.ulysses_parallel_sizes), + ringattn_parallel_sizes=_parse_int_list(args.ringattn_parallel_sizes), + topology_sweep=args.topology_sweep, + device_memory_limit_gb=args.device_memory_limit_gb, + memory_safety_factor=args.memory_safety_factor, + ) + rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/experiments/local_benchmark/training_sim/schemas.py b/experiments/local_benchmark/training_sim/schemas.py new file mode 100644 index 00000000..4536edf3 --- /dev/null +++ b/experiments/local_benchmark/training_sim/schemas.py @@ -0,0 +1,322 @@ +"""Dataclasses shared by the local training-engine simulator.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field, is_dataclass +from typing import Any + + +def to_jsonable(value: Any) -> Any: + """Convert simulator dataclasses into plain JSON-compatible containers.""" + if is_dataclass(value): + return {key: to_jsonable(item) for key, item in asdict(value).items()} + if isinstance(value, dict): + return {str(key): to_jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [to_jsonable(item) for item in value] + return value + + +@dataclass(frozen=True) +class ModelMetadata: + model_path: str | None + config_path: str | None + source: str + num_experts: int | None = None + top_k: int | None = None + num_hidden_layers: int | None = None + hidden_size: int | None = None + intermediate_size: int | None = None + moe_intermediate_size: int | None = None + shared_expert_intermediate_size: int | None = None + num_attention_heads: int | None = None + num_key_value_heads: int | None = None + head_dim: int | None = None + vocab_size: int | None = None + tie_word_embeddings: bool | None = None + + +@dataclass(frozen=True) +class Topology: + world_size: int + local_world_size: int + node_count: int + data_parallel_size: int + data_parallel_replicate_size: int + data_parallel_shard_size: int + tensor_parallel_size: int + pipeline_parallel_size: int + expert_parallel_size: int + ep_fsdp_size: int | None + ulysses_parallel_size: int + ringattn_parallel_size: int + micro_batch_size: int + gradient_accumulation_steps: int + global_batch_size: int + sample_packing_sequence_len: int | None + num_experts: int | None = None + top_k: int | None = None + + @property + def sequence_parallel_size(self) -> int: + return self.ulysses_parallel_size * self.ringattn_parallel_size + + +@dataclass(frozen=True) +class RunFingerprint: + config_path: str + config_sha256: str + config_name: str + repo_commit: str | None + balanced_routing: bool + topology: Topology + model_metadata: ModelMetadata + + +@dataclass(frozen=True) +class BalancedRoutingLedger: + total_slots: int + num_experts: int + counts_by_expert: list[int] + max_slots_per_expert: int + min_slots_per_expert: int + imbalance_slots: int + + +@dataclass(frozen=True) +class ShapeLedger: + microbatch_tokens_per_dp_rank: int | None + global_tokens_per_microbatch: int | None + global_tokens_per_train_step: int | None + tokens_per_gpu_per_train_step: float | None + sequence_parallel_size: int + tokens_per_model_rank_per_microbatch: int | None + routed_slots_per_model_rank_microbatch: int | None + routed_slots_per_train_step_model_rank: int | None + balanced_routing: BalancedRoutingLedger | None + ep_rank_slots_per_microbatch: list[int] | None = None + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class StepObservation: + source: str + step: int + max_steps: str + loss: float | None = None + grad_norm: float | None = None + lr: float | None = None + tflops_per_gpu: float | None = None + mfu: float | None = None + tokens_per_sec: float | None = None + step_time_s: float | None = None + peak_mem_gb: float | None = None + phase_memory_gb: dict[str, float] = field(default_factory=dict) + extra: dict[str, float] = field(default_factory=dict) + + +@dataclass(frozen=True) +class PhaseObservation: + source: str + prefix: str + step: int + max_steps: str + metrics: dict[str, float] + + +@dataclass(frozen=True) +class MemoryPhaseObservation: + source: str + prefix: str + step: int + max_steps: str + metrics: dict[str, float] + + +@dataclass(frozen=True) +class ObservedRun: + sources: list[str] + steps: list[StepObservation] + phases: list[PhaseObservation] = field(default_factory=list) + memory_phases: list[MemoryPhaseObservation] = field(default_factory=list) + + +@dataclass(frozen=True) +class MemoryBucket: + name: str + gb: float + source: str + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class MemoryLedger: + deepep_buffer_size_gb: float | None + observed_peak_mem_gb_max: float | None + observed_phase_peak_gb: dict[str, float] + estimated_total_params_b: float | None = None + estimated_local_params_b: float | None = None + persistent_model_state_gb: float | None = None + gradient_state_gb: float | None = None + optimizer_state_gb: float | None = None + analytic_peak_floor_gb: float | None = None + top_memory_buckets: list[MemoryBucket] = field(default_factory=list) + unsupported_buckets: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class BenchmarkBehaviorPoint: + label: str + source: str + micro_batch_size: int | None + global_batch_size: int | None + tokens_per_sec: float | None + step_time_sec: float | None + mfu_percent: float | None = None + tflops_per_gpu: float | None = None + peak_mem_gb: float | None = None + allocator_retries: int | None = None + measured_steps: int | None = None + warmup_steps: int | None = None + gpu_count: int | None = None + sample_packing_sequence_len: int | None = None + tensor_parallel_size: int | None = None + pipeline_parallel_size: int | None = None + ulysses_parallel_size: int | None = None + ringattn_parallel_size: int | None = None + expert_parallel_size: int | None = None + ep_fsdp_size: int | None = None + deepep_async_combine: bool | None = None + deepep_num_sms: int | None = None + deepep_buffer_size_gb: float | None = None + enable_compile: bool | None = None + gradient_checkpointing_method: str | None = None + enable_activation_offload: bool | None = None + activation_offload_prefetch_count: int | None = None + status: str = "observed" + correctness_status: str | None = None + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class BenchmarkBehaviorPrediction: + status: str + matched_label: str | None + source: str | None + tokens_per_sec: float | None + tokens_per_sec_per_gpu: float | None + step_time_sec: float | None + mfu_percent: float | None + tflops_per_gpu: float | None + promised_tflops_per_gpu: float | None + peak_mem_gb: float | None + allocator_retries: int | None + derived_global_tokens_per_step: int | None + correctness_status: str | None = None + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class PredictionReport: + fingerprint: RunFingerprint + shape: ShapeLedger + memory: MemoryLedger + benchmark_behavior: BenchmarkBehaviorPrediction | None = None + observed_summary: dict[str, Any] | None = None + calibration_sources: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class TradeoffCandidate: + label: str + config_path: str | None + behavior_source: str + topology: Topology | None + behavior: BenchmarkBehaviorPrediction + promotable: bool + score_tokens_per_sec: float | None + score_tflops_per_gpu: float | None + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class TradeoffReport: + benchmark_dir: str + status: str + candidate_count: int + best_raw: TradeoffCandidate | None + best_promotable: TradeoffCandidate | None + candidates: list[TradeoffCandidate] + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ScenarioCandidate: + label: str + config_path: str | None + topology: Topology + behavior: BenchmarkBehaviorPrediction + prediction_confidence: str + promotable: bool + feasibility_status: str + score_tokens_per_sec: float | None + score_tokens_per_gpu_per_sec: float | None + score_risk_adjusted_tokens_per_sec: float | None + analytic_peak_floor_gb: float | None + estimated_peak_mem_gb: float | None + memory_basis: str + memory_headroom_gb: float | None + max_ep_rank_slots_per_microbatch: int | None + calibration_scope: str + recommendation: str + risk_flags: list[str] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ScenarioReport: + base_config_path: str + benchmark_dir: str | None + device_memory_limit_gb: float + memory_safety_factor: float + topology_sweep: str + candidate_count: int + feasible_count: int + best_raw: ScenarioCandidate | None + best_risk_adjusted: ScenarioCandidate | None + best_next_measurement: ScenarioCandidate | None + best_promotable: ScenarioCandidate | None + candidates: list[ScenarioCandidate] + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class CalibrationHoldout: + label: str + source: str + topology_label: str + actual_tokens_per_sec: float + predicted_tokens_per_sec: float | None + prediction_status: str + matched_label: str | None + absolute_error_tokens_per_sec: float | None + absolute_percentage_error: float | None + calibrated_from_count: int + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class CalibrationReport: + base_config_path: str + benchmark_dir: str + status: str + measured_point_count: int + evaluated_count: int + skipped_count: int + mean_absolute_percentage_error: float | None + median_absolute_percentage_error: float | None + max_absolute_percentage_error: float | None + prediction_status_counts: dict[str, int] + holdouts: list[CalibrationHoldout] + warnings: list[str] = field(default_factory=list) diff --git a/experiments/local_benchmark/training_sim/shape_engine.py b/experiments/local_benchmark/training_sim/shape_engine.py new file mode 100644 index 00000000..36f17de7 --- /dev/null +++ b/experiments/local_benchmark/training_sim/shape_engine.py @@ -0,0 +1,109 @@ +"""Static token and balanced-routing shape calculations.""" + +from __future__ import annotations + +import math + + +try: + from .schemas import BalancedRoutingLedger, ShapeLedger, Topology +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import BalancedRoutingLedger, ShapeLedger, Topology + + +def balanced_counts(total_slots: int, num_experts: int) -> list[int]: + """Counts produced by round-robin balanced synthetic routing over expert slots.""" + if total_slots < 0: + raise ValueError("total_slots must be non-negative") + if num_experts <= 0: + raise ValueError("num_experts must be positive") + base, remainder = divmod(total_slots, num_experts) + return [base + (1 if expert_idx < remainder else 0) for expert_idx in range(num_experts)] + + +def balanced_routing_ledger(total_slots: int, num_experts: int) -> BalancedRoutingLedger: + counts = balanced_counts(total_slots, num_experts) + max_slots = max(counts) if counts else 0 + min_slots = min(counts) if counts else 0 + return BalancedRoutingLedger( + total_slots=total_slots, + num_experts=num_experts, + counts_by_expert=counts, + max_slots_per_expert=max_slots, + min_slots_per_expert=min_slots, + imbalance_slots=max_slots - min_slots, + ) + + +def ep_rank_counts(counts_by_expert: list[int], expert_parallel_size: int) -> list[int]: + if expert_parallel_size <= 0: + raise ValueError("expert_parallel_size must be positive") + if not counts_by_expert: + return [] + if len(counts_by_expert) % expert_parallel_size != 0: + raise ValueError("num_experts must be divisible by expert_parallel_size for contiguous EP ownership") + experts_per_rank = len(counts_by_expert) // expert_parallel_size + return [ + sum(counts_by_expert[start : start + experts_per_rank]) + for start in range(0, len(counts_by_expert), experts_per_rank) + ] + + +def build_shape_ledger(topology: Topology, *, balanced_routing: bool) -> ShapeLedger: + warnings: list[str] = [] + seq_len = topology.sample_packing_sequence_len + if seq_len is None: + warnings.append("data.sample_packing_sequence_len is not set; routed token counts are unavailable") + return ShapeLedger( + microbatch_tokens_per_dp_rank=None, + global_tokens_per_microbatch=None, + global_tokens_per_train_step=None, + tokens_per_gpu_per_train_step=None, + sequence_parallel_size=topology.sequence_parallel_size, + tokens_per_model_rank_per_microbatch=None, + routed_slots_per_model_rank_microbatch=None, + routed_slots_per_train_step_model_rank=None, + balanced_routing=None, + ep_rank_slots_per_microbatch=None, + warnings=warnings, + ) + + microbatch_tokens = topology.micro_batch_size * seq_len + global_tokens_per_microbatch = microbatch_tokens * topology.data_parallel_size + global_tokens_per_train_step = topology.global_batch_size * seq_len + tokens_per_gpu_per_train_step = global_tokens_per_train_step / topology.world_size + sequence_parallel_size = max(topology.sequence_parallel_size, 1) + model_rank_tokens = math.ceil(microbatch_tokens / sequence_parallel_size) + + routing_ledger = None + routed_slots_per_microbatch = None + routed_slots_per_step = None + ep_counts = None + + if topology.top_k is None or topology.num_experts is None: + warnings.append("num_experts/top_k are unknown; pass --num-experts and --top-k for MoE routing counts") + else: + routed_slots_per_microbatch = model_rank_tokens * topology.top_k + routed_slots_per_step = routed_slots_per_microbatch * topology.gradient_accumulation_steps + if balanced_routing: + routing_ledger = balanced_routing_ledger(routed_slots_per_microbatch, topology.num_experts) + try: + ep_counts = ep_rank_counts(routing_ledger.counts_by_expert, topology.expert_parallel_size) + except ValueError as exc: + warnings.append(str(exc)) + else: + warnings.append("balanced routing is disabled; expert-local slot counts are intentionally omitted") + + return ShapeLedger( + microbatch_tokens_per_dp_rank=microbatch_tokens, + global_tokens_per_microbatch=global_tokens_per_microbatch, + global_tokens_per_train_step=global_tokens_per_train_step, + tokens_per_gpu_per_train_step=tokens_per_gpu_per_train_step, + sequence_parallel_size=sequence_parallel_size, + tokens_per_model_rank_per_microbatch=model_rank_tokens, + routed_slots_per_model_rank_microbatch=routed_slots_per_microbatch, + routed_slots_per_train_step_model_rank=routed_slots_per_step, + balanced_routing=routing_ledger, + ep_rank_slots_per_microbatch=ep_counts, + warnings=warnings, + ) diff --git a/experiments/local_benchmark/training_sim/tradeoff_ranker.py b/experiments/local_benchmark/training_sim/tradeoff_ranker.py new file mode 100644 index 00000000..e0807cc7 --- /dev/null +++ b/experiments/local_benchmark/training_sim/tradeoff_ranker.py @@ -0,0 +1,202 @@ +"""Rank calibrated benchmark tradeoffs across parallelism/config choices.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +try: + from .benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from .config_fingerprint import load_training_config, resolve_topology + from .schemas import BenchmarkBehaviorPoint, Topology, TradeoffCandidate, TradeoffReport, to_jsonable + from .shape_engine import build_shape_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from config_fingerprint import load_training_config, resolve_topology + from schemas import BenchmarkBehaviorPoint, Topology, TradeoffCandidate, TradeoffReport, to_jsonable + from shape_engine import build_shape_ledger + + +def _config_output_basename(config_path: Path) -> str | None: + raw_config = load_training_config(config_path) + output_dir = raw_config.get("train", {}).get("output_dir") + return Path(output_dir).name if output_dir else None + + +def _config_topology(config_path: Path, *, world_size: int | None, local_world_size: int | None) -> Topology: + raw_config = load_training_config(config_path) + return resolve_topology(raw_config, world_size=world_size, local_world_size=local_world_size) + + +def _behavior_key(point: BenchmarkBehaviorPoint) -> str: + if point.label.startswith("best_by_mfu:"): + return point.label.split(":", 1)[1] + if ":" in point.label: + return point.label.split(":", 1)[1] + return point.label + + +def _find_matching_config( + point: BenchmarkBehaviorPoint, + configs: list[Path], + *, + world_size: int | None, + local_world_size: int | None, +) -> tuple[Path | None, Topology | None]: + behavior_key = _behavior_key(point) + output_names = {path: _config_output_basename(path) for path in configs} + for path, output_name in output_names.items(): + if output_name and output_name == behavior_key: + return path, _config_topology(path, world_size=world_size, local_world_size=local_world_size) + + for path in configs: + topology = _config_topology(path, world_size=world_size, local_world_size=local_world_size) + if ( + point.micro_batch_size == topology.micro_batch_size + and point.global_batch_size == topology.global_batch_size + ): + if point.expert_parallel_size is None or point.expert_parallel_size == topology.expert_parallel_size: + return path, topology + return None, None + + +def _is_promotable(point: BenchmarkBehaviorPoint) -> bool: + return point.correctness_status == "k3_pass" + + +def _candidate_sort_key(candidate: TradeoffCandidate) -> tuple[float, float]: + return ( + candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), + candidate.score_tflops_per_gpu if candidate.score_tflops_per_gpu is not None else float("-inf"), + ) + + +def _fallback_topology( + point: BenchmarkBehaviorPoint, + *, + world_size: int | None, + local_world_size: int | None, +) -> Topology: + resolved_world_size = world_size or point.gpu_count or 1 + resolved_local_world_size = local_world_size or (8 if resolved_world_size % 8 == 0 else resolved_world_size) + global_batch_size = point.global_batch_size or point.micro_batch_size or 1 + micro_batch_size = point.micro_batch_size or 1 + return Topology( + world_size=resolved_world_size, + local_world_size=resolved_local_world_size, + node_count=max(resolved_world_size // resolved_local_world_size, 1), + data_parallel_size=1, + data_parallel_replicate_size=1, + data_parallel_shard_size=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + expert_parallel_size=point.expert_parallel_size or 1, + ep_fsdp_size=point.ep_fsdp_size, + ulysses_parallel_size=1, + ringattn_parallel_size=1, + micro_batch_size=micro_batch_size, + gradient_accumulation_steps=max(global_batch_size // micro_batch_size, 1), + global_batch_size=global_batch_size, + sample_packing_sequence_len=None, + num_experts=None, + top_k=None, + ) + + +def rank_benchmark_tradeoffs( + benchmark_dir: str | Path, + *, + world_size: int | None = None, + local_world_size: int | None = None, +) -> TradeoffReport: + benchmark_path = Path(benchmark_dir) + configs = sorted((benchmark_path / "configs").glob("*.yaml")) + behavior_points = load_benchmark_behavior_points(benchmark_path) + warnings: list[str] = [] + candidates: list[TradeoffCandidate] = [] + + for point in behavior_points: + config_path, topology = _find_matching_config( + point, + configs, + world_size=world_size, + local_world_size=local_world_size, + ) + if topology is None: + warnings.append(f"no matching config found for behavior point {point.label}") + fallback_topology = _fallback_topology( + point, + world_size=world_size, + local_world_size=local_world_size, + ) + behavior = predict_benchmark_behavior( + [point], + fallback_topology, + build_shape_ledger(fallback_topology, balanced_routing=True), + ) + else: + shape = build_shape_ledger(topology, balanced_routing=True) + behavior = predict_benchmark_behavior([point], topology, shape) + + notes = list(point.notes) + if point.correctness_status and point.correctness_status != "k3_pass": + notes.append(f"not promotable: {point.correctness_status}") + candidates.append( + TradeoffCandidate( + label=point.label, + config_path=str(config_path) if config_path else None, + behavior_source=point.source, + topology=topology, + behavior=behavior, + promotable=_is_promotable(point), + score_tokens_per_sec=behavior.tokens_per_sec, + score_tflops_per_gpu=behavior.tflops_per_gpu, + notes=notes, + ) + ) + + candidates = sorted(candidates, key=_candidate_sort_key, reverse=True) + best_raw = candidates[0] if candidates else None + promotable = [candidate for candidate in candidates if candidate.promotable] + best_promotable = promotable[0] if promotable else None + if best_raw is not None and not best_raw.promotable: + warnings.append(f"best raw candidate {best_raw.label} is not correctness-promotable") + if best_promotable is None: + warnings.append("no correctness-promotable candidate found") + + return TradeoffReport( + benchmark_dir=str(benchmark_path), + status="ok" if candidates else "no_behavior_points", + candidate_count=len(candidates), + best_raw=best_raw, + best_promotable=best_promotable, + candidates=candidates, + warnings=warnings, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("benchmark_dir", type=Path) + parser.add_argument("--world-size", type=int, default=None) + parser.add_argument("--local-world-size", type=int, default=None) + parser.add_argument("--output", type=Path, default=None) + args = parser.parse_args() + + report = rank_benchmark_tradeoffs( + args.benchmark_dir, + world_size=args.world_size, + local_world_size=args.local_world_size, + ) + rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/experiments/local_benchmark/training_sim/validate_benchmarks.py b/experiments/local_benchmark/training_sim/validate_benchmarks.py new file mode 100644 index 00000000..4b8f5a00 --- /dev/null +++ b/experiments/local_benchmark/training_sim/validate_benchmarks.py @@ -0,0 +1,457 @@ +"""Validate simulator output against checked-in throughput benchmark recipes.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import subprocess +import sys +import tempfile +from dataclasses import replace +from pathlib import Path +from typing import Any + + +try: + from .benchmark_behavior import ( + H100_BF16_PROMISED_TFLOPS_PER_GPU, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from .config_fingerprint import load_training_config + from .predict import build_report + from .schemas import to_jsonable + from .shape_engine import build_shape_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from benchmark_behavior import ( + H100_BF16_PROMISED_TFLOPS_PER_GPU, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from config_fingerprint import load_training_config + from predict import build_report + from schemas import to_jsonable + from shape_engine import build_shape_ledger + + +def _human_number(value: str) -> float: + cleaned = value.strip().replace(",", "").lstrip("~") + multiplier = 1.0 + if cleaned.endswith(("K", "k")): + cleaned = cleaned[:-1] + multiplier = 1_000.0 + elif cleaned.endswith(("M", "m")): + cleaned = cleaned[:-1] + multiplier = 1_000_000.0 + return float(cleaned) * multiplier + + +def _parse_readme_metrics(readme_text: str) -> dict[str, Any]: + metrics: dict[str, Any] = {} + if match := re.search(r"Hardware:\s*(?P\d+)\s*nodes?\s*x\s*(?P\d+)\s*H100", readme_text): + metrics["nodes"] = int(match.group("nodes")) + metrics["gpus_per_node"] = int(match.group("gpus")) + metrics["world_size"] = metrics["nodes"] * metrics["gpus_per_node"] + if match := re.search(r"sample_packing_sequence_len:\s*(?P\d+)", readme_text): + metrics["sample_packing_sequence_len"] = int(match.group("seq")) + if match := re.search(r"\|\s*tokens/sec\s*\|\s*(?P~?[0-9.]+[KkMm]?)\s*\|", readme_text): + metrics["tokens_per_sec"] = _human_number(match.group("value")) + if match := re.search(r"\|\s*step time\s*\|\s*(?P~?[0-9.]+)s\s*\|", readme_text): + metrics["step_time_sec"] = float(match.group("value").lstrip("~")) + if match := re.search(r"\|\s*MFU\s*\|\s*(?P~?[0-9.]+)%", readme_text): + metrics["mfu_percent"] = float(match.group("value").lstrip("~")) + if match := re.search(r"\|\s*allocated memory\s*\|\s*(?P~?[0-9.]+)GB\s*\|", readme_text): + metrics["allocated_memory_gb"] = float(match.group("value").lstrip("~")) + if match := re.search(r"`mbs=10`[^~]+~(?P[0-9.]+)K tok/s", readme_text): + metrics["mbs10_tokens_per_sec"] = _human_number(match.group("value") + "K") + return metrics + + +def _check(name: str, expected: Any, actual: Any, *, tolerance: float | None = None) -> dict[str, Any]: + if tolerance is None: + passed = expected == actual + else: + passed = expected is not None and actual is not None and abs(float(expected) - float(actual)) <= tolerance + return { + "name": name, + "status": "pass" if passed else "fail", + "expected": expected, + "actual": actual, + "tolerance": tolerance, + } + + +def _benchmark_config_path(benchmark_dir: Path) -> Path: + configs = sorted((benchmark_dir / "configs").glob("*.yaml")) + if not configs: + raise FileNotFoundError(f"no benchmark configs found under {benchmark_dir / 'configs'}") + if len(configs) > 1: + raise ValueError(f"expected one config under {benchmark_dir / 'configs'}, found {len(configs)}") + return configs[0] + + +def _render_script_text(benchmark_dir: Path) -> str: + script_path = benchmark_dir / "scripts" / "render_k8s_manifest.sh" + return script_path.read_text(encoding="utf-8") if script_path.is_file() else "" + + +def _render_manifest_text(benchmark_dir: Path) -> str: + script_path = benchmark_dir / "scripts" / "render_k8s_manifest.sh" + if not script_path.is_file(): + return "" + with tempfile.TemporaryDirectory(prefix="xorl-training-sim-") as tmpdir: + output_path = Path(tmpdir) / "manifest.yaml" + env = os.environ.copy() + env["OUTPUT"] = str(output_path) + result = subprocess.run( + [str(script_path)], + check=False, + cwd=Path.cwd(), + env=env, + capture_output=True, + text=True, + ) + if result.returncode != 0: + return result.stdout + result.stderr + return output_path.read_text(encoding="utf-8") + + +def _validate_config_behavior( + benchmark_dir: Path, + readme_metrics: dict[str, Any], +) -> tuple[dict[str, Any], list[dict[str, Any]], dict[str, dict[str, Any]]]: + config_path = _benchmark_config_path(benchmark_dir) + raw_config = load_training_config(config_path) + world_size = readme_metrics.get("world_size") + local_world_size = readme_metrics.get("gpus_per_node") + report = build_report( + config_path, + world_size=world_size, + local_world_size=local_world_size, + balanced_routing=True, + num_experts=None, + top_k=None, + benchmark_dir=benchmark_dir, + ) + topology = report.fingerprint.topology + shape = report.shape + model = raw_config.get("model", {}) + data = raw_config.get("data", {}) + train = raw_config.get("train", {}) + + checks = [ + _check("readme_reference_tokens_per_sec", 261000.0, readme_metrics.get("tokens_per_sec")), + _check("readme_reference_step_time_sec", 8.04, readme_metrics.get("step_time_sec")), + _check("readme_reference_mfu_percent", 16.2, readme_metrics.get("mfu_percent")), + _check("readme_reference_allocated_memory_gb", 56.4, readme_metrics.get("allocated_memory_gb")), + _check("readme_mbs10_tokens_per_sec", 133700.0, readme_metrics.get("mbs10_tokens_per_sec")), + _check( + "readme_mbs10_allocator_pressure_slowdown", + True, + readme_metrics.get("mbs10_tokens_per_sec", 0) < readme_metrics.get("tokens_per_sec", 0) * 0.6, + ), + _check("world_size", world_size, topology.world_size), + _check("local_world_size", local_world_size, topology.local_world_size), + _check("pipeline_parallel_size", 1, topology.pipeline_parallel_size), + _check("tensor_parallel_size", 1, topology.tensor_parallel_size), + _check("ringattn_parallel_size", 1, topology.ringattn_parallel_size), + _check("ulysses_parallel_size", 1, topology.ulysses_parallel_size), + _check( + "sample_packing_sequence_len", + readme_metrics.get("sample_packing_sequence_len"), + topology.sample_packing_sequence_len, + ), + _check("micro_batch_size", 8, topology.micro_batch_size), + _check("global_batch_size", 256, topology.global_batch_size), + _check("data_parallel_replicate_size", 1, topology.data_parallel_replicate_size), + _check("expert_parallel_size", 8, topology.expert_parallel_size), + _check("ep_fsdp_size", 4, topology.ep_fsdp_size), + _check("data_parallel_shard_size", 32, topology.data_parallel_shard_size), + _check("num_experts", 256, topology.num_experts), + _check("top_k", 8, topology.top_k), + _check("dataset_path", "dummy", data.get("datasets", [{}])[0].get("path")), + _check("dataset_type", "tokenized", data.get("datasets", [{}])[0].get("type")), + _check("sample_packing_method", "sequential", data.get("sample_packing_method")), + _check("moe_implementation", "quack", model.get("moe_implementation")), + _check("ep_dispatch", "deepep", model.get("ep_dispatch")), + _check("train_router", False, model.get("train_router")), + _check("deepep_buffer_size_gb", 2.0, model.get("deepep_buffer_size_gb")), + _check("deepep_num_sms", 72, model.get("deepep_num_sms")), + _check("deepep_async_combine", True, model.get("deepep_async_combine")), + _check("data_parallel_mode", "fsdp2", train.get("data_parallel_mode")), + _check("gradient_checkpointing_method", "recompute_full_layer", train.get("gradient_checkpointing_method")), + _check("optimizer", "adamw", train.get("optimizer")), + _check("enable_mixed_precision", True, train.get("enable_mixed_precision")), + _check("enable_full_shard", True, train.get("enable_full_shard")), + _check("init_device", "meta", train.get("init_device")), + _check("load_weights_mode", "grouped", train.get("load_weights_mode")), + _check("enable_compile", True, train.get("enable_compile")), + _check("empty_cache_steps", 10, train.get("empty_cache_steps")), + _check("gc_steps", 10, train.get("gc_steps")), + _check("save_steps", 0, train.get("save_steps")), + _check("save_epochs", 0, train.get("save_epochs")), + _check("log_format", "structured", train.get("log_format")), + _check("global_tokens_per_train_step", 2_097_408, shape.global_tokens_per_train_step), + ] + behavior_points = load_benchmark_behavior_points(benchmark_dir) + behavior_prediction = predict_benchmark_behavior(behavior_points, topology, shape, raw_config) + behavior_labels = sorted(point.label for point in behavior_points) + behavior_predictions = _predict_all_behavior_points(behavior_points, topology) + checks.extend( + [ + _check( + "benchmark_behavior_points", + [ + "qwen36_static_k3_summary_20260519:q36-main-af98064-deepepenv-05190533", + "readme_adjacent_mbs10_allocator_pressure", + "readme_reference_mbs8", + ], + behavior_labels, + ), + _check("benchmark_behavior_prediction_status", "calibrated", behavior_prediction.status), + _check("benchmark_behavior_prediction_label", "readme_reference_mbs8", behavior_prediction.matched_label), + _check( + "benchmark_behavior_tokens_per_sec", + readme_metrics.get("tokens_per_sec"), + behavior_prediction.tokens_per_sec, + ), + _check( + "benchmark_behavior_step_time_sec", + readme_metrics.get("step_time_sec"), + behavior_prediction.step_time_sec, + ), + _check( + "benchmark_behavior_peak_mem_gb", + readme_metrics.get("allocated_memory_gb"), + behavior_prediction.peak_mem_gb, + ), + _check("benchmark_behavior_allocator_retries", 0, behavior_prediction.allocator_retries), + _check( + "benchmark_behavior_promised_tflops_per_gpu", + H100_BF16_PROMISED_TFLOPS_PER_GPU, + behavior_prediction.promised_tflops_per_gpu, + ), + _check("benchmark_behavior_tflops_per_gpu", 160.218, behavior_prediction.tflops_per_gpu, tolerance=0.001), + ] + ) + for label, prediction in behavior_predictions.items(): + point = next(point for point in behavior_points if point.label == label) + variant_shape = prediction["shape"] + variant_behavior = prediction["behavior"] + checks.extend( + [ + _check(f"behavior_matrix:{label}:prediction_status", "calibrated", variant_behavior["status"]), + _check(f"behavior_matrix:{label}:prediction_label", label, variant_behavior["matched_label"]), + _check( + f"behavior_matrix:{label}:tokens_per_sec", point.tokens_per_sec, variant_behavior["tokens_per_sec"] + ), + _check( + f"behavior_matrix:{label}:global_tokens_per_step", + (point.global_batch_size or 0) * (topology.sample_packing_sequence_len or 0), + variant_shape["global_tokens_per_train_step"], + ), + ] + ) + if point.step_time_sec is not None: + checks.append( + _check( + f"behavior_matrix:{label}:step_time_sec", + point.step_time_sec, + variant_behavior["step_time_sec"], + tolerance=0.01, + ) + ) + if point.mfu_percent is not None: + checks.extend( + [ + _check( + f"behavior_matrix:{label}:mfu_percent", + point.mfu_percent, + variant_behavior["mfu_percent"], + ), + _check( + f"behavior_matrix:{label}:promised_tflops_per_gpu", + H100_BF16_PROMISED_TFLOPS_PER_GPU, + variant_behavior["promised_tflops_per_gpu"], + ), + _check( + f"behavior_matrix:{label}:tflops_per_gpu", + H100_BF16_PROMISED_TFLOPS_PER_GPU * point.mfu_percent / 100.0, + variant_behavior["tflops_per_gpu"], + tolerance=0.001, + ), + ] + ) + if point.tokens_per_sec and variant_shape["global_tokens_per_train_step"]: + checks.append( + _check( + f"behavior_matrix:{label}:tokens_imply_step_time", + variant_behavior["step_time_sec"], + variant_shape["global_tokens_per_train_step"] / point.tokens_per_sec, + tolerance=0.05, + ) + ) + if readme_metrics.get("tokens_per_sec") and shape.global_tokens_per_train_step: + derived_step_time = shape.global_tokens_per_train_step / readme_metrics["tokens_per_sec"] + checks.append( + _check( + "readme_tokens_per_sec_implies_step_time", + readme_metrics.get("step_time_sec"), + derived_step_time, + tolerance=0.05, + ) + ) + if report.shape.balanced_routing is not None: + counts = report.shape.balanced_routing.counts_by_expert + checks.extend( + [ + _check("balanced_routing_imbalance_slots", 1, report.shape.balanced_routing.imbalance_slots), + _check("balanced_routing_count_sum", report.shape.balanced_routing.total_slots, sum(counts)), + _check("experts_per_ep_rank", 32, topology.num_experts // topology.expert_parallel_size), + ] + ) + script_text = _render_script_text(benchmark_dir) + manifest_text = _render_manifest_text(benchmark_dir) + checks.append( + _check( + "render_script_sets_balanced_synthetic_routing_env", + True, + "XORL_MOE_SYNTHETIC_ROUTING" in script_text and "balanced" in script_text, + ) + ) + checks.extend( + [ + _check("rendered_manifest_sets_team_turbo", True, "team: turbo" in manifest_text), + _check( + "rendered_manifest_sets_balanced_synthetic_routing_env", + True, + "name: XORL_MOE_SYNTHETIC_ROUTING" in manifest_text and 'value: "balanced"' in manifest_text, + ), + _check( + "rendered_manifest_sets_nccl_socket_ifname", + True, + "name: NCCL_SOCKET_IFNAME" in manifest_text and 'value: "bond0"' in manifest_text, + ), + _check("rendered_manifest_sets_runtime_class", True, "runtimeClassName: nvidia" in manifest_text), + ] + ) + return to_jsonable(report), checks, behavior_predictions + + +def _predict_all_behavior_points(behavior_points, topology) -> dict[str, dict[str, Any]]: + predictions: dict[str, dict[str, Any]] = {} + for point in behavior_points: + if not point.micro_batch_size or not point.global_batch_size: + continue + denom = point.micro_batch_size * topology.data_parallel_size + if point.global_batch_size % denom != 0: + continue + variant_topology = replace( + topology, + micro_batch_size=point.micro_batch_size, + gradient_accumulation_steps=point.global_batch_size // denom, + global_batch_size=point.global_batch_size, + ) + variant_shape = build_shape_ledger(variant_topology, balanced_routing=True) + predictions[point.label] = { + "shape": to_jsonable(variant_shape), + "behavior": to_jsonable(predict_benchmark_behavior(behavior_points, variant_topology, variant_shape)), + } + return predictions + + +def _validate_result_json(benchmark_dir: Path, readme_metrics: dict[str, Any]) -> list[dict[str, Any]]: + checks: list[dict[str, Any]] = [] + seq_len = int(readme_metrics.get("sample_packing_sequence_len") or 0) + for result_path in sorted((benchmark_dir / "results").glob("*.json")): + result = json.loads(result_path.read_text(encoding="utf-8")) + throughput = result.get("throughput", {}) + if throughput and seq_len: + global_tokens = throughput.get("global_batch_size", 0) * seq_len + derived_step_time = ( + global_tokens / throughput["tokens_per_sec"] if throughput.get("tokens_per_sec") else None + ) + checks.extend( + [ + _check( + f"{result_path.name}:throughput_candidate", + "q36-main-af98064-deepepenv-05190533", + throughput.get("candidate"), + ), + _check(f"{result_path.name}:throughput_gpus", 32, throughput.get("gpus")), + _check(f"{result_path.name}:throughput_tokens_per_sec", 254600.0, throughput.get("tokens_per_sec")), + _check( + f"{result_path.name}:derived_step_time_sec", + throughput.get("step_time_sec"), + derived_step_time, + tolerance=0.01, + ), + ] + ) + k3_gate = result.get("k3_gate", {}) + if k3_gate: + checks.extend( + [ + _check(f"{result_path.name}:k3_gate_status", "fail", k3_gate.get("status")), + _check(f"{result_path.name}:k3_total_tokens", 192, k3_gate.get("k3", {}).get("total_tokens")), + _check( + f"{result_path.name}:k3_primary_failure", "k3.mean <= 0.001", k3_gate.get("primary_failure") + ), + ] + ) + diagnostics = result.get("diagnostic_replays", []) + if diagnostics: + checks.append(_check(f"{result_path.name}:diagnostic_replay_count", 3, len(diagnostics))) + checks.append( + _check( + f"{result_path.name}:diagnostic_low_k3_rows", + 2, + sum(1 for row in diagnostics if row.get("status") == "diagnostic_low_k3"), + ) + ) + return checks + + +def validate_benchmark_dir(benchmark_dir: Path) -> dict[str, Any]: + readme_path = benchmark_dir / "README.md" + if not readme_path.is_file(): + raise FileNotFoundError(f"missing README.md in {benchmark_dir}") + readme_metrics = _parse_readme_metrics(readme_path.read_text(encoding="utf-8")) + report, config_checks, behavior_predictions = _validate_config_behavior(benchmark_dir, readme_metrics) + result_checks = _validate_result_json(benchmark_dir, readme_metrics) + checks = config_checks + result_checks + status = "pass" if all(check["status"] == "pass" for check in checks) else "fail" + return { + "benchmark_dir": str(benchmark_dir), + "status": status, + "readme_metrics": readme_metrics, + "simulator_report": report, + "behavior_points": to_jsonable(load_benchmark_behavior_points(benchmark_dir)), + "behavior_predictions": behavior_predictions, + "checks": checks, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--benchmarks-root", type=Path, required=True) + parser.add_argument("--model", required=True, help="Benchmark model subdirectory to validate") + parser.add_argument("--output", type=Path, default=None) + parser.add_argument("--no-fail-on-error", action="store_true", help="Always exit 0 after writing the report") + args = parser.parse_args() + + payload = validate_benchmark_dir(args.benchmarks_root / args.model) + rendered = json.dumps(payload, indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + if payload["status"] != "pass" and not args.no_fail_on_error: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/experiments/test_training_sim.py b/tests/experiments/test_training_sim.py new file mode 100644 index 00000000..dcaf145a --- /dev/null +++ b/tests/experiments/test_training_sim.py @@ -0,0 +1,624 @@ +from pathlib import Path + +import yaml + +from experiments.local_benchmark.training_sim.benchmark_behavior import load_benchmark_behavior_points +from experiments.local_benchmark.training_sim.calibration_evaluator import evaluate_calibration +from experiments.local_benchmark.training_sim.collect_calibration import parse_log_text, summarize_observed_run +from experiments.local_benchmark.training_sim.config_fingerprint import build_fingerprint, resolve_topology +from experiments.local_benchmark.training_sim.model_metadata import resolve_model_metadata +from experiments.local_benchmark.training_sim.scenario_planner import plan_scenario +from experiments.local_benchmark.training_sim.shape_engine import balanced_counts, build_shape_ledger + + +def test_balanced_counts_round_robin_distribution() -> None: + assert balanced_counts(20, 6) == [4, 4, 3, 3, 3, 3] + + +def test_resolve_topology_matches_training_arguments_dp_formula() -> None: + raw_config = { + "train": { + "micro_batch_size": 1, + "gradient_accumulation_steps": 16, + "ulysses_parallel_size": 4, + "tensor_parallel_size": 1, + "ringattn_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 8, + "data_parallel_replicate_size": 2, + }, + "data": {"sample_packing_sequence_len": 2048}, + "model": {"num_experts": 16, "num_experts_per_tok": 4}, + } + + topology = resolve_topology(raw_config, world_size=16, local_world_size=8) + + assert topology.data_parallel_size == 4 + assert topology.data_parallel_replicate_size == 2 + assert topology.data_parallel_shard_size == 2 + assert topology.global_batch_size == 64 + assert topology.sequence_parallel_size == 4 + assert topology.ep_fsdp_size == 2 + assert topology.num_experts == 16 + assert topology.top_k == 4 + + +def test_shape_ledger_uses_sequence_parallel_local_tokens() -> None: + raw_config = { + "train": { + "micro_batch_size": 1, + "gradient_accumulation_steps": 2, + "ulysses_parallel_size": 4, + "expert_parallel_size": 4, + }, + "data": {"sample_packing_sequence_len": 2048}, + "model": {"num_experts": 16, "num_experts_per_tok": 4}, + } + topology = resolve_topology(raw_config, world_size=16, local_world_size=8) + + ledger = build_shape_ledger(topology, balanced_routing=True) + + assert ledger.microbatch_tokens_per_dp_rank == 2048 + assert ledger.global_tokens_per_microbatch == 8192 + assert ledger.global_tokens_per_train_step == 16384 + assert ledger.tokens_per_gpu_per_train_step == 1024 + assert ledger.tokens_per_model_rank_per_microbatch == 512 + assert ledger.routed_slots_per_model_rank_microbatch == 2048 + assert ledger.routed_slots_per_train_step_model_rank == 4096 + assert ledger.balanced_routing is not None + assert ledger.balanced_routing.counts_by_expert == [128] * 16 + assert ledger.ep_rank_slots_per_microbatch == [512, 512, 512, 512] + + +def test_parse_structured_step_phase_and_memory_logs() -> None: + log_text = """ + [STEP 4/9] loss=1.0 grad_norm=0.1 lr=1.0e-5 tflops=100.2 mfu=0.1010 tokens_per_sec=53414 time=72.100s peak_mem=39.8GB fwd=20.1GB bwd=39.8GB optim=10.0GB + [STEP_PHASES 4/9] dataloader_max_s=0.100000 dataloader_mean_s=0.050000 model_forward_max_s=10.000000 model_forward_mean_s=9.000000 + [STEP_MEMORY 4/9] model_forward_after_allocated_max_gb=20.100 model_forward_phase_peak_allocated_max_gb=39.800 + """ + + observed = parse_log_text(log_text, source="sample.log") + summary = summarize_observed_run(observed, warmup_steps=0, world_size=16) + + assert len(observed.steps) == 1 + assert observed.steps[0].tokens_per_sec == 53414 + assert observed.steps[0].phase_memory_gb == {"fwd": 20.1, "bwd": 39.8, "optim": 10.0} + assert observed.phases[0].metrics["model_forward_max_s"] == 10.0 + assert observed.memory_phases[0].metrics["model_forward_phase_peak_allocated_max_gb"] == 39.8 + assert summary["tokens_per_sec_per_gpu_mean"] == 3338.375 + + +def _write_resolved_run_fixture(root: Path) -> Path: + config = { + "model": { + "model_path": "Qwen/Qwen3.6-35B-A3B", + "deepep_async_combine": False, + "deepep_num_sms": 24, + "deepep_buffer_size_gb": 1.0, + }, + "data": {"sample_packing_sequence_len": 1024}, + "train": { + "data_parallel_mode": "fsdp2", + "data_parallel_replicate_size": 1, + "data_parallel_shard_size": 4, + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "ulysses_parallel_size": 1, + "ringattn_parallel_size": 1, + "expert_parallel_size": 2, + "micro_batch_size": 2, + "gradient_accumulation_steps": 3, + "enable_compile": True, + "gradient_checkpointing_method": "recompute_full_layer", + "enable_activation_offload": True, + "activation_offload_prefetch_count": 4, + "optimizer": "muon", + "optimizer_dtype": "bf16", + "muon_momentum": 0.0, + }, + } + run_dir = root / "resolved" / "fit" + run_dir.mkdir(parents=True) + config_path = run_dir / "xorl_cli.yaml" + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + (run_dir / "startup_metrics.json").write_text( + """ + { + "repo_commit": "abc123", + "metrics": { + "startup/master_addr": "fit-master", + "startup/node_count": 1, + "startup/total_train_steps": 4 + } + } + """, + encoding="utf-8", + ) + log_dir = root / "fit" + log_dir.mkdir() + (log_dir / "node-0.log").write_text( + """ + [STEP 1/4] loss=1.0 grad_norm=0.1 lr=1.0e-4 tflops=1.0 mfu=0.001 tokens_per_sec=100 time=10.0s peak_mem=60.0GB + [STEP 2/4] loss=1.0 grad_norm=0.1 lr=1.0e-4 tflops=10.0 mfu=0.010 tokens_per_sec=1000 time=8.0s peak_mem=68.0GB + [STEP 3/4] loss=1.0 grad_norm=0.1 lr=1.0e-4 tflops=12.0 mfu=0.012 tokens_per_sec=1200 time=7.0s peak_mem=70.0GB + [STEP 4/4] loss=1.0 grad_norm=0.1 lr=1.0e-4 tflops=10.0 mfu=0.010 tokens_per_sec=1000 time=9.0s peak_mem=69.0GB + """, + encoding="utf-8", + ) + + oom_config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + oom_config["train"]["micro_batch_size"] = 3 + oom_dir = root / "resolved" / "oom" + oom_dir.mkdir() + (oom_dir / "xorl_cli.yaml").write_text(yaml.safe_dump(oom_config), encoding="utf-8") + (oom_dir / "startup_metrics.json").write_text( + """ + { + "repo_commit": "abc123", + "metrics": { + "startup/master_addr": "oom-master", + "startup/node_count": 1 + } + } + """, + encoding="utf-8", + ) + oom_log_dir = root / "oom" + oom_log_dir.mkdir() + (oom_log_dir / "node-0.log").write_text( + """ + [rank0]: Traceback (most recent call last): + [rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB. GPU 0 has a total capacity of 79.18 GiB of which 59.62 MiB is free. Including non-PyTorch memory, this process has 79.09 GiB memory in use. + torch.distributed.elastic.multiprocessing.errors.ChildFailedError: + """, + encoding="utf-8", + ) + return config_path + + +def test_benchmark_behavior_loader_ingests_resolved_run_logs_and_ooms(tmp_path: Path) -> None: + config_path = _write_resolved_run_fixture(tmp_path) + + points = load_benchmark_behavior_points(tmp_path) + by_label = {point.label: point for point in points} + + fit = by_label[f"resolved_run:{config_path.parent.relative_to(tmp_path)}"] + assert fit.status == "observed_log_summary" + assert fit.correctness_status == "not_promoted" + assert fit.tokens_per_sec == 1_100.0 + assert fit.step_time_sec == 8.0 + assert fit.tflops_per_gpu == 11.0 + assert fit.mfu_percent == 1.1 + assert fit.peak_mem_gb == 70.0 + assert fit.measured_steps == 2 + assert fit.warmup_steps == 2 + assert fit.gpu_count == 4 + assert fit.global_batch_size == 24 + assert fit.expert_parallel_size == 2 + assert fit.ep_fsdp_size == 2 + assert fit.deepep_num_sms == 24 + assert fit.enable_activation_offload is True + assert fit.activation_offload_prefetch_count == 4 + + oom = by_label["resolved_run:resolved/oom"] + assert oom.status == "observed_log_oom" + assert oom.correctness_status == "oom" + assert oom.tokens_per_sec is None + assert oom.peak_mem_gb == 79.09 + assert oom.micro_batch_size == 3 + + +def test_scenario_planner_keeps_observed_fit_feasible_when_safety_margin_is_tight(tmp_path: Path) -> None: + config_path = _write_resolved_run_fixture(tmp_path) + + report = plan_scenario( + config_path, + benchmark_dir=tmp_path, + micro_batch_sizes=[2], + expert_parallel_sizes=[2], + device_memory_limit_gb=80.0, + memory_safety_factor=1.15, + ) + + assert report.candidate_count == 1 + assert report.feasible_count == 1 + assert report.best_raw is not None + assert report.best_raw.label == "mbs2-gb24-ep2-efsdp2-tp1-pp1-u1-r1:resolved_run:resolved/fit" + assert report.best_raw.score_tokens_per_sec == 1_100.0 + assert report.best_raw.feasibility_status == "feasible_calibrated_peak_high_pressure" + assert report.best_raw.memory_headroom_gb == -0.5 + assert report.best_raw.recommendation == "correctness_gate_required" + + +def test_build_fingerprint_reads_config_file(tmp_path: Path) -> None: + config_path = tmp_path / "config.yaml" + config = { + "train": { + "micro_batch_size": 2, + "gradient_accumulation_steps": 3, + "ulysses_parallel_size": 2, + "data_parallel_shard_size": 4, + }, + "data": {"sample_packing_sequence_len": 1024}, + } + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + + fingerprint = build_fingerprint( + config_path, + world_size=8, + local_world_size=8, + balanced_routing=True, + num_experts=8, + top_k=2, + ) + + assert fingerprint.config_name == "config.yaml" + assert fingerprint.balanced_routing is True + assert fingerprint.topology.data_parallel_size == 4 + assert fingerprint.topology.data_parallel_replicate_size == 1 + assert fingerprint.topology.data_parallel_shard_size == 4 + assert fingerprint.topology.global_batch_size == 24 + assert len(fingerprint.config_sha256) == 64 + + +def test_resolve_model_metadata_from_hf_cache(tmp_path: Path) -> None: + config_dir = tmp_path / "models--Example--MoE" / "snapshots" / "abc123" + config_dir.mkdir(parents=True) + (config_dir / "config.json").write_text( + """ + { + "model_type": "example", + "text_config": { + "num_experts": 12, + "num_experts_per_tok": 3, + "num_hidden_layers": 7, + "hidden_size": 128, + "moe_intermediate_size": 32, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 32, + "vocab_size": 4096, + "tie_word_embeddings": false + } + } + """, + encoding="utf-8", + ) + + metadata = resolve_model_metadata({"model": {"model_path": "Example/MoE"}}, hf_cache_roots=[tmp_path]) + + assert metadata.source == "hf_config" + assert metadata.num_experts == 12 + assert metadata.top_k == 3 + assert metadata.num_hidden_layers == 7 + assert metadata.moe_intermediate_size == 32 + assert metadata.num_attention_heads == 4 + assert metadata.num_key_value_heads == 2 + assert metadata.head_dim == 32 + assert metadata.tie_word_embeddings is False + assert metadata.config_path is not None + + +def test_resolve_known_qwen235_metadata_without_hf_cache() -> None: + metadata = resolve_model_metadata( + {"model": {"model_path": "Qwen/Qwen3-235B-A22B"}}, + hf_cache_roots=[], + ) + + assert metadata.source == "known_model" + assert metadata.num_experts == 128 + assert metadata.top_k == 8 + assert metadata.num_hidden_layers == 94 + assert metadata.hidden_size == 4096 + assert metadata.moe_intermediate_size == 1536 + assert metadata.num_attention_heads == 64 + assert metadata.num_key_value_heads == 4 + assert metadata.head_dim == 128 + assert metadata.vocab_size == 151936 + + +def _write_q235_results_fixture(benchmark_dir: Path) -> None: + benchmark_dir.mkdir() + (benchmark_dir / "RESULTS.md").write_text( + """ +# Qwen3-235B-A22B @ 2k context + +Measured: 4 nodes / 32xH100, u1/dp_shard32/EP8/ep_fsdp4. + +| run | gcm | pack | mbs | tok/step | step s | MFU | tok/s tot | tok/s/GPU | peak GB | status | +|-----|-----|-----:|----:|---------:|-------:|----:|----------:|----------:|--------:|--------| +| n4_ep8_bd_pk4096 | before_dispatch | 4096 | 1 | 131,072 | ~18.4 | **~3.0%** | **~6,800** | ~213 | 68.3 | OK | +| n4_ep8_bd_pk4096_ga2 | before_dispatch | 4096 | 1 | 262,144 | ~31.3 | **~3.7%** | **~8,400** | ~263 | 68.3 | NEW BEST | +| n4_ep8_bd_pk16k | before_dispatch | 16384 | 1 | 524,288 | -- | -- | -- | -- | OOM | FAIL | +""", + encoding="utf-8", + ) + + +def test_qwen235_markdown_loader_extracts_pack_and_ga_rows(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + + points = load_benchmark_behavior_points(benchmark_dir) + by_label = {point.label: point for point in points} + + assert set(by_label) == { + "q235_markdown:n4_ep8_bd_pk4096", + "q235_markdown:n4_ep8_bd_pk4096_ga2", + "q235_markdown:n4_ep8_bd_pk16k", + } + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].global_batch_size == 32 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].tokens_per_sec == 6_800.0 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].mfu_percent == 3.0 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].gpu_count == 32 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].sample_packing_sequence_len == 4096 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].tensor_parallel_size == 1 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].pipeline_parallel_size == 1 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].ulysses_parallel_size == 1 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].ringattn_parallel_size == 1 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].expert_parallel_size == 8 + assert by_label["q235_markdown:n4_ep8_bd_pk4096"].ep_fsdp_size == 4 + assert by_label["q235_markdown:n4_ep8_bd_pk4096_ga2"].global_batch_size == 64 + assert by_label["q235_markdown:n4_ep8_bd_pk4096_ga2"].tokens_per_sec == 8_400.0 + assert by_label["q235_markdown:n4_ep8_bd_pk16k"].sample_packing_sequence_len == 16384 + assert by_label["q235_markdown:n4_ep8_bd_pk16k"].tokens_per_sec is None + assert by_label["q235_markdown:n4_ep8_bd_pk16k"].correctness_status == "oom" + + +def _write_q235_config_fixture(config_path: Path) -> None: + config = { + "model": { + "model_path": "Qwen/Qwen3-235B-A22B", + "ep_dispatch": "deepep", + "deepep_buffer_size_gb": 2.0, + }, + "data": { + "sample_packing_sequence_len": 4096, + }, + "train": { + "data_parallel_mode": "fsdp2", + "ulysses_parallel_size": 1, + "ringattn_parallel_size": 1, + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "expert_parallel_size": 8, + "data_parallel_replicate_size": 1, + "data_parallel_shard_size": 32, + "micro_batch_size": 1, + "gradient_accumulation_steps": 1, + "optimizer": "muon", + "optimizer_dtype": "bf16", + "muon_momentum": 0.0, + "enable_mixed_precision": True, + }, + } + config_path.write_text(yaml.safe_dump(config), encoding="utf-8") + + +def test_qwen235_scenario_planner_uses_markdown_calibration_for_ga_tradeoff(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + config_path = tmp_path / "q235.yaml" + _write_q235_config_fixture(config_path) + + report = plan_scenario( + config_path, + benchmark_dir=benchmark_dir, + world_size=32, + local_world_size=8, + micro_batch_sizes=[1], + gradient_accumulation_steps=[1, 2], + expert_parallel_sizes=[8], + ) + + assert report.candidate_count == 2 + assert report.best_raw is not None + assert report.best_raw.label == "mbs1-gb64-ep8-efsdp4-tp1-pp1-u1-r1:q235_markdown:n4_ep8_bd_pk4096_ga2" + assert report.best_raw.prediction_confidence == "calibrated" + assert report.best_raw.score_tokens_per_sec == 8_400.0 + assert report.best_raw.behavior.tokens_per_sec_per_gpu == 262.5 + assert report.best_raw.analytic_peak_floor_gb == 29.363 + assert report.best_raw.estimated_peak_mem_gb == 68.3 + assert report.best_raw.memory_basis == "calibrated_peak" + assert report.best_raw.feasibility_status == "feasible_calibrated_peak_high_pressure" + assert report.best_promotable is None + + +def test_qwen235_scenario_planner_extrapolates_ga_asymptote_from_step_time_fit(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + config_path = tmp_path / "q235.yaml" + _write_q235_config_fixture(config_path) + + report = plan_scenario( + config_path, + benchmark_dir=benchmark_dir, + world_size=32, + local_world_size=8, + micro_batch_sizes=[1], + gradient_accumulation_steps=[1, 2, 4, 8], + expert_parallel_sizes=[8], + ) + + assert report.candidate_count == 4 + assert report.best_raw is not None + assert report.best_raw.label == "mbs1-gb256-ep8-efsdp4-tp1-pp1-u1-r1:extrapolated" + assert report.best_raw.prediction_confidence == "extrapolated_step_time_fit" + assert report.best_raw.score_tokens_per_sec == 9646.513 + assert report.best_raw.behavior.step_time_sec == 108.7 + assert report.best_raw.estimated_peak_mem_gb == 68.3 + assert report.best_raw.memory_basis == "extrapolated_peak" + assert report.best_raw.feasibility_status == "feasible_extrapolated_peak_high_pressure" + assert report.best_raw.calibration_scope == "outside_measured_envelope" + assert report.best_raw.score_risk_adjusted_tokens_per_sec == 4796.729 + assert report.best_raw.recommendation == "remeasure_before_ranking" + assert "outside_measured_envelope" in report.best_raw.risk_flags + assert "requires_remeasurement" in report.best_raw.risk_flags + assert report.best_raw.promotable is False + assert report.best_risk_adjusted is not None + assert report.best_risk_adjusted.label == ("mbs1-gb64-ep8-efsdp4-tp1-pp1-u1-r1:q235_markdown:n4_ep8_bd_pk4096_ga2") + assert report.best_risk_adjusted.score_risk_adjusted_tokens_per_sec == 6783.0 + assert report.best_next_measurement is not None + assert report.best_next_measurement.label == "mbs1-gb256-ep8-efsdp4-tp1-pp1-u1-r1:extrapolated" + by_label = {candidate.label: candidate for candidate in report.candidates} + ga4 = by_label["mbs1-gb128-ep8-efsdp4-tp1-pp1-u1-r1:extrapolated"] + assert ga4.prediction_confidence == "extrapolated_step_time_fit" + assert ga4.score_tokens_per_sec == 9181.926 + + +def test_qwen235_calibration_evaluator_reports_leave_one_out_ga_error(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + config_path = tmp_path / "q235.yaml" + _write_q235_config_fixture(config_path) + + report = evaluate_calibration( + config_path, + benchmark_dir=benchmark_dir, + world_size=32, + local_world_size=8, + ) + + assert report.status == "ok" + assert report.measured_point_count == 2 + assert report.evaluated_count == 2 + assert report.skipped_count == 0 + assert report.prediction_status_counts == {"extrapolated": 2} + assert report.mean_absolute_percentage_error == 21.288 + by_label = {holdout.label: holdout for holdout in report.holdouts} + ga1 = by_label["q235_markdown:n4_ep8_bd_pk4096"] + ga2 = by_label["q235_markdown:n4_ep8_bd_pk4096_ga2"] + assert ga1.topology_label == "mbs1-gb32-ep8-efsdp4-tp1-pp1-u1-r1" + assert ga1.predicted_tokens_per_sec == 8_400.0 + assert ga1.absolute_percentage_error == 23.529 + assert ga2.predicted_tokens_per_sec == 6_800.0 + assert ga2.absolute_percentage_error == 19.048 + + +def test_qwen235_scenario_planner_does_not_exact_match_observed_row_to_tp_what_if(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + config_path = tmp_path / "q235.yaml" + _write_q235_config_fixture(config_path) + + report = plan_scenario( + config_path, + benchmark_dir=benchmark_dir, + world_size=32, + local_world_size=8, + micro_batch_sizes=[1], + gradient_accumulation_steps=[2], + expert_parallel_sizes=[8], + tensor_parallel_sizes=[1, 2], + ) + + by_label = {candidate.label: candidate for candidate in report.candidates} + exact = by_label["mbs1-gb64-ep8-efsdp4-tp1-pp1-u1-r1:q235_markdown:n4_ep8_bd_pk4096_ga2"] + tp2 = by_label["mbs1-gb32-ep8-efsdp4-tp2-pp1-u1-r1:extrapolated"] + assert exact.prediction_confidence == "calibrated" + assert exact.score_tokens_per_sec == 8_400.0 + assert tp2.prediction_confidence == "extrapolated" + assert tp2.behavior.matched_label == "q235_markdown:n4_ep8_bd_pk4096_ga2" + assert "TP extrapolation uses conservative communication penalty" in tp2.behavior.warnings + assert tp2.score_tokens_per_sec == 7_560.0 + + +def test_qwen235_scenario_planner_auto_sweeps_parallelism_strategy_space(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + config_path = tmp_path / "q235.yaml" + _write_q235_config_fixture(config_path) + + report = plan_scenario( + config_path, + benchmark_dir=benchmark_dir, + world_size=32, + local_world_size=8, + micro_batch_sizes=[1], + gradient_accumulation_steps=[2], + topology_sweep="auto", + ) + + assert report.topology_sweep == "auto" + assert {1, 2, 4, 8}.issubset({candidate.topology.tensor_parallel_size for candidate in report.candidates}) + assert 2 in {candidate.topology.pipeline_parallel_size for candidate in report.candidates} + assert {8, 16, 32}.issubset({candidate.topology.expert_parallel_size for candidate in report.candidates}) + assert report.best_raw is not None + assert report.best_raw.label == "mbs1-gb64-ep8-efsdp4-tp1-pp1-u1-r1:q235_markdown:n4_ep8_bd_pk4096_ga2" + tp_candidates = [ + candidate + for candidate in report.candidates + if candidate.topology.tensor_parallel_size == 2 + and candidate.topology.pipeline_parallel_size == 1 + and candidate.topology.expert_parallel_size == 8 + ] + assert tp_candidates + assert all(candidate.prediction_confidence == "extrapolated" for candidate in tp_candidates) + + +def test_qwen235_auto_sweep_includes_long_context_cp_without_cross_seq_calibration(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + config_path = tmp_path / "q235.yaml" + _write_q235_config_fixture(config_path) + raw_config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + raw_config["data"]["sample_packing_sequence_len"] = 64_000 + config_path.write_text(yaml.safe_dump(raw_config), encoding="utf-8") + + report = plan_scenario( + config_path, + benchmark_dir=benchmark_dir, + world_size=32, + local_world_size=8, + micro_batch_sizes=[1], + gradient_accumulation_steps=[1], + expert_parallel_sizes=[8], + tensor_parallel_sizes=[1], + pipeline_parallel_sizes=[1], + topology_sweep="auto", + ) + + cp_pairs = { + (candidate.topology.ulysses_parallel_size, candidate.topology.ringattn_parallel_size) + for candidate in report.candidates + } + assert {(2, 1), (4, 1), (1, 2)}.issubset(cp_pairs) + assert report.feasible_count == 0 + assert {candidate.prediction_confidence for candidate in report.candidates} == {"unscored"} + assert {candidate.calibration_scope for candidate in report.candidates} == {"outside_sequence_calibration_envelope"} + base_cp = next( + candidate + for candidate in report.candidates + if candidate.topology.ulysses_parallel_size == 1 and candidate.topology.ringattn_parallel_size == 1 + ) + assert "observed_oom_boundary:q235_markdown:n4_ep8_bd_pk16k" in base_cp.risk_flags + + +def test_qwen235_scenario_planner_marks_matching_oom_pack_infeasible(tmp_path: Path) -> None: + benchmark_dir = tmp_path / "q235" + _write_q235_results_fixture(benchmark_dir) + config_path = tmp_path / "q235.yaml" + _write_q235_config_fixture(config_path) + raw_config = yaml.safe_load(config_path.read_text(encoding="utf-8")) + raw_config["data"]["sample_packing_sequence_len"] = 16384 + config_path.write_text(yaml.safe_dump(raw_config), encoding="utf-8") + + report = plan_scenario( + config_path, + benchmark_dir=benchmark_dir, + world_size=32, + local_world_size=8, + micro_batch_sizes=[1], + gradient_accumulation_steps=[1], + expert_parallel_sizes=[8], + ) + + assert report.candidate_count == 1 + assert report.feasible_count == 0 + candidate = report.candidates[0] + assert candidate.label == "mbs1-gb32-ep8-efsdp4-tp1-pp1-u1-r1:q235_markdown:n4_ep8_bd_pk16k" + assert candidate.behavior.status == "calibrated_failure" + assert candidate.feasibility_status == "observed_oom" + assert candidate.score_tokens_per_sec is None + assert candidate.calibration_scope == "exact_calibrated" + assert "observed_oom_boundary:q235_markdown:n4_ep8_bd_pk16k" in candidate.risk_flags From 54e0c7911d33b4f9d8c83c1eca5276c7cbb11605 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Jul 2026 03:08:07 +0000 Subject: [PATCH 2/4] Package simulator with portable Qwen calibrations --- .../local_benchmark/training_sim/README.md | 185 - .../local_benchmark/training_sim/__init__.py | 1 - .../training_sim/benchmark_behavior.py | 1026 -- .../training_sim/calibration_evaluator.py | 256 - .../training_sim/collect_calibration.py | 224 - .../training_sim/k8s/README.md | 13 - .../training_sim/memory_ledger.py | 246 - .../training_sim/scenario_planner.py | 1024 -- .../local_benchmark/training_sim/schemas.py | 322 - .../training_sim/validate_benchmarks.py | 457 - pyproject.toml | 16 + src/xorl/sim/README.md | 137 + src/xorl/sim/__init__.py | 1 + src/xorl/sim/analytical_ledgers.py | 1510 +++ src/xorl/sim/benchmark_behavior.py | 2544 +++++ src/xorl/sim/calibration_evaluator.py | 1972 ++++ src/xorl/sim/calibration_packs.py | 175 + .../qwen3_235b_a22b/README.md | 4 + .../qwen3_235b_a22b/RESULTS.md | 9 + .../configs/qwen3_235b_a22b_2k_4node_ep8.yaml | 51 + ...35b_a22b_muon_8node_ep8_efsdp8_deepep.yaml | 56 + .../qwen3_235b_a22b/manifest.json | 30 + .../qwen3_5_397b_a17b/README.md | 21 + .../configs/qwen35_ep64_reference.yaml | 67 + .../configs/qwen35_r69_mbs4.yaml | 72 + .../configs/qwen35_r73_sync.yaml | 72 + .../configs/qwen35_r75_async.yaml | 72 + .../qwen3_5_397b_a17b/manifest.json | 32 + .../shortctx_8node_mfu_summary_20260519.json | 155 + .../qwen3_6_35b_a3b/README.md | 23 + ...k_4node_ep8_mbs8_fullrecompute_deepep.yaml | 62 + .../qwen3_6_35b_a3b/manifest.json | 29 + .../qwen36_static_k3_summary_20260519.json | 68 + src/xorl/sim/collect_calibration.py | 438 + .../xorl/sim}/config_fingerprint.py | 20 +- src/xorl/sim/feasibility_evaluator.py | 428 + src/xorl/sim/kernel_variants.py | 108 + src/xorl/sim/memory_ledger.py | 771 ++ .../xorl/sim}/model_metadata.py | 207 +- .../training_sim => src/xorl/sim}/predict.py | 47 +- src/xorl/sim/runtime_config.py | 18 + src/xorl/sim/scenario_planner.py | 9085 +++++++++++++++++ src/xorl/sim/schemas.py | 1912 ++++ .../xorl/sim}/shape_engine.py | 0 src/xorl/sim/simulator_support.py | 211 + src/xorl/sim/timing_ledger.py | 300 + .../xorl/sim}/tradeoff_ranker.py | 6 +- src/xorl/sim/validate.py | 143 + tests/experiments/test_training_sim.py | 175 +- 49 files changed, 21010 insertions(+), 3791 deletions(-) delete mode 100644 experiments/local_benchmark/training_sim/README.md delete mode 100644 experiments/local_benchmark/training_sim/__init__.py delete mode 100644 experiments/local_benchmark/training_sim/benchmark_behavior.py delete mode 100644 experiments/local_benchmark/training_sim/calibration_evaluator.py delete mode 100644 experiments/local_benchmark/training_sim/collect_calibration.py delete mode 100644 experiments/local_benchmark/training_sim/k8s/README.md delete mode 100644 experiments/local_benchmark/training_sim/memory_ledger.py delete mode 100644 experiments/local_benchmark/training_sim/scenario_planner.py delete mode 100644 experiments/local_benchmark/training_sim/schemas.py delete mode 100644 experiments/local_benchmark/training_sim/validate_benchmarks.py create mode 100644 src/xorl/sim/README.md create mode 100644 src/xorl/sim/__init__.py create mode 100644 src/xorl/sim/analytical_ledgers.py create mode 100644 src/xorl/sim/benchmark_behavior.py create mode 100644 src/xorl/sim/calibration_evaluator.py create mode 100644 src/xorl/sim/calibration_packs.py create mode 100644 src/xorl/sim/calibration_packs/qwen3_235b_a22b/README.md create mode 100644 src/xorl/sim/calibration_packs/qwen3_235b_a22b/RESULTS.md create mode 100644 src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_2k_4node_ep8.yaml create mode 100644 src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_muon_8node_ep8_efsdp8_deepep.yaml create mode 100644 src/xorl/sim/calibration_packs/qwen3_235b_a22b/manifest.json create mode 100644 src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/README.md create mode 100644 src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_ep64_reference.yaml create mode 100644 src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r69_mbs4.yaml create mode 100644 src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r73_sync.yaml create mode 100644 src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r75_async.yaml create mode 100644 src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/manifest.json create mode 100644 src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/results/shortctx_8node_mfu_summary_20260519.json create mode 100644 src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/README.md create mode 100644 src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/configs/qwen3_6_35b_a3b_8k_4node_ep8_mbs8_fullrecompute_deepep.yaml create mode 100644 src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/manifest.json create mode 100644 src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/results/qwen36_static_k3_summary_20260519.json create mode 100644 src/xorl/sim/collect_calibration.py rename {experiments/local_benchmark/training_sim => src/xorl/sim}/config_fingerprint.py (93%) create mode 100644 src/xorl/sim/feasibility_evaluator.py create mode 100644 src/xorl/sim/kernel_variants.py create mode 100644 src/xorl/sim/memory_ledger.py rename {experiments/local_benchmark/training_sim => src/xorl/sim}/model_metadata.py (53%) rename {experiments/local_benchmark/training_sim => src/xorl/sim}/predict.py (73%) create mode 100644 src/xorl/sim/runtime_config.py create mode 100644 src/xorl/sim/scenario_planner.py create mode 100644 src/xorl/sim/schemas.py rename {experiments/local_benchmark/training_sim => src/xorl/sim}/shape_engine.py (100%) create mode 100644 src/xorl/sim/simulator_support.py create mode 100644 src/xorl/sim/timing_ledger.py rename {experiments/local_benchmark/training_sim => src/xorl/sim}/tradeoff_ranker.py (96%) create mode 100644 src/xorl/sim/validate.py diff --git a/experiments/local_benchmark/training_sim/README.md b/experiments/local_benchmark/training_sim/README.md deleted file mode 100644 index ca8d3139..00000000 --- a/experiments/local_benchmark/training_sim/README.md +++ /dev/null @@ -1,185 +0,0 @@ -# XoRL Training-Engine Simulator - -This is a CPU-only first slice of the local training-engine simulator. It resolves the launch topology from a -XoRL YAML config, computes deterministic balanced-routing token shapes, estimates a sharded persistent model-state -memory floor, and parses structured trainer logs into calibration summaries. - -It does not yet model activation, attention workspace, MoE kernel workspace, FSDP transient, or allocator slack -memory. Those are left as explicit `unsupported_buckets` in the JSON report until calibrated formulas are added. - -## Predict From A Config - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.predict \ - --config path/to/xorl_config.yaml \ - --world-size 16 \ - --local-world-size 8 \ - --balanced-routing -``` - -## Add Log Calibration - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.predict \ - --config path/to/xorl_config.yaml \ - --world-size 16 \ - --local-world-size 8 \ - --balanced-routing \ - --num-experts 128 \ - --top-k 8 \ - --logs /shared/path/to/trainer-head/logs/run.log \ - --warmup-steps 3 \ - --output experiments/local_benchmark/training_sim/calibration/report.json -``` - -Pass `--benchmark-dir` to include empirical behavior calibrated from a recipe README and result JSON: - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.predict \ - --config path/to/benchmark/configs/xorl_cli.yaml \ - --balanced-routing \ - --benchmark-dir path/to/benchmark -``` - -Empirical matches are config-specific. The simulator checks both topology/workload shape and runtime knobs such as -`deepep_async_combine`, `deepep_num_sms`, `deepep_buffer_size_gb`, `enable_compile`, -`gradient_checkpointing_method`, activation offload, and prefetch count before treating a benchmark row as an exact -calibration point. - -`--benchmark-dir` can also point at a results root containing resolved run directories. Any subdirectory with -`xorl_cli.yaml` is treated as a candidate calibration source. If a matching `node-0.log` is available directly beside -the config or through `startup_metrics.json`'s `startup/master_addr`, the loader parses measured `[STEP ...]` rows -with two warmup steps excluded. OOM logs become calibrated failure boundaries, and runs that report throughput before -crashing are kept as partial-failure calibration points rather than clean promotion candidates. - -## Rank Benchmark Tradeoffs - -Use the tradeoff ranker to compare autotune rows in a benchmark folder. It keeps the fastest raw -candidate separate from the fastest correctness-promotable candidate, so a raw-speed win is not promoted unless it -has a matching `k3_pass` gate. - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.tradeoff_ranker \ - path/to/benchmark -``` - -The report keeps a faster raw candidate separate from a slower promotable candidate when only the latter has a -matching correctness gate. - -## Plan What-If Scenarios - -Use the scenario planner when the question is not just "what already won?" but "what should we try next?". It -mutates a base config across micro-batch and parallelism choices, computes a topology and sharded model-state memory -floor for each candidate, then ranks exact calibrated matches ahead of lower-confidence extrapolations. Extrapolated -candidates are never marked promotable; they need a fresh K3 gate. - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.scenario_planner \ - --config path/to/benchmark/configs/xorl_cli.yaml \ - --benchmark-dir path/to/benchmark \ - --micro-batch-sizes 5 \ - --expert-parallel-sizes 32,64 -``` - -The planner compares concrete parallelism tradeoffs while preserving correctness, runtime-compatibility, and -memory-feasibility caveats. - -For wider topology searches, add `--topology-sweep auto`. Auto mode derives legal candidate values for EP, TP, PP, -Ulysses, and Ring from the resolved world size and model metadata, while explicit comma lists still override any -individual dimension. Exact empirical matches are conservative: an observation only matches TP/PP/Ulysses/Ring values -known from that artifact, and legacy artifacts with missing topology dimensions only exact-match the default value of -1. Non-default TP/PP/CP candidates therefore remain extrapolated unless there is a measured row for that exact -topology. - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.scenario_planner \ - --config path/to/benchmark/configs/xorl_cli.yaml \ - --benchmark-dir path/to/benchmark \ - --gradient-accumulation-steps 1,2,4,8 \ - --topology-sweep auto -``` - -The planner also understands markdown result tables with `tok/s tot`, `tok/step`, and `peak GB` columns, such as the -Qwen3-235B 2k-context sweep. Observed peak memory overrides the analytic floor for feasibility checks, and OOM rows -are kept as calibrated failures when their topology and pack length match a scenario. When two or more -global-batch/GA points are calibrated for the same topology, it fits a simple -`step_time = fixed_overhead + token_slope * tokens` curve for larger GA what-ifs: - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.scenario_planner \ - --config path/to/benchmark/configs/xorl_cli.yaml \ - --benchmark-dir path/to/benchmark \ - --micro-batch-sizes 1 \ - --gradient-accumulation-steps 1,2,4,8 \ - --expert-parallel-sizes 8 -``` - -For that q235 scenario, GA2 is a calibrated measured point and GA4/GA8 are step-time-fit extrapolations. They are -useful next-run candidates, not correctness-promotable results. - -Planner candidates include `calibration_scope` and `risk_flags` fields. `exact_calibrated` means an empirical row -matched the full scenario topology. `inside_measured_envelope` and `outside_measured_envelope` describe whether an -extrapolated candidate stays inside the measured micro-batch/global-batch/parallelism range for that sequence length. -Risk flags call out cases like `requires_remeasurement`, `matched_allocator_pressure_slowdown`, an -`allocator_pressure_boundary:*`, an `observed_oom_boundary:*`, `correctness_runtime_failure_after_steps`, or -`runtime_mismatch:*` when extrapolation had to fall back to a row with different runtime knobs. Treat these flags as -launch-planning constraints: they do not erase the raw score, but they mean the row needs a fresh measurement or debug -pass before it can be used as an optimum. - -Scenario reports keep `best_raw` as the fastest feasible throughput hypothesis, then add `best_risk_adjusted` and -`best_next_measurement`. The risk-adjusted score penalizes extrapolation, memory pressure, missing correctness gates, -allocator-pressure slowdowns, and observed-OOM boundaries. This makes the planner useful as an optimizer loop: launch -the best next measurement when it is a hypothesis, but prefer the risk-adjusted or promotable row when choosing what is -already defensible. - -For exact calibrated rows, an observed peak below the device limit remains feasible even when the configured safety -factor would reserve slightly more than the device capacity. Those rows are marked `feasible_calibrated_peak_high_pressure`. -The safety margin still gates extrapolated peaks and analytic floors. - -## Validate Prediction Fidelity - -Use the calibration evaluator before trusting a scenario sweep as an optimizer. It runs leave-one-out validation over -measured benchmark rows, rebuilds the held-out topology, predicts it from the remaining calibration points, and reports -actual-vs-predicted error. - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.calibration_evaluator \ - --config path/to/benchmark/configs/xorl_cli.yaml \ - --benchmark-dir path/to/benchmark -``` - -Treat a large holdout error as a sign that the relevant lever needs a new calibration point or a more specific -simulator feature before promotion decisions rely on it. - -## Parse Logs Only - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.collect_calibration \ - /shared/path/to/trainer-head/logs/run.log \ - --warmup-steps 3 \ - --world-size 16 -``` - -The log parser recognizes `[STEP ...]`, `[STEP_PHASES ...]`, and `[STEP_MEMORY ...]` lines emitted by -`src/xorl/trainers/trainer.py`. - -## Validate Checked-In Benchmarks - -```bash -cd "$(git rev-parse --show-toplevel)" -python -m experiments.local_benchmark.training_sim.validate_benchmarks \ - --benchmarks-root path/to/benchmarks \ - --model benchmark_name -``` - -The validator checks benchmark YAML, README target metrics, synthetic-routing render scripts, stored throughput -summaries, and static-K3 gate status when those artifacts are present. diff --git a/experiments/local_benchmark/training_sim/__init__.py b/experiments/local_benchmark/training_sim/__init__.py deleted file mode 100644 index fd63eca3..00000000 --- a/experiments/local_benchmark/training_sim/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Local training-engine simulator helpers for benchmark planning.""" diff --git a/experiments/local_benchmark/training_sim/benchmark_behavior.py b/experiments/local_benchmark/training_sim/benchmark_behavior.py deleted file mode 100644 index d9fd6462..00000000 --- a/experiments/local_benchmark/training_sim/benchmark_behavior.py +++ /dev/null @@ -1,1026 +0,0 @@ -"""Empirical benchmark behavior calibration for checked-in benchmark recipes.""" - -from __future__ import annotations - -import argparse -import json -import re -from pathlib import Path -from typing import Any - - -try: - from .collect_calibration import parse_log_path, summarize_observed_run - from .config_fingerprint import load_training_config, resolve_topology - from .schemas import BenchmarkBehaviorPoint, BenchmarkBehaviorPrediction, ShapeLedger, Topology, to_jsonable -except ImportError: # pragma: no cover - exercised by direct script execution - from collect_calibration import parse_log_path, summarize_observed_run - from config_fingerprint import load_training_config, resolve_topology - from schemas import BenchmarkBehaviorPoint, BenchmarkBehaviorPrediction, ShapeLedger, Topology, to_jsonable - - -H100_BF16_PROMISED_TFLOPS_PER_GPU = 989.0 - - -def _gpu_count_from_text(text: str) -> int | None: - match = re.search(r"(?P\d+)\s+nodes?\s+x\s+(?P\d+)\s+H100", text, re.IGNORECASE) - if match: - return int(match.group("nodes")) * int(match.group("gpus")) - match = re.search(r"(?P\d+)\s*[x×]\s*H100", text, re.IGNORECASE) - if match: - return int(match.group("gpus")) - match = re.search(r"(?P\d+)\s+GPUs?", text, re.IGNORECASE) - return int(match.group("gpus")) if match else None - - -def human_number(value: str) -> float: - cleaned = value.strip().replace(",", "").lstrip("~") - multiplier = 1.0 - if cleaned.endswith(("K", "k")): - cleaned = cleaned[:-1] - multiplier = 1_000.0 - elif cleaned.endswith(("M", "m")): - cleaned = cleaned[:-1] - multiplier = 1_000_000.0 - return float(cleaned) * multiplier - - -def _first_non_none(*values: Any) -> Any: - for value in values: - if value is not None: - return value - return None - - -def _readme_point(readme_text: str, *, source: str) -> BenchmarkBehaviorPoint | None: - tps_match = re.search(r"\|\s*tokens/sec\s*\|\s*(?P~?[0-9.]+[KkMm]?)\s*\|", readme_text) - step_match = re.search(r"\|\s*step time\s*\|\s*(?P~?[0-9.]+)s\s*\|", readme_text) - mfu_match = re.search(r"\|\s*MFU\s*\|\s*(?P~?[0-9.]+)%", readme_text) - memory_match = re.search(r"\|\s*allocated memory\s*\|\s*(?P~?[0-9.]+)GB\s*\|", readme_text) - retries_match = re.search(r"\|\s*allocator retries\s*\|\s*(?P\d+)\s*\|", readme_text) - mbs_match = re.search(r"micro_batch_size:\s*(?P\d+)", readme_text) - global_batch_match = re.search(r"global_batch_size:\s*(?P\d+)", readme_text) - if not tps_match: - return None - return BenchmarkBehaviorPoint( - label="readme_reference_mbs8", - source=source, - micro_batch_size=int(mbs_match.group("value")) if mbs_match else None, - global_batch_size=int(global_batch_match.group("value")) if global_batch_match else None, - tokens_per_sec=human_number(tps_match.group("value")), - step_time_sec=float(step_match.group("value").lstrip("~")) if step_match else None, - mfu_percent=float(mfu_match.group("value").lstrip("~")) if mfu_match else None, - tflops_per_gpu=None, - peak_mem_gb=float(memory_match.group("value").lstrip("~")) if memory_match else None, - allocator_retries=int(retries_match.group("value")) if retries_match else None, - gpu_count=_gpu_count_from_text(readme_text), - sample_packing_sequence_len=_seq_len_from_readme(readme_text), - tensor_parallel_size=_readme_parallel_int( - readme_text, - "tensor_parallel_size", - (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), - ), - pipeline_parallel_size=_readme_parallel_int( - readme_text, - "pipeline_parallel_size", - (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), - ), - ulysses_parallel_size=_readme_parallel_int( - readme_text, - "ulysses_parallel_size", - (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), - ), - ringattn_parallel_size=_readme_parallel_int( - readme_text, - "ringattn_parallel_size", - (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), - ), - expert_parallel_size=_readme_parallel_int( - readme_text, - "expert_parallel_size", - (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), - ), - ep_fsdp_size=_readme_parallel_int( - readme_text, - "ep_fsdp_size", - (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), - ), - deepep_async_combine=_readme_bool_from_text(readme_text, "deepep_async_combine"), - deepep_num_sms=_readme_parallel_int( - readme_text, - "deepep_num_sms", - (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), - ), - deepep_buffer_size_gb=_readme_float_from_text(readme_text, "deepep_buffer_size_gb"), - enable_compile=_readme_bool_from_text(readme_text, "enable_compile"), - gradient_checkpointing_method=_checkpointing_method_from_text(readme_text), - status="reference_speed", - correctness_status="raw_speed_not_promoted_without_matching_k3_pass", - notes=["current-main logical FLOPs accounting", "balanced synthetic routing", "deepep_async_combine true"], - ) - - -def _readme_adjacent_mbs10_point( - readme_text: str, *, source: str, seq_len: int | None -) -> BenchmarkBehaviorPoint | None: - match = re.search(r"`mbs=10`[^~]+~(?P[0-9.]+)K tok/s", readme_text) - if not match: - return None - tokens_per_sec = human_number(match.group("value") + "K") - global_batch_size = 320 - step_time_sec = (global_batch_size * seq_len / tokens_per_sec) if seq_len else None - return BenchmarkBehaviorPoint( - label="readme_adjacent_mbs10_allocator_pressure", - source=source, - micro_batch_size=10, - global_batch_size=global_batch_size, - tokens_per_sec=tokens_per_sec, - step_time_sec=step_time_sec, - mfu_percent=None, - tflops_per_gpu=None, - peak_mem_gb=None, - allocator_retries=None, - gpu_count=_gpu_count_from_text(readme_text), - sample_packing_sequence_len=seq_len, - tensor_parallel_size=_readme_parallel_int( - readme_text, - "tensor_parallel_size", - (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), - ), - pipeline_parallel_size=_readme_parallel_int( - readme_text, - "pipeline_parallel_size", - (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), - ), - ulysses_parallel_size=_readme_parallel_int( - readme_text, - "ulysses_parallel_size", - (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), - ), - ringattn_parallel_size=_readme_parallel_int( - readme_text, - "ringattn_parallel_size", - (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), - ), - expert_parallel_size=_readme_parallel_int( - readme_text, - "expert_parallel_size", - (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), - ), - ep_fsdp_size=_readme_parallel_int( - readme_text, - "ep_fsdp_size", - (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), - ), - deepep_async_combine=_readme_bool_from_text(readme_text, "deepep_async_combine"), - deepep_num_sms=_readme_parallel_int( - readme_text, - "deepep_num_sms", - (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), - ), - deepep_buffer_size_gb=_readme_float_from_text(readme_text, "deepep_buffer_size_gb"), - enable_compile=_readme_bool_from_text(readme_text, "enable_compile"), - gradient_checkpointing_method=_checkpointing_method_from_text(readme_text), - status="allocator_pressure_slowdown", - correctness_status="not_promoted", - notes=["fit but slowed with allocator retries"], - ) - - -def _result_throughput_point( - result_path: Path, - result: dict[str, Any], - *, - topology_defaults: dict[str, int | float | bool | str], -) -> BenchmarkBehaviorPoint: - throughput = result["throughput"] - candidate = ( - throughput.get("candidate") - or result.get("candidate") - or (result.get("replay_candidate", {}) if isinstance(result.get("replay_candidate"), dict) else {}).get( - "candidate" - ) - or "throughput" - ) - return BenchmarkBehaviorPoint( - label=f"{result_path.stem}:{candidate}", - source=str(result_path), - micro_batch_size=throughput.get("micro_batch_size"), - global_batch_size=throughput.get("global_batch_size"), - tokens_per_sec=throughput.get("tokens_per_sec"), - step_time_sec=throughput.get("step_time_sec"), - mfu_percent=throughput.get("mfu_percent"), - tflops_per_gpu=throughput.get("mean_tflops_per_gpu"), - peak_mem_gb=throughput.get("gpu_alloc_gb"), - allocator_retries=None, - measured_steps=throughput.get("measured_steps"), - warmup_steps=throughput.get("warmup_steps"), - gpu_count=throughput.get("gpus"), - sample_packing_sequence_len=throughput.get("sample_packing_sequence_len"), - tensor_parallel_size=_first_non_none( - throughput.get("tensor_parallel_size"), topology_defaults.get("tensor_parallel_size") - ), - pipeline_parallel_size=_first_non_none( - throughput.get("pipeline_parallel_size"), topology_defaults.get("pipeline_parallel_size") - ), - ulysses_parallel_size=_first_non_none( - throughput.get("ulysses_parallel_size"), topology_defaults.get("ulysses_parallel_size") - ), - ringattn_parallel_size=_first_non_none( - throughput.get("ringattn_parallel_size"), topology_defaults.get("ringattn_parallel_size") - ), - expert_parallel_size=_first_non_none( - throughput.get("expert_parallel_size"), topology_defaults.get("expert_parallel_size") - ), - ep_fsdp_size=_first_non_none( - throughput.get("ep_fsdp"), throughput.get("ep_fsdp_size"), topology_defaults.get("ep_fsdp_size") - ), - deepep_async_combine=_first_non_none( - throughput.get("deepep_async_combine"), topology_defaults.get("deepep_async_combine") - ), - deepep_num_sms=_first_non_none(throughput.get("deepep_num_sms"), topology_defaults.get("deepep_num_sms")), - deepep_buffer_size_gb=_first_non_none( - throughput.get("deepep_buffer_size_gb"), topology_defaults.get("deepep_buffer_size_gb") - ), - enable_compile=_first_non_none(throughput.get("enable_compile"), topology_defaults.get("enable_compile")), - gradient_checkpointing_method=_first_non_none( - throughput.get("gradient_checkpointing_method"), topology_defaults.get("gradient_checkpointing_method") - ), - enable_activation_offload=_first_non_none( - throughput.get("enable_activation_offload"), topology_defaults.get("enable_activation_offload") - ), - activation_offload_prefetch_count=_first_non_none( - throughput.get("activation_offload_prefetch_count"), - topology_defaults.get("activation_offload_prefetch_count"), - ), - status="historical_throughput_artifact", - correctness_status=None, - notes=[f"commit={throughput.get('commit')}"] if throughput.get("commit") else [], - ) - - -def _with_k3_status(point: BenchmarkBehaviorPoint, result: dict[str, Any]) -> BenchmarkBehaviorPoint: - k3_gate = result.get("k3_gate", {}) - if not k3_gate or k3_gate.get("candidate") not in (None, point.label.split(":", 1)[-1]): - return point - notes = list(point.notes) - if k3_gate.get("primary_failure"): - notes.append(f"k3_primary_failure={k3_gate['primary_failure']}") - return BenchmarkBehaviorPoint( - label=point.label, - source=point.source, - micro_batch_size=point.micro_batch_size, - global_batch_size=point.global_batch_size, - tokens_per_sec=point.tokens_per_sec, - step_time_sec=point.step_time_sec, - mfu_percent=point.mfu_percent, - tflops_per_gpu=point.tflops_per_gpu, - peak_mem_gb=point.peak_mem_gb, - allocator_retries=point.allocator_retries, - measured_steps=point.measured_steps, - warmup_steps=point.warmup_steps, - gpu_count=point.gpu_count, - sample_packing_sequence_len=point.sample_packing_sequence_len, - tensor_parallel_size=point.tensor_parallel_size, - pipeline_parallel_size=point.pipeline_parallel_size, - ulysses_parallel_size=point.ulysses_parallel_size, - ringattn_parallel_size=point.ringattn_parallel_size, - expert_parallel_size=point.expert_parallel_size, - ep_fsdp_size=point.ep_fsdp_size, - deepep_async_combine=point.deepep_async_combine, - deepep_num_sms=point.deepep_num_sms, - deepep_buffer_size_gb=point.deepep_buffer_size_gb, - enable_compile=point.enable_compile, - gradient_checkpointing_method=point.gradient_checkpointing_method, - enable_activation_offload=point.enable_activation_offload, - activation_offload_prefetch_count=point.activation_offload_prefetch_count, - status=point.status, - correctness_status=f"k3_{k3_gate.get('status')}", - notes=notes, - ) - - -def _seq_len_from_readme(readme_text: str) -> int | None: - match = re.search(r"sample_packing_sequence_len:\s*(?P\d+)", readme_text) - return int(match.group("seq")) if match else None - - -def _config_int_from_text(text: str, key: str) -> int | None: - match = re.search(rf"{re.escape(key)}:\s*(?P\d+)", text) - return int(match.group("value")) if match else None - - -def _readme_float_from_text(text: str, key: str) -> float | None: - match = re.search(rf"{re.escape(key)}:\s*(?P\d+(?:\.\d+)?)", text) - return float(match.group("value")) if match else None - - -def _readme_bool_from_text(text: str, key: str) -> bool | None: - match = re.search(rf"{re.escape(key)}:\s*(?Ptrue|false)", text, re.IGNORECASE) - if not match: - return None - return match.group("value").lower() == "true" - - -def _checkpointing_method_from_text(text: str) -> str | None: - lowered = text.lower() - if "recompute_before_dispatch" in lowered or "before_dispatch" in lowered: - return "recompute_before_dispatch" - if "recompute_full_layer" in lowered or "full-layer recompute" in lowered or "fullrecompute" in lowered: - return "recompute_full_layer" - if "no_recompute" in lowered or "no recompute" in lowered: - return "no_recompute" - return None - - -def _trial_checkpointing_method(trial: str) -> str | None: - return _checkpointing_method_from_text(trial) - - -def _trial_activation_offload(trial: str) -> bool | None: - if "noactivationoffload" in trial: - return False - if "activationoffload" in trial: - return True - return None - - -def _trial_prefetch_count(trial: str) -> int | None: - match = re.search(r"prefetch(?P\d+)", trial) - return int(match.group("value")) if match else None - - -def _trial_compile_enabled(trial: str) -> bool | None: - if "nocompile" in trial: - return False - if "compile" in trial: - return True - return None - - -def _trial_deepep_async_combine(trial: str) -> bool | None: - if "noasync" in trial: - return False - if "async" in trial: - return True - return None - - -def _trial_sms_count(trial: str) -> int | None: - match = re.search(r"sms(?P\d+)", trial) - return int(match.group("value")) if match else None - - -def _trial_buffer_size_gb(trial: str) -> float | None: - match = re.search(r"buf(?P\d+)", trial) - if not match: - return None - raw = match.group("value") - if len(raw) == 1: - return float(raw) - return float(f"{raw[:-1]}.{raw[-1]}") - - -def _last_regex_int(line: str, patterns: tuple[str, ...]) -> int | None: - value = None - for pattern in patterns: - for match in re.finditer(pattern, line, re.IGNORECASE): - groupdict = match.groupdict() - for key in ("value", "tp", "pp", "u", "ring"): - if groupdict.get(key) is not None: - value = int(groupdict[key]) - break - return value - - -def _readme_parallel_int(readme_text: str, config_key: str, patterns: tuple[str, ...]) -> int | None: - if value := _config_int_from_text(readme_text, config_key): - return value - for line in readme_text.splitlines(): - if value := _last_regex_int(line, patterns): - return value - return None - - -def _readme_topology_defaults(readme_text: str) -> dict[str, int | float | bool | str]: - defaults: dict[str, int | float | bool | str] = {} - field_patterns = { - "tensor_parallel_size": (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), - "pipeline_parallel_size": (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), - "ulysses_parallel_size": (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), - "ringattn_parallel_size": (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), - "expert_parallel_size": (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), - "ep_fsdp_size": (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), - } - for field, patterns in field_patterns.items(): - if value := _readme_parallel_int(readme_text, field, patterns): - defaults[field] = value - if (value := _readme_bool_from_text(readme_text, "deepep_async_combine")) is not None: - defaults["deepep_async_combine"] = value - if value := _readme_parallel_int( - readme_text, - "deepep_num_sms", - (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), - ): - defaults["deepep_num_sms"] = value - if (value := _readme_float_from_text(readme_text, "deepep_buffer_size_gb")) is not None: - defaults["deepep_buffer_size_gb"] = value - if (value := _readme_bool_from_text(readme_text, "enable_compile")) is not None: - defaults["enable_compile"] = value - if value := _checkpointing_method_from_text(readme_text): - defaults["gradient_checkpointing_method"] = value - return defaults - - -def _first_markdown_number(value: str) -> float | None: - match = re.search(r"~?\s*(?P[0-9][0-9,.]*)(?P[KkMm]?)", value.replace("*", "")) - if not match: - return None - return human_number(match.group("value") + match.group("suffix")) - - -def _markdown_value(values: dict[str, str], key_substring: str) -> str: - for key, value in values.items(): - if key_substring in key: - return value - return "" - - -def _markdown_peak_gb(value: str) -> float | None: - if "oom" in value.lower(): - return None - return _first_markdown_number(value) - - -def _q235_markdown_points(readme_text: str, *, source: str) -> list[BenchmarkBehaviorPoint]: - if "Qwen3-235B" not in readme_text or "tok/s tot" not in readme_text: - return [] - - points: list[BenchmarkBehaviorPoint] = [] - current_header: list[str] | None = None - current_gpu_count: int | None = None - current_ep_size: int | None = None - current_ep_fsdp_size: int | None = None - current_tensor_parallel_size = 1 - current_pipeline_parallel_size = 1 - current_ulysses_parallel_size = 1 - current_ringattn_parallel_size = 1 - for line in readme_text.splitlines(): - if gpu_count := _gpu_count_from_text(line): - current_gpu_count = gpu_count - ep_matches = list(re.finditer(r"\bEP(?P\d+)\b", line)) - if ep_matches: - current_ep_size = int(ep_matches[-1].group("ep")) - efsdp_matches = list(re.finditer(r"(?:ep_fsdp|eFSDP)(?:[= ]|)(?P\d+)", line)) - if efsdp_matches: - current_ep_fsdp_size = int(efsdp_matches[-1].group("efsdp")) - if tp := _last_regex_int( - line, - ( - r"\bTP(?P\d+)\b", - r"\btensor_parallel_size[:= ]+(?P\d+)\b", - r"\btp[=_](?P\d+)\b", - ), - ): - current_tensor_parallel_size = tp - if pp := _last_regex_int( - line, - ( - r"\bPP(?P\d+)\b", - r"\bpipeline_parallel_size[:= ]+(?P\d+)\b", - r"\bpp[=_](?P\d+)\b", - ), - ): - current_pipeline_parallel_size = pp - if ulysses := _last_regex_int( - line, - ( - r"\bU(?P\d+)\b", - r"\bul[y]?sses_parallel_size[:= ]+(?P\d+)\b", - r"\bu[=_]?(?P\d+)\b", - ), - ): - current_ulysses_parallel_size = ulysses - if ringattn := _last_regex_int( - line, - ( - r"\bR(?P\d+)\b", - r"\bringattn_parallel_size[:= ]+(?P\d+)\b", - r"\bring[=_]?(?P\d+)\b", - ), - ): - current_ringattn_parallel_size = ringattn - if not line.startswith("|"): - continue - cells = [cell.strip() for cell in line.strip().strip("|").split("|")] - lowered = [cell.lower() for cell in cells] - if "run" in lowered and "tok/s tot" in lowered: - current_header = lowered - continue - if current_header is None or set(cells) == {"---"} or not cells: - continue - values = dict(zip(current_header, cells, strict=False)) - run = values.get("run", "").replace("*", "").strip("` ") - if not run or run.lower() in {"run", "-----"}: - continue - status_text = _markdown_value(values, "status") - is_failure = "oom" in status_text.lower() or "fail" in status_text.lower() - tokens_per_sec = _first_markdown_number(_markdown_value(values, "tok/s tot")) - tok_step = _first_markdown_number(_markdown_value(values, "tok/step")) - pack = _first_markdown_number(_markdown_value(values, "pack")) - if tok_step is None or pack in (None, 0): - continue - if tokens_per_sec is None and not is_failure: - continue - global_batch_size = int(round(tok_step / pack)) - step_time_sec = _first_markdown_number(_markdown_value(values, "step s")) - mfu_percent = _first_markdown_number(_markdown_value(values, "mfu")) - peak_mem_gb = _markdown_peak_gb(_markdown_value(values, "peak gb")) - points.append( - BenchmarkBehaviorPoint( - label=f"q235_markdown:{run}", - source=source, - micro_batch_size=int(_first_markdown_number(values.get("mbs", "")) or 1), - global_batch_size=global_batch_size, - tokens_per_sec=tokens_per_sec, - step_time_sec=step_time_sec, - mfu_percent=mfu_percent, - peak_mem_gb=peak_mem_gb, - allocator_retries=None, - gpu_count=current_gpu_count, - sample_packing_sequence_len=int(pack), - tensor_parallel_size=current_tensor_parallel_size, - pipeline_parallel_size=current_pipeline_parallel_size, - ulysses_parallel_size=current_ulysses_parallel_size, - ringattn_parallel_size=current_ringattn_parallel_size, - expert_parallel_size=current_ep_size, - ep_fsdp_size=current_ep_fsdp_size, - status="historical_q235_markdown_oom" if is_failure else "historical_q235_markdown", - correctness_status="oom" if is_failure else "not_promoted", - notes=[status_text] if status_text else [], - ) - ) - return points - - -def _best_by_mfu_point( - result_path: Path, - result: dict[str, Any], - row: dict[str, Any], - *, - topology_defaults: dict[str, int | float | bool | str], -) -> BenchmarkBehaviorPoint: - trial = str(row["trial"]) - caveat = row.get("caveat") - k3_gate = row.get("k3_gate") - notes = [] - if caveat: - notes.append(str(caveat)) - if k3_gate: - notes.append(str(k3_gate)) - correctness_status = None - if k3_gate and str(k3_gate).startswith("pass"): - correctness_status = "k3_pass" - elif _first_non_none(row.get("deepep_async_combine"), _trial_deepep_async_combine(trial)): - correctness_status = "raw_speed_not_promoted_without_matching_k3_pass" - return BenchmarkBehaviorPoint( - label=f"best_by_mfu:{trial}", - source=str(result_path), - micro_batch_size=row.get("micro_batch_size"), - global_batch_size=row.get("global_batch_size"), - tokens_per_sec=row.get("tokens_per_sec"), - step_time_sec=row.get("step_time_sec"), - mfu_percent=row.get("mfu_percent"), - tflops_per_gpu=row.get("mean_tflops_per_gpu"), - peak_mem_gb=None, - allocator_retries=None, - measured_steps=row.get("measured_steps"), - warmup_steps=row.get("warmup_steps"), - gpu_count=row.get("gpus") or _gpu_count_from_text(str(result.get("workload", ""))), - sample_packing_sequence_len=row.get("sample_packing_sequence_len"), - tensor_parallel_size=_first_non_none( - row.get("tensor_parallel_size"), topology_defaults.get("tensor_parallel_size") - ), - pipeline_parallel_size=_first_non_none( - row.get("pipeline_parallel_size"), topology_defaults.get("pipeline_parallel_size") - ), - ulysses_parallel_size=_first_non_none( - row.get("ulysses_parallel_size"), topology_defaults.get("ulysses_parallel_size") - ), - ringattn_parallel_size=_first_non_none( - row.get("ringattn_parallel_size"), topology_defaults.get("ringattn_parallel_size") - ), - expert_parallel_size=_first_non_none( - row.get("expert_parallel_size"), topology_defaults.get("expert_parallel_size") - ), - ep_fsdp_size=_first_non_none(row.get("ep_fsdp"), topology_defaults.get("ep_fsdp_size")), - deepep_async_combine=_first_non_none( - row.get("deepep_async_combine"), - _trial_deepep_async_combine(trial), - topology_defaults.get("deepep_async_combine"), - ), - deepep_num_sms=_first_non_none( - row.get("deepep_num_sms"), _trial_sms_count(trial), topology_defaults.get("deepep_num_sms") - ), - deepep_buffer_size_gb=_first_non_none( - row.get("deepep_buffer_size_gb"), - _trial_buffer_size_gb(trial), - topology_defaults.get("deepep_buffer_size_gb"), - ), - enable_compile=_first_non_none( - row.get("enable_compile"), _trial_compile_enabled(trial), topology_defaults.get("enable_compile") - ), - gradient_checkpointing_method=_first_non_none( - row.get("gradient_checkpointing_method"), - _trial_checkpointing_method(trial), - topology_defaults.get("gradient_checkpointing_method"), - ), - enable_activation_offload=_first_non_none( - row.get("enable_activation_offload"), - _trial_activation_offload(trial), - topology_defaults.get("enable_activation_offload"), - ), - activation_offload_prefetch_count=_first_non_none( - row.get("activation_offload_prefetch_count"), - _trial_prefetch_count(trial), - topology_defaults.get("activation_offload_prefetch_count"), - ), - status="autotune_result", - correctness_status=correctness_status, - notes=notes, - ) - - -def _load_startup_metrics(run_dir: Path) -> dict[str, Any]: - startup_path = run_dir / "startup_metrics.json" - if not startup_path.is_file(): - return {} - return json.loads(startup_path.read_text(encoding="utf-8")) - - -def _startup_master_log_path(benchmark_path: Path, startup_metrics: dict[str, Any]) -> Path | None: - metrics = startup_metrics.get("metrics", {}) - master_addr = metrics.get("startup/master_addr") - if not isinstance(master_addr, str) or not master_addr: - return None - run_name = master_addr.removesuffix("-master") - return benchmark_path / run_name / "node-0.log" - - -def _resolved_run_log_path(benchmark_path: Path, run_dir: Path, startup_metrics: dict[str, Any]) -> Path | None: - candidates = [ - run_dir / "node-0.log", - _startup_master_log_path(benchmark_path, startup_metrics), - ] - for candidate in candidates: - if candidate is not None and candidate.is_file(): - return candidate - return None - - -def _log_failure_status(text: str) -> str | None: - lowered = text.lower() - if "outofmemoryerror" in lowered or "cuda out of memory" in lowered: - return "oom" - if "childfailederror" in lowered or "traceback" in lowered: - return "runtime_failure_after_steps" - return None - - -def _oom_peak_mem_gb(text: str) -> float | None: - values = [ - float(match.group("value")) - for match in re.finditer(r"process has (?P\d+(?:\.\d+)?)\s+GiB memory in use", text) - ] - return max(values) if values else None - - -def _round_or_none(value: Any, ndigits: int) -> float | None: - return round(float(value), ndigits) if value is not None else None - - -def _resolved_run_behavior_point(benchmark_path: Path, config_path: Path) -> BenchmarkBehaviorPoint | None: - run_dir = config_path.parent - raw_config = load_training_config(config_path) - try: - topology = resolve_topology(raw_config) - except ValueError: - return None - - startup_metrics = _load_startup_metrics(run_dir) - log_path = _resolved_run_log_path(benchmark_path, run_dir, startup_metrics) - log_text = log_path.read_text(encoding="utf-8", errors="replace") if log_path is not None else "" - failure_status = _log_failure_status(log_text) - observed_summary: dict[str, Any] = {} - if log_path is not None: - observed = parse_log_path(log_path) - warmup_steps = 2 if len(observed.steps) > 2 else 0 - observed_summary = summarize_observed_run(observed, warmup_steps=warmup_steps, world_size=topology.world_size) - - tokens_per_sec = _round_or_none(observed_summary.get("tokens_per_sec_mean"), 3) - peak_mem_gb = _round_or_none(observed_summary.get("peak_mem_gb_max"), 3) - if peak_mem_gb is None and failure_status == "oom": - peak_mem_gb = _round_or_none(_oom_peak_mem_gb(log_text), 3) - if tokens_per_sec is None and failure_status is None: - return None - - if failure_status == "oom" and tokens_per_sec is None: - status = "observed_log_oom" - correctness_status = "oom" - elif failure_status is not None: - status = "observed_log_partial_failure" - correctness_status = failure_status - else: - status = "observed_log_summary" - correctness_status = "not_promoted" - - metrics = startup_metrics.get("metrics", {}) - notes = [ - f"warmup_excluded={observed_summary.get('warmup_excluded', 0)}", - f"parsed_steps={observed_summary.get('parsed_step_count', 0)}", - ] - if startup_metrics.get("repo_commit"): - notes.append(f"commit={startup_metrics['repo_commit']}") - if isinstance(metrics.get("startup/master_addr"), str): - notes.append(f"master_addr={metrics['startup/master_addr']}") - if failure_status is not None: - notes.append(f"log_failure_status={failure_status}") - - return BenchmarkBehaviorPoint( - label=f"resolved_run:{config_path.parent.relative_to(benchmark_path)}", - source=str(log_path or config_path), - micro_batch_size=topology.micro_batch_size, - global_batch_size=topology.global_batch_size, - tokens_per_sec=tokens_per_sec, - step_time_sec=_round_or_none(observed_summary.get("step_time_s_mean"), 6), - mfu_percent=_round_or_none((observed_summary.get("mfu_mean") or 0.0) * 100.0, 3) - if observed_summary.get("mfu_mean") is not None - else None, - tflops_per_gpu=_round_or_none(observed_summary.get("tflops_per_gpu_mean"), 3), - peak_mem_gb=peak_mem_gb, - allocator_retries=None, - measured_steps=observed_summary.get("measured_steps"), - warmup_steps=observed_summary.get("warmup_excluded"), - gpu_count=topology.world_size, - sample_packing_sequence_len=topology.sample_packing_sequence_len, - tensor_parallel_size=topology.tensor_parallel_size, - pipeline_parallel_size=topology.pipeline_parallel_size, - ulysses_parallel_size=topology.ulysses_parallel_size, - ringattn_parallel_size=topology.ringattn_parallel_size, - expert_parallel_size=topology.expert_parallel_size, - ep_fsdp_size=topology.ep_fsdp_size, - deepep_async_combine=_config_bool(raw_config, "model", "deepep_async_combine", False), - deepep_num_sms=_config_int(raw_config, "model", "deepep_num_sms"), - deepep_buffer_size_gb=_config_float(raw_config, "model", "deepep_buffer_size_gb"), - enable_compile=_config_bool(raw_config, "train", "enable_compile", False), - gradient_checkpointing_method=_config_str(raw_config, "train", "gradient_checkpointing_method"), - enable_activation_offload=_config_bool(raw_config, "train", "enable_activation_offload", False), - activation_offload_prefetch_count=_config_int(raw_config, "train", "activation_offload_prefetch_count"), - status=status, - correctness_status=correctness_status, - notes=notes, - ) - - -def _resolved_run_points(benchmark_path: Path) -> list[BenchmarkBehaviorPoint]: - points: list[BenchmarkBehaviorPoint] = [] - for config_path in sorted(benchmark_path.rglob("xorl_cli.yaml")): - if not config_path.is_file(): - continue - point = _resolved_run_behavior_point(benchmark_path, config_path) - if point is not None: - points.append(point) - return points - - -def load_benchmark_behavior_points(benchmark_dir: str | Path) -> list[BenchmarkBehaviorPoint]: - benchmark_path = Path(benchmark_dir) - points: list[BenchmarkBehaviorPoint] = [] - topology_defaults: dict[str, int | float | bool | str] = {} - - for readme_path in (benchmark_path / "README.md", benchmark_path / "RESULTS.md"): - if not readme_path.is_file(): - continue - readme_text = readme_path.read_text(encoding="utf-8") - topology_defaults.update(_readme_topology_defaults(readme_text)) - seq_len = _seq_len_from_readme(readme_text) - readme_reference = _readme_point(readme_text, source=str(readme_path)) - if readme_reference is not None: - points.append(readme_reference) - adjacent_mbs10 = _readme_adjacent_mbs10_point(readme_text, source=str(readme_path), seq_len=seq_len) - if adjacent_mbs10 is not None: - points.append(adjacent_mbs10) - points.extend(_q235_markdown_points(readme_text, source=str(readme_path))) - - for result_path in sorted((benchmark_path / "results").glob("*.json")): - result = json.loads(result_path.read_text(encoding="utf-8")) - for row in result.get("best_by_mfu", []): - if isinstance(row, dict) and row.get("trial"): - points.append(_best_by_mfu_point(result_path, result, row, topology_defaults=topology_defaults)) - throughput = result.get("throughput") - if isinstance(throughput, dict): - points.append( - _with_k3_status( - _result_throughput_point(result_path, result, topology_defaults=topology_defaults), result - ) - ) - - points.extend(_resolved_run_points(benchmark_path)) - return points - - -def behavior_point_matches_topology(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: - if point.micro_batch_size != topology.micro_batch_size or point.global_batch_size != topology.global_batch_size: - return False - if not _point_parallel_size_matches(point.tensor_parallel_size, topology.tensor_parallel_size): - return False - if not _point_parallel_size_matches(point.pipeline_parallel_size, topology.pipeline_parallel_size): - return False - if not _point_parallel_size_matches(point.ulysses_parallel_size, topology.ulysses_parallel_size): - return False - if not _point_parallel_size_matches(point.ringattn_parallel_size, topology.ringattn_parallel_size): - return False - if point.expert_parallel_size is None: - if topology.expert_parallel_size != 1: - return False - elif point.expert_parallel_size != topology.expert_parallel_size: - return False - if point.ep_fsdp_size is not None and point.ep_fsdp_size != topology.ep_fsdp_size: - return False - if ( - point.sample_packing_sequence_len is not None - and topology.sample_packing_sequence_len is not None - and point.sample_packing_sequence_len != topology.sample_packing_sequence_len - ): - return False - return True - - -def _section(raw_config: dict[str, Any], name: str) -> dict[str, Any]: - value = raw_config.get(name, {}) - return value if isinstance(value, dict) else {} - - -def _config_bool(raw_config: dict[str, Any], section_name: str, key: str, default: bool | None = None) -> bool | None: - section = _section(raw_config, section_name) - value = section.get(key, default) - if value is None: - return None - if isinstance(value, str): - return value.strip().lower() in {"1", "true", "yes", "on"} - return bool(value) - - -def _config_int(raw_config: dict[str, Any], section_name: str, key: str) -> int | None: - section = _section(raw_config, section_name) - value = section.get(key) - return int(value) if value is not None else None - - -def _config_float(raw_config: dict[str, Any], section_name: str, key: str) -> float | None: - section = _section(raw_config, section_name) - value = section.get(key) - return float(value) if value is not None else None - - -def _config_str(raw_config: dict[str, Any], section_name: str, key: str) -> str | None: - section = _section(raw_config, section_name) - value = section.get(key) - return str(value) if value is not None else None - - -def behavior_point_workload_mismatches(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> list[str]: - checks: tuple[tuple[str, Any, Any], ...] = ( - ( - "deepep_async_combine", - point.deepep_async_combine, - _config_bool(raw_config, "model", "deepep_async_combine", False), - ), - ("deepep_num_sms", point.deepep_num_sms, _config_int(raw_config, "model", "deepep_num_sms")), - ( - "deepep_buffer_size_gb", - point.deepep_buffer_size_gb, - _config_float(raw_config, "model", "deepep_buffer_size_gb"), - ), - ("enable_compile", point.enable_compile, _config_bool(raw_config, "train", "enable_compile", False)), - ( - "gradient_checkpointing_method", - point.gradient_checkpointing_method, - _config_str(raw_config, "train", "gradient_checkpointing_method"), - ), - ( - "enable_activation_offload", - point.enable_activation_offload, - _config_bool(raw_config, "train", "enable_activation_offload", False), - ), - ( - "activation_offload_prefetch_count", - point.activation_offload_prefetch_count, - _config_int(raw_config, "train", "activation_offload_prefetch_count"), - ), - ) - mismatches: list[str] = [] - for field_name, point_value, config_value in checks: - if point_value is None: - continue - if isinstance(point_value, float): - if config_value is None or abs(float(point_value) - float(config_value)) > 1e-9: - mismatches.append(field_name) - elif point_value != config_value: - mismatches.append(field_name) - return mismatches - - -def behavior_point_matches_workload(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> bool: - return not behavior_point_workload_mismatches(point, raw_config) - - -def _point_parallel_size_matches(point_value: int | None, topology_value: int) -> bool: - if point_value is None: - return topology_value == 1 - return point_value == topology_value - - -def predict_benchmark_behavior( - points: list[BenchmarkBehaviorPoint], - topology: Topology, - shape: ShapeLedger, - raw_config: dict[str, Any] | None = None, -) -> BenchmarkBehaviorPrediction: - matches = [ - point - for point in points - if behavior_point_matches_topology(point, topology) - and (raw_config is None or behavior_point_matches_workload(point, raw_config)) - ] - warnings: list[str] = [] - if not matches: - known = ", ".join( - f"{point.label}(mbs={point.micro_batch_size},gb={point.global_batch_size})" for point in points - ) - return BenchmarkBehaviorPrediction( - status="no_calibrated_match", - matched_label=None, - source=None, - tokens_per_sec=None, - tokens_per_sec_per_gpu=None, - step_time_sec=None, - mfu_percent=None, - tflops_per_gpu=None, - promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, - peak_mem_gb=None, - allocator_retries=None, - derived_global_tokens_per_step=shape.global_tokens_per_train_step, - warnings=[ - f"no empirical behavior point for mbs={topology.micro_batch_size}, gb={topology.global_batch_size}; known: {known}" - ], - ) - - point = matches[0] - tokens_per_sec_per_gpu = None - if point.tokens_per_sec is not None and topology.world_size: - tokens_per_sec_per_gpu = point.tokens_per_sec / topology.world_size - step_time_sec = point.step_time_sec - if step_time_sec is None and shape.global_tokens_per_train_step and point.tokens_per_sec: - step_time_sec = shape.global_tokens_per_train_step / point.tokens_per_sec - tflops_per_gpu = None - if point.tflops_per_gpu is not None: - tflops_per_gpu = point.tflops_per_gpu - elif point.mfu_percent is not None: - tflops_per_gpu = H100_BF16_PROMISED_TFLOPS_PER_GPU * point.mfu_percent / 100.0 - - if point.status == "allocator_pressure_slowdown": - warnings.append("matched behavior point is an allocator-pressure slowdown, not a promotable speed target") - if point.correctness_status and point.correctness_status != "k3_pass": - warnings.append(f"correctness status is {point.correctness_status}") - - prediction_status = "calibrated_failure" if point.correctness_status == "oom" else "calibrated" - - return BenchmarkBehaviorPrediction( - status=prediction_status, - matched_label=point.label, - source=point.source, - tokens_per_sec=point.tokens_per_sec, - tokens_per_sec_per_gpu=tokens_per_sec_per_gpu, - step_time_sec=step_time_sec, - mfu_percent=point.mfu_percent, - tflops_per_gpu=tflops_per_gpu, - promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, - peak_mem_gb=point.peak_mem_gb, - allocator_retries=point.allocator_retries, - derived_global_tokens_per_step=shape.global_tokens_per_train_step, - correctness_status=point.correctness_status, - warnings=warnings, - ) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("benchmark_dir", type=Path) - args = parser.parse_args() - points = load_benchmark_behavior_points(args.benchmark_dir) - print(json.dumps(to_jsonable({"points": points}), indent=2, sort_keys=True)) - - -if __name__ == "__main__": - main() diff --git a/experiments/local_benchmark/training_sim/calibration_evaluator.py b/experiments/local_benchmark/training_sim/calibration_evaluator.py deleted file mode 100644 index cd82fd24..00000000 --- a/experiments/local_benchmark/training_sim/calibration_evaluator.py +++ /dev/null @@ -1,256 +0,0 @@ -"""Evaluate scenario-prediction fidelity with leave-one-out benchmark holdouts.""" - -from __future__ import annotations - -import argparse -import json -import statistics -from pathlib import Path -from typing import Any - - -try: - from .benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior - from .config_fingerprint import load_training_config, resolve_topology - from .memory_ledger import build_memory_ledger - from .model_metadata import resolve_model_metadata - from .scenario_planner import _extrapolate_behavior, _mutated_config, _topology_label - from .schemas import ( - BenchmarkBehaviorPoint, - CalibrationHoldout, - CalibrationReport, - Topology, - to_jsonable, - ) - from .shape_engine import build_shape_ledger -except ImportError: # pragma: no cover - exercised by direct script execution - from benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior - from config_fingerprint import load_training_config, resolve_topology - from memory_ledger import build_memory_ledger - from model_metadata import resolve_model_metadata - from scenario_planner import _extrapolate_behavior, _mutated_config, _topology_label - from schemas import BenchmarkBehaviorPoint, CalibrationHoldout, CalibrationReport, Topology, to_jsonable - from shape_engine import build_shape_ledger - - -def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: - value = raw.get(name, {}) - if isinstance(value, dict): - return value - raw[name] = {} - return raw[name] - - -def _point_parallel_size(value: int | None, fallback: int) -> int: - return value if value is not None else fallback - - -def _set_if_known(section: dict[str, Any], key: str, value: Any) -> None: - if value is not None: - section[key] = value - - -def _apply_point_runtime_signature(raw_config: dict[str, Any], point: BenchmarkBehaviorPoint) -> None: - model = _section(raw_config, "model") - train = _section(raw_config, "train") - _set_if_known(model, "deepep_async_combine", point.deepep_async_combine) - _set_if_known(model, "deepep_num_sms", point.deepep_num_sms) - _set_if_known(model, "deepep_buffer_size_gb", point.deepep_buffer_size_gb) - _set_if_known(train, "enable_compile", point.enable_compile) - _set_if_known(train, "gradient_checkpointing_method", point.gradient_checkpointing_method) - _set_if_known(train, "enable_activation_offload", point.enable_activation_offload) - _set_if_known(train, "activation_offload_prefetch_count", point.activation_offload_prefetch_count) - - -def _topology_for_point( - base_config: dict[str, Any], - base_topology: Topology, - point: BenchmarkBehaviorPoint, - *, - world_size: int | None, - local_world_size: int | None, -) -> tuple[dict[str, Any] | None, Topology | None, str | None]: - if point.micro_batch_size is None or point.global_batch_size is None: - return None, None, "missing micro_batch_size/global_batch_size" - if point.tokens_per_sec is None: - return None, None, "missing tokens_per_sec" - - resolved_world_size = point.gpu_count or world_size or base_topology.world_size - resolved_local_world_size = local_world_size or base_topology.local_world_size - tensor_parallel = _point_parallel_size(point.tensor_parallel_size, base_topology.tensor_parallel_size) - pipeline_parallel = _point_parallel_size(point.pipeline_parallel_size, base_topology.pipeline_parallel_size) - ulysses_parallel = _point_parallel_size(point.ulysses_parallel_size, base_topology.ulysses_parallel_size) - ringattn_parallel = _point_parallel_size(point.ringattn_parallel_size, base_topology.ringattn_parallel_size) - expert_parallel = _point_parallel_size(point.expert_parallel_size, base_topology.expert_parallel_size) - non_dp = tensor_parallel * pipeline_parallel * ulysses_parallel * ringattn_parallel - if non_dp <= 0 or resolved_world_size % non_dp: - return None, None, "world_size is not divisible by heldout non-DP topology" - data_parallel_size = resolved_world_size // non_dp - denominator = point.micro_batch_size * data_parallel_size - if denominator <= 0 or point.global_batch_size % denominator: - return None, None, "global_batch_size is not divisible by micro_batch_size * data_parallel_size" - gradient_accumulation_steps = point.global_batch_size // denominator - - raw_config = _mutated_config( - base_config, - world_size=resolved_world_size, - micro_batch_size=point.micro_batch_size, - gradient_accumulation_steps=gradient_accumulation_steps, - expert_parallel_size=expert_parallel, - tensor_parallel_size=tensor_parallel, - pipeline_parallel_size=pipeline_parallel, - ulysses_parallel_size=ulysses_parallel, - ringattn_parallel_size=ringattn_parallel, - ) - if point.sample_packing_sequence_len is not None: - _section(raw_config, "data")["sample_packing_sequence_len"] = point.sample_packing_sequence_len - _apply_point_runtime_signature(raw_config, point) - try: - topology = resolve_topology( - raw_config, - world_size=resolved_world_size, - local_world_size=resolved_local_world_size, - ) - except ValueError as exc: - return None, None, str(exc) - if point.ep_fsdp_size is not None and topology.ep_fsdp_size != point.ep_fsdp_size: - return None, None, "heldout ep_fsdp_size does not match resolved topology" - return raw_config, topology, None - - -def _without_point( - behavior_points: list[BenchmarkBehaviorPoint], - heldout: BenchmarkBehaviorPoint, -) -> list[BenchmarkBehaviorPoint]: - return [point for point in behavior_points if not (point.label == heldout.label and point.source == heldout.source)] - - -def evaluate_calibration( - base_config_path: str | Path, - *, - benchmark_dir: str | Path, - world_size: int | None = None, - local_world_size: int | None = None, - device_memory_limit_gb: float = 80.0, - memory_safety_factor: float = 1.15, -) -> CalibrationReport: - base_path = Path(base_config_path) - benchmark_path = Path(benchmark_dir) - base_config = load_training_config(base_path) - base_topology = resolve_topology(base_config, world_size=world_size, local_world_size=local_world_size) - metadata = resolve_model_metadata(base_config) - behavior_points = load_benchmark_behavior_points(benchmark_path) - measured_points = [point for point in behavior_points if point.tokens_per_sec is not None] - - holdouts: list[CalibrationHoldout] = [] - warnings: list[str] = [] - skipped_count = 0 - for heldout in measured_points: - raw_config, topology, skip_reason = _topology_for_point( - base_config, - base_topology, - heldout, - world_size=world_size, - local_world_size=local_world_size, - ) - if raw_config is None or topology is None: - skipped_count += 1 - warnings.append(f"skipped {heldout.label}: {skip_reason}") - continue - - training_points = _without_point(behavior_points, heldout) - shape = build_shape_ledger(topology, balanced_routing=True) - exact_prediction = predict_benchmark_behavior(training_points, topology, shape, raw_config) - if exact_prediction.status == "calibrated": - prediction = exact_prediction - else: - memory = build_memory_ledger(raw_config, topology=topology, model_metadata=metadata) - prediction, _ = _extrapolate_behavior( - training_points, - topology, - shape, - raw_config=raw_config, - device_memory_limit_gb=device_memory_limit_gb, - memory_safety_factor=memory_safety_factor, - analytic_peak_floor_gb=memory.analytic_peak_floor_gb, - ) - - predicted = prediction.tokens_per_sec - absolute_error = None - absolute_percentage_error = None - if predicted is not None: - absolute_error = abs(predicted - heldout.tokens_per_sec) - absolute_percentage_error = 100.0 * absolute_error / heldout.tokens_per_sec - holdouts.append( - CalibrationHoldout( - label=heldout.label, - source=heldout.source, - topology_label=_topology_label(topology), - actual_tokens_per_sec=heldout.tokens_per_sec, - predicted_tokens_per_sec=predicted, - prediction_status=prediction.status, - matched_label=prediction.matched_label, - absolute_error_tokens_per_sec=round(absolute_error, 3) if absolute_error is not None else None, - absolute_percentage_error=round(absolute_percentage_error, 3) - if absolute_percentage_error is not None - else None, - calibrated_from_count=len(training_points), - warnings=prediction.warnings, - ) - ) - - errors = [ - holdout.absolute_percentage_error for holdout in holdouts if holdout.absolute_percentage_error is not None - ] - status_counts: dict[str, int] = {} - for holdout in holdouts: - status_counts[holdout.prediction_status] = status_counts.get(holdout.prediction_status, 0) + 1 - status = "ok" if errors else "insufficient_data" - if holdouts and not errors: - warnings.append("all holdouts were unscored") - - return CalibrationReport( - base_config_path=str(base_path), - benchmark_dir=str(benchmark_path), - status=status, - measured_point_count=len(measured_points), - evaluated_count=len(holdouts), - skipped_count=skipped_count, - mean_absolute_percentage_error=round(statistics.fmean(errors), 3) if errors else None, - median_absolute_percentage_error=round(statistics.median(errors), 3) if errors else None, - max_absolute_percentage_error=round(max(errors), 3) if errors else None, - prediction_status_counts=dict(sorted(status_counts.items())), - holdouts=holdouts, - warnings=warnings, - ) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config", type=Path, required=True) - parser.add_argument("--benchmark-dir", type=Path, required=True) - parser.add_argument("--world-size", type=int, default=None) - parser.add_argument("--local-world-size", type=int, default=None) - parser.add_argument("--device-memory-limit-gb", type=float, default=80.0) - parser.add_argument("--memory-safety-factor", type=float, default=1.15) - parser.add_argument("--output", type=Path, default=None) - args = parser.parse_args() - - report = evaluate_calibration( - args.config, - benchmark_dir=args.benchmark_dir, - world_size=args.world_size, - local_world_size=args.local_world_size, - device_memory_limit_gb=args.device_memory_limit_gb, - memory_safety_factor=args.memory_safety_factor, - ) - rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" - if args.output: - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(rendered, encoding="utf-8") - else: - print(rendered, end="") - - -if __name__ == "__main__": - main() diff --git a/experiments/local_benchmark/training_sim/collect_calibration.py b/experiments/local_benchmark/training_sim/collect_calibration.py deleted file mode 100644 index d203be37..00000000 --- a/experiments/local_benchmark/training_sim/collect_calibration.py +++ /dev/null @@ -1,224 +0,0 @@ -"""Parse XoRL trainer structured logs into calibration observations.""" - -from __future__ import annotations - -import argparse -import json -import re -import statistics -from pathlib import Path -from typing import Any, Iterable - - -try: - from .schemas import MemoryPhaseObservation, ObservedRun, PhaseObservation, StepObservation, to_jsonable -except ImportError: # pragma: no cover - exercised by direct script execution - from schemas import MemoryPhaseObservation, ObservedRun, PhaseObservation, StepObservation, to_jsonable - - -STEP_RE = re.compile(r"\[STEP\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") -PHASE_RE = re.compile(r"\[(?PSTEP_PHASES(?:_PARTIAL)?)\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") -MEMORY_RE = re.compile(r"\[(?PSTEP_MEMORY(?:_PARTIAL)?)\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") -KV_RE = re.compile(r"(?P[A-Za-z0-9_+./-]+)=(?P\S+)") - - -def _float_or_none(value: str | None) -> float | None: - if value is None: - return None - cleaned = value.strip().rstrip(",") - for suffix in ("GB", "gb", "s"): - if cleaned.endswith(suffix): - cleaned = cleaned[: -len(suffix)] - break - try: - return float(cleaned) - except ValueError: - return None - - -def _parse_metric_body(body: str) -> dict[str, float]: - metrics: dict[str, float] = {} - for match in KV_RE.finditer(body): - numeric = _float_or_none(match.group("value")) - if numeric is not None: - metrics[match.group("key")] = numeric - return metrics - - -def _step_from_match(match: re.Match[str], source: str) -> StepObservation: - metrics = _parse_metric_body(match.group("body")) - phase_memory: dict[str, float] = {} - for key in ("fwd", "bwd", "optim", "fwd+bwd", "offload"): - if key in metrics: - phase_memory[key] = metrics[key] - - known_keys = { - "loss", - "grad_norm", - "lr", - "tflops", - "mfu", - "tokens_per_sec", - "time", - "peak_mem", - "fwd", - "bwd", - "optim", - "fwd+bwd", - "offload", - } - extra = {key: value for key, value in metrics.items() if key not in known_keys} - return StepObservation( - source=source, - step=int(match.group("step")), - max_steps=match.group("max"), - loss=metrics.get("loss"), - grad_norm=metrics.get("grad_norm"), - lr=metrics.get("lr"), - tflops_per_gpu=metrics.get("tflops"), - mfu=metrics.get("mfu"), - tokens_per_sec=metrics.get("tokens_per_sec"), - step_time_s=metrics.get("time"), - peak_mem_gb=metrics.get("peak_mem"), - phase_memory_gb=phase_memory, - extra=extra, - ) - - -def parse_log_text(text: str, *, source: str = "") -> ObservedRun: - steps: list[StepObservation] = [] - phases: list[PhaseObservation] = [] - memory_phases: list[MemoryPhaseObservation] = [] - - for line in text.splitlines(): - if phase_match := PHASE_RE.search(line): - phases.append( - PhaseObservation( - source=source, - prefix=phase_match.group("prefix"), - step=int(phase_match.group("step")), - max_steps=phase_match.group("max"), - metrics=_parse_metric_body(phase_match.group("body")), - ) - ) - continue - - if memory_match := MEMORY_RE.search(line): - memory_phases.append( - MemoryPhaseObservation( - source=source, - prefix=memory_match.group("prefix"), - step=int(memory_match.group("step")), - max_steps=memory_match.group("max"), - metrics=_parse_metric_body(memory_match.group("body")), - ) - ) - continue - - if step_match := STEP_RE.search(line): - steps.append(_step_from_match(step_match, source)) - - return ObservedRun(sources=[source], steps=steps, phases=phases, memory_phases=memory_phases) - - -def parse_log_path(path: str | Path) -> ObservedRun: - log_path = Path(path) - text = log_path.read_text(encoding="utf-8", errors="replace") - return parse_log_text(text, source=str(log_path)) - - -def merge_observed_runs(runs: Iterable[ObservedRun]) -> ObservedRun: - sources: list[str] = [] - steps: list[StepObservation] = [] - phases: list[PhaseObservation] = [] - memory_phases: list[MemoryPhaseObservation] = [] - for run in runs: - sources.extend(run.sources) - steps.extend(run.steps) - phases.extend(run.phases) - memory_phases.extend(run.memory_phases) - return ObservedRun(sources=sources, steps=steps, phases=phases, memory_phases=memory_phases) - - -def _mean(values: list[float]) -> float | None: - return statistics.fmean(values) if values else None - - -def _median(values: list[float]) -> float | None: - return statistics.median(values) if values else None - - -def summarize_observed_run( - run: ObservedRun, - *, - warmup_steps: int = 0, - world_size: int | None = None, -) -> dict[str, Any]: - ordered_steps = sorted(run.steps, key=lambda row: (row.source, row.step)) - measured = ordered_steps[warmup_steps:] - tps = [row.tokens_per_sec for row in measured if row.tokens_per_sec is not None] - tflops = [row.tflops_per_gpu for row in measured if row.tflops_per_gpu is not None] - mfu = [row.mfu for row in measured if row.mfu is not None] - step_time = [row.step_time_s for row in measured if row.step_time_s is not None] - peaks = [row.peak_mem_gb for row in measured if row.peak_mem_gb is not None] - - summary: dict[str, Any] = { - "sources": run.sources, - "parsed_step_count": len(run.steps), - "parsed_phase_count": len(run.phases), - "parsed_memory_phase_count": len(run.memory_phases), - "warmup_excluded": warmup_steps, - "measured_steps": len(measured), - "tokens_per_sec_mean": _mean(tps), - "tokens_per_sec_median": _median(tps), - "tflops_per_gpu_mean": _mean(tflops), - "mfu_mean": _mean(mfu), - "step_time_s_mean": _mean(step_time), - "peak_mem_gb_max": max(peaks) if peaks else None, - } - if world_size and summary["tokens_per_sec_mean"] is not None: - summary["tokens_per_sec_per_gpu_mean"] = summary["tokens_per_sec_mean"] / world_size - if measured: - summary["first_measured_step"] = measured[0].step - summary["last_measured_step"] = measured[-1].step - summary["loss_last"] = measured[-1].loss - return summary - - -def _expand_paths(paths: list[Path]) -> list[Path]: - expanded: list[Path] = [] - for path in paths: - if path.is_dir(): - expanded.extend(sorted(child for child in path.rglob("*") if child.is_file())) - else: - expanded.append(path) - return expanded - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("paths", nargs="+", type=Path, help="Log files or directories to parse") - parser.add_argument("--warmup-steps", type=int, default=0, help="Drop this many parsed [STEP] rows from summary") - parser.add_argument("--world-size", type=int, default=None, help="Optional GPU count for per-GPU throughput") - parser.add_argument("--output", type=Path, default=None, help="Write JSON output to this path") - parser.add_argument("--include-rows", action="store_true", help="Include parsed row details, not just the summary") - args = parser.parse_args() - - runs = [parse_log_path(path) for path in _expand_paths(args.paths)] - observed = merge_observed_runs(runs) - payload: dict[str, Any] = { - "summary": summarize_observed_run(observed, warmup_steps=args.warmup_steps, world_size=args.world_size) - } - if args.include_rows: - payload["observed"] = to_jsonable(observed) - - rendered = json.dumps(to_jsonable(payload), indent=2, sort_keys=True) + "\n" - if args.output: - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(rendered, encoding="utf-8") - else: - print(rendered, end="") - - -if __name__ == "__main__": - main() diff --git a/experiments/local_benchmark/training_sim/k8s/README.md b/experiments/local_benchmark/training_sim/k8s/README.md deleted file mode 100644 index c53e3687..00000000 --- a/experiments/local_benchmark/training_sim/k8s/README.md +++ /dev/null @@ -1,13 +0,0 @@ -# Simulator Calibration K8s Notes - -Calibration jobs should be normal XoRL training benchmark jobs with these trainer flags enabled: - -```yaml -train: - enable_step_phase_timing: true - enable_per_component_timing: true - enable_step_memory_profiling: true -``` - -Any pod requesting GPUs on the research-common-h100 cluster must set `team: turbo` on the pod template labels. -Keep the run short, preserve the trainer-head log, and feed that log to `collect_calibration.py`. diff --git a/experiments/local_benchmark/training_sim/memory_ledger.py b/experiments/local_benchmark/training_sim/memory_ledger.py deleted file mode 100644 index e24b6cd4..00000000 --- a/experiments/local_benchmark/training_sim/memory_ledger.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Initial memory ledger built from config constants and observed structured logs.""" - -from __future__ import annotations - -from typing import Any - - -try: - from .schemas import MemoryBucket, MemoryLedger, ModelMetadata, ObservedRun, Topology -except ImportError: # pragma: no cover - exercised by direct script execution - from schemas import MemoryBucket, MemoryLedger, ModelMetadata, ObservedRun, Topology - - -BYTES_PER_GIB = 1024**3 - - -def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: - value = raw.get(name, {}) - return value if isinstance(value, dict) else {} - - -def _float_field(section: dict[str, Any], key: str) -> float | None: - value = section.get(key) - if value is None: - return None - return float(value) - - -def _dtype_bytes(dtype: Any, *, default: int) -> int: - if dtype is None: - return default - normalized = str(dtype).lower() - if normalized in {"bf16", "bfloat16", "fp16", "float16", "half"}: - return 2 - if normalized in {"fp32", "float32", "float"}: - return 4 - if normalized in {"fp8", "float8", "e4m3", "e5m2"}: - return 1 - return default - - -def _gb(byte_count: float) -> float: - return byte_count / BYTES_PER_GIB - - -def _round_gb(value: float | None) -> float | None: - return round(value, 3) if value is not None else None - - -def _estimate_param_counts(metadata: ModelMetadata) -> tuple[float, float, float] | None: - hidden = metadata.hidden_size - layers = metadata.num_hidden_layers - vocab = metadata.vocab_size - if hidden is None or layers is None or vocab is None: - return None - - head_dim = metadata.head_dim - if head_dim is None and metadata.num_attention_heads: - head_dim = hidden // metadata.num_attention_heads - attention_heads = metadata.num_attention_heads or 1 - key_value_heads = metadata.num_key_value_heads or attention_heads - if head_dim is None: - return None - - q_proj = hidden * attention_heads * head_dim - k_proj = hidden * key_value_heads * head_dim - v_proj = hidden * key_value_heads * head_dim - o_proj = attention_heads * head_dim * hidden - attention_params = layers * (q_proj + k_proj + v_proj + o_proj) - - dense_mlp_params = 0 - has_routed_experts = metadata.num_experts is not None and metadata.moe_intermediate_size is not None - if metadata.intermediate_size is not None and not has_routed_experts: - dense_mlp_params = layers * 3 * hidden * metadata.intermediate_size - - shared_expert_params = 0 - if metadata.shared_expert_intermediate_size is not None: - shared_expert_params = layers * 3 * hidden * metadata.shared_expert_intermediate_size - - expert_params = 0 - if has_routed_experts and metadata.num_experts is not None and metadata.moe_intermediate_size is not None: - expert_params = layers * metadata.num_experts * 3 * hidden * metadata.moe_intermediate_size - - embedding_params = vocab * hidden - lm_head_params = 0 if metadata.tie_word_embeddings else vocab * hidden - norm_params = (2 * layers + 1) * hidden - non_expert_params = attention_params + dense_mlp_params + shared_expert_params + embedding_params + lm_head_params - non_expert_params += norm_params - return float(non_expert_params + expert_params), float(non_expert_params), float(expert_params) - - -def _model_state_buckets( - raw_config: dict[str, Any], - topology: Topology | None, - metadata: ModelMetadata | None, - deepep_buffer_size_gb: float | None, -) -> tuple[float | None, float | None, float | None, float | None, float | None, list[MemoryBucket], list[str]]: - if topology is None or metadata is None: - return None, None, None, None, None, [], ["parameter_and_optimizer_bytes"] - - counts = _estimate_param_counts(metadata) - if counts is None: - return None, None, None, None, None, [], ["parameter_and_optimizer_bytes"] - - total_params, non_expert_params, expert_params = counts - train = _section(raw_config, "train") - expert_shard_size = topology.expert_parallel_size * (topology.ep_fsdp_size or 1) - local_non_expert_params = non_expert_params / max(topology.data_parallel_shard_size, 1) - local_expert_params = expert_params / max(expert_shard_size, 1) - local_params = local_non_expert_params + local_expert_params - - weight_bytes = _dtype_bytes(train.get("param_dtype"), default=2 if train.get("enable_mixed_precision") else 4) - optimizer = str(train.get("optimizer", "")).lower() - optimizer_dtype_bytes = _dtype_bytes(train.get("optimizer_dtype"), default=4) - gradient_dtype = train.get("gradient_dtype") or train.get("fsdp_reduce_dtype") or train.get("optimizer_dtype") - gradient_bytes = _dtype_bytes(gradient_dtype, default=4) - - sharded_param_gb = _gb(local_params * weight_bytes) - master_param_gb = 0.0 - if optimizer == "adamw": - master_param_gb = _gb(local_params * optimizer_dtype_bytes) - persistent_model_state_gb = sharded_param_gb + master_param_gb - - gradient_state_gb = _gb(local_params * gradient_bytes) - optimizer_state_gb = 0.0 - if optimizer == "adamw": - optimizer_state_gb = _gb(local_params * 2 * optimizer_dtype_bytes) - elif optimizer == "muon" and float(train.get("muon_momentum", 0.0) or 0.0) > 0: - optimizer_state_gb = _gb(local_params * optimizer_dtype_bytes) - - buckets = [ - MemoryBucket( - name="sharded_trainable_params", - gb=_round_gb(sharded_param_gb) or 0.0, - source="analytic_model_metadata", - notes=[ - f"weight_bytes={weight_bytes}", - f"local_non_expert_params={local_non_expert_params:.0f}", - f"local_expert_params={local_expert_params:.0f}", - ], - ), - MemoryBucket( - name="gradient_storage", - gb=_round_gb(gradient_state_gb) or 0.0, - source="analytic_dtype_policy", - notes=[f"gradient_bytes={gradient_bytes}"], - ), - ] - if master_param_gb: - buckets.append( - MemoryBucket( - name="optimizer_master_params", - gb=_round_gb(master_param_gb) or 0.0, - source="analytic_optimizer_policy", - notes=[f"optimizer={optimizer}", f"optimizer_dtype_bytes={optimizer_dtype_bytes}"], - ) - ) - if optimizer_state_gb: - buckets.append( - MemoryBucket( - name=f"{optimizer}_optimizer_state", - gb=_round_gb(optimizer_state_gb) or 0.0, - source="analytic_optimizer_policy", - notes=[f"optimizer_dtype_bytes={optimizer_dtype_bytes}"], - ) - ) - if deepep_buffer_size_gb: - buckets.append( - MemoryBucket( - name="deepep_static_buffer", - gb=deepep_buffer_size_gb, - source="config", - ) - ) - - unsupported = [ - "activation_recompute_schedule", - "attention_workspace", - "moe_kernel_workspace", - "fsdp_unshard_and_reduce_scatter_transients", - "allocator_reserved_slack", - ] - return ( - total_params / 1_000_000_000, - local_params / 1_000_000_000, - persistent_model_state_gb, - gradient_state_gb, - optimizer_state_gb, - sorted(buckets, key=lambda bucket: bucket.gb, reverse=True), - unsupported, - ) - - -def build_memory_ledger( - raw_config: dict[str, Any], - observed: ObservedRun | None = None, - *, - topology: Topology | None = None, - model_metadata: ModelMetadata | None = None, -) -> MemoryLedger: - model = _section(raw_config, "model") - train = _section(raw_config, "train") - observed_peak = None - observed_phase_peak: dict[str, float] = {} - - if observed is not None: - peaks = [row.peak_mem_gb for row in observed.steps if row.peak_mem_gb is not None] - observed_peak = max(peaks) if peaks else None - for row in observed.steps: - for phase, value in row.phase_memory_gb.items(): - observed_phase_peak[phase] = max(value, observed_phase_peak.get(phase, value)) - for memory_row in observed.memory_phases: - for key, value in memory_row.metrics.items(): - observed_phase_peak[key] = max(value, observed_phase_peak.get(key, value)) - - deepep_buffer_size_gb = _float_field(model, "deepep_buffer_size_gb") - if deepep_buffer_size_gb is None: - deepep_buffer_size_gb = _float_field(train, "deepep_buffer_size_gb") - - ( - estimated_total_params_b, - estimated_local_params_b, - persistent_model_state_gb, - gradient_state_gb, - optimizer_state_gb, - top_memory_buckets, - unsupported_buckets, - ) = _model_state_buckets(raw_config, topology, model_metadata, deepep_buffer_size_gb) - analytic_peak_floor_gb = None - if persistent_model_state_gb is not None and gradient_state_gb is not None and optimizer_state_gb is not None: - analytic_peak_floor_gb = persistent_model_state_gb + gradient_state_gb + optimizer_state_gb - analytic_peak_floor_gb += deepep_buffer_size_gb or 0.0 - - return MemoryLedger( - deepep_buffer_size_gb=deepep_buffer_size_gb, - observed_peak_mem_gb_max=observed_peak, - observed_phase_peak_gb=observed_phase_peak, - estimated_total_params_b=_round_gb(estimated_total_params_b), - estimated_local_params_b=_round_gb(estimated_local_params_b), - persistent_model_state_gb=_round_gb(persistent_model_state_gb), - gradient_state_gb=_round_gb(gradient_state_gb), - optimizer_state_gb=_round_gb(optimizer_state_gb), - analytic_peak_floor_gb=_round_gb(analytic_peak_floor_gb), - top_memory_buckets=top_memory_buckets, - unsupported_buckets=unsupported_buckets, - ) diff --git a/experiments/local_benchmark/training_sim/scenario_planner.py b/experiments/local_benchmark/training_sim/scenario_planner.py deleted file mode 100644 index a7859492..00000000 --- a/experiments/local_benchmark/training_sim/scenario_planner.py +++ /dev/null @@ -1,1024 +0,0 @@ -"""Plan and score topology scenarios from a base XoRL training config.""" - -from __future__ import annotations - -import argparse -import copy -import json -import math -from pathlib import Path -from typing import Any - - -try: - from .benchmark_behavior import ( - H100_BF16_PROMISED_TFLOPS_PER_GPU, - behavior_point_matches_topology, - behavior_point_matches_workload, - behavior_point_workload_mismatches, - load_benchmark_behavior_points, - predict_benchmark_behavior, - ) - from .config_fingerprint import load_training_config, resolve_topology - from .memory_ledger import build_memory_ledger - from .model_metadata import resolve_model_metadata - from .schemas import ( - BenchmarkBehaviorPoint, - BenchmarkBehaviorPrediction, - ModelMetadata, - ScenarioCandidate, - ScenarioReport, - Topology, - to_jsonable, - ) - from .shape_engine import ShapeLedger, build_shape_ledger -except ImportError: # pragma: no cover - exercised by direct script execution - from benchmark_behavior import ( - H100_BF16_PROMISED_TFLOPS_PER_GPU, - behavior_point_matches_topology, - behavior_point_matches_workload, - behavior_point_workload_mismatches, - load_benchmark_behavior_points, - predict_benchmark_behavior, - ) - from config_fingerprint import load_training_config, resolve_topology - from memory_ledger import build_memory_ledger - from model_metadata import resolve_model_metadata - from schemas import ( - BenchmarkBehaviorPoint, - BenchmarkBehaviorPrediction, - ModelMetadata, - ScenarioCandidate, - ScenarioReport, - Topology, - to_jsonable, - ) - from shape_engine import ShapeLedger, build_shape_ledger - - -def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: - value = raw.get(name, {}) - if isinstance(value, dict): - return value - raw[name] = {} - return raw[name] - - -def _parse_int_list(raw: str | None) -> list[int] | None: - if raw is None or raw == "auto": - return None - values = sorted({int(part.strip()) for part in raw.split(",") if part.strip()}) - if not values: - raise ValueError("expected at least one integer") - return values - - -def _divisors(value: int) -> list[int]: - return [candidate for candidate in range(1, value + 1) if value % candidate == 0] - - -def _power_of_two_divisors(value: int, *, max_value: int | None = None) -> list[int]: - limit = value if max_value is None else min(value, max_value) - return [candidate for candidate in _divisors(value) if candidate <= limit and candidate & (candidate - 1) == 0] - - -def _dedupe_sorted(values: list[int] | set[int]) -> list[int]: - return sorted(value for value in set(values) if value > 0) - - -def _default_micro_batch_sizes( - base_topology: Topology, - behavior_points: list[BenchmarkBehaviorPoint], -) -> list[int]: - values = {base_topology.micro_batch_size} - values.update(point.micro_batch_size for point in behavior_points if point.micro_batch_size is not None) - return sorted(values) - - -def _default_ep_sizes(base_topology: Topology) -> list[int]: - if base_topology.num_experts is None: - return [base_topology.expert_parallel_size] - ranks_per_pipeline_stage = base_topology.world_size // base_topology.pipeline_parallel_size - values = { - value for value in _divisors(ranks_per_pipeline_stage) if value > 0 and base_topology.num_experts % value == 0 - } - if base_topology.expert_parallel_size in values: - return [base_topology.expert_parallel_size] - return sorted(values) or [base_topology.expert_parallel_size] - - -def _auto_ep_sizes(base_topology: Topology) -> list[int]: - if base_topology.num_experts is None: - return [base_topology.expert_parallel_size] - values = { - value - for value in _divisors(base_topology.world_size) - if base_topology.num_experts % value == 0 and value <= base_topology.world_size - } - values.add(base_topology.expert_parallel_size) - return _dedupe_sorted(values) - - -def _auto_tensor_parallel_sizes(base_topology: Topology, metadata: ModelMetadata) -> list[int]: - values = set(_power_of_two_divisors(base_topology.world_size, max_value=base_topology.local_world_size)) - values.add(base_topology.tensor_parallel_size) - if metadata.hidden_size is not None: - values = {value for value in values if metadata.hidden_size % value == 0} - if metadata.num_attention_heads is not None: - values = {value for value in values if metadata.num_attention_heads % value == 0} - return _dedupe_sorted(values) or [base_topology.tensor_parallel_size] - - -def _auto_pipeline_parallel_sizes(base_topology: Topology, metadata: ModelMetadata) -> list[int]: - values = set(_power_of_two_divisors(base_topology.world_size, max_value=4)) - values.add(base_topology.pipeline_parallel_size) - if metadata.num_hidden_layers is not None: - values = {value for value in values if metadata.num_hidden_layers % value == 0} - values.add(base_topology.pipeline_parallel_size) - return _dedupe_sorted(values) or [base_topology.pipeline_parallel_size] - - -def _auto_ulysses_parallel_sizes(base_topology: Topology) -> list[int]: - values = {base_topology.ulysses_parallel_size, 1} - seq_len = base_topology.sample_packing_sequence_len or 0 - if seq_len >= 16_384: - values.update(_power_of_two_divisors(base_topology.world_size, max_value=64)) - return _dedupe_sorted(values) - - -def _auto_ringattn_parallel_sizes(base_topology: Topology) -> list[int]: - values = {base_topology.ringattn_parallel_size, 1} - seq_len = base_topology.sample_packing_sequence_len or 0 - if seq_len >= 64_000: - values.update(_power_of_two_divisors(base_topology.world_size, max_value=4)) - return _dedupe_sorted(values) - - -def _known_or_default_parallel_size(point_value: int | None) -> int: - return point_value if point_value is not None else 1 - - -def _point_matches_topology_parallel_dims(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: - return ( - _known_or_default_parallel_size(point.tensor_parallel_size) == topology.tensor_parallel_size - and _known_or_default_parallel_size(point.pipeline_parallel_size) == topology.pipeline_parallel_size - and _known_or_default_parallel_size(point.ulysses_parallel_size) == topology.ulysses_parallel_size - and _known_or_default_parallel_size(point.ringattn_parallel_size) == topology.ringattn_parallel_size - ) - - -def _point_matches_parallel_dims_for_risk(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: - if point.expert_parallel_size is not None and point.expert_parallel_size != topology.expert_parallel_size: - return False - if point.ep_fsdp_size is not None and point.ep_fsdp_size != topology.ep_fsdp_size: - return False - return _point_matches_topology_parallel_dims(point, topology) - - -def _calibration_scope( - behavior_points: list[BenchmarkBehaviorPoint], - topology: Topology, - *, - prediction_confidence: str, -) -> str: - if prediction_confidence == "calibrated": - return "exact_calibrated" - - throughput_points = [point for point in behavior_points if point.tokens_per_sec is not None] - if not throughput_points: - return "no_calibration" - - same_sequence_points = [ - point - for point in throughput_points - if point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) - ] - if not same_sequence_points: - return "outside_sequence_calibration_envelope" - - dimensions: tuple[tuple[str, int], ...] = ( - ("micro_batch_size", topology.micro_batch_size), - ("global_batch_size", topology.global_batch_size), - ("expert_parallel_size", topology.expert_parallel_size), - ("ep_fsdp_size", topology.ep_fsdp_size or 0), - ("tensor_parallel_size", topology.tensor_parallel_size), - ("pipeline_parallel_size", topology.pipeline_parallel_size), - ("ulysses_parallel_size", topology.ulysses_parallel_size), - ("ringattn_parallel_size", topology.ringattn_parallel_size), - ) - for field_name, topology_value in dimensions: - observed_values = [ - _known_or_default_parallel_size(getattr(point, field_name)) - if field_name.endswith("_parallel_size") - else getattr(point, field_name) - for point in same_sequence_points - if getattr(point, field_name) is not None - ] - if not observed_values: - continue - if topology_value < min(observed_values) or topology_value > max(observed_values): - return "outside_measured_envelope" - return "inside_measured_envelope" - - -def _candidate_risk_flags( - behavior_points: list[BenchmarkBehaviorPoint], - topology: Topology, - behavior: BenchmarkBehaviorPrediction, - *, - raw_config: dict[str, Any] | None, - calibration_scope: str, - prediction_confidence: str, -) -> list[str]: - flags: list[str] = [] - if prediction_confidence != "calibrated": - flags.append("requires_remeasurement") - if calibration_scope.startswith("outside"): - flags.append(calibration_scope) - if behavior.correctness_status and behavior.correctness_status != "k3_pass": - flags.append(f"correctness_{behavior.correctness_status}") - - matched_labels = {part.strip() for part in (behavior.matched_label or "").split(",") if part.strip()} - for point in behavior_points: - if raw_config is not None and point.label in matched_labels: - for mismatch in behavior_point_workload_mismatches(point, raw_config): - flags.append(f"runtime_mismatch:{mismatch}") - same_sequence = point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) - if not same_sequence and point.sample_packing_sequence_len is not None: - same_sequence = ( - topology.sample_packing_sequence_len is not None - and topology.sample_packing_sequence_len >= point.sample_packing_sequence_len - ) - if not same_sequence or not _point_matches_parallel_dims_for_risk(point, topology): - continue - - if point.status == "allocator_pressure_slowdown": - at_or_beyond_mbs = ( - point.micro_batch_size is not None and topology.micro_batch_size >= point.micro_batch_size - ) - at_or_beyond_global_batch = ( - point.global_batch_size is not None and topology.global_batch_size >= point.global_batch_size - ) - if point.label in matched_labels: - flags.append("matched_allocator_pressure_slowdown") - elif at_or_beyond_mbs or at_or_beyond_global_batch: - flags.append(f"allocator_pressure_boundary:{point.label}") - - if point.correctness_status == "oom": - at_or_beyond_sequence = ( - point.sample_packing_sequence_len is not None - and topology.sample_packing_sequence_len is not None - and topology.sample_packing_sequence_len >= point.sample_packing_sequence_len - ) - at_or_beyond_mbs = ( - point.micro_batch_size is not None and topology.micro_batch_size >= point.micro_batch_size - ) - at_or_beyond_global_batch = ( - point.global_batch_size is not None and topology.global_batch_size >= point.global_batch_size - ) - if point.label in matched_labels or ( - at_or_beyond_sequence and (at_or_beyond_mbs or at_or_beyond_global_batch) - ): - flags.append(f"observed_oom_boundary:{point.label}") - - return sorted(set(flags)) - - -def _risk_adjusted_score( - score_tokens_per_sec: float | None, - *, - calibration_scope: str, - risk_flags: list[str], - feasibility_status: str, -) -> float | None: - if score_tokens_per_sec is None: - return None - - multiplier = 1.0 - if calibration_scope == "inside_measured_envelope": - multiplier *= 0.85 - elif calibration_scope == "outside_measured_envelope": - multiplier *= 0.65 - elif calibration_scope == "outside_sequence_calibration_envelope": - multiplier *= 0.35 - elif calibration_scope == "no_calibration": - multiplier *= 0.20 - - if "matched_allocator_pressure_slowdown" in risk_flags: - multiplier *= 0.35 - elif any(flag.startswith("allocator_pressure_boundary:") for flag in risk_flags): - multiplier *= 0.50 - if any(flag.startswith("observed_oom_boundary:") for flag in risk_flags): - multiplier *= 0.25 - - for flag in risk_flags: - if flag == "correctness_k3_fail": - multiplier *= 0.50 - elif flag in {"correctness_not_promoted", "correctness_raw_speed_not_promoted_without_matching_k3_pass"}: - multiplier *= 0.95 - elif flag == "correctness_not_promoted_extrapolated": - multiplier *= 0.90 - elif flag == "correctness_runtime_failure_after_steps": - multiplier *= 0.45 - elif flag == "correctness_missing_calibration": - multiplier *= 0.50 - elif flag.startswith("correctness_"): - multiplier *= 0.75 - - if feasibility_status.endswith("_high_pressure"): - multiplier *= 0.85 - elif feasibility_status.endswith("_moderate_pressure"): - multiplier *= 0.95 - - return round(score_tokens_per_sec * multiplier, 3) - - -def _recommendation( - *, - feasible: bool, - promotable: bool, - feasibility_status: str, - risk_flags: list[str], -) -> str: - if feasibility_status == "observed_oom": - return "avoid_observed_oom" - if not feasible: - return "do_not_launch_unscored" - if promotable: - return "promote_candidate" - if "matched_allocator_pressure_slowdown" in risk_flags or any( - flag.startswith("allocator_pressure_boundary:") for flag in risk_flags - ): - return "measure_allocator_boundary" - if any(flag.startswith("observed_oom_boundary:") for flag in risk_flags): - return "remeasure_after_memory_fix" - if "requires_remeasurement" in risk_flags: - return "remeasure_before_ranking" - if "correctness_runtime_failure_after_steps" in risk_flags: - return "debug_runtime_failure" - if any(flag.startswith("correctness_") for flag in risk_flags): - return "correctness_gate_required" - return "review_candidate" - - -def _mutated_config( - base_config: dict[str, Any], - *, - world_size: int, - micro_batch_size: int, - gradient_accumulation_steps: int, - expert_parallel_size: int, - tensor_parallel_size: int, - pipeline_parallel_size: int, - ulysses_parallel_size: int, - ringattn_parallel_size: int, -) -> dict[str, Any]: - raw_config = copy.deepcopy(base_config) - train = _section(raw_config, "train") - train["micro_batch_size"] = micro_batch_size - train["gradient_accumulation_steps"] = gradient_accumulation_steps - train["expert_parallel_size"] = expert_parallel_size - train["tensor_parallel_size"] = tensor_parallel_size - train["pipeline_parallel_size"] = pipeline_parallel_size - train["ulysses_parallel_size"] = ulysses_parallel_size - train["ringattn_parallel_size"] = ringattn_parallel_size - - non_dp_size = tensor_parallel_size * pipeline_parallel_size * ulysses_parallel_size * ringattn_parallel_size - if non_dp_size <= 0 or world_size % non_dp_size != 0: - raise ValueError("world_size is not divisible by non-DP parallelism product") - data_parallel_size = world_size // non_dp_size - train["data_parallel_replicate_size"] = 1 - train["data_parallel_shard_size"] = data_parallel_size - if pipeline_parallel_size > 1: - train["gradient_accumulation_steps"] = max( - int(train.get("gradient_accumulation_steps", 1) or 1), pipeline_parallel_size - ) - return raw_config - - -def _topology_label(topology: Topology) -> str: - return ( - f"mbs{topology.micro_batch_size}-gb{topology.global_batch_size}-" - f"ep{topology.expert_parallel_size}-efsdp{topology.ep_fsdp_size}-" - f"tp{topology.tensor_parallel_size}-pp{topology.pipeline_parallel_size}-" - f"u{topology.ulysses_parallel_size}-r{topology.ringattn_parallel_size}" - ) - - -def _reference_tokens_per_gpu(point: BenchmarkBehaviorPoint, topology: Topology) -> float | None: - if point.tokens_per_sec is None: - return None - gpu_count = point.gpu_count or topology.world_size - if gpu_count <= 0: - return None - return point.tokens_per_sec / gpu_count - - -def _select_reference_point( - behavior_points: list[BenchmarkBehaviorPoint], - topology: Topology, - raw_config: dict[str, Any] | None = None, -) -> BenchmarkBehaviorPoint | None: - usable = [ - point - for point in behavior_points - if point.tokens_per_sec is not None - and point.micro_batch_size is not None - and point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) - and _reference_tokens_per_gpu(point, topology) is not None - ] - if not usable: - return None - - workload_compatible = ( - [point for point in usable if behavior_point_matches_workload(point, raw_config)] - if raw_config is not None - else usable - ) - same_ep = [ - point for point in workload_compatible if point.expert_parallel_size in (None, topology.expert_parallel_size) - ] - candidates = same_ep or workload_compatible or usable - - def key(point: BenchmarkBehaviorPoint) -> tuple[float, float, float]: - mismatch_count = len(behavior_point_workload_mismatches(point, raw_config)) if raw_config is not None else 0 - mbs_distance = abs((point.micro_batch_size or 1) - topology.micro_batch_size) - per_gpu = _reference_tokens_per_gpu(point, topology) or 0.0 - return (-mismatch_count, -mbs_distance, per_gpu) - - return max(candidates, key=key) - - -def _parallelism_factor(reference: BenchmarkBehaviorPoint, topology: Topology) -> tuple[float, list[str]]: - notes: list[str] = [] - factor = 1.0 - if reference.micro_batch_size: - mbs_ratio = topology.micro_batch_size / reference.micro_batch_size - factor *= min(1.15, max(0.55, mbs_ratio**0.20)) - if reference.expert_parallel_size and reference.expert_parallel_size != topology.expert_parallel_size: - ep_ratio = topology.expert_parallel_size / reference.expert_parallel_size - factor *= max(0.70, 1.0 - 0.04 * abs(math.log2(ep_ratio))) - notes.append(f"EP extrapolated from {reference.expert_parallel_size} to {topology.expert_parallel_size}") - reference_tp = _known_or_default_parallel_size(reference.tensor_parallel_size) - if reference_tp != topology.tensor_parallel_size: - tp_ratio = topology.tensor_parallel_size / reference_tp - factor *= 0.90 ** abs(math.log2(tp_ratio)) - notes.append("TP extrapolation uses conservative communication penalty") - reference_pp = _known_or_default_parallel_size(reference.pipeline_parallel_size) - if reference_pp != topology.pipeline_parallel_size: - pp_delta = abs(topology.pipeline_parallel_size - reference_pp) - factor *= 0.88**pp_delta - notes.append("PP extrapolation uses conservative bubble penalty") - reference_cp = _known_or_default_parallel_size(reference.ulysses_parallel_size) * _known_or_default_parallel_size( - reference.ringattn_parallel_size - ) - if reference_cp != topology.sequence_parallel_size: - cp_ratio = topology.sequence_parallel_size / reference_cp - if topology.sample_packing_sequence_len and topology.sample_packing_sequence_len >= 32768: - factor *= min(1.10, 1.0 + 0.04 * abs(math.log2(cp_ratio))) - else: - factor *= 0.94 ** abs(math.log2(cp_ratio)) - notes.append("SP/CP extrapolation penalized for short-context workload") - return factor, notes - - -def _step_time_fit_prediction( - behavior_points: list[BenchmarkBehaviorPoint], - topology: Topology, - shape: ShapeLedger, - raw_config: dict[str, Any] | None = None, -) -> BenchmarkBehaviorPrediction | None: - if topology.sample_packing_sequence_len is None or shape.global_tokens_per_train_step is None: - return None - compatible = [ - point - for point in behavior_points - if point.tokens_per_sec is not None - and point.global_batch_size is not None - and point.micro_batch_size == topology.micro_batch_size - and point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) - and point.expert_parallel_size in (None, topology.expert_parallel_size) - and point.ep_fsdp_size in (None, topology.ep_fsdp_size) - and _point_matches_topology_parallel_dims(point, topology) - and (raw_config is None or behavior_point_matches_workload(point, raw_config)) - ] - best_by_global_batch: dict[int, BenchmarkBehaviorPoint] = {} - for point in compatible: - current = best_by_global_batch.get(point.global_batch_size) - if current is None or (point.tokens_per_sec or 0.0) > (current.tokens_per_sec or 0.0): - best_by_global_batch[point.global_batch_size] = point - fit_points = sorted(best_by_global_batch.values(), key=lambda point: point.global_batch_size or 0) - if len(fit_points) < 2: - return None - - x_values: list[float] = [] - y_values: list[float] = [] - for point in fit_points: - tokens = point.global_batch_size * topology.sample_packing_sequence_len - step_time = point.step_time_sec - if step_time is None and point.tokens_per_sec: - step_time = tokens / point.tokens_per_sec - if step_time is None: - continue - x_values.append(float(tokens)) - y_values.append(float(step_time)) - if len(x_values) < 2 or len(set(x_values)) < 2: - return None - - x_mean = sum(x_values) / len(x_values) - y_mean = sum(y_values) / len(y_values) - denominator = sum((x_value - x_mean) ** 2 for x_value in x_values) - if denominator == 0: - return None - slope = sum((x_value - x_mean) * (y_value - y_mean) for x_value, y_value in zip(x_values, y_values, strict=False)) - slope /= denominator - intercept = y_mean - slope * x_mean - predicted_step = intercept + slope * shape.global_tokens_per_train_step - if predicted_step <= 0: - return None - tokens_per_sec = shape.global_tokens_per_train_step / predicted_step - tokens_per_sec_per_gpu = tokens_per_sec / topology.world_size - labels = ", ".join(point.label for point in fit_points) - peak_mem_gb = max((point.peak_mem_gb for point in fit_points if point.peak_mem_gb is not None), default=None) - return BenchmarkBehaviorPrediction( - status="extrapolated_step_time_fit", - matched_label=labels, - source="step_time_fit", - tokens_per_sec=round(tokens_per_sec, 3), - tokens_per_sec_per_gpu=round(tokens_per_sec_per_gpu, 3), - step_time_sec=round(predicted_step, 6), - mfu_percent=None, - tflops_per_gpu=None, - promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, - peak_mem_gb=peak_mem_gb, - allocator_retries=None, - derived_global_tokens_per_step=shape.global_tokens_per_train_step, - correctness_status="not_promoted_extrapolated", - warnings=[ - f"extrapolated step time from calibrated global batches: {labels}", - f"fit_intercept_sec={intercept:.6f}", - f"fit_sec_per_token={slope:.12f}", - "correctness must be re-gated before promotion", - ], - ) - - -def _memory_factor( - memory_estimate_gb: float | None, - *, - memory_basis: str, - device_memory_limit_gb: float, - memory_safety_factor: float, -) -> tuple[float, float | None, str]: - if memory_estimate_gb is None: - return 0.0, None, "unknown_memory_estimate" - reserved_memory = memory_estimate_gb * memory_safety_factor - headroom = device_memory_limit_gb - reserved_memory - status_basis = "floor" if memory_basis == "analytic_floor" else memory_basis - if headroom < 0: - if memory_basis == "calibrated_peak" and memory_estimate_gb <= device_memory_limit_gb: - return 0.75, headroom, f"feasible_{status_basis}_high_pressure" - if memory_basis == "analytic_floor": - return 0.0, headroom, "memory_floor_exceeds_limit" - return 0.0, headroom, f"{status_basis}_exceeds_limit" - utilization = reserved_memory / device_memory_limit_gb if device_memory_limit_gb else 1.0 - if utilization >= 0.90: - return 0.75, headroom, f"feasible_{status_basis}_high_pressure" - if utilization >= 0.80: - return 0.90, headroom, f"feasible_{status_basis}_moderate_pressure" - return 1.0, headroom, f"feasible_{status_basis}" - - -def _extrapolate_behavior( - behavior_points: list[BenchmarkBehaviorPoint], - topology: Topology, - shape: ShapeLedger, - *, - raw_config: dict[str, Any] | None = None, - device_memory_limit_gb: float, - memory_safety_factor: float, - analytic_peak_floor_gb: float | None, -) -> tuple[BenchmarkBehaviorPrediction, list[str]]: - step_fit = _step_time_fit_prediction(behavior_points, topology, shape, raw_config=raw_config) - if step_fit is not None: - return step_fit, ["step_time_fit_extrapolation"] - - reference = _select_reference_point(behavior_points, topology, raw_config=raw_config) - if reference is None: - return ( - BenchmarkBehaviorPrediction( - status="unscored", - matched_label=None, - source=None, - tokens_per_sec=None, - tokens_per_sec_per_gpu=None, - step_time_sec=None, - mfu_percent=None, - tflops_per_gpu=None, - promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, - peak_mem_gb=None, - allocator_retries=None, - derived_global_tokens_per_step=shape.global_tokens_per_train_step, - correctness_status="missing_calibration", - warnings=["no benchmark behavior point is available for extrapolation"], - ), - [], - ) - - ref_per_gpu = _reference_tokens_per_gpu(reference, topology) or 0.0 - parallel_factor, notes = _parallelism_factor(reference, topology) - memory_factor, _, memory_status = _memory_factor( - analytic_peak_floor_gb, - memory_basis="analytic_floor", - device_memory_limit_gb=device_memory_limit_gb, - memory_safety_factor=memory_safety_factor, - ) - tokens_per_sec_per_gpu = ref_per_gpu * parallel_factor * memory_factor - tokens_per_sec = tokens_per_sec_per_gpu * topology.world_size - step_time_sec = None - if shape.global_tokens_per_train_step and tokens_per_sec: - step_time_sec = shape.global_tokens_per_train_step / tokens_per_sec - tflops_per_gpu = reference.tflops_per_gpu - if tflops_per_gpu is None and reference.mfu_percent is not None: - tflops_per_gpu = H100_BF16_PROMISED_TFLOPS_PER_GPU * reference.mfu_percent / 100.0 - if tflops_per_gpu is not None and ref_per_gpu: - tflops_per_gpu *= tokens_per_sec_per_gpu / ref_per_gpu - - warnings = [ - f"extrapolated from {reference.label}; correctness must be re-gated before promotion", - f"memory feasibility status is {memory_status}", - ] - if raw_config is not None: - mismatches = behavior_point_workload_mismatches(reference, raw_config) - if mismatches: - warnings.append(f"reference runtime knobs differ: {', '.join(mismatches)}") - warnings.extend(notes) - return ( - BenchmarkBehaviorPrediction( - status="extrapolated", - matched_label=reference.label, - source=reference.source, - tokens_per_sec=round(tokens_per_sec, 3), - tokens_per_sec_per_gpu=round(tokens_per_sec_per_gpu, 3), - step_time_sec=round(step_time_sec, 6) if step_time_sec is not None else None, - mfu_percent=None, - tflops_per_gpu=round(tflops_per_gpu, 3) if tflops_per_gpu is not None else None, - promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, - peak_mem_gb=reference.peak_mem_gb, - allocator_retries=None, - derived_global_tokens_per_step=shape.global_tokens_per_train_step, - correctness_status="not_promoted_extrapolated", - warnings=warnings, - ), - notes, - ) - - -def _candidate_from_prediction( - *, - label: str, - config_path: str | None, - topology: Topology, - shape: ShapeLedger, - behavior: BenchmarkBehaviorPrediction, - prediction_confidence: str, - promotable: bool, - behavior_points: list[BenchmarkBehaviorPoint], - raw_config: dict[str, Any] | None, - device_memory_limit_gb: float, - memory_safety_factor: float, - analytic_peak_floor_gb: float | None, - notes: list[str], -) -> ScenarioCandidate: - estimated_peak_mem_gb = analytic_peak_floor_gb - memory_basis = "analytic_floor" - if behavior.peak_mem_gb is not None: - if analytic_peak_floor_gb is None or behavior.peak_mem_gb >= analytic_peak_floor_gb: - estimated_peak_mem_gb = behavior.peak_mem_gb - memory_basis = "calibrated_peak" if prediction_confidence == "calibrated" else "extrapolated_peak" - - _, headroom, feasibility_status = _memory_factor( - estimated_peak_mem_gb, - memory_basis=memory_basis, - device_memory_limit_gb=device_memory_limit_gb, - memory_safety_factor=memory_safety_factor, - ) - if behavior.status == "calibrated_failure" or behavior.correctness_status == "oom": - feasibility_status = "observed_oom" - if behavior.tokens_per_sec is None and feasibility_status.startswith("feasible"): - feasibility_status = "unscored" - feasible = feasibility_status.startswith("feasible") and behavior.tokens_per_sec is not None - score_tokens_per_sec = behavior.tokens_per_sec if feasible else None - score_tokens_per_gpu_per_sec = behavior.tokens_per_sec_per_gpu if feasible else None - max_ep_slots = max(shape.ep_rank_slots_per_microbatch) if shape.ep_rank_slots_per_microbatch else None - calibration_scope = _calibration_scope( - behavior_points, - topology, - prediction_confidence=prediction_confidence, - ) - risk_flags = _candidate_risk_flags( - behavior_points, - topology, - behavior, - raw_config=raw_config, - calibration_scope=calibration_scope, - prediction_confidence=prediction_confidence, - ) - score_risk_adjusted_tokens_per_sec = _risk_adjusted_score( - score_tokens_per_sec, - calibration_scope=calibration_scope, - risk_flags=risk_flags, - feasibility_status=feasibility_status, - ) - recommendation = _recommendation( - feasible=feasible, - promotable=promotable and feasible, - feasibility_status=feasibility_status, - risk_flags=risk_flags, - ) - return ScenarioCandidate( - label=label, - config_path=config_path, - topology=topology, - behavior=behavior, - prediction_confidence=prediction_confidence, - promotable=promotable and feasible, - feasibility_status=feasibility_status, - score_tokens_per_sec=score_tokens_per_sec, - score_tokens_per_gpu_per_sec=score_tokens_per_gpu_per_sec, - score_risk_adjusted_tokens_per_sec=score_risk_adjusted_tokens_per_sec, - analytic_peak_floor_gb=analytic_peak_floor_gb, - estimated_peak_mem_gb=estimated_peak_mem_gb, - memory_basis=memory_basis, - memory_headroom_gb=round(headroom, 3) if headroom is not None else None, - max_ep_rank_slots_per_microbatch=max_ep_slots, - calibration_scope=calibration_scope, - recommendation=recommendation, - risk_flags=risk_flags, - notes=notes, - ) - - -def _candidate_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: - return ( - candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), - candidate.score_tokens_per_gpu_per_sec if candidate.score_tokens_per_gpu_per_sec is not None else float("-inf"), - ) - - -def _risk_adjusted_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: - return ( - candidate.score_risk_adjusted_tokens_per_sec - if candidate.score_risk_adjusted_tokens_per_sec is not None - else float("-inf"), - candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), - ) - - -def plan_scenario( - base_config_path: str | Path, - *, - benchmark_dir: str | Path | None = None, - world_size: int | None = None, - local_world_size: int | None = None, - micro_batch_sizes: list[int] | None = None, - gradient_accumulation_steps: list[int] | None = None, - expert_parallel_sizes: list[int] | None = None, - tensor_parallel_sizes: list[int] | None = None, - pipeline_parallel_sizes: list[int] | None = None, - ulysses_parallel_sizes: list[int] | None = None, - ringattn_parallel_sizes: list[int] | None = None, - topology_sweep: str = "base", - device_memory_limit_gb: float = 80.0, - memory_safety_factor: float = 1.15, -) -> ScenarioReport: - if topology_sweep not in {"base", "auto"}: - raise ValueError("topology_sweep must be 'base' or 'auto'") - base_path = Path(base_config_path) - base_config = load_training_config(base_path) - base_topology = resolve_topology(base_config, world_size=world_size, local_world_size=local_world_size) - resolved_world_size = world_size or base_topology.world_size - resolved_local_world_size = local_world_size or base_topology.local_world_size - behavior_points = load_benchmark_behavior_points(benchmark_dir) if benchmark_dir is not None else [] - metadata = resolve_model_metadata(base_config) - - micro_batch_values = micro_batch_sizes or _default_micro_batch_sizes(base_topology, behavior_points) - gradient_accumulation_values = gradient_accumulation_steps or [base_topology.gradient_accumulation_steps] - if topology_sweep == "auto": - ep_values = expert_parallel_sizes or _auto_ep_sizes(base_topology) - tp_values = tensor_parallel_sizes or _auto_tensor_parallel_sizes(base_topology, metadata) - pp_values = pipeline_parallel_sizes or _auto_pipeline_parallel_sizes(base_topology, metadata) - ulysses_values = ulysses_parallel_sizes or _auto_ulysses_parallel_sizes(base_topology) - ring_values = ringattn_parallel_sizes or _auto_ringattn_parallel_sizes(base_topology) - else: - ep_values = expert_parallel_sizes or [base_topology.expert_parallel_size] - tp_values = tensor_parallel_sizes or [base_topology.tensor_parallel_size] - pp_values = pipeline_parallel_sizes or [base_topology.pipeline_parallel_size] - ulysses_values = ulysses_parallel_sizes or [base_topology.ulysses_parallel_size] - ring_values = ringattn_parallel_sizes or [base_topology.ringattn_parallel_size] - - candidates: list[ScenarioCandidate] = [] - warnings: list[str] = [] - seen: set[tuple[str, str]] = set() - for pp in pp_values: - for tp in tp_values: - for ulysses in ulysses_values: - for ringattn in ring_values: - for ep in ep_values: - for micro_batch_size in micro_batch_values: - for gradient_accumulation_step in gradient_accumulation_values: - try: - raw_config = _mutated_config( - base_config, - world_size=resolved_world_size, - micro_batch_size=micro_batch_size, - gradient_accumulation_steps=gradient_accumulation_step, - expert_parallel_size=ep, - tensor_parallel_size=tp, - pipeline_parallel_size=pp, - ulysses_parallel_size=ulysses, - ringattn_parallel_size=ringattn, - ) - topology = resolve_topology( - raw_config, - world_size=resolved_world_size, - local_world_size=resolved_local_world_size, - ) - except ValueError as exc: - warnings.append( - f"skipped mbs={micro_batch_size}, ga={gradient_accumulation_step}, " - f"ep={ep}, tp={tp}, pp={pp}, u={ulysses}, r={ringattn}: {exc}" - ) - continue - if topology.ep_fsdp_size is None: - warnings.append(f"skipped {_topology_label(topology)}: ep_fsdp is not integral") - continue - if ( - topology.num_experts is not None - and topology.num_experts % topology.expert_parallel_size - ): - warnings.append( - f"skipped {_topology_label(topology)}: EP does not divide num_experts" - ) - continue - - shape = build_shape_ledger(topology, balanced_routing=True) - memory = build_memory_ledger( - raw_config, - topology=topology, - model_metadata=metadata, - ) - exact_points = [ - point - for point in behavior_points - if behavior_point_matches_topology(point, topology) - and behavior_point_matches_workload(point, raw_config) - ] - if exact_points: - for point in exact_points: - behavior = predict_benchmark_behavior([point], topology, shape, raw_config) - label = f"{_topology_label(topology)}:{point.label}" - key = (label, point.source) - if key in seen: - continue - seen.add(key) - candidates.append( - _candidate_from_prediction( - label=label, - config_path=str(base_path), - topology=topology, - shape=shape, - behavior=behavior, - prediction_confidence="calibrated", - promotable=point.correctness_status == "k3_pass", - behavior_points=behavior_points, - raw_config=raw_config, - device_memory_limit_gb=device_memory_limit_gb, - memory_safety_factor=memory_safety_factor, - analytic_peak_floor_gb=memory.analytic_peak_floor_gb, - notes=list(point.notes), - ) - ) - continue - - behavior, extrapolation_notes = _extrapolate_behavior( - behavior_points, - topology, - shape, - raw_config=raw_config, - device_memory_limit_gb=device_memory_limit_gb, - memory_safety_factor=memory_safety_factor, - analytic_peak_floor_gb=memory.analytic_peak_floor_gb, - ) - label = f"{_topology_label(topology)}:extrapolated" - key = (label, behavior.source or "") - if key in seen: - continue - seen.add(key) - candidates.append( - _candidate_from_prediction( - label=label, - config_path=None, - topology=topology, - shape=shape, - behavior=behavior, - prediction_confidence=behavior.status, - promotable=False, - behavior_points=behavior_points, - raw_config=raw_config, - device_memory_limit_gb=device_memory_limit_gb, - memory_safety_factor=memory_safety_factor, - analytic_peak_floor_gb=memory.analytic_peak_floor_gb, - notes=extrapolation_notes, - ) - ) - - candidates = sorted(candidates, key=_candidate_sort_key, reverse=True) - feasible = [candidate for candidate in candidates if candidate.score_tokens_per_sec is not None] - best_raw = feasible[0] if feasible else None - risk_adjusted = [candidate for candidate in feasible if candidate.score_risk_adjusted_tokens_per_sec is not None] - best_risk_adjusted = max(risk_adjusted, key=_risk_adjusted_sort_key) if risk_adjusted else None - next_measurement = [candidate for candidate in risk_adjusted if "requires_remeasurement" in candidate.risk_flags] - best_next_measurement = max(next_measurement, key=_risk_adjusted_sort_key) if next_measurement else None - promotable = [candidate for candidate in feasible if candidate.promotable] - best_promotable = promotable[0] if promotable else None - if best_raw is not None and not best_raw.promotable: - warnings.append(f"best raw scenario {best_raw.label} is not correctness-promotable") - if best_raw is not None and best_risk_adjusted is not None and best_raw.label != best_risk_adjusted.label: - warnings.append( - f"best raw scenario {best_raw.label} differs from risk-adjusted choice {best_risk_adjusted.label}" - ) - if best_promotable is None: - warnings.append("no correctness-promotable scenario found") - - return ScenarioReport( - base_config_path=str(base_path), - benchmark_dir=str(benchmark_dir) if benchmark_dir is not None else None, - device_memory_limit_gb=device_memory_limit_gb, - memory_safety_factor=memory_safety_factor, - topology_sweep=topology_sweep, - candidate_count=len(candidates), - feasible_count=len(feasible), - best_raw=best_raw, - best_risk_adjusted=best_risk_adjusted, - best_next_measurement=best_next_measurement, - best_promotable=best_promotable, - candidates=candidates, - warnings=warnings, - ) - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config", type=Path, required=True) - parser.add_argument("--benchmark-dir", type=Path, default=None) - parser.add_argument("--world-size", type=int, default=None) - parser.add_argument("--local-world-size", type=int, default=None) - parser.add_argument("--micro-batch-sizes", default=None, help="Comma list, or auto when omitted") - parser.add_argument( - "--gradient-accumulation-steps", default=None, help="Comma list, or base config GA when omitted" - ) - parser.add_argument("--expert-parallel-sizes", default=None, help="Comma list, or base config EP when omitted") - parser.add_argument("--tensor-parallel-sizes", default=None, help="Comma list, or base config TP when omitted") - parser.add_argument("--pipeline-parallel-sizes", default=None, help="Comma list, or base config PP when omitted") - parser.add_argument( - "--ulysses-parallel-sizes", default=None, help="Comma list, or base config Ulysses when omitted" - ) - parser.add_argument("--ringattn-parallel-sizes", default=None, help="Comma list, or base config Ring when omitted") - parser.add_argument( - "--topology-sweep", - choices=("base", "auto"), - default="base", - help="Use base topology dimensions, or derive an automatic TP/PP/CP/EP sweep", - ) - parser.add_argument("--device-memory-limit-gb", type=float, default=80.0) - parser.add_argument("--memory-safety-factor", type=float, default=1.15) - parser.add_argument("--output", type=Path, default=None) - args = parser.parse_args() - - report = plan_scenario( - args.config, - benchmark_dir=args.benchmark_dir, - world_size=args.world_size, - local_world_size=args.local_world_size, - micro_batch_sizes=_parse_int_list(args.micro_batch_sizes), - gradient_accumulation_steps=_parse_int_list(args.gradient_accumulation_steps), - expert_parallel_sizes=_parse_int_list(args.expert_parallel_sizes), - tensor_parallel_sizes=_parse_int_list(args.tensor_parallel_sizes), - pipeline_parallel_sizes=_parse_int_list(args.pipeline_parallel_sizes), - ulysses_parallel_sizes=_parse_int_list(args.ulysses_parallel_sizes), - ringattn_parallel_sizes=_parse_int_list(args.ringattn_parallel_sizes), - topology_sweep=args.topology_sweep, - device_memory_limit_gb=args.device_memory_limit_gb, - memory_safety_factor=args.memory_safety_factor, - ) - rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" - if args.output: - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(rendered, encoding="utf-8") - else: - print(rendered, end="") - - -if __name__ == "__main__": - main() diff --git a/experiments/local_benchmark/training_sim/schemas.py b/experiments/local_benchmark/training_sim/schemas.py deleted file mode 100644 index 4536edf3..00000000 --- a/experiments/local_benchmark/training_sim/schemas.py +++ /dev/null @@ -1,322 +0,0 @@ -"""Dataclasses shared by the local training-engine simulator.""" - -from __future__ import annotations - -from dataclasses import asdict, dataclass, field, is_dataclass -from typing import Any - - -def to_jsonable(value: Any) -> Any: - """Convert simulator dataclasses into plain JSON-compatible containers.""" - if is_dataclass(value): - return {key: to_jsonable(item) for key, item in asdict(value).items()} - if isinstance(value, dict): - return {str(key): to_jsonable(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [to_jsonable(item) for item in value] - return value - - -@dataclass(frozen=True) -class ModelMetadata: - model_path: str | None - config_path: str | None - source: str - num_experts: int | None = None - top_k: int | None = None - num_hidden_layers: int | None = None - hidden_size: int | None = None - intermediate_size: int | None = None - moe_intermediate_size: int | None = None - shared_expert_intermediate_size: int | None = None - num_attention_heads: int | None = None - num_key_value_heads: int | None = None - head_dim: int | None = None - vocab_size: int | None = None - tie_word_embeddings: bool | None = None - - -@dataclass(frozen=True) -class Topology: - world_size: int - local_world_size: int - node_count: int - data_parallel_size: int - data_parallel_replicate_size: int - data_parallel_shard_size: int - tensor_parallel_size: int - pipeline_parallel_size: int - expert_parallel_size: int - ep_fsdp_size: int | None - ulysses_parallel_size: int - ringattn_parallel_size: int - micro_batch_size: int - gradient_accumulation_steps: int - global_batch_size: int - sample_packing_sequence_len: int | None - num_experts: int | None = None - top_k: int | None = None - - @property - def sequence_parallel_size(self) -> int: - return self.ulysses_parallel_size * self.ringattn_parallel_size - - -@dataclass(frozen=True) -class RunFingerprint: - config_path: str - config_sha256: str - config_name: str - repo_commit: str | None - balanced_routing: bool - topology: Topology - model_metadata: ModelMetadata - - -@dataclass(frozen=True) -class BalancedRoutingLedger: - total_slots: int - num_experts: int - counts_by_expert: list[int] - max_slots_per_expert: int - min_slots_per_expert: int - imbalance_slots: int - - -@dataclass(frozen=True) -class ShapeLedger: - microbatch_tokens_per_dp_rank: int | None - global_tokens_per_microbatch: int | None - global_tokens_per_train_step: int | None - tokens_per_gpu_per_train_step: float | None - sequence_parallel_size: int - tokens_per_model_rank_per_microbatch: int | None - routed_slots_per_model_rank_microbatch: int | None - routed_slots_per_train_step_model_rank: int | None - balanced_routing: BalancedRoutingLedger | None - ep_rank_slots_per_microbatch: list[int] | None = None - warnings: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class StepObservation: - source: str - step: int - max_steps: str - loss: float | None = None - grad_norm: float | None = None - lr: float | None = None - tflops_per_gpu: float | None = None - mfu: float | None = None - tokens_per_sec: float | None = None - step_time_s: float | None = None - peak_mem_gb: float | None = None - phase_memory_gb: dict[str, float] = field(default_factory=dict) - extra: dict[str, float] = field(default_factory=dict) - - -@dataclass(frozen=True) -class PhaseObservation: - source: str - prefix: str - step: int - max_steps: str - metrics: dict[str, float] - - -@dataclass(frozen=True) -class MemoryPhaseObservation: - source: str - prefix: str - step: int - max_steps: str - metrics: dict[str, float] - - -@dataclass(frozen=True) -class ObservedRun: - sources: list[str] - steps: list[StepObservation] - phases: list[PhaseObservation] = field(default_factory=list) - memory_phases: list[MemoryPhaseObservation] = field(default_factory=list) - - -@dataclass(frozen=True) -class MemoryBucket: - name: str - gb: float - source: str - notes: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class MemoryLedger: - deepep_buffer_size_gb: float | None - observed_peak_mem_gb_max: float | None - observed_phase_peak_gb: dict[str, float] - estimated_total_params_b: float | None = None - estimated_local_params_b: float | None = None - persistent_model_state_gb: float | None = None - gradient_state_gb: float | None = None - optimizer_state_gb: float | None = None - analytic_peak_floor_gb: float | None = None - top_memory_buckets: list[MemoryBucket] = field(default_factory=list) - unsupported_buckets: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class BenchmarkBehaviorPoint: - label: str - source: str - micro_batch_size: int | None - global_batch_size: int | None - tokens_per_sec: float | None - step_time_sec: float | None - mfu_percent: float | None = None - tflops_per_gpu: float | None = None - peak_mem_gb: float | None = None - allocator_retries: int | None = None - measured_steps: int | None = None - warmup_steps: int | None = None - gpu_count: int | None = None - sample_packing_sequence_len: int | None = None - tensor_parallel_size: int | None = None - pipeline_parallel_size: int | None = None - ulysses_parallel_size: int | None = None - ringattn_parallel_size: int | None = None - expert_parallel_size: int | None = None - ep_fsdp_size: int | None = None - deepep_async_combine: bool | None = None - deepep_num_sms: int | None = None - deepep_buffer_size_gb: float | None = None - enable_compile: bool | None = None - gradient_checkpointing_method: str | None = None - enable_activation_offload: bool | None = None - activation_offload_prefetch_count: int | None = None - status: str = "observed" - correctness_status: str | None = None - notes: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class BenchmarkBehaviorPrediction: - status: str - matched_label: str | None - source: str | None - tokens_per_sec: float | None - tokens_per_sec_per_gpu: float | None - step_time_sec: float | None - mfu_percent: float | None - tflops_per_gpu: float | None - promised_tflops_per_gpu: float | None - peak_mem_gb: float | None - allocator_retries: int | None - derived_global_tokens_per_step: int | None - correctness_status: str | None = None - warnings: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class PredictionReport: - fingerprint: RunFingerprint - shape: ShapeLedger - memory: MemoryLedger - benchmark_behavior: BenchmarkBehaviorPrediction | None = None - observed_summary: dict[str, Any] | None = None - calibration_sources: list[str] = field(default_factory=list) - warnings: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class TradeoffCandidate: - label: str - config_path: str | None - behavior_source: str - topology: Topology | None - behavior: BenchmarkBehaviorPrediction - promotable: bool - score_tokens_per_sec: float | None - score_tflops_per_gpu: float | None - notes: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class TradeoffReport: - benchmark_dir: str - status: str - candidate_count: int - best_raw: TradeoffCandidate | None - best_promotable: TradeoffCandidate | None - candidates: list[TradeoffCandidate] - warnings: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class ScenarioCandidate: - label: str - config_path: str | None - topology: Topology - behavior: BenchmarkBehaviorPrediction - prediction_confidence: str - promotable: bool - feasibility_status: str - score_tokens_per_sec: float | None - score_tokens_per_gpu_per_sec: float | None - score_risk_adjusted_tokens_per_sec: float | None - analytic_peak_floor_gb: float | None - estimated_peak_mem_gb: float | None - memory_basis: str - memory_headroom_gb: float | None - max_ep_rank_slots_per_microbatch: int | None - calibration_scope: str - recommendation: str - risk_flags: list[str] = field(default_factory=list) - notes: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class ScenarioReport: - base_config_path: str - benchmark_dir: str | None - device_memory_limit_gb: float - memory_safety_factor: float - topology_sweep: str - candidate_count: int - feasible_count: int - best_raw: ScenarioCandidate | None - best_risk_adjusted: ScenarioCandidate | None - best_next_measurement: ScenarioCandidate | None - best_promotable: ScenarioCandidate | None - candidates: list[ScenarioCandidate] - warnings: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class CalibrationHoldout: - label: str - source: str - topology_label: str - actual_tokens_per_sec: float - predicted_tokens_per_sec: float | None - prediction_status: str - matched_label: str | None - absolute_error_tokens_per_sec: float | None - absolute_percentage_error: float | None - calibrated_from_count: int - warnings: list[str] = field(default_factory=list) - - -@dataclass(frozen=True) -class CalibrationReport: - base_config_path: str - benchmark_dir: str - status: str - measured_point_count: int - evaluated_count: int - skipped_count: int - mean_absolute_percentage_error: float | None - median_absolute_percentage_error: float | None - max_absolute_percentage_error: float | None - prediction_status_counts: dict[str, int] - holdouts: list[CalibrationHoldout] - warnings: list[str] = field(default_factory=list) diff --git a/experiments/local_benchmark/training_sim/validate_benchmarks.py b/experiments/local_benchmark/training_sim/validate_benchmarks.py deleted file mode 100644 index 4b8f5a00..00000000 --- a/experiments/local_benchmark/training_sim/validate_benchmarks.py +++ /dev/null @@ -1,457 +0,0 @@ -"""Validate simulator output against checked-in throughput benchmark recipes.""" - -from __future__ import annotations - -import argparse -import json -import os -import re -import subprocess -import sys -import tempfile -from dataclasses import replace -from pathlib import Path -from typing import Any - - -try: - from .benchmark_behavior import ( - H100_BF16_PROMISED_TFLOPS_PER_GPU, - load_benchmark_behavior_points, - predict_benchmark_behavior, - ) - from .config_fingerprint import load_training_config - from .predict import build_report - from .schemas import to_jsonable - from .shape_engine import build_shape_ledger -except ImportError: # pragma: no cover - exercised by direct script execution - from benchmark_behavior import ( - H100_BF16_PROMISED_TFLOPS_PER_GPU, - load_benchmark_behavior_points, - predict_benchmark_behavior, - ) - from config_fingerprint import load_training_config - from predict import build_report - from schemas import to_jsonable - from shape_engine import build_shape_ledger - - -def _human_number(value: str) -> float: - cleaned = value.strip().replace(",", "").lstrip("~") - multiplier = 1.0 - if cleaned.endswith(("K", "k")): - cleaned = cleaned[:-1] - multiplier = 1_000.0 - elif cleaned.endswith(("M", "m")): - cleaned = cleaned[:-1] - multiplier = 1_000_000.0 - return float(cleaned) * multiplier - - -def _parse_readme_metrics(readme_text: str) -> dict[str, Any]: - metrics: dict[str, Any] = {} - if match := re.search(r"Hardware:\s*(?P\d+)\s*nodes?\s*x\s*(?P\d+)\s*H100", readme_text): - metrics["nodes"] = int(match.group("nodes")) - metrics["gpus_per_node"] = int(match.group("gpus")) - metrics["world_size"] = metrics["nodes"] * metrics["gpus_per_node"] - if match := re.search(r"sample_packing_sequence_len:\s*(?P\d+)", readme_text): - metrics["sample_packing_sequence_len"] = int(match.group("seq")) - if match := re.search(r"\|\s*tokens/sec\s*\|\s*(?P~?[0-9.]+[KkMm]?)\s*\|", readme_text): - metrics["tokens_per_sec"] = _human_number(match.group("value")) - if match := re.search(r"\|\s*step time\s*\|\s*(?P~?[0-9.]+)s\s*\|", readme_text): - metrics["step_time_sec"] = float(match.group("value").lstrip("~")) - if match := re.search(r"\|\s*MFU\s*\|\s*(?P~?[0-9.]+)%", readme_text): - metrics["mfu_percent"] = float(match.group("value").lstrip("~")) - if match := re.search(r"\|\s*allocated memory\s*\|\s*(?P~?[0-9.]+)GB\s*\|", readme_text): - metrics["allocated_memory_gb"] = float(match.group("value").lstrip("~")) - if match := re.search(r"`mbs=10`[^~]+~(?P[0-9.]+)K tok/s", readme_text): - metrics["mbs10_tokens_per_sec"] = _human_number(match.group("value") + "K") - return metrics - - -def _check(name: str, expected: Any, actual: Any, *, tolerance: float | None = None) -> dict[str, Any]: - if tolerance is None: - passed = expected == actual - else: - passed = expected is not None and actual is not None and abs(float(expected) - float(actual)) <= tolerance - return { - "name": name, - "status": "pass" if passed else "fail", - "expected": expected, - "actual": actual, - "tolerance": tolerance, - } - - -def _benchmark_config_path(benchmark_dir: Path) -> Path: - configs = sorted((benchmark_dir / "configs").glob("*.yaml")) - if not configs: - raise FileNotFoundError(f"no benchmark configs found under {benchmark_dir / 'configs'}") - if len(configs) > 1: - raise ValueError(f"expected one config under {benchmark_dir / 'configs'}, found {len(configs)}") - return configs[0] - - -def _render_script_text(benchmark_dir: Path) -> str: - script_path = benchmark_dir / "scripts" / "render_k8s_manifest.sh" - return script_path.read_text(encoding="utf-8") if script_path.is_file() else "" - - -def _render_manifest_text(benchmark_dir: Path) -> str: - script_path = benchmark_dir / "scripts" / "render_k8s_manifest.sh" - if not script_path.is_file(): - return "" - with tempfile.TemporaryDirectory(prefix="xorl-training-sim-") as tmpdir: - output_path = Path(tmpdir) / "manifest.yaml" - env = os.environ.copy() - env["OUTPUT"] = str(output_path) - result = subprocess.run( - [str(script_path)], - check=False, - cwd=Path.cwd(), - env=env, - capture_output=True, - text=True, - ) - if result.returncode != 0: - return result.stdout + result.stderr - return output_path.read_text(encoding="utf-8") - - -def _validate_config_behavior( - benchmark_dir: Path, - readme_metrics: dict[str, Any], -) -> tuple[dict[str, Any], list[dict[str, Any]], dict[str, dict[str, Any]]]: - config_path = _benchmark_config_path(benchmark_dir) - raw_config = load_training_config(config_path) - world_size = readme_metrics.get("world_size") - local_world_size = readme_metrics.get("gpus_per_node") - report = build_report( - config_path, - world_size=world_size, - local_world_size=local_world_size, - balanced_routing=True, - num_experts=None, - top_k=None, - benchmark_dir=benchmark_dir, - ) - topology = report.fingerprint.topology - shape = report.shape - model = raw_config.get("model", {}) - data = raw_config.get("data", {}) - train = raw_config.get("train", {}) - - checks = [ - _check("readme_reference_tokens_per_sec", 261000.0, readme_metrics.get("tokens_per_sec")), - _check("readme_reference_step_time_sec", 8.04, readme_metrics.get("step_time_sec")), - _check("readme_reference_mfu_percent", 16.2, readme_metrics.get("mfu_percent")), - _check("readme_reference_allocated_memory_gb", 56.4, readme_metrics.get("allocated_memory_gb")), - _check("readme_mbs10_tokens_per_sec", 133700.0, readme_metrics.get("mbs10_tokens_per_sec")), - _check( - "readme_mbs10_allocator_pressure_slowdown", - True, - readme_metrics.get("mbs10_tokens_per_sec", 0) < readme_metrics.get("tokens_per_sec", 0) * 0.6, - ), - _check("world_size", world_size, topology.world_size), - _check("local_world_size", local_world_size, topology.local_world_size), - _check("pipeline_parallel_size", 1, topology.pipeline_parallel_size), - _check("tensor_parallel_size", 1, topology.tensor_parallel_size), - _check("ringattn_parallel_size", 1, topology.ringattn_parallel_size), - _check("ulysses_parallel_size", 1, topology.ulysses_parallel_size), - _check( - "sample_packing_sequence_len", - readme_metrics.get("sample_packing_sequence_len"), - topology.sample_packing_sequence_len, - ), - _check("micro_batch_size", 8, topology.micro_batch_size), - _check("global_batch_size", 256, topology.global_batch_size), - _check("data_parallel_replicate_size", 1, topology.data_parallel_replicate_size), - _check("expert_parallel_size", 8, topology.expert_parallel_size), - _check("ep_fsdp_size", 4, topology.ep_fsdp_size), - _check("data_parallel_shard_size", 32, topology.data_parallel_shard_size), - _check("num_experts", 256, topology.num_experts), - _check("top_k", 8, topology.top_k), - _check("dataset_path", "dummy", data.get("datasets", [{}])[0].get("path")), - _check("dataset_type", "tokenized", data.get("datasets", [{}])[0].get("type")), - _check("sample_packing_method", "sequential", data.get("sample_packing_method")), - _check("moe_implementation", "quack", model.get("moe_implementation")), - _check("ep_dispatch", "deepep", model.get("ep_dispatch")), - _check("train_router", False, model.get("train_router")), - _check("deepep_buffer_size_gb", 2.0, model.get("deepep_buffer_size_gb")), - _check("deepep_num_sms", 72, model.get("deepep_num_sms")), - _check("deepep_async_combine", True, model.get("deepep_async_combine")), - _check("data_parallel_mode", "fsdp2", train.get("data_parallel_mode")), - _check("gradient_checkpointing_method", "recompute_full_layer", train.get("gradient_checkpointing_method")), - _check("optimizer", "adamw", train.get("optimizer")), - _check("enable_mixed_precision", True, train.get("enable_mixed_precision")), - _check("enable_full_shard", True, train.get("enable_full_shard")), - _check("init_device", "meta", train.get("init_device")), - _check("load_weights_mode", "grouped", train.get("load_weights_mode")), - _check("enable_compile", True, train.get("enable_compile")), - _check("empty_cache_steps", 10, train.get("empty_cache_steps")), - _check("gc_steps", 10, train.get("gc_steps")), - _check("save_steps", 0, train.get("save_steps")), - _check("save_epochs", 0, train.get("save_epochs")), - _check("log_format", "structured", train.get("log_format")), - _check("global_tokens_per_train_step", 2_097_408, shape.global_tokens_per_train_step), - ] - behavior_points = load_benchmark_behavior_points(benchmark_dir) - behavior_prediction = predict_benchmark_behavior(behavior_points, topology, shape, raw_config) - behavior_labels = sorted(point.label for point in behavior_points) - behavior_predictions = _predict_all_behavior_points(behavior_points, topology) - checks.extend( - [ - _check( - "benchmark_behavior_points", - [ - "qwen36_static_k3_summary_20260519:q36-main-af98064-deepepenv-05190533", - "readme_adjacent_mbs10_allocator_pressure", - "readme_reference_mbs8", - ], - behavior_labels, - ), - _check("benchmark_behavior_prediction_status", "calibrated", behavior_prediction.status), - _check("benchmark_behavior_prediction_label", "readme_reference_mbs8", behavior_prediction.matched_label), - _check( - "benchmark_behavior_tokens_per_sec", - readme_metrics.get("tokens_per_sec"), - behavior_prediction.tokens_per_sec, - ), - _check( - "benchmark_behavior_step_time_sec", - readme_metrics.get("step_time_sec"), - behavior_prediction.step_time_sec, - ), - _check( - "benchmark_behavior_peak_mem_gb", - readme_metrics.get("allocated_memory_gb"), - behavior_prediction.peak_mem_gb, - ), - _check("benchmark_behavior_allocator_retries", 0, behavior_prediction.allocator_retries), - _check( - "benchmark_behavior_promised_tflops_per_gpu", - H100_BF16_PROMISED_TFLOPS_PER_GPU, - behavior_prediction.promised_tflops_per_gpu, - ), - _check("benchmark_behavior_tflops_per_gpu", 160.218, behavior_prediction.tflops_per_gpu, tolerance=0.001), - ] - ) - for label, prediction in behavior_predictions.items(): - point = next(point for point in behavior_points if point.label == label) - variant_shape = prediction["shape"] - variant_behavior = prediction["behavior"] - checks.extend( - [ - _check(f"behavior_matrix:{label}:prediction_status", "calibrated", variant_behavior["status"]), - _check(f"behavior_matrix:{label}:prediction_label", label, variant_behavior["matched_label"]), - _check( - f"behavior_matrix:{label}:tokens_per_sec", point.tokens_per_sec, variant_behavior["tokens_per_sec"] - ), - _check( - f"behavior_matrix:{label}:global_tokens_per_step", - (point.global_batch_size or 0) * (topology.sample_packing_sequence_len or 0), - variant_shape["global_tokens_per_train_step"], - ), - ] - ) - if point.step_time_sec is not None: - checks.append( - _check( - f"behavior_matrix:{label}:step_time_sec", - point.step_time_sec, - variant_behavior["step_time_sec"], - tolerance=0.01, - ) - ) - if point.mfu_percent is not None: - checks.extend( - [ - _check( - f"behavior_matrix:{label}:mfu_percent", - point.mfu_percent, - variant_behavior["mfu_percent"], - ), - _check( - f"behavior_matrix:{label}:promised_tflops_per_gpu", - H100_BF16_PROMISED_TFLOPS_PER_GPU, - variant_behavior["promised_tflops_per_gpu"], - ), - _check( - f"behavior_matrix:{label}:tflops_per_gpu", - H100_BF16_PROMISED_TFLOPS_PER_GPU * point.mfu_percent / 100.0, - variant_behavior["tflops_per_gpu"], - tolerance=0.001, - ), - ] - ) - if point.tokens_per_sec and variant_shape["global_tokens_per_train_step"]: - checks.append( - _check( - f"behavior_matrix:{label}:tokens_imply_step_time", - variant_behavior["step_time_sec"], - variant_shape["global_tokens_per_train_step"] / point.tokens_per_sec, - tolerance=0.05, - ) - ) - if readme_metrics.get("tokens_per_sec") and shape.global_tokens_per_train_step: - derived_step_time = shape.global_tokens_per_train_step / readme_metrics["tokens_per_sec"] - checks.append( - _check( - "readme_tokens_per_sec_implies_step_time", - readme_metrics.get("step_time_sec"), - derived_step_time, - tolerance=0.05, - ) - ) - if report.shape.balanced_routing is not None: - counts = report.shape.balanced_routing.counts_by_expert - checks.extend( - [ - _check("balanced_routing_imbalance_slots", 1, report.shape.balanced_routing.imbalance_slots), - _check("balanced_routing_count_sum", report.shape.balanced_routing.total_slots, sum(counts)), - _check("experts_per_ep_rank", 32, topology.num_experts // topology.expert_parallel_size), - ] - ) - script_text = _render_script_text(benchmark_dir) - manifest_text = _render_manifest_text(benchmark_dir) - checks.append( - _check( - "render_script_sets_balanced_synthetic_routing_env", - True, - "XORL_MOE_SYNTHETIC_ROUTING" in script_text and "balanced" in script_text, - ) - ) - checks.extend( - [ - _check("rendered_manifest_sets_team_turbo", True, "team: turbo" in manifest_text), - _check( - "rendered_manifest_sets_balanced_synthetic_routing_env", - True, - "name: XORL_MOE_SYNTHETIC_ROUTING" in manifest_text and 'value: "balanced"' in manifest_text, - ), - _check( - "rendered_manifest_sets_nccl_socket_ifname", - True, - "name: NCCL_SOCKET_IFNAME" in manifest_text and 'value: "bond0"' in manifest_text, - ), - _check("rendered_manifest_sets_runtime_class", True, "runtimeClassName: nvidia" in manifest_text), - ] - ) - return to_jsonable(report), checks, behavior_predictions - - -def _predict_all_behavior_points(behavior_points, topology) -> dict[str, dict[str, Any]]: - predictions: dict[str, dict[str, Any]] = {} - for point in behavior_points: - if not point.micro_batch_size or not point.global_batch_size: - continue - denom = point.micro_batch_size * topology.data_parallel_size - if point.global_batch_size % denom != 0: - continue - variant_topology = replace( - topology, - micro_batch_size=point.micro_batch_size, - gradient_accumulation_steps=point.global_batch_size // denom, - global_batch_size=point.global_batch_size, - ) - variant_shape = build_shape_ledger(variant_topology, balanced_routing=True) - predictions[point.label] = { - "shape": to_jsonable(variant_shape), - "behavior": to_jsonable(predict_benchmark_behavior(behavior_points, variant_topology, variant_shape)), - } - return predictions - - -def _validate_result_json(benchmark_dir: Path, readme_metrics: dict[str, Any]) -> list[dict[str, Any]]: - checks: list[dict[str, Any]] = [] - seq_len = int(readme_metrics.get("sample_packing_sequence_len") or 0) - for result_path in sorted((benchmark_dir / "results").glob("*.json")): - result = json.loads(result_path.read_text(encoding="utf-8")) - throughput = result.get("throughput", {}) - if throughput and seq_len: - global_tokens = throughput.get("global_batch_size", 0) * seq_len - derived_step_time = ( - global_tokens / throughput["tokens_per_sec"] if throughput.get("tokens_per_sec") else None - ) - checks.extend( - [ - _check( - f"{result_path.name}:throughput_candidate", - "q36-main-af98064-deepepenv-05190533", - throughput.get("candidate"), - ), - _check(f"{result_path.name}:throughput_gpus", 32, throughput.get("gpus")), - _check(f"{result_path.name}:throughput_tokens_per_sec", 254600.0, throughput.get("tokens_per_sec")), - _check( - f"{result_path.name}:derived_step_time_sec", - throughput.get("step_time_sec"), - derived_step_time, - tolerance=0.01, - ), - ] - ) - k3_gate = result.get("k3_gate", {}) - if k3_gate: - checks.extend( - [ - _check(f"{result_path.name}:k3_gate_status", "fail", k3_gate.get("status")), - _check(f"{result_path.name}:k3_total_tokens", 192, k3_gate.get("k3", {}).get("total_tokens")), - _check( - f"{result_path.name}:k3_primary_failure", "k3.mean <= 0.001", k3_gate.get("primary_failure") - ), - ] - ) - diagnostics = result.get("diagnostic_replays", []) - if diagnostics: - checks.append(_check(f"{result_path.name}:diagnostic_replay_count", 3, len(diagnostics))) - checks.append( - _check( - f"{result_path.name}:diagnostic_low_k3_rows", - 2, - sum(1 for row in diagnostics if row.get("status") == "diagnostic_low_k3"), - ) - ) - return checks - - -def validate_benchmark_dir(benchmark_dir: Path) -> dict[str, Any]: - readme_path = benchmark_dir / "README.md" - if not readme_path.is_file(): - raise FileNotFoundError(f"missing README.md in {benchmark_dir}") - readme_metrics = _parse_readme_metrics(readme_path.read_text(encoding="utf-8")) - report, config_checks, behavior_predictions = _validate_config_behavior(benchmark_dir, readme_metrics) - result_checks = _validate_result_json(benchmark_dir, readme_metrics) - checks = config_checks + result_checks - status = "pass" if all(check["status"] == "pass" for check in checks) else "fail" - return { - "benchmark_dir": str(benchmark_dir), - "status": status, - "readme_metrics": readme_metrics, - "simulator_report": report, - "behavior_points": to_jsonable(load_benchmark_behavior_points(benchmark_dir)), - "behavior_predictions": behavior_predictions, - "checks": checks, - } - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--benchmarks-root", type=Path, required=True) - parser.add_argument("--model", required=True, help="Benchmark model subdirectory to validate") - parser.add_argument("--output", type=Path, default=None) - parser.add_argument("--no-fail-on-error", action="store_true", help="Always exit 0 after writing the report") - args = parser.parse_args() - - payload = validate_benchmark_dir(args.benchmarks_root / args.model) - rendered = json.dumps(payload, indent=2, sort_keys=True) + "\n" - if args.output: - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text(rendered, encoding="utf-8") - else: - print(rendered, end="") - if payload["status"] != "pass" and not args.no_fail_on_error: - sys.exit(1) - - -if __name__ == "__main__": - main() diff --git a/pyproject.toml b/pyproject.toml index ca8affd5..60519e61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ dependencies = [ "numpy<2.4", "numba", "pydantic", + "pyyaml", "uvicorn", "pyzmq", "msgpack", @@ -80,6 +81,18 @@ dependencies = [ # uv pip install fast-hadamard-transform ] +[project.scripts] +xorl-sim-calibrate = "xorl.sim.calibration_evaluator:main" +xorl-sim-collect = "xorl.sim.collect_calibration:main" +xorl-sim-feasibility = "xorl.sim.feasibility_evaluator:main" +xorl-sim-kernels = "xorl.sim.kernel_variants:main" +xorl-sim-packs = "xorl.sim.calibration_packs:main" +xorl-sim-plan = "xorl.sim.scenario_planner:main" +xorl-sim-predict = "xorl.sim.predict:main" +xorl-sim-rank = "xorl.sim.tradeoff_ranker:main" +xorl-sim-validate = "xorl.sim.validate:main" + + [dependency-groups] # Follow the best practice in https://docs.astral.sh/uv/concepts/projects/dependencies/#development-dependencies # to manage dev dependencies (i.e., dependencies that are only used in development) like @@ -115,6 +128,9 @@ version = {attr = "xorl.__version__"} [tool.setuptools.packages.find] where = ["src"] +[tool.setuptools.package-data] +"xorl.sim" = ["README.md", "calibration_packs/**/*"] + [tool.ruff] target-version = "py312" line-length = 120 diff --git a/src/xorl/sim/README.md b/src/xorl/sim/README.md new file mode 100644 index 00000000..dcf34cb2 --- /dev/null +++ b/src/xorl/sim/README.md @@ -0,0 +1,137 @@ +# XoRL Training Simulator + +`xorl.sim` is the installable, CPU-only planning and calibration surface for XoRL local training. It resolves distributed +topology, calculates token and model shapes, builds analytical FLOP/activation/communication/memory ledgers, ingests +structured trainer logs, and ranks measured or extrapolated training scenarios. + +The package deliberately contains no cluster launcher, scheduler configuration, storage configuration, or workload +manifest. Execution infrastructure remains outside the simulator. + +## Installed Commands + +After installing XoRL, the following commands are available: + +| Command | Purpose | +| --- | --- | +| `xorl-sim-packs` | List, locate, or validate built-in calibration packs | +| `xorl-sim-predict` | Build a static report for one config | +| `xorl-sim-plan` | Compare micro-batch and parallelism scenarios | +| `xorl-sim-calibrate` | Run leave-one-out prediction validation | +| `xorl-sim-feasibility` | Replay measured fit and OOM boundaries | +| `xorl-sim-rank` | Rank raw and correctness-promotable measurements | +| `xorl-sim-kernels` | Rank portable kernel-variant measurements | +| `xorl-sim-validate` | Run pack sanitation and analytical golden gates | + +## Built-In Qwen Packs + +```bash +xorl-sim-packs list +xorl-sim-packs validate +xorl-sim-predict --pack qwen3_6_35b_a3b --balanced-routing +xorl-sim-plan --pack qwen3_5_397b_a17b --micro-batch-sizes 4,5 +xorl-sim-feasibility --pack qwen3_235b_a22b +``` + +The built-in packs are: + +- `qwen3_235b_a22b`: 4-node 2k-context GA observations and an OOM boundary. +- `qwen3_5_397b_a17b`: 8-node short-context throughput rows with separate raw and static-K3-promotable winners. +- `qwen3_6_35b_a3b`: 4-node 8k throughput, allocator-pressure, and failed static-K3 evidence. + +Each pack contains `manifest.json`, portable YAML configs, summarized observations, limitations, and frozen golden values. +Pack validation rejects missing files, model/config mismatches, internal absolute paths, and infrastructure-specific content. + +## External Calibration Packs + +All APIs continue to accept a filesystem benchmark directory. A portable pack has this shape: + +```text +my_model/ + manifest.json + README.md or RESULTS.md + configs/ + training.yaml + results/ + measurements.json +``` + +The manifest schema version is currently `1`. `configs` and `results` are relative paths, `default_config` selects the +prediction baseline, and `golden` freezes the expected behavior count, topology, throughput, and analytical memory floor. +Use `xorl-sim-packs validate path/to/my_model` before relying on a new pack. + +The behavior reader also accepts results roots containing resolved run directories. A directory containing +`xorl_cli.yaml` plus a sibling `node-0.log` is ingested as observed evidence. Structured `[STEP ...]`, `[STEP_PHASES ...]`, +and `[STEP_MEMORY ...]` rows provide throughput, timing, and memory calibration. OOM logs become measured failure +boundaries rather than being discarded. + +## Prediction And Planning + +```bash +xorl-sim-predict \ + --config path/to/xorl_config.yaml \ + --balanced-routing \ + --benchmark-dir path/to/calibration_pack \ + --output prediction.json + +xorl-sim-plan \ + --config path/to/xorl_config.yaml \ + --benchmark-dir path/to/calibration_pack \ + --micro-batch-sizes 1,2,4 \ + --expert-parallel-sizes 8,16 \ + --topology-sweep auto \ + --output scenarios.json +``` + +Reports distinguish exact calibration, interpolation, extrapolation, observed failures, and missing support. Raw throughput, +risk-adjusted throughput, correctness promotion, memory feasibility, timing coverage, and communication scope remain +separate fields. An extrapolated candidate is never silently promoted as a measured winner. + +## Analytical Core + +The analytical surface includes: + +- Dense and MoE FLOP ledgers matching XoRL trainer accounting. +- Hybrid full-attention/GatedDeltaNet model metadata and parameter ownership. +- Activation lower bounds with recompute, CE-logit, attention, and MoE buffer terms. +- Per-rank FSDP, TP, PP, EP, and context-parallel communication bytes and cross-node scope. +- Parameter, gradient, optimizer, unshard, calibrated residual, and peak-memory attribution. +- Observed or benchmark-derived timing ledgers and explicit support/blocker reporting. + +Analytical memory remains a lower bound until a pack supplies calibrated peak or residual evidence. Timing requires measured +step or phase data; the simulator does not infer hardware timing from communication bytes alone. + +## Kernel Variants + +Kernel comparison consumes portable JSON measurements instead of loading files from a particular machine: + +```json +[ + { + "family": "attention", + "variant": "candidate_a", + "workload": "qwen35-seq4096", + "latency_ms": 8.5, + "correctness_status": "pass" + } +] +``` + +```bash +xorl-sim-kernels measurements.json +``` + +Measurements must share a kernel family and workload. By default, the fastest ungated row is reported but cannot become +the selected winner. + +## Validation + +```bash +xorl-sim-validate --output validation.json +``` + +The consolidated validator checks every installed pack, frozen behavior counts, default topology, calibrated throughput, +correctness-aware rankings, analytical memory goldens, and FLOP/activation/communication coverage. It exits nonzero on any +failed check and is intended to run in CPU CI and against an installed wheel. + +Built-in results are historical synthetic-data observations on H100 GPUs. They are regression fixtures and calibration +priors, not universal performance claims or substitutes for a fresh correctness gate on a changed runtime. diff --git a/src/xorl/sim/__init__.py b/src/xorl/sim/__init__.py new file mode 100644 index 00000000..703e8ba3 --- /dev/null +++ b/src/xorl/sim/__init__.py @@ -0,0 +1 @@ +"""Portable training-engine simulation, calibration, and planning APIs.""" diff --git a/src/xorl/sim/analytical_ledgers.py b/src/xorl/sim/analytical_ledgers.py new file mode 100644 index 00000000..8a0b3acc --- /dev/null +++ b/src/xorl/sim/analytical_ledgers.py @@ -0,0 +1,1510 @@ +"""Analytical-first ledgers for the XoRL training simulator (q35/q30 lane). + +These are *equation-based* predictions computed from model metadata + config/topology BEFORE any +measured log is consulted. Measured runs are used only as validation rows (predicted-vs-measured), +never as the model. Every term carries a ``status`` provenance field: + +- ``exact_analytic`` closed-form from shapes/config; no runtime coefficient. +- ``exact_analytic_lower_bound`` closed-form lower bound from shapes/config; no runtime coefficient. +- ``analytic_with_runtime_coefficient`` closed-form scaled by a documented runtime constant. +- ``calibrated_residual`` gap between analytic floor and a measured value. +- ``unsupported`` not yet modeled (named so the surface is auditable). +- ``not_applicable`` term is structurally absent for this config. + +The FLOPs ledger mirrors the trainer's ``src/xorl/utils/count_flops.py::XorlFlopsCounter`` +conventions, so analytical TFLOPS computed with a measured step time should match the logged TFLOPS. +Both q35 (Qwen3.5-35B-A3B) and q30 (Qwen3-30B-A3B) use the ``qwen3_moe`` path; dense Qwen3 models +such as Qwen3-8B use the trainer's qwen3/qwen2 dense-decoder path. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + + +try: + from .memory_ledger import ( + _dtype_bytes, + _estimate_param_breakdown, + _gradient_storage_bytes, + _local_param_ownership, + _muon_param_partition, + _param_storage_bytes, + ) + from .schemas import ModelMetadata, Topology +except ImportError: # pragma: no cover - exercised by direct script execution + from memory_ledger import ( + _dtype_bytes, + _estimate_param_breakdown, + _gradient_storage_bytes, + _local_param_ownership, + _muon_param_partition, + _param_storage_bytes, + ) + from schemas import ModelMetadata, Topology + + +# fwd+bwd multipliers WITHOUT recompute (matches XorlFlopsCounter._grad_ckpt_multipliers, which is +# deliberately checkpointing-independent so logged TFLOPS is stable across recompute strategies): +# 6 = 2 (multiply-add) x 3 (1 forward + 2 backward); attention score has two matmuls => 12. +_M_LINEAR = 6 +_M_ATTN_SCORE = 12 +_M_LM_HEAD = 6 +H100_BF16_PEAK_TFLOPS = 989.0 # matches src count_flops get_device_flops H100 (989e12) + MFU denom +# Exposed (non-overlapped) cross-node comm cost per GB of cross-node traffic, in ms. Calibrated from +# the measured 1-node->2-node step-time delta at identical per-rank work, cross-model-validated: +# q30=6.38, q35=5.82 ms/GB (agree within ~9%) -> mean ~6.1. Intra-node (NVLink) comm is overlapped. +# Recalibrated 2026-07-04 after the expert-FSDP byte-accounting fix (the old 6.1 was fitted against +# ~5x-overcounted expert collective bytes — two compensating errors). New per-model coefficients from +# the same measured 1->2-node step deltas over the CORRECTED exposed bytes: q30 29.08, q35 24.29 +# ms/GB (mean 26.7, cross-model spread +-9%). Independently confirmed by the standalone 2-node +# collective microbench (23.0-24.6 ms/GiB at MoE per-layer message sizes): the coefficient now +# decomposes as the real NCCL collective rate at message size with near-zero overlap for the exposed +# collectives, rather than an opaque fudge. +EXPOSED_CROSS_NODE_MS_PER_GB = 26.7 +H100_NVLINK_EFFECTIVE_GB_PER_S = 450.0 +H100_NDR400_UNIDIRECTIONAL_GB_PER_S = 50.0 +H100_NDR400_FULL_DUPLEX_GB_PER_S = 2 * H100_NDR400_UNIDIRECTIONAL_GB_PER_S +MEMORY_PEAK_VALIDATION_THRESHOLD = 0.05 + + +def comm_exposed_time_coefficient( + *, one_node_step_s: float, two_node_step_s: float, cross_node_gb: float +) -> dict[str, Any]: + """Calibrate the exposed cross-node comm coefficient (ms per cross-node GB) from node scaling. + + At identical per-rank work, the 1-node->2-node step-time increase is attributed to the cross-node + comm that was free NVLink at 1 node. Coefficient = delta_step_time / cross_node_GB. Validated by + agreement across models (q30/q35). Confounds (documented): ep_fsdp1->ep_fsdp2 and a small forward + growth are folded into the coefficient, so it is an upper-bound on pure cross-node-comm exposure. + """ + if not cross_node_gb or two_node_step_s <= 0: + return {"status": "uncomparable"} + delta = two_node_step_s - one_node_step_s + return { + "status": "calibrated", + "delta_step_time_s": round(delta, 4), + "cross_node_gb": round(cross_node_gb, 3), + "exposed_ms_per_gb": round(delta / cross_node_gb * 1000.0, 3), + "exposed_fraction_of_2node_step": round(delta / two_node_step_s, 4), + } + + +def static_cross_node_overlap_estimate( + terms: dict[str, dict[str, Any]], + *, + local_world_size: int, + num_experts: int | None = None, + top_k: int | None = None, + expert_parallel_size: int | None = None, +) -> dict[str, Any]: + """Static hardware comparison for exposed 2-node FSDP comm time. + + This is intentionally separate from the calibrated exposed-ms/GB model. It uses exact byte terms + plus a declared H100/NDR400 full-duplex link constant, and assumes the backward tail is dominated + by cross-node FSDP gradient reduce-scatter while parameter all-gather and intra-node collectives + are hidden by FSDP prefetch/layer compute. Validation decides whether that static assumption is + predictive for a measured topology. + """ + grad = terms.get("fsdp_grad_reduce_scatter", {}) + param = terms.get("fsdp_param_all_gather", {}) + expert_grad = terms.get("expert_fsdp_grad_reduce_scatter", {}) + expert_param = terms.get("expert_fsdp_param_all_gather", {}) + exposed_cross_gib = float(grad.get("cross_gb") or 0.0) + float(expert_grad.get("cross_gb") or 0.0) + param_passes = max(float(param.get("passes") or 1.0), 1.0) + prefetched_cross_gib_total = float(param.get("cross_gb") or 0.0) + float(expert_param.get("cross_gb") or 0.0) + prefetched_cross_gib = prefetched_cross_gib_total / param_passes + if exposed_cross_gib <= 0.0 and prefetched_cross_gib <= 0.0: + return { + "status": "intra_node_fully_overlapped", + "prediction_basis": "no cross-node FSDP traffic", + "predicted_backward_cross_node_exposed_s": 0.0, + } + predicted = exposed_cross_gib * _BYTES_PER_GIB / (H100_NDR400_FULL_DUPLEX_GB_PER_S * 1e9) + prefetched_serial = prefetched_cross_gib * _BYTES_PER_GIB / (H100_NDR400_FULL_DUPLEX_GB_PER_S * 1e9) + + def aggregate_two_node_transport_s(term_name: str, *, per_pass: bool = False) -> float | None: + term = terms.get(term_name, {}) + if int(term.get("nodes_spanned") or 0) != 2: + return None + logical_per_rank_gib = float(term.get("gb") or 0.0) + if per_pass: + logical_per_rank_gib /= max(float(term.get("passes") or 1.0), 1.0) + if logical_per_rank_gib <= 0.0 or local_world_size <= 0: + return None + # For a two-node collective, the node-pair transport lower bound is the bidirectional + # traffic crossing the node boundary divided by the aggregate full-duplex NIC bandwidth of + # the local GPUs. This is distinct from each rank's logical collective bytes. + aggregate_node_pair_gb_per_s = local_world_size * H100_NDR400_FULL_DUPLEX_GB_PER_S + bidirectional_node_pair_gib = 2.0 * logical_per_rank_gib + return bidirectional_node_pair_gib * _BYTES_PER_GIB / (aggregate_node_pair_gb_per_s * 1e9) + + aggregate_grad_s = aggregate_two_node_transport_s("fsdp_grad_reduce_scatter") + aggregate_param_s = aggregate_two_node_transport_s("fsdp_param_all_gather", per_pass=True) + aggregate_transport: dict[str, Any] + if aggregate_grad_s is None: + aggregate_transport = { + "status": "not_applicable", + "reason": "requires two-node FSDP collective terms", + } + else: + aggregate_transport = { + "status": "static_aggregate_node_pair_transport_bracket", + "prediction_basis": ( + "two-node FSDP physical transport lower/upper bracket from exact grad collective GiB, " + "one backward param-all-gather pass, and aggregate full-duplex H100/NDR400 bandwidth " + "across local GPUs" + ), + "local_world_size": int(local_world_size), + "aggregate_full_duplex_node_pair_gb_per_s": round( + local_world_size * H100_NDR400_FULL_DUPLEX_GB_PER_S, + 3, + ), + "fsdp_grad_reduce_scatter_transport_s": round(aggregate_grad_s, 4), + "fsdp_param_all_gather_transport_s": round(aggregate_param_s or 0.0, 4), + "fsdp_grad_only_lower_bound_s": round(aggregate_grad_s, 4), + "fsdp_grad_plus_param_upper_bound_s": round(aggregate_grad_s + (aggregate_param_s or 0.0), 4), + "term_status": "exact_analytic_bytes_plus_declared_aggregate_hardware_bandwidth", + } + topology_expert_overlap: dict[str, Any] + if ( + aggregate_grad_s is not None + and aggregate_param_s is not None + and num_experts is not None + and top_k is not None + and expert_parallel_size not in (None, 0) + ): + local_experts_per_ep_rank = float(num_experts) / float(expert_parallel_size) + active_expert_overlap_slots = float(2 * top_k) + param_exposure_fraction = min(1.0, active_expert_overlap_slots / local_experts_per_ep_rank) + predicted_exposed_s = aggregate_grad_s + aggregate_param_s * param_exposure_fraction + topology_expert_overlap = { + "status": "static_topology_expert_overlap_model", + "prediction_basis": ( + "two-node exposed FSDP time = aggregate grad reduce-scatter transport + " + "one aggregate backward param all-gather transport * min(1, 2 * top_k / " + "local_experts_per_ep_rank); the overlap fraction is topology-derived from expert " + "inventory per EP rank" + ), + "local_world_size": int(local_world_size), + "num_experts": int(num_experts), + "top_k": int(top_k), + "expert_parallel_size": int(expert_parallel_size), + "local_experts_per_ep_rank": round(local_experts_per_ep_rank, 3), + "active_expert_overlap_slots": round(active_expert_overlap_slots, 3), + "fsdp_grad_reduce_scatter_transport_s": round(aggregate_grad_s, 4), + "fsdp_param_all_gather_transport_s": round(aggregate_param_s, 4), + "param_all_gather_exposure_fraction": round(param_exposure_fraction, 6), + "predicted_backward_cross_node_exposed_s": round(predicted_exposed_s, 4), + "term_status": ("exact_analytic_bytes_plus_declared_aggregate_hardware_bandwidth_and_topology_overlap"), + } + else: + topology_expert_overlap = { + "status": "not_applicable", + "reason": "requires aggregate FSDP transport plus num_experts, top_k, and expert_parallel_size", + } + return { + "status": "static_hardware_overlap_model_compared", + "prediction_basis": ( + "exact cross-node FSDP grad reduce-scatter GiB divided by declared H100 NDR400 " + "full-duplex per-GPU bandwidth; FSDP param all-gather is treated as prefetched/overlapped" + ), + "exposed_collectives": ["fsdp_grad_reduce_scatter"], + "prefetched_or_overlapped_collectives": [ + "fsdp_param_all_gather_forward_pass_and_hidden_backward_fraction", + "ep_all_to_all_dispatch_combine", + "intra_node_collectives", + ], + "fsdp_grad_reduce_scatter_cross_gib": round(exposed_cross_gib, 4), + "fsdp_param_all_gather_cross_gib_treated_as_overlapped": round(prefetched_cross_gib, 4), + "fsdp_param_all_gather_cross_gib_total": round(prefetched_cross_gib_total, 4), + "fsdp_param_all_gather_passes": round(param_passes, 3), + "ndr400_unidirectional_gb_per_s": H100_NDR400_UNIDIRECTIONAL_GB_PER_S, + "full_duplex_cross_node_gb_per_s": H100_NDR400_FULL_DUPLEX_GB_PER_S, + "predicted_backward_cross_node_exposed_s": round(predicted, 4), + "prefetched_param_all_gather_serial_s_if_unhidden": round(prefetched_serial, 4), + "aggregate_node_pair_transport_estimate": aggregate_transport, + "topology_expert_overlap_estimate": topology_expert_overlap, + "term_status": "exact_analytic_bytes_plus_declared_hardware_bandwidth_assumption", + } + + +def _seq_len(topology: Topology, seq_len: int | None) -> int | None: + return seq_len if seq_len is not None else topology.sample_packing_sequence_len + + +def flops_ledger( + metadata: ModelMetadata, + topology: Topology, + *, + seq_len: int | None = None, + batch_seqlens: list[int] | None = None, +) -> dict[str, Any]: + """Per-step analytical FLOPs for Qwen3 dense/MoE models, mirroring the trainer's accounting. + + Global FLOPs = sum over all DP-replicated tokens in one optimizer step (global_batch_size + sequences of length seq_len). Per-GPU FLOPs = global / world_size (the trainer logs per-GPU). + """ + is_moe = ( + metadata.num_experts is not None and metadata.top_k is not None and metadata.moe_intermediate_size is not None + ) + fields = ( + metadata.hidden_size, + metadata.num_hidden_layers, + metadata.vocab_size, + metadata.num_attention_heads, + ) + seq = _seq_len(topology, seq_len) + if any(value is None for value in fields) or (seq is None and not batch_seqlens): + return {"status": "unsupported", "reason": "missing_model_metadata_or_seq_len", "components": {}} + if is_moe and metadata.moe_intermediate_size is None: + return {"status": "unsupported", "reason": "missing_moe_intermediate_size", "components": {}} + if not is_moe and metadata.intermediate_size is None: + return {"status": "unsupported", "reason": "missing_dense_intermediate_size", "components": {}} + + hidden = metadata.hidden_size + layers = metadata.num_hidden_layers + vocab = metadata.vocab_size + n_heads = metadata.num_attention_heads + n_kv = metadata.num_key_value_heads or n_heads + head_dim = metadata.head_dim or (hidden // n_heads) + num_experts = metadata.num_experts or 0 + top_k = metadata.top_k or 0 + moe_inter = metadata.moe_intermediate_size or 0 + dense_inter = metadata.intermediate_size or 0 + + if batch_seqlens is not None: + resolved_batch_seqlens = [int(value) for value in batch_seqlens if int(value) > 0] + num_sequences = len(resolved_batch_seqlens) + tokens_global = sum(resolved_batch_seqlens) + score_elements = sum(s * (s + 1) // 2 for s in resolved_batch_seqlens) + sequence_source = "batch_seqlens" + else: + resolved_batch_seqlens = None + num_sequences = topology.global_batch_size + tokens_global = num_sequences * seq + # causal attention score elements per packed sequence: S(S+1)/2 + # (matches _attention_score_elements) + score_elements = num_sequences * (seq * (seq + 1) // 2) + sequence_source = "seq_len" + + q_size = n_heads * head_dim + kv_size = n_kv * head_dim + o_size = n_heads * head_dim + attn_linear_n = hidden * (q_size + 2 * kv_size + o_size) + lm_head_n = vocab * hidden # embedding lookup is 0 FLOPs + + components = {} + if is_moe: + router_n = hidden * num_experts + gate_up_n = hidden * moe_inter * top_k * 2 # gate_proj + up_proj + down_n = hidden * moe_inter * top_k # down_proj + components.update( + { + "moe_router": _M_LINEAR * router_n * tokens_global * layers, + "moe_gate_up_proj": _M_LINEAR * gate_up_n * tokens_global * layers, + "moe_down_proj": _M_LINEAR * down_n * tokens_global * layers, + } + ) + source = "mirror:src/xorl/utils/count_flops.py::_estimate_qwen3_moe_flops" + else: + dense_mlp_n = hidden * dense_inter * 3 # gate_proj + up_proj + down_proj + components["dense_mlp"] = _M_LINEAR * dense_mlp_n * tokens_global * layers + source = "mirror:src/xorl/utils/count_flops.py::_estimate_qwen2_flops (qwen3 dense)" + + components.update( + { + "attn_qkvo_proj": _M_LINEAR * attn_linear_n * tokens_global * layers, + "attn_score_quadratic": _M_ATTN_SCORE * score_elements * head_dim * n_heads * layers, + "lm_head": _M_LM_HEAD * lm_head_n * tokens_global, + } + ) + total_flops = float(sum(components.values())) + world_size = max(topology.world_size, 1) + per_gpu_flops = total_flops / world_size + return { + "status": "exact_analytic", + "source": source, + "multipliers": { + "linear_fwd_bwd": _M_LINEAR, + "attn_score_fwd_bwd": _M_ATTN_SCORE, + "lm_head_fwd_bwd": _M_LM_HEAD, + }, + "recompute_in_flops": False, + "seq_len": seq, + "sequence_source": sequence_source, + "batch_seqlens": resolved_batch_seqlens, + "num_sequences": num_sequences, + "tokens_global": tokens_global, + "components_flops": {name: float(value) for name, value in components.items()}, + "total_flops": total_flops, + "per_gpu_flops": per_gpu_flops, + "mfu_denominator_tflops_per_gpu": H100_BF16_PEAK_TFLOPS, + "term_status": dict.fromkeys(components, "exact_analytic"), + } + + +_BYTES_PER_GIB = 1024**3 + + +def _act_bytes(train: dict[str, Any]) -> int: + # Activations are stored in the autocast/compute dtype under mixed precision (bf16). + return 2 if train.get("enable_mixed_precision") else 4 + + +def activation_ledger( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], + *, + seq_len: int | None = None, +) -> dict[str, Any]: + """Per-rank analytical activation-memory LOWER BOUND, by named term, with provenance. + + These are the tensors that must be live for one microbatch; they are an analytic lower bound, + not the allocator-reserved peak. The gap between this and the measured residual is attributed to + unmodeled allocator/NCCL/DeepEP workspace slack (see ``memory_residual_attribution``). + """ + needed = (metadata.hidden_size, metadata.num_hidden_layers, metadata.vocab_size) + seq = _seq_len(topology, seq_len) + if any(v is None for v in needed) or seq is None: + return {"status": "unsupported", "reason": "missing_metadata_or_seq_len", "terms": {}} + + hidden = metadata.hidden_size + layers = metadata.num_hidden_layers + vocab = metadata.vocab_size + top_k = metadata.top_k or 0 + moe_inter = metadata.moe_intermediate_size or 0 + dense_inter = metadata.intermediate_size or 0 + sp = max(topology.sequence_parallel_size, 1) + act = _act_bytes(train) + model_rank_tokens = (topology.micro_batch_size * seq + sp - 1) // sp # ceil; one microbatch live at a time + routed_slots_rank = model_rank_tokens * top_k + is_moe = top_k > 0 and moe_inter > 0 + + ckpt_method = str(train.get("gradient_checkpointing_method", "") or "") + full_recompute = bool(train.get("enable_gradient_checkpointing")) and ckpt_method in { + "recompute_full_layer", + "full", + } + # logit dtype for CE: lm_head_fp32 forces fp32 logits; chunked by ce_num_chunks. + logit_bytes = 4 if (_section_get(train, "lm_head_fp32") or _section_get(train, "router_fp32")) else act + ce_chunks = int(train.get("ce_num_chunks", 1) or 1) + + terms: dict[str, dict[str, Any]] = {} + + def add(name: str, byte_count: float, status: str, note: str) -> None: + terms[name] = {"gb": round(byte_count / _BYTES_PER_GIB, 4), "status": status, "note": note} + + activation_term_status = "exact_analytic_lower_bound" + if full_recompute: + add( + "saved_layer_inputs", + layers * model_rank_tokens * hidden * act, + activation_term_status, + "recompute_full_layer: only layer-boundary inputs are stashed for backward", + ) + if is_moe: + add( + "recompute_working_set_one_layer", + model_rank_tokens * top_k * moe_inter * act * 2, + activation_term_status, + "one layer recomputed at a time in backward; MoE gate/up intermediate dominates", + ) + else: + add( + "recompute_working_set_one_layer", + model_rank_tokens * dense_inter * act * 2, + activation_term_status, + "one dense layer recomputed at a time in backward; dense gate/up intermediate dominates", + ) + else: + add( + "saved_full_activations", + layers * model_rank_tokens * (hidden + top_k * moe_inter) * act, + activation_term_status, + "no full-layer recompute: per-layer activations retained", + ) + add( + "ce_logit_buffer", + (model_rank_tokens / max(ce_chunks, 1)) * vocab * logit_bytes, + activation_term_status, + f"chunked CE logits: tokens/{ce_chunks} x vocab x {logit_bytes}B (lm_head_fp32 -> 4B)", + ) + if is_moe: + add( + "moe_dispatch_combine_buffer", + routed_slots_rank * hidden * act * 2, + activation_term_status, + "alltoall dispatch + combine token buffers (routed slots x hidden)", + ) + else: + add( + "moe_dispatch_combine_buffer", + 0.0, + "not_applicable", + "dense model: no routed experts or EP dispatch/combine buffers", + ) + add( + "attention_workspace", + model_rank_tokens * hidden * act * 2, + activation_term_status, + "flash-attention working buffers (out + lse), conservative", + ) + total_gb = round(sum(t["gb"] for t in terms.values()), 4) + return { + "status": "exact_analytic_lower_bound", + "act_bytes": act, + "model_rank_tokens": model_rank_tokens, + "routed_slots_per_rank": routed_slots_rank, + "full_layer_recompute": full_recompute, + "terms": terms, + "analytic_activation_lower_bound_gb": total_gb, + "unmodeled_terms": [ + "allocator_reserved_slack", + "nccl_deepep_workspace" if is_moe else "nccl_workspace", + "cuda_context_and_fragmentation", + ], + } + + +def _section_get(train: dict[str, Any], key: str) -> Any: + return train.get(key) + + +def memory_residual_attribution( + *, + analytic_floor_gb: float | None, + measured_peak_gb: float | None, + activation_lower_bound_gb: float | None, +) -> dict[str, Any]: + """Split measured peak into: param/grad/opt floor + analytic activation lower bound + unmodeled slack.""" + if analytic_floor_gb is None or measured_peak_gb is None: + return {"status": "incomplete"} + residual = measured_peak_gb - analytic_floor_gb + act = activation_lower_bound_gb or 0.0 + unmodeled = residual - act + return { + "status": "attributed", + "measured_peak_gb": round(measured_peak_gb, 3), + "param_grad_opt_floor_gb": round(analytic_floor_gb, 3), + "param_grad_opt_floor_status": "exact_analytic", + "analytic_activation_lower_bound_gb": round(act, 3), + "analytic_activation_status": "exact_analytic_lower_bound", + "unmodeled_allocator_workspace_gb": round(unmodeled, 3), + "unmodeled_status": "calibrated_residual", + "residual_gb": round(residual, 3), + "residual_fraction_of_peak": round(residual / measured_peak_gb, 4) if measured_peak_gb else None, + "activation_explains_fraction_of_residual": round(act / residual, 4) if residual > 0 else None, + } + + +def _muon_update_dtype_bytes( + train: dict[str, Any], gradient_bytes: int, optimizer_dtype_bytes: int +) -> tuple[int, list[str]]: + update_dtype = train.get("muon_update_dtype") + if update_dtype is not None: + return _dtype_bytes(update_dtype, default=optimizer_dtype_bytes), [f"muon_update_dtype={update_dtype}"] + + momentum = float(train.get("muon_momentum", 0.95) or 0.0) + force_momentum = bool(train.get("muon_force_momentum_path")) + if momentum > 0 or force_momentum: + momentum_dtype = train.get("muon_momentum_dtype") + if momentum_dtype is not None: + return _dtype_bytes(momentum_dtype, default=optimizer_dtype_bytes), [ + f"muon_momentum_dtype={momentum_dtype}", + "update_dtype_inherits_momentum_buffer", + ] + optimizer_dtype = train.get("optimizer_dtype") + if optimizer_dtype is not None: + return optimizer_dtype_bytes, [ + f"optimizer_dtype={optimizer_dtype}", + "runtime_default_muon_momentum_dtype_inherits_optimizer_dtype", + "update_dtype_inherits_momentum_buffer", + ] + return gradient_bytes, ["update_dtype_inherits_gradient_dtype_without_explicit_momentum_dtype"] + + grad_dtype = train.get("muon_grad_dtype") + if grad_dtype is not None: + return _dtype_bytes(grad_dtype, default=gradient_bytes), [ + f"muon_grad_dtype={grad_dtype}", + "momentum_zero_update_dtype_inherits_muon_grad_dtype", + ] + return gradient_bytes, ["momentum_zero_update_dtype_inherits_gradient_dtype"] + + +def optimizer_step_ledger( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], +) -> dict[str, Any]: + """Optimizer-state and optimizer-step work terms by XoRL optimizer type. + + This is the optimizer-specific quantity used by the phase-time model. It is deliberately separate + from FLOPs: AdamW-like optimizers are memory/state scans, SignSGD is a state-free param/grad scan, + and Muon adds Newton-Schulz work that still needs optimizer-specific calibration. + """ + breakdown = _estimate_param_breakdown(metadata) + if breakdown is None: + return {"status": "unsupported", "reason": "param_breakdown_unavailable"} + + local_params, local_non_expert_params, local_expert_params, ownership_notes = _local_param_ownership( + breakdown, + train, + topology, + ) + optimizer = str(train.get("optimizer", "adamw") or "adamw").lower() + weight_bytes, weight_notes = _param_storage_bytes(train) + gradient_bytes, gradient_notes = _gradient_storage_bytes(train) + optimizer_dtype_bytes = _dtype_bytes(train.get("optimizer_dtype"), default=4) + param_gb = local_params * weight_bytes / _BYTES_PER_GIB + gradient_gb = local_params * gradient_bytes / _BYTES_PER_GIB + state_terms: dict[str, dict[str, Any]] = {} + transient_terms: dict[str, dict[str, Any]] = {} + notes = [*ownership_notes, *weight_notes, *gradient_notes] + + def add_state(name: str, params: float, bytes_per_param: int, status: str, note: str) -> None: + state_terms[name] = { + "gb": round(params * bytes_per_param / _BYTES_PER_GIB, 3), + "bytes_per_param": bytes_per_param, + "status": status, + "note": note, + } + + def add_transient(name: str, params: float, bytes_per_param: int, copies: int, status: str, note: str) -> None: + transient_terms[name] = { + "gb": round(params * bytes_per_param * copies / _BYTES_PER_GIB, 3), + "params": round(params), + "bytes_per_param": bytes_per_param, + "retained_copies": copies, + "status": status, + "note": note, + } + + status = "exact_analytic" + unsupported_terms: list[str] = [] + optimizer_family = optimizer + + if optimizer == "adamw": + if train.get("cautious_weight_decay"): + state_bytes = 4 + optimizer_family = "anyprecision_adamw" + notes.append("optimizer=adamw+cautious_weight_decay routes to AnyPrecisionAdamW fp32 state") + else: + state_bytes = weight_bytes + notes.append("torch.optim.AdamW states follow parameter dtype") + add_state("exp_avg", local_params, state_bytes, "exact_analytic", "AdamW first moment") + add_state("exp_avg_sq", local_params, state_bytes, "exact_analytic", "AdamW second moment") + elif optimizer == "anyprecision_adamw": + state_bytes = optimizer_dtype_bytes + reuse_grad = bool(train.get("anyprecision_adamw_reuse_grad_for_momentum")) + state_offload = bool(train.get("anyprecision_adamw_state_cpu_offload")) + if state_offload: + status = "analytic_with_runtime_coefficient" + notes.append("state_cpu_offload=true: GPU persistent state is offloaded between optimizer steps") + if reuse_grad: + add_state( + "exp_avg_aliases_gradient", + local_params, + 0, + "aliases_gradient_storage", + "reuse_grad_for_momentum=true stores first moment in the consumed grad tensor", + ) + else: + add_state("exp_avg", local_params, state_bytes, "exact_analytic", "AnyPrecisionAdamW first moment") + add_state("exp_avg_sq", local_params, state_bytes, "exact_analytic", "AnyPrecisionAdamW second moment") + if train.get("anyprecision_adamw_use_kahan_summation") or train.get("use_kahan_summation"): + add_state("compensation", local_params, state_bytes, "exact_analytic", "Kahan compensation buffer") + if state_offload: + for term in state_terms.values(): + if term["status"] == "exact_analytic": + term["gpu_persistent_gb"] = 0.0 + term["note"] += "; offloaded between steps" + elif optimizer == "sgd": + momentum = float(train.get("momentum", train.get("sgd_momentum", 0.0)) or 0.0) + if momentum > 0: + add_state("momentum_buffer", local_params, weight_bytes, "exact_analytic", f"SGD momentum={momentum}") + else: + notes.append("SGD momentum=0: state-free optimizer") + elif optimizer in {"signsgd", "distsignsgd"}: + notes.append(f"{optimizer} is state-free; distsignsgd signs gradients before FSDP reduce-scatter") + elif optimizer == "muon": + momentum = float(train.get("muon_momentum", 0.95) or 0.0) + force_momentum = bool(train.get("muon_force_momentum_path")) + fallback = str(train.get("muon_fallback_optimizer", "adamw") or "adamw").lower() + partition = _muon_param_partition(metadata, topology, train, local_params=local_params) + partition_exact = partition.get("status") in { + "exact_analytic_dense_muon_partition", + "exact_analytic_moe_muon_partition", + } + if partition_exact: + muon_matrix_params = float(partition["local_muon_matrix_params"]) + fallback_params = float(partition["local_fallback_params"]) + if momentum > 0 or force_momentum: + add_state( + "muon_momentum_buffer", + muon_matrix_params, + optimizer_dtype_bytes, + "exact_analytic", + ( + f"Muon momentum={momentum}; ndim>=2 dense projection matrix params use Muon " + "per runtime classifier" + ), + ) + else: + notes.append("Muon momentum=0 and force_momentum_path=false: no Muon momentum buffer") + if fallback == "adamw": + add_state( + "muon_fallback_exp_avg", + fallback_params, + optimizer_dtype_bytes, + "exact_analytic", + "Muon fallback AdamW first moment for non-Muon dense params", + ) + add_state( + "muon_fallback_exp_avg_sq", + fallback_params, + optimizer_dtype_bytes, + "exact_analytic", + "Muon fallback AdamW second moment for non-Muon dense params", + ) + elif fallback == "sgd": + notes.append("Muon fallback SGD is state-free for non-Muon dense params") + else: + unsupported_terms.append(f"muon_fallback_optimizer:{fallback}") + ns_algorithm = str(train.get("muon_ns_algorithm", "gram_newton_schulz") or "gram_newton_schulz") + update_bytes, update_notes = _muon_update_dtype_bytes(train, gradient_bytes, optimizer_dtype_bytes) + if ns_algorithm == "gram_newton_schulz": + retained_copies = 2 + transient_note = ( + "grouped Gram-NS retains the pre-orthogonalization update entries and the " + "orthogonalized output pieces until the parameter update loop" + ) + transient_status = "exact_analytic" + else: + retained_copies = 1 + transient_note = "standard Muon retains one orthogonalized update tensor per Muon parameter" + transient_status = "exact_analytic_lower_bound" + add_transient( + "muon_retained_update_tensors", + muon_matrix_params, + update_bytes, + retained_copies, + transient_status, + transient_note, + ) + notes.extend( + [ + f"Muon partition exact: {partition['status']}; ndim>=2 projection/expert matrices use Muon; " + "embeddings/lm_head/norms/router gates use fallback where present", + f"muon_matrix_params={partition['local_muon_matrix_params']}", + f"muon_fallback_params={partition['local_fallback_params']}", + f"muon_ns_algorithm={ns_algorithm}", + *update_notes, + ] + ) + status = "exact_analytic_lower_bound" + unsupported_terms.append("muon_newton_schulz_compute_and_kernel_time") + elif momentum > 0 or force_momentum: + add_state( + "muon_momentum_buffer", + local_params, + optimizer_dtype_bytes, + "analytic_with_runtime_coefficient", + f"Muon momentum={momentum}; metadata-level model treats all local trainable params as Muon-state scanned", + ) + status = "analytic_with_runtime_coefficient" + else: + notes.append("Muon momentum=0 and force_momentum_path=false: no Muon momentum buffer") + if partition.get("status") != "exact_analytic_dense_muon_partition": + if partition.get("status") == "exact_analytic_moe_muon_partition": + pass + else: + unsupported_terms.append(f"muon_fallback_param_partition:{fallback}") + notes.append( + "Muon fallback AdamW/SGD partition is parameter-name based at runtime; " + "metadata ledger uses aggregate local params" + ) + else: + status = "unsupported" + unsupported_terms.append(f"optimizer:{optimizer}") + + optimizer_state_gb = round(sum(term["gb"] for term in state_terms.values()), 3) + optimizer_peak_transient_gb = round(sum(term["gb"] for term in transient_terms.values()), 3) + persistent_scan_gb = round(param_gb + gradient_gb + optimizer_state_gb, 3) + # Conservative memory-traffic lower bound: read params/grads/states and write params/states. + step_read_write_gb = round(param_gb * 2 + gradient_gb + optimizer_state_gb * 2, 3) + if optimizer == "muon": + step_work_basis = ( + "param+grad+optimizer persistent state scan lower bound; " + "Muon Newton-Schulz compute/time remains unsupported" + ) + elif optimizer in {"signsgd", "distsignsgd", "sgd"} and optimizer_state_gb == 0.0: + step_work_basis = "state-free param+grad scan lower bound" + else: + step_work_basis = "param+grad+optimizer persistent state scan lower bound" + + return { + "status": status, + "optimizer": optimizer, + "optimizer_family": optimizer_family, + "local_params": round(local_params), + "local_non_expert_params": round(local_non_expert_params), + "local_expert_params": round(local_expert_params), + "param_gb": round(param_gb, 3), + "gradient_gb": round(gradient_gb, 3), + "optimizer_state_gb": optimizer_state_gb, + "persistent_scan_gb": persistent_scan_gb, + "step_read_write_lower_bound_gb": step_read_write_gb, + "phase_quantity": { + "name": "optimizer_step_work", + "value": persistent_scan_gb, + "unit": "gb", + "basis": step_work_basis, + }, + "state_terms": state_terms, + "optimizer_peak_transient_gb": optimizer_peak_transient_gb, + "transient_terms": transient_terms, + "unsupported_terms": unsupported_terms, + "notes": notes, + } + + +def muon_optimizer_peak_memory_attribution( + *, + analytic_floor_gb: float | None, + measured_peak_gb: float | None, + activation_lower_bound_gb: float | None, + optimizer_step: dict[str, Any], + validation_threshold: float = MEMORY_PEAK_VALIDATION_THRESHOLD, +) -> dict[str, Any]: + if optimizer_step.get("optimizer") != "muon": + return {"status": "not_applicable", "reason": "optimizer_is_not_muon"} + if analytic_floor_gb is None or measured_peak_gb is None: + return {"status": "incomplete"} + + transient_gb = optimizer_step.get("optimizer_peak_transient_gb") + if transient_gb is None or transient_gb <= 0: + return {"status": "not_applicable", "reason": "muon_optimizer_transient_unavailable"} + + act = activation_lower_bound_gb or 0.0 + activation_peak_gb = analytic_floor_gb + act + optimizer_peak_gb = analytic_floor_gb + float(transient_gb) + predicted_peak_gb = max(activation_peak_gb, optimizer_peak_gb) + relative_error = abs(predicted_peak_gb - measured_peak_gb) / measured_peak_gb if measured_peak_gb else None + matches = relative_error is not None and relative_error <= validation_threshold + transient_terms = optimizer_step.get("transient_terms", {}) + transient_statuses = { + term.get("status") for term in transient_terms.values() if isinstance(term, dict) and term.get("status") + } + transient_status = ( + "exact_analytic" + if transient_statuses == {"exact_analytic"} + else "exact_analytic_lower_bound" + if transient_statuses + else "unsupported" + ) + return { + "status": ( + "muon_optimizer_transient_peak_formula_matches_step_peak" + if matches + else "muon_optimizer_transient_peak_formula_compared" + ), + "model": ( + "peak = exact param/grad/optimizer floor + max(exact activation lower bound, " + "exact Muon grouped-update transient)" + ), + "validation_threshold": validation_threshold, + "measured_peak_gb": round(measured_peak_gb, 3), + "predicted_peak_gb": round(predicted_peak_gb, 3), + "relative_error": round(relative_error, 6) if relative_error is not None else None, + "validation_residual_gb": round(measured_peak_gb - predicted_peak_gb, 3), + "residual_status": "validation_error_not_fit_coefficient", + "param_grad_opt_floor_gb": round(analytic_floor_gb, 3), + "param_grad_opt_floor_status": "exact_analytic", + "analytic_activation_lower_bound_gb": round(act, 3), + "analytic_activation_status": "exact_analytic_lower_bound", + "activation_peak_gb": round(activation_peak_gb, 3), + "muon_optimizer_transient_gb": round(float(transient_gb), 3), + "muon_optimizer_transient_status": transient_status, + "optimizer_peak_gb": round(optimizer_peak_gb, 3), + "peak_driver": "optimizer_step" if optimizer_peak_gb >= activation_peak_gb else "activation", + "optimizer_transient_terms": transient_terms, + "optimizer_step_status": optimizer_step.get("status"), + "optimizer_step_unsupported_terms": optimizer_step.get("unsupported_terms", []), + } + + +def fsdp_unshard_transient_ledger( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], +) -> dict[str, Any]: + """Per-rank FSDP full-param transient memory for the largest dense transformer layer. + + The persistent memory floor counts only sharded parameter storage. During FSDP2 forward/backward, + a fully-sharded layer also materializes an unsharded full-parameter buffer. With forward prefetch, + the steady-state peak can hold the current layer and one prefetched layer. + """ + if topology.data_parallel_shard_size <= 1 and not train.get("enable_full_shard"): + return {"status": "not_applicable", "reason": "not_full_shard"} + needed = ( + metadata.hidden_size, + metadata.intermediate_size, + metadata.num_attention_heads, + metadata.num_hidden_layers, + ) + if any(value is None for value in needed): + return {"status": "unsupported", "reason": "missing_dense_layer_metadata"} + if metadata.moe_intermediate_size is not None or metadata.num_experts is not None: + return {"status": "unsupported", "reason": "moe_layer_unshard_formula_not_enabled_for_dense_q8_gate"} + + hidden = metadata.hidden_size + intermediate = metadata.intermediate_size + n_heads = metadata.num_attention_heads + n_kv = metadata.num_key_value_heads or n_heads + head_dim = metadata.head_dim or (hidden // n_heads) + q_size = n_heads * head_dim + kv_size = n_kv * head_dim + o_size = n_heads * head_dim + attn_params = hidden * (q_size + 2 * kv_size + o_size) + dense_mlp_params = hidden * intermediate * 3 + norm_params = 2 * hidden + 2 * head_dim + largest_layer_params = attn_params + dense_mlp_params + norm_params + param_bytes, param_notes = _param_storage_bytes(train) + prefetch_window_layers = 2 if train.get("enable_forward_prefetch") else 1 + layer_unshard_gb = largest_layer_params * param_bytes / _BYTES_PER_GIB + predicted_gb = prefetch_window_layers * layer_unshard_gb + return { + "status": "exact_analytic_dense_fsdp_prefetch", + "largest_layer_params": int(largest_layer_params), + "param_bytes": param_bytes, + "prefetch_window_layers": prefetch_window_layers, + "one_layer_unshard_gb": round(layer_unshard_gb, 3), + "predicted_prefetch_window_gb": round(predicted_gb, 3), + "term_status": { + "attn_params": "exact_analytic", + "dense_mlp_params": "exact_analytic", + "norm_params": "exact_analytic", + "prefetch_window_layers": "config_exact", + }, + "notes": [ + *param_notes, + f"enable_forward_prefetch={bool(train.get('enable_forward_prefetch'))}", + "persistent floor counts sharded params; FSDP unshard materializes full layer params transiently", + ], + } + + +def communication_ledger( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], + *, + seq_len: int | None = None, + param_bytes: int = 2, + grad_bytes: int = 4, +) -> dict[str, Any]: + """Per-optimizer-step analytical communication BYTES per rank, by collective, with provenance. + + Bytes are exact_analytic from shapes/topology. Communication TIME is NOT predicted here (it needs + a calibrated per-link bandwidth + overlap coefficient); for the q35/q30 1-node runs the measured + comm phases (e.g. sync_sp_gradients) are negligible, consistent with intra-node NVLink + overlap. + """ + needed = (metadata.hidden_size, metadata.num_hidden_layers) + seq = _seq_len(topology, seq_len) + if any(v is None for v in needed) or seq is None: + return {"status": "unsupported", "reason": "missing_metadata_or_seq_len", "terms": {}} + + from_breakdown = _param_split(metadata) + if from_breakdown is None: + return {"status": "unsupported", "reason": "param_breakdown_unavailable", "terms": {}} + non_expert_params, expert_params = from_breakdown + hidden = metadata.hidden_size + layers = metadata.num_hidden_layers + sp = max(topology.sequence_parallel_size, 1) + act = 2 if train.get("enable_mixed_precision") else 4 + node_count = max(topology.node_count, 1) + + # Sharded ownership group sizes (mirror memory_ledger ownership). + dp_shard = max(topology.data_parallel_shard_size, 1) + tp = max(topology.tensor_parallel_size, 1) + non_expert_group = dp_shard * tp * (sp if str(train.get("cp_fsdp_mode", "all") or "all") == "all" else 1) + + # all-gather/reduce-scatter move (G-1)/G of the full (unsharded) tensor per rank. + def ag_factor(group: int) -> float: + return (group - 1) / group if group > 1 else 0.0 + + raw_reshard_after_forward = train.get("reshard_after_forward") + if raw_reshard_after_forward is None: + # Mirror src/xorl/distributed/torch_parallelize.py: + # None means "auto"; PP disables resharding, non-PP uses FSDP's default + # reshard_after_forward=True. + reshard_after_forward = topology.pipeline_parallel_size <= 1 + reshard_after_forward_source = "auto_non_pp_true" if reshard_after_forward else "auto_pp_false" + else: + reshard_after_forward = bool(raw_reshard_after_forward) + reshard_after_forward_source = "explicit_config" + ag_passes = 2.0 if reshard_after_forward else 1.0 # re-gather in backward if resharded after fwd + local_ws = max(topology.local_world_size, 1) + + def node_span(group_size: int) -> int: + # how many distinct nodes a collective group spans (contiguous rank->node mapping). + return max(1, min(-(-group_size // local_ws), node_count)) + + def cross_fraction(group_size: int) -> float: + # ring/tree approximation: inter-node share of a collective whose group spans `ns` nodes. + ns = node_span(group_size) + return (ns - 1) / ns if ns > 1 else 0.0 + + terms: dict[str, dict[str, Any]] = {} + + def add( + name: str, + byte_count: float, + status: str, + group_size: int, + note: str, + *, + extra: dict[str, Any] | None = None, + ) -> None: + gb = byte_count / _BYTES_PER_GIB + cf = cross_fraction(group_size) + ns = node_span(group_size) + scope = "intra_node" if ns == 1 else ("cross_node" if cf == 1.0 else "intra_and_cross_node") + term = { + "gb": round(gb, 4), + "intra_gb": round(gb * (1 - cf), 4), + "cross_gb": round(gb * cf, 4), + "group_size": group_size, + "nodes_spanned": ns, + "status": status, + "scope": scope, + "note": note, + } + if extra: + term.update(extra) + terms[name] = term + + def na(name: str, note: str) -> None: + terms[name] = { + "gb": 0.0, + "intra_gb": 0.0, + "cross_gb": 0.0, + "nodes_spanned": 0, + "group_size": 0, + "status": "not_applicable", + "scope": "n/a", + "note": note, + } + + # FSDP all-gather of params (forward, +backward if resharded), NON-EXPERT shard group only. + # Expert params are EP-owned: each rank only ever materializes its own num_experts/ep slice, so the + # expert FSDP collectives move (per-rank slice) x (ep_fsdp-1)/ep_fsdp — NOT global experts — and are + # ZERO at ep_fsdp=1. They get their own terms below because the ep_fsdp mesh has a different + # node mapping (1 rank per node under ep_intranode multi-node -> every ring link is inter-node). + add( + "fsdp_param_all_gather", + ag_passes * non_expert_params * ag_factor(non_expert_group) * param_bytes, + "exact_analytic", + non_expert_group, + f"all-gather sharded non-expert params; passes={ag_passes} (reshard_after_forward={reshard_after_forward})", + extra={ + "passes": ag_passes, + "per_pass_gb": round((non_expert_params * ag_factor(non_expert_group) * param_bytes) / _BYTES_PER_GIB, 4), + "effective_reshard_after_forward": reshard_after_forward, + "raw_reshard_after_forward": raw_reshard_after_forward, + "reshard_after_forward_source": reshard_after_forward_source, + }, + ) + # FSDP reduce-scatter of gradients (once per step), NON-EXPERT shard group only. + add( + "fsdp_grad_reduce_scatter", + non_expert_params * ag_factor(non_expert_group) * grad_bytes, + "exact_analytic", + non_expert_group, + "reduce-scatter non-expert gradients across shard group", + ) + # Expert FSDP collectives over the ep_fsdp mesh (per-rank EP slice, zero when ep_fsdp == 1). + ep_size_for_split = max(topology.expert_parallel_size, 1) + ep_fsdp = max(topology.ep_fsdp_size or 1, 1) + expert_slice_params = expert_params / ep_size_for_split if expert_params else 0.0 + if expert_params and ep_fsdp > 1: + # ep_intranode packs EP groups on consecutive intra-node ranks, so the ep_fsdp dimension + # strides by ep: mesh {i, i+ep, i+2*ep, ...} with local_ws/ep members per node. Ring order + # keeps same-node members adjacent, so exactly ns of the ep_fsdp ring links cross nodes and + # the exact per-rank cross share is ns/ep_fsdp (1.0 at one rank per node — the previously + # handled case; 2026-07-06 stride-aware fix prices epnode convention wrongly scored intra-node, e.g. ep4 x ep_fsdp4 at world 16). + add( + "expert_fsdp_param_all_gather", + ag_passes * expert_slice_params * ag_factor(ep_fsdp) * param_bytes, + "exact_analytic", + ep_fsdp, + f"all-gather per-rank expert slice across ep_fsdp={ep_fsdp}; passes={ag_passes}", + extra={"passes": ag_passes}, + ) + add( + "expert_fsdp_grad_reduce_scatter", + expert_slice_params * ag_factor(ep_fsdp) * grad_bytes, + "exact_analytic", + ep_fsdp, + f"reduce-scatter expert gradients across ep_fsdp={ep_fsdp}", + ) + if bool(train.get("ep_intranode", True)) and node_count > 1: + _members_per_node = max(local_ws // ep_size_for_split, 1) + _strided_ns = min(-(-ep_fsdp // _members_per_node), node_count) + if _strided_ns > 1: + _strided_share = _strided_ns / ep_fsdp + for _expert_term in ("expert_fsdp_param_all_gather", "expert_fsdp_grad_reduce_scatter"): + _t = terms[_expert_term] + _t["cross_gb"] = round(_t["gb"] * _strided_share, 4) + _t["intra_gb"] = round(_t["gb"] * (1.0 - _strided_share), 4) + _t["nodes_spanned"] = _strided_ns + _t["scope"] = "cross_node" if _strided_share == 1.0 else "intra_and_cross_node" + _t["note"] += ( + f"; stride-aware node mapping: ep_fsdp strides by ep{ep_size_for_split} " + f"({_members_per_node} members/node, spans {_strided_ns} nodes), exact ring " + f"share {_strided_ns}/{ep_fsdp} of per-rank bytes crosses nodes" + ) + else: + na( + "expert_fsdp_param_all_gather", + "no expert FSDP sharding (dense model or ep_fsdp=1: experts are per-rank local)", + ) + na( + "expert_fsdp_grad_reduce_scatter", + "no expert FSDP sharding (dense model or ep_fsdp=1: experts are per-rank local)", + ) + # EP all-to-all dispatch + combine, per MoE layer (global routed tokens routed to expert ranks). + routed_slots_global = topology.global_batch_size * seq * (metadata.top_k or 0) + ep = max(topology.expert_parallel_size, 1) + a2a_per_rank = (routed_slots_global / max(topology.world_size, 1)) * hidden * act * 2 * layers * ag_factor(ep) + ep_intranode = bool(train.get("ep_intranode")) + if metadata.num_experts is not None and metadata.top_k is not None and metadata.moe_intermediate_size is not None: + # The EP group spans ceil(ep/local_ws) nodes whenever ep > local_ws regardless of + # ep_intranode (the flag only packs groups contiguously). All-to-all traffic is uniform over + # destinations, so the inter-node share is 1 - (intra-node group ranks)/ep, not the ring + # (ns-1)/ns convention. ep_intranode=False with ep <= local_ws strides EP groups across nodes; + # that layout has no measured row and keeps the contiguous accounting with a note. + add( + "ep_all_to_all_dispatch_combine", + a2a_per_rank, + "exact_analytic", + ep, + f"alltoall dispatch+combine over EP={ep} (ep_intranode={ep_intranode}); per-layer x{layers}", + ) + _a2a = terms["ep_all_to_all_dispatch_combine"] + _a2a_cf = 1.0 - (min(local_ws, ep) / ep) + _a2a["cross_gb"] = round(_a2a["gb"] * _a2a_cf, 4) + _a2a["intra_gb"] = round(_a2a["gb"] * (1.0 - _a2a_cf), 4) + _a2a["scope"] = "intra_node" if _a2a_cf == 0.0 else ("cross_node" if _a2a_cf == 1.0 else "intra_and_cross_node") + _a2a["note"] += "; uniform-destination inter-node share (not ring convention)" + if not ep_intranode and ep <= local_ws: + _a2a["note"] += "; WARNING ep_intranode=False strides EP across nodes (unmeasured layout)" + else: + na("ep_all_to_all_dispatch_combine", "dense model: no expert-parallel all-to-all") + # HSDP DP all-reduce across replicas (only if dp_replicate>1; replicas live on distinct nodes). + # Per-rank bytes: each rank all-reduces its own NON-EXPERT grad shard (P_ne x grad_bytes / + # non_expert_group) with its replica peers; ring all-reduce moves 2 x (r-1)/r of the message per + # rank (reduce-scatter + all-gather phases). Expert grads are EXCLUDED: the ep x ep_fsdp meshes + # span the full PP stage regardless of dp_replicate, so expert tensors have no replicate + # dimension and never all-reduce across replicas. (Corrected 2026-07-05: the earlier term + # charged (non_expert + expert) x full-model bytes x (r-1)/r — a per-global convention ~55x the + # honest per-rank shard bytes at the 65k replicate2 x shard_sp8 layout.) + if topology.data_parallel_replicate_size > 1: + replicate = topology.data_parallel_replicate_size + add( + "dp_grad_all_reduce_hsdp", + non_expert_params * grad_bytes / max(non_expert_group, 1) * 2.0 * ag_factor(replicate), + "exact_analytic", + node_count * local_ws, # replicas span nodes -> force cross-node share + ( + f"HSDP non-expert grad-shard all-reduce across {replicate} replicas " + f"(per-rank shard = non_expert/{non_expert_group}; ring 2x(r-1)/r; experts have no " + "replicate dim: ep x ep_fsdp spans the full stage)" + ), + ) + else: + na("dp_grad_all_reduce_hsdp", "dp_replicate=1") + # Ulysses / Ring sequence-parallel attention collectives. + if sp > 1: + add( + "sequence_parallel_attention_collective", + (topology.micro_batch_size * seq / sp) * hidden * act * 4 * layers, + "analytic_with_runtime_coefficient", + sp, + "Ulysses all-to-all (q,k,v,o) or Ring attention per layer", + ) + else: + na("sequence_parallel_attention_collective", f"sp={sp}") + # Pipeline sends (PP stages typically span nodes when multi-node). + if topology.pipeline_parallel_size > 1: + add( + "pipeline_activation_sends", + (topology.micro_batch_size * seq / sp) * hidden * act * 2 * topology.gradient_accumulation_steps, + "analytic_with_runtime_coefficient", + node_count * local_ws, + "1F1B activation send/recv between PP stages", + ) + else: + na("pipeline_activation_sends", "pp=1") + + total = round(sum(t["gb"] for t in terms.values()), 4) + cross = round(sum(t.get("cross_gb", 0.0) for t in terms.values()), 4) + intra = round(total - cross, 4) + # Time estimate: bytes / (link bandwidth x overlap efficiency). Bandwidths are nominal coefficients; + # they are NOT calibrated here because 1-node comm is fully overlapped/hidden (cannot isolate from + # measured phase timing). Treat as a lower bound on serial comm time, not a step-time predictor. + nvlink_gbps = H100_NVLINK_EFFECTIVE_GB_PER_S # H100 NVLink ~900 GB/s bidir; ~450 effective serial + crossnode_gbps = H100_NDR400_UNIDIRECTIONAL_GB_PER_S # nominal per-GPU NDR400 direction + serial_comm_time_s = intra * _BYTES_PER_GIB / (nvlink_gbps * 1e9) + cross * _BYTES_PER_GIB / (crossnode_gbps * 1e9) + # CALIBRATED exposed cross-node comm: from the 1-node->2-node step-time delta at identical per-rank + # work, cross-model-validated on q30 (6.38 ms/GB) and q35 (5.82 ms/GB) -> ~6.1 ms per cross-node GB + # is the EXPOSED (non-overlapped) step-time cost. Intra-node (NVLink) comm is treated as fully + # overlapped (~0 exposed), consistent with negligible 1-node comm phases. + exposed_cross_node_step_gb = 0.0 + for term_name, term in terms.items(): + term_cross_gb = float(term.get("cross_gb") or 0.0) + if term_name == "fsdp_param_all_gather": + # The byte ledger records every logical all-gather pass. The calibrated + # step-time coefficient was fit to the step-visible param traffic: one + # pass, plus grad reduce-scatter, when non-PP FSDP auto-reshards. + term_cross_gb /= max(float(term.get("passes") or 1.0), 1.0) + exposed_cross_node_step_gb += term_cross_gb + exposed_cross_node_step_gb = round(exposed_cross_node_step_gb, 4) + exposed_cross_node_step_time_s = exposed_cross_node_step_gb * (EXPOSED_CROSS_NODE_MS_PER_GB / 1000.0) + static_overlap_estimate = static_cross_node_overlap_estimate( + terms, + local_world_size=local_ws, + num_experts=metadata.num_experts, + top_k=metadata.top_k, + expert_parallel_size=topology.expert_parallel_size, + ) + return { + "status": "exact_analytic_bytes", + "param_bytes": param_bytes, + "grad_bytes": grad_bytes, + "node_count": node_count, + "terms": terms, + "total_per_rank_gb": total, + "intra_node_per_rank_gb": intra, + "cross_node_per_rank_gb": cross, + "time_estimate": { + "status": "calibrated_cross_model" if cross > 0 else "intra_node_fully_overlapped", + "exposed_cross_node_ms_per_gb": EXPOSED_CROSS_NODE_MS_PER_GB, + "exact_cross_node_per_rank_gb": cross, + "exposed_cross_node_step_gb": exposed_cross_node_step_gb, + "exposed_cross_node_step_time_s": round(exposed_cross_node_step_time_s, 4), + "calibration_source": "1node->2node step delta, cross-model-validated q30=6.38 q35=5.82 ms/GB", + "serial_comm_time_lower_bound_s": round(serial_comm_time_s, 4), + "nvlink_effective_gbps": nvlink_gbps, + "crossnode_effective_gbps": crossnode_gbps, + "static_hardware_overlap_estimate": static_overlap_estimate, + "note": ( + "exposed_cross_node_step_time_s = exposed_cross_node_step_gb x calibrated coefficient; " + "cross_node_per_rank_gb remains the exact byte ledger. intra-node comm is overlapped (~0 exposed). " + "serial_comm_time_lower_bound_s uses nominal bandwidths. static_hardware_overlap_estimate is a " + "separate non-calibrated comparison." + ), + }, + "time_status": "calibrated_cross_model" if cross > 0 else "intra_node_overlapped", + "note": "1-node: comm overlapped (cross=0). 2-node: cross-node exposed comm ~26-30% of step (calibrated).", + } + + +def _param_split(metadata: ModelMetadata) -> tuple[float, float] | None: + """(non_expert_params, expert_params) reusing the memory-ledger param breakdown.""" + breakdown = _estimate_param_breakdown(metadata) + if breakdown is None: + return None + return breakdown["non_expert_params"], breakdown["expert_params"] + + +def build_model_analytical_coverage( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], + *, + seq_len: int | None, + analytic_floor_gb: float | None, + measured_peak_gb: float | None, + measured_step_time_s: float | None, + measured_tflops_per_gpu: float | None = None, + measured_mfu: float | None = None, + flops_batch_seqlens: list[int] | None = None, +) -> dict[str, Any]: + """Full analytical coverage (FLOPs + memory + comm) with predicted-vs-measured for one run.""" + flops = flops_ledger(metadata, topology, seq_len=seq_len, batch_seqlens=flops_batch_seqlens) + hardware = hardware_flops_ledger(flops, train) + act = activation_ledger(metadata, topology, train, seq_len=seq_len) + comm = communication_ledger(metadata, topology, train, seq_len=seq_len) + optimizer_step = optimizer_step_ledger(metadata, topology, train) + efficiency = flops_consistency_and_efficiency( + flops, + measured_step_time_s=measured_step_time_s, + measured_tflops_per_gpu=measured_tflops_per_gpu, + measured_mfu=measured_mfu, + ) + if ( + efficiency.get("status") == "compared" + and hardware.get("status") + in { + "analytic_with_runtime_coefficient", + "exact_analytic", + } + and measured_step_time_s + ): + hw_mfu = hardware["hardware_per_gpu_flops"] / measured_step_time_s / 1e12 / H100_BF16_PEAK_TFLOPS + efficiency["hardware_achieved_mfu"] = round(hw_mfu, 5) + efficiency["hardware_vs_logical_note"] = ( + "logical MFU excludes recompute; hardware MFU counts the recomputed forward (true utilization)" + ) + attribution = memory_residual_attribution( + analytic_floor_gb=analytic_floor_gb, + measured_peak_gb=measured_peak_gb, + activation_lower_bound_gb=act.get("analytic_activation_lower_bound_gb"), + ) + muon_peak_attribution = muon_optimizer_peak_memory_attribution( + analytic_floor_gb=analytic_floor_gb, + measured_peak_gb=measured_peak_gb, + activation_lower_bound_gb=act.get("analytic_activation_lower_bound_gb"), + optimizer_step=optimizer_step, + ) + calibrated_peak_model: dict[str, Any] = {"status": "calibrated_residual_coefficient"} + is_moe = metadata.num_experts is not None and metadata.moe_intermediate_size is not None + if is_moe and muon_peak_attribution.get("status") in { + "muon_optimizer_transient_peak_formula_matches_step_peak", + "muon_optimizer_transient_peak_formula_compared", + }: + attribution = muon_peak_attribution + calibrated_peak_model = { + "status": attribution["status"], + "form": attribution["model"], + "predicted_peak_gb": attribution["predicted_peak_gb"], + "relative_error": attribution["relative_error"], + "validation_threshold": attribution["validation_threshold"], + "note": ( + "No residual fraction is fit here; the dominant q30/q35 peak is the Muon " + "grouped-update transient visible in src/xorl/optim/muon.py." + ), + } + elif attribution.get("status") == "attributed" and analytic_floor_gb: + rf = attribution.get("residual_fraction_of_peak") + calibrated_peak_model = { + "status": "calibrated_residual_coefficient", + "form": "predicted_peak = analytic_floor / (1 - residual_fraction)", + "residual_fraction_coefficient": rf, + "predicted_peak_gb": round(analytic_floor_gb / (1 - rf), 3) if rf is not None and rf < 1 else None, + "cross_model_validation": "q30-calibrated coefficient predicts q35 peak to ~1.1% (see q30_fit_boundary_predictions)", + } + return { + "calibrated_peak_model": calibrated_peak_model, + "flops_ledger": flops, + "hardware_flops_ledger": hardware, + "flops_efficiency_vs_measured": efficiency, + "activation_ledger": act, + "optimizer_step_ledger": optimizer_step, + "memory_residual_attribution": attribution, + "communication_ledger": comm, + } + + +def reference_counter_total_flops( + metadata: ModelMetadata, + topology: Topology, + *, + seq_len: int | None = None, + batch_seqlens: list[int] | None = None, +) -> float | None: + """Total FLOPs from the *actual* trainer ``XorlFlopsCounter`` (transcription ground truth). + + Returns None if xorl is not importable. Used by tests to assert the analytical ledger reproduces + the trainer's convention EXACTLY, which is the only honest validation of FLOPs (they are a logged + convention, not a hardware measurement). + """ + seq = _seq_len(topology, seq_len) + if seq is None and not batch_seqlens: + return None + try: + from xorl.utils.count_flops import XorlFlopsCounter # noqa: PLC0415 (lazy: xorl optional/heavy) + except Exception: # pragma: no cover - xorl not importable in this context + return None + cfg = SimpleNamespace( + model_type="qwen3_moe" if metadata.moe_intermediate_size is not None else "qwen3", + hidden_size=metadata.hidden_size, + vocab_size=metadata.vocab_size, + intermediate_size=metadata.intermediate_size, + moe_intermediate_size=metadata.moe_intermediate_size, + num_hidden_layers=metadata.num_hidden_layers, + num_key_value_heads=metadata.num_key_value_heads, + num_attention_heads=metadata.num_attention_heads, + num_experts=metadata.num_experts, + num_experts_per_tok=metadata.top_k, + head_dim=metadata.head_dim, + ) + counter = XorlFlopsCounter(cfg, gradient_checkpointing_enabled=False) + if batch_seqlens is None: + batch_seqlens = [seq] * topology.global_batch_size + tokens_sum = sum(batch_seqlens) + if cfg.model_type == "qwen3_moe": + tflops = counter._estimate_qwen3_moe_flops(tokens_sum, batch_seqlens, delta_time=1.0) + else: + tflops = counter._estimate_qwen2_flops(tokens_sum, batch_seqlens, delta_time=1.0) + return float(tflops) * 1e12 + + +def hardware_flops_ledger(ledger: dict[str, Any], train: dict[str, Any]) -> dict[str, Any]: + """Recompute-aware HARDWARE FLOPs (what actually runs on the GPU and determines step time). + + The logged ``flops_ledger`` is the *logical* convention (multiplier 6 = 2 MAC x [1 fwd + 2 bwd]; + recompute excluded). Under activation checkpointing the forward of the checkpointed scope is run a + SECOND time in backward, so the recomputed components do 4 passes (fwd + recompute-fwd + 2 bwd) => + multiplier 8 (linear) / 16 (attn score). lm_head/CE sit outside the recomputed transformer layers. + This explains why the *logical* MFU understates true hardware utilization. + """ + if ledger.get("status") != "exact_analytic": + return {"status": "unsupported"} + method = str(train.get("gradient_checkpointing_method", "") or "") + enabled = bool(train.get("enable_gradient_checkpointing")) + comp = ledger["components_flops"] + # recomputed scope: the transformer-layer components (not lm_head). + recomputed_keys = { + "dense_mlp", + "moe_router", + "moe_gate_up_proj", + "moe_down_proj", + "attn_qkvo_proj", + "attn_score_quadratic", + } + if enabled and method in {"recompute_full_layer", "full"}: + recompute_factor = 8.0 / 6.0 # extra forward over the base 6 multiplier + status = "exact_analytic" + note = "recompute_full_layer: checkpointed forward re-run once in backward (4 passes vs 3)" + elif enabled and method == "recompute_before_dispatch": + # only the pre-dispatch portion (attention + router + gate_up) is recomputed; down is not. + recompute_factor = 8.0 / 6.0 + recomputed_keys = {"moe_router", "moe_gate_up_proj", "attn_qkvo_proj", "attn_score_quadratic"} + status = "exact_analytic" + note = "recompute_before_dispatch: pre-dispatch scope recomputed" + else: + recompute_factor = 1.0 + status = "exact_analytic" + note = "no activation recompute" + actual_recomputed_keys = {name for name in recomputed_keys if name in comp} + hardware_components = { + name: (value * recompute_factor if name in actual_recomputed_keys else value) for name, value in comp.items() + } + hardware_total = float(sum(hardware_components.values())) + logical_total = ledger["total_flops"] + world = 1 if logical_total == 0 else (logical_total / ledger["per_gpu_flops"]) + return { + "status": status, + "note": note, + "recompute_factor_on_scope": recompute_factor, + "recomputed_components": sorted(actual_recomputed_keys) if recompute_factor != 1.0 else [], + "logical_total_flops": logical_total, + "hardware_total_flops": hardware_total, + "hardware_per_gpu_flops": hardware_total / world, + "recompute_overhead_fraction": round((hardware_total - logical_total) / logical_total, 4) + if logical_total + else None, + } + + +def flops_consistency_and_efficiency( + ledger: dict[str, Any], + *, + measured_step_time_s: float | None, + measured_tflops_per_gpu: float | None = None, + measured_mfu: float | None = None, +) -> dict[str, Any]: + """Recover the *real* achieved-FLOPS rate (and MFU) from analytical FLOPs + measured step time. + + NOTE: logged ``tflops``/``mfu`` are computed by the trainer from the SAME FLOPs formula + (``formula_FLOPs / step_time / world``), so matching them is a consistency check on the formula + inputs (step_time, token count), NOT an independent validation of FLOPs. The genuinely measured + quantity is ``step_time``; ``achieved_flops_rate`` (= per-GPU FLOPs / step_time) and ``mfu`` are + real hardware-efficiency numbers given the FLOPs convention. + """ + if ledger.get("status") != "exact_analytic" or not measured_step_time_s: + return {"status": "uncomparable", "reason": "missing_analytic_flops_or_step_time"} + per_gpu = ledger["per_gpu_flops"] + achieved_rate = per_gpu / measured_step_time_s # FLOP/s/GPU (real: step_time is measured) + achieved_tflops = achieved_rate / 1e12 + mfu = achieved_tflops / H100_BF16_PEAK_TFLOPS + out: dict[str, Any] = { + "status": "compared", + "interpretation": "logged_tflops_is_formula_derived_not_independent_measurement", + "measured_step_time_s": measured_step_time_s, + "achieved_flops_rate_per_gpu": achieved_rate, + "achieved_mfu": round(mfu, 5), # the calibrated efficiency coefficient (real) + "logged_tflops_per_gpu": measured_tflops_per_gpu, + "logged_mfu": measured_mfu, + } + # Formula-consistency: our FLOPs/step_time should reproduce the logged tflops (tautological up to + # our token-count vs the trainer's actual padded batch_seqlens). A large gap flags a token-count bug. + if measured_tflops_per_gpu: + out["tflops_formula_consistency_rel_error"] = round( + abs(achieved_tflops - measured_tflops_per_gpu) / measured_tflops_per_gpu, 4 + ) + return out + + +def step_time_leave_one_out(rows: list[dict[str, Any]]) -> dict[str, Any]: + """Leave-one-out step-time prediction across runs of one model (non-circular validation). + + Each row: {label, per_gpu_hardware_flops, measured_step_time_s}. For each run, calibrate the + achieved hardware-FLOPS rate from the OTHER runs and predict this run's step time; the residual is + the genuine per-config efficiency variation (mbs/ep/ga). Quantifies how predictive a single + calibrated rate is across the workload/parallelism axes. + """ + usable = [r for r in rows if r.get("per_gpu_hardware_flops") and r.get("measured_step_time_s")] + if len(usable) < 2: + return {"status": "insufficient_rows", "row_count": len(usable)} + predictions = [] + abs_errors = [] + for i, row in enumerate(usable): + others = [usable[j] for j in range(len(usable)) if j != i] + rates = [o["per_gpu_hardware_flops"] / o["measured_step_time_s"] for o in others] + cal_rate = sum(rates) / len(rates) + predicted = row["per_gpu_hardware_flops"] / cal_rate + measured = row["measured_step_time_s"] + rel = abs(predicted - measured) / measured + abs_errors.append(rel) + predictions.append( + { + "label": row.get("label"), + "measured_step_time_s": round(measured, 4), + "predicted_step_time_s": round(predicted, 4), + "rel_error": round(rel, 4), + } + ) + return { + "status": "validated", + "row_count": len(usable), + "predictions": predictions, + "mean_abs_rel_error": round(sum(abs_errors) / len(abs_errors), 4), + "max_abs_rel_error": round(max(abs_errors), 4), + "interpretation": "residual = per-config efficiency variation a single calibrated rate cannot capture", + } + + +def predict_step_time_from_calibrated_mfu(ledger: dict[str, Any], *, calibrated_mfu: float) -> dict[str, Any]: + """Held-out (non-circular) use of analytical FLOPs: predict step time from a calibrated MFU. + + ``predicted_step_time = per_gpu_flops / (calibrated_mfu * peak_flops_per_gpu)``. Calibrate + ``calibrated_mfu`` on a REFERENCE run and validate the predicted step time against a DIFFERENT + run; the residual is the real cross-run efficiency variation, not a tautology. + """ + if ledger.get("status") != "exact_analytic" or not calibrated_mfu: + return {"status": "uncomparable"} + per_gpu = ledger["per_gpu_flops"] + peak = H100_BF16_PEAK_TFLOPS * 1e12 + predicted = per_gpu / (calibrated_mfu * peak) + return { + "status": "predicted", + "calibrated_mfu": calibrated_mfu, + "predicted_step_time_s": round(predicted, 4), + } diff --git a/src/xorl/sim/benchmark_behavior.py b/src/xorl/sim/benchmark_behavior.py new file mode 100644 index 00000000..8ff0173d --- /dev/null +++ b/src/xorl/sim/benchmark_behavior.py @@ -0,0 +1,2544 @@ +"""Empirical benchmark behavior calibration for checked-in benchmark recipes.""" + +from __future__ import annotations + +import argparse +import glob +import json +import re +import statistics +from dataclasses import fields, replace +from pathlib import Path +from typing import Any + +import yaml + + +try: + from .calibration_packs import resolve_calibration_pack + from .collect_calibration import parse_log_text, summarize_observed_run + from .config_fingerprint import load_training_config, resolve_topology + from .model_metadata import model_ref_from_config + from .runtime_config import runtime_training_config + from .schemas import BenchmarkBehaviorPoint, BenchmarkBehaviorPrediction, ShapeLedger, Topology, to_jsonable +except ImportError: # pragma: no cover - exercised by direct script execution + from calibration_packs import resolve_calibration_pack + from collect_calibration import parse_log_text, summarize_observed_run + from config_fingerprint import load_training_config, resolve_topology + from model_metadata import model_ref_from_config + from runtime_config import runtime_training_config + from schemas import BenchmarkBehaviorPoint, BenchmarkBehaviorPrediction, ShapeLedger, Topology, to_jsonable + + +H100_BF16_PROMISED_TFLOPS_PER_GPU = 989.0 +BENCHMARK_SOURCE_MANIFESTS = ("benchmark_sources.yaml", "benchmark_sources.yml", "benchmark_sources.json") +BENCHMARK_BEHAVIOR_OVERRIDE_MANIFESTS = ( + "benchmark_behavior_overrides.yaml", + "benchmark_behavior_overrides.yml", + "benchmark_behavior_overrides.json", +) + + +def _gpu_count_from_text(text: str) -> int | None: + match = re.search(r"(?P\d+)\s+nodes?\s+x\s+(?P\d+)\s+H100", text, re.IGNORECASE) + if match: + return int(match.group("nodes")) * int(match.group("gpus")) + match = re.search(r"(?P\d+)\s*[x×]\s*H100", text, re.IGNORECASE) + if match: + return int(match.group("gpus")) + match = re.search(r"(?P\d+)\s+GPUs?", text, re.IGNORECASE) + return int(match.group("gpus")) if match else None + + +def human_number(value: str) -> float: + cleaned = value.strip().replace(",", "").lstrip("~") + multiplier = 1.0 + if cleaned.endswith(("K", "k")): + cleaned = cleaned[:-1] + multiplier = 1_000.0 + elif cleaned.endswith(("M", "m")): + cleaned = cleaned[:-1] + multiplier = 1_000_000.0 + return float(cleaned) * multiplier + + +def _first_non_none(*values: Any) -> Any: + for value in values: + if value is not None: + return value + return None + + +def _float_dict(value: Any) -> dict[str, float]: + if not isinstance(value, dict): + return {} + result: dict[str, float] = {} + for key, item in value.items(): + try: + result[str(key)] = float(item) + except (TypeError, ValueError): + continue + return result + + +def _model_ref_from_text_or_path(text: str, source: str | Path) -> str | None: + lowered = f"{text}\n{source}".lower() + if "qwen3.6-35b-a3b-fp8" in lowered: + return "Qwen/Qwen3.6-35B-A3B-FP8" + if "qwen3.6-35b-a3b" in lowered or "qwen36" in lowered: + return "Qwen/Qwen3.6-35B-A3B" + if "qwen3.5-397b-a17b" in lowered or "q397b" in lowered: + return "Qwen/Qwen3.5-397B-A17B" + if "qwen3-235b-a22b-instruct-2507" in lowered: + return "Qwen/Qwen3-235B-A22B-Instruct-2507" + if "qwen3-235b-a22b" in lowered or "q235" in lowered: + return "Qwen/Qwen3-235B-A22B" + if "qwen3.5-35b-a3b" in lowered or "q35_2node_005_050" in lowered: + return "Qwen/Qwen3.5-35B-A3B" + if "qwen3-coder-30b-a3b" in lowered: + return "Qwen/Qwen3-Coder-30B-A3B" + # Non-Coder Qwen3-30B-A3B. Checked after the Coder variant so a coder label/path never + # falls through here; "qwen3-coder-30b-a3b" does not contain the substring "qwen3-30b-a3b". + if "qwen3-30b-a3b" in lowered: + return "Qwen/Qwen3-30B-A3B" + if "qwen3-8b" in lowered: + return "Qwen/Qwen3-8B" + return None + + +def _model_family_key(model_ref: str | None) -> str | None: + if not model_ref: + return None + lowered = model_ref.strip().lower() + if "qwen3.6-35b-a3b-fp8" in lowered: + return "qwen3.6-35b-a3b-fp8" + if "qwen3.6-35b-a3b" in lowered: + return "qwen3.6-35b-a3b" + if "qwen3.5-397b-a17b" in lowered: + return "qwen3.5-397b-a17b" + if "qwen3.5-35b-a3b" in lowered: + return "qwen3.5-35b-a3b" + if "qwen3-235b-a22b" in lowered: + return "qwen3-235b-a22b" + return lowered.rstrip("/") + + +def behavior_point_model_mismatches(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> bool: + point_key = _model_family_key(point.model_ref) + config_key = _model_family_key(model_ref_from_config(raw_config)) + return point_key is not None and config_key is not None and point_key != config_key + + +def _readme_point(readme_text: str, *, source: str, model_ref: str | None) -> BenchmarkBehaviorPoint | None: + tps_match = re.search(r"\|\s*tokens/sec\s*\|\s*(?P~?[0-9.]+[KkMm]?)\s*\|", readme_text) + step_match = re.search(r"\|\s*step time\s*\|\s*(?P~?[0-9.]+)s\s*\|", readme_text) + mfu_match = re.search(r"\|\s*MFU\s*\|\s*(?P~?[0-9.]+)%", readme_text) + memory_match = re.search(r"\|\s*allocated memory\s*\|\s*(?P~?[0-9.]+)GB\s*\|", readme_text) + retries_match = re.search(r"\|\s*allocator retries\s*\|\s*(?P\d+)\s*\|", readme_text) + mbs_match = re.search(r"micro_batch_size:\s*(?P\d+)", readme_text) + global_batch_match = re.search(r"global_batch_size:\s*(?P\d+)", readme_text) + gradient_accumulation_match = re.search(r"gradient_accumulation_steps:\s*(?P\d+)", readme_text) + if not tps_match: + return None + return BenchmarkBehaviorPoint( + label="readme_reference_mbs8", + source=source, + micro_batch_size=int(mbs_match.group("value")) if mbs_match else None, + global_batch_size=int(global_batch_match.group("value")) if global_batch_match else None, + gradient_accumulation_steps=( + int(gradient_accumulation_match.group("value")) if gradient_accumulation_match else None + ), + tokens_per_sec=human_number(tps_match.group("value")), + step_time_sec=float(step_match.group("value").lstrip("~")) if step_match else None, + phase_time_sec=_float_dict({}), + phase_time_share=_float_dict({}), + mfu_percent=float(mfu_match.group("value").lstrip("~")) if mfu_match else None, + tflops_per_gpu=None, + peak_mem_gb=float(memory_match.group("value").lstrip("~")) if memory_match else None, + allocator_retries=int(retries_match.group("value")) if retries_match else None, + gpu_count=_gpu_count_from_text(readme_text), + model_ref=model_ref, + sample_packing_sequence_len=_seq_len_from_readme(readme_text), + tensor_parallel_size=_readme_parallel_int( + readme_text, + "tensor_parallel_size", + (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), + ), + pipeline_parallel_size=_readme_parallel_int( + readme_text, + "pipeline_parallel_size", + (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), + ), + ulysses_parallel_size=_readme_parallel_int( + readme_text, + "ulysses_parallel_size", + (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), + ), + ringattn_parallel_size=_readme_parallel_int( + readme_text, + "ringattn_parallel_size", + (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), + ), + expert_parallel_size=_readme_parallel_int( + readme_text, + "expert_parallel_size", + (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), + ), + ep_fsdp_size=_readme_parallel_int( + readme_text, + "ep_fsdp_size", + (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), + ), + deepep_async_combine=_readme_bool_from_text(readme_text, "deepep_async_combine"), + deepep_num_sms=_readme_parallel_int( + readme_text, + "deepep_num_sms", + (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), + ), + deepep_buffer_size_gb=_readme_float_from_text(readme_text, "deepep_buffer_size_gb"), + enable_compile=_readme_bool_from_text(readme_text, "enable_compile"), + gradient_checkpointing_method=_checkpointing_method_from_text(readme_text), + status="reference_speed", + correctness_status="raw_speed_not_promoted_without_matching_k3_pass", + notes=["current-main logical FLOPs accounting", "balanced synthetic routing", "deepep_async_combine true"], + ) + + +def _readme_adjacent_mbs10_point( + readme_text: str, *, source: str, seq_len: int | None, model_ref: str | None +) -> BenchmarkBehaviorPoint | None: + match = re.search(r"`mbs=10`[^~]+~(?P[0-9.]+)K tok/s", readme_text) + if not match: + return None + tokens_per_sec = human_number(match.group("value") + "K") + global_batch_size = 320 + step_time_sec = (global_batch_size * seq_len / tokens_per_sec) if seq_len else None + return BenchmarkBehaviorPoint( + label="readme_adjacent_mbs10_allocator_pressure", + source=source, + micro_batch_size=10, + global_batch_size=global_batch_size, + tokens_per_sec=tokens_per_sec, + step_time_sec=step_time_sec, + phase_time_sec=_float_dict({}), + phase_time_share=_float_dict({}), + mfu_percent=None, + tflops_per_gpu=None, + peak_mem_gb=None, + allocator_retries=None, + gpu_count=_gpu_count_from_text(readme_text), + model_ref=model_ref, + sample_packing_sequence_len=seq_len, + tensor_parallel_size=_readme_parallel_int( + readme_text, + "tensor_parallel_size", + (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), + ), + pipeline_parallel_size=_readme_parallel_int( + readme_text, + "pipeline_parallel_size", + (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), + ), + ulysses_parallel_size=_readme_parallel_int( + readme_text, + "ulysses_parallel_size", + (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), + ), + ringattn_parallel_size=_readme_parallel_int( + readme_text, + "ringattn_parallel_size", + (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), + ), + expert_parallel_size=_readme_parallel_int( + readme_text, + "expert_parallel_size", + (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), + ), + ep_fsdp_size=_readme_parallel_int( + readme_text, + "ep_fsdp_size", + (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), + ), + deepep_async_combine=_readme_bool_from_text(readme_text, "deepep_async_combine"), + deepep_num_sms=_readme_parallel_int( + readme_text, + "deepep_num_sms", + (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), + ), + deepep_buffer_size_gb=_readme_float_from_text(readme_text, "deepep_buffer_size_gb"), + enable_compile=_readme_bool_from_text(readme_text, "enable_compile"), + gradient_checkpointing_method=_checkpointing_method_from_text(readme_text), + status="allocator_pressure_slowdown", + correctness_status="not_promoted", + notes=["fit but slowed with allocator retries"], + ) + + +def _result_throughput_point( + result_path: Path, + result: dict[str, Any], + *, + topology_defaults: dict[str, int | float | bool | str], +) -> BenchmarkBehaviorPoint: + throughput = result["throughput"] + candidate = ( + throughput.get("candidate") + or result.get("candidate") + or (result.get("replay_candidate", {}) if isinstance(result.get("replay_candidate"), dict) else {}).get( + "candidate" + ) + or "throughput" + ) + return BenchmarkBehaviorPoint( + label=f"{result_path.stem}:{candidate}", + source=str(result_path), + micro_batch_size=throughput.get("micro_batch_size"), + global_batch_size=throughput.get("global_batch_size"), + gradient_accumulation_steps=_first_non_none( + throughput.get("gradient_accumulation_steps"), topology_defaults.get("gradient_accumulation_steps") + ), + tokens_per_sec=throughput.get("tokens_per_sec"), + step_time_sec=throughput.get("step_time_sec"), + tokens_per_sec_std=_first_non_none( + throughput.get("tokens_per_sec_std"), throughput.get("tokens_per_sec_stdev") + ), + tokens_per_sec_cv=throughput.get("tokens_per_sec_cv"), + step_time_sec_std=_first_non_none(throughput.get("step_time_sec_std"), throughput.get("step_time_s_std")), + step_time_sec_cv=_first_non_none(throughput.get("step_time_sec_cv"), throughput.get("step_time_s_cv")), + phase_time_sec=_float_dict(throughput.get("phase_time_sec")), + phase_time_share=_float_dict(throughput.get("phase_time_share")), + phase_memory_peak_gb=_float_dict(throughput.get("phase_memory_peak_gb")), + mfu_percent=throughput.get("mfu_percent"), + tflops_per_gpu=throughput.get("mean_tflops_per_gpu"), + peak_mem_gb=throughput.get("gpu_alloc_gb"), + allocator_retries=None, + measured_steps=throughput.get("measured_steps"), + warmup_steps=throughput.get("warmup_steps"), + gpu_count=throughput.get("gpus"), + model_ref=_first_non_none(throughput.get("model_ref"), topology_defaults.get("model_ref")), + sample_packing_sequence_len=_first_non_none( + throughput.get("sample_packing_sequence_len"), topology_defaults.get("sample_packing_sequence_len") + ), + data_parallel_replicate_size=_first_non_none( + throughput.get("data_parallel_replicate_size"), topology_defaults.get("data_parallel_replicate_size") + ), + data_parallel_shard_size=_first_non_none( + throughput.get("data_parallel_shard_size"), topology_defaults.get("data_parallel_shard_size") + ), + tensor_parallel_size=_first_non_none( + throughput.get("tensor_parallel_size"), topology_defaults.get("tensor_parallel_size") + ), + pipeline_parallel_size=_first_non_none( + throughput.get("pipeline_parallel_size"), topology_defaults.get("pipeline_parallel_size") + ), + ulysses_parallel_size=_first_non_none( + throughput.get("ulysses_parallel_size"), topology_defaults.get("ulysses_parallel_size") + ), + ringattn_parallel_size=_first_non_none( + throughput.get("ringattn_parallel_size"), topology_defaults.get("ringattn_parallel_size") + ), + expert_parallel_size=_first_non_none( + throughput.get("expert_parallel_size"), topology_defaults.get("expert_parallel_size") + ), + ep_fsdp_size=_first_non_none( + throughput.get("ep_fsdp"), throughput.get("ep_fsdp_size"), topology_defaults.get("ep_fsdp_size") + ), + deepep_async_combine=_first_non_none( + throughput.get("deepep_async_combine"), topology_defaults.get("deepep_async_combine") + ), + deepep_num_sms=_first_non_none(throughput.get("deepep_num_sms"), topology_defaults.get("deepep_num_sms")), + deepep_buffer_size_gb=_first_non_none( + throughput.get("deepep_buffer_size_gb"), topology_defaults.get("deepep_buffer_size_gb") + ), + enable_compile=_first_non_none(throughput.get("enable_compile"), topology_defaults.get("enable_compile")), + gradient_checkpointing_method=_first_non_none( + throughput.get("gradient_checkpointing_method"), topology_defaults.get("gradient_checkpointing_method") + ), + enable_activation_offload=_first_non_none( + throughput.get("enable_activation_offload"), topology_defaults.get("enable_activation_offload") + ), + activation_offload_prefetch_count=_first_non_none( + throughput.get("activation_offload_prefetch_count"), + topology_defaults.get("activation_offload_prefetch_count"), + ), + fsdp_reduce_dtype=_first_non_none( + throughput.get("fsdp_reduce_dtype"), topology_defaults.get("fsdp_reduce_dtype") + ), + ce_mode=_first_non_none(throughput.get("ce_mode"), topology_defaults.get("ce_mode")), + moe_implementation=_first_non_none( + throughput.get("moe_implementation"), topology_defaults.get("moe_implementation") + ), + moe_checkpoint_method=_first_non_none( + throughput.get("moe_checkpoint_method"), topology_defaults.get("moe_checkpoint_method") + ), + muon_update_dtype=_first_non_none( + throughput.get("muon_update_dtype"), topology_defaults.get("muon_update_dtype") + ), + attention_backend=_first_non_none( + throughput.get("attention_backend"), topology_defaults.get("attention_backend") + ), + balanced_routing=_first_non_none( + throughput.get("balanced_routing"), + throughput.get("synthetic_balanced_routing"), + topology_defaults.get("balanced_routing"), + ), + status="historical_throughput_artifact", + correctness_status=None, + notes=[f"commit={throughput.get('commit')}"] if throughput.get("commit") else [], + ) + + +def _with_k3_status(point: BenchmarkBehaviorPoint, result: dict[str, Any]) -> BenchmarkBehaviorPoint: + k3_gate = result.get("k3_gate", {}) + if not k3_gate or k3_gate.get("candidate") not in (None, point.label.split(":", 1)[-1]): + return point + notes = list(point.notes) + if k3_gate.get("primary_failure"): + notes.append(f"k3_primary_failure={k3_gate['primary_failure']}") + return replace( + point, + correctness_status=f"k3_{k3_gate.get('status')}", + notes=notes, + ) + + +def _seq_len_from_readme(readme_text: str) -> int | None: + match = re.search(r"sample_packing_sequence_len:\s*(?P\d+)", readme_text) + if match: + return int(match.group("seq")) + match = re.search(r"max_seq_len[=:]\s*(?P\d+)", readme_text) + return int(match.group("seq")) if match else None + + +def _config_int_from_text(text: str, key: str) -> int | None: + match = re.search(rf"{re.escape(key)}:\s*(?P\d+)", text) + return int(match.group("value")) if match else None + + +def _readme_float_from_text(text: str, key: str) -> float | None: + match = re.search(rf"{re.escape(key)}:\s*(?P\d+(?:\.\d+)?)", text) + return float(match.group("value")) if match else None + + +def _readme_bool_from_text(text: str, key: str) -> bool | None: + match = re.search(rf"{re.escape(key)}:\s*(?Ptrue|false)", text, re.IGNORECASE) + if not match: + return None + return match.group("value").lower() == "true" + + +def _checkpointing_method_from_text(text: str) -> str | None: + lowered = text.lower() + if "recompute_before_dispatch" in lowered or "before_dispatch" in lowered: + return "recompute_before_dispatch" + if "recompute_full_layer" in lowered or "full-layer recompute" in lowered or "fullrecompute" in lowered: + return "recompute_full_layer" + if "no_recompute" in lowered or "no recompute" in lowered: + return "no_recompute" + return None + + +def _trial_checkpointing_method(trial: str) -> str | None: + return _checkpointing_method_from_text(trial) + + +def _trial_activation_offload(trial: str) -> bool | None: + if "noactivationoffload" in trial: + return False + if "activationoffload" in trial: + return True + return None + + +def _trial_prefetch_count(trial: str) -> int | None: + match = re.search(r"prefetch(?P\d+)", trial) + return int(match.group("value")) if match else None + + +def _trial_compile_enabled(trial: str) -> bool | None: + if "nocompile" in trial: + return False + if "compile" in trial: + return True + return None + + +def _trial_deepep_async_combine(trial: str) -> bool | None: + if "noasync" in trial: + return False + if "async" in trial: + return True + return None + + +def _trial_sms_count(trial: str) -> int | None: + match = re.search(r"sms(?P\d+)", trial) + return int(match.group("value")) if match else None + + +def _trial_buffer_size_gb(trial: str) -> float | None: + match = re.search(r"buf(?P\d+)", trial) + if not match: + return None + raw = match.group("value") + if len(raw) == 1: + return float(raw) + return float(f"{raw[:-1]}.{raw[-1]}") + + +def _trial_ce_mode(trial: str) -> str | None: + lowered = trial.lower() + if "quackce" in lowered or "quack_linear" in lowered: + return "quack_linear" + if "compiledce" in lowered: + return "compiled" + return None + + +def _trial_moe_implementation(trial: str) -> str | None: + lowered = trial.lower() + if "untunedquack" in lowered or "quack" in lowered: + return "quack" + if "tritonmoe" in lowered or "triton_moe" in lowered: + return "triton" + return None + + +def _trial_muon_update_dtype(trial: str) -> str | None: + lowered = trial.lower() + if "bf16update" in lowered or "updatebf16" in lowered: + return "bf16" + if "fp32update" in lowered or "updatefp32" in lowered: + return "fp32" + return None + + +def _last_regex_int(line: str, patterns: tuple[str, ...]) -> int | None: + value = None + for pattern in patterns: + for match in re.finditer(pattern, line, re.IGNORECASE): + groupdict = match.groupdict() + for key in ("value", "tp", "pp", "u", "ring"): + if groupdict.get(key) is not None: + value = int(groupdict[key]) + break + return value + + +def _readme_parallel_int(readme_text: str, config_key: str, patterns: tuple[str, ...]) -> int | None: + if value := _config_int_from_text(readme_text, config_key): + return value + for line in readme_text.splitlines(): + if value := _last_regex_int(line, patterns): + return value + return None + + +def _readme_default_balanced_routing(readme_text: str) -> bool | None: + in_fixed_block = False + for line in readme_text.splitlines(): + lowered = line.strip().lower() + if not lowered or lowered.startswith("#"): + in_fixed_block = False + continue + if re.search(r"\bfixed\s+(path|flags?)\b", lowered): + in_fixed_block = True + elif lowered.startswith(("-", "*")): + in_fixed_block = False + if not in_fixed_block: + continue + if ( + "balanced synthetic routing" in lowered + or "synthetic balanced routing" in lowered + or "xorl_moe_synthetic_routing=balanced" in lowered + ): + return True + if "real imbalanced routing" in lowered or "real routing" in lowered: + return False + return None + + +def _readme_topology_defaults(readme_text: str) -> dict[str, int | float | bool | str]: + defaults: dict[str, int | float | bool | str] = {} + field_patterns = { + "data_parallel_replicate_size": ( + r"\bdata_parallel_replicate_size[=: ]+(?P\d+)\b", + r"\bdp[_-]?rep(?:licate)?[=_ ](?P\d+)\b", + ), + "data_parallel_shard_size": ( + r"\bdata_parallel_shard_size[=: ]+(?P\d+)\b", + r"\bdp[_-]?shard[=_ ](?P\d+)\b", + ), + "tensor_parallel_size": (r"\btp[=_ ](?P\d+)\b", r"\bTP(?P\d+)\b"), + "pipeline_parallel_size": (r"\bpp[=_ ](?P\d+)\b", r"\bPP(?P\d+)\b"), + "ulysses_parallel_size": (r"\bulysses[=_ ](?P\d+)\b", r"\bU(?P\d+)\b"), + "ringattn_parallel_size": (r"\bring[=_ ](?P\d+)\b", r"\bR(?P\d+)\b"), + "expert_parallel_size": (r"\bep[=_ ](?P\d+)\b", r"\bEP(?P\d+)\b"), + "ep_fsdp_size": (r"\bep_fsdp[=_ ](?P\d+)\b", r"\beFSDP(?P\d+)\b"), + } + for field, patterns in field_patterns.items(): + if value := _readme_parallel_int(readme_text, field, patterns): + defaults[field] = value + if value := _readme_parallel_int( + readme_text, + "gradient_accumulation_steps", + (r"\bgradient_accumulation_steps[=: ]+(?P\d+)\b", r"\bga[=_ ](?P\d+)\b"), + ): + defaults["gradient_accumulation_steps"] = value + if (value := _readme_bool_from_text(readme_text, "deepep_async_combine")) is not None: + defaults["deepep_async_combine"] = value + if value := _readme_parallel_int( + readme_text, + "deepep_num_sms", + (r"\bdeepep_num_sms[=: ]+(?P\d+)\b", r"\bSMS(?P\d+)\b"), + ): + defaults["deepep_num_sms"] = value + if (value := _readme_float_from_text(readme_text, "deepep_buffer_size_gb")) is not None: + defaults["deepep_buffer_size_gb"] = value + if (value := _readme_bool_from_text(readme_text, "enable_compile")) is not None: + defaults["enable_compile"] = value + if value := _checkpointing_method_from_text(readme_text): + defaults["gradient_checkpointing_method"] = value + if (value := _readme_default_balanced_routing(readme_text)) is not None: + defaults["balanced_routing"] = value + return defaults + + +def _first_markdown_number(value: str) -> float | None: + match = re.search(r"~?\s*(?P[0-9][0-9,.]*)(?P[KkMm]?)", value.replace("*", "")) + if not match: + return None + return human_number(match.group("value") + match.group("suffix")) + + +def _markdown_value(values: dict[str, str], key_substring: str) -> str: + for key, value in values.items(): + if key_substring in key: + return value + return "" + + +def _markdown_peak_gb(value: str) -> float | None: + if "oom" in value.lower(): + return None + return _first_markdown_number(value) + + +def _q235_fp32_master_64k_points( + readme_text: str, *, source: str, model_ref: str | None +) -> list[BenchmarkBehaviorPoint]: + if "MEMORY-DOWN WITH FP32 MASTER" not in readme_text or "pr83_opt_mom0_bf16grad" not in readme_text: + return [] + + base_kwargs: dict[str, Any] = { + "source": source, + "micro_batch_size": 1, + "global_batch_size": 8, + "gpu_count": 64, + "model_ref": model_ref, + "sample_packing_sequence_len": 64_000, + "data_parallel_replicate_size": 8, + "data_parallel_shard_size": 1, + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "ulysses_parallel_size": 8, + "ringattn_parallel_size": 1, + "expert_parallel_size": 8, + "ep_fsdp_size": 8, + "deepep_async_combine": False, + "deepep_num_sms": 48, + "enable_compile": False, + "gradient_checkpointing_method": "recompute_full_layer", + "skip_param_upcast": False, + "ce_mode": "compiled", + "muon_momentum": 0.0, + "balanced_routing": False, + "status": "historical_q235_64k_fp32_master", + } + return [ + BenchmarkBehaviorPoint( + label="q235_markdown:pr83_mom0_fp32grad_64k", + tokens_per_sec=28_500.0, + step_time_sec=18.4, + mfu_percent=18.8, + tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU * 0.188, + peak_mem_gb=60.8, + allocator_retries=None, + deepep_buffer_size_gb=1.0, + enable_activation_offload=False, + fsdp_reduce_dtype="fp32", + correctness_status="not_promoted", + notes=["8-node 64K fp32-master mom0 control", "not K3-gated"], + **base_kwargs, + ), + BenchmarkBehaviorPoint( + label="q235_markdown:pr83_mom0_activation_offload_64k", + tokens_per_sec=27_600.0, + step_time_sec=19.0, + mfu_percent=18.2, + tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU * 0.182, + peak_mem_gb=55.0, + allocator_retries=None, + deepep_buffer_size_gb=1.0, + enable_activation_offload=True, + fsdp_reduce_dtype="fp32", + correctness_status="not_promoted", + notes=["8-node 64K fp32-master mom0 with activation offload", "not K3-gated"], + **base_kwargs, + ), + BenchmarkBehaviorPoint( + label="q235_markdown:pr83_mom0_bf16grad_64k", + tokens_per_sec=33_300.0, + step_time_sec=15.7, + mfu_percent=22.0, + tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU * 0.22, + peak_mem_gb=62.3, + allocator_retries=None, + deepep_buffer_size_gb=2.0, + enable_activation_offload=False, + fsdp_reduce_dtype="bf16", + correctness_status="requires_k3", + notes=["8-node 64K fp32-master mom0 with bf16 grad-reduce", "numeric change requires K3 before promotion"], + **base_kwargs, + ), + ] + + +def _q235_markdown_points(readme_text: str, *, source: str, model_ref: str | None) -> list[BenchmarkBehaviorPoint]: + if "Qwen3-235B" not in readme_text or "tok/s tot" not in readme_text: + return [] + + points: list[BenchmarkBehaviorPoint] = _q235_fp32_master_64k_points(readme_text, source=source, model_ref=model_ref) + current_header: list[str] | None = None + current_gpu_count: int | None = None + current_ep_size: int | None = None + current_ep_fsdp_size: int | None = None + current_tensor_parallel_size = 1 + current_pipeline_parallel_size = 1 + current_ulysses_parallel_size = 1 + current_ringattn_parallel_size = 1 + for line in readme_text.splitlines(): + if gpu_count := _gpu_count_from_text(line): + current_gpu_count = gpu_count + ep_matches = list(re.finditer(r"\bEP(?P\d+)\b", line)) + if ep_matches: + current_ep_size = int(ep_matches[-1].group("ep")) + efsdp_matches = list(re.finditer(r"(?:ep_fsdp|eFSDP)(?:[= ]|)(?P\d+)", line)) + if efsdp_matches: + current_ep_fsdp_size = int(efsdp_matches[-1].group("efsdp")) + if tp := _last_regex_int( + line, + ( + r"\bTP(?P\d+)\b", + r"\btensor_parallel_size[:= ]+(?P\d+)\b", + r"\btp[=_](?P\d+)\b", + ), + ): + current_tensor_parallel_size = tp + if pp := _last_regex_int( + line, + ( + r"\bPP(?P\d+)\b", + r"\bpipeline_parallel_size[:= ]+(?P\d+)\b", + r"\bpp[=_](?P\d+)\b", + ), + ): + current_pipeline_parallel_size = pp + if ulysses := _last_regex_int( + line, + ( + r"\bU(?P\d+)\b", + r"\bul[y]?sses_parallel_size[:= ]+(?P\d+)\b", + r"\bu[=_]?(?P\d+)\b", + ), + ): + current_ulysses_parallel_size = ulysses + if ringattn := _last_regex_int( + line, + ( + r"\bR(?P\d+)\b", + r"\bringattn_parallel_size[:= ]+(?P\d+)\b", + r"\bring[=_]?(?P\d+)\b", + ), + ): + current_ringattn_parallel_size = ringattn + if not line.startswith("|"): + continue + cells = [cell.strip() for cell in line.strip().strip("|").split("|")] + lowered = [cell.lower() for cell in cells] + if "run" in lowered and "tok/s tot" in lowered: + current_header = lowered + continue + if current_header is None or set(cells) == {"---"} or not cells: + continue + values = dict(zip(current_header, cells, strict=False)) + run = values.get("run", "").replace("*", "").strip("` ") + if not run or run.lower() in {"run", "-----"}: + continue + status_text = _markdown_value(values, "status") + is_failure = "oom" in status_text.lower() or "fail" in status_text.lower() + tokens_per_sec = _first_markdown_number(_markdown_value(values, "tok/s tot")) + tok_step = _first_markdown_number(_markdown_value(values, "tok/step")) + pack = _first_markdown_number(_markdown_value(values, "pack")) + if tok_step is None or pack in (None, 0): + continue + if tokens_per_sec is None and not is_failure: + continue + global_batch_size = int(round(tok_step / pack)) + step_time_sec = _first_markdown_number(_markdown_value(values, "step s")) + mfu_percent = _first_markdown_number(_markdown_value(values, "mfu")) + peak_mem_gb = _markdown_peak_gb(_markdown_value(values, "peak gb")) + skip_param_upcast = "no skipupcast" not in run.lower() + deepep_buffer_size_gb = 4.0 if "buffer 4" in status_text.lower() or "pk8192_fix" in run.lower() else 2.0 + points.append( + BenchmarkBehaviorPoint( + label=f"q235_markdown:{run}", + source=source, + micro_batch_size=int(_first_markdown_number(values.get("mbs", "")) or 1), + global_batch_size=global_batch_size, + tokens_per_sec=tokens_per_sec, + step_time_sec=step_time_sec, + phase_time_sec=_float_dict({}), + phase_time_share=_float_dict({}), + mfu_percent=mfu_percent, + peak_mem_gb=peak_mem_gb, + allocator_retries=None, + gpu_count=current_gpu_count, + model_ref=model_ref, + sample_packing_sequence_len=int(pack), + tensor_parallel_size=current_tensor_parallel_size, + pipeline_parallel_size=current_pipeline_parallel_size, + ulysses_parallel_size=current_ulysses_parallel_size, + ringattn_parallel_size=current_ringattn_parallel_size, + expert_parallel_size=current_ep_size, + ep_fsdp_size=current_ep_fsdp_size, + deepep_async_combine=False, + deepep_num_sms=24, + deepep_buffer_size_gb=deepep_buffer_size_gb, + enable_compile=False, + gradient_checkpointing_method="recompute_before_dispatch", + skip_param_upcast=skip_param_upcast, + fsdp_reduce_dtype="fp32", + ce_mode="quack_linear", + moe_implementation="quack", + muon_momentum=0.95, + balanced_routing=False, + status="historical_q235_markdown_oom" if is_failure else "historical_q235_markdown", + correctness_status="oom" if is_failure else "not_promoted", + notes=[status_text] if status_text else [], + ) + ) + return points + + +def _result_metric_value(readme_text: str, metric: str) -> str | None: + pattern = rf"\|\s*(?:\*\*)?{re.escape(metric)}(?:\*\*)?[^|]*\|\s*(?P[^|]+)\|" + match = re.search(pattern, readme_text, re.IGNORECASE) + return match.group("value").strip() if match else None + + +def _first_int(value: str | None) -> int | None: + if value is None: + return None + match = re.search(r"\d+", value.replace(",", "")) + return int(match.group(0)) if match else None + + +def _q35_headroom_phase_shares(readme_text: str) -> dict[str, float]: + lowered = readme_text.lower() + if "comm/data-movement-bound" not in lowered: + return {} + comm_match = re.search(r"~(?P[0-9.]+)%\s*:\s*moe a2a", lowered) + gemm_match = re.search(r"only\s*~(?P[0-9.]+)%\s*gemm", lowered) + shares: dict[str, float] = {} + if comm_match: + shares["communication_data_movement"] = round(float(comm_match.group("value")) / 100.0, 3) + if gemm_match: + shares["model_forward_gemm"] = round(float(gemm_match.group("value")) / 100.0, 3) + attributed = sum(shares.values()) + if shares and attributed < 1.0: + shares["other_unattributed"] = round(1.0 - attributed, 3) + return shares + + +def _q35_markdown_points(readme_text: str, *, source: str, model_ref: str | None) -> list[BenchmarkBehaviorPoint]: + if not readme_text.lstrip().startswith("# Qwen3.5-35B-A3B"): + return [] + + gpu_count = _gpu_count_from_text(readme_text) or 16 + seq_len = 65_536 + headroom_phase_share = _q35_headroom_phase_shares(readme_text) + baseline_step_time_sec = float( + _first_markdown_number(_result_metric_value(readme_text, "mean step time") or "71.8") or 71.8 + ) + base_kwargs: dict[str, Any] = { + "gpu_count": gpu_count, + "model_ref": model_ref, + "sample_packing_sequence_len": seq_len, + "data_parallel_replicate_size": 2, + "data_parallel_shard_size": 2, + "tensor_parallel_size": 1, + "pipeline_parallel_size": 1, + "ulysses_parallel_size": 4, + "ringattn_parallel_size": 1, + "expert_parallel_size": 8, + "ep_fsdp_size": 2, + "deepep_async_combine": False, + "deepep_num_sms": 24, + "deepep_buffer_size_gb": 2.0, + "enable_compile": True, + "ce_mode": "quack_linear", + "moe_implementation": "quack", + } + points = [ + BenchmarkBehaviorPoint( + label="q35_markdown:n2_u4_ga16_compile", + source=source, + micro_batch_size=1, + global_batch_size=64, + tokens_per_sec=float( + _first_markdown_number(_result_metric_value(readme_text, "mean tokens/s") or "47352") or 47_352.0 + ), + step_time_sec=baseline_step_time_sec, + phase_time_share=headroom_phase_share, + mfu_percent=float(_first_markdown_number(_result_metric_value(readme_text, "mean MFU") or "8.97") or 8.97), + tflops_per_gpu=float( + _first_markdown_number(_result_metric_value(readme_text, "mean TFLOPS/GPU") or "88.7") or 88.7 + ), + peak_mem_gb=float( + _first_markdown_number(_result_metric_value(readme_text, "peak memory / rank") or "51.2") or 51.2 + ), + allocator_retries=None, + measured_steps=_first_int(_result_metric_value(readme_text, "measured steps")) or 11, + warmup_steps=_first_int(_result_metric_value(readme_text, "warmup excluded")) or 3, + gradient_checkpointing_method="recompute_before_dispatch", + fsdp_reduce_dtype="fp32", + balanced_routing=False, + correctness_status="not_promoted", + status="historical_q35_markdown", + notes=[ + "real CoderForge 65k packs", + "clean completed 2-node reproduction", + "markdown headroom analysis: ~60% communication/data movement, ~18% GEMM", + ], + **base_kwargs, + ), + BenchmarkBehaviorPoint( + label="q35_markdown:no_recompute", + source=source, + micro_batch_size=1, + global_batch_size=64, + tokens_per_sec=46_358.0, + step_time_sec=(64 * seq_len / 46_358.0), + phase_time_share=headroom_phase_share, + mfu_percent=None, + tflops_per_gpu=87.0, + peak_mem_gb=51.2, + allocator_retries=None, + gradient_checkpointing_method="no_recompute", + fsdp_reduce_dtype="fp32", + balanced_routing=False, + correctness_status="not_promoted", + status="historical_q35_headroom", + notes=[ + "no_recompute was throughput-neutral versus before_dispatch", + "markdown headroom analysis: ~60% communication/data movement, ~18% GEMM", + ], + **base_kwargs, + ), + BenchmarkBehaviorPoint( + label="q35_markdown:mbs2_ga8_oom", + source=source, + micro_batch_size=2, + global_batch_size=64, + tokens_per_sec=None, + step_time_sec=None, + mfu_percent=None, + tflops_per_gpu=None, + peak_mem_gb=None, + allocator_retries=None, + gradient_checkpointing_method="recompute_before_dispatch", + fsdp_reduce_dtype="fp32", + balanced_routing=False, + correctness_status="oom", + status="historical_q35_markdown_oom", + notes=["ragged real pack spikes past 80GB under mbs2"], + **base_kwargs, + ), + BenchmarkBehaviorPoint( + label="q35_markdown:bf16red_not_loss_safe", + source=source, + micro_batch_size=1, + global_batch_size=64, + tokens_per_sec=48_488.0, + step_time_sec=(64 * seq_len / 48_488.0), + phase_time_share=headroom_phase_share, + mfu_percent=None, + tflops_per_gpu=90.9, + peak_mem_gb=50.5, + allocator_retries=None, + gradient_checkpointing_method="recompute_before_dispatch", + fsdp_reduce_dtype="bf16", + balanced_routing=False, + correctness_status="not_loss_safe", + status="historical_q35_headroom_not_loss_safe", + notes=[ + "fsdp_reduce_dtype=bf16 underflowed grad_norm; throughput curiosity only", + "markdown headroom analysis: ~60% communication/data movement, ~18% GEMM", + ], + **base_kwargs, + ), + BenchmarkBehaviorPoint( + label="q35_markdown:balanced_mbs1", + source=source, + micro_batch_size=1, + global_batch_size=64, + tokens_per_sec=53_414.0, + step_time_sec=(64 * seq_len / 53_414.0), + mfu_percent=10.1, + tflops_per_gpu=100.2, + peak_mem_gb=39.8, + allocator_retries=None, + gradient_checkpointing_method="recompute_before_dispatch", + fsdp_reduce_dtype="fp32", + balanced_routing=True, + correctness_status="synthetic_routing_not_loss_valid", + status="historical_q35_balanced_routing", + notes=["XORL_MOE_SYNTHETIC_ROUTING=balanced; throughput ceiling only"], + **base_kwargs, + ), + BenchmarkBehaviorPoint( + label="q35_markdown:balanced_mbs2", + source=source, + micro_batch_size=2, + global_batch_size=64, + tokens_per_sec=52_416.0, + step_time_sec=(64 * seq_len / 52_416.0), + mfu_percent=10.0, + tflops_per_gpu=99.0, + peak_mem_gb=52.6, + allocator_retries=None, + gradient_checkpointing_method="recompute_before_dispatch", + fsdp_reduce_dtype="fp32", + balanced_routing=True, + correctness_status="synthetic_routing_not_loss_valid", + status="historical_q35_balanced_routing", + notes=["balanced routing makes mbs2 fit; throughput-neutral versus balanced mbs1"], + **base_kwargs, + ), + ] + return points + + +def _best_by_mfu_point( + result_path: Path, + result: dict[str, Any], + row: dict[str, Any], + *, + topology_defaults: dict[str, int | float | bool | str], +) -> BenchmarkBehaviorPoint: + trial = str(row["trial"]) + caveat = row.get("caveat") + k3_gate = row.get("k3_gate") + notes = [] + if caveat: + notes.append(str(caveat)) + if k3_gate: + notes.append(str(k3_gate)) + correctness_status = None + if k3_gate and str(k3_gate).startswith("pass"): + correctness_status = "k3_pass" + elif _first_non_none(row.get("deepep_async_combine"), _trial_deepep_async_combine(trial)): + correctness_status = "raw_speed_not_promoted_without_matching_k3_pass" + return BenchmarkBehaviorPoint( + label=f"best_by_mfu:{trial}", + source=str(result_path), + micro_batch_size=row.get("micro_batch_size"), + global_batch_size=row.get("global_batch_size"), + tokens_per_sec=row.get("tokens_per_sec"), + step_time_sec=row.get("step_time_sec"), + tokens_per_sec_std=_first_non_none(row.get("tokens_per_sec_std"), row.get("tokens_per_sec_stdev")), + tokens_per_sec_cv=row.get("tokens_per_sec_cv"), + step_time_sec_std=_first_non_none(row.get("step_time_sec_std"), row.get("step_time_s_std")), + step_time_sec_cv=_first_non_none(row.get("step_time_sec_cv"), row.get("step_time_s_cv")), + phase_time_sec=_float_dict(row.get("phase_time_sec")), + phase_time_rank_mean_sec=_float_dict(row.get("phase_time_rank_mean_sec")), + phase_time_share=_float_dict(row.get("phase_time_share")), + mfu_percent=row.get("mfu_percent"), + tflops_per_gpu=row.get("mean_tflops_per_gpu"), + peak_mem_gb=None, + allocator_retries=None, + measured_steps=row.get("measured_steps"), + warmup_steps=row.get("warmup_steps"), + gpu_count=row.get("gpus") or _gpu_count_from_text(str(result.get("workload", ""))), + model_ref=_first_non_none(row.get("model_ref"), topology_defaults.get("model_ref")), + sample_packing_sequence_len=_first_non_none( + row.get("sample_packing_sequence_len"), topology_defaults.get("sample_packing_sequence_len") + ), + data_parallel_replicate_size=_first_non_none( + row.get("data_parallel_replicate_size"), topology_defaults.get("data_parallel_replicate_size") + ), + data_parallel_shard_size=_first_non_none( + row.get("data_parallel_shard_size"), topology_defaults.get("data_parallel_shard_size") + ), + tensor_parallel_size=_first_non_none( + row.get("tensor_parallel_size"), topology_defaults.get("tensor_parallel_size") + ), + pipeline_parallel_size=_first_non_none( + row.get("pipeline_parallel_size"), topology_defaults.get("pipeline_parallel_size") + ), + ulysses_parallel_size=_first_non_none( + row.get("ulysses_parallel_size"), topology_defaults.get("ulysses_parallel_size") + ), + ringattn_parallel_size=_first_non_none( + row.get("ringattn_parallel_size"), topology_defaults.get("ringattn_parallel_size") + ), + expert_parallel_size=_first_non_none( + row.get("expert_parallel_size"), topology_defaults.get("expert_parallel_size") + ), + ep_fsdp_size=_first_non_none(row.get("ep_fsdp"), topology_defaults.get("ep_fsdp_size")), + deepep_async_combine=_first_non_none( + row.get("deepep_async_combine"), + _trial_deepep_async_combine(trial), + topology_defaults.get("deepep_async_combine"), + ), + deepep_num_sms=_first_non_none( + row.get("deepep_num_sms"), _trial_sms_count(trial), topology_defaults.get("deepep_num_sms") + ), + deepep_buffer_size_gb=_first_non_none( + row.get("deepep_buffer_size_gb"), + _trial_buffer_size_gb(trial), + topology_defaults.get("deepep_buffer_size_gb"), + ), + enable_compile=_first_non_none( + row.get("enable_compile"), _trial_compile_enabled(trial), topology_defaults.get("enable_compile") + ), + gradient_checkpointing_method=_first_non_none( + row.get("gradient_checkpointing_method"), + _trial_checkpointing_method(trial), + topology_defaults.get("gradient_checkpointing_method"), + ), + enable_activation_offload=_first_non_none( + row.get("enable_activation_offload"), + _trial_activation_offload(trial), + topology_defaults.get("enable_activation_offload"), + ), + activation_offload_prefetch_count=_first_non_none( + row.get("activation_offload_prefetch_count"), + _trial_prefetch_count(trial), + topology_defaults.get("activation_offload_prefetch_count"), + ), + fsdp_reduce_dtype=_first_non_none(row.get("fsdp_reduce_dtype"), topology_defaults.get("fsdp_reduce_dtype")), + ce_mode=_first_non_none(row.get("ce_mode"), _trial_ce_mode(trial), topology_defaults.get("ce_mode")), + moe_implementation=_first_non_none( + row.get("moe_implementation"), + _trial_moe_implementation(trial), + topology_defaults.get("moe_implementation"), + ), + moe_checkpoint_method=_first_non_none( + row.get("moe_checkpoint_method"), topology_defaults.get("moe_checkpoint_method") + ), + muon_update_dtype=_first_non_none( + row.get("muon_update_dtype"), + _trial_muon_update_dtype(trial), + topology_defaults.get("muon_update_dtype"), + ), + attention_backend=_first_non_none(row.get("attention_backend"), topology_defaults.get("attention_backend")), + balanced_routing=_first_non_none( + row.get("balanced_routing"), + row.get("synthetic_balanced_routing"), + topology_defaults.get("balanced_routing"), + ), + status="autotune_result", + correctness_status=correctness_status, + notes=notes, + ) + + +def _load_startup_metrics(run_dir: Path) -> dict[str, Any]: + startup_path = run_dir / "startup_metrics.json" + if not startup_path.is_file(): + return {} + return json.loads(startup_path.read_text(encoding="utf-8")) + + +def _startup_master_log_path(benchmark_path: Path, startup_metrics: dict[str, Any]) -> Path | None: + metrics = startup_metrics.get("metrics", {}) + master_addr = metrics.get("startup/master_addr") + if not isinstance(master_addr, str) or not master_addr: + return None + run_name = master_addr.removesuffix("-master") + return benchmark_path / run_name / "node-0.log" + + +def _resolved_run_log_path( + benchmark_path: Path, + run_dir: Path, + startup_metrics: dict[str, Any], + *, + log_paths: list[Path] | None = None, +) -> Path | None: + paths = _resolved_run_log_paths(benchmark_path, run_dir, startup_metrics, log_paths=log_paths) + return paths[0] if paths else None + + +def _resolved_run_log_paths( + benchmark_path: Path, + run_dir: Path, + startup_metrics: dict[str, Any], + *, + log_paths: list[Path] | None = None, +) -> list[Path]: + explicit_paths = [path for path in log_paths or [] if path.is_file()] + if explicit_paths: + return explicit_paths + candidates = [ + run_dir / "node-0.log", + _startup_master_log_path(benchmark_path, startup_metrics), + ] + for candidate in candidates: + if candidate is not None and candidate.is_file(): + return [candidate] + return [] + + +def _read_benchmark_source_manifest(benchmark_path: Path) -> tuple[Path, dict[str, Any]] | None: + for filename in BENCHMARK_SOURCE_MANIFESTS: + manifest_path = benchmark_path / filename + if not manifest_path.is_file(): + continue + if manifest_path.suffix == ".json": + payload = json.loads(manifest_path.read_text(encoding="utf-8")) + else: + payload = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + return manifest_path, {} + return manifest_path, payload + return None + + +def _read_manifest_from_candidates( + benchmark_path: Path, filenames: tuple[str, ...] +) -> tuple[Path, dict[str, Any]] | None: + for filename in filenames: + manifest_path = benchmark_path / filename + if not manifest_path.is_file(): + continue + if manifest_path.suffix == ".json": + payload = json.loads(manifest_path.read_text(encoding="utf-8")) + else: + payload = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) + if not isinstance(payload, dict): + return manifest_path, {} + return manifest_path, payload + return None + + +def _as_list(value: Any) -> list[Any]: + if value is None: + return [] + if isinstance(value, list): + return value + return [value] + + +def _resolve_manifest_path(base_path: Path, value: Any) -> Path | None: + if not isinstance(value, str) or not value.strip(): + return None + path = Path(value).expanduser() + if path.is_absolute(): + return path.resolve(strict=False) + return (base_path / path).resolve(strict=False) + + +def _expand_manifest_path_patterns(base_path: Path, value: Any) -> list[Path]: + paths: list[Path] = [] + for item in _as_list(value): + path = _resolve_manifest_path(base_path, item) + if path is None: + continue + pattern = str(path) + if any(token in pattern for token in ("*", "?", "[")): + paths.extend(Path(match) for match in sorted(glob.glob(pattern, recursive=True))) + else: + paths.append(path) + return paths + + +def _resolved_run_label(config_path: Path, label_root: Path, label_prefix: str | None) -> str: + try: + relative = config_path.parent.relative_to(label_root) + except ValueError: + relative = config_path.parent.name + if label_prefix: + return f"resolved_run:{label_prefix}/{relative}" + return f"resolved_run:{relative}" + + +def _log_paths_for_manifest_run( + manifest_dir: Path, + root: Path, + config_path: Path, + source: dict[str, Any], +) -> list[Path]: + log_root = _resolve_manifest_path(manifest_dir, source.get("log_root")) + log_base = log_root if log_root is not None else manifest_dir + log_paths_by_run = source.get("log_paths_by_run", source.get("logs_by_run", {})) + if not isinstance(log_paths_by_run, dict): + return [] + try: + relative_key = config_path.parent.relative_to(root).as_posix() + except ValueError: + relative_key = config_path.parent.name + run_keys = tuple(dict.fromkeys((relative_key, config_path.parent.name))) + paths: list[Path] = [] + for run_key in run_keys: + if run_key in log_paths_by_run: + paths.extend(_expand_manifest_path_patterns(log_base, log_paths_by_run[run_key])) + return paths + + +def _manifest_run_keys(root: Path, config_path: Path) -> tuple[str, str]: + try: + relative_key = config_path.parent.relative_to(root).as_posix() + except ValueError: + relative_key = config_path.parent.name + return relative_key, config_path.parent.name + + +def _manifest_run_int( + root: Path, + config_path: Path, + source: dict[str, Any], + *, + field_name: str, + by_run_field_name: str, +) -> int | None: + by_run = source.get(by_run_field_name, {}) + if isinstance(by_run, dict): + for run_key in _manifest_run_keys(root, config_path): + value = by_run.get(run_key) + if value is not None: + return int(value) + value = source.get(field_name) + return int(value) if value is not None else None + + +def _manifest_run_value( + root: Path, + config_path: Path, + source: dict[str, Any], + *, + field_name: str, + by_run_field_name: str, +) -> Any | None: + by_run = source.get(by_run_field_name, {}) + if isinstance(by_run, dict): + for run_key in _manifest_run_keys(root, config_path): + if run_key in by_run: + return by_run[run_key] + return source.get(field_name) + + +def _manifest_run_metrics_only_reason( + root: Path, + config_path: Path, + source: dict[str, Any], +) -> str | None: + value = _manifest_run_value( + root, + config_path, + source, + field_name="metrics_only", + by_run_field_name="metrics_only_by_run", + ) + if value is None: + value = _manifest_run_value( + root, + config_path, + source, + field_name="exclude_throughput", + by_run_field_name="exclude_throughput_by_run", + ) + if value in (None, False): + return None + if value is True: + return "metrics_only" + return str(value) + + +def _log_failure_status(text: str) -> str | None: + lowered = text.lower() + if "outofmemoryerror" in lowered or "cuda out of memory" in lowered: + return "oom" + if ( + "childfailederror" in lowered + or "traceback" in lowered + or "distbackenderror" in lowered + or "watchdog caught collective operation timeout" in lowered + or "indentationerror" in lowered + ): + return "runtime_failure_after_steps" + return None + + +def _oom_peak_mem_gb(text: str) -> float | None: + values = [ + float(match.group("value")) + for match in re.finditer(r"process has (?P\d+(?:\.\d+)?)\s+GiB memory in use", text) + ] + return max(values) if values else None + + +def _round_or_none(value: Any, ndigits: int) -> float | None: + return round(float(value), ndigits) if value is not None else None + + +RESOLVED_CONFIG_RE = re.compile(r"(?m)^resolved_config=(?P\S+)\s*$") + + +def _same_resolved_config_path(raw_path: str, config_path: Path) -> bool: + marker_path = Path(raw_path).expanduser() + if marker_path == config_path: + return True + try: + if marker_path.resolve() == config_path.resolve(): + return True + except OSError: + pass + if marker_path.is_file() and config_path.is_file(): + try: + if runtime_training_config(load_training_config(marker_path)) == runtime_training_config( + load_training_config(config_path) + ): + return True + except (OSError, ValueError): + pass + marker_stem = _normalized_config_token(marker_path.stem) + run_dir_name = _normalized_config_token(config_path.parent.name) + if marker_stem and run_dir_name and (run_dir_name in marker_stem or marker_stem in run_dir_name): + return True + return str(marker_path) == str(config_path) + + +def _normalized_config_token(value: str) -> str: + return re.sub(r"[^a-z0-9]+", "", value.lower()) + + +def _log_segment_for_resolved_config(log_text: str, config_path: Path) -> tuple[str, list[str]]: + markers = list(RESOLVED_CONFIG_RE.finditer(log_text)) + if not markers: + return log_text, [] + + matching_indexes = [ + index for index, marker in enumerate(markers) if _same_resolved_config_path(marker.group("path"), config_path) + ] + if not matching_indexes: + return "", [f"log_segments_total={len(markers)}", "log_segment_selected=none_for_resolved_config"] + + selected_index = matching_indexes[-1] + segment_end = markers[selected_index + 1].start() if selected_index + 1 < len(markers) else len(log_text) + notes = [ + f"log_segments_total={len(markers)}", + f"log_matching_segments={len(matching_indexes)}", + f"log_segment_selected={matching_indexes.index(selected_index) + 1}_of_{len(matching_indexes)}", + ] + return log_text[markers[selected_index].start() : segment_end], notes + + +def _summary_values(summaries: list[dict[str, Any]], key: str) -> list[float]: + return [float(summary[key]) for summary in summaries if isinstance(summary.get(key), (int, float))] + + +def _summary_mean(summaries: list[dict[str, Any]], key: str) -> float | None: + values = _summary_values(summaries, key) + return statistics.fmean(values) if values else None + + +def _summary_max(summaries: list[dict[str, Any]], key: str) -> float | None: + values = _summary_values(summaries, key) + return max(values) if values else None + + +def _summary_min(summaries: list[dict[str, Any]], key: str) -> int | None: + values = [int(summary[key]) for summary in summaries if isinstance(summary.get(key), int)] + return min(values) if values else None + + +def _summary_dict_metric(summaries: list[dict[str, Any]], key: str, *, method: str = "mean") -> dict[str, float]: + values_by_name: dict[str, list[float]] = {} + for summary in summaries: + values = summary.get(key) + if not isinstance(values, dict): + continue + for name, value in values.items(): + if isinstance(value, (int, float)): + values_by_name.setdefault(str(name), []).append(float(value)) + result: dict[str, float] = {} + for name, values in sorted(values_by_name.items()): + if not values: + continue + result[name] = max(values) if method == "max" else statistics.fmean(values) + return result + + +def _aggregate_resolved_run_summaries(summaries: list[dict[str, Any]]) -> dict[str, Any]: + if not summaries: + return {} + if len(summaries) == 1: + return summaries[0] + return { + "tokens_per_sec_mean": _summary_mean(summaries, "tokens_per_sec_mean"), + "step_time_s_mean": _summary_mean(summaries, "step_time_s_mean"), + "tokens_per_sec_std": _summary_mean(summaries, "tokens_per_sec_std"), + "tokens_per_sec_cv": _summary_mean(summaries, "tokens_per_sec_cv"), + "step_time_s_std": _summary_mean(summaries, "step_time_s_std"), + "step_time_s_cv": _summary_mean(summaries, "step_time_s_cv"), + "tokens_per_step_median": _summary_mean(summaries, "tokens_per_step_median"), + "phase_time_sec": _summary_dict_metric(summaries, "phase_time_sec"), + "phase_time_rank_mean_sec": _summary_dict_metric(summaries, "phase_time_rank_mean_sec"), + "phase_time_share": _summary_dict_metric(summaries, "phase_time_share"), + "phase_memory_peak_gb": _summary_dict_metric(summaries, "phase_memory_peak_gb", method="max"), + "peak_mem_gb_max": _summary_max(summaries, "peak_mem_gb_max"), + "mfu_mean": _summary_mean(summaries, "mfu_mean"), + "tflops_per_gpu_mean": _summary_mean(summaries, "tflops_per_gpu_mean"), + "measured_steps": _summary_min(summaries, "measured_steps"), + "warmup_excluded": _summary_min(summaries, "warmup_excluded"), + "parsed_step_count": _summary_min(summaries, "parsed_step_count"), + } + + +def _combined_failure_status(log_texts: list[str]) -> str | None: + statuses = [status for text in log_texts if (status := _log_failure_status(text)) is not None] + if "oom" in statuses: + return "oom" + return statuses[0] if statuses else None + + +def _resolved_run_behavior_point( + benchmark_path: Path, + config_path: Path, + *, + label_root: Path | None = None, + label_prefix: str | None = None, + label: str | None = None, + log_paths: list[Path] | None = None, + warmup_steps: int | None = None, + metrics_only_reason: str | None = None, + notes: list[str] | None = None, +) -> BenchmarkBehaviorPoint | None: + run_dir = config_path.parent + raw_config = load_training_config(config_path) + try: + topology = resolve_topology(raw_config) + except ValueError: + return None + + startup_metrics = _load_startup_metrics(run_dir) + resolved_log_paths = _resolved_run_log_paths(benchmark_path, run_dir, startup_metrics, log_paths=log_paths) + log_texts: list[str] = [] + log_segment_notes: list[str] = [] + observed_summaries: list[dict[str, Any]] = [] + for log_path in resolved_log_paths: + full_log_text = log_path.read_text(encoding="utf-8", errors="replace") + log_text, segment_notes = _log_segment_for_resolved_config(full_log_text, config_path) + log_texts.append(log_text) + log_segment_notes.extend(segment_notes) + observed = parse_log_text(log_text, source=str(log_path)) + selected_warmup_steps = warmup_steps if warmup_steps is not None else 2 if len(observed.steps) > 2 else 0 + observed_summaries.append( + summarize_observed_run( + observed, + warmup_steps=selected_warmup_steps, + world_size=topology.world_size, + ) + ) + failure_status = _combined_failure_status(log_texts) + observed_summary = _aggregate_resolved_run_summaries(observed_summaries) + + tokens_per_sec = _round_or_none(observed_summary.get("tokens_per_sec_mean"), 3) + step_time_sec = _round_or_none(observed_summary.get("step_time_s_mean"), 6) + tokens_per_step = _round_or_none(observed_summary.get("tokens_per_step_median"), 1) + tokens_per_sec_std = _round_or_none(observed_summary.get("tokens_per_sec_std"), 3) + tokens_per_sec_cv = _round_or_none(observed_summary.get("tokens_per_sec_cv"), 6) + step_time_sec_std = _round_or_none(observed_summary.get("step_time_s_std"), 6) + step_time_sec_cv = _round_or_none(observed_summary.get("step_time_s_cv"), 6) + phase_time_sec = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_time_sec")).items() + } + phase_time_rank_mean_sec = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_time_rank_mean_sec")).items() + } + phase_time_share = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_time_share")).items() + } + phase_memory_peak_gb = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_memory_peak_gb")).items() + } + peak_mem_gb = _round_or_none(observed_summary.get("peak_mem_gb_max"), 3) + if failure_status == "oom": + oom_peak_mem_gb = _round_or_none(max((_oom_peak_mem_gb(text) or 0.0) for text in log_texts), 3) + if oom_peak_mem_gb is not None: + peak_mem_gb = max(value for value in (peak_mem_gb, oom_peak_mem_gb) if value is not None) + observed_has_metrics = peak_mem_gb is not None or bool(phase_time_sec) or bool(phase_memory_peak_gb) + if metrics_only_reason is not None and failure_status is None: + tokens_per_sec = None + step_time_sec = None + tokens_per_sec_std = None + tokens_per_sec_cv = None + step_time_sec_std = None + step_time_sec_cv = None + tokens_per_step = None + if ( + tokens_per_sec is None + and failure_status is None + and not (metrics_only_reason is not None and observed_has_metrics) + ): + return None + + if failure_status == "oom" and tokens_per_sec is None: + status = "observed_log_oom" + correctness_status = "oom" + elif failure_status is not None: + status = "observed_log_partial_failure" + correctness_status = failure_status + elif metrics_only_reason is not None: + status = "observed_log_metrics_only" + correctness_status = "not_promoted" + else: + status = "observed_log_summary" + correctness_status = "not_promoted" + + metrics = startup_metrics.get("metrics", {}) + source_paths = [str(path) for path in resolved_log_paths] + point_notes = [ + f"warmup_excluded={observed_summary.get('warmup_excluded', 0)}", + f"parsed_steps={observed_summary.get('parsed_step_count', 0)}", + f"resolved_log_count={len(resolved_log_paths)}", + *dict.fromkeys(log_segment_notes), + *(notes or []), + ] + if startup_metrics.get("repo_commit"): + point_notes.append(f"commit={startup_metrics['repo_commit']}") + if isinstance(metrics.get("startup/master_addr"), str): + point_notes.append(f"master_addr={metrics['startup/master_addr']}") + if failure_status is not None: + point_notes.append(f"log_failure_status={failure_status}") + if metrics_only_reason is not None and failure_status is None: + point_notes.append(f"metrics_only={metrics_only_reason}") + + return BenchmarkBehaviorPoint( + label=label or _resolved_run_label(config_path, label_root or benchmark_path, label_prefix), + source=";".join(source_paths) if source_paths else str(config_path), + micro_batch_size=topology.micro_batch_size, + global_batch_size=topology.global_batch_size, + gradient_accumulation_steps=topology.gradient_accumulation_steps, + tokens_per_sec=tokens_per_sec, + step_time_sec=step_time_sec, + tokens_per_sec_std=tokens_per_sec_std, + tokens_per_sec_cv=tokens_per_sec_cv, + step_time_sec_std=step_time_sec_std, + step_time_sec_cv=step_time_sec_cv, + tokens_per_step=tokens_per_step, + phase_time_sec=phase_time_sec, + phase_time_rank_mean_sec=phase_time_rank_mean_sec, + phase_time_share=phase_time_share, + phase_memory_peak_gb=phase_memory_peak_gb, + mfu_percent=_round_or_none((observed_summary.get("mfu_mean") or 0.0) * 100.0, 3) + if observed_summary.get("mfu_mean") is not None and metrics_only_reason is None + else None, + tflops_per_gpu=_round_or_none(observed_summary.get("tflops_per_gpu_mean"), 3) + if metrics_only_reason is None + else None, + peak_mem_gb=peak_mem_gb, + allocator_retries=None, + measured_steps=observed_summary.get("measured_steps"), + warmup_steps=observed_summary.get("warmup_excluded"), + gpu_count=topology.world_size, + model_ref=_first_non_none( + _model_ref_from_text_or_path(json.dumps(raw_config, sort_keys=True), config_path), + model_ref_from_config(raw_config), + ), + sample_packing_sequence_len=topology.sample_packing_sequence_len, + data_parallel_replicate_size=topology.data_parallel_replicate_size, + data_parallel_shard_size=topology.data_parallel_shard_size, + tensor_parallel_size=topology.tensor_parallel_size, + pipeline_parallel_size=topology.pipeline_parallel_size, + ulysses_parallel_size=topology.ulysses_parallel_size, + ringattn_parallel_size=topology.ringattn_parallel_size, + expert_parallel_size=topology.expert_parallel_size, + ep_fsdp_size=topology.ep_fsdp_size, + deepep_async_combine=_config_bool(raw_config, "model", "deepep_async_combine", False), + deepep_num_sms=_config_int(raw_config, "model", "deepep_num_sms"), + deepep_buffer_size_gb=_config_float(raw_config, "model", "deepep_buffer_size_gb"), + enable_compile=_config_bool(raw_config, "train", "enable_compile", False), + gradient_checkpointing_method=_config_str(raw_config, "train", "gradient_checkpointing_method"), + enable_activation_offload=_config_bool(raw_config, "train", "enable_activation_offload", False), + activation_offload_prefetch_count=_config_int(raw_config, "train", "activation_offload_prefetch_count"), + skip_param_upcast=_config_bool(raw_config, "train", "skip_param_upcast", False), + fsdp_reduce_dtype=_config_fsdp_reduce_dtype(raw_config), + ce_mode=_config_ce_mode(raw_config), + moe_implementation=_config_str(raw_config, "model", "moe_implementation"), + moe_checkpoint_method=_config_str(raw_config, "train", "moe_checkpoint_method"), + muon_momentum=_config_float(raw_config, "train", "muon_momentum"), + muon_update_dtype=_config_str(raw_config, "train", "muon_update_dtype"), + attention_backend=_config_attention_backend(raw_config), + balanced_routing=_config_balanced_routing(raw_config), + status=status, + correctness_status=correctness_status, + notes=point_notes, + ) + + +def _resolved_run_points(benchmark_path: Path) -> list[BenchmarkBehaviorPoint]: + points: list[BenchmarkBehaviorPoint] = [] + for config_path in sorted(benchmark_path.rglob("xorl_cli.yaml")): + if not config_path.is_file(): + continue + point = _resolved_run_behavior_point(benchmark_path, config_path) + if point is not None: + points.append(point) + return points + + +def _logged_training_config(log_text: str) -> dict[str, Any] | None: + decoder = json.JSONDecoder() + starts = [match.start("brace") for match in re.finditer(r"xorl\.trainers\.trainer\s+-\s+(?P\{)", log_text)] + for start in reversed(starts): + try: + parsed, _ = decoder.raw_decode(log_text[start:]) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict) and {"model", "train", "data"}.issubset(parsed): + return parsed + return None + + +def _standalone_log_label(log_path: Path, benchmark_path: Path) -> str: + try: + relative = log_path.parent.relative_to(benchmark_path) + except ValueError: + relative = log_path.parent.name + return f"observed_log:{relative.as_posix() if isinstance(relative, Path) else relative}" + + +def _standalone_log_behavior_point( + benchmark_path: Path, + log_path: Path, +) -> BenchmarkBehaviorPoint | None: + full_log_text = log_path.read_text(encoding="utf-8", errors="replace") + raw_config = _logged_training_config(full_log_text) + if raw_config is None: + return None + try: + topology = resolve_topology(raw_config) + except ValueError: + return None + + failure_status = _log_failure_status(full_log_text) + observed = parse_log_text(full_log_text, source=str(log_path)) + warmup_steps = 2 if len(observed.steps) > 2 else 0 + observed_summary = summarize_observed_run( + observed, + warmup_steps=warmup_steps, + world_size=topology.world_size, + ) + tokens_per_sec = _round_or_none(observed_summary.get("tokens_per_sec_mean"), 3) + peak_mem_gb = _round_or_none(observed_summary.get("peak_mem_gb_max"), 3) + if failure_status == "oom": + oom_peak_mem_gb = _round_or_none(_oom_peak_mem_gb(full_log_text), 3) + if oom_peak_mem_gb is not None: + peak_mem_gb = max(value for value in (peak_mem_gb, oom_peak_mem_gb) if value is not None) + if tokens_per_sec is None and failure_status is None and peak_mem_gb is None: + return None + + if failure_status == "oom" and tokens_per_sec is None: + status = "observed_log_oom" + correctness_status = "oom" + elif failure_status is not None: + status = "observed_log_partial_failure" + correctness_status = failure_status + else: + status = "observed_log_summary" + correctness_status = "not_promoted" + + phase_time_sec = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_time_sec")).items() + } + phase_time_rank_mean_sec = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_time_rank_mean_sec")).items() + } + phase_time_share = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_time_share")).items() + } + phase_memory_peak_gb = { + key: round(value, 6) for key, value in _float_dict(observed_summary.get("phase_memory_peak_gb")).items() + } + notes = [ + "source=standalone_node_log", + f"warmup_excluded={observed_summary.get('warmup_excluded', 0)}", + f"parsed_steps={observed_summary.get('parsed_step_count', 0)}", + ] + resolved_config_match = RESOLVED_CONFIG_RE.search(full_log_text) + if resolved_config_match: + notes.append(f"resolved_config={resolved_config_match.group('path')}") + git_match = re.search(r"(?m)^git\s+(?P[0-9a-f]{7,40})\s*$", full_log_text) + if git_match: + notes.append(f"commit={git_match.group('commit')}") + if failure_status is not None: + notes.append(f"log_failure_status={failure_status}") + + return BenchmarkBehaviorPoint( + label=_standalone_log_label(log_path, benchmark_path), + source=str(log_path), + micro_batch_size=topology.micro_batch_size, + global_batch_size=topology.global_batch_size, + gradient_accumulation_steps=topology.gradient_accumulation_steps, + tokens_per_sec=tokens_per_sec, + step_time_sec=_round_or_none(observed_summary.get("step_time_s_mean"), 6), + tokens_per_sec_std=_round_or_none(observed_summary.get("tokens_per_sec_std"), 3), + tokens_per_sec_cv=_round_or_none(observed_summary.get("tokens_per_sec_cv"), 6), + step_time_sec_std=_round_or_none(observed_summary.get("step_time_s_std"), 6), + step_time_sec_cv=_round_or_none(observed_summary.get("step_time_s_cv"), 6), + phase_time_sec=phase_time_sec, + phase_time_rank_mean_sec=phase_time_rank_mean_sec, + phase_time_share=phase_time_share, + phase_memory_peak_gb=phase_memory_peak_gb, + mfu_percent=_round_or_none((observed_summary.get("mfu_mean") or 0.0) * 100.0, 3) + if observed_summary.get("mfu_mean") is not None + else None, + tflops_per_gpu=_round_or_none(observed_summary.get("tflops_per_gpu_mean"), 3), + peak_mem_gb=peak_mem_gb, + allocator_retries=None, + measured_steps=observed_summary.get("measured_steps"), + warmup_steps=observed_summary.get("warmup_excluded"), + gpu_count=topology.world_size, + model_ref=_first_non_none( + _model_ref_from_text_or_path(json.dumps(raw_config, sort_keys=True), log_path), + model_ref_from_config(raw_config), + ), + sample_packing_sequence_len=topology.sample_packing_sequence_len, + data_parallel_replicate_size=topology.data_parallel_replicate_size, + data_parallel_shard_size=topology.data_parallel_shard_size, + tensor_parallel_size=topology.tensor_parallel_size, + pipeline_parallel_size=topology.pipeline_parallel_size, + ulysses_parallel_size=topology.ulysses_parallel_size, + ringattn_parallel_size=topology.ringattn_parallel_size, + expert_parallel_size=topology.expert_parallel_size, + ep_fsdp_size=topology.ep_fsdp_size, + deepep_async_combine=_config_bool(raw_config, "model", "deepep_async_combine", False), + deepep_num_sms=_config_int(raw_config, "model", "deepep_num_sms"), + deepep_buffer_size_gb=_config_float(raw_config, "model", "deepep_buffer_size_gb"), + enable_compile=_config_bool(raw_config, "train", "enable_compile", False), + gradient_checkpointing_method=_config_str(raw_config, "train", "gradient_checkpointing_method"), + enable_activation_offload=_config_bool(raw_config, "train", "enable_activation_offload", False), + activation_offload_prefetch_count=_config_int(raw_config, "train", "activation_offload_prefetch_count"), + skip_param_upcast=_config_bool(raw_config, "train", "skip_param_upcast", False), + fsdp_reduce_dtype=_config_fsdp_reduce_dtype(raw_config), + ce_mode=_config_ce_mode(raw_config), + moe_implementation=_config_str(raw_config, "model", "moe_implementation"), + moe_checkpoint_method=_config_str(raw_config, "train", "moe_checkpoint_method"), + muon_momentum=_config_float(raw_config, "train", "muon_momentum"), + muon_update_dtype=_config_str(raw_config, "train", "muon_update_dtype"), + attention_backend=_config_attention_backend(raw_config), + balanced_routing=_config_balanced_routing(raw_config), + status=status, + correctness_status=correctness_status, + notes=notes, + ) + + +def _standalone_log_points( + benchmark_path: Path, + *, + used_sources: set[str], +) -> list[BenchmarkBehaviorPoint]: + points: list[BenchmarkBehaviorPoint] = [] + for log_path in sorted(benchmark_path.rglob("node-0.log")): + if str(log_path) in used_sources or str(log_path.resolve(strict=False)) in used_sources: + continue + point = _standalone_log_behavior_point(benchmark_path, log_path) + if point is not None: + points.append(point) + return points + + +_FLASHQLA_PROFILE_PHASE_KEYS = { + "train.forward_loss": "forward", + "train.backward": "backward", + "train.optimizer_step": "optimizer", + "train.reduce_metrics": "reduce_metrics", +} + + +def _flashqla_profile_phase_time_sec(profile: Any) -> dict[str, float]: + if not isinstance(profile, dict): + return {} + phases: dict[str, float] = {} + for profile_key, phase_name in _FLASHQLA_PROFILE_PHASE_KEYS.items(): + row = profile.get(profile_key) + if not isinstance(row, dict): + continue + value_ms = _first_non_none(row.get("max_ms_rank"), row.get("mean_ms_per_rank")) + if value_ms is None: + continue + phases[phase_name] = round(float(value_ms) / 1000.0, 6) + return phases + + +def _phase_time_share_from_step(phase_time_sec: dict[str, float], step_time_sec: float | None) -> dict[str, float]: + if not phase_time_sec or step_time_sec is None or step_time_sec <= 0: + return {} + return {phase: round(value / step_time_sec, 6) for phase, value in phase_time_sec.items()} + + +def _flashqla_summary_config_path(summary_path: Path, backend: str, source: dict[str, Any]) -> Path: + config_paths_by_backend = source.get("config_paths_by_backend", source.get("configs_by_backend", {})) + if isinstance(config_paths_by_backend, dict) and backend in config_paths_by_backend: + resolved = _resolve_manifest_path(summary_path.parent, config_paths_by_backend[backend]) + if resolved is not None: + return resolved + config_root = _resolve_manifest_path(summary_path.parent, source.get("config_root")) + base = config_root if config_root is not None else summary_path.parent + return base / f"out_{backend}" / "xorl_cli.yaml" + + +def _flashqla_summary_point( + *, + summary_path: Path, + item: dict[str, Any], + backend: str, + backend_entry: dict[str, Any], + source: dict[str, Any], + source_notes: list[str], +) -> BenchmarkBehaviorPoint | None: + steps = backend_entry.get("steps") + if not isinstance(steps, dict): + return None + tokens_per_sec = steps.get("mean_tokens_per_sec") + if tokens_per_sec is None: + return None + + config_path = _flashqla_summary_config_path(summary_path, backend, source) + if not config_path.is_file(): + return None + raw_config = load_training_config(config_path) + try: + topology = resolve_topology(raw_config) + except ValueError: + return None + + run_name = str(item.get("run") or summary_path.parent.name) + label_prefix = source.get("label_prefix") + prefix = f"{label_prefix}/" if label_prefix else "" + step_time_sec = _round_or_none(steps.get("mean_step_time_s"), 6) + phase_time_sec = _flashqla_profile_phase_time_sec(backend_entry.get("profile")) + measured_steps = steps.get("measured_steps") + total_steps = steps.get("steps") + warmup_steps = None + if total_steps is not None and measured_steps is not None: + warmup_steps = max(0, int(total_steps) - int(measured_steps)) + + notes = [ + *source_notes, + f"summary_run={run_name}", + f"attention_backend={backend}", + ] + if total_steps is not None: + notes.append(f"summary_steps={total_steps}") + + return BenchmarkBehaviorPoint( + label=f"flashqla_summary:{prefix}{run_name}/{backend}", + source=str(summary_path), + micro_batch_size=topology.micro_batch_size, + global_batch_size=topology.global_batch_size, + gradient_accumulation_steps=topology.gradient_accumulation_steps, + tokens_per_sec=_round_or_none(tokens_per_sec, 3), + step_time_sec=step_time_sec, + phase_time_sec=phase_time_sec, + phase_time_share=_phase_time_share_from_step(phase_time_sec, step_time_sec), + mfu_percent=None, + tflops_per_gpu=_round_or_none(steps.get("mean_tflops_per_gpu"), 3), + peak_mem_gb=None, + allocator_retries=None, + measured_steps=int(measured_steps) if measured_steps is not None else None, + warmup_steps=warmup_steps, + gpu_count=topology.world_size, + model_ref=_first_non_none( + _model_ref_from_text_or_path(json.dumps(raw_config, sort_keys=True), config_path), + model_ref_from_config(raw_config), + ), + sample_packing_sequence_len=topology.sample_packing_sequence_len, + data_parallel_replicate_size=topology.data_parallel_replicate_size, + data_parallel_shard_size=topology.data_parallel_shard_size, + tensor_parallel_size=topology.tensor_parallel_size, + pipeline_parallel_size=topology.pipeline_parallel_size, + ulysses_parallel_size=topology.ulysses_parallel_size, + ringattn_parallel_size=topology.ringattn_parallel_size, + expert_parallel_size=topology.expert_parallel_size, + ep_fsdp_size=topology.ep_fsdp_size, + deepep_async_combine=_config_bool(raw_config, "model", "deepep_async_combine", False), + deepep_num_sms=_config_int(raw_config, "model", "deepep_num_sms"), + deepep_buffer_size_gb=_config_float(raw_config, "model", "deepep_buffer_size_gb"), + enable_compile=_config_bool(raw_config, "train", "enable_compile", False), + gradient_checkpointing_method=_config_str(raw_config, "train", "gradient_checkpointing_method"), + enable_activation_offload=_config_bool(raw_config, "train", "enable_activation_offload", False), + activation_offload_prefetch_count=_config_int(raw_config, "train", "activation_offload_prefetch_count"), + skip_param_upcast=_config_bool(raw_config, "train", "skip_param_upcast", False), + fsdp_reduce_dtype=_config_fsdp_reduce_dtype(raw_config), + ce_mode=_config_ce_mode(raw_config), + moe_implementation=_config_str(raw_config, "model", "moe_implementation"), + moe_checkpoint_method=_config_str(raw_config, "train", "moe_checkpoint_method"), + muon_momentum=_config_float(raw_config, "train", "muon_momentum"), + muon_update_dtype=_config_str(raw_config, "train", "muon_update_dtype"), + attention_backend=backend, + balanced_routing=_config_balanced_routing(raw_config), + status="historical_flashqla_profile_summary", + correctness_status="not_promoted", + notes=notes, + ) + + +def _manifest_flashqla_summary_points(benchmark_path: Path) -> list[BenchmarkBehaviorPoint]: + manifest = _read_benchmark_source_manifest(benchmark_path) + if manifest is None: + return [] + manifest_path, payload = manifest + manifest_dir = manifest_path.parent + points: list[BenchmarkBehaviorPoint] = [] + + for source in _as_list(payload.get("flashqla_profile_summaries")): + if isinstance(source, str): + source = {"path": source} + if not isinstance(source, dict): + continue + summary_paths = _expand_manifest_path_patterns(manifest_dir, source.get("path", source.get("summary"))) + selected_backends = {str(item) for item in _as_list(source.get("backends"))} + source_notes = [f"benchmark_source_manifest={manifest_path.name}", "source=flashqla_profile_summaries"] + label_prefix = source.get("label_prefix") + if label_prefix: + source_notes.append(f"source_label_prefix={label_prefix}") + for summary_path in summary_paths: + if not summary_path.is_file(): + continue + payload = json.loads(summary_path.read_text(encoding="utf-8")) + items = payload if isinstance(payload, list) else [payload] + for item in items: + if not isinstance(item, dict): + continue + backends = item.get("backends", {}) + if not isinstance(backends, dict): + continue + for backend, backend_entry in sorted(backends.items()): + if selected_backends and backend not in selected_backends: + continue + if not isinstance(backend_entry, dict): + continue + point = _flashqla_summary_point( + summary_path=summary_path, + item=item, + backend=str(backend), + backend_entry=backend_entry, + source=source, + source_notes=source_notes, + ) + if point is not None: + points.append(point) + return points + + +def _manifest_resolved_run_points(benchmark_path: Path) -> list[BenchmarkBehaviorPoint]: + manifest = _read_benchmark_source_manifest(benchmark_path) + if manifest is None: + return [] + manifest_path, payload = manifest + manifest_dir = manifest_path.parent + points: list[BenchmarkBehaviorPoint] = [] + + for source in _as_list(payload.get("resolved_run_roots")): + if isinstance(source, str): + source = {"path": source} + if not isinstance(source, dict): + continue + root = _resolve_manifest_path(manifest_dir, source.get("path")) + if root is None or not root.is_dir(): + continue + label_prefix = source.get("label_prefix") + label_prefix = str(label_prefix) if label_prefix is not None else None + source_notes = [f"benchmark_source_manifest={manifest_path.name}"] + if label_prefix: + source_notes.append(f"source_label_prefix={label_prefix}") + for config_path in sorted(root.rglob("xorl_cli.yaml")): + if not config_path.is_file(): + continue + point = _resolved_run_behavior_point( + benchmark_path, + config_path, + label_root=root, + label_prefix=label_prefix, + log_paths=_log_paths_for_manifest_run(manifest_dir, root, config_path, source), + warmup_steps=_manifest_run_int( + root, + config_path, + source, + field_name="warmup_steps", + by_run_field_name="warmup_steps_by_run", + ), + metrics_only_reason=_manifest_run_metrics_only_reason(root, config_path, source), + notes=source_notes, + ) + if point is not None: + points.append(point) + + for source in _as_list(payload.get("resolved_runs")): + if not isinstance(source, dict): + continue + config_path = _resolve_manifest_path(manifest_dir, source.get("config")) + if config_path is None or not config_path.is_file(): + continue + label = source.get("label") + if isinstance(label, str) and not label.startswith("resolved_run:"): + label = f"resolved_run:{label}" + elif label is not None: + label = str(label) + log_root = _resolve_manifest_path(manifest_dir, source.get("log_root")) + log_base = log_root if log_root is not None else manifest_dir + log_paths = _expand_manifest_path_patterns(log_base, source.get("logs")) + source_notes = [f"benchmark_source_manifest={manifest_path.name}", "source=resolved_runs"] + point = _resolved_run_behavior_point( + benchmark_path, + config_path, + label=label, + log_paths=log_paths, + warmup_steps=int(source["warmup_steps"]) if source.get("warmup_steps") is not None else None, + metrics_only_reason=( + None + if source.get("metrics_only") in (None, False) + else "metrics_only" + if source.get("metrics_only") is True + else str(source.get("metrics_only")) + ), + notes=source_notes, + ) + if point is not None: + points.append(point) + return points + + +def _behavior_override_entries(payload: dict[str, Any]) -> dict[str, dict[str, Any]]: + raw_points = payload.get("points", payload.get("overrides", {})) + if isinstance(raw_points, dict): + return {str(label): dict(value) for label, value in raw_points.items() if isinstance(value, dict)} + entries: dict[str, dict[str, Any]] = {} + if isinstance(raw_points, list): + for item in raw_points: + if not isinstance(item, dict) or not item.get("label"): + continue + label = str(item["label"]) + entry = dict(item) + entry.pop("label", None) + entries[label] = entry + return entries + + +def _coerce_behavior_override_field(field_name: str, value: Any) -> Any: + if field_name in {"phase_time_sec", "phase_time_rank_mean_sec", "phase_time_share", "phase_memory_peak_gb"}: + return _float_dict(value) + if field_name == "notes": + return [str(item) for item in _as_list(value)] + return value + + +def _apply_behavior_point_overrides( + benchmark_path: Path, + points: list[BenchmarkBehaviorPoint], +) -> list[BenchmarkBehaviorPoint]: + manifest = _read_manifest_from_candidates(benchmark_path, BENCHMARK_BEHAVIOR_OVERRIDE_MANIFESTS) + if manifest is None: + return points + manifest_path, payload = manifest + overrides = _behavior_override_entries(payload) + if not overrides: + return points + + valid_fields = {field.name for field in fields(BenchmarkBehaviorPoint)} + result: list[BenchmarkBehaviorPoint] = [] + for point in points: + override = overrides.get(point.label) + if override is None: + result.append(point) + continue + replace_kwargs: dict[str, Any] = {} + for field_name, value in override.items(): + if field_name in {"label", "notes_append"} or field_name not in valid_fields: + continue + replace_kwargs[field_name] = _coerce_behavior_override_field(field_name, value) + notes = list(replace_kwargs.pop("notes", point.notes)) + notes.append(f"behavior_override_manifest={manifest_path.name}") + for item in _as_list(override.get("notes_append")): + notes.append(str(item)) + result.append(replace(point, notes=notes, **replace_kwargs)) + return result + + +def load_benchmark_behavior_points(benchmark_dir: str | Path) -> list[BenchmarkBehaviorPoint]: + benchmark_path = resolve_calibration_pack(benchmark_dir) + points: list[BenchmarkBehaviorPoint] = [] + topology_defaults: dict[str, int | float | bool | str] = {} + + for readme_path in (benchmark_path / "README.md", benchmark_path / "RESULTS.md"): + if not readme_path.is_file(): + continue + readme_text = readme_path.read_text(encoding="utf-8") + model_ref = _model_ref_from_text_or_path(readme_text, readme_path) + if model_ref is not None: + topology_defaults["model_ref"] = model_ref + topology_defaults.update(_readme_topology_defaults(readme_text)) + seq_len = _seq_len_from_readme(readme_text) + if seq_len is not None: + topology_defaults["sample_packing_sequence_len"] = seq_len + readme_reference = _readme_point(readme_text, source=str(readme_path), model_ref=model_ref) + if readme_reference is not None: + points.append(readme_reference) + adjacent_mbs10 = _readme_adjacent_mbs10_point( + readme_text, source=str(readme_path), seq_len=seq_len, model_ref=model_ref + ) + if adjacent_mbs10 is not None: + points.append(adjacent_mbs10) + points.extend(_q235_markdown_points(readme_text, source=str(readme_path), model_ref=model_ref)) + points.extend(_q35_markdown_points(readme_text, source=str(readme_path), model_ref=model_ref)) + + for result_path in sorted((benchmark_path / "results").glob("*.json")): + result = json.loads(result_path.read_text(encoding="utf-8")) + result_model_ref = _model_ref_from_text_or_path(json.dumps(result, sort_keys=True), result_path) + result_defaults = dict(topology_defaults) + if result_model_ref is not None: + result_defaults["model_ref"] = result_model_ref + for row in result.get("best_by_mfu", []): + if isinstance(row, dict) and row.get("trial"): + points.append(_best_by_mfu_point(result_path, result, row, topology_defaults=result_defaults)) + throughput = result.get("throughput") + if isinstance(throughput, dict): + points.append( + _with_k3_status( + _result_throughput_point(result_path, result, topology_defaults=result_defaults), result + ) + ) + + points.extend(_manifest_resolved_run_points(benchmark_path)) + points.extend(_manifest_flashqla_summary_points(benchmark_path)) + points.extend(_resolved_run_points(benchmark_path)) + used_sources = {point.source for point in points} + used_sources.update(str(Path(point.source).resolve(strict=False)) for point in points if point.source) + points.extend(_standalone_log_points(benchmark_path, used_sources=used_sources)) + return _apply_behavior_point_overrides(benchmark_path, points) + + +def behavior_point_matches_topology(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + if point.micro_batch_size != topology.micro_batch_size or point.global_batch_size != topology.global_batch_size: + return False + if ( + point.gradient_accumulation_steps is not None + and point.gradient_accumulation_steps != topology.gradient_accumulation_steps + ): + return False + if not _point_parallel_size_matches(point.tensor_parallel_size, topology.tensor_parallel_size): + return False + if not _point_parallel_size_matches(point.pipeline_parallel_size, topology.pipeline_parallel_size): + return False + if not _point_parallel_size_matches(point.ulysses_parallel_size, topology.ulysses_parallel_size): + return False + if not _point_parallel_size_matches(point.ringattn_parallel_size, topology.ringattn_parallel_size): + return False + if point.expert_parallel_size is None: + if topology.expert_parallel_size != 1: + return False + elif point.expert_parallel_size != topology.expert_parallel_size: + return False + if point.ep_fsdp_size is not None and point.ep_fsdp_size != topology.ep_fsdp_size: + return False + if ( + point.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len is not None + and point.sample_packing_sequence_len != topology.sample_packing_sequence_len + ): + return False + return True + + +def _section(raw_config: dict[str, Any], name: str) -> dict[str, Any]: + value = raw_config.get(name, {}) + return value if isinstance(value, dict) else {} + + +def _config_bool(raw_config: dict[str, Any], section_name: str, key: str, default: bool | None = None) -> bool | None: + section = _section(raw_config, section_name) + value = section.get(key, default) + if value is None: + return None + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + + +def _config_int(raw_config: dict[str, Any], section_name: str, key: str) -> int | None: + section = _section(raw_config, section_name) + value = section.get(key) + return int(value) if value is not None else None + + +def _config_float(raw_config: dict[str, Any], section_name: str, key: str) -> float | None: + section = _section(raw_config, section_name) + value = section.get(key) + return float(value) if value is not None else None + + +def _config_str(raw_config: dict[str, Any], section_name: str, key: str) -> str | None: + section = _section(raw_config, section_name) + value = section.get(key) + return str(value) if value is not None else None + + +def _boolish(value: Any) -> bool | None: + if value is None: + return None + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on", "balanced"} + return bool(value) + + +def _config_balanced_routing(raw_config: dict[str, Any]) -> bool: + for section_name in ("simulator", "_simulator", "train", "model", "data"): + section = _section(raw_config, section_name) + if "balanced_routing" in section: + value = _boolish(section.get("balanced_routing")) + if value is not None: + return value + for key in ("synthetic_routing", "synthetic_routing_mode", "moe_synthetic_routing"): + value = section.get(key) + if isinstance(value, str) and value.strip().lower() == "balanced": + return True + return False + + +def _config_attention_backend(raw_config: dict[str, Any]) -> str | None: + for section_name in ("simulator", "_simulator"): + section = _section(raw_config, section_name) + value = section.get("attention_backend") + if value is not None: + return str(value) + return None + + +def _config_fsdp_reduce_dtype(raw_config: dict[str, Any]) -> str: + return _config_str(raw_config, "train", "fsdp_reduce_dtype") or "fp32" + + +def _config_ce_mode(raw_config: dict[str, Any]) -> str | None: + return _config_str(raw_config, "train", "ce_mode") or "compiled" + + +def behavior_point_workload_mismatches(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> list[str]: + config_activation_offload = _config_bool(raw_config, "train", "enable_activation_offload", False) + checks: tuple[tuple[str, Any, Any], ...] = ( + ("balanced_routing", point.balanced_routing, _config_balanced_routing(raw_config)), + ("attention_backend", point.attention_backend, _config_attention_backend(raw_config)), + ( + "deepep_async_combine", + point.deepep_async_combine, + _config_bool(raw_config, "model", "deepep_async_combine", False), + ), + ("deepep_num_sms", point.deepep_num_sms, _config_int(raw_config, "model", "deepep_num_sms")), + ( + "deepep_buffer_size_gb", + point.deepep_buffer_size_gb, + _config_float(raw_config, "model", "deepep_buffer_size_gb"), + ), + ("enable_compile", point.enable_compile, _config_bool(raw_config, "train", "enable_compile", False)), + ( + "gradient_checkpointing_method", + point.gradient_checkpointing_method, + _config_str(raw_config, "train", "gradient_checkpointing_method"), + ), + ( + "enable_activation_offload", + point.enable_activation_offload, + config_activation_offload, + ), + ( + "activation_offload_prefetch_count", + point.activation_offload_prefetch_count, + _config_int(raw_config, "train", "activation_offload_prefetch_count"), + ), + ("skip_param_upcast", point.skip_param_upcast, _config_bool(raw_config, "train", "skip_param_upcast", False)), + ("fsdp_reduce_dtype", point.fsdp_reduce_dtype, _config_fsdp_reduce_dtype(raw_config)), + ("ce_mode", point.ce_mode, _config_ce_mode(raw_config)), + ("moe_implementation", point.moe_implementation, _config_str(raw_config, "model", "moe_implementation")), + ( + "moe_checkpoint_method", + point.moe_checkpoint_method, + _config_str(raw_config, "train", "moe_checkpoint_method"), + ), + ("muon_momentum", point.muon_momentum, _config_float(raw_config, "train", "muon_momentum")), + ("muon_update_dtype", point.muon_update_dtype, _config_str(raw_config, "train", "muon_update_dtype")), + ) + mismatches: list[str] = [] + if behavior_point_model_mismatches(point, raw_config): + mismatches.append("model_ref") + for field_name, point_value, config_value in checks: + if point_value is None: + continue + if ( + field_name == "activation_offload_prefetch_count" + and point.enable_activation_offload is False + and config_activation_offload is False + ): + continue + if isinstance(point_value, float): + if config_value is None or abs(float(point_value) - float(config_value)) > 1e-9: + mismatches.append(field_name) + elif point_value != config_value: + mismatches.append(field_name) + return mismatches + + +def behavior_point_matches_workload(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> bool: + return not behavior_point_workload_mismatches(point, raw_config) + + +def _point_parallel_size_matches(point_value: int | None, topology_value: int) -> bool: + if point_value is None: + return topology_value == 1 + return point_value == topology_value + + +def _fsdp_reduce_dtype_selection_penalty(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any] | None) -> int: + if raw_config is None or point.fsdp_reduce_dtype is None: + return 0 + config_value = _config_str(raw_config, "train", "fsdp_reduce_dtype") or "fp32" + return 0 if point.fsdp_reduce_dtype == config_value else 1 + + +def _behavior_point_evidence_rank(point: BenchmarkBehaviorPoint) -> int: + if point.label.startswith("resolved_run:") and point.status.startswith("observed_log"): + return 0 + return 1 + + +def _select_behavior_point_match( + matches: list[BenchmarkBehaviorPoint], + raw_config: dict[str, Any] | None, +) -> BenchmarkBehaviorPoint: + indexed_matches = list(enumerate(matches)) + _, point = min( + indexed_matches, + key=lambda item: ( + _fsdp_reduce_dtype_selection_penalty(item[1], raw_config), + _behavior_point_evidence_rank(item[1]), + item[0], + ), + ) + return point + + +def predict_benchmark_behavior( + points: list[BenchmarkBehaviorPoint], + topology: Topology, + shape: ShapeLedger, + raw_config: dict[str, Any] | None = None, +) -> BenchmarkBehaviorPrediction: + matches = [ + point + for point in points + if behavior_point_matches_topology(point, topology) + and (raw_config is None or behavior_point_matches_workload(point, raw_config)) + and point.status != "observed_log_metrics_only" + ] + warnings: list[str] = [] + if not matches: + known = ", ".join( + f"{point.label}(mbs={point.micro_batch_size},gb={point.global_batch_size})" for point in points + ) + return BenchmarkBehaviorPrediction( + status="no_calibrated_match", + matched_label=None, + source=None, + tokens_per_sec=None, + tokens_per_sec_per_gpu=None, + step_time_sec=None, + mfu_percent=None, + tflops_per_gpu=None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=None, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + balanced_routing=_config_balanced_routing(raw_config) if raw_config is not None else None, + warnings=[ + f"no empirical behavior point for mbs={topology.micro_batch_size}, gb={topology.global_batch_size}; known: {known}" + ], + ) + + point = _select_behavior_point_match(matches, raw_config) + tokens_per_sec_per_gpu = None + if point.tokens_per_sec is not None and topology.world_size: + tokens_per_sec_per_gpu = point.tokens_per_sec / topology.world_size + step_time_sec = point.step_time_sec + if step_time_sec is None and shape.global_tokens_per_train_step and point.tokens_per_sec: + step_time_sec = shape.global_tokens_per_train_step / point.tokens_per_sec + warnings.append( + "step_time_sec derived as NOMINAL global tokens (gbs x seq) / measured tokens_per_sec; " + "real-dataset packs may not fill bins (q35 65k realized ~0.81x nominal), which would " + "overstate the derived step time by the same factor" + ) + tflops_per_gpu = None + if point.tflops_per_gpu is not None: + tflops_per_gpu = point.tflops_per_gpu + elif point.mfu_percent is not None: + tflops_per_gpu = H100_BF16_PROMISED_TFLOPS_PER_GPU * point.mfu_percent / 100.0 + + if point.status == "allocator_pressure_slowdown": + warnings.append("matched behavior point is an allocator-pressure slowdown, not a promotable speed target") + if point.correctness_status and point.correctness_status != "k3_pass": + warnings.append(f"correctness status is {point.correctness_status}") + + prediction_status = "calibrated_failure" if point.correctness_status == "oom" else "calibrated" + + return BenchmarkBehaviorPrediction( + status=prediction_status, + matched_label=point.label, + source=point.source, + tokens_per_sec=point.tokens_per_sec, + tokens_per_sec_per_gpu=tokens_per_sec_per_gpu, + step_time_sec=step_time_sec, + tokens_per_sec_std=point.tokens_per_sec_std, + tokens_per_sec_cv=point.tokens_per_sec_cv, + step_time_sec_std=point.step_time_sec_std, + step_time_sec_cv=point.step_time_sec_cv, + phase_time_sec=point.phase_time_sec, + phase_time_share=point.phase_time_share, + phase_memory_peak_gb=point.phase_memory_peak_gb, + measured_steps=point.measured_steps, + warmup_steps=point.warmup_steps, + model_ref=point.model_ref, + balanced_routing=point.balanced_routing, + mfu_percent=point.mfu_percent, + tflops_per_gpu=tflops_per_gpu, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=point.peak_mem_gb, + allocator_retries=point.allocator_retries, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + correctness_status=point.correctness_status, + warnings=warnings, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("benchmark_dir", help="Path or built-in calibration-pack name") + args = parser.parse_args() + points = load_benchmark_behavior_points(args.benchmark_dir) + print(json.dumps(to_jsonable({"points": points}), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/src/xorl/sim/calibration_evaluator.py b/src/xorl/sim/calibration_evaluator.py new file mode 100644 index 00000000..b3c48682 --- /dev/null +++ b/src/xorl/sim/calibration_evaluator.py @@ -0,0 +1,1972 @@ +"""Evaluate scenario-prediction fidelity with leave-one-out benchmark holdouts.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml + + +try: + from .benchmark_behavior import ( + behavior_point_model_mismatches, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from .calibration_packs import resolve_calibration_pack, resolve_pack_inputs + from .config_fingerprint import load_training_config, resolve_topology + from .memory_ledger import build_memory_ledger + from .model_metadata import resolve_model_metadata + from .runtime_config import runtime_training_config + from .scenario_planner import ( + _apply_config_override, + _calibrated_memory_peak_estimate, + _calibration_distance, + _calibration_scope, + _candidate_risk_flags, + _extrapolate_behavior, + _measurement_config_filename, + _memory_coverage_for_candidate, + _memory_factor, + _mutated_config, + _phase_bucket, + _prediction_interval, + _prediction_uncertainty_fraction, + _topology_label, + _workload_design_variants, + ) + from .schemas import ( + BenchmarkBehaviorPoint, + CalibrationHoldout, + CalibrationReport, + CalibrationValidationGap, + ScenarioMeasurementConfig, + Topology, + to_jsonable, + ) + from .shape_engine import build_shape_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from benchmark_behavior import ( + behavior_point_model_mismatches, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from calibration_packs import resolve_calibration_pack, resolve_pack_inputs + from config_fingerprint import load_training_config, resolve_topology + from memory_ledger import build_memory_ledger + from model_metadata import resolve_model_metadata + from runtime_config import runtime_training_config + from scenario_planner import ( + _apply_config_override, + _calibrated_memory_peak_estimate, + _calibration_distance, + _calibration_scope, + _candidate_risk_flags, + _extrapolate_behavior, + _measurement_config_filename, + _memory_coverage_for_candidate, + _memory_factor, + _mutated_config, + _phase_bucket, + _prediction_interval, + _prediction_uncertainty_fraction, + _topology_label, + _workload_design_variants, + ) + from schemas import ( + BenchmarkBehaviorPoint, + CalibrationHoldout, + CalibrationReport, + CalibrationValidationGap, + ScenarioMeasurementConfig, + Topology, + to_jsonable, + ) + from shape_engine import build_shape_ledger + + +_THROUGHPUT_MAPE_TARGET_PERCENT = 8.0 +_MEMORY_ABS_ERROR_TARGET_GB = 2.0 +_MEMORY_RELATIVE_ERROR_TARGET_FRACTION = 0.03 +_PHASE_SHARE_ABS_ERROR_TARGET = 0.10 + +_CALIBRATION_COMPONENT_TIMING_OVERRIDES = ( + "train.enable_step_phase_timing=true", + "train.enable_per_component_timing=true", + "train.step_phase_timing_sync_cuda=true", +) + +_CALIBRATION_MEMORY_PROFILE_OVERRIDES = ( + "train.enable_step_phase_timing=true", + "train.enable_step_memory_profiling=true", +) + +_CALIBRATION_REPLAY_OVERRIDES_BY_MEASUREMENT = { + "collect_phase_timing_for_calibration_holdouts": _CALIBRATION_COMPONENT_TIMING_OVERRIDES, + "replay_high_error_holdouts_with_component_timing": _CALIBRATION_COMPONENT_TIMING_OVERRIDES, + "replay_phase_bottleneck_holdouts_with_component_timing": _CALIBRATION_COMPONENT_TIMING_OVERRIDES, + "replay_phase_top3_holdouts_with_component_timing": _CALIBRATION_COMPONENT_TIMING_OVERRIDES, + "replay_phase_share_holdouts_with_component_timing": _CALIBRATION_COMPONENT_TIMING_OVERRIDES, + "collect_peak_memory_for_calibration_holdouts": _CALIBRATION_MEMORY_PROFILE_OVERRIDES, + "replay_high_memory_error_holdouts_with_memory_profile": _CALIBRATION_MEMORY_PROFILE_OVERRIDES, + "replay_memory_bottleneck_holdouts_with_phase_memory_profile": _CALIBRATION_MEMORY_PROFILE_OVERRIDES, +} + +_CALIBRATION_REPLAY_KIND_BY_MEASUREMENT = { + "collect_phase_timing_for_calibration_holdouts": "phase_timing", + "replay_high_error_holdouts_with_component_timing": "component_timing", + "replay_phase_bottleneck_holdouts_with_component_timing": "component_timing", + "replay_phase_top3_holdouts_with_component_timing": "component_timing", + "replay_phase_share_holdouts_with_component_timing": "component_timing", + "collect_peak_memory_for_calibration_holdouts": "memory_profile", + "replay_high_memory_error_holdouts_with_memory_profile": "memory_profile", + "replay_memory_bottleneck_holdouts_with_phase_memory_profile": "memory_profile", +} + +_CALIBRATION_HOLDOUT_REPLAY_KIND_BY_MEASUREMENT = { + "add_scored_same_context_calibration_holdouts": "scored_same_context", + "add_supported_same_context_calibration_holdouts": "supported_same_context", +} + +_CALIBRATION_NEARBY_HOLDOUT_KIND_BY_MEASUREMENT = { + "add_holdouts_to_recalibrate_prediction_interval": "interval_recalibration", + "add_nearby_calibration_holdouts_to_tighten_uncertainty": "nearby_holdout", +} + + +@dataclass(frozen=True) +class _CalibrationMemoryPeakEstimate: + peak_gb: float + overhead_gb: float + source_label: str + notes: list[str] + basis: str = "calibration_residual_floor_peak" + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + if isinstance(value, dict): + return value + raw[name] = {} + return raw[name] + + +def _point_parallel_size(value: int | None, fallback: int) -> int: + return value if value is not None else fallback + + +def _set_if_known(section: dict[str, Any], key: str, value: Any) -> None: + if value is not None: + section[key] = value + + +def _apply_point_runtime_signature(raw_config: dict[str, Any], point: BenchmarkBehaviorPoint) -> None: + model = _section(raw_config, "model") + train = _section(raw_config, "train") + simulator = raw_config.setdefault("simulator", {}) + if point.model_ref is not None: + for key in ("config_path", "model_path", "model_name"): + if key in model: + model[key] = point.model_ref + model.setdefault("model_path", point.model_ref) + _set_if_known(model, "deepep_async_combine", point.deepep_async_combine) + _set_if_known(model, "deepep_num_sms", point.deepep_num_sms) + _set_if_known(model, "deepep_buffer_size_gb", point.deepep_buffer_size_gb) + _set_if_known(train, "enable_compile", point.enable_compile) + _set_if_known(train, "gradient_checkpointing_method", point.gradient_checkpointing_method) + _set_if_known(train, "enable_activation_offload", point.enable_activation_offload) + _set_if_known(train, "activation_offload_prefetch_count", point.activation_offload_prefetch_count) + _set_if_known(train, "skip_param_upcast", point.skip_param_upcast) + _set_if_known(train, "fsdp_reduce_dtype", point.fsdp_reduce_dtype) + _set_if_known(train, "ce_mode", point.ce_mode) + _set_if_known(model, "moe_implementation", point.moe_implementation) + _set_if_known(train, "moe_checkpoint_method", point.moe_checkpoint_method) + _set_if_known(train, "muon_momentum", point.muon_momentum) + _set_if_known(train, "muon_update_dtype", point.muon_update_dtype) + if isinstance(simulator, dict): + _set_if_known(simulator, "balanced_routing", point.balanced_routing) + _set_if_known(simulator, "attention_backend", point.attention_backend) + + +def _topology_for_point( + base_config: dict[str, Any], + base_topology: Topology, + point: BenchmarkBehaviorPoint, + *, + world_size: int | None, + local_world_size: int | None, + require_tokens: bool = True, +) -> tuple[dict[str, Any] | None, Topology | None, str | None]: + if point.micro_batch_size is None or point.global_batch_size is None: + return None, None, "missing micro_batch_size/global_batch_size" + if require_tokens and point.tokens_per_sec is None: + return None, None, "missing tokens_per_sec" + + resolved_world_size = point.gpu_count or world_size or base_topology.world_size + resolved_local_world_size = local_world_size or base_topology.local_world_size + tensor_parallel = _point_parallel_size(point.tensor_parallel_size, base_topology.tensor_parallel_size) + pipeline_parallel = _point_parallel_size(point.pipeline_parallel_size, base_topology.pipeline_parallel_size) + ulysses_parallel = _point_parallel_size(point.ulysses_parallel_size, base_topology.ulysses_parallel_size) + ringattn_parallel = _point_parallel_size(point.ringattn_parallel_size, base_topology.ringattn_parallel_size) + expert_parallel = _point_parallel_size(point.expert_parallel_size, base_topology.expert_parallel_size) + non_dp = tensor_parallel * pipeline_parallel * ulysses_parallel * ringattn_parallel + if non_dp <= 0 or resolved_world_size % non_dp: + return None, None, "world_size is not divisible by heldout non-DP topology" + data_parallel_size = resolved_world_size // non_dp + denominator = point.micro_batch_size * data_parallel_size + if denominator <= 0: + return None, None, "micro_batch_size * data_parallel_size must be positive" + if point.gradient_accumulation_steps is not None: + gradient_accumulation_steps = point.gradient_accumulation_steps + expected_global_batch_size = denominator * gradient_accumulation_steps + if expected_global_batch_size != point.global_batch_size: + return None, None, "explicit gradient_accumulation_steps does not match global_batch_size" + else: + if point.global_batch_size % denominator: + return None, None, "global_batch_size is not divisible by micro_batch_size * data_parallel_size" + gradient_accumulation_steps = point.global_batch_size // denominator + + raw_config = _mutated_config( + base_config, + world_size=resolved_world_size, + micro_batch_size=point.micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + expert_parallel_size=expert_parallel, + tensor_parallel_size=tensor_parallel, + pipeline_parallel_size=pipeline_parallel, + ulysses_parallel_size=ulysses_parallel, + ringattn_parallel_size=ringattn_parallel, + data_parallel_replicate_size=point.data_parallel_replicate_size, + data_parallel_shard_size=point.data_parallel_shard_size, + ) + if point.sample_packing_sequence_len is not None: + _section(raw_config, "data")["sample_packing_sequence_len"] = point.sample_packing_sequence_len + _apply_point_runtime_signature(raw_config, point) + try: + topology = resolve_topology( + raw_config, + world_size=resolved_world_size, + local_world_size=resolved_local_world_size, + ) + except ValueError as exc: + return None, None, str(exc) + if point.ep_fsdp_size is not None and topology.ep_fsdp_size != point.ep_fsdp_size: + return None, None, "heldout ep_fsdp_size does not match resolved topology" + return raw_config, topology, None + + +def _without_point( + behavior_points: list[BenchmarkBehaviorPoint], + heldout: BenchmarkBehaviorPoint, +) -> list[BenchmarkBehaviorPoint]: + return [point for point in behavior_points if not (point.label == heldout.label and point.source == heldout.source)] + + +def _memory_prediction( + *, + prediction_peak_mem_gb: float | None, + prediction_status: str, + analytic_peak_floor_gb: float | None, + memory_peak_estimate: Any | None, + device_memory_limit_gb: float, + memory_safety_factor: float, +) -> tuple[float | None, str, str, str, float | None, float | None, str | None, list[str]]: + predicted_peak_mem_gb = analytic_peak_floor_gb + memory_basis = "analytic_floor" + memory_calibration_source = None + memory_calibration_notes: list[str] = [] + if prediction_peak_mem_gb is not None: + if prediction_status == "calibrated" and ( + analytic_peak_floor_gb is None or prediction_peak_mem_gb >= analytic_peak_floor_gb + ): + predicted_peak_mem_gb = prediction_peak_mem_gb + memory_basis = "calibrated_peak" + elif prediction_status != "calibrated" and memory_peak_estimate is not None: + predicted_peak_mem_gb = memory_peak_estimate.peak_gb + memory_basis = getattr(memory_peak_estimate, "basis", "calibrated_overhead_peak") + memory_calibration_source = memory_peak_estimate.source_label + memory_calibration_notes = memory_peak_estimate.notes + elif analytic_peak_floor_gb is None or prediction_peak_mem_gb >= analytic_peak_floor_gb: + predicted_peak_mem_gb = prediction_peak_mem_gb + memory_basis = "extrapolated_peak" + elif memory_peak_estimate is not None: + predicted_peak_mem_gb = memory_peak_estimate.peak_gb + memory_basis = getattr(memory_peak_estimate, "basis", "calibrated_overhead_peak") + memory_calibration_source = memory_peak_estimate.source_label + memory_calibration_notes = memory_peak_estimate.notes + + _, _, feasibility_status = _memory_factor( + predicted_peak_mem_gb, + memory_basis=memory_basis, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + memory_coverage_status, predicted_residual_gb, predicted_residual_fraction = _memory_coverage_for_candidate( + analytic_peak_floor_gb=analytic_peak_floor_gb, + estimated_peak_mem_gb=predicted_peak_mem_gb, + memory_basis=memory_basis, + ) + return ( + predicted_peak_mem_gb, + memory_basis, + feasibility_status, + memory_coverage_status, + predicted_residual_gb, + predicted_residual_fraction, + memory_calibration_source, + memory_calibration_notes, + ) + + +def _actual_memory_residual( + actual_peak_mem_gb: float | None, + analytic_peak_floor_gb: float | None, +) -> tuple[float | None, float | None]: + if actual_peak_mem_gb is None or analytic_peak_floor_gb is None or actual_peak_mem_gb <= 0: + return None, None + residual = max(actual_peak_mem_gb - analytic_peak_floor_gb, 0.0) + return round(residual, 3), round(residual / actual_peak_mem_gb, 3) + + +def _calibration_residual_memory_peak_estimate( + *, + training_points: list[BenchmarkBehaviorPoint], + base_config: dict[str, Any], + base_topology: Topology, + raw_config: dict[str, Any], + target_topology: Topology, + metadata: Any, + analytic_peak_floor_gb: float | None, + world_size: int | None, + local_world_size: int | None, +) -> _CalibrationMemoryPeakEstimate | None: + if analytic_peak_floor_gb is None: + return None + + estimates: list[tuple[tuple[float, float, float, float, float], _CalibrationMemoryPeakEstimate]] = [] + for point in training_points: + if point.peak_mem_gb is None or point.correctness_status == "oom": + continue + if behavior_point_model_mismatches(point, raw_config): + continue + reference_config, reference_topology, _ = _topology_for_point( + base_config, + base_topology, + point, + world_size=world_size, + local_world_size=local_world_size, + require_tokens=False, + ) + if reference_config is None or reference_topology is None: + continue + reference_memory = build_memory_ledger( + reference_config, + topology=reference_topology, + model_metadata=metadata, + ) + reference_floor = reference_memory.analytic_peak_floor_gb + if reference_floor is None or point.peak_mem_gb < reference_floor: + continue + residual_gb = point.peak_mem_gb - reference_floor + if residual_gb <= 0: + continue + sequence_distance = abs( + math.log( + (target_topology.sample_packing_sequence_len or 1) + / (reference_topology.sample_packing_sequence_len or 1) + ) + ) + parallel_distance = sum( + abs(math.log2(max(target_value, 1) / max(reference_value, 1))) + for target_value, reference_value in ( + (target_topology.expert_parallel_size, reference_topology.expert_parallel_size), + (target_topology.ep_fsdp_size or 1, reference_topology.ep_fsdp_size or 1), + (target_topology.tensor_parallel_size, reference_topology.tensor_parallel_size), + (target_topology.pipeline_parallel_size, reference_topology.pipeline_parallel_size), + (target_topology.sequence_parallel_size, reference_topology.sequence_parallel_size), + ) + ) + batch_distance = abs( + math.log2(max(target_topology.micro_batch_size, 1) / max(reference_topology.micro_batch_size, 1)) + ) + workload_mismatches = behavior_point_workload_mismatches(point, raw_config) + notes = [ + "calibration_residual_prior=minimum_same_model_measured_residual", + f"memory_residual_reference={point.label}", + f"reference_peak_gb={point.peak_mem_gb:.3f}", + f"reference_floor_gb={reference_floor:.3f}", + f"reference_residual_gb={residual_gb:.3f}", + f"estimated_residual_gb={residual_gb:.3f}", + f"sequence_distance_log={sequence_distance:.3f}", + f"parallel_distance_log2={parallel_distance:.3f}", + f"batch_distance_log2={batch_distance:.3f}", + ] + if workload_mismatches: + notes.append("reference_workload_mismatches=" + ",".join(workload_mismatches)) + estimates.append( + ( + ( + residual_gb, + float(len(workload_mismatches)), + sequence_distance, + parallel_distance, + batch_distance, + ), + _CalibrationMemoryPeakEstimate( + peak_gb=round(analytic_peak_floor_gb + residual_gb, 3), + overhead_gb=round(residual_gb, 3), + source_label=point.label, + notes=notes, + ), + ) + ) + + if not estimates: + return None + return min(estimates, key=lambda item: item[0])[1] + + +def _phase_bottleneck_details(phase_time_share: dict[str, float]) -> tuple[str | None, str | None, float | None]: + visible = {phase: share for phase, share in phase_time_share.items() if phase != "train_step_total"} + if not visible: + return None, None, None + phase = max(visible, key=visible.get) + share = visible[phase] + return phase, _phase_bucket(phase), round(share, 6) + + +def _memory_bottleneck_details( + phase_memory_peak_gb: dict[str, float], + peak_mem_gb: float | None, +) -> tuple[str | None, str | None, float | None, float | None]: + visible = {phase: peak for phase, peak in phase_memory_peak_gb.items() if peak > 0} + if not visible: + return None, None, None, None + phase, peak = max(visible.items(), key=lambda item: (item[1], item[0])) + denominator = peak_mem_gb if peak_mem_gb is not None and peak_mem_gb > 0 else peak + return phase, _phase_bucket(phase), round(peak, 3), round(peak / denominator, 3) + + +def _parallel_size_for_attribution(point: BenchmarkBehaviorPoint, field_name: str, topology: Topology) -> int: + value = getattr(point, field_name) + if value is not None: + return int(value) + return int(getattr(topology, field_name)) + + +def _memory_attribution_distance(point: BenchmarkBehaviorPoint, topology: Topology) -> tuple[float, float, float]: + point_sequence = point.sample_packing_sequence_len or topology.sample_packing_sequence_len or 1 + target_sequence = topology.sample_packing_sequence_len or point_sequence + sequence_distance = abs(math.log(max(target_sequence, 1) / max(point_sequence, 1))) + parallel_pairs = ( + (topology.expert_parallel_size, _parallel_size_for_attribution(point, "expert_parallel_size", topology)), + (topology.ep_fsdp_size or 1, point.ep_fsdp_size or topology.ep_fsdp_size or 1), + (topology.tensor_parallel_size, _parallel_size_for_attribution(point, "tensor_parallel_size", topology)), + (topology.pipeline_parallel_size, _parallel_size_for_attribution(point, "pipeline_parallel_size", topology)), + (topology.ulysses_parallel_size, _parallel_size_for_attribution(point, "ulysses_parallel_size", topology)), + (topology.ringattn_parallel_size, _parallel_size_for_attribution(point, "ringattn_parallel_size", topology)), + ) + parallel_distance = sum( + abs(math.log2(max(target_value, 1) / max(reference_value, 1))) + for target_value, reference_value in parallel_pairs + ) + batch_distance = abs(math.log2(max(topology.micro_batch_size, 1) / max(point.micro_batch_size or 1, 1))) + return sequence_distance, parallel_distance, batch_distance + + +def _select_memory_attribution_point( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + raw_config: dict[str, Any], +) -> BenchmarkBehaviorPoint | None: + usable = [ + point + for point in behavior_points + if point.phase_memory_peak_gb + and point.peak_mem_gb is not None + and point.peak_mem_gb > 0 + and point.correctness_status != "oom" + and not behavior_point_model_mismatches(point, raw_config) + ] + if not usable: + return None + + def key(point: BenchmarkBehaviorPoint) -> tuple[int, float, float, float, int, float]: + runtime_mismatch_count = len(behavior_point_workload_mismatches(point, raw_config)) + sequence_distance, parallel_distance, batch_distance = _memory_attribution_distance(point, topology) + return ( + -runtime_mismatch_count, + -sequence_distance, + -parallel_distance, + -batch_distance, + 1 if point.tokens_per_sec is not None else 0, + point.peak_mem_gb or 0.0, + ) + + return max(usable, key=key) + + +def _select_phase_timing_attribution_point( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + raw_config: dict[str, Any], +) -> BenchmarkBehaviorPoint | None: + usable = [ + point + for point in behavior_points + if point.phase_time_share + and point.correctness_status != "oom" + and not behavior_point_model_mismatches(point, raw_config) + ] + if not usable: + return None + + def key(point: BenchmarkBehaviorPoint) -> tuple[int, float, float, float, int, int]: + runtime_mismatch_count = len(behavior_point_workload_mismatches(point, raw_config)) + sequence_distance, parallel_distance, batch_distance = _memory_attribution_distance(point, topology) + return ( + -runtime_mismatch_count, + -sequence_distance, + -parallel_distance, + -batch_distance, + 1 if point.tokens_per_sec is not None else 0, + point.measured_steps or 0, + ) + + return max(usable, key=key) + + +def _scaled_phase_memory_peak_gb( + point: BenchmarkBehaviorPoint, + predicted_peak_mem_gb: float | None, +) -> dict[str, float]: + scale = 1.0 + if predicted_peak_mem_gb is not None and point.peak_mem_gb is not None and point.peak_mem_gb > 0: + scale = max(predicted_peak_mem_gb, 0.0) / point.peak_mem_gb + return {phase: round(max(0.0, peak * scale), 6) for phase, peak in point.phase_memory_peak_gb.items()} + + +def _phase_top_items(phase_time_share: dict[str, float], *, limit: int = 3) -> list[tuple[str, str, float]]: + visible = [(phase, share) for phase, share in phase_time_share.items() if phase != "train_step_total"] + ordered = sorted(visible, key=lambda item: (-item[1], item[0])) + return [(phase, _phase_bucket(phase), round(share, 6)) for phase, share in ordered[:limit]] + + +def _ordered_unique(values: list[str]) -> list[str]: + result: list[str] = [] + seen: set[str] = set() + for value in values: + if value not in seen: + seen.add(value) + result.append(value) + return result + + +def _top_overlap(actual: list[str], predicted: list[str]) -> tuple[int | None, float | None]: + actual_unique = _ordered_unique(actual) + if not actual_unique: + return None, None + predicted_unique = set(predicted) + overlap = sum(1 for value in actual_unique if value in predicted_unique) + return overlap, round(overlap / len(actual_unique), 3) + + +def _max_holdout_by_field(holdouts: list[CalibrationHoldout], field_name: str) -> CalibrationHoldout | None: + return max( + (holdout for holdout in holdouts if getattr(holdout, field_name) is not None), + key=lambda holdout: (getattr(holdout, field_name), holdout.label), + default=None, + ) + + +def _min_holdout_by_field(holdouts: list[CalibrationHoldout], field_name: str) -> CalibrationHoldout | None: + return min( + (holdout for holdout in holdouts if getattr(holdout, field_name) is not None), + key=lambda holdout: (getattr(holdout, field_name), holdout.label), + default=None, + ) + + +def _memory_error_exceeds_target(holdout: CalibrationHoldout) -> bool: + if holdout.memory_absolute_error_gb is None or holdout.actual_peak_mem_gb is None: + return False + threshold = max(_MEMORY_ABS_ERROR_TARGET_GB, holdout.actual_peak_mem_gb * _MEMORY_RELATIVE_ERROR_TARGET_FRACTION) + return holdout.memory_absolute_error_gb > threshold + + +def _empirical_required_uncertainty_fraction(holdout: CalibrationHoldout) -> float | None: + if holdout.predicted_tokens_per_sec is None or holdout.predicted_tokens_per_sec <= 0: + return None + if holdout.actual_tokens_per_sec is None: + return None + return abs(holdout.actual_tokens_per_sec - holdout.predicted_tokens_per_sec) / holdout.predicted_tokens_per_sec + + +def _eligible_calibration_point(point: BenchmarkBehaviorPoint) -> bool: + if point.tokens_per_sec is not None: + return True + if point.correctness_status in {"oom", "runtime_failure_after_steps"}: + return False + return point.status == "observed_log_metrics_only" and ( + point.peak_mem_gb is not None or bool(point.phase_memory_peak_gb) or bool(point.phase_time_share) + ) + + +def _load_calibration_support_points( + support_benchmark_dirs: list[str | Path] | tuple[str | Path, ...], +) -> list[BenchmarkBehaviorPoint]: + points: list[BenchmarkBehaviorPoint] = [] + for support_dir in support_benchmark_dirs: + points.extend(load_benchmark_behavior_points(support_dir)) + return points + + +def _prediction_uncertainty_calibration_status( + *, + errors: list[float], + interval_covered_count: int, + uncertainty_fractions: list[float], + empirical_required_uncertainty_fractions: list[float], +) -> str: + if not errors or not empirical_required_uncertainty_fractions: + return "no_scored_uncertainty_holdouts" + if interval_covered_count < len(errors): + return "prediction_interval_undercoverage" + max_required = max(empirical_required_uncertainty_fractions) + if max_required <= _THROUGHPUT_MAPE_TARGET_PERCENT / 100.0: + if uncertainty_fractions and max(uncertainty_fractions) >= 0.50: + return "conservative_uncertainty_empirically_supported" + return "prediction_uncertainty_empirically_supported" + return "empirical_uncertainty_exceeds_target" + + +def _calibration_fidelity_support( + holdouts: list[CalibrationHoldout], + *, + errors: list[float], + interval_covered_count: int, + uncertainty_fractions: list[float], + empirical_required_uncertainty_fractions: list[float], + memory_percentage_errors: list[float], + memory_bottleneck_evaluated_count: int, + memory_bottleneck_bucket_match_count: int, + phase_bottleneck_evaluated_count: int, + phase_bottleneck_bucket_match_count: int, + phase_top3_evaluated_count: int, + phase_bucket_top3_overlap_rates: list[float], + phase_bottleneck_share_errors: list[float], +) -> tuple[str, list[str]]: + if not holdouts: + return "no_calibration_holdouts", ["no_evaluated_holdouts"] + if not errors: + return "no_scored_calibration_holdouts", ["no_scored_holdouts"] + + blockers: set[str] = set() + throughput_holdout_count = sum(1 for holdout in holdouts if holdout.actual_tokens_per_sec is not None) + if len(errors) < throughput_holdout_count: + blockers.add("unscored_holdouts") + if max(errors) > _THROUGHPUT_MAPE_TARGET_PERCENT: + blockers.add("max_throughput_mape_exceeds_8_percent") + if interval_covered_count < len(errors): + blockers.add("prediction_interval_misses") + high_uncertainty_empirically_supported = ( + max(errors) <= _THROUGHPUT_MAPE_TARGET_PERCENT + and interval_covered_count == len(errors) + and bool(empirical_required_uncertainty_fractions) + and max(empirical_required_uncertainty_fractions) <= _THROUGHPUT_MAPE_TARGET_PERCENT / 100.0 + ) + if uncertainty_fractions and max(uncertainty_fractions) >= 0.50 and not high_uncertainty_empirically_supported: + blockers.add("high_prediction_uncertainty") + + if not memory_percentage_errors: + blockers.add("missing_memory_holdouts") + elif any(_memory_error_exceeds_target(holdout) for holdout in holdouts): + blockers.add("memory_error_exceeds_3_percent_or_2gb") + missing_memory_bottleneck_prediction_count = sum( + 1 + for holdout in holdouts + if holdout.actual_memory_bottleneck_bucket is not None and holdout.predicted_memory_bottleneck_bucket is None + ) + if missing_memory_bottleneck_prediction_count > 0: + blockers.add("missing_memory_bottleneck_predictions") + elif ( + memory_bottleneck_evaluated_count > 0 + and memory_bottleneck_bucket_match_count < memory_bottleneck_evaluated_count + ): + blockers.add("memory_bottleneck_bucket_mismatch") + + if phase_bottleneck_evaluated_count == 0: + blockers.add("missing_phase_bottleneck_holdouts") + elif phase_bottleneck_bucket_match_count < phase_bottleneck_evaluated_count: + blockers.add("phase_bottleneck_bucket_mismatch") + if phase_top3_evaluated_count > 0 and phase_bucket_top3_overlap_rates: + if min(phase_bucket_top3_overlap_rates) < 1.0: + blockers.add("phase_bucket_top3_mismatch") + if phase_bottleneck_share_errors and max(phase_bottleneck_share_errors) > _PHASE_SHARE_ABS_ERROR_TARGET: + blockers.add("phase_bottleneck_share_error_exceeds_10_percent") + + blocker_list = sorted(blockers) + if "prediction_interval_misses" in blockers: + return "calibration_interval_coverage_failed", blocker_list + if "max_throughput_mape_exceeds_8_percent" in blockers: + return "calibration_error_attribution_needed", blocker_list + if "memory_error_exceeds_3_percent_or_2gb" in blockers: + return "calibration_memory_error_exceeds_target", blocker_list + if any(blocker.startswith("memory_bottleneck_") for blocker in blockers): + return "calibration_memory_attribution_mismatch", blocker_list + if any(blocker.startswith("phase_") for blocker in blockers): + return "calibration_phase_attribution_mismatch", blocker_list + missing_blockers = { + "missing_memory_holdouts", + "missing_memory_bottleneck_predictions", + "missing_phase_bottleneck_holdouts", + "unscored_holdouts", + "high_prediction_uncertainty", + } + if blockers and blockers <= missing_blockers: + return "partial_calibration_fidelity_missing_attribution", blocker_list + if blockers: + return "partial_calibration_fidelity", blocker_list + return "calibration_fidelity_supported", [] + + +def _calibration_gap_status_counts(gaps: list[CalibrationValidationGap]) -> dict[str, int]: + counts: dict[str, int] = {} + for gap in gaps: + counts[gap.gap_status] = counts.get(gap.gap_status, 0) + 1 + return dict(sorted(counts.items())) + + +def _unique_required_measurements(gaps: list[CalibrationValidationGap]) -> list[str]: + seen: set[str] = set() + required: list[str] = [] + for gap in gaps: + if gap.required_measurement in seen: + continue + seen.add(gap.required_measurement) + required.append(gap.required_measurement) + return required + + +def _labels(holdouts: list[CalibrationHoldout]) -> list[str]: + return [holdout.label for holdout in holdouts] + + +def _max_percentage_error(holdouts: list[CalibrationHoldout]) -> tuple[float | None, str | None]: + holdout = _max_holdout_by_field(holdouts, "absolute_percentage_error") + if holdout is None: + return None, None + return holdout.absolute_percentage_error, holdout.label + + +def _max_memory_error(holdouts: list[CalibrationHoldout]) -> tuple[float | None, str | None]: + holdout = _max_holdout_by_field(holdouts, "memory_absolute_error_gb") + if holdout is None: + return None, None + return holdout.memory_absolute_error_gb, holdout.label + + +def _max_phase_share_error(holdouts: list[CalibrationHoldout]) -> tuple[float | None, str | None]: + holdout = _max_holdout_by_field(holdouts, "phase_bottleneck_share_absolute_error") + if holdout is None: + return None, None + return holdout.phase_bottleneck_share_absolute_error, holdout.label + + +def _calibration_validation_gap( + *, + gap_status: str, + priority: int, + required_measurement: str, + reason: str, + affected_holdouts: list[CalibrationHoldout], + blocker_names: list[str], +) -> CalibrationValidationGap: + max_percentage_error, max_percentage_error_label = _max_percentage_error(affected_holdouts) + max_memory_error, max_memory_error_label = _max_memory_error(affected_holdouts) + max_phase_share_error, max_phase_share_error_label = _max_phase_share_error(affected_holdouts) + return CalibrationValidationGap( + gap_status=gap_status, + priority=priority, + required_measurement=required_measurement, + reason=reason, + affected_holdout_count=len(affected_holdouts), + affected_holdout_labels=_labels(affected_holdouts), + blocker_names=sorted(set(blocker_names)), + max_absolute_percentage_error=(round(max_percentage_error, 3) if max_percentage_error is not None else None), + max_absolute_percentage_error_label=max_percentage_error_label, + max_memory_absolute_error_gb=round(max_memory_error, 3) if max_memory_error is not None else None, + max_memory_absolute_error_label=max_memory_error_label, + max_phase_bottleneck_share_absolute_error=( + round(max_phase_share_error, 6) if max_phase_share_error is not None else None + ), + max_phase_bottleneck_share_absolute_error_label=max_phase_share_error_label, + missing_memory_count=sum(1 for holdout in affected_holdouts if holdout.memory_absolute_error_gb is None), + missing_phase_bottleneck_count=sum( + 1 for holdout in affected_holdouts if holdout.phase_bottleneck_bucket_match is None + ), + ) + + +def _calibration_validation_gap_portfolio( + holdouts: list[CalibrationHoldout], + blockers: list[str], +) -> list[CalibrationValidationGap]: + blocker_set = set(blockers) + gaps: list[CalibrationValidationGap] = [] + if "no_evaluated_holdouts" in blocker_set: + gaps.append( + _calibration_validation_gap( + gap_status="no_calibration_holdouts_need_replay", + priority=130, + required_measurement="add_calibration_holdouts", + reason="no measured holdouts were available for calibration validation", + affected_holdouts=[], + blocker_names=blockers, + ) + ) + if "no_scored_holdouts" in blocker_set: + gaps.append( + _calibration_validation_gap( + gap_status="no_scored_holdouts_need_supported_replay", + priority=125, + required_measurement="add_scored_same_context_calibration_holdouts", + reason="measured holdouts exist but none received a scored simulator prediction", + affected_holdouts=holdouts, + blocker_names=blockers, + ) + ) + if "max_throughput_mape_exceeds_8_percent" in blocker_set: + affected = [ + holdout + for holdout in holdouts + if holdout.absolute_percentage_error is not None + and holdout.absolute_percentage_error > _THROUGHPUT_MAPE_TARGET_PERCENT + ] + gaps.append( + _calibration_validation_gap( + gap_status="throughput_error_needs_attribution", + priority=120, + required_measurement="replay_high_error_holdouts_with_component_timing", + reason="at least one scored holdout exceeds the throughput MAPE target", + affected_holdouts=affected, + blocker_names=["max_throughput_mape_exceeds_8_percent"], + ) + ) + if "prediction_interval_misses" in blocker_set: + affected = [holdout for holdout in holdouts if holdout.actual_tokens_in_prediction_interval is False] + gaps.append( + _calibration_validation_gap( + gap_status="prediction_interval_miss_needs_uncertainty_recalibration", + priority=115, + required_measurement="add_holdouts_to_recalibrate_prediction_interval", + reason="one or more actual holdout throughputs fall outside the predicted interval", + affected_holdouts=affected, + blocker_names=["prediction_interval_misses"], + ) + ) + if "memory_error_exceeds_3_percent_or_2gb" in blocker_set: + affected = [holdout for holdout in holdouts if _memory_error_exceeds_target(holdout)] + gaps.append( + _calibration_validation_gap( + gap_status="memory_error_needs_residual_attribution", + priority=110, + required_measurement="replay_high_memory_error_holdouts_with_memory_profile", + reason="one or more memory predictions exceed the absolute or relative memory error target", + affected_holdouts=affected, + blocker_names=["memory_error_exceeds_3_percent_or_2gb"], + ) + ) + if "memory_bottleneck_bucket_mismatch" in blocker_set: + affected = [holdout for holdout in holdouts if holdout.memory_bottleneck_bucket_match is False] + gaps.append( + _calibration_validation_gap( + gap_status="memory_bottleneck_mismatch_needs_phase_memory_replay", + priority=105, + required_measurement="replay_memory_bottleneck_holdouts_with_phase_memory_profile", + reason="predicted memory bottleneck bucket does not match observed phase memory attribution", + affected_holdouts=affected, + blocker_names=["memory_bottleneck_bucket_mismatch"], + ) + ) + if "missing_memory_bottleneck_predictions" in blocker_set: + affected = [ + holdout + for holdout in holdouts + if holdout.actual_memory_bottleneck_bucket is not None + and holdout.predicted_memory_bottleneck_bucket is None + ] + gaps.append( + _calibration_validation_gap( + gap_status="missing_memory_bottleneck_predictions_need_phase_memory_replay", + priority=105, + required_measurement="replay_memory_bottleneck_holdouts_with_phase_memory_profile", + reason="observed phase memory attribution exists but the prediction has no memory bottleneck bucket", + affected_holdouts=affected, + blocker_names=["missing_memory_bottleneck_predictions"], + ) + ) + if "missing_memory_holdouts" in blocker_set: + affected = [holdout for holdout in holdouts if holdout.memory_absolute_error_gb is None] + gaps.append( + _calibration_validation_gap( + gap_status="missing_memory_holdouts_need_memory_profile", + priority=100, + required_measurement="collect_peak_memory_for_calibration_holdouts", + reason="calibration fidelity cannot validate memory without observed peak memory holdouts", + affected_holdouts=affected, + blocker_names=["missing_memory_holdouts"], + ) + ) + if "phase_bottleneck_bucket_mismatch" in blocker_set: + affected = [holdout for holdout in holdouts if holdout.phase_bottleneck_bucket_match is False] + gaps.append( + _calibration_validation_gap( + gap_status="phase_bottleneck_mismatch_needs_component_timing_replay", + priority=95, + required_measurement="replay_phase_bottleneck_holdouts_with_component_timing", + reason="predicted phase bottleneck bucket does not match observed phase timing attribution", + affected_holdouts=affected, + blocker_names=["phase_bottleneck_bucket_mismatch"], + ) + ) + if "phase_bucket_top3_mismatch" in blocker_set: + affected = [ + holdout + for holdout in holdouts + if holdout.phase_bucket_top3_overlap_rate is not None and holdout.phase_bucket_top3_overlap_rate < 1.0 + ] + gaps.append( + _calibration_validation_gap( + gap_status="phase_top3_mismatch_needs_component_timing_replay", + priority=90, + required_measurement="replay_phase_top3_holdouts_with_component_timing", + reason="predicted phase bucket top-3 attribution misses observed timing buckets", + affected_holdouts=affected, + blocker_names=["phase_bucket_top3_mismatch"], + ) + ) + if "phase_bottleneck_share_error_exceeds_10_percent" in blocker_set: + affected = [ + holdout + for holdout in holdouts + if holdout.phase_bottleneck_share_absolute_error is not None + and holdout.phase_bottleneck_share_absolute_error > _PHASE_SHARE_ABS_ERROR_TARGET + ] + gaps.append( + _calibration_validation_gap( + gap_status="phase_share_error_needs_component_timing_replay", + priority=85, + required_measurement="replay_phase_share_holdouts_with_component_timing", + reason="predicted phase bottleneck share exceeds the attribution error target", + affected_holdouts=affected, + blocker_names=["phase_bottleneck_share_error_exceeds_10_percent"], + ) + ) + if "missing_phase_bottleneck_holdouts" in blocker_set: + affected = [holdout for holdout in holdouts if holdout.phase_bottleneck_bucket_match is None] + gaps.append( + _calibration_validation_gap( + gap_status="missing_phase_bottleneck_holdouts_need_phase_timing", + priority=80, + required_measurement="collect_phase_timing_for_calibration_holdouts", + reason="calibration fidelity cannot validate phase bottlenecks without observed phase timing holdouts", + affected_holdouts=affected, + blocker_names=["missing_phase_bottleneck_holdouts"], + ) + ) + if "high_prediction_uncertainty" in blocker_set: + affected = [ + holdout + for holdout in holdouts + if holdout.prediction_uncertainty_fraction is not None and holdout.prediction_uncertainty_fraction >= 0.50 + ] + gaps.append( + _calibration_validation_gap( + gap_status="high_prediction_uncertainty_needs_nearby_holdouts", + priority=70, + required_measurement="add_nearby_calibration_holdouts_to_tighten_uncertainty", + reason="predicted uncertainty remains high and is not empirically justified by observed errors", + affected_holdouts=affected, + blocker_names=["high_prediction_uncertainty"], + ) + ) + if "unscored_holdouts" in blocker_set: + affected = [ + holdout + for holdout in holdouts + if holdout.actual_tokens_per_sec is not None and holdout.absolute_percentage_error is None + ] + gaps.append( + _calibration_validation_gap( + gap_status="unscored_holdouts_need_supported_replay", + priority=60, + required_measurement="add_supported_same_context_calibration_holdouts", + reason="some measured holdouts could not be scored by the simulator", + affected_holdouts=affected, + blocker_names=["unscored_holdouts"], + ) + ) + return sorted(gaps, key=lambda gap: (-gap.priority, gap.gap_status)) + + +def _append_calibration_design_config( + rendered: list[ScenarioMeasurementConfig], + seen: set[tuple[str, str]], + design: ScenarioMeasurementConfig | None, +) -> bool: + if design is None: + return False + required_measurement = design.label.split(":", 3)[1] if ":" in design.label else design.label + key = (required_measurement, yaml.safe_dump(design.config, sort_keys=True)) + if key in seen: + return False + seen.add(key) + rendered.append(design) + return True + + +def _calibration_design_config_from_point( + *, + base_config: dict[str, Any], + base_topology: Topology, + point: BenchmarkBehaviorPoint, + required_measurement: str, + design_kind: str, + index: int, + world_size: int | None, + local_world_size: int | None, + config_overrides: tuple[str, ...], +) -> ScenarioMeasurementConfig | None: + raw_config, topology, _ = _topology_for_point( + base_config, + base_topology, + point, + world_size=world_size, + local_world_size=local_world_size, + require_tokens=False, + ) + if raw_config is None or topology is None: + return None + for override in config_overrides: + _apply_config_override(raw_config, override) + label = f"design:{required_measurement}:{design_kind}_{index:02d}:{point.label}:{_topology_label(topology)}" + return ScenarioMeasurementConfig( + label=label, + filename=_measurement_config_filename(index, label), + config=raw_config, + ) + + +def _calibration_nearby_design_config_from_topology( + *, + base_config: dict[str, Any], + source_point: BenchmarkBehaviorPoint, + source_topology: Topology, + required_measurement: str, + design_kind: str, + variant_label: str, + workload_values: dict[str, int], + index: int, + config_overrides: tuple[str, ...], +) -> ScenarioMeasurementConfig | None: + try: + raw_config = _mutated_config( + base_config, + world_size=source_topology.world_size, + micro_batch_size=workload_values["micro_batch_size"], + gradient_accumulation_steps=workload_values["gradient_accumulation_steps"], + expert_parallel_size=source_topology.expert_parallel_size, + tensor_parallel_size=source_topology.tensor_parallel_size, + pipeline_parallel_size=source_topology.pipeline_parallel_size, + ulysses_parallel_size=source_topology.ulysses_parallel_size, + ringattn_parallel_size=source_topology.ringattn_parallel_size, + data_parallel_replicate_size=source_topology.data_parallel_replicate_size, + data_parallel_shard_size=source_topology.data_parallel_shard_size, + ) + _section(raw_config, "data")["sample_packing_sequence_len"] = workload_values["sample_packing_sequence_len"] + _apply_point_runtime_signature(raw_config, source_point) + for override in config_overrides: + _apply_config_override(raw_config, override) + topology = resolve_topology( + raw_config, + world_size=source_topology.world_size, + local_world_size=source_topology.local_world_size, + ) + except (KeyError, TypeError, ValueError): + return None + + label = ( + f"design:{required_measurement}:{design_kind}_{index:02d}_{variant_label}:" + f"{source_point.label}:{_topology_label(topology)}" + ) + return ScenarioMeasurementConfig( + label=label, + filename=_measurement_config_filename(index, label), + config=raw_config, + ) + + +def _materialize_calibration_measurement_design_configs_from_context( + *, + base_config: dict[str, Any], + base_topology: Topology, + behavior_points: list[BenchmarkBehaviorPoint], + calibration_validation_gaps: list[CalibrationValidationGap], + world_size: int | None, + local_world_size: int | None, + max_configs_per_measurement: int, +) -> list[ScenarioMeasurementConfig]: + points_by_label = {point.label: point for point in behavior_points} + rendered: list[ScenarioMeasurementConfig] = [] + seen: set[tuple[str, str]] = set() + index = 1 + for gap in sorted(calibration_validation_gaps, key=lambda item: (-item.priority, item.gap_status)): + count_for_measurement = 0 + + def add_limited(design: ScenarioMeasurementConfig | None) -> None: + nonlocal count_for_measurement, index + if count_for_measurement >= max_configs_per_measurement: + return + if _append_calibration_design_config(rendered, seen, design): + index += 1 + count_for_measurement += 1 + + overrides = _CALIBRATION_REPLAY_OVERRIDES_BY_MEASUREMENT.get(gap.required_measurement) + if overrides is not None: + design_kind = _CALIBRATION_REPLAY_KIND_BY_MEASUREMENT[gap.required_measurement] + for label in gap.affected_holdout_labels: + if count_for_measurement >= max_configs_per_measurement: + break + point = points_by_label.get(label) + if point is None: + continue + add_limited( + _calibration_design_config_from_point( + base_config=base_config, + base_topology=base_topology, + point=point, + required_measurement=gap.required_measurement, + design_kind=design_kind, + index=index, + world_size=world_size, + local_world_size=local_world_size, + config_overrides=overrides, + ) + ) + continue + + design_kind = _CALIBRATION_HOLDOUT_REPLAY_KIND_BY_MEASUREMENT.get(gap.required_measurement) + if design_kind is not None: + for label in gap.affected_holdout_labels: + if count_for_measurement >= max_configs_per_measurement: + break + point = points_by_label.get(label) + if point is None: + continue + add_limited( + _calibration_design_config_from_point( + base_config=base_config, + base_topology=base_topology, + point=point, + required_measurement=gap.required_measurement, + design_kind=design_kind, + index=index, + world_size=world_size, + local_world_size=local_world_size, + config_overrides=(), + ) + ) + continue + + design_kind = _CALIBRATION_NEARBY_HOLDOUT_KIND_BY_MEASUREMENT.get(gap.required_measurement) + if design_kind is None: + continue + for label in gap.affected_holdout_labels: + if count_for_measurement >= max_configs_per_measurement: + break + point = points_by_label.get(label) + if point is None: + continue + _, source_topology, _ = _topology_for_point( + base_config, + base_topology, + point, + world_size=world_size, + local_world_size=local_world_size, + require_tokens=False, + ) + if source_topology is None: + continue + for variant_label, workload_values in _workload_design_variants(source_topology): + if count_for_measurement >= max_configs_per_measurement: + break + add_limited( + _calibration_nearby_design_config_from_topology( + base_config=base_config, + source_point=point, + source_topology=source_topology, + required_measurement=gap.required_measurement, + design_kind=design_kind, + variant_label=variant_label, + workload_values=workload_values, + index=index, + config_overrides=(), + ) + ) + return rendered + + +def materialize_calibration_measurement_design_configs( + report: CalibrationReport, + *, + max_configs_per_measurement: int = 4, +) -> list[ScenarioMeasurementConfig]: + """Render bounded YAML design rows for calibration replay/profile gaps.""" + base_config = load_training_config(report.base_config_path) + base_topology = resolve_topology(base_config) + behavior_points = [ + *load_benchmark_behavior_points(report.benchmark_dir), + *_load_calibration_support_points(tuple(report.calibration_support_benchmark_dirs)), + ] + return _materialize_calibration_measurement_design_configs_from_context( + base_config=base_config, + base_topology=base_topology, + behavior_points=behavior_points, + calibration_validation_gaps=report.calibration_validation_gaps, + world_size=None, + local_world_size=None, + max_configs_per_measurement=max_configs_per_measurement, + ) + + +def materialize_calibration_measurement_configs(report: CalibrationReport) -> list[ScenarioMeasurementConfig]: + """Render calibration design YAML payloads with their on-disk filenames.""" + return [ + ScenarioMeasurementConfig( + label=item.label, + filename=f"design_{item.filename}", + config=item.config, + ) + for item in materialize_calibration_measurement_design_configs(report) + ] + + +def write_measurement_configs(report: CalibrationReport, output_dir: str | Path) -> list[ScenarioMeasurementConfig]: + """Write calibration replay/profile design configs as YAML files.""" + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + rendered = materialize_calibration_measurement_configs(report) + for item in rendered: + (output_path / item.filename).write_text( + yaml.safe_dump(runtime_training_config(item.config), sort_keys=False), + encoding="utf-8", + ) + return rendered + + +def evaluate_calibration( + base_config_path: str | Path, + *, + benchmark_dir: str | Path, + calibration_support_benchmark_dirs: list[str | Path] | tuple[str | Path, ...] = (), + world_size: int | None = None, + local_world_size: int | None = None, + device_memory_limit_gb: float = 80.0, + memory_safety_factor: float = 1.15, +) -> CalibrationReport: + base_path = Path(base_config_path) + benchmark_path = resolve_calibration_pack(benchmark_dir) + base_config = load_training_config(base_path) + base_topology = resolve_topology(base_config, world_size=world_size, local_world_size=local_world_size) + behavior_points = load_benchmark_behavior_points(benchmark_path) + support_points = _load_calibration_support_points(calibration_support_benchmark_dirs) + prediction_points = [*behavior_points, *support_points] + measured_points = [ + point + for point in behavior_points + if _eligible_calibration_point(point) and not behavior_point_model_mismatches(point, base_config) + ] + + holdouts: list[CalibrationHoldout] = [] + warnings: list[str] = [] + skipped_count = 0 + for heldout in measured_points: + raw_config, topology, skip_reason = _topology_for_point( + base_config, + base_topology, + heldout, + world_size=world_size, + local_world_size=local_world_size, + require_tokens=False, + ) + if raw_config is None or topology is None: + skipped_count += 1 + warnings.append(f"skipped {heldout.label}: {skip_reason}") + continue + + training_points = _without_point(prediction_points, heldout) + shape = build_shape_ledger(topology, balanced_routing=True) + metadata = resolve_model_metadata(raw_config) + memory = build_memory_ledger(raw_config, topology=topology, model_metadata=metadata) + exact_prediction = predict_benchmark_behavior(training_points, topology, shape, raw_config) + if exact_prediction.status == "calibrated": + prediction = exact_prediction + memory_peak_estimate = None + else: + memory_peak_estimate = _calibrated_memory_peak_estimate( + training_points, + base_config, + raw_config, + topology, + shape, + metadata, + default_world_size=base_topology.world_size, + default_local_world_size=base_topology.local_world_size, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + ) + prediction, _ = _extrapolate_behavior( + training_points, + topology, + shape, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + ) + + if memory_peak_estimate is None and ( + prediction.peak_mem_gb is None + or (memory.analytic_peak_floor_gb is not None and prediction.peak_mem_gb <= memory.analytic_peak_floor_gb) + ): + memory_peak_estimate = _calibration_residual_memory_peak_estimate( + training_points=training_points, + base_config=base_config, + base_topology=base_topology, + raw_config=raw_config, + target_topology=topology, + metadata=metadata, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + world_size=world_size, + local_world_size=local_world_size, + ) + + predicted = prediction.tokens_per_sec + absolute_error = None + absolute_percentage_error = None + if predicted is not None and heldout.tokens_per_sec is not None: + absolute_error = abs(predicted - heldout.tokens_per_sec) + absolute_percentage_error = 100.0 * absolute_error / heldout.tokens_per_sec + ( + predicted_peak_mem_gb, + memory_basis, + memory_feasibility_status, + memory_coverage_status, + predicted_memory_residual_gb, + predicted_memory_residual_fraction, + memory_calibration_source, + memory_calibration_notes, + ) = _memory_prediction( + prediction_peak_mem_gb=prediction.peak_mem_gb, + prediction_status=prediction.status, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + actual_memory_residual_gb, actual_memory_residual_fraction = _actual_memory_residual( + heldout.peak_mem_gb, + memory.analytic_peak_floor_gb, + ) + calibration_scope = _calibration_scope( + training_points, + topology, + prediction_confidence=prediction.status, + raw_config=raw_config, + ) + calibration_distance, _ = _calibration_distance(training_points, topology, prediction) + risk_flags = _candidate_risk_flags( + training_points, + topology, + prediction, + raw_config=raw_config, + calibration_scope=calibration_scope, + prediction_confidence=prediction.status, + communication=None, + ) + if memory_basis in {"calibrated_overhead_peak", "calibration_residual_floor_peak"}: + risk_flags = sorted({*risk_flags, "memory_extrapolated_overhead"}) + prediction_uncertainty = _prediction_uncertainty_fraction( + prediction, + prediction_confidence=prediction.status, + calibration_scope=calibration_scope, + calibration_distance=calibration_distance, + risk_flags=risk_flags, + memory_coverage_status=memory_coverage_status, + ) + prediction_interval_lower, prediction_interval_upper = _prediction_interval(predicted, prediction_uncertainty) + actual_in_prediction_interval = ( + ( + prediction_interval_lower is not None + and prediction_interval_upper is not None + and prediction_interval_lower <= heldout.tokens_per_sec <= prediction_interval_upper + ) + if predicted is not None and heldout.tokens_per_sec is not None + else None + ) + memory_absolute_error_gb = None + memory_absolute_percentage_error = None + if heldout.peak_mem_gb is not None and predicted_peak_mem_gb is not None and heldout.peak_mem_gb > 0: + memory_absolute_error_gb = abs(predicted_peak_mem_gb - heldout.peak_mem_gb) + memory_absolute_percentage_error = 100.0 * memory_absolute_error_gb / heldout.peak_mem_gb + ( + actual_memory_phase, + actual_memory_bucket, + actual_memory_peak_gb, + actual_memory_fraction, + ) = _memory_bottleneck_details(heldout.phase_memory_peak_gb, heldout.peak_mem_gb) + predicted_phase_memory_peak_gb = prediction.phase_memory_peak_gb + memory_attribution_warnings: list[str] = [] + if not predicted_phase_memory_peak_gb: + memory_attribution_point = _select_memory_attribution_point(training_points, topology, raw_config) + if memory_attribution_point is not None: + predicted_phase_memory_peak_gb = _scaled_phase_memory_peak_gb( + memory_attribution_point, + predicted_peak_mem_gb, + ) + memory_attribution_warnings.append( + f"memory_bottleneck_attribution_source={memory_attribution_point.label}" + ) + ( + predicted_memory_phase, + predicted_memory_bucket, + predicted_memory_peak_gb, + predicted_memory_fraction, + ) = _memory_bottleneck_details(predicted_phase_memory_peak_gb, predicted_peak_mem_gb) + memory_bottleneck_phase_match = None + memory_bottleneck_bucket_match = None + if actual_memory_phase is not None and predicted_memory_phase is not None: + memory_bottleneck_phase_match = actual_memory_phase == predicted_memory_phase + if actual_memory_bucket is not None and predicted_memory_bucket is not None: + memory_bottleneck_bucket_match = actual_memory_bucket == predicted_memory_bucket + memory_bottleneck_peak_error = ( + round(abs(actual_memory_peak_gb - predicted_memory_peak_gb), 3) + if actual_memory_peak_gb is not None and predicted_memory_peak_gb is not None + else None + ) + memory_bottleneck_fraction_error = ( + round(abs(actual_memory_fraction - predicted_memory_fraction), 3) + if actual_memory_fraction is not None and predicted_memory_fraction is not None + else None + ) + actual_phase, actual_bucket, actual_share = _phase_bottleneck_details(heldout.phase_time_share) + predicted_phase_time_share = prediction.phase_time_share + phase_attribution_warnings: list[str] = [] + if not predicted_phase_time_share: + phase_attribution_point = _select_phase_timing_attribution_point(training_points, topology, raw_config) + if phase_attribution_point is not None: + predicted_phase_time_share = phase_attribution_point.phase_time_share + phase_attribution_warnings.append(f"phase_timing_attribution_source={phase_attribution_point.label}") + predicted_phase, predicted_bucket, predicted_share = _phase_bottleneck_details(predicted_phase_time_share) + phase_bottleneck_phase_match = None + phase_bottleneck_bucket_match = None + if actual_phase is not None: + phase_bottleneck_phase_match = actual_phase == predicted_phase + if actual_bucket is not None: + phase_bottleneck_bucket_match = actual_bucket == predicted_bucket + phase_bottleneck_share_error = ( + round(abs(actual_share - predicted_share), 6) + if actual_share is not None and predicted_share is not None + else None + ) + actual_top_items = _phase_top_items(heldout.phase_time_share) + predicted_top_items = _phase_top_items(predicted_phase_time_share) + actual_phase_top3 = [phase for phase, _, _ in actual_top_items] + predicted_phase_top3 = [phase for phase, _, _ in predicted_top_items] + actual_phase_bucket_top3 = _ordered_unique([bucket for _, bucket, _ in actual_top_items]) + predicted_phase_bucket_top3 = _ordered_unique([bucket for _, bucket, _ in predicted_top_items]) + phase_top3_overlap_count, phase_top3_overlap_rate = _top_overlap(actual_phase_top3, predicted_phase_top3) + phase_bucket_top3_overlap_count, phase_bucket_top3_overlap_rate = _top_overlap( + actual_phase_bucket_top3, + predicted_phase_bucket_top3, + ) + holdouts.append( + CalibrationHoldout( + label=heldout.label, + source=heldout.source, + topology_label=_topology_label(topology), + actual_tokens_per_sec=heldout.tokens_per_sec, + predicted_tokens_per_sec=predicted, + prediction_uncertainty_fraction=prediction_uncertainty, + prediction_interval_lower_tokens_per_sec=prediction_interval_lower, + prediction_interval_upper_tokens_per_sec=prediction_interval_upper, + actual_tokens_in_prediction_interval=actual_in_prediction_interval, + actual_peak_mem_gb=heldout.peak_mem_gb, + predicted_peak_mem_gb=round(predicted_peak_mem_gb, 3) if predicted_peak_mem_gb is not None else None, + prediction_status=prediction.status, + matched_label=prediction.matched_label, + absolute_error_tokens_per_sec=round(absolute_error, 3) if absolute_error is not None else None, + absolute_percentage_error=round(absolute_percentage_error, 3) + if absolute_percentage_error is not None + else None, + analytic_peak_floor_gb=( + round(memory.analytic_peak_floor_gb, 3) if memory.analytic_peak_floor_gb is not None else None + ), + memory_prediction_basis=memory_basis, + memory_coverage_status=memory_coverage_status, + memory_feasibility_status=memory_feasibility_status, + predicted_memory_residual_gb=predicted_memory_residual_gb, + predicted_memory_residual_fraction_of_peak=predicted_memory_residual_fraction, + actual_memory_residual_gb=actual_memory_residual_gb, + actual_memory_residual_fraction_of_peak=actual_memory_residual_fraction, + memory_absolute_error_gb=round(memory_absolute_error_gb, 3) + if memory_absolute_error_gb is not None + else None, + memory_absolute_percentage_error=round(memory_absolute_percentage_error, 3) + if memory_absolute_percentage_error is not None + else None, + actual_memory_bottleneck_phase=actual_memory_phase, + actual_memory_bottleneck_bucket=actual_memory_bucket, + actual_memory_bottleneck_peak_gb=actual_memory_peak_gb, + actual_memory_bottleneck_fraction_of_peak=actual_memory_fraction, + predicted_memory_bottleneck_phase=predicted_memory_phase, + predicted_memory_bottleneck_bucket=predicted_memory_bucket, + predicted_memory_bottleneck_peak_gb=predicted_memory_peak_gb, + predicted_memory_bottleneck_fraction_of_peak=predicted_memory_fraction, + memory_bottleneck_phase_match=memory_bottleneck_phase_match, + memory_bottleneck_bucket_match=memory_bottleneck_bucket_match, + memory_bottleneck_peak_absolute_error_gb=memory_bottleneck_peak_error, + memory_bottleneck_fraction_absolute_error=memory_bottleneck_fraction_error, + actual_phase_bottleneck_phase=actual_phase, + actual_phase_bottleneck_bucket=actual_bucket, + actual_phase_bottleneck_share=actual_share, + predicted_phase_bottleneck_phase=predicted_phase, + predicted_phase_bottleneck_bucket=predicted_bucket, + predicted_phase_bottleneck_share=predicted_share, + phase_bottleneck_phase_match=phase_bottleneck_phase_match, + phase_bottleneck_bucket_match=phase_bottleneck_bucket_match, + phase_bottleneck_share_absolute_error=phase_bottleneck_share_error, + actual_phase_top3=actual_phase_top3, + predicted_phase_top3=predicted_phase_top3, + actual_phase_bucket_top3=actual_phase_bucket_top3, + predicted_phase_bucket_top3=predicted_phase_bucket_top3, + phase_top3_overlap_count=phase_top3_overlap_count, + phase_top3_overlap_rate=phase_top3_overlap_rate, + phase_bucket_top3_overlap_count=phase_bucket_top3_overlap_count, + phase_bucket_top3_overlap_rate=phase_bucket_top3_overlap_rate, + memory_calibration_source=memory_calibration_source, + calibrated_from_count=len(training_points), + memory_calibration_notes=memory_calibration_notes, + warnings=[*prediction.warnings, *memory_attribution_warnings, *phase_attribution_warnings], + ) + ) + + errors = [ + holdout.absolute_percentage_error for holdout in holdouts if holdout.absolute_percentage_error is not None + ] + interval_covered_count = sum(1 for holdout in holdouts if holdout.actual_tokens_in_prediction_interval) + uncertainty_fractions = [ + holdout.prediction_uncertainty_fraction + for holdout in holdouts + if holdout.prediction_uncertainty_fraction is not None + ] + empirical_required_uncertainty_pairs = [ + (required, holdout) + for holdout in holdouts + if (required := _empirical_required_uncertainty_fraction(holdout)) is not None + ] + empirical_required_uncertainty_fractions = [required for required, _ in empirical_required_uncertainty_pairs] + status_counts: dict[str, int] = {} + memory_basis_counts: dict[str, int] = {} + memory_coverage_counts: dict[str, int] = {} + memory_feasibility_counts: dict[str, int] = {} + for holdout in holdouts: + status_counts[holdout.prediction_status] = status_counts.get(holdout.prediction_status, 0) + 1 + memory_basis_counts[holdout.memory_prediction_basis] = ( + memory_basis_counts.get(holdout.memory_prediction_basis, 0) + 1 + ) + memory_coverage_counts[holdout.memory_coverage_status] = ( + memory_coverage_counts.get(holdout.memory_coverage_status, 0) + 1 + ) + memory_feasibility_counts[holdout.memory_feasibility_status] = ( + memory_feasibility_counts.get(holdout.memory_feasibility_status, 0) + 1 + ) + memory_absolute_errors = [ + holdout.memory_absolute_error_gb for holdout in holdouts if holdout.memory_absolute_error_gb is not None + ] + memory_percentage_errors = [ + holdout.memory_absolute_percentage_error + for holdout in holdouts + if holdout.memory_absolute_percentage_error is not None + ] + predicted_memory_residuals = [ + holdout.predicted_memory_residual_gb for holdout in holdouts if holdout.predicted_memory_residual_gb is not None + ] + predicted_memory_residual_fractions = [ + holdout.predicted_memory_residual_fraction_of_peak + for holdout in holdouts + if holdout.predicted_memory_residual_fraction_of_peak is not None + ] + actual_memory_residuals = [ + holdout.actual_memory_residual_gb for holdout in holdouts if holdout.actual_memory_residual_gb is not None + ] + actual_memory_residual_fractions = [ + holdout.actual_memory_residual_fraction_of_peak + for holdout in holdouts + if holdout.actual_memory_residual_fraction_of_peak is not None + ] + memory_bottleneck_holdouts = [holdout for holdout in holdouts if holdout.memory_bottleneck_bucket_match is not None] + memory_bottleneck_phase_match_count = sum( + 1 for holdout in memory_bottleneck_holdouts if holdout.memory_bottleneck_phase_match + ) + memory_bottleneck_bucket_match_count = sum( + 1 for holdout in memory_bottleneck_holdouts if holdout.memory_bottleneck_bucket_match + ) + memory_bottleneck_peak_errors = [ + holdout.memory_bottleneck_peak_absolute_error_gb + for holdout in memory_bottleneck_holdouts + if holdout.memory_bottleneck_peak_absolute_error_gb is not None + ] + memory_bottleneck_fraction_errors = [ + holdout.memory_bottleneck_fraction_absolute_error + for holdout in memory_bottleneck_holdouts + if holdout.memory_bottleneck_fraction_absolute_error is not None + ] + phase_bottleneck_holdouts = [holdout for holdout in holdouts if holdout.phase_bottleneck_bucket_match is not None] + phase_bottleneck_phase_match_count = sum( + 1 for holdout in phase_bottleneck_holdouts if holdout.phase_bottleneck_phase_match + ) + phase_bottleneck_bucket_match_count = sum( + 1 for holdout in phase_bottleneck_holdouts if holdout.phase_bottleneck_bucket_match + ) + phase_bottleneck_share_errors = [ + holdout.phase_bottleneck_share_absolute_error + for holdout in phase_bottleneck_holdouts + if holdout.phase_bottleneck_share_absolute_error is not None + ] + phase_top3_holdouts = [holdout for holdout in holdouts if holdout.phase_top3_overlap_rate is not None] + phase_top3_overlap_rates = [holdout.phase_top3_overlap_rate for holdout in phase_top3_holdouts] + phase_bucket_top3_overlap_rates = [ + holdout.phase_bucket_top3_overlap_rate + for holdout in phase_top3_holdouts + if holdout.phase_bucket_top3_overlap_rate is not None + ] + max_error_holdout = _max_holdout_by_field(holdouts, "absolute_percentage_error") + max_uncertainty_holdout = _max_holdout_by_field(holdouts, "prediction_uncertainty_fraction") + max_empirical_required_uncertainty_pair = ( + max(empirical_required_uncertainty_pairs, key=lambda item: (item[0], item[1].label)) + if empirical_required_uncertainty_pairs + else None + ) + max_memory_error_holdout = _max_holdout_by_field(holdouts, "memory_absolute_error_gb") + max_memory_percentage_error_holdout = _max_holdout_by_field(holdouts, "memory_absolute_percentage_error") + max_predicted_memory_residual_holdout = _max_holdout_by_field(holdouts, "predicted_memory_residual_gb") + max_predicted_memory_residual_fraction_holdout = _max_holdout_by_field( + holdouts, + "predicted_memory_residual_fraction_of_peak", + ) + max_actual_memory_residual_holdout = _max_holdout_by_field(holdouts, "actual_memory_residual_gb") + max_actual_memory_residual_fraction_holdout = _max_holdout_by_field( + holdouts, + "actual_memory_residual_fraction_of_peak", + ) + max_memory_bottleneck_peak_error_holdout = _max_holdout_by_field( + holdouts, + "memory_bottleneck_peak_absolute_error_gb", + ) + max_memory_bottleneck_fraction_error_holdout = _max_holdout_by_field( + holdouts, + "memory_bottleneck_fraction_absolute_error", + ) + max_phase_bottleneck_share_error_holdout = _max_holdout_by_field( + holdouts, + "phase_bottleneck_share_absolute_error", + ) + min_phase_top3_overlap_holdout = _min_holdout_by_field(holdouts, "phase_top3_overlap_rate") + min_phase_bucket_top3_overlap_holdout = _min_holdout_by_field(holdouts, "phase_bucket_top3_overlap_rate") + status = "ok" if errors else "insufficient_data" + if holdouts and not errors: + warnings.append("all holdouts were unscored") + calibration_fidelity_status, calibration_fidelity_blockers = _calibration_fidelity_support( + holdouts, + errors=errors, + interval_covered_count=interval_covered_count, + uncertainty_fractions=uncertainty_fractions, + empirical_required_uncertainty_fractions=empirical_required_uncertainty_fractions, + memory_percentage_errors=memory_percentage_errors, + memory_bottleneck_evaluated_count=len(memory_bottleneck_holdouts), + memory_bottleneck_bucket_match_count=memory_bottleneck_bucket_match_count, + phase_bottleneck_evaluated_count=len(phase_bottleneck_holdouts), + phase_bottleneck_bucket_match_count=phase_bottleneck_bucket_match_count, + phase_top3_evaluated_count=len(phase_top3_holdouts), + phase_bucket_top3_overlap_rates=phase_bucket_top3_overlap_rates, + phase_bottleneck_share_errors=phase_bottleneck_share_errors, + ) + calibration_validation_gaps = _calibration_validation_gap_portfolio(holdouts, calibration_fidelity_blockers) + calibration_design_configs = _materialize_calibration_measurement_design_configs_from_context( + base_config=base_config, + base_topology=base_topology, + behavior_points=behavior_points, + calibration_validation_gaps=calibration_validation_gaps, + world_size=world_size, + local_world_size=local_world_size, + max_configs_per_measurement=4, + ) + + return CalibrationReport( + base_config_path=str(base_path), + benchmark_dir=str(benchmark_path), + status=status, + measured_point_count=len(measured_points), + evaluated_count=len(holdouts), + skipped_count=skipped_count, + mean_absolute_percentage_error=round(statistics.fmean(errors), 3) if errors else None, + median_absolute_percentage_error=round(statistics.median(errors), 3) if errors else None, + max_absolute_percentage_error=round(max(errors), 3) if errors else None, + max_absolute_percentage_error_label=max_error_holdout.label if max_error_holdout is not None else None, + max_absolute_percentage_error_prediction_status=( + max_error_holdout.prediction_status if max_error_holdout is not None else None + ), + max_absolute_percentage_error_in_prediction_interval=( + max_error_holdout.actual_tokens_in_prediction_interval if max_error_holdout is not None else None + ), + prediction_interval_coverage_count=interval_covered_count, + prediction_interval_coverage_rate=round(interval_covered_count / len(errors), 3) if errors else None, + mean_prediction_uncertainty_fraction=( + round(statistics.fmean(uncertainty_fractions), 3) if uncertainty_fractions else None + ), + max_prediction_uncertainty_fraction=round(max(uncertainty_fractions), 3) if uncertainty_fractions else None, + max_prediction_uncertainty_label=( + max_uncertainty_holdout.label if max_uncertainty_holdout is not None else None + ), + memory_evaluated_count=len(memory_percentage_errors), + mean_memory_absolute_error_gb=round(statistics.fmean(memory_absolute_errors), 3) + if memory_absolute_errors + else None, + max_memory_absolute_error_gb=round(max(memory_absolute_errors), 3) if memory_absolute_errors else None, + max_memory_absolute_error_label=( + max_memory_error_holdout.label if max_memory_error_holdout is not None else None + ), + mean_memory_absolute_percentage_error=round(statistics.fmean(memory_percentage_errors), 3) + if memory_percentage_errors + else None, + max_memory_absolute_percentage_error=round(max(memory_percentage_errors), 3) + if memory_percentage_errors + else None, + max_memory_absolute_percentage_error_label=( + max_memory_percentage_error_holdout.label if max_memory_percentage_error_holdout is not None else None + ), + memory_prediction_basis_counts=dict(sorted(memory_basis_counts.items())), + memory_coverage_status_counts=dict(sorted(memory_coverage_counts.items())), + memory_feasibility_status_counts=dict(sorted(memory_feasibility_counts.items())), + max_predicted_memory_residual_gb=( + round(max(predicted_memory_residuals), 3) if predicted_memory_residuals else None + ), + max_predicted_memory_residual_gb_label=( + max_predicted_memory_residual_holdout.label if max_predicted_memory_residual_holdout is not None else None + ), + max_predicted_memory_residual_fraction_of_peak=( + round(max(predicted_memory_residual_fractions), 3) if predicted_memory_residual_fractions else None + ), + max_predicted_memory_residual_fraction_of_peak_label=( + max_predicted_memory_residual_fraction_holdout.label + if max_predicted_memory_residual_fraction_holdout is not None + else None + ), + max_actual_memory_residual_gb=round(max(actual_memory_residuals), 3) if actual_memory_residuals else None, + max_actual_memory_residual_gb_label=( + max_actual_memory_residual_holdout.label if max_actual_memory_residual_holdout is not None else None + ), + max_actual_memory_residual_fraction_of_peak=( + round(max(actual_memory_residual_fractions), 3) if actual_memory_residual_fractions else None + ), + max_actual_memory_residual_fraction_of_peak_label=( + max_actual_memory_residual_fraction_holdout.label + if max_actual_memory_residual_fraction_holdout is not None + else None + ), + memory_bottleneck_evaluated_count=len(memory_bottleneck_holdouts), + memory_bottleneck_phase_match_count=memory_bottleneck_phase_match_count, + memory_bottleneck_phase_match_rate=( + round(memory_bottleneck_phase_match_count / len(memory_bottleneck_holdouts), 3) + if memory_bottleneck_holdouts + else None + ), + memory_bottleneck_bucket_match_count=memory_bottleneck_bucket_match_count, + memory_bottleneck_bucket_match_rate=( + round(memory_bottleneck_bucket_match_count / len(memory_bottleneck_holdouts), 3) + if memory_bottleneck_holdouts + else None + ), + mean_memory_bottleneck_peak_absolute_error_gb=( + round(statistics.fmean(memory_bottleneck_peak_errors), 3) if memory_bottleneck_peak_errors else None + ), + max_memory_bottleneck_peak_absolute_error_gb=( + round(max(memory_bottleneck_peak_errors), 3) if memory_bottleneck_peak_errors else None + ), + max_memory_bottleneck_peak_absolute_error_label=( + max_memory_bottleneck_peak_error_holdout.label + if max_memory_bottleneck_peak_error_holdout is not None + else None + ), + mean_memory_bottleneck_fraction_absolute_error=( + round(statistics.fmean(memory_bottleneck_fraction_errors), 3) if memory_bottleneck_fraction_errors else None + ), + max_memory_bottleneck_fraction_absolute_error=( + round(max(memory_bottleneck_fraction_errors), 3) if memory_bottleneck_fraction_errors else None + ), + max_memory_bottleneck_fraction_absolute_error_label=( + max_memory_bottleneck_fraction_error_holdout.label + if max_memory_bottleneck_fraction_error_holdout is not None + else None + ), + memory_bottleneck_phase_mismatch_labels=[ + holdout.label for holdout in memory_bottleneck_holdouts if holdout.memory_bottleneck_phase_match is False + ], + memory_bottleneck_bucket_mismatch_labels=[ + holdout.label for holdout in memory_bottleneck_holdouts if holdout.memory_bottleneck_bucket_match is False + ], + phase_bottleneck_evaluated_count=len(phase_bottleneck_holdouts), + phase_bottleneck_phase_match_count=phase_bottleneck_phase_match_count, + phase_bottleneck_phase_match_rate=( + round(phase_bottleneck_phase_match_count / len(phase_bottleneck_holdouts), 3) + if phase_bottleneck_holdouts + else None + ), + phase_bottleneck_bucket_match_count=phase_bottleneck_bucket_match_count, + phase_bottleneck_bucket_match_rate=( + round(phase_bottleneck_bucket_match_count / len(phase_bottleneck_holdouts), 3) + if phase_bottleneck_holdouts + else None + ), + mean_phase_bottleneck_share_absolute_error=( + round(statistics.fmean(phase_bottleneck_share_errors), 6) if phase_bottleneck_share_errors else None + ), + max_phase_bottleneck_share_absolute_error=( + round(max(phase_bottleneck_share_errors), 6) if phase_bottleneck_share_errors else None + ), + max_phase_bottleneck_share_absolute_error_label=( + max_phase_bottleneck_share_error_holdout.label + if max_phase_bottleneck_share_error_holdout is not None + else None + ), + phase_bottleneck_phase_mismatch_labels=[ + holdout.label for holdout in phase_bottleneck_holdouts if holdout.phase_bottleneck_phase_match is False + ], + phase_bottleneck_bucket_mismatch_labels=[ + holdout.label for holdout in phase_bottleneck_holdouts if holdout.phase_bottleneck_bucket_match is False + ], + phase_top3_evaluated_count=len(phase_top3_holdouts), + mean_phase_top3_overlap_rate=( + round(statistics.fmean(phase_top3_overlap_rates), 3) if phase_top3_overlap_rates else None + ), + min_phase_top3_overlap_rate=round(min(phase_top3_overlap_rates), 3) if phase_top3_overlap_rates else None, + min_phase_top3_overlap_rate_label=( + min_phase_top3_overlap_holdout.label if min_phase_top3_overlap_holdout is not None else None + ), + mean_phase_bucket_top3_overlap_rate=( + round(statistics.fmean(phase_bucket_top3_overlap_rates), 3) if phase_bucket_top3_overlap_rates else None + ), + min_phase_bucket_top3_overlap_rate=( + round(min(phase_bucket_top3_overlap_rates), 3) if phase_bucket_top3_overlap_rates else None + ), + min_phase_bucket_top3_overlap_rate_label=( + min_phase_bucket_top3_overlap_holdout.label if min_phase_bucket_top3_overlap_holdout is not None else None + ), + calibration_fidelity_status=calibration_fidelity_status, + calibration_fidelity_blockers=calibration_fidelity_blockers, + calibration_validation_gap_count=len(calibration_validation_gaps), + calibration_validation_gap_status_counts=_calibration_gap_status_counts(calibration_validation_gaps), + calibration_validation_gap_required_measurements=_unique_required_measurements(calibration_validation_gaps), + calibration_validation_gaps=calibration_validation_gaps, + prediction_status_counts=dict(sorted(status_counts.items())), + holdouts=holdouts, + warnings=warnings, + prediction_uncertainty_calibration_status=_prediction_uncertainty_calibration_status( + errors=errors, + interval_covered_count=interval_covered_count, + uncertainty_fractions=uncertainty_fractions, + empirical_required_uncertainty_fractions=empirical_required_uncertainty_fractions, + ), + mean_empirical_required_uncertainty_fraction=( + round(statistics.fmean(empirical_required_uncertainty_fractions), 3) + if empirical_required_uncertainty_fractions + else None + ), + max_empirical_required_uncertainty_fraction=( + round(max(empirical_required_uncertainty_fractions), 3) + if empirical_required_uncertainty_fractions + else None + ), + max_empirical_required_uncertainty_label=( + max_empirical_required_uncertainty_pair[1].label + if max_empirical_required_uncertainty_pair is not None + else None + ), + measurement_design_config_count=len(calibration_design_configs), + measurement_design_config_labels=[item.label for item in calibration_design_configs], + measurement_design_config_filenames=[f"design_{item.filename}" for item in calibration_design_configs], + calibration_support_benchmark_dirs=[str(path) for path in calibration_support_benchmark_dirs], + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pack", help="Built-in calibration-pack name") + parser.add_argument("--config", type=Path, default=None) + parser.add_argument("--benchmark-dir", type=Path, default=None) + parser.add_argument("--world-size", type=int, default=None) + parser.add_argument("--local-world-size", type=int, default=None) + parser.add_argument("--device-memory-limit-gb", type=float, default=80.0) + parser.add_argument("--memory-safety-factor", type=float, default=1.15) + parser.add_argument( + "--calibration-support-benchmark-dir", + dest="calibration_support_benchmark_dirs", + action="append", + type=Path, + default=[], + help="Additional benchmark dir whose rows can support predictions but are not target holdouts", + ) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument( + "--write-measurement-configs", + type=Path, + default=None, + help="Write bounded calibration replay/profile measurement configs in this directory", + ) + args = parser.parse_args() + + args.config, args.benchmark_dir = resolve_pack_inputs(args.pack, args.config, args.benchmark_dir) + if args.config is None or args.benchmark_dir is None: + parser.error("provide --pack, or both --config and --benchmark-dir") + + report = evaluate_calibration( + args.config, + benchmark_dir=args.benchmark_dir, + world_size=args.world_size, + local_world_size=args.local_world_size, + calibration_support_benchmark_dirs=tuple(args.calibration_support_benchmark_dirs), + device_memory_limit_gb=args.device_memory_limit_gb, + memory_safety_factor=args.memory_safety_factor, + ) + if args.write_measurement_configs: + write_measurement_configs(report, args.write_measurement_configs) + rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/src/xorl/sim/calibration_packs.py b/src/xorl/sim/calibration_packs.py new file mode 100644 index 00000000..d2309175 --- /dev/null +++ b/src/xorl/sim/calibration_packs.py @@ -0,0 +1,175 @@ +"""Discovery and validation for portable simulator calibration packs.""" + +from __future__ import annotations + +import argparse +import json +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import yaml + + +PACK_SCHEMA_VERSION = 1 +_PACK_ROOT = Path(__file__).with_name("calibration_packs") +_FORBIDDEN_CONTENT = { + "absolute home path": re.compile(r"/home/"), + "workspace mount": re.compile(r"/workspace(?:/|\b)"), + "shared mount": re.compile(r"/shared(?:/|\b)"), + "internal repository name": re.compile(r"\bxorl-internal\b"), + "Kubernetes command": re.compile(r"\bkubectl\b", re.IGNORECASE), + "Kubernetes service address": re.compile(r"\.svc\.cluster\.local\b"), + "PVC setting": re.compile(r"\b(?:home|shared)[_-]?pvc\b", re.IGNORECASE), + "Volcano setting": re.compile(r"\bvolcano\b", re.IGNORECASE), + "team scheduling label": re.compile(r"\bteam:\s*(?:turbo|shaping)\b", re.IGNORECASE), +} + + +@dataclass(frozen=True) +class CalibrationPack: + name: str + path: Path + manifest: dict[str, Any] + + @property + def default_config(self) -> Path: + return self.path / str(self.manifest["default_config"]) + + +def calibration_pack_root() -> Path: + """Return the installed directory containing built-in calibration packs.""" + + return _PACK_ROOT + + +def list_calibration_packs() -> list[str]: + if not _PACK_ROOT.is_dir(): + return [] + return sorted(path.name for path in _PACK_ROOT.iterdir() if (path / "manifest.json").is_file()) + + +def resolve_calibration_pack(value: str | Path) -> Path: + """Resolve a filesystem path or ``builtin:`` calibration-pack reference.""" + + raw = str(value) + name = raw.removeprefix("builtin:") + builtin = _PACK_ROOT / name + if raw.startswith("builtin:") or (not Path(raw).exists() and (builtin / "manifest.json").is_file()): + if not (builtin / "manifest.json").is_file(): + available = ", ".join(list_calibration_packs()) or "none" + raise ValueError(f"unknown built-in calibration pack {name!r}; available: {available}") + return builtin + return Path(raw) + + +def load_calibration_pack(value: str | Path) -> CalibrationPack: + path = resolve_calibration_pack(value) + manifest_path = path / "manifest.json" + if not manifest_path.is_file(): + raise FileNotFoundError(f"missing calibration-pack manifest: {manifest_path}") + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + return CalibrationPack(name=str(manifest.get("name", path.name)), path=path, manifest=manifest) + + +def resolve_pack_inputs( + pack_name: str | None, + config: Path | None, + benchmark_dir: Path | None, +) -> tuple[Path | None, Path | None]: + """Fill omitted config and benchmark paths from a built-in pack.""" + + if pack_name is None: + return config, benchmark_dir + pack = load_calibration_pack(pack_name) + return config or pack.default_config, benchmark_dir or pack.path + + +def _check(name: str, passed: bool, detail: str) -> dict[str, str]: + return {"name": name, "status": "pass" if passed else "fail", "detail": detail} + + +def validate_calibration_pack(value: str | Path) -> dict[str, Any]: + pack = load_calibration_pack(value) + manifest = pack.manifest + checks = [ + _check( + "schema_version", + manifest.get("schema_version") == PACK_SCHEMA_VERSION, + f"expected {PACK_SCHEMA_VERSION}, found {manifest.get('schema_version')!r}", + ), + _check("manifest_name", manifest.get("name") == pack.path.name, f"manifest name={manifest.get('name')!r}"), + _check("model", isinstance(manifest.get("model"), str), f"model={manifest.get('model')!r}"), + ] + + declared = [manifest.get("default_config"), *manifest.get("configs", []), *manifest.get("results", [])] + relative_paths = [Path(str(item)) for item in declared if item] + for relative in relative_paths: + is_safe = not relative.is_absolute() and ".." not in relative.parts + checks.append(_check(f"portable_path:{relative}", is_safe, str(relative))) + checks.append(_check(f"file_exists:{relative}", is_safe and (pack.path / relative).is_file(), str(relative))) + + config_paths = sorted({Path(str(item)) for item in manifest.get("configs", [])}) + for relative in config_paths: + path = pack.path / relative + if not path.is_file(): + continue + config = yaml.safe_load(path.read_text(encoding="utf-8")) + model_path = config.get("model", {}).get("model_path") if isinstance(config, dict) else None + checks.append( + _check( + f"config_model:{relative}", + model_path == manifest.get("model"), + f"expected {manifest.get('model')!r}, found {model_path!r}", + ) + ) + + for path in sorted(item for item in pack.path.rglob("*") if item.is_file()): + text = path.read_text(encoding="utf-8") + relative = path.relative_to(pack.path) + for label, pattern in _FORBIDDEN_CONTENT.items(): + match = pattern.search(text) + checks.append( + _check( + f"sanitized:{relative}:{label}", + match is None, + "not present" if match is None else f"matched {match.group(0)!r}", + ) + ) + + return { + "name": pack.name, + "path": str(pack.path), + "schema_version": manifest.get("schema_version"), + "status": "pass" if all(check["status"] == "pass" for check in checks) else "fail", + "checks": checks, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + subparsers = parser.add_subparsers(dest="command", required=True) + subparsers.add_parser("list", help="List installed calibration packs") + path_parser = subparsers.add_parser("path", help="Print the path to an installed calibration pack") + path_parser.add_argument("pack") + validate_parser = subparsers.add_parser("validate", help="Validate one or all installed calibration packs") + validate_parser.add_argument("pack", nargs="?", default=None) + args = parser.parse_args() + + if args.command == "list": + print("\n".join(list_calibration_packs())) + return + if args.command == "path": + print(resolve_calibration_pack(args.pack)) + return + + names = [args.pack] if args.pack else list_calibration_packs() + reports = [validate_calibration_pack(name) for name in names] + print(json.dumps(reports, indent=2, sort_keys=True)) + if any(report["status"] != "pass" for report in reports): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/src/xorl/sim/calibration_packs/qwen3_235b_a22b/README.md b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/README.md new file mode 100644 index 00000000..a4154a8c --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/README.md @@ -0,0 +1,4 @@ +# Qwen3-235B-A22B Calibration Pack + +This portable pack contains one topology/configuration reference and summarized historical observations used to test +gradient-accumulation interpolation and observed OOM boundaries. It contains no launcher or infrastructure settings. diff --git a/src/xorl/sim/calibration_packs/qwen3_235b_a22b/RESULTS.md b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/RESULTS.md new file mode 100644 index 00000000..d4a0065a --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/RESULTS.md @@ -0,0 +1,9 @@ +# Qwen3-235B-A22B at 2k Context + +Measured: 4 nodes x 8 H100, U1, DP shard 32, EP8, eFSDP4. + +| run | gcm | pack | mbs | tok/step | step s | MFU | tok/s tot | tok/s/GPU | peak GB | status | +|-----|-----|-----:|----:|---------:|-------:|----:|----------:|----------:|--------:|--------| +| n4_ep8_bd_pk4096 | before_dispatch | 4096 | 1 | 131,072 | ~18.4 | ~3.0% | ~6,800 | ~213 | 68.3 | OK | +| n4_ep8_bd_pk4096_ga2 | before_dispatch | 4096 | 1 | 262,144 | ~31.3 | ~3.7% | ~8,400 | ~263 | 68.3 | NEW BEST | +| n4_ep8_bd_pk16k | before_dispatch | 16384 | 1 | 524,288 | -- | -- | -- | -- | OOM | FAIL | diff --git a/src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_2k_4node_ep8.yaml b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_2k_4node_ep8.yaml new file mode 100644 index 00000000..167d2e5e --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_2k_4node_ep8.yaml @@ -0,0 +1,51 @@ +model: + model_path: Qwen/Qwen3-235B-A22B + tokenizer_path: Qwen/Qwen3-235B-A22B + attn_implementation: flash_attention_3 + moe_implementation: quack + ep_dispatch: deepep + deepep_buffer_size_gb: 2.0 + deepep_num_sms: 24 + deepep_async_combine: false +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 4096 + select_columns: + - input_ids + - labels + sample_packing_method: sequential + sample_packing_sequence_len: 4096 + dataloader_num_workers: 0 + dataloader_pin_memory: false + pad_to_multiple_of: 128 +train: + output_dir: outputs/xorl-sim/qwen3_235b_a22b_2k_4node_ep8 + data_parallel_mode: fsdp2 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 32 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + expert_parallel_size: 8 + micro_batch_size: 1 + gradient_accumulation_steps: 1 + optimizer: muon + optimizer_dtype: bf16 + muon_momentum: 0.95 + enable_mixed_precision: true + skip_param_upcast: true + fsdp_reduce_dtype: fp32 + ce_mode: quack_linear + enable_gradient_checkpointing: true + gradient_checkpointing_method: recompute_before_dispatch + enable_full_shard: true + init_device: meta + max_steps: 8 + save_steps: 0 + save_epochs: 0 + save_hf_weights: false + log_format: structured + use_wandb: false diff --git a/src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_muon_8node_ep8_efsdp8_deepep.yaml b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_muon_8node_ep8_efsdp8_deepep.yaml new file mode 100644 index 00000000..3d778759 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/configs/qwen3_235b_a22b_muon_8node_ep8_efsdp8_deepep.yaml @@ -0,0 +1,56 @@ +model: + model_path: Qwen/Qwen3-235B-A22B + attn_implementation: flash_attention_3 + ep_dispatch: deepep + rmsnorm_mode: compile + deepep_buffer_size_gb: 2.0 + deepep_num_sms: 48 +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 64000 + select_columns: + - input_ids + - labels + sample_packing_method: sequential + sample_packing_sequence_len: 64000 + dataloader_num_workers: 0 + dataloader_pin_memory: false + pad_to_multiple_of: 4096 +train: + output_dir: outputs/xorl-sim/qwen3_235b_a22b_muon_8node_ep8_efsdp8_deepep + data_parallel_mode: fsdp2 + ulysses_parallel_size: 8 + expert_parallel_size: 8 + data_parallel_replicate_size: 8 + data_parallel_shard_size: 1 + num_train_epochs: 1 + max_steps: 20 + micro_batch_size: 1 + gradient_accumulation_steps: 1 + optimizer: muon + optimizer_dtype: bf16 + muon_lr: 0.0001 + muon_momentum: 0.95 + lr: 1.0e-05 + lr_warmup_ratio: 0.0 + lr_decay_style: constant + lr_decay_ratio: 1.0 + weight_decay: 0.01 + max_grad_norm: 1.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + enable_full_shard: true + enable_activation_offload: false + init_device: meta + load_weights_mode: all_ranks + enable_full_determinism: false + seed: 42 + empty_cache_steps: 500 + ckpt_manager: dcp + save_steps: 0 + save_epochs: 0 + save_hf_weights: false + log_format: structured + use_wandb: false diff --git a/src/xorl/sim/calibration_packs/qwen3_235b_a22b/manifest.json b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/manifest.json new file mode 100644 index 00000000..7a06e143 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_235b_a22b/manifest.json @@ -0,0 +1,30 @@ +{ + "schema_version": 1, + "name": "qwen3_235b_a22b", + "model": "Qwen/Qwen3-235B-A22B", + "description": "Portable Qwen3-235B topology config and summarized 2k-context GA/OOM calibration rows.", + "captured_at": "2026-06-26", + "hardware": "32 NVIDIA H100 GPUs", + "balanced_routing": false, + "default_config": "configs/qwen3_235b_a22b_2k_4node_ep8.yaml", + "configs": [ + "configs/qwen3_235b_a22b_2k_4node_ep8.yaml", + "configs/qwen3_235b_a22b_muon_8node_ep8_efsdp8_deepep.yaml" + ], + "results": [ + "RESULTS.md" + ], + "golden": { + "behavior_point_count": 3, + "world_size": 32, + "global_batch_size": 32, + "default_tokens_per_sec": 6800.0, + "best_raw_tokens_per_sec": 8400.0, + "best_promotable_tokens_per_sec": null, + "analytic_peak_floor_gb": 56.812 + }, + "limitations": [ + "Historical observations are intended for regression and calibration, not hardware-independent performance claims.", + "No correctness gate is attached to these throughput rows." + ] +} diff --git a/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/README.md b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/README.md new file mode 100644 index 00000000..2adcf591 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/README.md @@ -0,0 +1,21 @@ +# Qwen3.5-397B-A17B Short-Context Calibration + +- Model: `Qwen/Qwen3.5-397B-A17B` +- Data: synthetic tokenized data, sequential packing, `max_seq_len=4096` +- Hardware: 8 nodes x 8 H100 GPUs +- Topology: `pp=1`, `tp=1`, `ring=1`, `ulysses=1`, `dp_shard=64`, `ep=32`, `ep_fsdp=2` +- Runtime: Quack MoE, DeepEP dispatch, SMS48, balanced synthetic routing + +MFU uses 989 BF16 TFLOPS/GPU as the H100 denominator. + +| trial | tok/s | tok/s/GPU | TFLOPS/GPU | MFU | step | correctness | +| --- | ---: | ---: | ---: | ---: | ---: | --- | +| R75 | 59.217K | 925.3 | 93.4 | 9.44% | 22.299s | raw-speed only | +| R73 | 59.188K | 924.8 | 93.4 | 9.44% | 22.300s | static K3 pass | +| R70 | 54.227K | 847.3 | 85.5 | 8.65% | 24.561s | not promoted | +| R69 | 52.545K | 821.0 | 82.9 | 8.38% | 20.004s | not promoted | +| R67 | 48.474K | 757.4 | 76.5 | 7.74% | 21.804s | not promoted | +| R66 | 43.770K | 683.9 | 69.0 | 6.98% | 18.204s | not promoted | + +R73 passed its static replay over 129 output tokens with mean K3 `0.000475`, p95 `0.001335`, and max `0.028437`. +R75 changes DeepEP combine to asynchronous mode and is not promotable without a matching gate. diff --git a/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_ep64_reference.yaml b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_ep64_reference.yaml new file mode 100644 index 00000000..b4db2a7f --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_ep64_reference.yaml @@ -0,0 +1,67 @@ +model: + model_path: Qwen/Qwen3.5-397B-A17B + tokenizer_path: Qwen/Qwen3.5-397B-A17B + attn_implementation: flash_attention_3 + moe_implementation: quack + ep_dispatch: deepep + train_router: false + record_routing_weights: false + rmsnorm_mode: compile + deepep_buffer_size_gb: 1.0 + deepep_num_sms: 48 +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 4096 + select_columns: + - input_ids + - labels + sample_packing_method: sequential + sample_packing_sequence_len: 4096 + dataloader_num_workers: 0 + dataloader_pin_memory: false + pad_to_multiple_of: 128 +train: + output_dir: outputs/xorl-sim/ep64_reference + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + expert_parallel_size: 64 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 64 + ep_intranode: true + num_train_epochs: 1 + max_steps: 13 + micro_batch_size: 1 + gradient_accumulation_steps: 1 + optimizer: muon + optimizer_dtype: bf16 + muon_lr: 0.0001 + muon_momentum: 0.95 + muon_ns_algorithm: gram_newton_schulz + muon_ns_use_quack_kernels: true + lr: 1.0e-05 + lr_warmup_ratio: 0.0 + lr_decay_style: constant + lr_decay_ratio: 1.0 + weight_decay: 0.01 + max_grad_norm: 1.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + gradient_checkpointing_method: recompute_before_dispatch + enable_full_shard: true + enable_activation_offload: false + init_device: meta + load_weights_mode: all_ranks + enable_full_determinism: false + empty_cache_steps: 500 + gc_steps: 500 + ckpt_manager: dcp + save_steps: 0 + save_epochs: 0 + save_hf_weights: false + log_format: structured + use_wandb: false diff --git a/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r69_mbs4.yaml b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r69_mbs4.yaml new file mode 100644 index 00000000..15144577 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r69_mbs4.yaml @@ -0,0 +1,72 @@ +model: + model_path: Qwen/Qwen3.5-397B-A17B + tokenizer_path: Qwen/Qwen3.5-397B-A17B + attn_implementation: flash_attention_3 + moe_implementation: quack + ep_dispatch: deepep + train_router: false + record_routing_weights: false + rmsnorm_mode: compile + deepep_buffer_size_gb: 1.0 + deepep_num_sms: 48 + deepep_async_combine: false +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 4096 + select_columns: + - input_ids + - labels + sample_packing_method: sequential + sample_packing_sequence_len: 4096 + dataloader_num_workers: 0 + dataloader_pin_memory: false + pad_to_multiple_of: 128 +train: + output_dir: outputs/xorl-sim/r69 + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + expert_parallel_size: 32 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 64 + ep_intranode: true + num_train_epochs: 1 + max_steps: 8 + micro_batch_size: 4 + gradient_accumulation_steps: 1 + optimizer: muon + optimizer_dtype: bf16 + muon_lr: 0.0001 + muon_momentum: 0.0 + muon_update_dtype: bf16 + muon_ns_algorithm: gram_newton_schulz + muon_ns_use_quack_kernels: true + lr: 1.0e-05 + lr_warmup_ratio: 0.0 + lr_decay_style: constant + lr_decay_ratio: 1.0 + weight_decay: 0.01 + max_grad_norm: 0.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + gradient_checkpointing_method: recompute_full_layer + enable_full_shard: true + enable_activation_offload: true + activation_gpu_limit: 0.0 + activation_offload_prefetch_count: 2 + init_device: meta + load_weights_mode: all_ranks + enable_full_determinism: false + enable_compile: true + empty_cache_steps: 10 + gc_steps: 10 + ckpt_manager: dcp + save_steps: 0 + save_epochs: 0 + save_hf_weights: false + log_format: structured + use_wandb: false diff --git a/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r73_sync.yaml b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r73_sync.yaml new file mode 100644 index 00000000..1567286e --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r73_sync.yaml @@ -0,0 +1,72 @@ +model: + model_path: Qwen/Qwen3.5-397B-A17B + tokenizer_path: Qwen/Qwen3.5-397B-A17B + attn_implementation: flash_attention_3 + moe_implementation: quack + ep_dispatch: deepep + train_router: false + record_routing_weights: false + rmsnorm_mode: compile + deepep_buffer_size_gb: 0.5 + deepep_num_sms: 48 + deepep_async_combine: false +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 4096 + select_columns: + - input_ids + - labels + sample_packing_method: sequential + sample_packing_sequence_len: 4096 + dataloader_num_workers: 0 + dataloader_pin_memory: false + pad_to_multiple_of: 128 +train: + output_dir: outputs/xorl-sim/r73 + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + expert_parallel_size: 32 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 64 + ep_intranode: true + num_train_epochs: 1 + max_steps: 8 + micro_batch_size: 5 + gradient_accumulation_steps: 1 + optimizer: muon + optimizer_dtype: bf16 + muon_lr: 0.0001 + muon_momentum: 0.0 + muon_update_dtype: bf16 + muon_ns_algorithm: gram_newton_schulz + muon_ns_use_quack_kernels: true + lr: 1.0e-05 + lr_warmup_ratio: 0.0 + lr_decay_style: constant + lr_decay_ratio: 1.0 + weight_decay: 0.01 + max_grad_norm: 0.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + gradient_checkpointing_method: recompute_full_layer + enable_full_shard: true + enable_activation_offload: true + activation_gpu_limit: 0.0 + activation_offload_prefetch_count: 4 + init_device: meta + load_weights_mode: all_ranks + enable_full_determinism: false + enable_compile: true + empty_cache_steps: 10 + gc_steps: 10 + ckpt_manager: dcp + save_steps: 0 + save_epochs: 0 + save_hf_weights: false + log_format: structured + use_wandb: false diff --git a/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r75_async.yaml b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r75_async.yaml new file mode 100644 index 00000000..56e4d2b0 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/configs/qwen35_r75_async.yaml @@ -0,0 +1,72 @@ +model: + model_path: Qwen/Qwen3.5-397B-A17B + tokenizer_path: Qwen/Qwen3.5-397B-A17B + attn_implementation: flash_attention_3 + moe_implementation: quack + ep_dispatch: deepep + train_router: false + record_routing_weights: false + rmsnorm_mode: compile + deepep_buffer_size_gb: 0.5 + deepep_num_sms: 48 + deepep_async_combine: true +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 4096 + select_columns: + - input_ids + - labels + sample_packing_method: sequential + sample_packing_sequence_len: 4096 + dataloader_num_workers: 0 + dataloader_pin_memory: false + pad_to_multiple_of: 128 +train: + output_dir: outputs/xorl-sim/r75 + data_parallel_mode: fsdp2 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 + expert_parallel_size: 32 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 64 + ep_intranode: true + num_train_epochs: 1 + max_steps: 8 + micro_batch_size: 5 + gradient_accumulation_steps: 1 + optimizer: muon + optimizer_dtype: bf16 + muon_lr: 0.0001 + muon_momentum: 0.0 + muon_update_dtype: bf16 + muon_ns_algorithm: gram_newton_schulz + muon_ns_use_quack_kernels: true + lr: 1.0e-05 + lr_warmup_ratio: 0.0 + lr_decay_style: constant + lr_decay_ratio: 1.0 + weight_decay: 0.01 + max_grad_norm: 0.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + gradient_checkpointing_method: recompute_full_layer + enable_full_shard: true + enable_activation_offload: true + activation_gpu_limit: 0.0 + activation_offload_prefetch_count: 4 + init_device: meta + load_weights_mode: all_ranks + enable_full_determinism: false + enable_compile: true + empty_cache_steps: 10 + gc_steps: 10 + ckpt_manager: dcp + save_steps: 0 + save_epochs: 0 + save_hf_weights: false + log_format: structured + use_wandb: false diff --git a/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/manifest.json b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/manifest.json new file mode 100644 index 00000000..a4591d3b --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/manifest.json @@ -0,0 +1,32 @@ +{ + "schema_version": 1, + "name": "qwen3_5_397b_a17b", + "model": "Qwen/Qwen3.5-397B-A17B", + "description": "Portable Qwen3.5-397B short-context throughput and correctness calibration pack.", + "captured_at": "2026-05-20", + "hardware": "64 NVIDIA H100 GPUs", + "balanced_routing": true, + "default_config": "configs/qwen35_r73_sync.yaml", + "configs": [ + "configs/qwen35_r69_mbs4.yaml", + "configs/qwen35_r73_sync.yaml", + "configs/qwen35_r75_async.yaml", + "configs/qwen35_ep64_reference.yaml" + ], + "results": [ + "results/shortctx_8node_mfu_summary_20260519.json" + ], + "golden": { + "behavior_point_count": 6, + "world_size": 64, + "global_batch_size": 320, + "default_tokens_per_sec": 59188.0, + "best_raw_tokens_per_sec": 59217.0, + "best_promotable_tokens_per_sec": 59188.0, + "analytic_peak_floor_gb": 46.767 + }, + "limitations": [ + "The rows are short synthetic-data measurements.", + "The synchronous R73 row has a matching static K3 pass; asynchronous rows require their own correctness gate." + ] +} diff --git a/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/results/shortctx_8node_mfu_summary_20260519.json b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/results/shortctx_8node_mfu_summary_20260519.json new file mode 100644 index 00000000..1d2739a1 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_5_397b_a17b/results/shortctx_8node_mfu_summary_20260519.json @@ -0,0 +1,155 @@ +{ + "best_by_mfu": [ + { + "activation_offload_prefetch_count": 4, + "caveat": "Raw speed high-water only; async combine still needs a passing static K3 gate before promotion.", + "deepep_async_combine": true, + "deepep_buffer_size_gb": 0.5, + "deepep_num_sms": 48, + "enable_activation_offload": true, + "enable_compile": true, + "ep_fsdp": 2, + "expert_parallel_size": 32, + "global_batch_size": 320, + "gradient_checkpointing_method": "recompute_full_layer", + "mean_tflops_per_gpu": 93.4, + "measured_steps": 6, + "mfu_percent": 9.44, + "micro_batch_size": 5, + "model_ref": "Qwen/Qwen3.5-397B-A17B", + "sample_packing_sequence_len": 4096, + "step_time_sec": 22.299, + "tokens_per_sec": 59217.0, + "tokens_per_sec_per_gpu": 925.3, + "trial": "r75", + "warmup_steps": 2 + }, + { + "activation_offload_prefetch_count": 4, + "caveat": "Correctness-promoted synchronous-combine row; practical tie with R75 by rounded MFU.", + "deepep_async_combine": false, + "deepep_buffer_size_gb": 0.5, + "deepep_num_sms": 48, + "enable_activation_offload": true, + "enable_compile": true, + "ep_fsdp": 2, + "expert_parallel_size": 32, + "global_batch_size": 320, + "gradient_checkpointing_method": "recompute_full_layer", + "k3_gate": "pass: mean=0.0004746597891399218, p95=0.00133467541659926, max=0.02843669572903562 over 129 tokens", + "mean_tflops_per_gpu": 93.4, + "measured_steps": 6, + "mfu_percent": 9.44, + "micro_batch_size": 5, + "model_ref": "Qwen/Qwen3.5-397B-A17B", + "sample_packing_sequence_len": 4096, + "step_time_sec": 22.3, + "tokens_per_sec": 59188.0, + "tokens_per_sec_per_gpu": 924.8, + "trial": "r73", + "warmup_steps": 2 + }, + { + "activation_offload_prefetch_count": 2, + "deepep_async_combine": false, + "deepep_buffer_size_gb": 0.5, + "deepep_num_sms": 48, + "enable_activation_offload": true, + "enable_compile": true, + "ep_fsdp": 2, + "expert_parallel_size": 32, + "global_batch_size": 320, + "gradient_checkpointing_method": "recompute_full_layer", + "mean_tflops_per_gpu": 85.5, + "measured_steps": 6, + "mfu_percent": 8.65, + "micro_batch_size": 5, + "model_ref": "Qwen/Qwen3.5-397B-A17B", + "sample_packing_sequence_len": 4096, + "step_time_sec": 24.561, + "tokens_per_sec": 54227.0, + "tokens_per_sec_per_gpu": 847.3, + "trial": "r70", + "warmup_steps": 2 + }, + { + "activation_offload_prefetch_count": 2, + "deepep_async_combine": false, + "deepep_buffer_size_gb": 1.0, + "deepep_num_sms": 48, + "enable_activation_offload": true, + "enable_compile": true, + "ep_fsdp": 2, + "expert_parallel_size": 32, + "global_batch_size": 256, + "gradient_checkpointing_method": "recompute_full_layer", + "mean_tflops_per_gpu": 82.9, + "measured_steps": 6, + "mfu_percent": 8.38, + "micro_batch_size": 4, + "model_ref": "Qwen/Qwen3.5-397B-A17B", + "sample_packing_sequence_len": 4096, + "step_time_sec": 20.004, + "tokens_per_sec": 52545.0, + "tokens_per_sec_per_gpu": 821.0, + "trial": "r69", + "warmup_steps": 2 + }, + { + "activation_offload_prefetch_count": 2, + "deepep_async_combine": false, + "deepep_buffer_size_gb": 1.0, + "deepep_num_sms": 48, + "enable_activation_offload": true, + "enable_compile": false, + "ep_fsdp": 2, + "expert_parallel_size": 32, + "global_batch_size": 256, + "gradient_checkpointing_method": "recompute_full_layer", + "mean_tflops_per_gpu": 76.5, + "measured_steps": 6, + "mfu_percent": 7.74, + "micro_batch_size": 4, + "model_ref": "Qwen/Qwen3.5-397B-A17B", + "sample_packing_sequence_len": 4096, + "step_time_sec": 21.804, + "tokens_per_sec": 48474.0, + "tokens_per_sec_per_gpu": 757.4, + "trial": "r67", + "warmup_steps": 2 + }, + { + "activation_offload_prefetch_count": 2, + "deepep_async_combine": false, + "deepep_buffer_size_gb": 1.0, + "deepep_num_sms": 48, + "enable_activation_offload": true, + "enable_compile": false, + "ep_fsdp": 2, + "expert_parallel_size": 32, + "global_batch_size": 192, + "gradient_checkpointing_method": "recompute_full_layer", + "mean_tflops_per_gpu": 69.0, + "measured_steps": 6, + "mfu_percent": 6.98, + "micro_batch_size": 3, + "model_ref": "Qwen/Qwen3.5-397B-A17B", + "sample_packing_sequence_len": 4096, + "step_time_sec": 18.204, + "tokens_per_sec": 43770.0, + "tokens_per_sec_per_gpu": 683.9, + "trial": "r66", + "warmup_steps": 2 + } + ], + "k3_status": "R73 synchronous-combine row passed Qwen3.5-specific static K3 replay on 2026-05-20; R75 remains raw-speed-only until async combine has its own K3 gate.", + "mfu_convention": "H100 denominator is 989 TFLOPS/GPU; mfu_percent = mean_tflops_per_gpu / 989 * 100.", + "model": "Qwen/Qwen3.5-397B-A17B", + "notes": [ + "The top-MFU rows are short smokes with 6 measured steps after 2 warmup steps.", + "Activation offload with prefetch4 and DeepEP buffer 0.5GB is the key difference between the top rows and earlier mbs4/mbs3 rows.", + "R73 and R75 are a practical tie by rounded TFLOPS/MFU; R73 is the safer correctness candidate because it keeps DeepEP combine synchronous." + ], + "source": "sanitized historical benchmark summary", + "workload": "dummy packed short context, max_seq_len=4096, 8 H100 nodes / 64 GPUs" +} diff --git a/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/README.md b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/README.md new file mode 100644 index 00000000..9277bf11 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/README.md @@ -0,0 +1,23 @@ +# Qwen3.6-35B-A3B 8k Calibration + +- Model: `Qwen/Qwen3.6-35B-A3B` +- Data: synthetic tokenized data, sequential packing, `sample_packing_sequence_len: 8193` +- Hardware: 4 nodes x 8 H100 GPUs +- Topology: `pp=1`, `tp=1`, `ring=1`, `ulysses=1`, `dp_shard=32`, `ep=8`, `ep_fsdp=4` +- Training: AdamW, BF16 mixed precision, FSDP2, full-layer recompute, DeepEP +- Runtime: `deepep_num_sms: 72`, `deepep_buffer_size_gb: 2.0`, `deepep_async_combine: true` +- Compiler: `enable_compile: true`, `gradient_checkpointing_method: recompute_full_layer` + +Reference shape: `micro_batch_size: 8`, `global_batch_size: 256`. + +| metric | value | +| --- | ---: | +| tokens/sec | ~261.0K | +| step time | ~8.04s | +| MFU | ~16.2% | +| allocated memory | ~56.4GB | +| allocator retries | 0 | + +The adjacent `mbs=10` observation fit but slowed to ~133.7K tok/s with allocator retries. The separate measured mbs6 +row reached 254.6K tok/s but failed its matching static K3 gate: mean `2.1339`, p95 `0.0834`, max `394.1871`. +These rows are useful calibration and failure-boundary evidence, not a correctness-promotable recipe. diff --git a/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/configs/qwen3_6_35b_a3b_8k_4node_ep8_mbs8_fullrecompute_deepep.yaml b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/configs/qwen3_6_35b_a3b_8k_4node_ep8_mbs8_fullrecompute_deepep.yaml new file mode 100644 index 00000000..cf95f92f --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/configs/qwen3_6_35b_a3b_8k_4node_ep8_mbs8_fullrecompute_deepep.yaml @@ -0,0 +1,62 @@ +model: + model_path: Qwen/Qwen3.6-35B-A3B + attn_implementation: flash_attention_3 + moe_implementation: quack + ep_dispatch: deepep + train_router: false + deepep_buffer_size_gb: 2.0 + deepep_num_sms: 72 + deepep_async_combine: true +data: + datasets: + - path: dummy + type: tokenized + max_seq_len: 8193 + select_columns: + - input_ids + - labels + dataset_prepared_path: last_prepared_dataset + sample_packing_method: sequential + sample_packing_sequence_len: 8193 + dataloader_num_workers: 4 + dataloader_prefetch_factor: 2 + dataloader_pin_memory: true + dataloader_drop_last: true + pad_to_multiple_of: 128 +train: + output_dir: outputs/xorl-sim/qwen3_6_35b_a3b_8k_4node_ep8_mbs8_fullrecompute_deepep + data_parallel_mode: fsdp2 + pipeline_parallel_size: 1 + ulysses_parallel_size: 1 + ringattn_parallel_size: 1 + tensor_parallel_size: 1 + data_parallel_replicate_size: 1 + data_parallel_shard_size: 32 + expert_parallel_size: 8 + num_train_epochs: 1 + max_steps: 12 + micro_batch_size: 8 + gradient_accumulation_steps: 1 + empty_cache_steps: 10 + gc_steps: 10 + gradient_checkpointing_method: recompute_full_layer + optimizer: adamw + lr: 1.0e-05 + lr_warmup_ratio: 0.0 + lr_decay_style: cosine + lr_decay_ratio: 1.0 + weight_decay: 0.01 + max_grad_norm: 1.0 + enable_mixed_precision: true + enable_gradient_checkpointing: true + enable_compile: true + enable_full_shard: true + init_device: meta + load_weights_mode: grouped + enable_full_determinism: false + ckpt_manager: dcp + save_steps: 0 + save_epochs: 0 + save_hf_weights: false + log_format: structured + use_wandb: false diff --git a/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/manifest.json b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/manifest.json new file mode 100644 index 00000000..5077c0c8 --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/manifest.json @@ -0,0 +1,29 @@ +{ + "schema_version": 1, + "name": "qwen3_6_35b_a3b", + "model": "Qwen/Qwen3.6-35B-A3B", + "description": "Portable Qwen3.6-35B 8k throughput, allocator-pressure, and static-K3 calibration pack.", + "captured_at": "2026-05-20", + "hardware": "32 NVIDIA H100 GPUs", + "balanced_routing": true, + "default_config": "configs/qwen3_6_35b_a3b_8k_4node_ep8_mbs8_fullrecompute_deepep.yaml", + "configs": [ + "configs/qwen3_6_35b_a3b_8k_4node_ep8_mbs8_fullrecompute_deepep.yaml" + ], + "results": [ + "results/qwen36_static_k3_summary_20260519.json" + ], + "golden": { + "behavior_point_count": 3, + "world_size": 32, + "global_batch_size": 256, + "default_tokens_per_sec": 261000.0, + "best_raw_tokens_per_sec": 261000.0, + "best_promotable_tokens_per_sec": null, + "analytic_peak_floor_gb": 18.14 + }, + "limitations": [ + "The measurements use synthetic packed data and balanced routing.", + "The fastest recorded path fails its matching static K3 gate and is not promotable." + ] +} diff --git a/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/results/qwen36_static_k3_summary_20260519.json b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/results/qwen36_static_k3_summary_20260519.json new file mode 100644 index 00000000..02a5393d --- /dev/null +++ b/src/xorl/sim/calibration_packs/qwen3_6_35b_a3b/results/qwen36_static_k3_summary_20260519.json @@ -0,0 +1,68 @@ +{ + "diagnostic_replays": [ + { + "candidate": "qwen36_diagnostic_1", + "k3_max": 0.3917743468870727, + "k3_mean": 0.007526155813053868, + "k3_p95": 0.02364128327561078, + "status": "diagnostic_low_k3", + "total_tokens": 192 + }, + { + "candidate": "qwen36_diagnostic_2", + "k3_max": 0.3917743468870727, + "k3_mean": 0.007526155813053868, + "k3_p95": 0.02364128327561078, + "status": "diagnostic_low_k3", + "total_tokens": 192 + }, + { + "candidate": "qwen36_diagnostic_3", + "k3_max": 4.682859666612269, + "k3_mean": 0.039442088877642674, + "k3_p95": 0.043214317136827184, + "status": "diagnostic", + "total_tokens": 192 + } + ], + "k3_gate": { + "candidate": "qwen36_deepep_mbs6", + "k3": { + "max": 394.1870587177588, + "mean": 2.1339466950283534, + "median": 2.0777892217216376e-05, + "p95": 0.08342470559875832, + "p99": 1.6374857623690962, + "total_prompts": 3, + "total_tokens": 192 + }, + "outlier_note": "The worst token accounts for most of the total K3; SGLang generation logprob and prefill top-logprob disagree on that token, so this artifact is useful for debugging but is not a passing promotion gate.", + "primary_failure": "k3.mean <= 0.001", + "source": "sanitized static replay summary", + "status": "fail", + "thresholds": { + "mean": 0.001, + "p95": 0.01 + } + }, + "model": "Qwen/Qwen3.6-35B-A3B", + "static_traces": { + "note": "These traces are Qwen3.6-specific and must not be reused for Qwen3.5-397B.", + "num_traces": 3, + "source": "sanitized Qwen3.6 static replay summary", + "total_tokens": 192, + "trace_mode": "sglang_generation" + }, + "throughput": { + "candidate": "qwen36_deepep_mbs6", + "global_batch_size": 192, + "gpu_alloc_gb": 46.08, + "gpus": 32, + "mfu_percent": 15.84, + "micro_batch_size": 6, + "step_time_sec": 6.178, + "tokens_per_sec": 254600.0, + "tokens_per_sec_per_gpu": 7956.25 + }, + "workload": "synthetic packed 8k full training, 4 nodes x 8 H100" +} diff --git a/src/xorl/sim/collect_calibration.py b/src/xorl/sim/collect_calibration.py new file mode 100644 index 00000000..3c21160c --- /dev/null +++ b/src/xorl/sim/collect_calibration.py @@ -0,0 +1,438 @@ +"""Parse XoRL trainer structured logs into calibration observations.""" + +from __future__ import annotations + +import argparse +import json +import re +import statistics +from pathlib import Path +from typing import Any, Iterable + + +try: + from .schemas import MemoryPhaseObservation, ObservedRun, PhaseObservation, StepObservation, to_jsonable +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import MemoryPhaseObservation, ObservedRun, PhaseObservation, StepObservation, to_jsonable + + +STEP_RE = re.compile(r"\[STEP\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") +PHASE_RE = re.compile(r"\[(?PSTEP_PHASES(?:_PARTIAL)?)\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") +MEMORY_RE = re.compile(r"\[(?PSTEP_MEMORY(?:_PARTIAL)?)\s+(?P\d+)/(?P[^\]]+)\]\s+(?P.*)") +KV_RE = re.compile(r"(?P[A-Za-z0-9_+./-]+)=(?P\S+)") + + +def _float_or_none(value: str | None) -> float | None: + if value is None: + return None + cleaned = value.strip().rstrip(",") + for suffix in ("GB", "gb", "s"): + if cleaned.endswith(suffix): + cleaned = cleaned[: -len(suffix)] + break + try: + return float(cleaned) + except ValueError: + return None + + +def _parse_metric_body(body: str) -> dict[str, float]: + metrics: dict[str, float] = {} + for match in KV_RE.finditer(body): + numeric = _float_or_none(match.group("value")) + if numeric is not None: + metrics[match.group("key")] = numeric + return metrics + + +def _step_from_match(match: re.Match[str], source: str) -> StepObservation: + metrics = _parse_metric_body(match.group("body")) + phase_memory: dict[str, float] = {} + for key in ("fwd", "bwd", "optim", "fwd+bwd", "offload"): + if key in metrics: + phase_memory[key] = metrics[key] + + known_keys = { + "loss", + "grad_norm", + "lr", + "tflops", + "mfu", + "tokens_per_sec", + "time", + "peak_mem", + "fwd", + "bwd", + "optim", + "fwd+bwd", + "offload", + } + extra = {key: value for key, value in metrics.items() if key not in known_keys} + return StepObservation( + source=source, + step=int(match.group("step")), + max_steps=match.group("max"), + loss=metrics.get("loss"), + grad_norm=metrics.get("grad_norm"), + lr=metrics.get("lr"), + tflops_per_gpu=metrics.get("tflops"), + mfu=metrics.get("mfu"), + tokens_per_sec=metrics.get("tokens_per_sec"), + step_time_s=metrics.get("time"), + peak_mem_gb=metrics.get("peak_mem"), + phase_memory_gb=phase_memory, + extra=extra, + ) + + +def parse_log_text(text: str, *, source: str = "") -> ObservedRun: + steps: list[StepObservation] = [] + phases: list[PhaseObservation] = [] + memory_phases: list[MemoryPhaseObservation] = [] + + for line in text.splitlines(): + if phase_match := PHASE_RE.search(line): + phases.append( + PhaseObservation( + source=source, + prefix=phase_match.group("prefix"), + step=int(phase_match.group("step")), + max_steps=phase_match.group("max"), + metrics=_parse_metric_body(phase_match.group("body")), + ) + ) + continue + + if memory_match := MEMORY_RE.search(line): + memory_phases.append( + MemoryPhaseObservation( + source=source, + prefix=memory_match.group("prefix"), + step=int(memory_match.group("step")), + max_steps=memory_match.group("max"), + metrics=_parse_metric_body(memory_match.group("body")), + ) + ) + continue + + if step_match := STEP_RE.search(line): + steps.append(_step_from_match(step_match, source)) + + return ObservedRun(sources=[source], steps=steps, phases=phases, memory_phases=memory_phases) + + +def parse_log_path(path: str | Path) -> ObservedRun: + log_path = Path(path) + text = log_path.read_text(encoding="utf-8", errors="replace") + return parse_log_text(text, source=str(log_path)) + + +def merge_observed_runs(runs: Iterable[ObservedRun]) -> ObservedRun: + sources: list[str] = [] + steps: list[StepObservation] = [] + phases: list[PhaseObservation] = [] + memory_phases: list[MemoryPhaseObservation] = [] + for run in runs: + sources.extend(run.sources) + steps.extend(run.steps) + phases.extend(run.phases) + memory_phases.extend(run.memory_phases) + return ObservedRun(sources=sources, steps=steps, phases=phases, memory_phases=memory_phases) + + +def _mean(values: list[float]) -> float | None: + return statistics.fmean(values) if values else None + + +def _median(values: list[float]) -> float | None: + return statistics.median(values) if values else None + + +def _stdev(values: list[float]) -> float | None: + return statistics.stdev(values) if len(values) >= 2 else None + + +def _coefficient_of_variation(values: list[float]) -> float | None: + mean = _mean(values) + stdev = _stdev(values) + if mean in (None, 0.0) or stdev is None: + return None + return stdev / mean + + +def _phase_metric_name(key: str, suffix: str) -> str | None: + if not key.endswith(suffix): + return None + return key[: -len(suffix)] + + +def _is_composite_phase_for_bottleneck(phase: str, phases: set[str]) -> bool: + lowered = phase.lower() + lowered_phases = {item.lower() for item in phases} + if lowered == "train_step_total": + return True + if lowered in {"forward_backward", "forward_backward_total", "fwd_bwd", "fwd_bwd_total"}: + return bool( + lowered_phases + & { + "model_forward", + "forward", + "fwd", + "loss_compute", + "loss", + "backward", + "model_backward", + "bwd", + } + ) + if lowered == "clip_and_step_total": + return bool( + lowered_phases + & { + "clip_gradients", + "optimizer_step", + "optimizer", + "optim", + "lr_scheduler_step", + } + ) + return False + + +def _phase_timing_summary( + run: ObservedRun, + *, + measured_step_keys: set[tuple[str, int]], + warmup_steps: int, +) -> dict[str, Any]: + if measured_step_keys: + phase_rows = [row for row in run.phases if (row.source, row.step) in measured_step_keys] + else: + phase_rows = sorted(run.phases, key=lambda row: (row.source, row.step))[warmup_steps:] + max_values: dict[str, list[float]] = {} + mean_values: dict[str, list[float]] = {} + for row in phase_rows: + for key, value in row.metrics.items(): + if phase := _phase_metric_name(key, "_max_s"): + max_values.setdefault(phase, []).append(value) + elif phase := _phase_metric_name(key, "_mean_s"): + mean_values.setdefault(phase, []).append(value) + + phase_time_sec: dict[str, float] = {} + phase_time_max_sec: dict[str, float] = {} + phase_time_rank_mean_sec: dict[str, float] = {} + for phase in sorted(set(max_values) | set(mean_values)): + values = max_values.get(phase) or mean_values.get(phase) or [] + if not values: + continue + # Median over post-warmup steps (stable profile): straggler steps a 2-step warmup cannot + # catch (e.g. a 0.586s optimizer step in a 0.20s-steady run) contaminate a mean read and + # made the phase pins disagree with the stable-profile reader; max stays available below. + phase_time_sec[phase] = statistics.median(values) + phase_time_max_sec[phase] = max(values) + # Cross-rank MEAN companion (same median-over-steps profile): phase_time_sec is the + # cross-rank MAX convention, so balanced-rank term comparisons need the mean to separate + # rank asymmetry (routing-imbalance stragglers) from term error. + rank_mean_rows = mean_values.get(phase) or [] + if rank_mean_rows: + phase_time_rank_mean_sec[phase] = statistics.median(rank_mean_rows) + + denominator = phase_time_sec.get("train_step_total") + if denominator is None: + denominator = sum(value for phase, value in phase_time_sec.items() if phase != "train_step_total") + phase_time_share = { + phase: value / denominator + for phase, value in phase_time_sec.items() + if denominator and phase != "train_step_total" + } + bottleneck_phase = None + bottleneck_candidates = { + phase: value + for phase, value in phase_time_share.items() + if not _is_composite_phase_for_bottleneck(phase, set(phase_time_share)) + } + if bottleneck_candidates: + bottleneck_phase = max(bottleneck_candidates, key=bottleneck_candidates.get) + elif phase_time_share: + bottleneck_phase = max(phase_time_share, key=phase_time_share.get) + + return { + "phase_time_sec": phase_time_sec, + "phase_time_max_sec": phase_time_max_sec, + "phase_time_rank_mean_sec": phase_time_rank_mean_sec, + "phase_time_share": phase_time_share, + "phase_bottleneck": bottleneck_phase, + } + + +def _phase_memory_metric_name(key: str) -> str | None: + for suffix in ( + "_phase_peak_allocated_max_gb", + "_phase_peak_reserved_max_gb", + ): + if key.endswith(suffix): + return key[: -len(suffix)] + if "_after_" in key or "_delta_" in key: + return None + for suffix in ("_allocated_max_gb", "_reserved_max_gb"): + if key.endswith(suffix): + return key[: -len(suffix)] + return None + + +def _phase_memory_delta_allows_peak(metrics: dict[str, float], phase: str) -> bool: + delta_values = [ + value + for key, value in metrics.items() + if key.startswith(f"{phase}_delta_") and key.endswith(("_allocated_max_gb", "_reserved_max_gb")) + ] + return not delta_values or max(delta_values) > 0.0 + + +def _phase_memory_summary( + run: ObservedRun, + *, + measured_step_keys: set[tuple[str, int]], + warmup_steps: int, + peak_mem_gb_max: float | None, +) -> dict[str, Any]: + if measured_step_keys: + step_rows = [row for row in run.steps if (row.source, row.step) in measured_step_keys] + memory_rows = [row for row in run.memory_phases if (row.source, row.step) in measured_step_keys] + else: + step_rows = sorted(run.steps, key=lambda row: (row.source, row.step))[warmup_steps:] + memory_rows = sorted(run.memory_phases, key=lambda row: (row.source, row.step))[warmup_steps:] + + peak_values: dict[str, list[float]] = {} + for row in step_rows: + for phase, value in row.phase_memory_gb.items(): + peak_values.setdefault(phase, []).append(value) + for row in memory_rows: + for key, value in row.metrics.items(): + if phase := _phase_memory_metric_name(key): + if not _phase_memory_delta_allows_peak(row.metrics, phase): + continue + peak_values.setdefault(phase, []).append(value) + + phase_memory_peak_gb = {phase: max(values) for phase, values in sorted(peak_values.items()) if values} + phase_memory_fraction_of_peak = { + phase: value / peak_mem_gb_max + for phase, value in phase_memory_peak_gb.items() + if peak_mem_gb_max not in (None, 0.0) + } + memory_bottleneck_phase = None + memory_bottleneck_candidates = { + phase: value + for phase, value in phase_memory_peak_gb.items() + if not _is_composite_phase_for_bottleneck(phase, set(phase_memory_peak_gb)) + } + if memory_bottleneck_candidates: + memory_bottleneck_phase = max(memory_bottleneck_candidates.items(), key=lambda item: (item[1], item[0]))[0] + elif phase_memory_peak_gb: + memory_bottleneck_phase = max(phase_memory_peak_gb.items(), key=lambda item: (item[1], item[0]))[0] + return { + "phase_memory_peak_gb": phase_memory_peak_gb, + "phase_memory_fraction_of_peak": phase_memory_fraction_of_peak, + "memory_bottleneck_phase": memory_bottleneck_phase, + } + + +def summarize_observed_run( + run: ObservedRun, + *, + warmup_steps: int = 0, + world_size: int | None = None, +) -> dict[str, Any]: + ordered_steps = sorted(run.steps, key=lambda row: (row.source, row.step)) + measured = ordered_steps[warmup_steps:] + measured_step_keys = {(row.source, row.step) for row in measured} + tps = [row.tokens_per_sec for row in measured if row.tokens_per_sec is not None] + tflops = [row.tflops_per_gpu for row in measured if row.tflops_per_gpu is not None] + mfu = [row.mfu for row in measured if row.mfu is not None] + step_time = [row.step_time_s for row in measured if row.step_time_s is not None] + peaks = [row.peak_mem_gb for row in measured if row.peak_mem_gb is not None] + # Realized tokens per step, computed PER STEP (tokens_per_sec x step_time of the SAME step) before + # aggregating: mean(tps) x median(step) mixes bases and mean(tps) x mean(step) carries the + # tps/step-time anti-correlation bias (a straggler step drags mean tps down while inflating mean + # time) — both misprice the realized token load that per-step phase predictions consume. + tokens_per_step = [ + row.tokens_per_sec * row.step_time_s + for row in measured + if row.tokens_per_sec is not None and row.step_time_s is not None + ] + + summary: dict[str, Any] = { + "sources": run.sources, + "parsed_step_count": len(run.steps), + "parsed_phase_count": len(run.phases), + "parsed_memory_phase_count": len(run.memory_phases), + "warmup_excluded": warmup_steps, + "measured_steps": len(measured), + "tokens_per_sec_mean": _mean(tps), + "tokens_per_sec_median": _median(tps), + "tokens_per_sec_std": _stdev(tps), + "tokens_per_sec_cv": _coefficient_of_variation(tps), + "tflops_per_gpu_mean": _mean(tflops), + "mfu_mean": _mean(mfu), + "step_time_s_mean": _mean(step_time), + "step_time_s_std": _stdev(step_time), + "step_time_s_cv": _coefficient_of_variation(step_time), + "tokens_per_step_median": _median(tokens_per_step), + "tokens_per_step_mean": _mean(tokens_per_step), + "peak_mem_gb_max": max(peaks) if peaks else None, + } + summary.update(_phase_timing_summary(run, measured_step_keys=measured_step_keys, warmup_steps=warmup_steps)) + summary.update( + _phase_memory_summary( + run, + measured_step_keys=measured_step_keys, + warmup_steps=warmup_steps, + peak_mem_gb_max=summary["peak_mem_gb_max"], + ) + ) + if world_size and summary["tokens_per_sec_mean"] is not None: + summary["tokens_per_sec_per_gpu_mean"] = summary["tokens_per_sec_mean"] / world_size + if measured: + summary["first_measured_step"] = measured[0].step + summary["last_measured_step"] = measured[-1].step + summary["loss_last"] = measured[-1].loss + return summary + + +def _expand_paths(paths: list[Path]) -> list[Path]: + expanded: list[Path] = [] + for path in paths: + if path.is_dir(): + expanded.extend(sorted(child for child in path.rglob("*") if child.is_file())) + else: + expanded.append(path) + return expanded + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", nargs="+", type=Path, help="Log files or directories to parse") + parser.add_argument("--warmup-steps", type=int, default=0, help="Drop this many parsed [STEP] rows from summary") + parser.add_argument("--world-size", type=int, default=None, help="Optional GPU count for per-GPU throughput") + parser.add_argument("--output", type=Path, default=None, help="Write JSON output to this path") + parser.add_argument("--include-rows", action="store_true", help="Include parsed row details, not just the summary") + args = parser.parse_args() + + runs = [parse_log_path(path) for path in _expand_paths(args.paths)] + observed = merge_observed_runs(runs) + payload: dict[str, Any] = { + "summary": summarize_observed_run(observed, warmup_steps=args.warmup_steps, world_size=args.world_size) + } + if args.include_rows: + payload["observed"] = to_jsonable(observed) + + rendered = json.dumps(to_jsonable(payload), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/experiments/local_benchmark/training_sim/config_fingerprint.py b/src/xorl/sim/config_fingerprint.py similarity index 93% rename from experiments/local_benchmark/training_sim/config_fingerprint.py rename to src/xorl/sim/config_fingerprint.py index 428840dd..12e4c467 100644 --- a/experiments/local_benchmark/training_sim/config_fingerprint.py +++ b/src/xorl/sim/config_fingerprint.py @@ -14,9 +14,11 @@ try: from .model_metadata import resolve_model_metadata from .schemas import RunFingerprint, Topology + from .simulator_support import requested_simulator_surface except ImportError: # pragma: no cover - exercised by direct script execution from model_metadata import resolve_model_metadata from schemas import RunFingerprint, Topology + from simulator_support import requested_simulator_surface REPO_ROOT = Path(__file__).resolve().parents[3] @@ -48,6 +50,21 @@ def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: return value if isinstance(value, dict) else {} +def _topology_sections(raw_config: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: + train = _section(raw_config, "train") + data = _section(raw_config, "data") + if train: + return train, data + + if requested_simulator_surface(raw_config) == "server_forward_backward": + server = _section(raw_config, "server") + if server: + return server, server + return raw_config, raw_config + + return train, data + + def _int_value(section: dict[str, Any], key: str, default: int | None = None) -> int | None: value = section.get(key, default) if value is None: @@ -121,8 +138,7 @@ def resolve_topology( num_experts: int | None = None, top_k: int | None = None, ) -> Topology: - train = _section(raw_config, "train") - data = _section(raw_config, "data") + train, data = _topology_sections(raw_config) ulysses = _int_value(train, "ulysses_parallel_size", 1) or 1 ringattn = _int_value(train, "ringattn_parallel_size", 1) or 1 diff --git a/src/xorl/sim/feasibility_evaluator.py b/src/xorl/sim/feasibility_evaluator.py new file mode 100644 index 00000000..13b5c69a --- /dev/null +++ b/src/xorl/sim/feasibility_evaluator.py @@ -0,0 +1,428 @@ +"""Evaluate whether simulator memory feasibility predicts fit versus OOM rows.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + + +try: + from .benchmark_behavior import ( + behavior_point_model_mismatches, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from .calibration_evaluator import _actual_memory_residual, _topology_for_point, _without_point + from .calibration_packs import resolve_calibration_pack, resolve_pack_inputs + from .config_fingerprint import load_training_config, resolve_topology + from .memory_ledger import build_memory_ledger + from .model_metadata import resolve_model_metadata + from .scenario_planner import ( + _calibrated_memory_peak_estimate, + _candidate_from_prediction, + _communication_ledger, + _extrapolate_behavior, + _memory_ownership_notes, + _topology_label, + ) + from .schemas import BenchmarkBehaviorPoint, FeasibilityHoldout, FeasibilityReport, to_jsonable + from .shape_engine import build_shape_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from benchmark_behavior import ( + behavior_point_model_mismatches, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from calibration_evaluator import _actual_memory_residual, _topology_for_point, _without_point + from calibration_packs import resolve_calibration_pack, resolve_pack_inputs + from config_fingerprint import load_training_config, resolve_topology + from memory_ledger import build_memory_ledger + from model_metadata import resolve_model_metadata + from scenario_planner import ( + _calibrated_memory_peak_estimate, + _candidate_from_prediction, + _communication_ledger, + _extrapolate_behavior, + _memory_ownership_notes, + _topology_label, + ) + from schemas import BenchmarkBehaviorPoint, FeasibilityHoldout, FeasibilityReport, to_jsonable + from shape_engine import build_shape_ledger + + +def _actual_outcome(point: BenchmarkBehaviorPoint) -> str | None: + if point.correctness_status == "oom": + return "oom" + if point.tokens_per_sec is not None: + return "fit" + return None + + +def _has_simulator_only_runtime_mismatch(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> bool: + return "attention_backend" in behavior_point_workload_mismatches(point, raw_config) + + +def _predicted_outcome(feasibility_status: str, risk_flags: list[str]) -> str: + if feasibility_status.startswith("feasible"): + memory_sensitive_runtime_mismatch = any( + flag + in { + "runtime_mismatch:activation_offload_prefetch_count", + "runtime_mismatch:deepep_buffer_size_gb", + "runtime_mismatch:enable_activation_offload", + "runtime_mismatch:gradient_checkpointing_method", + "runtime_mismatch:muon_momentum", + "runtime_mismatch:skip_param_upcast", + } + for flag in risk_flags + ) + if memory_sensitive_runtime_mismatch and ( + feasibility_status.endswith("_high_pressure") + or feasibility_status.endswith("_moderate_pressure") + or "memory_extrapolated_overhead" in risk_flags + ): + return "unknown" + if "real_routing_outside_fit_envelope" in risk_flags and ( + feasibility_status.endswith("_high_pressure") + or feasibility_status.endswith("_moderate_pressure") + or "memory_extrapolated_overhead" in risk_flags + ): + return "unknown" + return "fit" + if feasibility_status == "observed_oom" or feasibility_status.endswith("_exceeds_limit"): + return "blocked" + if feasibility_status in {"memory_floor_exceeds_limit", "memory_floor_exceeds_safety_margin"}: + return "blocked" + return "unknown" + + +def _max_holdout_by_field(holdouts: list[FeasibilityHoldout], field_name: str) -> FeasibilityHoldout | None: + return max( + (holdout for holdout in holdouts if getattr(holdout, field_name) is not None), + key=lambda holdout: (getattr(holdout, field_name), holdout.label), + default=None, + ) + + +def _count_values(values: list[str]) -> dict[str, int]: + counts: dict[str, int] = {} + for value in values: + counts[value] = counts.get(value, 0) + 1 + return dict(sorted(counts.items())) + + +def _classify_correct(actual_outcome: str, predicted_outcome: str) -> bool: + if actual_outcome == "fit": + return predicted_outcome == "fit" + if actual_outcome == "oom": + return predicted_outcome == "blocked" + return False + + +def _fit_recall(holdouts: list[FeasibilityHoldout]) -> float | None: + fit_holdouts = [holdout for holdout in holdouts if holdout.actual_outcome == "fit"] + if not fit_holdouts: + return None + correct = sum(1 for holdout in fit_holdouts if holdout.classified_correctly) + return round(correct / len(fit_holdouts), 3) + + +def _oom_recall(holdouts: list[FeasibilityHoldout]) -> float | None: + oom_holdouts = [holdout for holdout in holdouts if holdout.actual_outcome == "oom"] + if not oom_holdouts: + return None + correct = sum(1 for holdout in oom_holdouts if holdout.classified_correctly) + return round(correct / len(oom_holdouts), 3) + + +def _feasibility_holdout( + heldout: BenchmarkBehaviorPoint, + *, + behavior_points: list[BenchmarkBehaviorPoint], + base_config: dict[str, Any], + base_topology, + world_size: int | None, + local_world_size: int | None, + device_memory_limit_gb: float, + memory_safety_factor: float, +) -> tuple[FeasibilityHoldout | None, str | None]: + actual_outcome = _actual_outcome(heldout) + if actual_outcome is None: + return None, "missing tokens_per_sec and not an OOM row" + raw_config, topology, skip_reason = _topology_for_point( + base_config, + base_topology, + heldout, + world_size=world_size, + local_world_size=local_world_size, + require_tokens=False, + ) + if raw_config is None or topology is None: + return None, skip_reason + + training_points = _without_point(behavior_points, heldout) + shape = build_shape_ledger(topology, balanced_routing=True) + metadata = resolve_model_metadata(raw_config) + memory = build_memory_ledger(raw_config, topology=topology, model_metadata=metadata) + exact_prediction = predict_benchmark_behavior(training_points, topology, shape, raw_config) + if exact_prediction.status in {"calibrated", "calibrated_failure"}: + behavior = exact_prediction + prediction_confidence = exact_prediction.status + memory_peak_estimate = None + notes: list[str] = [] + else: + memory_peak_estimate = _calibrated_memory_peak_estimate( + training_points, + base_config, + raw_config, + topology, + shape, + metadata, + default_world_size=base_topology.world_size, + default_local_world_size=base_topology.local_world_size, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + ) + behavior, notes = _extrapolate_behavior( + training_points, + topology, + shape, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + ) + prediction_confidence = behavior.status + + candidate = _candidate_from_prediction( + label=heldout.label, + config_path=None, + topology=topology, + shape=shape, + behavior=behavior, + prediction_confidence=prediction_confidence, + promotable=False, + behavior_points=training_points, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + memory_ownership_notes=_memory_ownership_notes(memory), + communication=_communication_ledger(topology), + notes=notes, + ) + predicted_outcome = _predicted_outcome(candidate.feasibility_status, candidate.risk_flags) + actual_memory_residual_gb, actual_memory_residual_fraction = _actual_memory_residual( + heldout.peak_mem_gb, + candidate.analytic_peak_floor_gb, + ) + return ( + FeasibilityHoldout( + label=heldout.label, + source=heldout.source, + topology_label=_topology_label(topology), + actual_outcome=actual_outcome, + predicted_outcome=predicted_outcome, + actual_tokens_per_sec=heldout.tokens_per_sec, + actual_peak_mem_gb=heldout.peak_mem_gb, + predicted_tokens_per_sec=behavior.tokens_per_sec, + predicted_peak_mem_gb=candidate.estimated_peak_mem_gb, + prediction_status=behavior.status, + matched_label=behavior.matched_label, + memory_prediction_basis=candidate.memory_basis, + analytic_peak_floor_gb=candidate.analytic_peak_floor_gb, + memory_coverage_status=candidate.memory_coverage_status, + predicted_memory_residual_gb=candidate.estimated_memory_residual_gb, + predicted_memory_residual_fraction_of_peak=candidate.estimated_memory_residual_fraction_of_peak, + actual_memory_residual_gb=actual_memory_residual_gb, + actual_memory_residual_fraction_of_peak=actual_memory_residual_fraction, + memory_calibration_source=candidate.memory_calibration_source, + predicted_feasibility_status=candidate.feasibility_status, + classified_correctly=_classify_correct(actual_outcome, predicted_outcome), + calibrated_from_count=len(training_points), + memory_calibration_notes=candidate.memory_calibration_notes, + risk_flags=candidate.risk_flags, + warnings=behavior.warnings, + ), + None, + ) + + +def evaluate_feasibility( + base_config_path: str | Path, + *, + benchmark_dir: str | Path, + world_size: int | None = None, + local_world_size: int | None = None, + device_memory_limit_gb: float = 80.0, + memory_safety_factor: float = 1.15, +) -> FeasibilityReport: + base_path = Path(base_config_path) + benchmark_path = resolve_calibration_pack(benchmark_dir) + base_config = load_training_config(base_path) + base_topology = resolve_topology(base_config, world_size=world_size, local_world_size=local_world_size) + behavior_points = load_benchmark_behavior_points(benchmark_path) + observed_points = [ + point + for point in behavior_points + if _actual_outcome(point) is not None and not _has_simulator_only_runtime_mismatch(point, base_config) + ] + + holdouts: list[FeasibilityHoldout] = [] + warnings: list[str] = [] + skipped_count = 0 + for heldout in observed_points: + if behavior_point_model_mismatches(heldout, base_config): + skipped_count += 1 + warnings.append(f"skipped {heldout.label}: model_ref mismatch") + continue + holdout, skip_reason = _feasibility_holdout( + heldout, + behavior_points=behavior_points, + base_config=base_config, + base_topology=base_topology, + world_size=world_size, + local_world_size=local_world_size, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + if holdout is None: + skipped_count += 1 + warnings.append(f"skipped {heldout.label}: {skip_reason}") + continue + holdouts.append(holdout) + + correct_count = sum(1 for holdout in holdouts if holdout.classified_correctly) + actual_fit_count = sum(1 for holdout in holdouts if holdout.actual_outcome == "fit") + actual_oom_count = sum(1 for holdout in holdouts if holdout.actual_outcome == "oom") + predicted_fit_count = sum(1 for holdout in holdouts if holdout.predicted_outcome == "fit") + predicted_blocked_count = sum(1 for holdout in holdouts if holdout.predicted_outcome == "blocked") + predicted_unknown_count = sum(1 for holdout in holdouts if holdout.predicted_outcome == "unknown") + false_fit_count = sum( + 1 for holdout in holdouts if holdout.actual_outcome == "oom" and holdout.predicted_outcome == "fit" + ) + false_blocked_count = sum( + 1 for holdout in holdouts if holdout.actual_outcome == "fit" and holdout.predicted_outcome == "blocked" + ) + if predicted_unknown_count: + warnings.append(f"{predicted_unknown_count} feasibility holdouts were predicted unknown") + predicted_memory_residuals = [ + holdout.predicted_memory_residual_gb for holdout in holdouts if holdout.predicted_memory_residual_gb is not None + ] + predicted_memory_residual_fractions = [ + holdout.predicted_memory_residual_fraction_of_peak + for holdout in holdouts + if holdout.predicted_memory_residual_fraction_of_peak is not None + ] + actual_memory_residuals = [ + holdout.actual_memory_residual_gb for holdout in holdouts if holdout.actual_memory_residual_gb is not None + ] + actual_memory_residual_fractions = [ + holdout.actual_memory_residual_fraction_of_peak + for holdout in holdouts + if holdout.actual_memory_residual_fraction_of_peak is not None + ] + max_predicted_memory_residual_holdout = _max_holdout_by_field(holdouts, "predicted_memory_residual_gb") + max_predicted_memory_residual_fraction_holdout = _max_holdout_by_field( + holdouts, + "predicted_memory_residual_fraction_of_peak", + ) + max_actual_memory_residual_holdout = _max_holdout_by_field(holdouts, "actual_memory_residual_gb") + max_actual_memory_residual_fraction_holdout = _max_holdout_by_field( + holdouts, + "actual_memory_residual_fraction_of_peak", + ) + + return FeasibilityReport( + base_config_path=str(base_path), + benchmark_dir=str(benchmark_path), + status="ok" if holdouts else "insufficient_data", + observed_point_count=len(observed_points), + evaluated_count=len(holdouts), + skipped_count=skipped_count, + actual_fit_count=actual_fit_count, + actual_oom_count=actual_oom_count, + predicted_fit_count=predicted_fit_count, + predicted_blocked_count=predicted_blocked_count, + predicted_unknown_count=predicted_unknown_count, + correct_count=correct_count, + false_fit_count=false_fit_count, + false_blocked_count=false_blocked_count, + accuracy=round(correct_count / len(holdouts), 3) if holdouts else None, + fit_recall=_fit_recall(holdouts), + oom_recall=_oom_recall(holdouts), + prediction_status_counts=_count_values([holdout.prediction_status for holdout in holdouts]), + memory_prediction_basis_counts=_count_values([holdout.memory_prediction_basis for holdout in holdouts]), + memory_coverage_status_counts=_count_values([holdout.memory_coverage_status for holdout in holdouts]), + feasibility_status_counts=_count_values([holdout.predicted_feasibility_status for holdout in holdouts]), + max_predicted_memory_residual_gb=( + round(max(predicted_memory_residuals), 3) if predicted_memory_residuals else None + ), + max_predicted_memory_residual_gb_label=( + max_predicted_memory_residual_holdout.label if max_predicted_memory_residual_holdout is not None else None + ), + max_predicted_memory_residual_fraction_of_peak=( + round(max(predicted_memory_residual_fractions), 3) if predicted_memory_residual_fractions else None + ), + max_predicted_memory_residual_fraction_of_peak_label=( + max_predicted_memory_residual_fraction_holdout.label + if max_predicted_memory_residual_fraction_holdout is not None + else None + ), + max_actual_memory_residual_gb=round(max(actual_memory_residuals), 3) if actual_memory_residuals else None, + max_actual_memory_residual_gb_label=( + max_actual_memory_residual_holdout.label if max_actual_memory_residual_holdout is not None else None + ), + max_actual_memory_residual_fraction_of_peak=( + round(max(actual_memory_residual_fractions), 3) if actual_memory_residual_fractions else None + ), + max_actual_memory_residual_fraction_of_peak_label=( + max_actual_memory_residual_fraction_holdout.label + if max_actual_memory_residual_fraction_holdout is not None + else None + ), + risk_flag_counts=_count_values([flag for holdout in holdouts for flag in holdout.risk_flags]), + holdouts=holdouts, + warnings=warnings, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pack", help="Built-in calibration-pack name") + parser.add_argument("--config", type=Path, default=None) + parser.add_argument("--benchmark-dir", type=Path, default=None) + parser.add_argument("--world-size", type=int, default=None) + parser.add_argument("--local-world-size", type=int, default=None) + parser.add_argument("--device-memory-limit-gb", type=float, default=80.0) + parser.add_argument("--memory-safety-factor", type=float, default=1.15) + parser.add_argument("--output", type=Path, default=None) + args = parser.parse_args() + + args.config, args.benchmark_dir = resolve_pack_inputs(args.pack, args.config, args.benchmark_dir) + if args.config is None or args.benchmark_dir is None: + parser.error("provide --pack, or both --config and --benchmark-dir") + + report = evaluate_feasibility( + args.config, + benchmark_dir=args.benchmark_dir, + world_size=args.world_size, + local_world_size=args.local_world_size, + device_memory_limit_gb=args.device_memory_limit_gb, + memory_safety_factor=args.memory_safety_factor, + ) + rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + + +if __name__ == "__main__": + main() diff --git a/src/xorl/sim/kernel_variants.py b/src/xorl/sim/kernel_variants.py new file mode 100644 index 00000000..3c281448 --- /dev/null +++ b/src/xorl/sim/kernel_variants.py @@ -0,0 +1,108 @@ +"""Portable comparison helpers for measured kernel variants.""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class KernelVariantMeasurement: + family: str + variant: str + workload: str + latency_ms: float + correctness_status: str + tokens: int | None = None + peak_memory_gb: float | None = None + source: str | None = None + notes: list[str] | None = None + + @property + def promotable(self) -> bool: + return self.correctness_status in {"pass", "k3_pass", "validated"} + + +def _measurement(value: KernelVariantMeasurement | dict[str, Any]) -> KernelVariantMeasurement: + return value if isinstance(value, KernelVariantMeasurement) else KernelVariantMeasurement(**value) + + +def rank_kernel_variants( + measurements: list[KernelVariantMeasurement | dict[str, Any]], + *, + require_correctness: bool = True, +) -> dict[str, Any]: + """Rank like-for-like measured variants without loading infrastructure-specific artifacts.""" + + rows = [_measurement(value) for value in measurements] + if not rows: + return {"status": "no_measurements", "best": None, "measurements": []} + families = {row.family for row in rows} + workloads = {row.workload for row in rows} + if len(families) != 1 or len(workloads) != 1: + raise ValueError("kernel variants must share one family and one workload") + if any(row.latency_ms <= 0 for row in rows): + raise ValueError("kernel-variant latency must be positive") + + eligible = [row for row in rows if row.promotable or not require_correctness] + ranked = sorted(rows, key=lambda row: row.latency_ms) + best = min(eligible, key=lambda row: row.latency_ms) if eligible else None + baseline = max(row.latency_ms for row in rows) + rendered = [] + for row in ranked: + payload = asdict(row) + payload["promotable"] = row.promotable + payload["speedup_vs_slowest"] = round(baseline / row.latency_ms, 6) + rendered.append(payload) + return { + "status": "ok" if best is not None else "no_correctness_promotable_variant", + "family": next(iter(families)), + "workload": next(iter(workloads)), + "require_correctness": require_correctness, + "best": asdict(best) if best is not None else None, + "measurements": rendered, + } + + +def compare_kernel_variants( + baseline: KernelVariantMeasurement | dict[str, Any], + candidate: KernelVariantMeasurement | dict[str, Any], +) -> dict[str, Any]: + base = _measurement(baseline) + other = _measurement(candidate) + if (base.family, base.workload) != (other.family, other.workload): + raise ValueError("kernel variants must share one family and one workload") + return { + "family": base.family, + "workload": base.workload, + "baseline": base.variant, + "candidate": other.variant, + "latency_delta_ms": round(other.latency_ms - base.latency_ms, 6), + "latency_delta_percent": round((other.latency_ms / base.latency_ms - 1.0) * 100.0, 6), + "speedup": round(base.latency_ms / other.latency_ms, 6), + "peak_memory_delta_gb": ( + round(other.peak_memory_gb - base.peak_memory_gb, 6) + if base.peak_memory_gb is not None and other.peak_memory_gb is not None + else None + ), + "candidate_promotable": other.promotable, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("measurements", type=Path, help="JSON list of kernel-variant measurements") + parser.add_argument("--allow-ungated", action="store_true") + args = parser.parse_args() + rows = json.loads(args.measurements.read_text(encoding="utf-8")) + report = rank_kernel_variants(rows, require_correctness=not args.allow_ungated) + print(json.dumps(report, indent=2, sort_keys=True)) + if report["status"] != "ok": + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/src/xorl/sim/memory_ledger.py b/src/xorl/sim/memory_ledger.py new file mode 100644 index 00000000..64b3581f --- /dev/null +++ b/src/xorl/sim/memory_ledger.py @@ -0,0 +1,771 @@ +"""Initial memory ledger built from config constants and observed structured logs.""" + +from __future__ import annotations + +from typing import Any + + +try: + from .schemas import MemoryBucket, MemoryLedger, ModelMetadata, ObservedRun, Topology +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import MemoryBucket, MemoryLedger, ModelMetadata, ObservedRun, Topology + + +BYTES_PER_GIB = 1024**3 + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + return value if isinstance(value, dict) else {} + + +def _float_field(section: dict[str, Any], key: str) -> float | None: + value = section.get(key) + if value is None: + return None + return float(value) + + +def _dtype_bytes(dtype: Any, *, default: int) -> int: + if dtype is None: + return default + normalized = str(dtype).lower() + if normalized in {"bf16", "bfloat16", "fp16", "float16", "half"}: + return 2 + if normalized in {"fp32", "float32", "float"}: + return 4 + if normalized in {"fp8", "float8", "e4m3", "e5m2"}: + return 1 + return default + + +def _gb(byte_count: float) -> float: + return byte_count / BYTES_PER_GIB + + +def _round_gb(value: float | None) -> float | None: + return round(value, 3) if value is not None else None + + +def _round_fraction(value: float | None) -> float | None: + return round(value, 3) if value is not None else None + + +def _param_storage_bytes(train: dict[str, Any]) -> tuple[int, list[str]]: + explicit_param_dtype = train.get("param_dtype") + if explicit_param_dtype is not None: + return _dtype_bytes(explicit_param_dtype, default=4), [f"param_dtype={explicit_param_dtype}"] + + if train.get("enable_mixed_precision"): + if train.get("skip_param_upcast"): + return 2, ["mixed_precision_checkpoint_native_params", "skip_param_upcast=true"] + return 4, ["mixed_precision_generic_fp32_param_upcast", "skip_param_upcast=false"] + + return 4, ["mixed_precision=false"] + + +def _gradient_storage_bytes(train: dict[str, Any]) -> tuple[int, list[str]]: + explicit_gradient_dtype = train.get("gradient_dtype") + if explicit_gradient_dtype is not None: + return _dtype_bytes(explicit_gradient_dtype, default=4), [f"gradient_dtype={explicit_gradient_dtype}"] + + notes = ["gradient_storage_default=fp32"] + if train.get("fsdp_reduce_dtype") is not None: + notes.append(f"fsdp_reduce_dtype={train.get('fsdp_reduce_dtype')}:comm_buffer_only") + if train.get("muon_grad_dtype") is not None: + notes.append(f"muon_grad_dtype={train.get('muon_grad_dtype')}:optimizer_update_only") + return 4, notes + + +def _estimate_param_breakdown(metadata: ModelMetadata) -> dict[str, float] | None: + hidden = metadata.hidden_size + layers = metadata.num_hidden_layers + vocab = metadata.vocab_size + if hidden is None or layers is None or vocab is None: + return None + + head_dim = metadata.head_dim + if head_dim is None and metadata.num_attention_heads: + head_dim = hidden // metadata.num_attention_heads + attention_heads = metadata.num_attention_heads or 1 + key_value_heads = metadata.num_key_value_heads or attention_heads + if head_dim is None: + return None + + # Hybrid GatedDeltaNet models (Qwen3.5/3.6): only every full_attention_interval-th layer is full + # attention; the rest are GatedDeltaNet linear-attention layers with their own projections, + # depthwise short convolutions, and decay params. Gated attention doubles q_proj (query + output + # gate). For non-hybrid models this reduces exactly to the previous all-layers-attention formula. + interval = metadata.full_attention_interval + gdn_dims = ( + metadata.linear_num_key_heads, + metadata.linear_num_value_heads, + metadata.linear_key_head_dim, + metadata.linear_value_head_dim, + metadata.linear_conv_kernel_dim, + ) + is_hybrid = bool(interval) and all(value is not None for value in gdn_dims) + num_full_attention_layers = layers // int(interval) if is_hybrid else layers + num_linear_attention_layers = layers - num_full_attention_layers + gated = bool(metadata.attn_output_gate) + + q_proj = hidden * attention_heads * head_dim * (2 if gated else 1) + k_proj = hidden * key_value_heads * head_dim + v_proj = hidden * key_value_heads * head_dim + o_proj = attention_heads * head_dim * hidden + attention_params = num_full_attention_layers * (q_proj + k_proj + v_proj + o_proj) + + linear_attention_muon_params = 0 + linear_attention_fallback_params = 0 + if is_hybrid and num_linear_attention_layers > 0: + num_k_heads = int(metadata.linear_num_key_heads) + num_v_heads = int(metadata.linear_num_value_heads) + key_dim = num_k_heads * int(metadata.linear_key_head_dim) + value_dim = num_v_heads * int(metadata.linear_value_head_dim) + conv_k = int(metadata.linear_conv_kernel_dim) + # Muon side (trainer ndim>=2 classifier, no exclusion match): q/k/v[/g]/a/b/o projections and + # the three depthwise conv weights ([C, 1, k], conv_bias=False). + gdn_proj = ( + key_dim * hidden # q_proj + + key_dim * hidden # k_proj + + value_dim * hidden # v_proj + + (value_dim * hidden if gated else 0) # g_proj + + num_v_heads * hidden # a_proj + + num_v_heads * hidden # b_proj + + value_dim * hidden # o_proj + ) + gdn_convs = (2 * key_dim + value_dim) * conv_k + linear_attention_muon_params = num_linear_attention_layers * (gdn_proj + gdn_convs) + # Fallback side: A_log + dt_bias (1D, fp32-pinned) and the gated output norm (head_v_dim). + gdn_small = 2 * num_v_heads + int(metadata.linear_value_head_dim) + linear_attention_fallback_params = num_linear_attention_layers * gdn_small + linear_attention_params = linear_attention_muon_params + linear_attention_fallback_params + + dense_mlp_params = 0 + has_routed_experts = metadata.num_experts is not None and metadata.moe_intermediate_size is not None + if metadata.intermediate_size is not None and not has_routed_experts: + dense_mlp_params = layers * 3 * hidden * metadata.intermediate_size + + shared_expert_params = 0 + if metadata.shared_expert_intermediate_size is not None: + shared_expert_params = layers * 3 * hidden * metadata.shared_expert_intermediate_size + + router_params = 0 + if has_routed_experts and metadata.num_experts is not None: + router_params = layers * hidden * metadata.num_experts + + expert_params = 0 + if has_routed_experts and metadata.num_experts is not None and metadata.moe_intermediate_size is not None: + expert_params = layers * metadata.num_experts * 3 * hidden * metadata.moe_intermediate_size + + embedding_params = vocab * hidden + lm_head_params = 0 if metadata.tie_word_embeddings else vocab * hidden + # q/k norms exist only on the full-attention layers; GDN layers carry their own output norm + # (already counted in linear_attention_fallback_params). + qk_norm_params = 2 * num_full_attention_layers * head_dim + layer_norm_params = 2 * layers * hidden + qk_norm_params + final_norm_params = hidden + non_expert_params = ( + attention_params + + linear_attention_params + + dense_mlp_params + + shared_expert_params + + router_params + + embedding_params + + lm_head_params + ) + non_expert_params += layer_norm_params + final_norm_params + return { + "total_params": float(non_expert_params + expert_params), + "non_expert_params": float(non_expert_params), + "expert_params": float(expert_params), + "layer_non_expert_params": float( + attention_params + + linear_attention_params + + dense_mlp_params + + shared_expert_params + + router_params + + layer_norm_params + ), + "attention_params": float(attention_params), + "linear_attention_params": float(linear_attention_params), + "linear_attention_muon_params": float(linear_attention_muon_params), + "linear_attention_fallback_params": float(linear_attention_fallback_params), + "num_full_attention_layers": float(num_full_attention_layers), + "num_linear_attention_layers": float(num_linear_attention_layers), + "dense_mlp_params": float(dense_mlp_params), + "shared_expert_params": float(shared_expert_params), + "router_params": float(router_params), + "embedding_params": float(embedding_params), + "lm_head_params": float(lm_head_params), + "qk_norm_params": float(qk_norm_params), + "layer_norm_params": float(layer_norm_params), + "final_norm_params": float(final_norm_params), + } + + +def _estimate_param_counts(metadata: ModelMetadata) -> tuple[float, float, float] | None: + breakdown = _estimate_param_breakdown(metadata) + if breakdown is None: + return None + return breakdown["total_params"], breakdown["non_expert_params"], breakdown["expert_params"] + + +def _non_expert_fsdp_shard_size(topology: Topology, train: dict[str, Any]) -> tuple[int, list[str]]: + shard_size = max(topology.data_parallel_shard_size, 1) + notes = [f"dp_shard_size={topology.data_parallel_shard_size}"] + cp_mode = str(train.get("cp_fsdp_mode", "all") or "all") + if cp_mode == "all": + shard_size *= max(topology.sequence_parallel_size, 1) + notes.append(f"cp_fsdp_mode=all:sequence_parallel_size={topology.sequence_parallel_size}") + elif cp_mode == "ulysses_only": + shard_size *= max(topology.ulysses_parallel_size, 1) + notes.append(f"cp_fsdp_mode=ulysses_only:ulysses_parallel_size={topology.ulysses_parallel_size}") + elif cp_mode == "ring_only": + shard_size *= max(topology.ringattn_parallel_size, 1) + notes.append(f"cp_fsdp_mode=ring_only:ringattn_parallel_size={topology.ringattn_parallel_size}") + else: + notes.append(f"cp_fsdp_mode={cp_mode}:no_cp_fsdp_fold") + return max(shard_size, 1), notes + + +def _local_param_ownership( + breakdown: dict[str, float], + train: dict[str, Any], + topology: Topology, +) -> tuple[float, float, float, list[str]]: + pp_size = max(topology.pipeline_parallel_size, 1) + tp_size = max(topology.tensor_parallel_size, 1) + if pp_size == 1: + stage_non_expert_params = breakdown["non_expert_params"] + stage_expert_params = breakdown["expert_params"] + endpoint_note = "pp_stage=max_all_layers_and_endpoints" + else: + layer_non_expert_stage = breakdown["layer_non_expert_params"] / pp_size + expert_stage = breakdown["expert_params"] / pp_size + first_stage_extra = breakdown["embedding_params"] + last_stage_extra = breakdown["lm_head_params"] + breakdown["final_norm_params"] + stage_non_expert_params = layer_non_expert_stage + max(first_stage_extra, last_stage_extra) + stage_expert_params = expert_stage + endpoint_note = ( + f"pp_stage=max(layer_non_expert/{pp_size}+embedding, layer_non_expert/{pp_size}+lm_head+final_norm)" + ) + + non_expert_fsdp_shard_size, fsdp_notes = _non_expert_fsdp_shard_size(topology, train) + non_expert_shard_size = non_expert_fsdp_shard_size * tp_size + expert_shard_size = topology.expert_parallel_size * (topology.ep_fsdp_size or 1) + local_non_expert_params = stage_non_expert_params / max(non_expert_shard_size, 1) + local_expert_params = stage_expert_params / max(expert_shard_size, 1) + notes = [ + endpoint_note, + f"tp_non_expert_shard_size={tp_size}", + f"non_expert_total_shard_size={non_expert_shard_size}", + f"expert_shard_size={expert_shard_size}", + *fsdp_notes, + ] + return local_non_expert_params + local_expert_params, local_non_expert_params, local_expert_params, notes + + +def _dense_muon_param_partition( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], + *, + local_params: float | None = None, +) -> dict[str, Any]: + """Exact dense-model split between Muon matrix params and fallback params.""" + needed = ( + metadata.hidden_size, + metadata.intermediate_size, + metadata.num_hidden_layers, + metadata.num_attention_heads, + ) + if any(value is None for value in needed): + return {"status": "unsupported", "reason": "missing_dense_muon_partition_metadata"} + if metadata.num_experts is not None or metadata.moe_intermediate_size is not None: + return {"status": "unsupported", "reason": "dense_muon_partition_only"} + + breakdown = _estimate_param_breakdown(metadata) + if breakdown is None: + return {"status": "unsupported", "reason": "param_breakdown_unavailable"} + if local_params is None: + local_params, _, _, _ = _local_param_ownership(breakdown, train, topology) + + hidden = int(metadata.hidden_size) + intermediate = int(metadata.intermediate_size) + layers = int(metadata.num_hidden_layers) + n_heads = int(metadata.num_attention_heads) + n_kv = int(metadata.num_key_value_heads or n_heads) + head_dim = int(metadata.head_dim or hidden // n_heads) + pp_size = max(int(topology.pipeline_parallel_size), 1) + tp_size = max(int(topology.tensor_parallel_size), 1) + non_expert_fsdp_shard_size, fsdp_notes = _non_expert_fsdp_shard_size(topology, train) + total_non_expert_shard_size = max(non_expert_fsdp_shard_size * tp_size, 1) + + per_layer_matrix_shapes = [ + ("q_proj", n_heads * head_dim, hidden), + ("k_proj", n_kv * head_dim, hidden), + ("v_proj", n_kv * head_dim, hidden), + ("o_proj", n_heads * head_dim, hidden), + ("gate_proj", intermediate, hidden), + ("up_proj", intermediate, hidden), + ("down_proj", hidden, intermediate), + ] + per_layer_matrix_params = sum(rows * cols for _, rows, cols in per_layer_matrix_shapes) + stage_matrix_params = per_layer_matrix_params * layers / pp_size + local_muon_matrix_params = stage_matrix_params / total_non_expert_shard_size + local_fallback_params = max(float(local_params) - local_muon_matrix_params, 0.0) + + return { + "status": "exact_analytic_dense_muon_partition", + "local_params": round(float(local_params)), + "local_muon_matrix_params": round(local_muon_matrix_params), + "local_fallback_params": round(local_fallback_params), + "pipeline_parallel_size": pp_size, + "tensor_parallel_size": tp_size, + "non_expert_fsdp_shard_size": non_expert_fsdp_shard_size, + "total_non_expert_shard_size": total_non_expert_shard_size, + "matrix_entry_count": round(len(per_layer_matrix_shapes) * layers / pp_size), + "matrix_shapes_per_layer": [ + {"name": name, "rows": rows, "cols": cols, "numel": rows * cols} + for name, rows, cols in per_layer_matrix_shapes + ], + "fallback_optimizer": str(train.get("muon_fallback_optimizer", "adamw") or "adamw"), + "notes": [ + "matches src/xorl/optim/optimizer.py Muon ndim>=2 classifier for dense Qwen projections", + "fallback bucket covers embeddings/lm_head/norms and other excluded or non-matrix params", + *fsdp_notes, + ], + } + + +def _moe_muon_param_partition( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], + *, + local_params: float | None = None, +) -> dict[str, Any]: + """Exact MoE split between Muon matrix params and fallback params for qwen-style MoE. + + This mirrors ``src/xorl/optim/optimizer.py::_classify_muon_params`` at the + metadata level for the q30/q35 lane: embeddings, lm_head, norms, and + router ``gate.weight`` use the fallback optimizer; attention projections, + shared-expert projections, and routed expert tensors use Muon. + """ + needed = ( + metadata.hidden_size, + metadata.num_hidden_layers, + metadata.num_attention_heads, + metadata.num_experts, + metadata.moe_intermediate_size, + ) + if any(value is None for value in needed): + return {"status": "unsupported", "reason": "missing_moe_muon_partition_metadata"} + if metadata.num_experts is None or metadata.moe_intermediate_size is None: + return {"status": "unsupported", "reason": "moe_muon_partition_only"} + + breakdown = _estimate_param_breakdown(metadata) + if breakdown is None: + return {"status": "unsupported", "reason": "param_breakdown_unavailable"} + + pp_size = max(int(topology.pipeline_parallel_size), 1) + if pp_size != 1: + return {"status": "unsupported", "reason": "moe_muon_partition_pp_stage_not_enabled"} + + if local_params is None: + local_params, _, _, _ = _local_param_ownership(breakdown, train, topology) + + tp_size = max(int(topology.tensor_parallel_size), 1) + non_expert_fsdp_shard_size, fsdp_notes = _non_expert_fsdp_shard_size(topology, train) + total_non_expert_shard_size = max(non_expert_fsdp_shard_size * tp_size, 1) + expert_shard_size = max(topology.expert_parallel_size * (topology.ep_fsdp_size or 1), 1) + + non_expert_muon_params = ( + breakdown["attention_params"] + + breakdown.get("linear_attention_muon_params", 0.0) + + breakdown["dense_mlp_params"] + + breakdown["shared_expert_params"] + ) + non_expert_fallback_params = ( + breakdown["embedding_params"] + + breakdown["lm_head_params"] + + breakdown["layer_norm_params"] + + breakdown["final_norm_params"] + + breakdown["router_params"] + + breakdown.get("linear_attention_fallback_params", 0.0) + ) + local_non_expert_muon_params = non_expert_muon_params / total_non_expert_shard_size + local_non_expert_fallback_params = non_expert_fallback_params / total_non_expert_shard_size + local_expert_muon_params = breakdown["expert_params"] / expert_shard_size + local_muon_matrix_params = local_non_expert_muon_params + local_expert_muon_params + local_fallback_params = local_non_expert_fallback_params + partitioned_local_params = local_muon_matrix_params + local_fallback_params + + hidden = int(metadata.hidden_size) + n_heads = int(metadata.num_attention_heads) + n_kv = int(metadata.num_key_value_heads or n_heads) + head_dim = int(metadata.head_dim or hidden // n_heads) + moe_intermediate = int(metadata.moe_intermediate_size) + shared_intermediate = int(metadata.shared_expert_intermediate_size or 0) + num_experts = int(metadata.num_experts) + + q_proj_rows = n_heads * head_dim * (2 if metadata.attn_output_gate else 1) + matrix_shapes_per_layer = [ + {"name": "q_proj", "rows": q_proj_rows, "cols": hidden, "numel": hidden * q_proj_rows}, + {"name": "k_proj", "rows": n_kv * head_dim, "cols": hidden, "numel": hidden * n_kv * head_dim}, + {"name": "v_proj", "rows": n_kv * head_dim, "cols": hidden, "numel": hidden * n_kv * head_dim}, + {"name": "o_proj", "rows": n_heads * head_dim, "cols": hidden, "numel": hidden * n_heads * head_dim}, + { + "name": "routed_expert_gate_up_down", + "rows": num_experts, + "cols": 3 * hidden * moe_intermediate, + "numel": num_experts * 3 * hidden * moe_intermediate, + }, + ] + if shared_intermediate: + matrix_shapes_per_layer.append( + { + "name": "shared_expert_gate_up_down", + "rows": 3 * shared_intermediate, + "cols": hidden, + "numel": 3 * hidden * shared_intermediate, + } + ) + + return { + "status": "exact_analytic_moe_muon_partition", + "local_params": round(float(local_params)), + "partitioned_local_params": round(partitioned_local_params), + "local_muon_matrix_params": round(local_muon_matrix_params), + "local_fallback_params": round(local_fallback_params), + "local_non_expert_muon_params": round(local_non_expert_muon_params), + "local_expert_muon_params": round(local_expert_muon_params), + "local_non_expert_fallback_params": round(local_non_expert_fallback_params), + "pipeline_parallel_size": pp_size, + "tensor_parallel_size": tp_size, + "non_expert_fsdp_shard_size": non_expert_fsdp_shard_size, + "total_non_expert_shard_size": total_non_expert_shard_size, + "expert_shard_size": expert_shard_size, + "matrix_shapes_per_layer": matrix_shapes_per_layer, + "fallback_optimizer": str(train.get("muon_fallback_optimizer", "adamw") or "adamw"), + "fallback_param_sources": [ + "embed_tokens", + "lm_head", + "norms", + "router_gate_weight", + "gdn_A_log_dt_bias_o_norm", + ], + "notes": [ + "matches src/xorl/optim/optimizer.py Muon ndim>=2 classifier for qwen-style MoE", + "routed expert tensors and shared-expert projections use Muon", + "embeddings/lm_head/norms/router gate.weight use fallback optimizer", + "hybrid GDN layers: q/k/v[/g]/a/b/o projections and 3D conv1d weights are Muon; " + "A_log/dt_bias/o_norm are fallback; attention shapes apply to the full-attention " + "layers only (gated q_proj at 2x width when attn_output_gate)", + *fsdp_notes, + ], + } + + +def _muon_param_partition( + metadata: ModelMetadata, + topology: Topology, + train: dict[str, Any], + *, + local_params: float | None = None, +) -> dict[str, Any]: + if metadata.num_experts is not None or metadata.moe_intermediate_size is not None: + return _moe_muon_param_partition(metadata, topology, train, local_params=local_params) + return _dense_muon_param_partition(metadata, topology, train, local_params=local_params) + + +def _model_state_buckets( + raw_config: dict[str, Any], + topology: Topology | None, + metadata: ModelMetadata | None, + deepep_buffer_size_gb: float | None, +) -> tuple[float | None, float | None, float | None, float | None, float | None, list[MemoryBucket], list[str]]: + if topology is None or metadata is None: + return None, None, None, None, None, [], ["parameter_and_optimizer_bytes"] + + breakdown = _estimate_param_breakdown(metadata) + if breakdown is None: + return None, None, None, None, None, [], ["parameter_and_optimizer_bytes"] + + train = _section(raw_config, "train") + local_params, local_non_expert_params, local_expert_params, ownership_notes = _local_param_ownership( + breakdown, + train, + topology, + ) + + weight_bytes, weight_notes = _param_storage_bytes(train) + optimizer = str(train.get("optimizer", "")).lower() + optimizer_dtype_bytes = _dtype_bytes(train.get("optimizer_dtype"), default=4) + gradient_bytes, gradient_notes = _gradient_storage_bytes(train) + + sharded_param_gb = _gb(local_params * weight_bytes) + master_param_gb = 0.0 + if optimizer == "adamw" and weight_bytes < optimizer_dtype_bytes: + master_param_gb = _gb(local_params * optimizer_dtype_bytes) + persistent_model_state_gb = sharded_param_gb + master_param_gb + + gradient_state_gb = _gb(local_params * gradient_bytes) + optimizer_state_gb = 0.0 + optimizer_state_notes: list[str] = [] + if optimizer == "adamw": + if train.get("cautious_weight_decay"): + optimizer_state_bytes = 4 + optimizer_state_notes.append("cautious_weight_decay_routes_to_anyprecision_adamw_fp32_state") + else: + optimizer_state_bytes = weight_bytes + optimizer_state_notes.append("torch_adamw_state_dtype_follows_param_dtype") + optimizer_state_gb = _gb(local_params * 2 * optimizer_state_bytes) + optimizer_dtype_bytes = optimizer_state_bytes + elif optimizer == "anyprecision_adamw": + exp_avg_factor = 0 if train.get("anyprecision_adamw_reuse_grad_for_momentum") else 1 + compensation_factor = ( + 1 if train.get("anyprecision_adamw_use_kahan_summation") or train.get("use_kahan_summation") else 0 + ) + optimizer_state_factor = exp_avg_factor + 1 + compensation_factor + optimizer_state_gb = _gb(local_params * optimizer_state_factor * optimizer_dtype_bytes) + optimizer_state_notes.extend( + [ + f"optimizer_state_factor={optimizer_state_factor}", + f"reuse_grad_for_momentum={bool(train.get('anyprecision_adamw_reuse_grad_for_momentum'))}", + f"state_cpu_offload={bool(train.get('anyprecision_adamw_state_cpu_offload'))}:step_peak_still_loads_state", + ] + ) + elif optimizer == "sgd": + momentum = float(train.get("momentum", train.get("sgd_momentum", 0.0)) or 0.0) + if momentum > 0: + optimizer_state_gb = _gb(local_params * weight_bytes) + optimizer_dtype_bytes = weight_bytes + optimizer_state_notes.append(f"sgd_momentum_buffer={momentum}") + elif optimizer == "muon": + momentum = float(train.get("muon_momentum", 0.0) or 0.0) + force_momentum = bool(train.get("muon_force_momentum_path")) + fallback = str(train.get("muon_fallback_optimizer", "adamw") or "adamw").lower() + partition = _muon_param_partition(metadata, topology, train, local_params=local_params) + if partition.get("status") in {"exact_analytic_dense_muon_partition", "exact_analytic_moe_muon_partition"}: + muon_matrix_params = float(partition["local_muon_matrix_params"]) + fallback_params = float(partition["local_fallback_params"]) + muon_state_gb = _gb(muon_matrix_params * optimizer_dtype_bytes) if momentum > 0 or force_momentum else 0.0 + if fallback == "adamw": + fallback_state_gb = _gb(fallback_params * 2 * optimizer_dtype_bytes) + fallback_note = "fallback_adamw_exp_avg_and_exp_avg_sq" + elif fallback == "sgd": + fallback_state_gb = 0.0 + fallback_note = "fallback_sgd_state_free" + else: + fallback_state_gb = 0.0 + fallback_note = f"unsupported_fallback_optimizer={fallback}" + optimizer_state_gb = muon_state_gb + fallback_state_gb + optimizer_state_notes.extend( + [ + f"muon_partition={partition['status']}", + f"muon_momentum_buffer_params={partition['local_muon_matrix_params']}", + f"muon_fallback_params={partition['local_fallback_params']}", + fallback_note, + ] + ) + elif momentum > 0 or force_momentum: + optimizer_state_gb = _gb(local_params * optimizer_dtype_bytes) + optimizer_state_notes.extend(["muon_momentum_buffer", f"partition_status={partition.get('status')}"]) + + buckets = [ + MemoryBucket( + name="sharded_trainable_params", + gb=_round_gb(sharded_param_gb) or 0.0, + source="analytic_model_metadata", + notes=[ + f"weight_bytes={weight_bytes}", + *weight_notes, + *ownership_notes, + f"local_non_expert_params={local_non_expert_params:.0f}", + f"local_expert_params={local_expert_params:.0f}", + ], + ), + MemoryBucket( + name="gradient_storage", + gb=_round_gb(gradient_state_gb) or 0.0, + source="analytic_dtype_policy", + notes=[f"gradient_bytes={gradient_bytes}", *gradient_notes], + ), + ] + if master_param_gb: + buckets.append( + MemoryBucket( + name="optimizer_master_params", + gb=_round_gb(master_param_gb) or 0.0, + source="analytic_optimizer_policy", + notes=[f"optimizer={optimizer}", f"optimizer_dtype_bytes={optimizer_dtype_bytes}"], + ) + ) + if optimizer_state_gb: + buckets.append( + MemoryBucket( + name=f"{optimizer}_optimizer_state", + gb=_round_gb(optimizer_state_gb) or 0.0, + source="analytic_optimizer_policy", + notes=[f"optimizer_dtype_bytes={optimizer_dtype_bytes}", *optimizer_state_notes], + ) + ) + if deepep_buffer_size_gb: + buckets.append( + MemoryBucket( + name="deepep_static_buffer", + gb=deepep_buffer_size_gb, + source="config", + ) + ) + + has_routed_experts = metadata.num_experts is not None and metadata.moe_intermediate_size is not None + unsupported = [ + "activation_recompute_schedule", + "attention_workspace", + "moe_kernel_workspace" if has_routed_experts else "dense_mlp_workspace", + "fsdp_unshard_and_reduce_scatter_transients", + "allocator_reserved_slack", + ] + return ( + breakdown["total_params"] / 1_000_000_000, + local_params / 1_000_000_000, + persistent_model_state_gb, + gradient_state_gb, + optimizer_state_gb, + sorted(buckets, key=lambda bucket: bucket.gb, reverse=True), + unsupported, + ) + + +def _deepep_static_buffer_applies(metadata: ModelMetadata | None, topology: Topology | None) -> bool: + if metadata is None or topology is None: + return False + has_routed_experts = metadata.num_experts is not None and metadata.moe_intermediate_size is not None + return has_routed_experts + + +def _coverage_status( + *, + analytic_peak_floor_gb: float | None, + calibrated_peak_mem_gb: float | None, + unsupported_buckets: list[str], +) -> tuple[str, float | None, float | None, float | None]: + if analytic_peak_floor_gb is None: + return "unresolved_analytic_floor", None, None, None + if calibrated_peak_mem_gb is None: + return "analytic_floor_only", None, None, None + if calibrated_peak_mem_gb <= 0: + return "invalid_calibrated_peak", None, None, None + + floor_fraction = analytic_peak_floor_gb / calibrated_peak_mem_gb + residual = calibrated_peak_mem_gb - analytic_peak_floor_gb + if residual < 0: + return "calibrated_peak_below_analytic_floor", floor_fraction, 0.0, 0.0 + residual_fraction = residual / calibrated_peak_mem_gb + if residual == 0: + return "analytic_floor_matches_calibrated_peak", floor_fraction, 0.0, 0.0 + if unsupported_buckets: + return "calibrated_peak_with_unmodeled_residual", floor_fraction, residual, residual_fraction + return "calibrated_peak_residual_without_unsupported_bucket", floor_fraction, residual, residual_fraction + + +def build_memory_ledger( + raw_config: dict[str, Any], + observed: ObservedRun | None = None, + *, + topology: Topology | None = None, + model_metadata: ModelMetadata | None = None, + calibrated_peak_mem_gb: float | None = None, + calibrated_peak_source: str | None = None, + calibrated_phase_peak_gb: dict[str, float] | None = None, +) -> MemoryLedger: + model = _section(raw_config, "model") + train = _section(raw_config, "train") + observed_peak = None + observed_phase_peak: dict[str, float] = {} + + if observed is not None: + peaks = [row.peak_mem_gb for row in observed.steps if row.peak_mem_gb is not None] + observed_peak = max(peaks) if peaks else None + for row in observed.steps: + for phase, value in row.phase_memory_gb.items(): + observed_phase_peak[phase] = max(value, observed_phase_peak.get(phase, value)) + for memory_row in observed.memory_phases: + for key, value in memory_row.metrics.items(): + observed_phase_peak[key] = max(value, observed_phase_peak.get(key, value)) + if calibrated_phase_peak_gb: + for phase, value in calibrated_phase_peak_gb.items(): + observed_phase_peak[phase] = max(value, observed_phase_peak.get(phase, value)) + + deepep_buffer_size_gb = _float_field(model, "deepep_buffer_size_gb") + if deepep_buffer_size_gb is None: + deepep_buffer_size_gb = _float_field(train, "deepep_buffer_size_gb") + effective_deepep_buffer_size_gb = ( + deepep_buffer_size_gb if _deepep_static_buffer_applies(model_metadata, topology) else None + ) + + ( + estimated_total_params_b, + estimated_local_params_b, + persistent_model_state_gb, + gradient_state_gb, + optimizer_state_gb, + top_memory_buckets, + unsupported_buckets, + ) = _model_state_buckets(raw_config, topology, model_metadata, effective_deepep_buffer_size_gb) + analytic_peak_floor_gb = None + if persistent_model_state_gb is not None and gradient_state_gb is not None and optimizer_state_gb is not None: + analytic_peak_floor_gb = persistent_model_state_gb + gradient_state_gb + optimizer_state_gb + analytic_peak_floor_gb += effective_deepep_buffer_size_gb or 0.0 + + coverage_peak = observed_peak if observed_peak is not None else calibrated_peak_mem_gb + coverage_source = "observed_log" if observed_peak is not None else calibrated_peak_source + ( + memory_coverage_status, + floor_fraction, + residual_peak_gb, + residual_fraction, + ) = _coverage_status( + analytic_peak_floor_gb=analytic_peak_floor_gb, + calibrated_peak_mem_gb=coverage_peak, + unsupported_buckets=unsupported_buckets, + ) + if residual_peak_gb is not None and residual_peak_gb > 0: + top_memory_buckets = [ + *top_memory_buckets, + MemoryBucket( + name="calibrated_unmodeled_peak_residual", + gb=_round_gb(residual_peak_gb) or 0.0, + source=coverage_source or "calibrated_peak", + notes=[ + f"calibrated_peak_gb={coverage_peak:.3f}", + f"analytic_peak_floor_gb={analytic_peak_floor_gb:.3f}", + f"residual_fraction_of_peak={residual_fraction:.3f}", + "covers_unsupported_buckets=" + ",".join(unsupported_buckets), + ], + ), + ] + top_memory_buckets = sorted(top_memory_buckets, key=lambda bucket: bucket.gb, reverse=True) + + return MemoryLedger( + deepep_buffer_size_gb=effective_deepep_buffer_size_gb, + observed_peak_mem_gb_max=observed_peak, + calibrated_peak_mem_gb=_round_gb(coverage_peak), + calibrated_peak_source=coverage_source, + observed_phase_peak_gb=observed_phase_peak, + estimated_total_params_b=_round_gb(estimated_total_params_b), + estimated_local_params_b=_round_gb(estimated_local_params_b), + persistent_model_state_gb=_round_gb(persistent_model_state_gb), + gradient_state_gb=_round_gb(gradient_state_gb), + optimizer_state_gb=_round_gb(optimizer_state_gb), + analytic_peak_floor_gb=_round_gb(analytic_peak_floor_gb), + analytic_floor_fraction_of_calibrated_peak=_round_fraction(floor_fraction), + calibrated_residual_peak_gb=_round_gb(residual_peak_gb), + calibrated_residual_fraction_of_peak=_round_fraction(residual_fraction), + memory_coverage_status=memory_coverage_status, + top_memory_buckets=top_memory_buckets, + unsupported_buckets=unsupported_buckets, + ) diff --git a/experiments/local_benchmark/training_sim/model_metadata.py b/src/xorl/sim/model_metadata.py similarity index 53% rename from experiments/local_benchmark/training_sim/model_metadata.py rename to src/xorl/sim/model_metadata.py index e2267d2e..4dedd41d 100644 --- a/experiments/local_benchmark/training_sim/model_metadata.py +++ b/src/xorl/sim/model_metadata.py @@ -28,6 +28,59 @@ "vocab_size": 151936, "tie_word_embeddings": False, }, + "Qwen/Qwen3-235B-A22B-Instruct-2507": { + "num_experts": 128, + "top_k": 8, + "num_hidden_layers": 94, + "hidden_size": 4096, + "intermediate_size": 12288, + "moe_intermediate_size": 1536, + "num_attention_heads": 64, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 151936, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3.5-397B-A17B": { + "num_experts": 512, + "top_k": 10, + "num_hidden_layers": 60, + "hidden_size": 4096, + "moe_intermediate_size": 1024, + "shared_expert_intermediate_size": 1024, + "num_attention_heads": 32, + "num_key_value_heads": 2, + "head_dim": 256, + "full_attention_interval": 4, + "attn_output_gate": True, + "linear_num_key_heads": 16, + "linear_num_value_heads": 64, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "vocab_size": 248320, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3.5-35B-A3B": { + "num_experts": 256, + "top_k": 8, + "num_hidden_layers": 40, + "hidden_size": 2048, + "moe_intermediate_size": 512, + "shared_expert_intermediate_size": 512, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "head_dim": 256, + "full_attention_interval": 4, + "attn_output_gate": True, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, + "vocab_size": 248320, + "tie_word_embeddings": False, + }, "Qwen/Qwen3.6-35B-A3B": { "num_experts": 256, "top_k": 8, @@ -38,6 +91,13 @@ "num_attention_heads": 16, "num_key_value_heads": 2, "head_dim": 256, + "full_attention_interval": 4, + "attn_output_gate": True, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, "vocab_size": 248320, "tie_word_embeddings": False, }, @@ -51,9 +111,83 @@ "num_attention_heads": 16, "num_key_value_heads": 2, "head_dim": 256, + "full_attention_interval": 4, + "attn_output_gate": True, + "linear_num_key_heads": 16, + "linear_num_value_heads": 32, + "linear_key_head_dim": 128, + "linear_value_head_dim": 128, + "linear_conv_kernel_dim": 4, "vocab_size": 248320, "tie_word_embeddings": False, }, + "Qwen/Qwen3-30B-A3B": { + # Non-Coder Qwen3-30B-A3B (model_type=qwen3_moe). Verified against the cached HF + # config.json: decoder_sparse_step=1 and mlp_only_layers=[] mean every layer is MoE, + # so the dense intermediate_size=6144 is unused for FFN params (the memory ledger + # ignores it once moe_intermediate_size is set). Shape params are identical to the + # Coder-30B variant; the two differ only in max_position_embeddings (40960 vs 262144) + # and rope_theta, neither of which changes training fwd/bwd shapes or param counts. + # Kept as a distinct identity so its evidence stream is not confused with Coder-30B. + "num_experts": 128, + "top_k": 8, + "num_hidden_layers": 48, + "hidden_size": 2048, + "intermediate_size": 6144, + "moe_intermediate_size": 768, + "num_attention_heads": 32, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 151936, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3-Coder-30B-A3B": { + "num_experts": 128, + "top_k": 8, + "num_hidden_layers": 48, + "hidden_size": 2048, + "intermediate_size": 6144, + "moe_intermediate_size": 768, + "num_attention_heads": 32, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 151936, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3-Coder-30B-A3B-Instruct": { + "num_experts": 128, + "top_k": 8, + "num_hidden_layers": 48, + "hidden_size": 2048, + "intermediate_size": 6144, + "moe_intermediate_size": 768, + "num_attention_heads": 32, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 151936, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3-32B": { + "num_hidden_layers": 64, + "hidden_size": 5120, + "intermediate_size": 25600, + "num_attention_heads": 64, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 151936, + "tie_word_embeddings": False, + }, + "Qwen/Qwen3-8B": { + # Dense Qwen3-8B (model_type=qwen3), verified against its Hugging Face config. + "num_hidden_layers": 36, + "hidden_size": 4096, + "intermediate_size": 12288, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 151936, + "tie_word_embeddings": False, + }, } @@ -86,12 +220,7 @@ def default_hf_cache_roots() -> list[Path]: hf_home = os.environ.get("HF_HOME") if hf_home: roots.append(Path(hf_home) / "hub") - roots.extend( - [ - Path("/shared/huggingface/hub"), - Path.home() / ".cache" / "huggingface" / "hub", - ] - ) + roots.append(Path.home() / ".cache" / "huggingface" / "hub") deduped: list[Path] = [] seen: set[Path] = set() for root in roots: @@ -157,11 +286,28 @@ def _read_metadata_file(config_path: Path, model_ref: str | None) -> ModelMetada head_dim=_find_int(sections, ("head_dim",)), vocab_size=_find_int(sections, ("vocab_size",)), tie_word_embeddings=_find_bool(sections, ("tie_word_embeddings",)), + full_attention_interval=_find_int(sections, ("full_attention_interval",)), + attn_output_gate=_find_bool(sections, ("attn_output_gate",)), + linear_num_key_heads=_find_int(sections, ("linear_num_key_heads",)), + linear_num_value_heads=_find_int(sections, ("linear_num_value_heads",)), + linear_key_head_dim=_find_int(sections, ("linear_key_head_dim",)), + linear_value_head_dim=_find_int(sections, ("linear_value_head_dim",)), + linear_conv_kernel_dim=_find_int(sections, ("linear_conv_kernel_dim",)), ) def _known_metadata(model_ref: str) -> ModelMetadata | None: values = KNOWN_MODEL_METADATA.get(model_ref) + if values is None: + lowered = model_ref.lower() + for known_ref in sorted(KNOWN_MODEL_METADATA, key=len, reverse=True): + known_lowered = known_ref.lower() + known_name = known_ref.rsplit("/", 1)[-1].lower() + cache_name = known_ref.replace("/", "--").lower() + if known_lowered in lowered or known_name in lowered or cache_name in lowered: + values = KNOWN_MODEL_METADATA[known_ref] + model_ref = known_ref + break if values is None: return None return ModelMetadata( @@ -180,6 +326,13 @@ def _known_metadata(model_ref: str) -> ModelMetadata | None: head_dim=values.get("head_dim"), vocab_size=values.get("vocab_size"), tie_word_embeddings=values.get("tie_word_embeddings"), + full_attention_interval=values.get("full_attention_interval"), + attn_output_gate=values.get("attn_output_gate"), + linear_num_key_heads=values.get("linear_num_key_heads"), + linear_num_value_heads=values.get("linear_num_value_heads"), + linear_key_head_dim=values.get("linear_key_head_dim"), + linear_value_head_dim=values.get("linear_value_head_dim"), + linear_conv_kernel_dim=values.get("linear_conv_kernel_dim"), ) @@ -212,6 +365,13 @@ def resolve_model_metadata( head_dim=_find_int(config_sections, ("head_dim",)), vocab_size=_find_int(config_sections, ("vocab_size",)), tie_word_embeddings=_find_bool(config_sections, ("tie_word_embeddings",)), + full_attention_interval=_find_int(config_sections, ("full_attention_interval",)), + attn_output_gate=_find_bool(config_sections, ("attn_output_gate",)), + linear_num_key_heads=_find_int(config_sections, ("linear_num_key_heads",)), + linear_num_value_heads=_find_int(config_sections, ("linear_num_value_heads",)), + linear_key_head_dim=_find_int(config_sections, ("linear_key_head_dim",)), + linear_value_head_dim=_find_int(config_sections, ("linear_value_head_dim",)), + linear_conv_kernel_dim=_find_int(config_sections, ("linear_conv_kernel_dim",)), ) file_metadata = None @@ -252,4 +412,39 @@ def resolve_model_metadata( head_dim=source_metadata.head_dim, vocab_size=source_metadata.vocab_size, tie_word_embeddings=source_metadata.tie_word_embeddings, + full_attention_interval=( + config_metadata.full_attention_interval + if config_metadata.full_attention_interval is not None + else source_metadata.full_attention_interval + ), + attn_output_gate=( + config_metadata.attn_output_gate + if config_metadata.attn_output_gate is not None + else source_metadata.attn_output_gate + ), + linear_num_key_heads=( + config_metadata.linear_num_key_heads + if config_metadata.linear_num_key_heads is not None + else source_metadata.linear_num_key_heads + ), + linear_num_value_heads=( + config_metadata.linear_num_value_heads + if config_metadata.linear_num_value_heads is not None + else source_metadata.linear_num_value_heads + ), + linear_key_head_dim=( + config_metadata.linear_key_head_dim + if config_metadata.linear_key_head_dim is not None + else source_metadata.linear_key_head_dim + ), + linear_value_head_dim=( + config_metadata.linear_value_head_dim + if config_metadata.linear_value_head_dim is not None + else source_metadata.linear_value_head_dim + ), + linear_conv_kernel_dim=( + config_metadata.linear_conv_kernel_dim + if config_metadata.linear_conv_kernel_dim is not None + else source_metadata.linear_conv_kernel_dim + ), ) diff --git a/experiments/local_benchmark/training_sim/predict.py b/src/xorl/sim/predict.py similarity index 73% rename from experiments/local_benchmark/training_sim/predict.py rename to src/xorl/sim/predict.py index d967856f..261160dd 100644 --- a/experiments/local_benchmark/training_sim/predict.py +++ b/src/xorl/sim/predict.py @@ -10,18 +10,24 @@ try: from .benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from .calibration_packs import resolve_pack_inputs from .collect_calibration import merge_observed_runs, parse_log_path, summarize_observed_run from .config_fingerprint import build_fingerprint, load_training_config from .memory_ledger import build_memory_ledger from .schemas import PredictionReport, to_jsonable from .shape_engine import build_shape_ledger + from .simulator_support import resolve_simulator_support + from .timing_ledger import build_timing_ledger except ImportError: # pragma: no cover - exercised by direct script execution from benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from calibration_packs import resolve_pack_inputs from collect_calibration import merge_observed_runs, parse_log_path, summarize_observed_run from config_fingerprint import build_fingerprint, load_training_config from memory_ledger import build_memory_ledger from schemas import PredictionReport, to_jsonable from shape_engine import build_shape_ledger + from simulator_support import resolve_simulator_support + from timing_ledger import build_timing_ledger def build_report( @@ -45,6 +51,9 @@ def build_report( top_k=top_k, ) raw_config = load_training_config(config_path) + simulator = raw_config.setdefault("simulator", {}) + if isinstance(simulator, dict): + simulator["balanced_routing"] = balanced_routing shape = build_shape_ledger(fingerprint.topology, balanced_routing=balanced_routing) observed = None observed_summary: dict[str, Any] | None = None @@ -58,26 +67,49 @@ def build_report( world_size=fingerprint.topology.world_size, ) + benchmark_behavior = None + if benchmark_dir is not None: + behavior_points = load_benchmark_behavior_points(benchmark_dir) + benchmark_behavior = predict_benchmark_behavior(behavior_points, fingerprint.topology, shape, raw_config) memory = build_memory_ledger( raw_config, observed, topology=fingerprint.topology, model_metadata=fingerprint.model_metadata, + calibrated_peak_mem_gb=benchmark_behavior.peak_mem_gb if benchmark_behavior is not None else None, + calibrated_peak_source=( + f"benchmark_behavior:{benchmark_behavior.matched_label}" + if benchmark_behavior is not None and benchmark_behavior.matched_label is not None + else None + ), + calibrated_phase_peak_gb=(benchmark_behavior.phase_memory_peak_gb if benchmark_behavior is not None else None), ) - benchmark_behavior = None - if benchmark_dir is not None: - behavior_points = load_benchmark_behavior_points(benchmark_dir) - benchmark_behavior = predict_benchmark_behavior(behavior_points, fingerprint.topology, shape, raw_config) warnings = list(shape.warnings) if memory.observed_peak_mem_gb_max is None: warnings.append("no observed memory calibration was supplied") if benchmark_behavior is not None: warnings.extend(benchmark_behavior.warnings) + timing = build_timing_ledger( + observed_summary, + benchmark_behavior, + calibration_sources=calibration_sources, + ) + support = resolve_simulator_support( + raw_config, + topology=fingerprint.topology, + memory=memory, + timing=timing, + ) + if support.support_blockers: + warnings.append(f"simulator support status: {support.support_status}") + warnings.extend(f"simulator support blocker: {blocker}" for blocker in support.support_blockers) return PredictionReport( fingerprint=fingerprint, shape=shape, memory=memory, + timing=timing, + support=support, benchmark_behavior=benchmark_behavior, observed_summary=observed_summary, calibration_sources=calibration_sources, @@ -87,7 +119,8 @@ def build_report( def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--config", type=Path, required=True, help="XoRL YAML config") + parser.add_argument("--pack", help="Built-in calibration-pack name") + parser.add_argument("--config", type=Path, default=None, help="XoRL YAML config") parser.add_argument("--world-size", type=int, default=None, help="Override WORLD_SIZE for config resolution") parser.add_argument("--local-world-size", type=int, default=None, help="Override LOCAL_WORLD_SIZE") parser.add_argument("--balanced-routing", action="store_true", help="Assume deterministic balanced MoE routing") @@ -101,6 +134,10 @@ def main() -> None: parser.add_argument("--output", type=Path, default=None, help="Write JSON report to this path") args = parser.parse_args() + args.config, args.benchmark_dir = resolve_pack_inputs(args.pack, args.config, args.benchmark_dir) + if args.config is None: + parser.error("provide --pack or --config") + report = build_report( args.config, world_size=args.world_size, diff --git a/src/xorl/sim/runtime_config.py b/src/xorl/sim/runtime_config.py new file mode 100644 index 00000000..40c36189 --- /dev/null +++ b/src/xorl/sim/runtime_config.py @@ -0,0 +1,18 @@ +"""Helpers for writing simulator configs that are runnable by XORL.""" + +from __future__ import annotations + +import copy +from typing import Any + + +SIMULATOR_ONLY_TOP_LEVEL_SECTIONS = ("simulator", "_simulator") + + +def runtime_training_config(config: dict[str, Any]) -> dict[str, Any]: + """Return a deep-copied config with simulator-only metadata removed.""" + + rendered = copy.deepcopy(config) + for section_name in SIMULATOR_ONLY_TOP_LEVEL_SECTIONS: + rendered.pop(section_name, None) + return rendered diff --git a/src/xorl/sim/scenario_planner.py b/src/xorl/sim/scenario_planner.py new file mode 100644 index 00000000..62f82911 --- /dev/null +++ b/src/xorl/sim/scenario_planner.py @@ -0,0 +1,9085 @@ +"""Plan and score topology scenarios from a base XoRL training config.""" + +from __future__ import annotations + +import argparse +import copy +import json +import math +import re +from collections import Counter +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Any + +import yaml + + +try: + from .analytical_ledgers import activation_ledger + from .benchmark_behavior import ( + H100_BF16_PROMISED_TFLOPS_PER_GPU, + _config_balanced_routing, + behavior_point_matches_topology, + behavior_point_matches_workload, + behavior_point_model_mismatches, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from .calibration_packs import resolve_pack_inputs + from .config_fingerprint import load_training_config, resolve_topology + from .memory_ledger import build_memory_ledger + from .model_metadata import model_ref_from_config, resolve_model_metadata + from .runtime_config import runtime_training_config + from .schemas import ( + BenchmarkBehaviorPoint, + BenchmarkBehaviorPrediction, + CommLedger, + MemoryLedger, + ModelMetadata, + ParallelismAxisComparison, + ParallelismBoundaryAxisCoverage, + ParallelismBoundaryGroup, + ScenarioBenchmarkSupport, + ScenarioCandidate, + ScenarioCaptureGap, + ScenarioDecisionSummary, + ScenarioMeasurementConfig, + ScenarioParallelismAxisCoverage, + ScenarioReadiness, + ScenarioReport, + ScenarioValidationAction, + Topology, + to_jsonable, + ) + from .shape_engine import ShapeLedger, build_shape_ledger + from .simulator_support import requested_simulator_surface, resolve_simulator_support + from .timing_ledger import build_timing_ledger +except ImportError: # pragma: no cover - exercised by direct script execution + from analytical_ledgers import activation_ledger + from benchmark_behavior import ( + H100_BF16_PROMISED_TFLOPS_PER_GPU, + _config_balanced_routing, + behavior_point_matches_topology, + behavior_point_matches_workload, + behavior_point_model_mismatches, + behavior_point_workload_mismatches, + load_benchmark_behavior_points, + predict_benchmark_behavior, + ) + from calibration_packs import resolve_pack_inputs + from config_fingerprint import load_training_config, resolve_topology + from memory_ledger import build_memory_ledger + from model_metadata import model_ref_from_config, resolve_model_metadata + from runtime_config import runtime_training_config + from schemas import ( + BenchmarkBehaviorPoint, + BenchmarkBehaviorPrediction, + CommLedger, + MemoryLedger, + ModelMetadata, + ParallelismAxisComparison, + ParallelismBoundaryAxisCoverage, + ParallelismBoundaryGroup, + ScenarioBenchmarkSupport, + ScenarioCandidate, + ScenarioCaptureGap, + ScenarioDecisionSummary, + ScenarioMeasurementConfig, + ScenarioParallelismAxisCoverage, + ScenarioReadiness, + ScenarioReport, + ScenarioValidationAction, + Topology, + to_jsonable, + ) + from shape_engine import ShapeLedger, build_shape_ledger + from simulator_support import requested_simulator_surface, resolve_simulator_support + from timing_ledger import build_timing_ledger + + +_EXACT_CROSS_MODEL_SEQUENCE_RATIO_WINDOW = (0.90, 1.10) +_EXTRAPOLATED_CROSS_MODEL_SEQUENCE_RATIO_WINDOW = (0.45, 2.20) +_MIN_ULYSSES_SEQUENCE_LEN = 16_384 +_MIN_RINGATTN_SEQUENCE_LEN = 64_000 +_PHASE_BOTTLENECK_HALFSPEED_SCALE = 0.5 + + +def _section(raw: dict[str, Any], name: str) -> dict[str, Any]: + value = raw.get(name, {}) + if isinstance(value, dict): + return value + raw[name] = {} + return raw[name] + + +def _parse_int_list(raw: str | None) -> list[int] | None: + if raw is None or raw == "auto": + return None + values = sorted({int(part.strip()) for part in raw.split(",") if part.strip()}) + if not values: + raise ValueError("expected at least one integer") + return values + + +_SCENARIO_PARALLELISM_DIMENSIONS = ( + "world_size", + "data_parallel_replicate_size", + "data_parallel_shard_size", + "tensor_parallel_size", + "pipeline_parallel_size", + "expert_parallel_size", + "ep_fsdp_size", + "ulysses_parallel_size", + "ringattn_parallel_size", +) + +_PARALLELISM_COMPARISON_DIMENSIONS = ( + "world_size", + "node_count", + "data_parallel_size", + "data_parallel_replicate_size", + "data_parallel_shard_size", + "tensor_parallel_size", + "pipeline_parallel_size", + "expert_parallel_size", + "ep_fsdp_size", + "ulysses_parallel_size", + "ringattn_parallel_size", +) + +_PARALLELISM_AXIS_DIMENSIONS = { + "world_size": ( + "world_size", + "node_count", + "data_parallel_size", + "data_parallel_shard_size", + "ep_fsdp_size", + ), + "dp_replicate": ( + "data_parallel_replicate_size", + "data_parallel_size", + "data_parallel_shard_size", + "ep_fsdp_size", + ), + "dp_shard": ( + "data_parallel_shard_size", + "data_parallel_size", + "ep_fsdp_size", + ), + "tensor_parallel": ( + "tensor_parallel_size", + "data_parallel_size", + "data_parallel_shard_size", + "ep_fsdp_size", + ), + "pipeline_parallel": ( + "pipeline_parallel_size", + "data_parallel_size", + "data_parallel_shard_size", + "ep_fsdp_size", + ), + "expert_parallel": ( + "expert_parallel_size", + "ep_fsdp_size", + ), + # expert_parallel co-varies with ep_fsdp at fixed dp (ep x ep_fsdp = dp_shard), mirroring the + # converse declaration above; without it the one physical ep<->ep_fsdp pair is double-counted as + # a confounded ep_fsdp-axis row even when it is a clean expert_parallel pair. + "ep_fsdp": ( + "ep_fsdp_size", + "expert_parallel_size", + ), + "ulysses": ( + "ulysses_parallel_size", + "data_parallel_size", + "data_parallel_shard_size", + "ep_fsdp_size", + ), + "ringattn": ( + "ringattn_parallel_size", + "data_parallel_size", + "data_parallel_shard_size", + "ep_fsdp_size", + ), +} + +_PARALLELISM_AXIS_PRIMARY_DIMENSIONS = { + "world_size": ("world_size",), + "dp_replicate": ("data_parallel_replicate_size",), + "dp_shard": ("data_parallel_shard_size",), + "tensor_parallel": ("tensor_parallel_size",), + "pipeline_parallel": ("pipeline_parallel_size",), + "expert_parallel": ("expert_parallel_size",), + "ep_fsdp": ("ep_fsdp_size",), + "ulysses": ("ulysses_parallel_size",), + "ringattn": ("ringattn_parallel_size",), +} + +_SCENARIO_WORKLOAD_DIMENSIONS = ( + "micro_batch_size", + "gradient_accumulation_steps", + "global_batch_size", + "sample_packing_sequence_len", +) + +_SCENARIO_RUNTIME_SIGNATURE_FIELDS = ( + ("train", "gradient_checkpointing_method"), + ("train", "enable_activation_offload"), + ("train", "activation_offload_prefetch_count"), + ("train", "skip_param_upcast"), + ("train", "fsdp_reduce_dtype"), + ("train", "ce_mode"), + ("model", "moe_implementation"), + ("train", "moe_checkpoint_method"), + ("train", "muon_momentum"), + ("train", "muon_update_dtype"), + ("model", "deepep_async_combine"), + ("model", "deepep_num_sms"), + ("model", "deepep_buffer_size_gb"), + ("train", "enable_compile"), + ("_simulator", "attention_backend"), +) + +_SCENARIO_RUNTIME_DIMENSIONS = tuple(field_name for _, field_name in _SCENARIO_RUNTIME_SIGNATURE_FIELDS) + +_RUNTIME_VARIANT_CONFIG_PATHS = { + field_name: (section_name, field_name) for section_name, field_name in _SCENARIO_RUNTIME_SIGNATURE_FIELDS +} + +_RUNTIME_VARIANT_BOOL_DEFAULTS = { + "enable_activation_offload": False, + "skip_param_upcast": False, + "deepep_async_combine": False, + "enable_compile": False, +} + +_SCENARIO_BOUNDARY_WORKLOAD_DIMENSIONS = ( + *_SCENARIO_WORKLOAD_DIMENSIONS, + "balanced_routing", +) + +_SCENARIO_BOUNDARY_SIGNATURE_DIMENSIONS = ( + "micro_batch_size", + "gradient_accumulation_steps", + "sample_packing_sequence_len", + "balanced_routing", +) + +_CALIBRATION_DISTANCE_DIMENSIONS = ( + "world_size", + "micro_batch_size", + "global_batch_size", + "sample_packing_sequence_len", + "expert_parallel_size", + "ep_fsdp_size", + "tensor_parallel_size", + "pipeline_parallel_size", + "ulysses_parallel_size", + "ringattn_parallel_size", +) + + +def _divisors(value: int) -> list[int]: + return [candidate for candidate in range(1, value + 1) if value % candidate == 0] + + +def _power_of_two_divisors(value: int, *, max_value: int | None = None) -> list[int]: + limit = value if max_value is None else min(value, max_value) + return [candidate for candidate in _divisors(value) if candidate <= limit and candidate & (candidate - 1) == 0] + + +def _dedupe_sorted(values: list[int] | set[int]) -> list[int]: + return sorted(value for value in set(values) if value > 0) + + +def _default_micro_batch_sizes( + base_topology: Topology, + behavior_points: list[BenchmarkBehaviorPoint], +) -> list[int]: + values = {base_topology.micro_batch_size} + values.update(point.micro_batch_size for point in behavior_points if point.micro_batch_size is not None) + return sorted(values) + + +def _default_ep_sizes(base_topology: Topology) -> list[int]: + if base_topology.num_experts is None: + return [base_topology.expert_parallel_size] + ranks_per_pipeline_stage = base_topology.world_size // base_topology.pipeline_parallel_size + values = { + value for value in _divisors(ranks_per_pipeline_stage) if value > 0 and base_topology.num_experts % value == 0 + } + if base_topology.expert_parallel_size in values: + return [base_topology.expert_parallel_size] + return sorted(values) or [base_topology.expert_parallel_size] + + +def _auto_ep_sizes(base_topology: Topology) -> list[int]: + if base_topology.num_experts is None: + return [base_topology.expert_parallel_size] + values = { + value + for value in _divisors(base_topology.world_size) + if base_topology.num_experts % value == 0 and value <= base_topology.world_size + } + values.add(base_topology.expert_parallel_size) + return _dedupe_sorted(values) + + +def _auto_tensor_parallel_sizes(base_topology: Topology, metadata: ModelMetadata) -> list[int]: + values = set(_power_of_two_divisors(base_topology.world_size, max_value=base_topology.local_world_size)) + values.add(base_topology.tensor_parallel_size) + if metadata.hidden_size is not None: + values = {value for value in values if metadata.hidden_size % value == 0} + if metadata.num_attention_heads is not None: + values = {value for value in values if metadata.num_attention_heads % value == 0} + return _dedupe_sorted(values) or [base_topology.tensor_parallel_size] + + +def _auto_pipeline_parallel_sizes(base_topology: Topology, metadata: ModelMetadata) -> list[int]: + values = set(_power_of_two_divisors(base_topology.world_size, max_value=4)) + values.add(base_topology.pipeline_parallel_size) + if metadata.num_hidden_layers is not None: + values = {value for value in values if metadata.num_hidden_layers % value == 0} + values.add(base_topology.pipeline_parallel_size) + return _dedupe_sorted(values) or [base_topology.pipeline_parallel_size] + + +def _auto_ulysses_parallel_sizes(base_topology: Topology, metadata: ModelMetadata) -> list[int]: + values = {base_topology.ulysses_parallel_size, 1} + seq_len = base_topology.sample_packing_sequence_len or 0 + if seq_len >= _MIN_ULYSSES_SEQUENCE_LEN: + max_ulysses = 64 + if metadata.num_key_value_heads is not None: + max_ulysses = max(1, metadata.num_key_value_heads, base_topology.ulysses_parallel_size) + values.update(_power_of_two_divisors(base_topology.world_size, max_value=max_ulysses)) + return _dedupe_sorted(values) + + +def _auto_ringattn_parallel_sizes(base_topology: Topology) -> list[int]: + values = {base_topology.ringattn_parallel_size, 1} + seq_len = base_topology.sample_packing_sequence_len or 0 + if seq_len >= _MIN_RINGATTN_SEQUENCE_LEN: + values.update(_power_of_two_divisors(base_topology.world_size, max_value=4)) + return _dedupe_sorted(values) + + +def _known_or_default_parallel_size(point_value: int | None) -> int: + return point_value if point_value is not None else 1 + + +def _point_matches_topology_parallel_dims(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + return ( + _known_or_default_parallel_size(point.tensor_parallel_size) == topology.tensor_parallel_size + and _known_or_default_parallel_size(point.pipeline_parallel_size) == topology.pipeline_parallel_size + and _known_or_default_parallel_size(point.ulysses_parallel_size) == topology.ulysses_parallel_size + and _known_or_default_parallel_size(point.ringattn_parallel_size) == topology.ringattn_parallel_size + ) + + +def _point_matches_parallel_dims_for_risk(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + if point.expert_parallel_size is not None and point.expert_parallel_size != topology.expert_parallel_size: + return False + if point.ep_fsdp_size is not None and point.ep_fsdp_size != topology.ep_fsdp_size: + return False + return _point_matches_topology_parallel_dims(point, topology) + + +def _phase_bucket(phase: str) -> str: + lowered = phase.lower() + if any(part in lowered for part in ("dataloader", "get_batch", "input", "tokenize", "collator")): + return "input" + if any(part in lowered for part in ("optimizer", "optim", "clip_and_step", "lr_scheduler", "zero_grad")): + return "optimizer" + if any( + part in lowered + for part in ( + "sync", + "all_reduce", + "reduce_scatter", + "all_gather", + "nccl", + "deepep", + "dispatch", + "combine", + "communication", + "data_movement", + "a2a", + "fsdp", + ) + ): + return "communication" + if lowered in {"fwd", "bwd", "fwd+bwd"} or any( + part in lowered for part in ("forward", "backward", "loss", "recompute", "moe", "attention") + ): + return "model_compute" + return "other" + + +def _is_composite_phase_for_bottleneck(phase: str, phases: set[str]) -> bool: + lowered = phase.lower() + lowered_phases = {item.lower() for item in phases} + if lowered == "train_step_total": + return True + if lowered in {"forward_backward", "forward_backward_total", "fwd_bwd", "fwd_bwd_total"}: + return bool( + lowered_phases + & { + "model_forward", + "forward", + "fwd", + "loss_compute", + "loss", + "backward", + "model_backward", + "bwd", + } + ) + if lowered == "clip_and_step_total": + return bool( + lowered_phases + & { + "clip_gradients", + "optimizer_step", + "optimizer", + "optim", + "lr_scheduler_step", + } + ) + return False + + +def _phase_bottleneck(phase_time_share: dict[str, float]) -> tuple[str, float] | None: + visible = { + phase: share + for phase, share in phase_time_share.items() + if not _is_composite_phase_for_bottleneck(phase, set(phase_time_share)) + } + if not visible: + visible = {phase: share for phase, share in phase_time_share.items() if phase != "train_step_total"} + if not visible: + return None + phase = max(visible, key=visible.get) + return phase, visible[phase] + + +def _phase_bottleneck_note(phase_time_share: dict[str, float]) -> str | None: + bottleneck = _phase_bottleneck(phase_time_share) + if bottleneck is None: + return None + phase, share = bottleneck + return f"phase_bottleneck={phase}:{share:.1%}" + + +def _phase_bottleneck_details( + phase_time_share: dict[str, float], + phase_time_sec: dict[str, float], +) -> tuple[str, str, float, float | None] | None: + bottleneck = _phase_bottleneck(phase_time_share) + if bottleneck is None: + return None + phase, share = bottleneck + time_sec = phase_time_sec.get(phase) + return phase, _phase_bucket(phase), round(share, 6), round(time_sec, 6) if time_sec is not None else None + + +def _phase_bottleneck_half_speedup_counterfactual( + score_tokens_per_sec: float | None, + phase_bottleneck_share: float | None, +) -> tuple[float | None, float | None]: + if score_tokens_per_sec is None or score_tokens_per_sec <= 0: + return None, None + if phase_bottleneck_share is None or not 0 < phase_bottleneck_share <= 1: + return None, None + counterfactual_time_fraction = ( + 1.0 - phase_bottleneck_share + (phase_bottleneck_share * _PHASE_BOTTLENECK_HALFSPEED_SCALE) + ) + if counterfactual_time_fraction <= 0: + return None, None + speedup = 1.0 / counterfactual_time_fraction + return round(score_tokens_per_sec * speedup, 3), round((speedup - 1.0) * 100.0, 3) + + +def _memory_bottleneck_details( + phase_memory_peak_gb: dict[str, float], + peak_mem_gb: float | None, +) -> tuple[str, str, float, float] | None: + visible = { + phase: peak + for phase, peak in phase_memory_peak_gb.items() + if peak > 0 and not _is_composite_phase_for_bottleneck(phase, set(phase_memory_peak_gb)) + } + if not visible: + visible = {phase: peak for phase, peak in phase_memory_peak_gb.items() if peak > 0} + if not visible: + return None + phase, peak = max(visible.items(), key=lambda item: (item[1], item[0])) + denominator = peak_mem_gb if peak_mem_gb is not None and peak_mem_gb > 0 else peak + return phase, _phase_bucket(phase), round(peak, 3), round(peak / denominator, 3) + + +def _scenario_timing_coverage_status( + behavior: BenchmarkBehaviorPrediction, + *, + prediction_confidence: str, + calibration_scope: str, + timing_coverage_status: str, +) -> str: + if timing_coverage_status == "no_timing_calibration": + return "no_timing_evidence" + has_phase_timing = timing_coverage_status.endswith("_phase_timing") + exact_calibrated = prediction_confidence == "calibrated" and calibration_scope == "exact_calibrated" + if exact_calibrated: + return "exact_phase_timing" if has_phase_timing else "exact_total_step_only" + if behavior.status == "extrapolated_step_time_fit": + return "step_time_fit_phase_timing" if has_phase_timing else "step_time_fit_total_step_only" + if behavior.status == "cross_model_extrapolated" or calibration_scope == "cross_model_analog": + return "cross_model_reference_phase_timing" if has_phase_timing else "cross_model_reference_total_step_only" + return "reference_phase_timing_extrapolated" if has_phase_timing else "reference_total_step_extrapolated" + + +def _matched_labels(behavior: BenchmarkBehaviorPrediction) -> set[str]: + return {part.strip() for part in (behavior.matched_label or "").split(",") if part.strip()} + + +def _point_dimension_value(point: BenchmarkBehaviorPoint, dimension: str) -> int | None: + if dimension == "world_size": + return point.gpu_count + value = getattr(point, dimension) + if value is None: + return None + if dimension.endswith("_parallel_size"): + return _known_or_default_parallel_size(value) + return value + + +def _topology_dimension_value(topology: Topology, dimension: str) -> int | None: + if dimension == "world_size": + return topology.world_size + value = getattr(topology, dimension) + if value is None: + return None + if dimension.endswith("_parallel_size"): + return _known_or_default_parallel_size(value) + return value + + +def _calibration_distance( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + behavior: BenchmarkBehaviorPrediction, +) -> tuple[float | None, list[str]]: + matched_labels = _matched_labels(behavior) + if not matched_labels: + return None, [] + matched_points = [point for point in behavior_points if point.label in matched_labels] + if not matched_points: + return None, [] + + options: list[tuple[float, str, list[tuple[float, str]]]] = [] + for point in matched_points: + factors: list[tuple[float, str]] = [] + total = 0.0 + for dimension in _CALIBRATION_DISTANCE_DIMENSIONS: + reference = _point_dimension_value(point, dimension) + target = _topology_dimension_value(topology, dimension) + if reference is None or target is None or reference <= 0 or target <= 0: + continue + distance = abs(math.log2(target / reference)) + if distance == 0.0: + continue + total += distance + factors.append((distance, f"{dimension}:{reference}->{target}:log2={distance:.3f}")) + options.append((total, point.label, factors)) + + if not options: + return None, [] + total, label, factors = min(options, key=lambda option: (option[0], option[1])) + factor_notes = [f"reference={label}"] + factor_notes.extend(note for _, note in sorted(factors, key=lambda factor: (-factor[0], factor[1]))) + return round(total, 3), factor_notes + + +def _calibration_scope( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + *, + prediction_confidence: str, + raw_config: dict[str, Any] | None, +) -> str: + if prediction_confidence == "calibrated": + return "exact_calibrated" + if prediction_confidence == "cross_model_extrapolated": + return "cross_model_analog" + + throughput_points = [ + point + for point in behavior_points + if point.tokens_per_sec is not None + and (raw_config is None or not behavior_point_model_mismatches(point, raw_config)) + ] + if not throughput_points: + return "no_calibration" + + same_sequence_points = [ + point + for point in throughput_points + if point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) + ] + if not same_sequence_points: + return "outside_sequence_calibration_envelope" + + dimensions: tuple[tuple[str, int], ...] = ( + ("micro_batch_size", topology.micro_batch_size), + ("global_batch_size", topology.global_batch_size), + ("expert_parallel_size", topology.expert_parallel_size), + ("ep_fsdp_size", topology.ep_fsdp_size or 0), + ("tensor_parallel_size", topology.tensor_parallel_size), + ("pipeline_parallel_size", topology.pipeline_parallel_size), + ("ulysses_parallel_size", topology.ulysses_parallel_size), + ("ringattn_parallel_size", topology.ringattn_parallel_size), + ) + for field_name, topology_value in dimensions: + observed_values = [ + _known_or_default_parallel_size(getattr(point, field_name)) + if field_name.endswith("_parallel_size") + else getattr(point, field_name) + for point in same_sequence_points + if getattr(point, field_name) is not None + ] + if not observed_values: + continue + if topology_value < min(observed_values) or topology_value > max(observed_values): + return "outside_measured_envelope" + return "inside_measured_envelope" + + +def _same_or_longer_sequence_for_risk(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + if point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len): + return True + return ( + point.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len >= point.sample_packing_sequence_len + ) + + +def _real_routing_boundary_flags( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + *, + raw_config: dict[str, Any] | None, + prediction_confidence: str, +) -> list[str]: + if raw_config is None or prediction_confidence == "calibrated" or _config_balanced_routing(raw_config): + return [] + + real_router_fits: list[BenchmarkBehaviorPoint] = [] + balanced_router_fits: list[BenchmarkBehaviorPoint] = [] + for point in behavior_points: + if point.tokens_per_sec is None or point.micro_batch_size is None: + continue + if behavior_point_model_mismatches(point, raw_config): + continue + if not _same_or_longer_sequence_for_risk(point, topology): + continue + if not _point_matches_parallel_dims_for_risk(point, topology): + continue + + mismatches = set(behavior_point_workload_mismatches(point, raw_config)) + if not mismatches: + real_router_fits.append(point) + elif point.balanced_routing is True and mismatches <= {"balanced_routing"}: + balanced_router_fits.append(point) + + if not real_router_fits or not balanced_router_fits: + return [] + + max_real_router_mbs = max(point.micro_batch_size or 0 for point in real_router_fits) + max_real_router_global_batch = max(point.global_batch_size or 0 for point in real_router_fits) + outside_real_router_fit_envelope = topology.micro_batch_size > max_real_router_mbs or ( + max_real_router_global_batch > 0 and topology.global_batch_size > max_real_router_global_batch + ) + if not outside_real_router_fit_envelope: + return [] + + balanced_fit_covers_target = any( + (point.micro_batch_size or 0) >= topology.micro_batch_size + and (point.global_batch_size or 0) >= topology.global_batch_size + for point in balanced_router_fits + ) + if not balanced_fit_covers_target: + return [] + return ["balanced_routing_only_fit_boundary", "real_routing_outside_fit_envelope"] + + +def _candidate_risk_flags( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + behavior: BenchmarkBehaviorPrediction, + *, + raw_config: dict[str, Any] | None, + calibration_scope: str, + prediction_confidence: str, + communication: CommLedger | None, +) -> list[str]: + flags: list[str] = [] + if prediction_confidence != "calibrated": + flags.append("requires_remeasurement") + if prediction_confidence == "cross_model_extrapolated": + flags.append("cross_model_analog") + if calibration_scope.startswith("outside"): + flags.append(calibration_scope) + if behavior.correctness_status and behavior.correctness_status != "k3_pass": + flags.append(f"correctness_{behavior.correctness_status}") + if communication is not None and prediction_confidence != "calibrated": + flags.extend( + _communication_risk_flags( + behavior_points, + topology, + behavior, + communication, + prediction_confidence=prediction_confidence, + ) + ) + flags.extend( + _real_routing_boundary_flags( + behavior_points, + topology, + raw_config=raw_config, + prediction_confidence=prediction_confidence, + ) + ) + + for phase, share in behavior.phase_time_share.items(): + bucket = _phase_bucket(phase) + if bucket == "input" and share >= 0.15: + flags.append("input_pipeline_bottleneck") + elif bucket == "optimizer" and share >= 0.25: + flags.append("optimizer_bottleneck") + elif bucket == "communication" and share >= 0.20: + flags.append("communication_bottleneck") + + matched_labels = {part.strip() for part in (behavior.matched_label or "").split(",") if part.strip()} + for point in behavior_points: + if raw_config is not None and point.label in matched_labels: + for mismatch in behavior_point_workload_mismatches(point, raw_config): + if prediction_confidence == "cross_model_extrapolated" and mismatch == "model_ref": + continue + flags.append(f"runtime_mismatch:{mismatch}") + if point.label in matched_labels: + if point.expert_parallel_size is not None and point.expert_parallel_size != topology.expert_parallel_size: + flags.append("parallelism_extrapolation:ep") + if point.ep_fsdp_size is not None and point.ep_fsdp_size != topology.ep_fsdp_size: + flags.append("parallelism_extrapolation:ep_fsdp") + if _known_or_default_parallel_size(point.tensor_parallel_size) != topology.tensor_parallel_size: + flags.append("parallelism_extrapolation:tp") + if _known_or_default_parallel_size(point.pipeline_parallel_size) != topology.pipeline_parallel_size: + flags.append("parallelism_extrapolation:pp") + if _known_or_default_parallel_size(point.ulysses_parallel_size) != topology.ulysses_parallel_size: + flags.append("parallelism_extrapolation:ulysses") + if _known_or_default_parallel_size(point.ringattn_parallel_size) != topology.ringattn_parallel_size: + flags.append("parallelism_extrapolation:ring") + if point.measured_steps is not None and point.measured_steps < 3: + flags.append("short_measurement_window") + if point.tokens_per_sec_cv is not None and point.tokens_per_sec_cv >= 0.15: + flags.append("high_throughput_variance") + same_sequence = _same_or_longer_sequence_for_risk(point, topology) + workload_compatible = raw_config is None or behavior_point_matches_workload(point, raw_config) + if not workload_compatible or not same_sequence or not _point_matches_parallel_dims_for_risk(point, topology): + continue + + if point.status == "allocator_pressure_slowdown": + at_or_beyond_mbs = ( + point.micro_batch_size is not None and topology.micro_batch_size >= point.micro_batch_size + ) + at_or_beyond_global_batch = ( + point.global_batch_size is not None and topology.global_batch_size >= point.global_batch_size + ) + if point.label in matched_labels: + flags.append("matched_allocator_pressure_slowdown") + elif at_or_beyond_mbs or at_or_beyond_global_batch: + flags.append(f"allocator_pressure_boundary:{point.label}") + + if point.correctness_status == "oom": + at_or_beyond_sequence = ( + point.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len is not None + and topology.sample_packing_sequence_len >= point.sample_packing_sequence_len + ) + if point.micro_batch_size is not None: + at_or_beyond_batch = topology.micro_batch_size >= point.micro_batch_size + else: + at_or_beyond_batch = ( + point.global_batch_size is not None and topology.global_batch_size >= point.global_batch_size + ) + if point.label in matched_labels or (at_or_beyond_sequence and at_or_beyond_batch): + flags.append(f"observed_oom_boundary:{point.label}") + + return sorted(set(flags)) + + +def _risk_adjusted_score( + score_tokens_per_sec: float | None, + *, + calibration_scope: str, + calibration_distance: float | None, + risk_flags: list[str], + feasibility_status: str, +) -> float | None: + if score_tokens_per_sec is None: + return None + + multiplier = 1.0 + if calibration_scope == "inside_measured_envelope": + multiplier *= 0.85 + elif calibration_scope == "outside_measured_envelope": + multiplier *= 0.65 + elif calibration_scope == "outside_sequence_calibration_envelope": + multiplier *= 0.35 + elif calibration_scope == "no_calibration": + multiplier *= 0.20 + elif calibration_scope == "cross_model_analog": + multiplier *= 0.25 + + if calibration_scope != "exact_calibrated" and calibration_distance is not None and calibration_distance > 0: + multiplier *= max(0.70, 1.0 - 0.04 * calibration_distance) + + if "matched_allocator_pressure_slowdown" in risk_flags: + multiplier *= 0.35 + elif any(flag.startswith("allocator_pressure_boundary:") for flag in risk_flags): + multiplier *= 0.50 + if any(flag.startswith("observed_oom_boundary:") for flag in risk_flags): + multiplier *= 0.25 + + for flag in risk_flags: + if flag == "correctness_k3_fail": + multiplier *= 0.50 + elif flag in {"correctness_not_promoted", "correctness_raw_speed_not_promoted_without_matching_k3_pass"}: + multiplier *= 0.95 + elif flag == "correctness_not_promoted_extrapolated": + multiplier *= 0.90 + elif flag == "correctness_runtime_failure_after_steps": + multiplier *= 0.45 + elif flag == "correctness_missing_calibration": + multiplier *= 0.50 + elif flag.startswith("correctness_"): + multiplier *= 0.75 + elif flag == "short_measurement_window": + multiplier *= 0.80 + elif flag == "high_throughput_variance": + multiplier *= 0.75 + elif flag == "input_pipeline_bottleneck": + multiplier *= 0.70 + elif flag == "optimizer_bottleneck": + multiplier *= 0.85 + elif flag == "real_routing_outside_fit_envelope": + multiplier *= 0.55 + elif flag == "balanced_routing_only_fit_boundary": + multiplier *= 0.80 + elif flag == "cross_model_analog": + multiplier *= 0.70 + elif flag == "parallelism_extrapolation:tp": + multiplier *= 0.80 + elif flag == "parallelism_extrapolation:pp": + multiplier *= 0.75 + elif flag in {"parallelism_extrapolation:ulysses", "parallelism_extrapolation:ring"}: + multiplier *= 0.85 + elif flag == "parallelism_extrapolation:ep": + multiplier *= 0.90 + elif flag == "parallelism_extrapolation:ep_fsdp": + multiplier *= 0.95 + elif flag.startswith("simulator_surface_unsupported:"): + multiplier *= 0.0 + elif flag.startswith("simulator_surface_partial:"): + multiplier *= 0.60 + elif flag in {"communication_cross_node:tp", "communication_cross_node:pp", "communication_cross_node:cp"}: + multiplier *= 0.85 + elif flag in { + "communication_cross_node:ep", + "communication_cross_node:fsdp", + "communication_cross_node:ep_fsdp", + }: + multiplier *= 0.92 + + if feasibility_status.endswith("_high_pressure"): + multiplier *= 0.85 + elif feasibility_status.endswith("_moderate_pressure"): + multiplier *= 0.95 + + return round(score_tokens_per_sec * multiplier, 3) + + +def _prediction_uncertainty_fraction( + behavior: BenchmarkBehaviorPrediction, + *, + prediction_confidence: str, + calibration_scope: str, + calibration_distance: float | None, + risk_flags: list[str], + memory_coverage_status: str, +) -> float | None: + if behavior.tokens_per_sec is None: + return None + + if prediction_confidence == "calibrated": + fraction = 0.05 + elif prediction_confidence == "cross_model_extrapolated": + fraction = 0.45 + else: + fraction = 0.25 + + if calibration_scope == "inside_measured_envelope": + fraction += 0.05 + elif calibration_scope == "outside_measured_envelope": + fraction += 0.15 + elif calibration_scope == "outside_sequence_calibration_envelope": + fraction += 0.25 + elif calibration_scope == "no_calibration": + fraction += 0.35 + elif calibration_scope == "cross_model_analog": + fraction += 0.25 + + if calibration_distance is not None and calibration_distance > 0: + fraction += min(0.25, 0.03 * calibration_distance) + if behavior.tokens_per_sec_cv is not None: + fraction = max(fraction, min(0.60, behavior.tokens_per_sec_cv)) + + for flag in risk_flags: + if flag == "high_throughput_variance": + fraction += 0.15 + elif flag == "short_measurement_window": + fraction += 0.10 + elif flag == "cross_model_analog": + fraction += 0.15 + elif flag.startswith("cross_model_support:"): + fraction += 0.10 + elif flag == "runtime_mismatch:gradient_checkpointing_method": + # Checkpointing-method extrapolation is the one runtime channel with measured leave-one-out + # MISSES (q35 no_recompute rows: 74% and 155% error vs 0.54-0.58 predicted uncertainty); + # every other runtime mismatch stays at the small generic bump. + fraction += 0.30 + elif flag.startswith("runtime_mismatch:"): + fraction += 0.04 + elif flag.startswith("communication_cross_node:"): + fraction += 0.04 + elif flag.startswith("parallelism_extrapolation:"): + fraction += 0.05 + elif flag.startswith("simulator_surface_unsupported:"): + fraction += 0.50 + elif flag.startswith("simulator_surface_partial:"): + fraction += 0.20 + elif flag == "memory_extrapolated_overhead": + fraction += 0.08 + elif flag == "allocator_pressure_risk_extrapolated_peak": + # Replaced the binary 0.5x throughput prior (falsified by the mbs3 prospective holdout): + # an extrapolated peak crossing 85% reserved widens the interval instead of halving the score. + fraction += 0.15 + elif flag.startswith("observed_oom_boundary:") or flag.startswith("allocator_pressure_boundary:"): + fraction += 0.12 + + if memory_coverage_status == "analytic_floor_only": + fraction += 0.10 + elif memory_coverage_status.startswith("calibrated_overhead"): + fraction += 0.06 + elif memory_coverage_status.startswith("extrapolated"): + fraction += 0.08 + + # Calibrated-class predictions are bounded by measured leave-one-out evidence, not the raw flag + # stack: across the q35/q235/q36/q397 calibration holdouts every calibrated row's error is <= 11.6% + # (most 1-6%), while flag stacking pushed several to 0.34-0.38 (20x+ over-conservative), turning + # real promotable-tie decisions into artificial interval overlaps. Cap at 0.20 (1.7x margin over + # the observed max); the CV floor above still lifts noisy references. + if prediction_confidence == "calibrated": + fraction = min(fraction, 0.20) + # Step-time-fit extrapolations: after the definition-consistent fit (2026-07-04) the q35 65k + # ga-family leave-one-out errors are 1.8-5.1% with the ga=1 regime-boundary row at 21.1%; the + # flag stack had pushed these to 0.72-0.79. Cap at 0.35 (1.65x margin over the observed max). + elif behavior.status == "extrapolated_step_time_fit": + fraction = min(fraction, 0.35) + + return round(min(max(fraction, 0.0), 0.95), 3) + + +def _prediction_interval(score: float | None, uncertainty_fraction: float | None) -> tuple[float | None, float | None]: + if score is None or uncertainty_fraction is None: + return None, None + return round(max(score * (1.0 - uncertainty_fraction), 0.0), 3), round(score * (1.0 + uncertainty_fraction), 3) + + +def _recommendation( + *, + feasible: bool, + promotable: bool, + feasibility_status: str, + risk_flags: list[str], +) -> str: + if feasibility_status == "observed_oom": + return "avoid_observed_oom" + if feasibility_status == "unsupported_simulator_surface": + return "build_simulator_backend_before_ranking" + if not feasible: + return "do_not_launch_unscored" + if any(flag.startswith("simulator_surface_unsupported:") for flag in risk_flags): + return "build_simulator_backend_before_ranking" + if any(flag.startswith("simulator_surface_partial:") for flag in risk_flags): + return "measure_partial_simulator_surface" + if promotable: + return "promote_candidate" + if "matched_allocator_pressure_slowdown" in risk_flags or any( + flag.startswith("allocator_pressure_boundary:") for flag in risk_flags + ): + return "measure_allocator_boundary" + if any(flag.startswith("observed_oom_boundary:") for flag in risk_flags): + return "remeasure_after_memory_fix" + if "real_routing_outside_fit_envelope" in risk_flags: + return "measure_real_routing_boundary" + if "correctness_runtime_failure_after_steps" in risk_flags: + return "debug_runtime_failure" + if "short_measurement_window" in risk_flags or "high_throughput_variance" in risk_flags: + return "remeasure_for_stability" + if "input_pipeline_bottleneck" in risk_flags: + return "fix_input_pipeline_before_ranking" + if "cross_model_analog" in risk_flags: + return "measure_cross_model_analog" + if "requires_remeasurement" in risk_flags: + return "remeasure_before_ranking" + if any(flag.startswith("correctness_") for flag in risk_flags): + return "correctness_gate_required" + return "review_candidate" + + +def _mutated_config( + base_config: dict[str, Any], + *, + world_size: int, + micro_batch_size: int, + gradient_accumulation_steps: int, + expert_parallel_size: int, + tensor_parallel_size: int, + pipeline_parallel_size: int, + ulysses_parallel_size: int, + ringattn_parallel_size: int, + data_parallel_replicate_size: int | None = None, + data_parallel_shard_size: int | None = None, +) -> dict[str, Any]: + raw_config = copy.deepcopy(base_config) + surface = requested_simulator_surface(raw_config) + if surface == "server_forward_backward": + nested_server = _section(raw_config, "server") + train = nested_server if nested_server else raw_config + else: + train = raw_config.setdefault("train", {}) + if not isinstance(train, dict): + train = {} + raw_config["train"] = train + + train["micro_batch_size"] = micro_batch_size + train["gradient_accumulation_steps"] = gradient_accumulation_steps + train["expert_parallel_size"] = expert_parallel_size + train["tensor_parallel_size"] = tensor_parallel_size + train["pipeline_parallel_size"] = pipeline_parallel_size + train["ulysses_parallel_size"] = ulysses_parallel_size + train["ringattn_parallel_size"] = ringattn_parallel_size + + non_dp_size = tensor_parallel_size * pipeline_parallel_size * ulysses_parallel_size * ringattn_parallel_size + if non_dp_size <= 0 or world_size % non_dp_size != 0: + raise ValueError("world_size is not divisible by non-DP parallelism product") + data_parallel_size = world_size // non_dp_size + replicate_size = data_parallel_replicate_size + shard_size = data_parallel_shard_size + if replicate_size is None and shard_size is None: + replicate_size = 1 + shard_size = data_parallel_size + elif replicate_size is None: + if shard_size is None or shard_size <= 0 or data_parallel_size % shard_size != 0: + raise ValueError("data_parallel_shard_size must divide data_parallel_size") + replicate_size = data_parallel_size // shard_size + elif shard_size is None: + if replicate_size <= 0 or data_parallel_size % replicate_size != 0: + raise ValueError("data_parallel_replicate_size must divide data_parallel_size") + shard_size = data_parallel_size // replicate_size + elif replicate_size <= 0 or shard_size <= 0 or replicate_size * shard_size != data_parallel_size: + raise ValueError("data_parallel_replicate_size * data_parallel_shard_size must equal data_parallel_size") + train["data_parallel_replicate_size"] = replicate_size + train["data_parallel_shard_size"] = shard_size + if pipeline_parallel_size > 1: + train["gradient_accumulation_steps"] = max( + int(train.get("gradient_accumulation_steps", 1) or 1), pipeline_parallel_size + ) + return raw_config + + +def _set_balanced_routing(raw_config: dict[str, Any], balanced_routing: bool) -> None: + simulator = raw_config.setdefault("simulator", {}) + if isinstance(simulator, dict): + simulator["balanced_routing"] = balanced_routing + + +def _topology_label(topology: Topology) -> str: + return ( + f"mbs{topology.micro_batch_size}-gb{topology.global_batch_size}-" + f"ep{topology.expert_parallel_size}-efsdp{topology.ep_fsdp_size}-" + f"tp{topology.tensor_parallel_size}-pp{topology.pipeline_parallel_size}-" + f"u{topology.ulysses_parallel_size}-r{topology.ringattn_parallel_size}" + ) + + +def _candidate_topology_label(topology: Topology, *, include_sequence_len: bool) -> str: + label = _topology_label(topology) + if include_sequence_len and topology.sample_packing_sequence_len is not None: + return f"{label}-seq{topology.sample_packing_sequence_len}" + return label + + +def _communication_ledger(topology: Topology) -> CommLedger: + local_world_size = max(topology.local_world_size, 1) + sequence_parallel_size = max(topology.sequence_parallel_size, 1) + ranks_per_pipeline_stage = max(topology.world_size // max(topology.pipeline_parallel_size, 1), 1) + ep_fsdp_size = topology.ep_fsdp_size or 1 + + tensor_cross = topology.tensor_parallel_size > local_world_size + pipeline_cross = topology.pipeline_parallel_size > local_world_size + expert_cross = topology.expert_parallel_size > local_world_size + context_cross = sequence_parallel_size > local_world_size + fsdp_cross = topology.data_parallel_shard_size > local_world_size + ep_fsdp_cross = ep_fsdp_size > local_world_size + + dimensions: list[str] = [] + notes = [ + f"node_count={topology.node_count}", + f"local_world_size={topology.local_world_size}", + f"ranks_per_pipeline_stage={ranks_per_pipeline_stage}", + ] + if tensor_cross: + dimensions.append("tp") + notes.append(f"tp_group_size={topology.tensor_parallel_size}:cross_node") + if pipeline_cross: + dimensions.append("pp") + notes.append(f"pp_group_size={topology.pipeline_parallel_size}:cross_node") + if expert_cross: + dimensions.append("ep") + notes.append(f"ep_group_size={topology.expert_parallel_size}:cross_node") + if context_cross: + dimensions.append("cp") + notes.append(f"cp_group_size={sequence_parallel_size}:cross_node") + if fsdp_cross: + dimensions.append("fsdp") + notes.append(f"dp_shard_size={topology.data_parallel_shard_size}:cross_node") + if ep_fsdp_cross: + dimensions.append("ep_fsdp") + notes.append(f"ep_fsdp_size={ep_fsdp_size}:cross_node") + + return CommLedger( + tensor_parallel_cross_node=tensor_cross, + pipeline_parallel_cross_node=pipeline_cross, + expert_parallel_cross_node=expert_cross, + context_parallel_cross_node=context_cross, + fsdp_cross_node=fsdp_cross or ep_fsdp_cross, + cross_node_dimensions=dimensions, + notes=notes, + ) + + +def _communication_risk_flags( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + behavior: BenchmarkBehaviorPrediction, + communication: CommLedger, + *, + prediction_confidence: str, +) -> list[str]: + if prediction_confidence == "calibrated": + return [] + matched_labels = {part.strip() for part in (behavior.matched_label or "").split(",") if part.strip()} + if not matched_labels: + return [f"communication_cross_node:{dimension}" for dimension in communication.cross_node_dimensions] + + local_world_size = max(topology.local_world_size, 1) + flags: set[str] = set() + for point in behavior_points: + if point.label not in matched_labels: + continue + reference_cp = _known_or_default_parallel_size(point.ulysses_parallel_size) * _known_or_default_parallel_size( + point.ringattn_parallel_size + ) + checks = ( + ("tp", topology.tensor_parallel_size, _known_or_default_parallel_size(point.tensor_parallel_size)), + ("pp", topology.pipeline_parallel_size, _known_or_default_parallel_size(point.pipeline_parallel_size)), + ("cp", topology.sequence_parallel_size, reference_cp), + ("ep", topology.expert_parallel_size, point.expert_parallel_size or 1), + ("ep_fsdp", topology.ep_fsdp_size or 1, point.ep_fsdp_size or 1), + ) + for dimension, target_size, reference_size in checks: + if target_size > local_world_size and reference_size <= local_world_size: + flags.add(f"communication_cross_node:{dimension}") + return sorted(flags) + + +def _reference_tokens_per_gpu(point: BenchmarkBehaviorPoint, topology: Topology) -> float | None: + if point.tokens_per_sec is None: + return None + gpu_count = point.gpu_count or topology.world_size + if gpu_count <= 0: + return None + return point.tokens_per_sec / gpu_count + + +def _runtime_mismatches_without_model(point: BenchmarkBehaviorPoint, raw_config: dict[str, Any]) -> list[str]: + return [mismatch for mismatch in behavior_point_workload_mismatches(point, raw_config) if mismatch != "model_ref"] + + +def _target_runtime_signature(raw_config: dict[str, Any] | None) -> str: + if raw_config is None: + return "unknown" + parts: list[str] = [] + for section_name, field_name in _SCENARIO_RUNTIME_SIGNATURE_FIELDS: + value = _section(raw_config, section_name).get(field_name) + if value is not None: + parts.append(f"{field_name}={value}") + return ",".join(parts) if parts else "unknown" + + +def _reference_safety_score(point: BenchmarkBehaviorPoint) -> int: + if point.correctness_status == "oom": + return -1 + if point.correctness_status == "k3_pass": + return 3 + if point.correctness_status in (None, "not_promoted", "raw_speed_not_promoted_without_matching_k3_pass"): + return 2 + if point.correctness_status == "not_promoted_extrapolated": + return 1 + return 0 + + +def _is_stable_cross_model_reference(point: BenchmarkBehaviorPoint) -> bool: + return point.correctness_status in { + None, + "k3_pass", + "not_promoted", + "raw_speed_not_promoted_without_matching_k3_pass", + } + + +def _reference_throughput_quality_tier(point: BenchmarkBehaviorPoint, topology: Topology) -> int: + per_gpu = _reference_tokens_per_gpu(point, topology) + if per_gpu is None or per_gpu <= 0: + return 0 + return int(math.log2(max(per_gpu, 1.0))) + + +def _sequence_len_ratio(point: BenchmarkBehaviorPoint, topology: Topology) -> float | None: + if point.sample_packing_sequence_len is None or topology.sample_packing_sequence_len is None: + return None + if point.sample_packing_sequence_len <= 0 or topology.sample_packing_sequence_len <= 0: + return None + return topology.sample_packing_sequence_len / point.sample_packing_sequence_len + + +def _sequence_len_matches_for_cross_model(point: BenchmarkBehaviorPoint, topology: Topology) -> bool: + ratio = _sequence_len_ratio(point, topology) + if ratio is None: + return False + lower, upper = _EXTRAPOLATED_CROSS_MODEL_SEQUENCE_RATIO_WINDOW + return lower <= ratio <= upper + + +def _active_param_proxy(metadata: ModelMetadata) -> float | None: + hidden = metadata.hidden_size + layers = metadata.num_hidden_layers + if hidden is None or layers is None: + return None + + head_dim = metadata.head_dim + if head_dim is None and metadata.num_attention_heads: + head_dim = hidden // metadata.num_attention_heads + if head_dim is None: + return None + + attention_heads = metadata.num_attention_heads or 1 + key_value_heads = metadata.num_key_value_heads or attention_heads + q_proj = hidden * attention_heads * head_dim + k_proj = hidden * key_value_heads * head_dim + v_proj = hidden * key_value_heads * head_dim + o_proj = attention_heads * head_dim * hidden + attention_params = layers * (q_proj + k_proj + v_proj + o_proj) + + dense_mlp_params = 0 + has_routed_experts = metadata.num_experts is not None and metadata.moe_intermediate_size is not None + if metadata.intermediate_size is not None and not has_routed_experts: + dense_mlp_params = layers * 3 * hidden * metadata.intermediate_size + + shared_expert_params = 0 + if metadata.shared_expert_intermediate_size is not None: + shared_expert_params = layers * 3 * hidden * metadata.shared_expert_intermediate_size + + active_expert_params = 0 + if has_routed_experts: + if metadata.top_k is None or metadata.moe_intermediate_size is None: + return None + active_expert_params = layers * metadata.top_k * 3 * hidden * metadata.moe_intermediate_size + + lm_head_params = 0 + if metadata.vocab_size is not None and not metadata.tie_word_embeddings: + lm_head_params = metadata.vocab_size * hidden + + return float(attention_params + dense_mlp_params + shared_expert_params + active_expert_params + lm_head_params) + + +def _metadata_for_model_ref(model_ref: str | None) -> ModelMetadata | None: + if not model_ref: + return None + return resolve_model_metadata({"model": {"model_path": model_ref}}, hf_cache_roots=[]) + + +@dataclass(frozen=True) +class _CrossModelScale: + factor: float + note: str + active_param_ratio: float + reference_active_params_b: float + target_active_params_b: float + + +def _cross_model_scale( + reference: BenchmarkBehaviorPoint, + raw_config: dict[str, Any], +) -> _CrossModelScale | None: + target_model_ref = model_ref_from_config(raw_config) + target_metadata = resolve_model_metadata(raw_config, hf_cache_roots=[]) + reference_metadata = _metadata_for_model_ref(reference.model_ref) + if target_model_ref is None or reference_metadata is None: + return None + + target_proxy = _active_param_proxy(target_metadata) + reference_proxy = _active_param_proxy(reference_metadata) + if target_proxy is None or reference_proxy is None or target_proxy <= 0 or reference_proxy <= 0: + return None + + raw_ratio = reference_proxy / target_proxy + scale = min(1.20, max(0.18, raw_ratio**0.90)) + reference_active_params_b = reference_proxy / 1_000_000_000 + target_active_params_b = target_proxy / 1_000_000_000 + note = ( + f"cross-model active-param scale {reference.model_ref}->{target_model_ref}: " + f"{reference_active_params_b:.3f}B/{target_active_params_b:.3f}B => {scale:.3f}" + ) + return _CrossModelScale( + factor=scale, + note=note, + active_param_ratio=raw_ratio, + reference_active_params_b=reference_active_params_b, + target_active_params_b=target_active_params_b, + ) + + +@dataclass(frozen=True) +class _MemoryPeakEstimate: + peak_gb: float + overhead_gb: float + source_label: str + notes: list[str] + + +def _point_int_or_default(point: BenchmarkBehaviorPoint, field_name: str, default: int = 1) -> int: + value = getattr(point, field_name) + return int(value) if value is not None else default + + +def _raw_config_for_behavior_point( + base_config: dict[str, Any], + point: BenchmarkBehaviorPoint, + *, + default_world_size: int, +) -> dict[str, Any] | None: + if point.micro_batch_size is None or point.global_batch_size is None: + return None + world_size = point.gpu_count or default_world_size + tp = _point_int_or_default(point, "tensor_parallel_size") + pp = _point_int_or_default(point, "pipeline_parallel_size") + ulysses = _point_int_or_default(point, "ulysses_parallel_size") + ring = _point_int_or_default(point, "ringattn_parallel_size") + ep = _point_int_or_default(point, "expert_parallel_size") + non_dp = tp * pp * ulysses * ring + if non_dp <= 0 or world_size % non_dp: + return None + dp = world_size // non_dp + ga_denominator = point.micro_batch_size * dp + if ga_denominator <= 0 or point.global_batch_size % ga_denominator: + return None + gradient_accumulation_steps = point.global_batch_size // ga_denominator + if gradient_accumulation_steps <= 0: + return None + + raw_config = _mutated_config( + base_config, + world_size=world_size, + micro_batch_size=point.micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + expert_parallel_size=ep, + tensor_parallel_size=tp, + pipeline_parallel_size=pp, + ulysses_parallel_size=ulysses, + ringattn_parallel_size=ring, + data_parallel_replicate_size=point.data_parallel_replicate_size, + data_parallel_shard_size=point.data_parallel_shard_size, + ) + if point.sample_packing_sequence_len is not None: + _section(raw_config, "data")["sample_packing_sequence_len"] = point.sample_packing_sequence_len + train = _section(raw_config, "train") + if point.skip_param_upcast is not None: + train["skip_param_upcast"] = point.skip_param_upcast + if point.muon_momentum is not None: + train["muon_momentum"] = point.muon_momentum + if point.balanced_routing is not None: + _set_balanced_routing(raw_config, point.balanced_routing) + return raw_config + + +def _memory_overhead_scale( + reference_topology: Topology, + reference_shape: ShapeLedger, + target_topology: Topology, + target_shape: ShapeLedger, +) -> tuple[float, list[str]]: + def ratio(target: float | int | None, reference: float | int | None) -> float: + if target is None or reference in (None, 0): + return 1.0 + return max(float(target) / float(reference), 0.01) + + token_ratio = ratio( + target_shape.tokens_per_model_rank_per_microbatch, reference_shape.tokens_per_model_rank_per_microbatch + ) + target_ep_slots = ( + max(target_shape.ep_rank_slots_per_microbatch) if target_shape.ep_rank_slots_per_microbatch else None + ) + reference_ep_slots = ( + max(reference_shape.ep_rank_slots_per_microbatch) if reference_shape.ep_rank_slots_per_microbatch else None + ) + routed_ratio = ratio(target_ep_slots, reference_ep_slots) + sequence_ratio = ratio(target_topology.sample_packing_sequence_len, reference_topology.sample_packing_sequence_len) + + same_parallel_sequence_shape = ( + reference_topology.expert_parallel_size == target_topology.expert_parallel_size + and reference_topology.ep_fsdp_size == target_topology.ep_fsdp_size + and reference_topology.tensor_parallel_size == target_topology.tensor_parallel_size + and reference_topology.pipeline_parallel_size == target_topology.pipeline_parallel_size + and reference_topology.sequence_parallel_size == target_topology.sequence_parallel_size + and reference_topology.sample_packing_sequence_len == target_topology.sample_packing_sequence_len + ) + if same_parallel_sequence_shape: + token_component = token_ratio + routed_component = routed_ratio + token_weight = 0.50 + routed_weight = 0.40 + sequence_weight = 0.10 + scaling_regime = "linear_same_topology" + else: + token_component = math.sqrt(token_ratio) + routed_component = math.sqrt(routed_ratio) + token_weight = 0.50 + routed_weight = 0.35 + sequence_weight = 0.15 + scaling_regime = "sqrt_cross_topology" + sequence_component = sequence_ratio**0.20 + scale = token_weight * token_component + routed_weight * routed_component + sequence_weight * sequence_component + + reference_cp = max(reference_topology.sequence_parallel_size, 1) + target_cp = max(target_topology.sequence_parallel_size, 1) + if target_cp > reference_cp: + scale *= max(0.65, (reference_cp / target_cp) ** 0.15) + elif target_cp < reference_cp: + scale *= min(1.35, (reference_cp / target_cp) ** 0.10) + + scale = min(4.0, max(0.20, scale)) + notes = [ + f"memory_overhead_scale={scale:.3f}", + f"token_ratio={token_ratio:.3f}", + f"routed_ratio={routed_ratio:.3f}", + f"sequence_ratio={sequence_ratio:.3f}", + f"memory_overhead_scaling_regime={scaling_regime}", + ] + if target_cp != reference_cp: + notes.append(f"cp_ratio={target_cp}/{reference_cp}") + return scale, notes + + +def _analytic_activation_fit_gb( + metadata: ModelMetadata, + topology: Topology, + train_config: dict[str, Any], +) -> tuple[float | None, float, list[str]]: + """Analytic activation lower bound (GB) + exact ep_fsdp expert-unshard transient for fit estimates. + + - The activation lower bound comes from the activation ledger. Under pipeline parallelism the + 1F1B in-flight depth (= pipeline_parallel_size) and the per-stage layer split (1/pp of the + layers) CANCEL on the saved-activation term, so no PP multiplier is applied — the measured + balanced PP2 row fits at ~64 GB, confirming the cancellation; the real-routing PP2 OOM was the + mbs2 routing-imbalance effect, not a PP term. + - With ``ep_fsdp_size > 1`` FSDP re-gathers each layer's full expert group for compute: the exact + transient is (param + grad) x gathered expert bytes per layer x prefetch depth 2. This is a + LOWER bound: the measured ep1/efsdp16 65k boundary sits ~13 GB above floor+activation+transient, + an unattributed re-gather-path residual named in the notes rather than fitted away. + """ + notes: list[str] = [] + try: + ledger = activation_ledger(metadata, topology, train_config, seq_len=topology.sample_packing_sequence_len) + except Exception: # pragma: no cover - defensive: fall back to heuristic scaling + return None, 0.0, ["activation_ledger_unavailable"] + terms = ledger.get("terms") or {} + total = 0.0 + for name, term in terms.items(): + gb = term.get("gb") if isinstance(term, dict) else None + if not gb: + continue + total += float(gb) + ep_fsdp = int(topology.ep_fsdp_size or 1) + transient = 0.0 + if ( + ep_fsdp > 1 + and metadata.num_experts + and metadata.moe_intermediate_size + and metadata.hidden_size + and metadata.num_hidden_layers + ): + local_experts = max(int(metadata.num_experts) // max(int(topology.expert_parallel_size), 1), 1) + gathered_layer_gb = ( + local_experts * 3 * int(metadata.hidden_size) * int(metadata.moe_intermediate_size) * 2 / 1024**3 + ) + transient = round(gathered_layer_gb * 2 * 2, 3) # (param + grad) x prefetch depth 2 + notes.append(f"ep_fsdp_unshard_transient_gb={transient}") + notes.append("ep_fsdp_unshard_transient_is_lower_bound_named_residual_at_ep1_65k") + return round(total, 3), transient, notes + + +def _calibrated_memory_peak_estimate( + behavior_points: list[BenchmarkBehaviorPoint], + base_config: dict[str, Any], + raw_config: dict[str, Any], + target_topology: Topology, + target_shape: ShapeLedger, + metadata: ModelMetadata, + *, + default_world_size: int, + default_local_world_size: int, + analytic_peak_floor_gb: float | None, +) -> _MemoryPeakEstimate | None: + if analytic_peak_floor_gb is None: + return None + + estimates: list[tuple[tuple[float, ...], _MemoryPeakEstimate]] = [] + for point in behavior_points: + if point.peak_mem_gb is None or point.correctness_status == "oom": + continue + if behavior_point_model_mismatches(point, raw_config) or not behavior_point_matches_workload(point, raw_config): + continue + reference_config = _raw_config_for_behavior_point( + base_config, + point, + default_world_size=default_world_size, + ) + if reference_config is None: + continue + point_world_size = point.gpu_count or default_world_size + point_local_world_size = min(default_local_world_size, point_world_size) + try: + reference_topology = resolve_topology( + reference_config, + world_size=point_world_size, + local_world_size=point_local_world_size, + ) + except ValueError: + continue + if point.ep_fsdp_size is not None and point.ep_fsdp_size != reference_topology.ep_fsdp_size: + continue + reference_memory = build_memory_ledger( + reference_config, + topology=reference_topology, + model_metadata=metadata, + ) + reference_floor = reference_memory.analytic_peak_floor_gb + if reference_floor is None or point.peak_mem_gb < reference_floor: + continue + overhead_gb = point.peak_mem_gb - reference_floor + reference_shape = build_shape_ledger(reference_topology, balanced_routing=True) + # Analytic-first overhead transfer: scale the measured torch-side residual by the ratio of the + # EXACT activation lower bounds (seq/topology/PP-aware), and add the exact ep_fsdp unshard + # transient delta. The old sqrt heuristic under-scaled the measured 65k boundaries (u1's real + # 4x activation growth became ~2x); it remains only as the fallback when the ledger cannot + # compute either side. + target_act, target_transient, target_act_notes = _analytic_activation_fit_gb( + metadata, target_topology, raw_config.get("train", {}) + ) + reference_act, reference_transient, _ = _analytic_activation_fit_gb( + metadata, reference_topology, reference_config.get("train", {}) + ) + + # The analytic ratio captures topology/sequence-driven activation changes ONLY; when the + # reference and target differ on runtime fields that move activations through other channels + # (offload, checkpointing method, reduce dtype, ce mode), the ratio is wrong-by-construction + # and the heuristic transfer stays in effect (validated on the q235 offload/dtype holdouts). + def _activation_runtime_signature(config: dict[str, Any]) -> tuple[Any, ...]: + train_section = config.get("train", {}) if isinstance(config.get("train"), dict) else {} + return ( + train_section.get("enable_activation_offload"), + train_section.get("gradient_checkpointing_method"), + train_section.get("fsdp_reduce_dtype"), + train_section.get("ce_mode"), + train_section.get("skip_param_upcast"), + ) + + same_activation_runtime = _activation_runtime_signature(raw_config) == _activation_runtime_signature( + reference_config + ) + if target_act and reference_act and same_activation_runtime: + scale = target_act / reference_act + scale_notes = [ + f"memory_overhead_scale={scale:.4f}", + "scaling_regime=analytic_activation_ratio", + f"target_activation_lower_bound_gb={target_act}", + f"reference_activation_lower_bound_gb={reference_act}", + *target_act_notes, + ] + transient_delta = max(0.0, target_transient - reference_transient) + else: + scale, scale_notes = _memory_overhead_scale( + reference_topology, reference_shape, target_topology, target_shape + ) + transient_delta = 0.0 + estimated_overhead = max(0.0, overhead_gb * scale) + transient_delta + estimated_peak = analytic_peak_floor_gb + estimated_overhead + sequence_distance = abs( + math.log( + (target_topology.sample_packing_sequence_len or 1) + / (reference_topology.sample_packing_sequence_len or 1) + ) + ) + parallel_distance = sum( + abs(math.log2(max(target_value, 1) / max(reference_value, 1))) + for target_value, reference_value in ( + (target_topology.expert_parallel_size, reference_topology.expert_parallel_size), + (target_topology.ep_fsdp_size or 1, reference_topology.ep_fsdp_size or 1), + (target_topology.tensor_parallel_size, reference_topology.tensor_parallel_size), + (target_topology.pipeline_parallel_size, reference_topology.pipeline_parallel_size), + (target_topology.sequence_parallel_size, reference_topology.sequence_parallel_size), + ) + ) + batch_distance = abs( + math.log2(max(target_topology.micro_batch_size, 1) / max(reference_topology.micro_batch_size, 1)) + ) + key = ( + -sequence_distance, + -parallel_distance, + -batch_distance, + float(point.peak_mem_gb), + ) + notes = [ + f"memory_overhead_reference={point.label}", + f"reference_peak_gb={point.peak_mem_gb:.3f}", + f"reference_floor_gb={reference_floor:.3f}", + f"reference_overhead_gb={overhead_gb:.3f}", + f"estimated_overhead_gb={estimated_overhead:.3f}", + *scale_notes, + ] + estimates.append( + ( + key, + _MemoryPeakEstimate( + peak_gb=round(estimated_peak, 3), + overhead_gb=round(estimated_overhead, 3), + source_label=point.label, + notes=notes, + ), + ) + ) + + if not estimates: + return None + return max(estimates, key=lambda item: item[0])[1] + + +def _memory_ownership_notes(memory: MemoryLedger) -> list[str]: + ownership_prefixes = ( + "pp_stage=", + "tp_non_expert_shard_size=", + "non_expert_total_shard_size=", + "expert_shard_size=", + "dp_shard_size=", + "cp_fsdp_mode=", + "local_non_expert_params=", + "local_expert_params=", + ) + for bucket in memory.top_memory_buckets: + if bucket.name == "sharded_trainable_params": + return [note for note in bucket.notes if note.startswith(ownership_prefixes)] + return [] + + +def _select_reference_point( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + raw_config: dict[str, Any] | None = None, +) -> BenchmarkBehaviorPoint | None: + usable = [ + point + for point in behavior_points + if point.tokens_per_sec is not None + and point.micro_batch_size is not None + and point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) + and _reference_tokens_per_gpu(point, topology) is not None + and (raw_config is None or not behavior_point_model_mismatches(point, raw_config)) + ] + if not usable: + return None + + workload_compatible = ( + [point for point in usable if behavior_point_matches_workload(point, raw_config)] + if raw_config is not None + else usable + ) + same_ep = [ + point for point in workload_compatible if point.expert_parallel_size in (None, topology.expert_parallel_size) + ] + candidates = same_ep or workload_compatible or usable + + def key(point: BenchmarkBehaviorPoint) -> tuple[float, float, float]: + mismatch_count = len(behavior_point_workload_mismatches(point, raw_config)) if raw_config is not None else 0 + mbs_distance = abs((point.micro_batch_size or 1) - topology.micro_batch_size) + per_gpu = _reference_tokens_per_gpu(point, topology) or 0.0 + return (-mismatch_count, -mbs_distance, per_gpu) + + return max(candidates, key=key) + + +def _select_cross_model_reference_point( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + raw_config: dict[str, Any] | None, +) -> BenchmarkBehaviorPoint | None: + if raw_config is None or model_ref_from_config(raw_config) is None: + return None + + usable = [ + point + for point in behavior_points + if point.tokens_per_sec is not None + and point.micro_batch_size is not None + and _is_stable_cross_model_reference(point) + and behavior_point_model_mismatches(point, raw_config) + and _sequence_len_matches_for_cross_model(point, topology) + and _reference_tokens_per_gpu(point, topology) is not None + and _cross_model_scale(point, raw_config) is not None + and "attention_backend" not in _runtime_mismatches_without_model(point, raw_config) + ] + if not usable: + return None + + def key(point: BenchmarkBehaviorPoint) -> tuple[int, int, float, float, float, float, float]: + runtime_mismatch_count = len(_runtime_mismatches_without_model(point, raw_config)) + mbs_distance = abs((point.micro_batch_size or 1) - topology.micro_batch_size) + ep_distance = 0.0 + if point.expert_parallel_size: + ep_distance = abs(math.log2(topology.expert_parallel_size / point.expert_parallel_size)) + seq_ratio = _sequence_len_ratio(point, topology) or 1.0 + seq_distance = abs(math.log(seq_ratio)) + per_gpu = _reference_tokens_per_gpu(point, topology) or 0.0 + return ( + _reference_safety_score(point), + _reference_throughput_quality_tier(point, topology), + -runtime_mismatch_count, + -seq_distance, + -ep_distance, + -mbs_distance, + per_gpu, + ) + + return max(usable, key=key) + + +def _parallelism_factor(reference: BenchmarkBehaviorPoint, topology: Topology) -> tuple[float, list[str]]: + notes: list[str] = [] + factor = 1.0 + if reference.micro_batch_size: + mbs_ratio = topology.micro_batch_size / reference.micro_batch_size + factor *= min(1.15, max(0.55, mbs_ratio**0.20)) + if reference.expert_parallel_size and reference.expert_parallel_size != topology.expert_parallel_size: + ep_ratio = topology.expert_parallel_size / reference.expert_parallel_size + factor *= max(0.70, 1.0 - 0.04 * abs(math.log2(ep_ratio))) + notes.append(f"EP extrapolated from {reference.expert_parallel_size} to {topology.expert_parallel_size}") + reference_tp = _known_or_default_parallel_size(reference.tensor_parallel_size) + if reference_tp != topology.tensor_parallel_size: + tp_ratio = topology.tensor_parallel_size / reference_tp + factor *= 0.90 ** abs(math.log2(tp_ratio)) + notes.append("TP extrapolation uses conservative communication penalty") + reference_pp = _known_or_default_parallel_size(reference.pipeline_parallel_size) + if reference_pp != topology.pipeline_parallel_size: + pp_delta = abs(topology.pipeline_parallel_size - reference_pp) + factor *= 0.88**pp_delta + notes.append("PP extrapolation uses conservative bubble penalty") + reference_cp = _known_or_default_parallel_size(reference.ulysses_parallel_size) * _known_or_default_parallel_size( + reference.ringattn_parallel_size + ) + if reference_cp != topology.sequence_parallel_size: + cp_ratio = topology.sequence_parallel_size / reference_cp + if topology.sample_packing_sequence_len and topology.sample_packing_sequence_len >= 32768: + factor *= min(1.10, 1.0 + 0.04 * abs(math.log2(cp_ratio))) + else: + factor *= 0.94 ** abs(math.log2(cp_ratio)) + notes.append("SP/CP extrapolation penalized for short-context workload") + return factor, notes + + +def _cross_model_behavior_prediction( + reference: BenchmarkBehaviorPoint, + topology: Topology, + shape: ShapeLedger, + *, + raw_config: dict[str, Any], + device_memory_limit_gb: float, + memory_safety_factor: float, + analytic_peak_floor_gb: float | None, +) -> tuple[BenchmarkBehaviorPrediction, list[str]] | None: + model_scale = _cross_model_scale(reference, raw_config) + if model_scale is None: + return None + model_factor = model_scale.factor + model_note = model_scale.note + ref_per_gpu = _reference_tokens_per_gpu(reference, topology) or 0.0 + if ref_per_gpu <= 0: + return None + + parallel_factor, notes = _parallelism_factor(reference, topology) + memory_factor, _, memory_status = _memory_factor( + analytic_peak_floor_gb, + memory_basis="analytic_floor", + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + seq_ratio = _sequence_len_ratio(reference, topology) or 1.0 + sequence_factor = min(1.08, max(0.88, seq_ratio**-0.12)) + tokens_per_sec_per_gpu = ref_per_gpu * model_factor * parallel_factor * memory_factor * sequence_factor + tokens_per_sec = tokens_per_sec_per_gpu * topology.world_size + step_time_sec = None + if shape.global_tokens_per_train_step and tokens_per_sec: + step_time_sec = shape.global_tokens_per_train_step / tokens_per_sec + + tflops_per_gpu = reference.tflops_per_gpu + if tflops_per_gpu is None and reference.mfu_percent is not None: + tflops_per_gpu = H100_BF16_PROMISED_TFLOPS_PER_GPU * reference.mfu_percent / 100.0 + if tflops_per_gpu is not None and ref_per_gpu: + tflops_per_gpu *= tokens_per_sec_per_gpu / ref_per_gpu + + target_model_ref = model_ref_from_config(raw_config) + runtime_mismatches = _runtime_mismatches_without_model(reference, raw_config) + warnings = [ + f"cross-model analog from {reference.model_ref} to {target_model_ref}; measure before ranking", + model_note, + f"sequence_length_factor={sequence_factor:.3f}", + f"memory feasibility status is {memory_status}", + "correctness must be re-gated before promotion", + ] + exact_sequence_lower, exact_sequence_upper = _EXACT_CROSS_MODEL_SEQUENCE_RATIO_WINDOW + if not exact_sequence_lower <= seq_ratio <= exact_sequence_upper: + warnings.append(f"cross-model sequence ratio outside exact-context window: {seq_ratio:.3f}") + if reference.correctness_status and reference.correctness_status not in {"k3_pass", "not_promoted"}: + warnings.append(f"reference correctness status is {reference.correctness_status}") + if runtime_mismatches: + warnings.append(f"reference runtime knobs differ: {', '.join(runtime_mismatches)}") + warnings.extend(notes) + + return ( + BenchmarkBehaviorPrediction( + status="cross_model_extrapolated", + matched_label=reference.label, + source=reference.source, + tokens_per_sec=round(tokens_per_sec, 3), + tokens_per_sec_per_gpu=round(tokens_per_sec_per_gpu, 3), + step_time_sec=round(step_time_sec, 6) if step_time_sec is not None else None, + mfu_percent=None, + tflops_per_gpu=round(tflops_per_gpu, 3) if tflops_per_gpu is not None else None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=None, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + phase_time_sec=reference.phase_time_sec, + phase_time_share=reference.phase_time_share, + phase_memory_peak_gb=reference.phase_memory_peak_gb, + model_ref=target_model_ref, + balanced_routing=reference.balanced_routing, + correctness_status="not_promoted_extrapolated", + cross_model_active_param_ratio=round(model_scale.active_param_ratio, 3), + cross_model_active_param_scale=round(model_scale.factor, 3), + cross_model_reference_active_params_b=round(model_scale.reference_active_params_b, 3), + cross_model_target_active_params_b=round(model_scale.target_active_params_b, 3), + cross_model_sequence_length_factor=round(sequence_factor, 3), + cross_model_parallelism_factor=round(parallel_factor, 3), + cross_model_memory_factor=round(memory_factor, 3), + warnings=warnings, + ), + ["cross_model_analog", model_note], + ) + + +def _step_time_fit_prediction( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + shape: ShapeLedger, + raw_config: dict[str, Any] | None = None, +) -> BenchmarkBehaviorPrediction | None: + if topology.sample_packing_sequence_len is None or shape.global_tokens_per_train_step is None: + return None + compatible = [ + point + for point in behavior_points + if point.tokens_per_sec is not None + and point.global_batch_size is not None + and point.micro_batch_size == topology.micro_batch_size + and point.sample_packing_sequence_len in (None, topology.sample_packing_sequence_len) + and point.expert_parallel_size in (None, topology.expert_parallel_size) + and point.ep_fsdp_size in (None, topology.ep_fsdp_size) + # The fit family must share the full sharding regime: rows at the same nominal global batch + # but a different world size or dp-replicate split have different per-microbatch step + # structure (measured: the repl=1 gate-lane gbs64 row steps 88.3s vs the repl=2 ga16 row's + # 101.6s), and best-by-gbs cherry-picking across regimes bent the line by ~20%. + and point.gpu_count in (None, topology.world_size) + and ( + point.data_parallel_replicate_size is None + or point.data_parallel_replicate_size == topology.data_parallel_replicate_size + ) + and _point_matches_topology_parallel_dims(point, topology) + and (raw_config is None or behavior_point_matches_workload(point, raw_config)) + ] + best_by_global_batch: dict[int, BenchmarkBehaviorPoint] = {} + for point in compatible: + current = best_by_global_batch.get(point.global_batch_size) + if current is None or (point.tokens_per_sec or 0.0) > (current.tokens_per_sec or 0.0): + best_by_global_batch[point.global_batch_size] = point + fit_points = sorted(best_by_global_batch.values(), key=lambda point: point.global_batch_size or 0) + if len(fit_points) < 2: + return None + + x_values: list[float] = [] + y_values: list[float] = [] + for point in fit_points: + tokens = point.global_batch_size * topology.sample_packing_sequence_len + # Definition-consistent y: the fit converts the predicted step BACK to tokens_per_sec via + # nominal tokens, so y must be nominal_tokens / tokens_per_sec. The measured step_time_sec + # field disagrees with that identity by 0.81x-1.07x across the q35 65k ga family (valid-token + # and window differences), which bent the line by ~20% (LOO APE 20-23% -> 2-5% with the + # consistent definition). + step_time = tokens / point.tokens_per_sec if point.tokens_per_sec else point.step_time_sec + if step_time is None: + continue + x_values.append(float(tokens)) + y_values.append(float(step_time)) + if len(x_values) < 2 or len(set(x_values)) < 2: + return None + + x_mean = sum(x_values) / len(x_values) + y_mean = sum(y_values) / len(y_values) + denominator = sum((x_value - x_mean) ** 2 for x_value in x_values) + if denominator == 0: + return None + slope = sum((x_value - x_mean) * (y_value - y_mean) for x_value, y_value in zip(x_values, y_values, strict=False)) + slope /= denominator + intercept = y_mean - slope * x_mean + predicted_step = intercept + slope * shape.global_tokens_per_train_step + if predicted_step <= 0: + return None + tokens_per_sec = shape.global_tokens_per_train_step / predicted_step + tokens_per_sec_per_gpu = tokens_per_sec / topology.world_size + labels = ", ".join(point.label for point in fit_points) + model_refs = {point.model_ref for point in fit_points if point.model_ref is not None} + peak_mem_gb = max((point.peak_mem_gb for point in fit_points if point.peak_mem_gb is not None), default=None) + phase_memory_peak_gb: dict[str, float] = {} + for point in fit_points: + for phase, peak in point.phase_memory_peak_gb.items(): + phase_memory_peak_gb[phase] = max(phase_memory_peak_gb.get(phase, 0.0), peak) + return BenchmarkBehaviorPrediction( + status="extrapolated_step_time_fit", + matched_label=labels, + source="step_time_fit", + tokens_per_sec=round(tokens_per_sec, 3), + tokens_per_sec_per_gpu=round(tokens_per_sec_per_gpu, 3), + step_time_sec=round(predicted_step, 6), + mfu_percent=None, + tflops_per_gpu=None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=peak_mem_gb, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + phase_memory_peak_gb=phase_memory_peak_gb, + model_ref=next(iter(model_refs)) if len(model_refs) == 1 else None, + balanced_routing=( + fit_points[0].balanced_routing if len({point.balanced_routing for point in fit_points}) == 1 else None + ), + correctness_status="not_promoted_extrapolated", + warnings=[ + f"extrapolated step time from calibrated global batches: {labels}", + f"fit_intercept_sec={intercept:.6f}", + f"fit_sec_per_token={slope:.12f}", + "correctness must be re-gated before promotion", + ], + ) + + +# Measured non-PyTorch device overhead at long context (NCCL buffers + CUDA context + triton cache): +# the 65k 2-node gate-lane OOMs died with 9.0-12.6 GB of non-torch memory resident (e.g. o-ep4mbs2: +# 65.9 GB torch-allocated but 78.3 GB device-used). Device fit must account for it; this is a +# calibrated_residual term (max measured at 65k/2-node) until decomposed. +LONG_CONTEXT_NON_TORCH_DEVICE_OVERHEAD_GB = 12.6 +LONG_CONTEXT_NON_TORCH_SEQ_THRESHOLD = 32768 + + +def _device_side_peak_gb(estimated_peak_mem_gb: float | None, topology: Topology | None) -> float | None: + """Torch-side peak estimate -> device-side estimate for the fit check.""" + if estimated_peak_mem_gb is None: + return None + seq_len = int(getattr(topology, "sample_packing_sequence_len", 0) or 0) if topology is not None else 0 + if seq_len >= LONG_CONTEXT_NON_TORCH_SEQ_THRESHOLD: + return estimated_peak_mem_gb + LONG_CONTEXT_NON_TORCH_DEVICE_OVERHEAD_GB + return estimated_peak_mem_gb + + +def _memory_factor( + memory_estimate_gb: float | None, + *, + memory_basis: str, + device_memory_limit_gb: float, + memory_safety_factor: float, +) -> tuple[float, float | None, str]: + if memory_estimate_gb is None: + return 0.0, None, "unknown_memory_estimate" + reserved_memory = memory_estimate_gb * memory_safety_factor + headroom = device_memory_limit_gb - reserved_memory + status_basis = "floor" if memory_basis == "analytic_floor" else memory_basis + if headroom < 0: + if memory_basis != "analytic_floor" and memory_estimate_gb <= device_memory_limit_gb: + return 0.75, headroom, f"feasible_{status_basis}_high_pressure" + if memory_basis == "analytic_floor": + if memory_estimate_gb <= device_memory_limit_gb: + return 0.0, headroom, "memory_floor_exceeds_safety_margin" + return 0.0, headroom, "memory_floor_exceeds_limit" + return 0.0, headroom, f"{status_basis}_exceeds_limit" + utilization = reserved_memory / device_memory_limit_gb if device_memory_limit_gb else 1.0 + if utilization >= 0.90: + return 0.75, headroom, f"feasible_{status_basis}_high_pressure" + if utilization >= 0.80: + return 0.90, headroom, f"feasible_{status_basis}_moderate_pressure" + return 1.0, headroom, f"feasible_{status_basis}" + + +def _throughput_memory_factor( + analytic_peak_floor_gb: float | None, + *, + memory_peak_estimate: _MemoryPeakEstimate | None, + reference: BenchmarkBehaviorPoint, + topology: Topology, + device_memory_limit_gb: float, + memory_safety_factor: float, +) -> tuple[float, str, list[str]]: + notes: list[str] = [] + factor, _, status = _memory_factor( + analytic_peak_floor_gb, + memory_basis="analytic_floor", + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + if memory_peak_estimate is not None and device_memory_limit_gb > 0: + reserved_utilization = memory_peak_estimate.peak_gb * memory_safety_factor / device_memory_limit_gb + notes.append(f"throughput_memory_peak_source={memory_peak_estimate.source_label}") + notes.append(f"throughput_memory_reserved_utilization={reserved_utilization:.3f}") + expands_microbatch = ( + reference.micro_batch_size is not None and topology.micro_batch_size > reference.micro_batch_size + ) + # Graded allocator-pressure prior, calibrated by one measured point on each side: the q35 + # mbs3 prospective blind holdout (estimated utilization 0.916) measured NO pressure — the old + # binary 0.5x at >=0.85 turned a 2.5%-accurate row into a 2x miss — while the q36 mbs10 row + # (estimated 0.948) measured a real ~2x slowdown. Onset ramps 1.0 -> 0.5 across 0.92 -> 0.94 + # estimated reserved utilization; the risk flag attaches from 0.85 up so intervals widen + # before the score moves. A reference whose OWN measured peak crosses 0.85 keeps the hard 0.5x. + reference_reserved = ( + reference.peak_mem_gb * memory_safety_factor / device_memory_limit_gb if reference.peak_mem_gb else None + ) + if reference.status != "allocator_pressure_slowdown" and expands_microbatch and reserved_utilization >= 0.85: + if reference_reserved is not None and reference_reserved >= 0.85: + factor *= 0.50 + notes.append("allocator_pressure_prior=larger_microbatch_measured_reference_peak_ge_85pct_reserved") + else: + ramp = min(max((reserved_utilization - 0.92) / 0.02, 0.0), 1.0) + graded = 1.0 - 0.5 * ramp + if graded < 1.0: + factor *= graded + notes.append( + f"allocator_pressure_prior=graded_extrapolated_peak_factor_{graded:.3f}" + f"_at_{reserved_utilization:.3f}_reserved" + ) + notes.append("allocator_pressure_risk=extrapolated_peak_ge_85pct_reserved_unmeasured") + return factor, status, notes + + +def _memory_coverage_for_candidate( + *, + analytic_peak_floor_gb: float | None, + estimated_peak_mem_gb: float | None, + memory_basis: str, +) -> tuple[str, float | None, float | None]: + if analytic_peak_floor_gb is None: + return "unresolved_analytic_floor", None, None + if memory_basis == "analytic_floor": + return "analytic_floor_only", None, None + if estimated_peak_mem_gb is None: + return "unresolved_estimated_peak", None, None + if estimated_peak_mem_gb <= 0: + return "invalid_estimated_peak", None, None + + residual = estimated_peak_mem_gb - analytic_peak_floor_gb + if residual < 0: + return f"{memory_basis}_below_analytic_floor", 0.0, 0.0 + if residual == 0: + if memory_basis == "calibrated_peak": + return "analytic_floor_matches_calibrated_peak", 0.0, 0.0 + return f"{memory_basis}_matches_analytic_floor", 0.0, 0.0 + + residual_fraction = residual / estimated_peak_mem_gb + if memory_basis == "calibrated_peak": + status = "calibrated_peak_with_unmodeled_residual" + elif memory_basis == "calibrated_overhead_peak": + status = "calibrated_overhead_peak_with_scaled_residual" + elif memory_basis == "extrapolated_peak": + status = "extrapolated_peak_with_unmodeled_residual" + else: + status = f"{memory_basis}_with_unmodeled_residual" + return status, round(residual, 3), round(residual_fraction, 3) + + +def _extrapolate_behavior( + behavior_points: list[BenchmarkBehaviorPoint], + topology: Topology, + shape: ShapeLedger, + *, + raw_config: dict[str, Any] | None = None, + device_memory_limit_gb: float, + memory_safety_factor: float, + analytic_peak_floor_gb: float | None, + memory_peak_estimate: _MemoryPeakEstimate | None = None, +) -> tuple[BenchmarkBehaviorPrediction, list[str]]: + step_fit = _step_time_fit_prediction(behavior_points, topology, shape, raw_config=raw_config) + if step_fit is not None: + return step_fit, ["step_time_fit_extrapolation"] + + reference = _select_reference_point(behavior_points, topology, raw_config=raw_config) + if reference is None: + cross_model_reference = _select_cross_model_reference_point(behavior_points, topology, raw_config) + if cross_model_reference is not None and raw_config is not None: + cross_model_prediction = _cross_model_behavior_prediction( + cross_model_reference, + topology, + shape, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=analytic_peak_floor_gb, + ) + if cross_model_prediction is not None: + return cross_model_prediction + return ( + BenchmarkBehaviorPrediction( + status="unscored", + matched_label=None, + source=None, + tokens_per_sec=None, + tokens_per_sec_per_gpu=None, + step_time_sec=None, + mfu_percent=None, + tflops_per_gpu=None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=None, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + balanced_routing=_config_balanced_routing(raw_config) if raw_config is not None else None, + correctness_status="missing_calibration", + warnings=["no benchmark behavior point is available for extrapolation"], + ), + [], + ) + + ref_per_gpu = _reference_tokens_per_gpu(reference, topology) or 0.0 + parallel_factor, notes = _parallelism_factor(reference, topology) + memory_factor, memory_status, memory_notes = _throughput_memory_factor( + analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + reference=reference, + topology=topology, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + notes.extend(memory_notes) + tokens_per_sec_per_gpu = ref_per_gpu * parallel_factor * memory_factor + tokens_per_sec = tokens_per_sec_per_gpu * topology.world_size + step_time_sec = None + if shape.global_tokens_per_train_step and tokens_per_sec: + step_time_sec = shape.global_tokens_per_train_step / tokens_per_sec + # ga fixed-tail amortization: a pure-ga extrapolation reuses the reference tok/s verbatim, which + # drops the per-step fixed cost (clip + optimizer + metrics) amortizing across microbatch chains. + # step(ga) = ga_ratio x (ref_step - fixed_tail) + fixed_tail from the reference's own phase reads + # predicted the ga2@mbs2 prospective blind holdout to 0.02% (4.3454 vs measured 4.3455). + ga_amortization_applicable = ( + step_time_sec is not None + and parallel_factor == 1.0 + and reference.micro_batch_size == topology.micro_batch_size + and reference.gradient_accumulation_steps + and topology.gradient_accumulation_steps != reference.gradient_accumulation_steps + and reference.step_time_sec + and reference.phase_time_sec + and "clip_and_step_total" in reference.phase_time_sec + ) + if ga_amortization_applicable: + target_ga = topology.gradient_accumulation_steps + ga_ratio = target_ga / reference.gradient_accumulation_steps + fixed_tail = float(reference.phase_time_sec["clip_and_step_total"]) + float( + reference.phase_time_sec.get("reduce_metrics", 0.0) + ) + chain_time = max(float(reference.step_time_sec) - fixed_tail, 0.0) + amortized_step = ga_ratio * chain_time + fixed_tail + if amortized_step > 0: + step_time_sec = amortized_step * (step_time_sec / (float(reference.step_time_sec) * ga_ratio)) + tokens_per_sec = shape.global_tokens_per_train_step / step_time_sec + tokens_per_sec_per_gpu = tokens_per_sec / topology.world_size + notes.append( + f"ga_fixed_tail_amortization: ga {reference.gradient_accumulation_steps}->{target_ga}, " + f"fixed_tail={fixed_tail:.4f}s from reference phase reads" + ) + tflops_per_gpu = reference.tflops_per_gpu + if tflops_per_gpu is None and reference.mfu_percent is not None: + tflops_per_gpu = H100_BF16_PROMISED_TFLOPS_PER_GPU * reference.mfu_percent / 100.0 + if tflops_per_gpu is not None and ref_per_gpu: + tflops_per_gpu *= tokens_per_sec_per_gpu / ref_per_gpu + + warnings = [ + f"extrapolated from {reference.label}; correctness must be re-gated before promotion", + f"memory feasibility status is {memory_status}", + ] + if raw_config is not None: + mismatches = behavior_point_workload_mismatches(reference, raw_config) + if mismatches: + warnings.append(f"reference runtime knobs differ: {', '.join(mismatches)}") + warnings.extend(notes) + return ( + BenchmarkBehaviorPrediction( + status="extrapolated", + matched_label=reference.label, + source=reference.source, + tokens_per_sec=round(tokens_per_sec, 3), + tokens_per_sec_per_gpu=round(tokens_per_sec_per_gpu, 3), + step_time_sec=round(step_time_sec, 6) if step_time_sec is not None else None, + mfu_percent=None, + tflops_per_gpu=round(tflops_per_gpu, 3) if tflops_per_gpu is not None else None, + promised_tflops_per_gpu=H100_BF16_PROMISED_TFLOPS_PER_GPU, + peak_mem_gb=reference.peak_mem_gb, + allocator_retries=None, + derived_global_tokens_per_step=shape.global_tokens_per_train_step, + phase_time_sec=reference.phase_time_sec, + phase_time_share=reference.phase_time_share, + phase_memory_peak_gb=reference.phase_memory_peak_gb, + model_ref=reference.model_ref, + balanced_routing=reference.balanced_routing, + correctness_status="not_promoted_extrapolated", + warnings=warnings, + ), + notes, + ) + + +def _format_factor_float(value: float | None, suffix: str = "") -> str: + if value is None: + return "unknown" + return f"{value:.3f}{suffix}" + + +def _candidate_decision_factors( + *, + behavior: BenchmarkBehaviorPrediction, + prediction_confidence: str, + calibration_scope: str, + calibration_distance: float | None, + calibration_distance_factors: list[str], + feasibility_status: str, + score_tokens_per_sec: float | None, + score_tokens_per_gpu_per_sec: float | None, + score_risk_adjusted_tokens_per_sec: float | None, + score_risk_adjusted_tokens_per_gpu_per_sec: float | None, + prediction_uncertainty_fraction: float | None, + prediction_interval_lower_tokens_per_sec: float | None, + prediction_interval_upper_tokens_per_sec: float | None, + risk_adjusted_prediction_interval_lower_tokens_per_sec: float | None, + risk_adjusted_prediction_interval_upper_tokens_per_sec: float | None, + analytic_peak_floor_gb: float | None, + estimated_peak_mem_gb: float | None, + memory_basis: str, + memory_coverage_status: str, + memory_residual_gb: float | None, + memory_residual_fraction: float | None, + headroom_gb: float | None, + promotable: bool, + recommendation: str, + simulator_surface: str, + simulator_support_status: str, + simulator_support_blockers: list[str], + risk_flags: list[str], + communication: CommLedger | None, +) -> list[str]: + factors = [ + f"matched={behavior.matched_label or behavior.status}", + f"calibration={calibration_scope}/{prediction_confidence}", + f"feasibility={feasibility_status}", + f"score_tokens_per_sec={_format_factor_float(score_tokens_per_sec)}", + f"score_tokens_per_gpu_per_sec={_format_factor_float(score_tokens_per_gpu_per_sec)}", + ] + if calibration_distance is not None: + factors.append(f"calibration_distance={_format_factor_float(calibration_distance)}") + if calibration_distance_factors: + factors.append(f"calibration_distance_factors={';'.join(calibration_distance_factors)}") + if score_risk_adjusted_tokens_per_sec is not None: + factors.append(f"risk_adjusted_tokens_per_sec={_format_factor_float(score_risk_adjusted_tokens_per_sec)}") + if score_risk_adjusted_tokens_per_gpu_per_sec is not None: + factors.append( + f"risk_adjusted_tokens_per_gpu_per_sec={_format_factor_float(score_risk_adjusted_tokens_per_gpu_per_sec)}" + ) + if prediction_uncertainty_fraction is not None: + factors.append(f"prediction_uncertainty_fraction={_format_factor_float(prediction_uncertainty_fraction)}") + if prediction_interval_lower_tokens_per_sec is not None and prediction_interval_upper_tokens_per_sec is not None: + factors.append( + "prediction_interval_tokens_per_sec=" + f"{_format_factor_float(prediction_interval_lower_tokens_per_sec)}.." + f"{_format_factor_float(prediction_interval_upper_tokens_per_sec)}" + ) + if ( + risk_adjusted_prediction_interval_lower_tokens_per_sec is not None + and risk_adjusted_prediction_interval_upper_tokens_per_sec is not None + ): + factors.append( + "risk_adjusted_prediction_interval_tokens_per_sec=" + f"{_format_factor_float(risk_adjusted_prediction_interval_lower_tokens_per_sec)}.." + f"{_format_factor_float(risk_adjusted_prediction_interval_upper_tokens_per_sec)}" + ) + factors.append( + "memory=" + f"{memory_basis}:floor={_format_factor_float(analytic_peak_floor_gb, 'GB')}," + f"peak={_format_factor_float(estimated_peak_mem_gb, 'GB')}," + f"headroom={_format_factor_float(headroom_gb, 'GB')}" + ) + factors.append( + "memory_coverage=" + f"{memory_coverage_status}:" + f"residual={_format_factor_float(memory_residual_gb, 'GB')}," + f"residual_fraction={_format_factor_float(memory_residual_fraction)}" + ) + if behavior.correctness_status: + factors.append(f"correctness={behavior.correctness_status}") + factors.append(f"promotable={str(promotable).lower()}") + if communication is not None and communication.cross_node_dimensions: + factors.append(f"cross_node={','.join(communication.cross_node_dimensions)}") + factors.append(f"simulator_support={simulator_surface}/{simulator_support_status}") + if simulator_support_blockers: + factors.append(f"simulator_support_blockers={','.join(simulator_support_blockers)}") + if risk_flags: + factors.append(f"risks={','.join(risk_flags)}") + factors.append(f"recommendation={recommendation}") + return factors + + +def _candidate_from_prediction( + *, + label: str, + config_path: str | None, + topology: Topology, + shape: ShapeLedger, + behavior: BenchmarkBehaviorPrediction, + prediction_confidence: str, + promotable: bool, + behavior_points: list[BenchmarkBehaviorPoint], + raw_config: dict[str, Any] | None, + device_memory_limit_gb: float, + memory_safety_factor: float, + analytic_peak_floor_gb: float | None, + memory_peak_estimate: _MemoryPeakEstimate | None, + memory_ownership_notes: list[str], + communication: CommLedger, + notes: list[str], +) -> ScenarioCandidate: + estimated_peak_mem_gb = analytic_peak_floor_gb + memory_basis = "analytic_floor" + peak_below_floor_note = None + memory_calibration_source = None + memory_calibration_notes: list[str] = [] + if behavior.peak_mem_gb is not None: + if prediction_confidence == "calibrated" and ( + analytic_peak_floor_gb is None or behavior.peak_mem_gb >= analytic_peak_floor_gb + ): + estimated_peak_mem_gb = behavior.peak_mem_gb + memory_basis = "calibrated_peak" + elif prediction_confidence != "calibrated" and memory_peak_estimate is not None: + estimated_peak_mem_gb = memory_peak_estimate.peak_gb + memory_basis = "calibrated_overhead_peak" + memory_calibration_source = memory_peak_estimate.source_label + memory_calibration_notes = memory_peak_estimate.notes + elif analytic_peak_floor_gb is None or behavior.peak_mem_gb >= analytic_peak_floor_gb: + estimated_peak_mem_gb = behavior.peak_mem_gb + memory_basis = "extrapolated_peak" + else: + peak_below_floor_note = ( + f"calibrated_peak_below_analytic_floor: peak={behavior.peak_mem_gb:.3f} " + f"floor={analytic_peak_floor_gb:.3f}; using analytic_floor" + ) + elif memory_peak_estimate is not None: + estimated_peak_mem_gb = memory_peak_estimate.peak_gb + memory_basis = "calibrated_overhead_peak" + memory_calibration_source = memory_peak_estimate.source_label + memory_calibration_notes = memory_peak_estimate.notes + + # Fit is a DEVICE-memory question: an ESTIMATED torch-side peak plus the measured long-context + # non-torch overhead (NCCL/context/triton) is what actually has to fit under the device limit. + # A MEASURED peak (the candidate's own completed run) is direct device-fit evidence and must not + # be re-blocked by an overhead estimate. + fit_check_peak_gb = ( + _device_side_peak_gb(estimated_peak_mem_gb, topology) + if memory_basis in ("calibrated_overhead_peak", "extrapolated_peak") + else estimated_peak_mem_gb + ) + _, headroom, feasibility_status = _memory_factor( + fit_check_peak_gb, + memory_basis=memory_basis, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ) + ( + memory_coverage_status, + estimated_memory_residual_gb, + estimated_memory_residual_fraction, + ) = _memory_coverage_for_candidate( + analytic_peak_floor_gb=analytic_peak_floor_gb, + estimated_peak_mem_gb=estimated_peak_mem_gb, + memory_basis=memory_basis, + ) + support = resolve_simulator_support(raw_config or {}, topology=topology) + if behavior.status == "calibrated_failure" or behavior.correctness_status == "oom": + feasibility_status = "observed_oom" + if support.support_status.startswith("unsupported_"): + feasibility_status = "unsupported_simulator_surface" + if behavior.tokens_per_sec is None and feasibility_status.startswith("feasible"): + feasibility_status = "unscored" + feasible = feasibility_status.startswith("feasible") and behavior.tokens_per_sec is not None + score_tokens_per_sec = behavior.tokens_per_sec if feasible else None + score_tokens_per_gpu_per_sec = behavior.tokens_per_sec_per_gpu if feasible else None + max_ep_slots = max(shape.ep_rank_slots_per_microbatch) if shape.ep_rank_slots_per_microbatch else None + calibration_scope = _calibration_scope( + behavior_points, + topology, + prediction_confidence=prediction_confidence, + raw_config=raw_config, + ) + risk_flags = _candidate_risk_flags( + behavior_points, + topology, + behavior, + raw_config=raw_config, + calibration_scope=calibration_scope, + prediction_confidence=prediction_confidence, + communication=communication, + ) + if peak_below_floor_note is not None: + risk_flags = sorted({*risk_flags, "calibrated_peak_below_analytic_floor"}) + if memory_basis == "calibrated_overhead_peak": + risk_flags = sorted({*risk_flags, "memory_extrapolated_overhead"}) + if any("allocator_pressure_risk=extrapolated_peak_ge_85pct_reserved_unmeasured" in w for w in behavior.warnings): + risk_flags = sorted({*risk_flags, "allocator_pressure_risk_extrapolated_peak"}) + if support.support_status.startswith("unsupported_"): + risk_flags = sorted({*risk_flags, f"simulator_surface_unsupported:{support.support_status}"}) + elif support.support_status != "supported_local_non_pp": + risk_flags = sorted({*risk_flags, f"simulator_surface_partial:{support.support_status}"}) + calibration_distance, calibration_distance_factors = _calibration_distance(behavior_points, topology, behavior) + score_risk_adjusted_tokens_per_sec = _risk_adjusted_score( + score_tokens_per_sec, + calibration_scope=calibration_scope, + calibration_distance=calibration_distance, + risk_flags=risk_flags, + feasibility_status=feasibility_status, + ) + score_risk_adjusted_tokens_per_gpu_per_sec = ( + round(score_risk_adjusted_tokens_per_sec / topology.world_size, 3) + if score_risk_adjusted_tokens_per_sec is not None and topology.world_size + else None + ) + prediction_uncertainty_fraction = _prediction_uncertainty_fraction( + behavior, + prediction_confidence=prediction_confidence, + calibration_scope=calibration_scope, + calibration_distance=calibration_distance, + risk_flags=risk_flags, + memory_coverage_status=memory_coverage_status, + ) + prediction_interval_lower, prediction_interval_upper = _prediction_interval( + score_tokens_per_sec, + prediction_uncertainty_fraction, + ) + risk_adjusted_interval_lower, risk_adjusted_interval_upper = _prediction_interval( + score_risk_adjusted_tokens_per_sec, + prediction_uncertainty_fraction, + ) + recommendation = _recommendation( + feasible=feasible, + promotable=promotable and feasible, + feasibility_status=feasibility_status, + risk_flags=risk_flags, + ) + candidate_notes = list(notes) + if peak_below_floor_note is not None: + candidate_notes.append(peak_below_floor_note) + candidate_notes.extend(memory_calibration_notes) + timing = build_timing_ledger(None, behavior) + phase_details = _phase_bottleneck_details(timing.phase_time_share, timing.phase_time_sec) + phase_bottleneck_share = phase_details[2] if phase_details is not None else None + ( + phase_bottleneck_half_speedup_tokens_per_sec, + phase_bottleneck_half_speedup_delta_pct, + ) = _phase_bottleneck_half_speedup_counterfactual( + score_tokens_per_sec, + phase_bottleneck_share, + ) + ( + phase_bottleneck_half_speedup_risk_adjusted_tokens_per_sec, + phase_bottleneck_half_speedup_risk_adjusted_delta_pct, + ) = _phase_bottleneck_half_speedup_counterfactual( + score_risk_adjusted_tokens_per_sec, + phase_bottleneck_share, + ) + memory_details = _memory_bottleneck_details(behavior.phase_memory_peak_gb, estimated_peak_mem_gb) + if phase_note := _phase_bottleneck_note(timing.phase_time_share): + candidate_notes.append(phase_note) + if phase_bottleneck_half_speedup_delta_pct is not None and phase_details is not None: + candidate_notes.append( + "fixed_schedule_phase_bottleneck_half_speedup=" + f"{phase_details[0]}:+{phase_bottleneck_half_speedup_delta_pct:.3f}%" + ) + if memory_details is not None: + candidate_notes.append( + f"memory_bottleneck={memory_details[0]}:{memory_details[2]:.3f}GB({memory_details[3]:.3f}x_peak)" + ) + timing_coverage_status = _scenario_timing_coverage_status( + behavior, + prediction_confidence=prediction_confidence, + calibration_scope=calibration_scope, + timing_coverage_status=timing.timing_coverage_status, + ) + decision_factors = _candidate_decision_factors( + behavior=behavior, + prediction_confidence=prediction_confidence, + calibration_scope=calibration_scope, + calibration_distance=calibration_distance, + calibration_distance_factors=calibration_distance_factors, + feasibility_status=feasibility_status, + score_tokens_per_sec=score_tokens_per_sec, + score_tokens_per_gpu_per_sec=score_tokens_per_gpu_per_sec, + score_risk_adjusted_tokens_per_sec=score_risk_adjusted_tokens_per_sec, + score_risk_adjusted_tokens_per_gpu_per_sec=score_risk_adjusted_tokens_per_gpu_per_sec, + prediction_uncertainty_fraction=prediction_uncertainty_fraction, + prediction_interval_lower_tokens_per_sec=prediction_interval_lower, + prediction_interval_upper_tokens_per_sec=prediction_interval_upper, + risk_adjusted_prediction_interval_lower_tokens_per_sec=risk_adjusted_interval_lower, + risk_adjusted_prediction_interval_upper_tokens_per_sec=risk_adjusted_interval_upper, + analytic_peak_floor_gb=analytic_peak_floor_gb, + estimated_peak_mem_gb=estimated_peak_mem_gb, + memory_basis=memory_basis, + memory_coverage_status=memory_coverage_status, + memory_residual_gb=estimated_memory_residual_gb, + memory_residual_fraction=estimated_memory_residual_fraction, + headroom_gb=headroom, + promotable=promotable and feasible, + recommendation=recommendation, + simulator_surface=support.requested_surface, + simulator_support_status=support.support_status, + simulator_support_blockers=support.support_blockers, + risk_flags=risk_flags, + communication=communication, + ) + decision_factors.append(f"timing_coverage={timing_coverage_status}") + if timing.step_time_s is not None: + decision_factors.append(f"timing_step_time_s={timing.step_time_s:.6f}") + if timing.forward_backward_s is not None: + decision_factors.append(f"timing_forward_backward_s={timing.forward_backward_s:.6f}") + if phase_bottleneck_half_speedup_delta_pct is not None: + decision_factors.append( + f"phase_bottleneck_half_speedup_delta_pct={phase_bottleneck_half_speedup_delta_pct:.3f}" + ) + if phase_bottleneck_half_speedup_tokens_per_sec is not None: + decision_factors.append( + f"phase_bottleneck_half_speedup_tokens_per_sec={phase_bottleneck_half_speedup_tokens_per_sec:.3f}" + ) + if phase_bottleneck_half_speedup_risk_adjusted_delta_pct is not None: + decision_factors.append( + "phase_bottleneck_half_speedup_risk_adjusted_delta_pct=" + f"{phase_bottleneck_half_speedup_risk_adjusted_delta_pct:.3f}" + ) + if memory_details is not None: + decision_factors.append( + "memory_bottleneck=" + f"{memory_details[0]}:{memory_details[2]:.3f}GB," + f"bucket={memory_details[1]},fraction_of_peak={memory_details[3]:.3f}" + ) + return ScenarioCandidate( + label=label, + config_path=config_path, + topology=topology, + behavior=behavior, + prediction_confidence=prediction_confidence, + promotable=promotable and feasible, + feasibility_status=feasibility_status, + score_tokens_per_sec=score_tokens_per_sec, + score_tokens_per_gpu_per_sec=score_tokens_per_gpu_per_sec, + score_risk_adjusted_tokens_per_sec=score_risk_adjusted_tokens_per_sec, + score_risk_adjusted_tokens_per_gpu_per_sec=score_risk_adjusted_tokens_per_gpu_per_sec, + prediction_uncertainty_fraction=prediction_uncertainty_fraction, + prediction_interval_lower_tokens_per_sec=prediction_interval_lower, + prediction_interval_upper_tokens_per_sec=prediction_interval_upper, + risk_adjusted_prediction_interval_lower_tokens_per_sec=risk_adjusted_interval_lower, + risk_adjusted_prediction_interval_upper_tokens_per_sec=risk_adjusted_interval_upper, + analytic_peak_floor_gb=analytic_peak_floor_gb, + estimated_peak_mem_gb=estimated_peak_mem_gb, + memory_basis=memory_basis, + memory_coverage_status=memory_coverage_status, + memory_headroom_gb=round(headroom, 3) if headroom is not None else None, + estimated_memory_residual_gb=estimated_memory_residual_gb, + estimated_memory_residual_fraction_of_peak=estimated_memory_residual_fraction, + max_ep_rank_slots_per_microbatch=max_ep_slots, + phase_bottleneck_phase=phase_details[0] if phase_details is not None else None, + phase_bottleneck_bucket=phase_details[1] if phase_details is not None else None, + phase_bottleneck_share=phase_details[2] if phase_details is not None else None, + phase_bottleneck_time_sec=phase_details[3] if phase_details is not None else None, + memory_bottleneck_phase=memory_details[0] if memory_details is not None else None, + memory_bottleneck_bucket=memory_details[1] if memory_details is not None else None, + memory_bottleneck_peak_gb=memory_details[2] if memory_details is not None else None, + memory_bottleneck_fraction_of_peak=memory_details[3] if memory_details is not None else None, + timing_coverage_status=timing_coverage_status, + timing_source_label=timing.source, + timing_step_time_s=timing.step_time_s, + timing_forward_backward_s=timing.forward_backward_s, + calibration_scope=calibration_scope, + recommendation=recommendation, + phase_bottleneck_half_speedup_scale=( + _PHASE_BOTTLENECK_HALFSPEED_SCALE if phase_bottleneck_half_speedup_delta_pct is not None else None + ), + phase_bottleneck_half_speedup_tokens_per_sec=phase_bottleneck_half_speedup_tokens_per_sec, + phase_bottleneck_half_speedup_delta_pct=phase_bottleneck_half_speedup_delta_pct, + phase_bottleneck_half_speedup_risk_adjusted_tokens_per_sec=( + phase_bottleneck_half_speedup_risk_adjusted_tokens_per_sec + ), + phase_bottleneck_half_speedup_risk_adjusted_delta_pct=(phase_bottleneck_half_speedup_risk_adjusted_delta_pct), + simulator_surface=support.requested_surface, + simulator_support_status=support.support_status, + simulator_support_blockers=support.support_blockers, + target_runtime_signature=_target_runtime_signature(raw_config), + calibration_distance=calibration_distance, + calibration_distance_factors=calibration_distance_factors, + memory_calibration_source=memory_calibration_source, + memory_calibration_notes=memory_calibration_notes, + memory_ownership_notes=memory_ownership_notes, + communication=communication, + decision_factors=decision_factors, + risk_flags=risk_flags, + notes=candidate_notes, + ) + + +def _candidate_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: + return ( + candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), + candidate.score_tokens_per_gpu_per_sec if candidate.score_tokens_per_gpu_per_sec is not None else float("-inf"), + ) + + +def _risk_adjusted_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: + return ( + candidate.score_risk_adjusted_tokens_per_sec + if candidate.score_risk_adjusted_tokens_per_sec is not None + else float("-inf"), + candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), + ) + + +def _efficiency_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: + return ( + candidate.score_tokens_per_gpu_per_sec if candidate.score_tokens_per_gpu_per_sec is not None else float("-inf"), + candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf"), + ) + + +def _risk_adjusted_efficiency_sort_key(candidate: ScenarioCandidate) -> tuple[float, float]: + return ( + candidate.score_risk_adjusted_tokens_per_gpu_per_sec + if candidate.score_risk_adjusted_tokens_per_gpu_per_sec is not None + else float("-inf"), + candidate.score_risk_adjusted_tokens_per_sec + if candidate.score_risk_adjusted_tokens_per_sec is not None + else float("-inf"), + ) + + +def _throughput_efficiency_frontier_labels( + candidates: list[ScenarioCandidate], + *, + throughput_attr: str, + efficiency_attr: str, +) -> list[str]: + scored: list[tuple[ScenarioCandidate, float, float]] = [] + for candidate in candidates: + throughput = getattr(candidate, throughput_attr) + efficiency = getattr(candidate, efficiency_attr) + if throughput is not None and efficiency is not None: + scored.append((candidate, throughput, efficiency)) + + frontier: list[tuple[ScenarioCandidate, float, float]] = [] + for candidate, throughput, efficiency in scored: + dominated = any( + other is not candidate + and other_throughput >= throughput + and other_efficiency >= efficiency + and (other_throughput > throughput or other_efficiency > efficiency) + for other, other_throughput, other_efficiency in scored + ) + if not dominated: + frontier.append((candidate, throughput, efficiency)) + + return [ + candidate.label + for candidate, _, _ in sorted( + frontier, + key=lambda item: (item[1], item[2], item[0].label), + reverse=True, + ) + ] + + +def _candidate_dominator( + candidate: ScenarioCandidate, + scored: list[tuple[ScenarioCandidate, float, float]], + *, + throughput: float, + efficiency: float, +) -> tuple[ScenarioCandidate, float, float] | None: + dominators = [ + (other, other_throughput, other_efficiency) + for other, other_throughput, other_efficiency in scored + if other is not candidate + and other_throughput >= throughput + and other_efficiency >= efficiency + and (other_throughput > throughput or other_efficiency > efficiency) + ] + if not dominators: + return None + return max( + dominators, + key=lambda item: ( + item[1] - throughput, + item[2] - efficiency, + item[1], + item[2], + item[0].label, + ), + ) + + +def _dominance_updates( + candidates: list[ScenarioCandidate], + *, + throughput_attr: str, + efficiency_attr: str, + frontier_member_field: str, + dominated_by_field: str, + throughput_margin_field: str, + efficiency_margin_field: str, + decision_factor_prefix: str, +) -> dict[str, ScenarioCandidate]: + scored: list[tuple[ScenarioCandidate, float, float]] = [] + for candidate in candidates: + throughput = getattr(candidate, throughput_attr) + efficiency = getattr(candidate, efficiency_attr) + if throughput is not None and efficiency is not None: + scored.append((candidate, throughput, efficiency)) + + updates: dict[str, ScenarioCandidate] = {} + for candidate, throughput, efficiency in scored: + dominator = _candidate_dominator(candidate, scored, throughput=throughput, efficiency=efficiency) + if dominator is None: + updates[candidate.label] = replace(candidate, **{frontier_member_field: True}) + continue + + dominator_candidate, dominator_throughput, dominator_efficiency = dominator + throughput_margin = round(dominator_throughput - throughput, 3) + efficiency_margin = round(dominator_efficiency - efficiency, 3) + updates[candidate.label] = replace( + candidate, + **{ + frontier_member_field: False, + dominated_by_field: dominator_candidate.label, + throughput_margin_field: throughput_margin, + efficiency_margin_field: efficiency_margin, + "decision_factors": [ + *candidate.decision_factors, + f"{decision_factor_prefix}_dominated_by={dominator_candidate.label}", + f"{decision_factor_prefix}_dominance_margin_tokens_per_sec={throughput_margin:.3f}", + f"{decision_factor_prefix}_dominance_margin_tokens_per_gpu_per_sec={efficiency_margin:.3f}", + ], + "notes": [ + *candidate.notes, + f"{decision_factor_prefix}_frontier_dominated_by={dominator_candidate.label}", + ], + }, + ) + return updates + + +def _apply_frontier_dominance(candidates: list[ScenarioCandidate]) -> list[ScenarioCandidate]: + raw_updates = _dominance_updates( + candidates, + throughput_attr="score_tokens_per_sec", + efficiency_attr="score_tokens_per_gpu_per_sec", + frontier_member_field="raw_frontier_member", + dominated_by_field="raw_dominated_by_label", + throughput_margin_field="raw_dominance_margin_tokens_per_sec", + efficiency_margin_field="raw_dominance_margin_tokens_per_gpu_per_sec", + decision_factor_prefix="raw", + ) + raw_candidates = [raw_updates.get(candidate.label, candidate) for candidate in candidates] + risk_updates = _dominance_updates( + raw_candidates, + throughput_attr="score_risk_adjusted_tokens_per_sec", + efficiency_attr="score_risk_adjusted_tokens_per_gpu_per_sec", + frontier_member_field="risk_adjusted_frontier_member", + dominated_by_field="risk_adjusted_dominated_by_label", + throughput_margin_field="risk_adjusted_dominance_margin_tokens_per_sec", + efficiency_margin_field="risk_adjusted_dominance_margin_tokens_per_gpu_per_sec", + decision_factor_prefix="risk_adjusted", + ) + return [risk_updates.get(candidate.label, candidate) for candidate in raw_candidates] + + +def _scaling_workload_signature(candidate: ScenarioCandidate) -> tuple[Any, ...]: + return tuple(getattr(candidate.topology, dimension) for dimension in _SCENARIO_WORKLOAD_DIMENSIONS) + + +def _apply_same_workload_scaling_metrics(candidates: list[ScenarioCandidate]) -> list[ScenarioCandidate]: + groups: dict[tuple[Any, ...], list[ScenarioCandidate]] = {} + for candidate in candidates: + if candidate.score_tokens_per_sec is None: + continue + groups.setdefault(_scaling_workload_signature(candidate), []).append(candidate) + + scaling_updates: dict[str, ScenarioCandidate] = {} + for group in groups.values(): + if len({candidate.topology.world_size for candidate in group}) < 2: + continue + baseline = max( + ( + candidate + for candidate in group + if candidate.topology.world_size == min(c.topology.world_size for c in group) + ), + key=_candidate_sort_key, + ) + if not baseline.topology.world_size or not baseline.score_tokens_per_sec: + continue + + for candidate in group: + if candidate.score_tokens_per_sec is None or not candidate.topology.world_size: + continue + gpu_ratio = candidate.topology.world_size / baseline.topology.world_size + if gpu_ratio <= 0: + continue + speedup = candidate.score_tokens_per_sec / baseline.score_tokens_per_sec + scaling_efficiency = speedup / gpu_ratio + risk_adjusted_speedup = None + risk_adjusted_scaling_efficiency = None + if ( + candidate.score_risk_adjusted_tokens_per_sec is not None + and baseline.score_risk_adjusted_tokens_per_sec is not None + and baseline.score_risk_adjusted_tokens_per_sec > 0 + ): + risk_adjusted_speedup = ( + candidate.score_risk_adjusted_tokens_per_sec / baseline.score_risk_adjusted_tokens_per_sec + ) + risk_adjusted_scaling_efficiency = risk_adjusted_speedup / gpu_ratio + + decision_factors = [ + *candidate.decision_factors, + f"scaling_baseline={baseline.label}", + f"scaling_gpu_ratio={gpu_ratio:.3f}", + f"scaling_speedup={speedup:.3f}", + f"scaling_efficiency={scaling_efficiency:.3f}", + ] + if risk_adjusted_speedup is not None and risk_adjusted_scaling_efficiency is not None: + decision_factors.extend( + [ + f"risk_adjusted_scaling_speedup={risk_adjusted_speedup:.3f}", + f"risk_adjusted_scaling_efficiency={risk_adjusted_scaling_efficiency:.3f}", + ] + ) + scaling_updates[candidate.label] = replace( + candidate, + scaling_baseline_label=baseline.label, + scaling_baseline_world_size=baseline.topology.world_size, + scaling_gpu_ratio=round(gpu_ratio, 3), + scaling_speedup=round(speedup, 3), + scaling_efficiency=round(scaling_efficiency, 3), + risk_adjusted_scaling_speedup=( + round(risk_adjusted_speedup, 3) if risk_adjusted_speedup is not None else None + ), + risk_adjusted_scaling_efficiency=( + round(risk_adjusted_scaling_efficiency, 3) if risk_adjusted_scaling_efficiency is not None else None + ), + decision_factors=decision_factors, + notes=[ + *candidate.notes, + f"same_workload_scaling_baseline={baseline.label}", + ], + ) + + return [scaling_updates.get(candidate.label, candidate) for candidate in candidates] + + +def _candidate_dimension_value(candidate: ScenarioCandidate, dimension: str) -> Any: + return getattr(candidate.topology, dimension) + + +def _varied_candidate_dimensions(candidates: list[ScenarioCandidate], dimensions: tuple[str, ...]) -> list[str]: + varied: list[str] = [] + for dimension in dimensions: + values = {_candidate_dimension_value(candidate, dimension) for candidate in candidates} + if len(values) > 1: + varied.append(dimension) + return varied + + +def _parallelism_strategy_key(candidate: ScenarioCandidate) -> tuple[Any, ...]: + return tuple(_candidate_dimension_value(candidate, dimension) for dimension in _SCENARIO_PARALLELISM_DIMENSIONS) + + +def _parallelism_strategy_counts(candidates: list[ScenarioCandidate]) -> tuple[int, int, int, int]: + unique_strategies = {_parallelism_strategy_key(candidate) for candidate in candidates} + scored_strategies = { + _parallelism_strategy_key(candidate) for candidate in candidates if candidate.score_tokens_per_sec is not None + } + promotable_strategies = { + _parallelism_strategy_key(candidate) + for candidate in candidates + if candidate.promotable and candidate.score_tokens_per_sec is not None + } + remeasurement_strategies = { + _parallelism_strategy_key(candidate) + for candidate in candidates + if candidate.score_tokens_per_sec is not None and "requires_remeasurement" in candidate.risk_flags + } + return ( + len(unique_strategies), + len(scored_strategies), + len(promotable_strategies), + len(remeasurement_strategies), + ) + + +def _signature_for_dimensions(candidate: ScenarioCandidate, dimensions: tuple[str, ...]) -> tuple[tuple[str, Any], ...]: + return tuple((dimension, getattr(candidate.topology, dimension)) for dimension in dimensions) + + +def _format_signature(signature: tuple[tuple[str, Any], ...]) -> str: + return ",".join(f"{dimension}={value}" for dimension, value in signature) + + +def _candidate_runtime_signature_values(candidate: ScenarioCandidate) -> dict[str, str]: + if candidate.target_runtime_signature == "unknown": + return {} + values: dict[str, str] = {} + for part in candidate.target_runtime_signature.split(","): + if "=" not in part: + continue + key, value = part.split("=", 1) + values[key] = value + return values + + +def _candidate_runtime_dimension_value(candidate: ScenarioCandidate, dimension: str) -> str: + return _candidate_runtime_signature_values(candidate).get(dimension, "unknown") + + +def _varied_candidate_runtime_dimensions(candidates: list[ScenarioCandidate]) -> list[str]: + varied: list[str] = [] + for dimension in _SCENARIO_RUNTIME_DIMENSIONS: + values = {_candidate_runtime_dimension_value(candidate, dimension) for candidate in candidates} + if len(values) > 1: + varied.append(dimension) + return varied + + +def _runtime_dimension_sort_key(dimension: str) -> tuple[int, str]: + try: + return _SCENARIO_RUNTIME_DIMENSIONS.index(dimension), dimension + except ValueError: + return len(_SCENARIO_RUNTIME_DIMENSIONS), dimension + + +def _candidate_runtime_mismatch_dimensions(candidates: list[ScenarioCandidate]) -> list[str]: + dimensions = { + flag.split(":", 1)[1] + for candidate in candidates + for flag in candidate.risk_flags + if flag.startswith("runtime_mismatch:") and flag.split(":", 1)[1] != "model_ref" + } + return sorted(dimensions, key=_runtime_dimension_sort_key) + + +def _runtime_signature_for_candidate(candidate: ScenarioCandidate) -> tuple[tuple[str, Any], ...]: + return (("target_runtime_signature", candidate.target_runtime_signature),) + + +def _benchmark_point_parallel_size(point: BenchmarkBehaviorPoint, field_name: str) -> int: + value = getattr(point, field_name) + return int(value) if value is not None else 1 + + +def _benchmark_point_topology_values( + point: BenchmarkBehaviorPoint, + base_topology: Topology, +) -> dict[str, Any]: + world_size = point.gpu_count or base_topology.world_size + local_world_size = base_topology.local_world_size + tensor_parallel_size = _benchmark_point_parallel_size(point, "tensor_parallel_size") + pipeline_parallel_size = _benchmark_point_parallel_size(point, "pipeline_parallel_size") + expert_parallel_size = _benchmark_point_parallel_size(point, "expert_parallel_size") + ulysses_parallel_size = _benchmark_point_parallel_size(point, "ulysses_parallel_size") + ringattn_parallel_size = _benchmark_point_parallel_size(point, "ringattn_parallel_size") + non_dp_size = tensor_parallel_size * pipeline_parallel_size * ulysses_parallel_size * ringattn_parallel_size + data_parallel_size = world_size // non_dp_size if non_dp_size and world_size % non_dp_size == 0 else None + if data_parallel_size is None: + replicate_size = None + shard_size = None + else: + if point.data_parallel_replicate_size is not None or point.data_parallel_shard_size is not None: + preferred_replicate_size = point.data_parallel_replicate_size + preferred_shard_size = point.data_parallel_shard_size + else: + preferred_replicate_size = base_topology.data_parallel_replicate_size + preferred_shard_size = base_topology.data_parallel_shard_size + replicate_size, shard_size = _valid_dp_split_for_size( + data_parallel_size, + preferred_replicate_size=preferred_replicate_size, + preferred_shard_size=preferred_shard_size, + ) + ep_fsdp_size = point.ep_fsdp_size + if ep_fsdp_size is None: + ranks_per_pipeline_stage = ( + world_size // pipeline_parallel_size if world_size % pipeline_parallel_size == 0 else None + ) + if ranks_per_pipeline_stage is not None and ranks_per_pipeline_stage % expert_parallel_size == 0: + ep_fsdp_size = ranks_per_pipeline_stage // expert_parallel_size + + return { + "world_size": world_size, + "local_world_size": local_world_size, + "node_count": world_size // local_world_size + if local_world_size and world_size % local_world_size == 0 + else None, + "data_parallel_size": data_parallel_size, + "data_parallel_replicate_size": replicate_size, + "data_parallel_shard_size": shard_size, + "tensor_parallel_size": tensor_parallel_size, + "pipeline_parallel_size": pipeline_parallel_size, + "expert_parallel_size": expert_parallel_size, + "ep_fsdp_size": ep_fsdp_size, + "ulysses_parallel_size": ulysses_parallel_size, + "ringattn_parallel_size": ringattn_parallel_size, + "micro_batch_size": point.micro_batch_size, + "gradient_accumulation_steps": ( + point.global_batch_size // point.micro_batch_size // data_parallel_size + if point.global_batch_size is not None + and point.micro_batch_size is not None + and data_parallel_size + and point.global_batch_size % (point.micro_batch_size * data_parallel_size) == 0 + else None + ), + "global_batch_size": point.global_batch_size, + "sample_packing_sequence_len": point.sample_packing_sequence_len, + } + + +def _benchmark_point_dimension_value( + point: BenchmarkBehaviorPoint, + dimension: str, + base_topology: Topology, +) -> Any: + return _benchmark_point_topology_values(point, base_topology).get(dimension) + + +def _benchmark_signature_for_dimensions( + point: BenchmarkBehaviorPoint, + dimensions: tuple[str, ...], + base_topology: Topology, +) -> tuple[tuple[str, Any], ...]: + return tuple( + (dimension, _benchmark_point_dimension_value(point, dimension, base_topology)) for dimension in dimensions + ) + + +def _benchmark_point_runtime_dimension_value(point: BenchmarkBehaviorPoint, dimension: str) -> str: + value = getattr(point, dimension, None) + if value is None: + return "unknown" + return str(value) + + +def _benchmark_runtime_signature_for_point(point: BenchmarkBehaviorPoint) -> tuple[tuple[str, Any], ...]: + return ( + ( + "target_runtime_signature", + ",".join( + f"{dimension}={_benchmark_point_runtime_dimension_value(point, dimension)}" + for dimension in _SCENARIO_RUNTIME_DIMENSIONS + ), + ), + ) + + +def _benchmark_varied_dimensions( + points: list[BenchmarkBehaviorPoint], + dimensions: tuple[str, ...], + base_topology: Topology, +) -> list[str]: + varied: list[str] = [] + for dimension in dimensions: + values = {_benchmark_point_dimension_value(point, dimension, base_topology) for point in points} + values.discard(None) + if len(values) > 1: + varied.append(dimension) + return varied + + +def _benchmark_varied_runtime_dimensions(points: list[BenchmarkBehaviorPoint]) -> list[str]: + varied: list[str] = [] + for dimension in _SCENARIO_RUNTIME_DIMENSIONS: + values = {_benchmark_point_runtime_dimension_value(point, dimension) for point in points} + if len(values) > 1: + varied.append(dimension) + return varied + + +def _benchmark_axis_value( + point: BenchmarkBehaviorPoint, + dimensions: tuple[str, ...], + base_topology: Topology, +) -> str: + return _format_signature(_benchmark_signature_for_dimensions(point, dimensions, base_topology)) + + +def _benchmark_point_is_memory_blocked(point: BenchmarkBehaviorPoint) -> bool: + return point.correctness_status == "oom" + + +def _benchmark_parallelism_axis_coverage( + points: list[BenchmarkBehaviorPoint], + base_topology: Topology, +) -> list[ScenarioParallelismAxisCoverage]: + coverage: list[ScenarioParallelismAxisCoverage] = [] + workload_dimensions = _SCENARIO_WORKLOAD_DIMENSIONS + for axis, axis_dimensions in _PARALLELISM_AXIS_DIMENSIONS.items(): + outside_dimensions = tuple( + dimension for dimension in _PARALLELISM_COMPARISON_DIMENSIONS if dimension not in axis_dimensions + ) + groups: dict[tuple[tuple[str, Any], ...], list[BenchmarkBehaviorPoint]] = {} + for point in points: + key = ( + *_benchmark_signature_for_dimensions(point, workload_dimensions, base_topology), + *_benchmark_signature_for_dimensions(point, outside_dimensions, base_topology), + *_benchmark_runtime_signature_for_point(point), + ) + groups.setdefault(key, []).append(point) + + varied_groups = [ + group + for group in groups.values() + if len({_benchmark_axis_value(point, axis_dimensions, base_topology) for point in group}) > 1 + ] + if not varied_groups: + varied_dimensions = _benchmark_varied_dimensions(points, axis_dimensions, base_topology) + primary_dimensions = [ + dimension for dimension in varied_dimensions if dimension in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + co_varied_dimensions = [ + dimension + for dimension in varied_dimensions + if dimension not in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + if primary_dimensions: + scored_points = [point for point in points if point.tokens_per_sec is not None] + blocked_points = [point for point in points if _benchmark_point_is_memory_blocked(point)] + unscored_points = [ + point + for point in points + if point.tokens_per_sec is None and not _benchmark_point_is_memory_blocked(point) + ] + if scored_points and blocked_points: + status = "confounded_benchmark_single_scored_axis_with_blocked_alternatives" + elif blocked_points: + status = "confounded_benchmark_blocked_parallelism_axis" + elif scored_points: + status = "confounded_benchmark_scored_parallelism_axis" + else: + status = "confounded_benchmark_unscored_parallelism_axis" + coverage.append( + ScenarioParallelismAxisCoverage( + axis=axis, + status=status, + candidate_group_count=0, + candidate_count=len(points), + scored_count=len(scored_points), + blocked_count=len(blocked_points), + unscored_count=len(unscored_points), + varied_dimensions=varied_dimensions, + primary_varied_dimensions=primary_dimensions, + co_varied_axis_dimensions=co_varied_dimensions, + confounded_runtime_dimensions=_benchmark_varied_runtime_dimensions(points), + feasibility_status_counts=_count_values( + [ + "observed_oom" if _benchmark_point_is_memory_blocked(point) else "observed_fit" + for point in points + ] + ), + ) + ) + continue + coverage.append( + ScenarioParallelismAxisCoverage( + axis=axis, + status="missing_benchmark_parallelism_axis", + candidate_group_count=0, + candidate_count=0, + scored_count=0, + blocked_count=0, + unscored_count=0, + confounded_runtime_dimensions=[], + ) + ) + continue + + grouped_points = [point for group in varied_groups for point in group] + scored_points = [point for point in grouped_points if point.tokens_per_sec is not None] + blocked_points = [point for point in grouped_points if _benchmark_point_is_memory_blocked(point)] + unscored_points = [ + point + for point in grouped_points + if point.tokens_per_sec is None and not _benchmark_point_is_memory_blocked(point) + ] + has_scored_axis_comparison = any( + len( + { + _benchmark_axis_value(point, axis_dimensions, base_topology) + for point in group + if point.tokens_per_sec is not None + } + ) + > 1 + for group in varied_groups + ) + has_single_scored_blocked_alternatives = any( + len([point for point in group if point.tokens_per_sec is not None]) == 1 + and any(point.tokens_per_sec is None for point in group) + and any(_benchmark_point_is_memory_blocked(point) for point in group) + for group in varied_groups + ) + if has_scored_axis_comparison: + status = "scored_benchmark_parallelism_axis" + elif has_single_scored_blocked_alternatives: + status = "benchmark_single_scored_axis_with_blocked_alternatives" + elif blocked_points: + status = "blocked_benchmark_parallelism_axis" + else: + status = "unscored_benchmark_parallelism_axis" + varied_dimensions = _benchmark_varied_dimensions(grouped_points, axis_dimensions, base_topology) + primary_dimensions = [ + dimension for dimension in varied_dimensions if dimension in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + co_varied_dimensions = [ + dimension for dimension in varied_dimensions if dimension not in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + coverage.append( + ScenarioParallelismAxisCoverage( + axis=axis, + status=status, + candidate_group_count=len(varied_groups), + candidate_count=len(grouped_points), + scored_count=len(scored_points), + blocked_count=len(blocked_points), + unscored_count=len(unscored_points), + varied_dimensions=varied_dimensions, + primary_varied_dimensions=primary_dimensions, + co_varied_axis_dimensions=co_varied_dimensions, + confounded_runtime_dimensions=[], + feasibility_status_counts=_count_values( + [ + "observed_oom" if _benchmark_point_is_memory_blocked(point) else "observed_fit" + for point in grouped_points + ] + ), + ) + ) + return coverage + + +def _benchmark_support_status( + *, + point_count: int, + scored_count: int, + memory_blocked_count: int, + varied_parallelism_dimensions: list[str], + varied_workload_dimensions: list[str], + varied_runtime_dimensions: list[str], + parallelism_axis_coverage_status_counts: dict[str, int], +) -> tuple[str, list[str]]: + if point_count == 0: + return "no_benchmark_support", ["no_benchmark_points"] + + has_parallelism_variation = bool(varied_parallelism_dimensions) + has_workload_variation = bool(varied_workload_dimensions) + clean_axis_count = parallelism_axis_coverage_status_counts.get("scored_benchmark_parallelism_axis", 0) + blockers: set[str] = set() + if point_count == 1: + blockers.add("single_benchmark_point") + if scored_count == 0: + blockers.add("no_scored_benchmark_points") + elif scored_count < point_count: + blockers.add("unscored_benchmark_points") + if memory_blocked_count == point_count: + blockers.add("memory_blocked_all_benchmark_points") + if not has_parallelism_variation: + blockers.add("missing_parallelism_benchmark_support") + if not has_workload_variation: + blockers.add("missing_workload_benchmark_support") + if varied_runtime_dimensions: + blockers.add("runtime_variant_benchmark_support") + if has_parallelism_variation and clean_axis_count == 0: + blockers.add("no_clean_benchmark_parallelism_axis_coverage") + if any("blocked" in status for status in parallelism_axis_coverage_status_counts): + blockers.add("blocked_benchmark_parallelism_axes") + if any(status.startswith("confounded_") for status in parallelism_axis_coverage_status_counts): + blockers.add("confounded_benchmark_parallelism_axes") + if any("unscored" in status for status in parallelism_axis_coverage_status_counts): + blockers.add("unscored_benchmark_parallelism_axes") + + blocker_list = sorted(blockers) + if not has_parallelism_variation and not has_workload_variation: + return "single_shape_benchmark_support", blocker_list + if has_workload_variation and not has_parallelism_variation: + return "workload_only_benchmark_support", blocker_list + if has_parallelism_variation and not has_workload_variation: + if scored_count == 0: + return "unscored_parallelism_benchmark_support", blocker_list + if clean_axis_count == 0: + return "parallelism_benchmark_support_without_clean_axis", blocker_list + if blockers - {"missing_workload_benchmark_support"}: + return "partial_parallelism_benchmark_support", blocker_list + return "parallelism_only_benchmark_support", blocker_list + + if scored_count == 0: + return "unscored_broad_benchmark_support", blocker_list + if clean_axis_count == 0: + return "confounded_broad_benchmark_support", blocker_list + if blockers: + return "partial_broad_benchmark_support", blocker_list + return "broad_benchmark_support", [] + + +def _scenario_benchmark_support( + behavior_points: list[BenchmarkBehaviorPoint], + *, + base_config: dict[str, Any], + base_topology: Topology, +) -> ScenarioBenchmarkSupport: + support_points = [ + point + for point in behavior_points + if not behavior_point_model_mismatches(point, base_config) + and (point.tokens_per_sec is not None or point.correctness_status == "oom") + and ( + point.sample_packing_sequence_len is None + or base_topology.sample_packing_sequence_len is None + or point.sample_packing_sequence_len == base_topology.sample_packing_sequence_len + ) + ] + if not support_points: + return ScenarioBenchmarkSupport() + + parallelism_axis_coverage = _benchmark_parallelism_axis_coverage(support_points, base_topology) + parallelism_axis_coverage_status_counts = _count_values([coverage.status for coverage in parallelism_axis_coverage]) + varied_parallelism_dimensions = _benchmark_varied_dimensions( + support_points, + _SCENARIO_PARALLELISM_DIMENSIONS, + base_topology, + ) + varied_workload_dimensions = _benchmark_varied_dimensions( + support_points, + _SCENARIO_WORKLOAD_DIMENSIONS, + base_topology, + ) + varied_runtime_dimensions = _benchmark_varied_runtime_dimensions(support_points) + scored_count = sum(1 for point in support_points if point.tokens_per_sec is not None) + memory_blocked_count = sum(1 for point in support_points if _benchmark_point_is_memory_blocked(point)) + support_status, support_blockers = _benchmark_support_status( + point_count=len(support_points), + scored_count=scored_count, + memory_blocked_count=memory_blocked_count, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + parallelism_axis_coverage_status_counts=parallelism_axis_coverage_status_counts, + ) + return ScenarioBenchmarkSupport( + support_status=support_status, + support_blockers=support_blockers, + point_count=len(support_points), + scored_count=scored_count, + memory_blocked_count=memory_blocked_count, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + parallelism_axis_coverage_status_counts=parallelism_axis_coverage_status_counts, + scored_parallelism_axis_names=[ + coverage.axis + for coverage in parallelism_axis_coverage + if coverage.status == "scored_benchmark_parallelism_axis" + ], + blocked_parallelism_axis_names=[ + coverage.axis for coverage in parallelism_axis_coverage if "blocked" in coverage.status + ], + confounded_parallelism_axis_names=[ + coverage.axis for coverage in parallelism_axis_coverage if coverage.status.startswith("confounded_") + ], + unscored_parallelism_axis_names=[ + coverage.axis + for coverage in parallelism_axis_coverage + if coverage.status == "unscored_benchmark_parallelism_axis" + ], + missing_parallelism_axis_names=[ + coverage.axis + for coverage in parallelism_axis_coverage + if coverage.status == "missing_benchmark_parallelism_axis" + ], + point_labels=sorted(point.label for point in support_points), + ) + + +def _axis_value(candidate: ScenarioCandidate, dimensions: tuple[str, ...]) -> str: + return _format_signature(_signature_for_dimensions(candidate, dimensions)) + + +def _axis_best_worst( + candidates: list[ScenarioCandidate], + *, + score_attr: str, +) -> tuple[ScenarioCandidate | None, ScenarioCandidate | None, float | None, float | None]: + scored = [candidate for candidate in candidates if getattr(candidate, score_attr) is not None] + if not scored: + return None, None, None, None + best = max(scored, key=lambda candidate: (getattr(candidate, score_attr), candidate.label)) + worst = min(scored, key=lambda candidate: (getattr(candidate, score_attr), candidate.label)) + best_score = getattr(best, score_attr) + worst_score = getattr(worst, score_attr) + spread = round(best_score - worst_score, 3) + ratio = round(best_score / worst_score, 3) if worst_score and worst_score > 0 else None + return best, worst, spread, ratio + + +def _axis_comparison_status( + raw_best: ScenarioCandidate | None, + raw_spread: float | None, + risk_adjusted_best: ScenarioCandidate | None, + risk_adjusted_spread: float | None, +) -> str: + if raw_best is None and risk_adjusted_best is None: + return "unscored_axis_comparison" + if raw_best is None: + return "risk_adjusted_only_axis_comparison" + if risk_adjusted_best is None: + return "raw_only_axis_comparison" + if raw_spread == 0 and risk_adjusted_spread == 0: + return "axis_tie" + if raw_best.label == risk_adjusted_best.label: + return "raw_and_risk_adjusted_agree" + return "risk_adjusted_changes_axis_winner" + + +def _axis_interval_overlap_summary( + group: list[ScenarioCandidate], + risk_adjusted_best: ScenarioCandidate | None, +) -> tuple[str, int, list[str], float | None]: + if risk_adjusted_best is None: + return "no_scored_interval", 0, [], None + best_lower = risk_adjusted_best.risk_adjusted_prediction_interval_lower_tokens_per_sec + best_upper = risk_adjusted_best.risk_adjusted_prediction_interval_upper_tokens_per_sec + if best_lower is None or best_upper is None: + return "no_scored_interval", 0, [], None + + overlap_labels: list[str] = [] + other_upper_bounds: list[float] = [] + for candidate in group: + if candidate.label == risk_adjusted_best.label: + continue + lower = candidate.risk_adjusted_prediction_interval_lower_tokens_per_sec + upper = candidate.risk_adjusted_prediction_interval_upper_tokens_per_sec + if lower is None or upper is None: + continue + other_upper_bounds.append(upper) + if lower <= best_upper and upper >= best_lower: + overlap_labels.append(candidate.label) + + if not other_upper_bounds: + return "single_scored_interval", 0, [], None + + margin = round(best_lower - max(other_upper_bounds), 3) + if overlap_labels: + return "overlapping_best_interval", len(overlap_labels), sorted(overlap_labels), margin + return "clear_best_interval", 0, [], margin + + +def _parallelism_axis_comparisons(candidates: list[ScenarioCandidate]) -> list[ParallelismAxisComparison]: + comparisons: list[ParallelismAxisComparison] = [] + workload_dimensions = _SCENARIO_WORKLOAD_DIMENSIONS + for axis, axis_dimensions in _PARALLELISM_AXIS_DIMENSIONS.items(): + outside_dimensions = tuple( + dimension for dimension in _PARALLELISM_COMPARISON_DIMENSIONS if dimension not in axis_dimensions + ) + groups: dict[tuple[tuple[str, Any], ...], list[ScenarioCandidate]] = {} + for candidate in candidates: + if candidate.score_tokens_per_sec is None: + continue + key = ( + *_signature_for_dimensions(candidate, workload_dimensions), + *_signature_for_dimensions(candidate, outside_dimensions), + *_runtime_signature_for_candidate(candidate), + ) + groups.setdefault(key, []).append(candidate) + + for group_key, group in groups.items(): + axis_values = {_axis_value(candidate, axis_dimensions) for candidate in group} + if len(group) < 2 or len(axis_values) < 2: + continue + + raw_best, raw_worst, raw_spread, raw_ratio = _axis_best_worst( + group, + score_attr="score_tokens_per_sec", + ) + risk_best, risk_worst, risk_spread, risk_ratio = _axis_best_worst( + group, + score_attr="score_risk_adjusted_tokens_per_sec", + ) + ( + interval_overlap_status, + interval_overlap_count, + interval_overlap_labels, + interval_margin, + ) = _axis_interval_overlap_summary(group, risk_best) + status = _axis_comparison_status(raw_best, raw_spread, risk_best, risk_spread) + varied_dimensions = _varied_candidate_dimensions(group, axis_dimensions) + primary_dimensions = [ + dimension for dimension in varied_dimensions if dimension in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + co_varied_dimensions = [ + dimension + for dimension in varied_dimensions + if dimension not in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + comparisons.append( + ParallelismAxisComparison( + axis=axis, + varied_dimensions=varied_dimensions, + primary_varied_dimensions=primary_dimensions, + co_varied_axis_dimensions=co_varied_dimensions, + coupling_status=("coupled_axis_comparison" if co_varied_dimensions else "isolated_axis_comparison"), + group_key=_format_signature(group_key), + candidate_count=len(group), + raw_best_label=raw_best.label if raw_best is not None else None, + raw_best_axis_value=_axis_value(raw_best, axis_dimensions) if raw_best is not None else None, + raw_best_score_tokens_per_sec=(raw_best.score_tokens_per_sec if raw_best is not None else None), + raw_worst_label=raw_worst.label if raw_worst is not None else None, + raw_worst_axis_value=_axis_value(raw_worst, axis_dimensions) if raw_worst is not None else None, + raw_worst_score_tokens_per_sec=(raw_worst.score_tokens_per_sec if raw_worst is not None else None), + raw_spread_tokens_per_sec=raw_spread, + raw_spread_ratio=raw_ratio, + risk_adjusted_best_label=risk_best.label if risk_best is not None else None, + risk_adjusted_best_axis_value=( + _axis_value(risk_best, axis_dimensions) if risk_best is not None else None + ), + risk_adjusted_best_score_tokens_per_sec=( + risk_best.score_risk_adjusted_tokens_per_sec if risk_best is not None else None + ), + risk_adjusted_worst_label=risk_worst.label if risk_worst is not None else None, + risk_adjusted_worst_axis_value=( + _axis_value(risk_worst, axis_dimensions) if risk_worst is not None else None + ), + risk_adjusted_worst_score_tokens_per_sec=( + risk_worst.score_risk_adjusted_tokens_per_sec if risk_worst is not None else None + ), + risk_adjusted_spread_tokens_per_sec=risk_spread, + risk_adjusted_spread_ratio=risk_ratio, + risk_adjusted_winner_matches_raw=( + raw_best.label == risk_best.label if raw_best is not None and risk_best is not None else None + ), + comparison_status=status, + risk_adjusted_best_interval_lower_tokens_per_sec=( + risk_best.risk_adjusted_prediction_interval_lower_tokens_per_sec + if risk_best is not None + else None + ), + risk_adjusted_best_interval_upper_tokens_per_sec=( + risk_best.risk_adjusted_prediction_interval_upper_tokens_per_sec + if risk_best is not None + else None + ), + risk_adjusted_worst_interval_lower_tokens_per_sec=( + risk_worst.risk_adjusted_prediction_interval_lower_tokens_per_sec + if risk_worst is not None + else None + ), + risk_adjusted_worst_interval_upper_tokens_per_sec=( + risk_worst.risk_adjusted_prediction_interval_upper_tokens_per_sec + if risk_worst is not None + else None + ), + risk_adjusted_interval_overlap_status=interval_overlap_status, + risk_adjusted_interval_overlap_candidate_count=interval_overlap_count, + risk_adjusted_interval_overlap_candidate_labels=interval_overlap_labels, + risk_adjusted_interval_margin_tokens_per_sec=interval_margin, + ) + ) + + return sorted( + comparisons, + key=lambda comparison: ( + comparison.risk_adjusted_spread_tokens_per_sec + if comparison.risk_adjusted_spread_tokens_per_sec is not None + else float("-inf"), + comparison.raw_spread_tokens_per_sec if comparison.raw_spread_tokens_per_sec is not None else float("-inf"), + comparison.axis, + comparison.group_key, + ), + reverse=True, + ) + + +def _parallelism_axis_applies_to_scenario_candidates(axis: str, candidates: list[ScenarioCandidate]) -> bool: + if not candidates: + return True + if axis == "ulysses": + return any( + candidate.topology.ulysses_parallel_size > 1 + or (candidate.topology.sample_packing_sequence_len or 0) >= _MIN_ULYSSES_SEQUENCE_LEN + for candidate in candidates + ) + if axis == "ringattn": + return any( + candidate.topology.ringattn_parallel_size > 1 + or (candidate.topology.sample_packing_sequence_len or 0) >= _MIN_RINGATTN_SEQUENCE_LEN + for candidate in candidates + ) + return True + + +def _scenario_parallelism_axis_coverage(candidates: list[ScenarioCandidate]) -> list[ScenarioParallelismAxisCoverage]: + coverage: list[ScenarioParallelismAxisCoverage] = [] + workload_dimensions = _SCENARIO_WORKLOAD_DIMENSIONS + for axis, axis_dimensions in _PARALLELISM_AXIS_DIMENSIONS.items(): + if not _parallelism_axis_applies_to_scenario_candidates(axis, candidates): + coverage.append( + ScenarioParallelismAxisCoverage( + axis=axis, + status="not_applicable_parallelism_axis", + candidate_group_count=0, + candidate_count=len(candidates), + scored_count=0, + blocked_count=0, + unscored_count=0, + confounded_runtime_dimensions=[], + ) + ) + continue + outside_dimensions = tuple( + dimension for dimension in _PARALLELISM_COMPARISON_DIMENSIONS if dimension not in axis_dimensions + ) + groups: dict[tuple[tuple[str, Any], ...], list[ScenarioCandidate]] = {} + for candidate in candidates: + key = ( + *_signature_for_dimensions(candidate, workload_dimensions), + *_signature_for_dimensions(candidate, outside_dimensions), + *_runtime_signature_for_candidate(candidate), + ) + groups.setdefault(key, []).append(candidate) + + varied_groups = [ + group + for group in groups.values() + if len({_axis_value(candidate, axis_dimensions) for candidate in group}) > 1 + ] + if not varied_groups: + varied_dimensions = _varied_candidate_dimensions(candidates, axis_dimensions) + primary_dimensions = [ + dimension for dimension in varied_dimensions if dimension in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + co_varied_dimensions = [ + dimension + for dimension in varied_dimensions + if dimension not in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + if primary_dimensions: + scored_candidates = [ + candidate for candidate in candidates if candidate.score_tokens_per_sec is not None + ] + blocked_candidates = [candidate for candidate in candidates if _is_memory_blocked(candidate)] + unscored_candidates = [ + candidate + for candidate in candidates + if candidate.score_tokens_per_sec is None and not _is_memory_blocked(candidate) + ] + if scored_candidates and blocked_candidates: + status = "confounded_single_scored_axis_with_blocked_alternatives" + elif blocked_candidates: + status = "confounded_blocked_parallelism_axis" + elif scored_candidates: + status = "confounded_scored_parallelism_axis" + else: + status = "confounded_unscored_parallelism_axis" + coverage.append( + ScenarioParallelismAxisCoverage( + axis=axis, + status=status, + candidate_group_count=0, + candidate_count=len(candidates), + scored_count=len(scored_candidates), + blocked_count=len(blocked_candidates), + unscored_count=len(unscored_candidates), + varied_dimensions=varied_dimensions, + primary_varied_dimensions=primary_dimensions, + co_varied_axis_dimensions=co_varied_dimensions, + confounded_runtime_dimensions=_varied_candidate_runtime_dimensions(candidates), + feasibility_status_counts=_count_values( + [candidate.feasibility_status for candidate in candidates] + ), + ) + ) + continue + coverage.append( + ScenarioParallelismAxisCoverage( + axis=axis, + status="missing_parallelism_axis", + candidate_group_count=0, + candidate_count=0, + scored_count=0, + blocked_count=0, + unscored_count=0, + confounded_runtime_dimensions=[], + ) + ) + continue + + grouped_candidates = [candidate for group in varied_groups for candidate in group] + scored_candidates = [ + candidate for candidate in grouped_candidates if candidate.score_tokens_per_sec is not None + ] + blocked_candidates = [candidate for candidate in grouped_candidates if _is_memory_blocked(candidate)] + unscored_candidates = [ + candidate + for candidate in grouped_candidates + if candidate.score_tokens_per_sec is None and not _is_memory_blocked(candidate) + ] + has_scored_axis_comparison = any( + len( + { + _axis_value(candidate, axis_dimensions) + for candidate in group + if candidate.score_tokens_per_sec is not None + } + ) + > 1 + for group in varied_groups + ) + has_single_scored_blocked_alternatives = any( + len([candidate for candidate in group if candidate.score_tokens_per_sec is not None]) == 1 + and any(candidate.score_tokens_per_sec is None for candidate in group) + and any(_is_memory_blocked(candidate) for candidate in group) + for group in varied_groups + ) + if has_scored_axis_comparison: + status = "scored_parallelism_axis" + elif has_single_scored_blocked_alternatives: + status = "single_scored_axis_with_blocked_alternatives" + elif blocked_candidates: + status = "blocked_parallelism_axis" + else: + status = "unscored_parallelism_axis" + varied_dimensions = _varied_candidate_dimensions(grouped_candidates, axis_dimensions) + primary_dimensions = [ + dimension for dimension in varied_dimensions if dimension in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + co_varied_dimensions = [ + dimension for dimension in varied_dimensions if dimension not in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + ] + coverage.append( + ScenarioParallelismAxisCoverage( + axis=axis, + status=status, + candidate_group_count=len(varied_groups), + candidate_count=len(grouped_candidates), + scored_count=len(scored_candidates), + blocked_count=len(blocked_candidates), + unscored_count=len(unscored_candidates), + varied_dimensions=varied_dimensions, + primary_varied_dimensions=primary_dimensions, + co_varied_axis_dimensions=co_varied_dimensions, + confounded_runtime_dimensions=[], + feasibility_status_counts=_count_values( + [candidate.feasibility_status for candidate in grouped_candidates] + ), + ) + ) + return coverage + + +def _scenario_boundary_dimension_value(candidate: ScenarioCandidate, dimension: str) -> Any: + if dimension == "balanced_routing": + return candidate.behavior.balanced_routing + return getattr(candidate.topology, dimension) + + +def _scenario_boundary_varied_dimensions( + candidates: list[ScenarioCandidate], + dimensions: tuple[str, ...], +) -> list[str]: + varied: list[str] = [] + for dimension in dimensions: + values = {_scenario_boundary_dimension_value(candidate, dimension) for candidate in candidates} + if len(values) > 1: + varied.append(dimension) + return varied + + +def _scenario_boundary_signature(candidate: ScenarioCandidate) -> tuple[tuple[str, Any], ...]: + return tuple( + (dimension, _scenario_boundary_dimension_value(candidate, dimension)) + for dimension in _SCENARIO_BOUNDARY_SIGNATURE_DIMENSIONS + ) + + +def _scenario_boundary_outcome(candidate: ScenarioCandidate) -> str | None: + if _is_memory_blocked(candidate): + return "failure" + if candidate.score_tokens_per_sec is not None: + return "fit" + return None + + +def _scenario_parallelism_boundary_groups(candidates: list[ScenarioCandidate]) -> list[ParallelismBoundaryGroup]: + grouped: dict[tuple[tuple[str, Any], ...], list[ScenarioCandidate]] = {} + for candidate in candidates: + if _scenario_boundary_outcome(candidate) is None: + continue + grouped.setdefault(_scenario_boundary_signature(candidate), []).append(candidate) + + boundary_groups: list[ParallelismBoundaryGroup] = [] + for signature, group in grouped.items(): + varied_parallelism = _varied_candidate_dimensions(group, _SCENARIO_PARALLELISM_DIMENSIONS) + if not varied_parallelism: + continue + outcomes = [_scenario_boundary_outcome(candidate) for candidate in group] + if "fit" not in outcomes or "failure" not in outcomes: + continue + fits = [candidate for candidate in group if _scenario_boundary_outcome(candidate) == "fit"] + failures = [candidate for candidate in group if _scenario_boundary_outcome(candidate) == "failure"] + best_fit = max(fits, key=_candidate_sort_key) if fits else None + boundary_groups.append( + ParallelismBoundaryGroup( + signature=_format_signature(signature), + candidate_count=len(group), + fit_count=len(fits), + failure_count=len(failures), + best_fit_label=best_fit.label if best_fit is not None else None, + best_fit_tokens_per_sec=best_fit.score_tokens_per_sec if best_fit is not None else None, + failure_labels=sorted(candidate.label for candidate in failures), + varied_parallelism_dimensions=varied_parallelism, + confounded_workload_dimensions=[ + dimension + for dimension in _scenario_boundary_varied_dimensions(group, _SCENARIO_BOUNDARY_WORKLOAD_DIMENSIONS) + if dimension not in _SCENARIO_BOUNDARY_SIGNATURE_DIMENSIONS + ], + confounded_runtime_dimensions=_varied_candidate_runtime_dimensions(group), + ) + ) + return sorted(boundary_groups, key=lambda group: group.signature) + + +def _scenario_parallelism_boundary_status( + boundary_groups: list[ParallelismBoundaryGroup], + candidates: list[ScenarioCandidate], +) -> str: + observed = [candidate for candidate in candidates if _scenario_boundary_outcome(candidate) is not None] + if not observed: + return "insufficient_data" + if not boundary_groups: + if len({_parallelism_strategy_key(candidate) for candidate in observed}) > 1: + return "no_fit_failure_parallelism_boundary" + return "no_measured_parallelism_variation" + if any( + not group.confounded_workload_dimensions and not group.confounded_runtime_dimensions + for group in boundary_groups + ): + return "measured_parallelism_fit_failure_boundary" + return "confounded_parallelism_fit_failure_boundary" + + +def _scenario_parallelism_boundary_axis_coverage( + boundary_groups: list[ParallelismBoundaryGroup], +) -> list[ParallelismBoundaryAxisCoverage]: + coverage: list[ParallelismBoundaryAxisCoverage] = [] + for axis in _PARALLELISM_AXIS_DIMENSIONS: + primary_dimensions = _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis] + axis_groups = [ + group + for group in boundary_groups + if any(dimension in group.varied_parallelism_dimensions for dimension in primary_dimensions) + ] + if not axis_groups: + coverage.append( + ParallelismBoundaryAxisCoverage( + axis=axis, + status="missing_parallelism_boundary_axis", + group_count=0, + candidate_count=0, + fit_count=0, + failure_count=0, + varied_parallelism_dimensions=[], + co_varied_parallelism_dimensions=[], + confounded_workload_dimensions=[], + confounded_runtime_dimensions=[], + ) + ) + continue + + varied_parallelism_dimensions = [ + dimension + for dimension in _SCENARIO_PARALLELISM_DIMENSIONS + if any(dimension in group.varied_parallelism_dimensions for group in axis_groups) + ] + co_varied_parallelism_dimensions = [ + dimension for dimension in varied_parallelism_dimensions if dimension not in primary_dimensions + ] + confounded_workload_dimensions = sorted( + {dimension for group in axis_groups for dimension in group.confounded_workload_dimensions} + ) + confounded_runtime_dimensions = sorted( + {dimension for group in axis_groups for dimension in group.confounded_runtime_dimensions}, + key=_SCENARIO_RUNTIME_DIMENSIONS.index, + ) + status = ( + "confounded_parallelism_boundary_axis" + if confounded_workload_dimensions or confounded_runtime_dimensions + else "measured_parallelism_boundary_axis" + ) + coverage.append( + ParallelismBoundaryAxisCoverage( + axis=axis, + status=status, + group_count=len(axis_groups), + candidate_count=sum(group.candidate_count for group in axis_groups), + fit_count=sum(group.fit_count for group in axis_groups), + failure_count=sum(group.failure_count for group in axis_groups), + varied_parallelism_dimensions=[ + dimension for dimension in varied_parallelism_dimensions if dimension in primary_dimensions + ], + co_varied_parallelism_dimensions=co_varied_parallelism_dimensions, + confounded_workload_dimensions=confounded_workload_dimensions, + confounded_runtime_dimensions=confounded_runtime_dimensions, + ) + ) + return coverage + + +def _scenario_parallelism_boundary_prediction_support( + *, + parallelism_boundary_status: str, + parallelism_boundary_fit_count: int, + parallelism_boundary_failure_count: int, + parallelism_boundary_measured_axis_names: list[str], + parallelism_boundary_confounded_axis_names: list[str], + parallelism_boundary_missing_axis_names: list[str], + parallelism_boundary_confounded_dimensions: list[str], +) -> tuple[str, list[str]]: + blockers: list[str] = [] + if parallelism_boundary_fit_count == 0: + blockers.append("no_fit_rows") + if parallelism_boundary_failure_count == 0: + blockers.append("no_failure_rows") + if parallelism_boundary_confounded_axis_names: + blockers.append("confounded_parallelism_boundary_axes") + if parallelism_boundary_missing_axis_names: + blockers.append("missing_parallelism_boundary_axes") + if parallelism_boundary_confounded_dimensions: + blockers.append("confounded_boundary_dimensions") + + blockers = sorted(set(blockers)) + if parallelism_boundary_status == "insufficient_data": + return "insufficient_parallelism_boundary_data", sorted({*blockers, "insufficient_data"}) + if parallelism_boundary_status == "no_measured_parallelism_variation": + return "no_parallelism_boundary_variation", sorted({*blockers, "no_measured_parallelism_variation"}) + if parallelism_boundary_status == "no_fit_failure_parallelism_boundary": + return "no_fit_failure_boundary_evidence", sorted({*blockers, "no_fit_failure_parallelism_boundary"}) + if parallelism_boundary_status == "confounded_parallelism_fit_failure_boundary": + return "confounded_parallelism_boundary_prediction", sorted( + {*blockers, "confounded_parallelism_fit_failure_boundary"} + ) + if parallelism_boundary_status == "measured_parallelism_fit_failure_boundary": + if not parallelism_boundary_measured_axis_names: + return "parallelism_boundary_without_axis_coverage", sorted( + {*blockers, "no_measured_parallelism_boundary_axes"} + ) + if blockers: + return "partial_parallelism_fit_failure_boundary", blockers + return "validated_parallelism_fit_failure_boundary", [] + return "unknown_parallelism_boundary_prediction", sorted({*blockers, parallelism_boundary_status}) + + +def _parallelism_tradeoff_status( + *, + unique_strategy_count: int, + scored_strategy_count: int, + promotable_strategy_count: int, + requires_remeasurement_strategy_count: int, +) -> str: + if unique_strategy_count == 0: + return "no_candidates" + if unique_strategy_count == 1: + return "single_parallelism_strategy" + if scored_strategy_count == 0: + return "unscored_parallelism_tradeoff" + if scored_strategy_count == 1: + return "single_scored_strategy_with_blocked_alternatives" + if requires_remeasurement_strategy_count: + return "scored_parallelism_tradeoff_requires_remeasurement" + if promotable_strategy_count >= 2: + return "promotable_parallelism_tradeoff" + if promotable_strategy_count == 1: + return "single_promotable_strategy_with_scored_alternatives" + return "scored_parallelism_tradeoff_not_promotable" + + +def _throughput_efficiency_tradeoff_status( + *, + best_raw: ScenarioCandidate | None, + best_risk_adjusted: ScenarioCandidate | None, + best_efficiency: ScenarioCandidate | None, + best_risk_adjusted_efficiency: ScenarioCandidate | None, +) -> str: + if best_raw is None and best_efficiency is None: + return "no_scored_candidates" + raw_differs = best_raw is not None and best_efficiency is not None and best_raw.label != best_efficiency.label + risk_differs = ( + best_risk_adjusted is not None + and best_risk_adjusted_efficiency is not None + and best_risk_adjusted.label != best_risk_adjusted_efficiency.label + ) + if raw_differs and risk_differs: + return "raw_and_risk_adjusted_efficiency_diverge" + if raw_differs: + return "raw_throughput_efficiency_diverge" + if risk_differs: + return "risk_adjusted_throughput_efficiency_diverge" + return "throughput_efficiency_aligned" + + +def _same_workload_scaling_status(scaling_candidates: list[ScenarioCandidate]) -> str: + if not scaling_candidates: + return "no_same_workload_scaling_comparison" + min_efficiency = min(candidate.scaling_efficiency or 0.0 for candidate in scaling_candidates) + if min_efficiency < 0.50: + return "poor_same_workload_scaling" + if min_efficiency < 0.80: + return "sublinear_same_workload_scaling" + return "strong_same_workload_scaling" + + +def _candidate_score_gap( + best_candidate: ScenarioCandidate | None, + fallback_candidate: ScenarioCandidate | None, + score_attr: str, +) -> float | None: + if best_candidate is None or fallback_candidate is None: + return None + best_score = getattr(best_candidate, score_attr) + fallback_score = getattr(fallback_candidate, score_attr) + if best_score is None or fallback_score is None: + return None + return round(max(float(best_score) - float(fallback_score), 0.0), 3) + + +def _candidate_score_gap_percentage( + best_candidate: ScenarioCandidate | None, + score_attr: str, + gap: float | None, +) -> float | None: + if best_candidate is None or gap is None: + return None + best_score = getattr(best_candidate, score_attr) + if best_score is None or best_score <= 0: + return None + return round(gap / float(best_score) * 100.0, 3) + + +def _scenario_promotion_readiness_status( + *, + best_raw: ScenarioCandidate | None, + best_risk_adjusted: ScenarioCandidate | None, + best_promotable: ScenarioCandidate | None, +) -> str: + selected = best_risk_adjusted or best_raw + if selected is None: + return "no_scored_candidate" + if selected.promotable: + if best_raw is not None and selected.label == best_raw.label: + return "promote_raw_and_risk_adjusted_winner" + return "promote_risk_adjusted_winner" + if "requires_remeasurement" in selected.risk_flags or selected.prediction_confidence != "calibrated": + return "remeasure_risk_adjusted_winner_before_promotion" + if selected.recommendation == "debug_runtime_failure": + return "debug_risk_adjusted_winner_before_promotion" + if selected.recommendation == "correctness_gate_required": + return "correctness_gate_risk_adjusted_winner_before_promotion" + if selected.recommendation.startswith("remeasure"): + return "remeasure_risk_adjusted_winner_before_promotion" + if best_promotable is not None: + return "promote_best_promotable_fallback" + return "no_promotable_candidate" + + +def _exact_timing_support_status(candidate: ScenarioCandidate | None) -> str: + if candidate is None: + return "missing" + if candidate.timing_coverage_status == "exact_phase_timing": + return "exact_phase" + if candidate.timing_coverage_status == "exact_total_step_only": + return "exact_total_step_only" + if candidate.timing_coverage_status == "no_timing_evidence": + return "missing" + return "reference_or_extrapolated" + + +def _parallelism_optimality_support( + *, + unique_strategy_count: int, + scored_strategy_count: int, + memory_blocked_count: int, + best_risk_adjusted: ScenarioCandidate | None, + best_promotable: ScenarioCandidate | None, + risk_adjusted_interval_overlap_status: str, + parallelism_tradeoff_status: str, + parallelism_axis_coverage_status_counts: dict[str, int], + varied_runtime_dimensions: list[str], + simulator_support_status_counts: dict[str, int], + interval_overlap_only_promotable_tie: bool = False, +) -> tuple[str, list[str]]: + blockers: list[str] = [] + if unique_strategy_count == 0: + return "no_candidates", ["no_candidates"] + if unique_strategy_count == 1: + return "single_parallelism_strategy_no_tradeoff", ["single_parallelism_strategy"] + if scored_strategy_count == 0: + blocker = "memory_blocked_all_candidates" if memory_blocked_count else "no_scored_parallelism_candidates" + return "unscored_parallelism_tradeoff", [blocker] + if best_risk_adjusted is None: + return "no_risk_adjusted_winner", ["no_risk_adjusted_winner"] + + if ( + best_risk_adjusted.prediction_confidence != "calibrated" + or "requires_remeasurement" in best_risk_adjusted.risk_flags + ): + blockers.append("winner_requires_measurement") + if parallelism_tradeoff_status == "scored_parallelism_tradeoff_requires_remeasurement": + blockers.append("unmeasured_parallelism_alternatives") + if not best_risk_adjusted.promotable: + blockers.append("winner_not_promotable") + if best_promotable is None: + blockers.append("no_promotable_candidate") + elif best_promotable.label != best_risk_adjusted.label: + blockers.append("promotable_candidate_not_selected_winner") + if risk_adjusted_interval_overlap_status == "overlapping_best_interval": + if interval_overlap_only_promotable_tie: + # A measured tie among K3-promotable strategies IS a resolved tradeoff ("either is + # optimal"), not selection uncertainty: every overlapping contender and the winner are + # promotable, so the choice cannot be wrong. Keep the tie visible without blocking. + blockers.append("risk_adjusted_interval_tie_between_promotable_strategies") + else: + blockers.append("risk_adjusted_interval_overlap") + elif risk_adjusted_interval_overlap_status in {"unknown", "no_scored_interval"}: + blockers.append("missing_risk_adjusted_interval") + if parallelism_axis_coverage_status_counts.get("scored_parallelism_axis", 0) == 0: + blockers.append("no_clean_parallelism_axis_coverage") + if any(status.startswith("confounded_") for status in parallelism_axis_coverage_status_counts): + blockers.append("confounded_parallelism_axes") + if varied_runtime_dimensions: + blockers.append("runtime_variant_variation") + if any(status.startswith("unsupported_") for status in simulator_support_status_counts): + blockers.append("unsupported_simulator_surface") + partial_surface_count = sum( + count + for status, count in simulator_support_status_counts.items() + if status != "supported_local_non_pp" and not status.startswith("unsupported_") + ) + if partial_surface_count: + blockers.append("partial_simulator_surface_support") + + timing_status = _exact_timing_support_status(best_risk_adjusted) + if timing_status == "missing": + blockers.append("winner_missing_timing_evidence") + elif timing_status == "reference_or_extrapolated": + blockers.append("winner_timing_is_reference_or_extrapolated") + elif timing_status == "exact_total_step_only": + blockers.append("winner_missing_phase_timing") + + blockers = sorted(set(blockers)) + if "winner_requires_measurement" in blockers or "unmeasured_parallelism_alternatives" in blockers: + return "requires_measurement_before_parallelism_optimality", blockers + if "risk_adjusted_interval_overlap" in blockers: + return "interval_overlap_parallelism_uncertain", blockers + if ( + "no_clean_parallelism_axis_coverage" in blockers + or "confounded_parallelism_axes" in blockers + or "runtime_variant_variation" in blockers + ): + return "confounded_parallelism_winner", blockers + if "unsupported_simulator_surface" in blockers: + return "unsupported_surface_parallelism_winner", blockers + if "partial_simulator_surface_support" in blockers: + return "partial_surface_parallelism_winner", blockers + if "winner_not_promotable" in blockers or "no_promotable_candidate" in blockers: + return "not_promotable_parallelism_winner", blockers + if "winner_timing_is_reference_or_extrapolated" in blockers or "winner_missing_timing_evidence" in blockers: + return "timing_unsupported_parallelism_winner", blockers + if "winner_missing_phase_timing" in blockers: + return "total_step_only_parallelism_winner", blockers + if parallelism_tradeoff_status == "promotable_parallelism_tradeoff": + return "supported_promotable_parallelism_tradeoff", blockers + return "supported_parallelism_winner", blockers + + +def _measurement_portfolio_sort_key(candidate: ScenarioCandidate) -> tuple[float, float, str]: + risk_adjusted = ( + candidate.score_risk_adjusted_tokens_per_sec + if candidate.score_risk_adjusted_tokens_per_sec is not None + else float("-inf") + ) + raw = candidate.score_tokens_per_sec if candidate.score_tokens_per_sec is not None else float("-inf") + return risk_adjusted, raw, candidate.label + + +_MEASUREMENT_REASON_WEIGHTS = { + "best_next_measurement": 4.0, + "best_risk_adjusted_candidate": 1.5, + "best_raw_candidate": 1.0, + "best_promotable_candidate": 0.75, + "best_gpu_efficiency_candidate": 1.0, + "throughput_efficiency_frontier": 1.5, + "risk_adjusted_efficiency_frontier": 1.25, + "scored_parallelism_tradeoff_requires_remeasurement": 2.5, + "promotable_parallelism_tradeoff": 1.5, + "scored_parallelism_tradeoff_not_promotable": 2.0, + "single_promotable_strategy_with_scored_alternatives": 1.25, + "poor_same_workload_scaling": 2.5, + "sublinear_same_workload_scaling": 1.75, + "partial_simulator_surface_support": 2.0, +} + +_PARALLELISM_EXTRAPOLATION_AXIS_BY_FLAG = { + "parallelism_extrapolation:ep": "expert_parallel", + "parallelism_extrapolation:ep_fsdp": "ep_fsdp", + "parallelism_extrapolation:tp": "tensor_parallel", + "parallelism_extrapolation:pp": "pipeline_parallel", + "parallelism_extrapolation:ulysses": "ulysses", + "parallelism_extrapolation:ring": "ringattn", +} + +_MEASUREMENT_PARALLELISM_AXIS_GAP_STATUSES = { + "single_scored_axis_with_blocked_alternatives", + "blocked_parallelism_axis", + "unscored_parallelism_axis", + "confounded_single_scored_axis_with_blocked_alternatives", + "confounded_blocked_parallelism_axis", + "confounded_scored_parallelism_axis", + "confounded_unscored_parallelism_axis", +} + +_PHASE_TIMING_MEASUREMENT_CONFIG_OVERRIDES = ( + "train.enable_step_phase_timing=true", + "train.enable_per_component_timing=true", + "train.step_phase_timing_sync_cuda=true", +) + +_MEMORY_PRESSURE_MEASUREMENT_CONFIG_OVERRIDES = ( + "train.enable_step_phase_timing=true", + "train.enable_step_memory_profiling=true", +) + +_FIT_BOUNDARY_MEASUREMENT_CONFIG_OVERRIDES = ( + "train.enable_step_phase_timing=true", + "train.enable_step_memory_profiling=true", + "train.enable_per_component_timing=true", + "train.step_phase_timing_sync_cuda=true", +) + +_PHASE_TIMING_SCENARIO_BLOCKERS = frozenset( + { + "missing_phase_timing", + "missing_phase_bottleneck_evidence", + "no_timing_evidence", + "reference_or_extrapolated_timing", + } +) + +_PHASE_TIMING_DESIGN_MEASUREMENTS = frozenset( + { + "add_fit_failure_boundary_near_blocked_axes", + "add_same_parallelism_runtime_workload_variants", + "add_same_parallelism_workload_runtime_variants", + "add_same_workload_same_runtime_axis_pairs", + "add_same_workload_same_runtime_parallelism_axis_variants", + "add_workload_and_parallelism_variants", + "score_unscored_capture_candidates", + } +) + + +def _reason_priority_weight(reason: str) -> float: + if reason.startswith("cross_model_analog_support:"): + return 2.5 + if reason.startswith("cross_model_prediction_interval_top:"): + return 2.0 + if reason.startswith("parallelism_axis_gap:"): + if ":confounded_" in reason: + return 1.5 + if ":blocked_" in reason or ":single_scored_axis_with_blocked_" in reason: + return 1.25 + if ":unscored_" in reason: + return 1.0 + return 0.75 + if reason.startswith("phase_timing_gap:"): + return 1.25 + if reason.startswith("memory_pressure_probe:"): + return 2.0 + return _MEASUREMENT_REASON_WEIGHTS.get(reason, 0.5) + + +def _add_measurement_priority( + factors: list[str], + score: float, + *, + name: str, + weight: float, +) -> float: + factors.append(f"{name}={weight:.3f}") + return score + weight + + +def _timing_coverage_measurement_weight(timing_coverage_status: str) -> float: + if timing_coverage_status == "no_timing_evidence": + return 1.50 + if timing_coverage_status.startswith("reference_") or timing_coverage_status.startswith("cross_model_reference_"): + return 1.25 + if timing_coverage_status.startswith("step_time_fit_"): + return 1.00 + if timing_coverage_status == "exact_total_step_only": + return 0.50 + return 0.0 + + +def _measurement_priority( + candidate: ScenarioCandidate, + reasons: set[str], +) -> tuple[float, float, int, list[str]]: + is_memory_pressure_probe = any(reason.startswith("memory_pressure_probe:") for reason in reasons) + if candidate.score_tokens_per_sec is None: + score = 0.25 + factors = ["unscored_candidate=0.250"] + else: + score = 1.0 + factors = ["scored_candidate=1.000"] + + for reason in sorted(reasons): + weight = _reason_priority_weight(reason) + score = _add_measurement_priority(factors, score, name=f"reason:{reason}", weight=weight) + + if "requires_remeasurement" in candidate.risk_flags: + score = _add_measurement_priority(factors, score, name="requires_remeasurement", weight=2.0) + if candidate.prediction_confidence == "cross_model_extrapolated": + score = _add_measurement_priority(factors, score, name="cross_model_extrapolated", weight=2.0) + elif candidate.prediction_confidence != "calibrated": + score = _add_measurement_priority(factors, score, name="extrapolated_prediction", weight=1.25) + if "cross_model_analog" in candidate.risk_flags: + score = _add_measurement_priority(factors, score, name="cross_model_analog", weight=1.5) + if candidate.simulator_support_status != "supported_local_non_pp": + if candidate.simulator_support_status.startswith("unsupported_"): + score = _add_measurement_priority(factors, score, name="unsupported_simulator_surface", weight=3.0) + else: + score = _add_measurement_priority(factors, score, name="partial_simulator_surface", weight=2.0) + + timing_weight = _timing_coverage_measurement_weight(candidate.timing_coverage_status) + if timing_weight: + score = _add_measurement_priority( + factors, + score, + name=f"timing_coverage:{candidate.timing_coverage_status}", + weight=timing_weight, + ) + + if candidate.calibration_distance is not None: + if is_memory_pressure_probe: + closeness_weight = max(0.0, 2.0 - min(2.0, candidate.calibration_distance * 0.25)) + if closeness_weight > 0: + score = _add_measurement_priority( + factors, + score, + name="calibration_closeness", + weight=closeness_weight, + ) + elif candidate.calibration_distance > 0: + distance_weight = min(2.0, candidate.calibration_distance * 0.25) + score = _add_measurement_priority( + factors, + score, + name="calibration_distance", + weight=distance_weight, + ) + + if candidate.memory_coverage_status == "analytic_floor_only": + score = _add_measurement_priority(factors, score, name="analytic_memory_floor_only", weight=1.75) + elif candidate.memory_coverage_status.startswith("calibrated_overhead"): + score = _add_measurement_priority(factors, score, name="calibrated_overhead_memory", weight=1.0) + elif candidate.memory_coverage_status.startswith("extrapolated"): + score = _add_measurement_priority(factors, score, name="extrapolated_memory_peak", weight=1.25) + + if candidate.estimated_memory_residual_fraction_of_peak is not None: + residual_weight = min(1.5, candidate.estimated_memory_residual_fraction_of_peak * 1.5) + if residual_weight > 0: + score = _add_measurement_priority(factors, score, name="memory_residual_fraction", weight=residual_weight) + + if candidate.recommendation == "debug_runtime_failure": + score = _add_measurement_priority(factors, score, name="debug_runtime_failure", weight=3.0) + elif candidate.recommendation.startswith("remeasure"): + score = _add_measurement_priority(factors, score, name="remeasure_recommendation", weight=2.0) + elif candidate.recommendation == "correctness_gate_required": + score = _add_measurement_priority(factors, score, name="correctness_gate_required", weight=1.25) + + if candidate.scaling_efficiency is not None and candidate.scaling_gpu_ratio is not None: + if candidate.scaling_gpu_ratio > 1.0 and candidate.scaling_efficiency < 0.5: + score = _add_measurement_priority(factors, score, name="poor_scaling_efficiency", weight=2.0) + elif candidate.scaling_gpu_ratio > 1.0 and candidate.scaling_efficiency < 0.8: + score = _add_measurement_priority(factors, score, name="sublinear_scaling_efficiency", weight=1.25) + + boundary_prefixes = ( + "allocator_pressure_boundary:", + "communication_cross_node:", + "observed_oom_boundary:", + "runtime_mismatch:", + ) + boundary_count = sum(1 for flag in candidate.risk_flags if flag.startswith(boundary_prefixes)) + if boundary_count: + score = _add_measurement_priority( + factors, + score, + name="boundary_or_runtime_mismatch_count", + weight=min(2.0, boundary_count * 0.5), + ) + + bottleneck_count = sum(1 for flag in candidate.risk_flags if flag.endswith("_bottleneck")) + if bottleneck_count: + score = _add_measurement_priority( + factors, + score, + name="phase_bottleneck_count", + weight=min(1.5, bottleneck_count * 0.5), + ) + + cost_gpus = max(candidate.topology.world_size, 1) + priority_score = round(score, 3) + return priority_score, round(priority_score / cost_gpus, 3), cost_gpus, factors + + +def _allow_unscored_axis_gap_candidate(candidate: ScenarioCandidate) -> bool: + if candidate.score_tokens_per_sec is not None: + return True + if candidate.simulator_support_status.startswith("unsupported_"): + return False + return candidate.memory_coverage_status != "analytic_floor_only" + + +def _candidate_axis_value_tuple(candidate: ScenarioCandidate, axis: str) -> tuple[Any, ...]: + return tuple(getattr(candidate.topology, dimension) for dimension in _PARALLELISM_AXIS_PRIMARY_DIMENSIONS[axis]) + + +def _candidate_contributes_axis_gap( + candidate: ScenarioCandidate, + *, + axis: str, + reference_candidate: ScenarioCandidate | None, + candidates: list[ScenarioCandidate], +) -> bool: + all_values = {_candidate_axis_value_tuple(row, axis) for row in candidates} + if len(all_values) <= 1: + return False + if reference_candidate is None: + return True + return _candidate_axis_value_tuple(candidate, axis) != _candidate_axis_value_tuple(reference_candidate, axis) + + +def _synthesize_axis_gap_reason( + candidate: ScenarioCandidate, + *, + cross_model_analog_support_status: str, +) -> bool: + if candidate.score_tokens_per_sec is None: + return _allow_unscored_axis_gap_candidate(candidate) + return ( + candidate.prediction_confidence == "cross_model_extrapolated" + and cross_model_analog_support_status == "single_reference_cannot_rank_parallelism_variants" + ) + + +def _phase_timing_gap_status(candidate: ScenarioCandidate) -> str | None: + status = candidate.timing_coverage_status + if status == "exact_phase_timing" or status.endswith("_phase_timing"): + return None + if status == "exact_total_step_only" or status.endswith("_total_step_only"): + return status + if status.startswith(("reference_", "cross_model_reference_")): + return status + return None + + +def _allow_memory_pressure_probe_candidate(candidate: ScenarioCandidate) -> bool: + if candidate.score_tokens_per_sec is not None: + return False + if candidate.feasibility_status != "memory_floor_exceeds_safety_margin": + return False + if candidate.simulator_support_status != "supported_local_non_pp": + return False + return candidate.memory_coverage_status == "analytic_floor_only" + + +def _measurement_portfolio( + *, + candidates: list[ScenarioCandidate], + throughput_efficiency_frontier_labels: list[str], + risk_adjusted_efficiency_frontier_labels: list[str], + parallelism_axis_coverage: list[ScenarioParallelismAxisCoverage], + parallelism_tradeoff_status: str, + cross_model_analog_support_status: str, + cross_model_analog_prediction_interval_selectivity_status: str, + cross_model_analog_prediction_interval_top_labels: list[str], + same_workload_scaling_status: str, + best_raw: ScenarioCandidate | None, + best_risk_adjusted: ScenarioCandidate | None, + best_efficiency: ScenarioCandidate | None, + best_risk_adjusted_efficiency: ScenarioCandidate | None, + best_next_measurement: ScenarioCandidate | None, + best_promotable: ScenarioCandidate | None, + max_candidates: int = 4, +) -> tuple[list[str], dict[str, list[str]], dict[str, float], dict[str, float], dict[str, int], dict[str, list[str]]]: + by_label = {candidate.label: candidate for candidate in candidates} + reasons: dict[str, set[str]] = {} + + def add(candidate: ScenarioCandidate | None, reason: str, *, allow_unscored: bool = False) -> None: + if candidate is None: + return + if candidate.score_tokens_per_sec is None and not allow_unscored: + return + reasons.setdefault(candidate.label, set()).add(reason) + + def add_label(label: str, reason: str) -> None: + add(by_label.get(label), reason) + + add(best_next_measurement, "best_next_measurement") + add(best_risk_adjusted, "best_risk_adjusted_candidate") + add(best_raw, "best_raw_candidate") + if best_promotable is not None and best_next_measurement is None: + add(best_promotable, "best_promotable_candidate") + + if len(throughput_efficiency_frontier_labels) > 1: + for label in throughput_efficiency_frontier_labels: + add_label(label, "throughput_efficiency_frontier") + if len(risk_adjusted_efficiency_frontier_labels) > 1: + for label in risk_adjusted_efficiency_frontier_labels: + add_label(label, "risk_adjusted_efficiency_frontier") + + if parallelism_tradeoff_status in { + "scored_parallelism_tradeoff_requires_remeasurement", + "promotable_parallelism_tradeoff", + "scored_parallelism_tradeoff_not_promotable", + "single_promotable_strategy_with_scored_alternatives", + }: + for candidate in sorted(candidates, key=_measurement_portfolio_sort_key, reverse=True): + if ( + parallelism_tradeoff_status == "scored_parallelism_tradeoff_not_promotable" + and candidate.score_tokens_per_sec is not None + ): + add(candidate, parallelism_tradeoff_status) + elif "requires_remeasurement" in candidate.risk_flags or candidate.promotable: + add(candidate, parallelism_tradeoff_status) + + axis_status_by_name = { + coverage.axis: coverage.status + for coverage in parallelism_axis_coverage + if coverage.status in _MEASUREMENT_PARALLELISM_AXIS_GAP_STATUSES + } + if axis_status_by_name: + reference_candidate = best_risk_adjusted or best_raw + for candidate in sorted(candidates, key=_measurement_portfolio_sort_key, reverse=True): + for flag in candidate.risk_flags: + axis = _PARALLELISM_EXTRAPOLATION_AXIS_BY_FLAG.get(flag) + status = axis_status_by_name.get(axis) if axis is not None else None + if status is not None: + add( + candidate, + f"parallelism_axis_gap:{axis}:{status}", + allow_unscored=_allow_unscored_axis_gap_candidate(candidate), + ) + for axis, status in axis_status_by_name.items(): + if not _synthesize_axis_gap_reason( + candidate, + cross_model_analog_support_status=cross_model_analog_support_status, + ): + continue + if _candidate_contributes_axis_gap( + candidate, + axis=axis, + reference_candidate=reference_candidate, + candidates=candidates, + ): + add( + candidate, + f"parallelism_axis_gap:{axis}:{status}", + allow_unscored=_allow_unscored_axis_gap_candidate(candidate), + ) + + if cross_model_analog_support_status not in {"not_used", "no_scored_cross_model_candidates"}: + for candidate in sorted( + _cross_model_analog_candidates(candidates), key=_measurement_portfolio_sort_key, reverse=True + ): + add(candidate, f"cross_model_analog_support:{cross_model_analog_support_status}") + + if cross_model_analog_prediction_interval_selectivity_status in { + "partial_prediction_interval_top", + "nonselective_prediction_interval_top", + }: + for label in cross_model_analog_prediction_interval_top_labels or []: + add_label( + label, + f"cross_model_prediction_interval_top:{cross_model_analog_prediction_interval_selectivity_status}", + ) + + if same_workload_scaling_status in {"poor_same_workload_scaling", "sublinear_same_workload_scaling"}: + for candidate in sorted(candidates, key=_measurement_portfolio_sort_key, reverse=True): + if candidate.scaling_gpu_ratio is not None and candidate.scaling_gpu_ratio > 1.0: + add(candidate, same_workload_scaling_status) + + for candidate in sorted(candidates, key=_measurement_portfolio_sort_key, reverse=True): + if candidate.score_tokens_per_sec is None: + continue + phase_gap_status = _phase_timing_gap_status(candidate) + if phase_gap_status is not None: + add(candidate, f"phase_timing_gap:{phase_gap_status}") + + for candidate in sorted(candidates, key=_measurement_portfolio_sort_key, reverse=True): + if _allow_memory_pressure_probe_candidate(candidate): + add(candidate, f"memory_pressure_probe:{candidate.feasibility_status}", allow_unscored=True) + + for candidate in sorted(candidates, key=_measurement_portfolio_sort_key, reverse=True): + if ( + candidate.simulator_support_status != "supported_local_non_pp" + and not candidate.simulator_support_status.startswith("unsupported_") + ): + add(candidate, "partial_simulator_surface_support", allow_unscored=True) + + for candidate in [best_efficiency, best_risk_adjusted_efficiency]: + add(candidate, "best_gpu_efficiency_candidate") + + priority_by_label: dict[str, tuple[float, float, int, list[str]]] = { + label: _measurement_priority(by_label[label], candidate_reasons) for label, candidate_reasons in reasons.items() + } + ordered_labels = sorted( + reasons, + key=lambda label: (*priority_by_label[label][:2], *_measurement_portfolio_sort_key(by_label[label])), + reverse=True, + ) + categories_by_label = { + label: sorted({_measurement_reason_category(reason) for reason in candidate_reasons}) + for label, candidate_reasons in reasons.items() + } + ordered_categories = sorted( + {category for categories in categories_by_label.values() for category in categories}, + key=lambda category: (-_VALIDATION_ACTION_PRIORITY_BY_CATEGORY.get(category, 10), category), + ) + selected_labels: list[str] = [] + for category in ordered_categories: + if len(selected_labels) >= max_candidates: + break + label = next( + (candidate_label for candidate_label in ordered_labels if category in categories_by_label[candidate_label]), + None, + ) + if label is not None and label not in selected_labels: + selected_labels.append(label) + for label in ordered_labels: + if len(selected_labels) >= max_candidates: + break + if label not in selected_labels: + selected_labels.append(label) + return ( + selected_labels, + {label: sorted(reasons[label]) for label in selected_labels}, + {label: priority_by_label[label][0] for label in selected_labels}, + {label: priority_by_label[label][1] for label in selected_labels}, + {label: priority_by_label[label][2] for label in selected_labels}, + {label: priority_by_label[label][3] for label in selected_labels}, + ) + + +def _measurement_reason_category(reason: str) -> str: + if reason == "best_next_measurement": + return "best_next_measurement" + if reason.startswith(("cross_model_analog_support:", "cross_model_prediction_interval_top:")): + return "cross_model_analog" + if reason.startswith("parallelism_axis_gap:"): + return "parallelism_axis_gap" + if reason.startswith("phase_timing_gap:"): + return "phase_timing_gap" + if reason.startswith("memory_pressure_probe:"): + return "memory_pressure_probe" + if reason in { + "scored_parallelism_tradeoff_requires_remeasurement", + "promotable_parallelism_tradeoff", + "scored_parallelism_tradeoff_not_promotable", + "single_promotable_strategy_with_scored_alternatives", + }: + return "parallelism_tradeoff" + if reason in { + "best_gpu_efficiency_candidate", + "throughput_efficiency_frontier", + "risk_adjusted_efficiency_frontier", + }: + return "efficiency_tradeoff" + if reason in {"poor_same_workload_scaling", "sublinear_same_workload_scaling"}: + return "same_workload_scaling" + if reason == "partial_simulator_surface_support": + return "simulator_surface" + if reason in {"best_raw_candidate", "best_risk_adjusted_candidate", "best_promotable_candidate"}: + return "winner_tracking" + return "other" + + +def _measurement_portfolio_reason_category_counts(candidate_reasons: dict[str, list[str]]) -> dict[str, int]: + counts: Counter[str] = Counter() + for reasons in candidate_reasons.values(): + for category in {_measurement_reason_category(reason) for reason in reasons}: + counts[category] += 1 + return dict(sorted(counts.items())) + + +def _measurement_portfolio_axis_gap_names(candidate_reasons: dict[str, list[str]]) -> list[str]: + axes = { + parts[1] + for reasons in candidate_reasons.values() + for reason in reasons + if reason.startswith("parallelism_axis_gap:") + for parts in [reason.split(":")] + if len(parts) >= 3 + } + return sorted(axes) + + +def _measurement_portfolio_required_categories( + *, + best_next_measurement: ScenarioCandidate | None, + parallelism_tradeoff_status: str, + throughput_efficiency_tradeoff_status: str, + same_workload_scaling_status: str, + cross_model_analog_support_status: str, + candidate_reasons: dict[str, list[str]], +) -> set[str]: + required: set[str] = set() + if best_next_measurement is not None: + required.add("best_next_measurement") + if parallelism_tradeoff_status in { + "scored_parallelism_tradeoff_requires_remeasurement", + "promotable_parallelism_tradeoff", + "scored_parallelism_tradeoff_not_promotable", + "single_promotable_strategy_with_scored_alternatives", + }: + required.add("parallelism_tradeoff") + if throughput_efficiency_tradeoff_status not in {"throughput_efficiency_aligned", "no_scored_candidates"}: + required.add("efficiency_tradeoff") + if same_workload_scaling_status in {"poor_same_workload_scaling", "sublinear_same_workload_scaling"}: + required.add("same_workload_scaling") + if cross_model_analog_support_status not in {"not_used", "no_scored_cross_model_candidates"}: + required.add("cross_model_analog") + if _measurement_portfolio_axis_gap_names(candidate_reasons): + required.add("parallelism_axis_gap") + if any(any(reason.startswith("phase_timing_gap:") for reason in reasons) for reasons in candidate_reasons.values()): + required.add("phase_timing_gap") + if any( + any(reason.startswith("memory_pressure_probe:") for reason in reasons) for reasons in candidate_reasons.values() + ): + required.add("memory_pressure_probe") + if any("partial_simulator_surface_support" in reasons for reasons in candidate_reasons.values()): + required.add("simulator_surface") + return required + + +def _measurement_portfolio_coverage_status( + *, + candidate_reasons: dict[str, list[str]], + best_next_measurement: ScenarioCandidate | None, + parallelism_tradeoff_status: str, + throughput_efficiency_tradeoff_status: str, + same_workload_scaling_status: str, + cross_model_analog_support_status: str, +) -> tuple[str, list[str], dict[str, int], list[str], int]: + category_counts = _measurement_portfolio_reason_category_counts(candidate_reasons) + axis_gap_names = _measurement_portfolio_axis_gap_names(candidate_reasons) + required = _measurement_portfolio_required_categories( + best_next_measurement=best_next_measurement, + parallelism_tradeoff_status=parallelism_tradeoff_status, + throughput_efficiency_tradeoff_status=throughput_efficiency_tradeoff_status, + same_workload_scaling_status=same_workload_scaling_status, + cross_model_analog_support_status=cross_model_analog_support_status, + candidate_reasons=candidate_reasons, + ) + present = set(category_counts) + missing = sorted(required - present) + blockers = [f"missing_measurement_category:{category}" for category in missing] + cross_model_count = category_counts.get("cross_model_analog", 0) + if not candidate_reasons and required: + return "missing_required_measurement_coverage", blockers, category_counts, axis_gap_names, cross_model_count + if missing: + return "partial_required_measurement_gap_coverage", blockers, category_counts, axis_gap_names, cross_model_count + if required: + return "covers_required_measurement_gaps", [], category_counts, axis_gap_names, cross_model_count + if candidate_reasons: + return "opportunistic_measurement_portfolio", [], category_counts, axis_gap_names, cross_model_count + return "no_measurement_portfolio_needed", [], category_counts, axis_gap_names, cross_model_count + + +_VALIDATION_ACTION_PRIORITY_BY_CATEGORY = { + "memory_pressure_probe": 120, + "parallelism_axis_gap": 110, + "cross_model_analog": 105, + "best_next_measurement": 100, + "parallelism_tradeoff": 95, + "phase_timing_gap": 90, + "same_workload_scaling": 85, + "efficiency_tradeoff": 80, + "simulator_surface": 70, + "winner_tracking": 60, + "other": 10, +} + + +def _reason_matches_category(reason: str, category: str) -> bool: + return _measurement_reason_category(reason) == category + + +def _reason_suffixes_for_category(reasons: set[str], category: str) -> list[str]: + suffixes: set[str] = set() + for reason in reasons: + if not _reason_matches_category(reason, category): + continue + parts = reason.split(":") + if category == "parallelism_axis_gap" and len(parts) >= 3: + suffixes.add(parts[2]) + elif category in {"cross_model_analog", "phase_timing_gap", "memory_pressure_probe"} and len(parts) >= 2: + suffixes.add(parts[1]) + else: + suffixes.add(reason) + return sorted(suffixes) + + +def _parallelism_axis_names_for_reasons(reasons: set[str]) -> list[str]: + axes = { + parts[1] + for reason in reasons + if reason.startswith("parallelism_axis_gap:") + for parts in [reason.split(":")] + if len(parts) >= 3 + } + return sorted(axes) + + +def _validation_action_status(category: str, reason_statuses: list[str]) -> str: + if category == "memory_pressure_probe": + return "memory_pressure_probe_needed" + if category == "parallelism_axis_gap": + if any("confounded" in status and "blocked" in status for status in reason_statuses): + return "confounded_parallelism_axis_boundary_action" + if any("confounded" in status for status in reason_statuses): + return "confounded_parallelism_axis_action" + if any("blocked" in status for status in reason_statuses): + return "blocked_parallelism_axis_action" + if any("unscored" in status for status in reason_statuses): + return "unscored_parallelism_axis_action" + return "parallelism_axis_action" + if category == "cross_model_analog": + if any(status.startswith("nonselective_") or status.startswith("partial_") for status in reason_statuses): + return "cross_model_interval_tiebreak_action" + return "cross_model_analog_action" + if category == "best_next_measurement": + return "best_next_measurement_action" + if category == "parallelism_tradeoff": + return "parallelism_tradeoff_action" + if category == "phase_timing_gap": + return "phase_timing_instrumentation_action" + if category == "same_workload_scaling": + return "same_workload_scaling_action" + if category == "efficiency_tradeoff": + return "throughput_efficiency_tradeoff_action" + if category == "simulator_surface": + return "simulator_surface_action" + if category == "winner_tracking": + return "winner_tracking_action" + return "measurement_followup_action" + + +def _validation_action_required_measurement(action_status: str) -> str: + if action_status == "memory_pressure_probe_needed": + return "launch_memory_pressure_fit_probe_with_memory_profiling" + if action_status == "confounded_parallelism_axis_boundary_action": + return "same_workload_same_runtime_axis_pair_with_fit_failure_boundary" + if action_status == "confounded_parallelism_axis_action": + return "same_workload_same_runtime_axis_pair" + if action_status == "blocked_parallelism_axis_action": + return "same_workload_axis_pair_with_fit_failure_boundary" + if action_status == "unscored_parallelism_axis_action": + return "score_unscored_parallelism_axis_candidate" + if action_status == "parallelism_axis_action": + return "add_same_workload_same_runtime_axis_pair" + if action_status == "cross_model_interval_tiebreak_action": + return "measure_target_cross_model_interval_top_candidates" + if action_status == "cross_model_analog_action": + return "measure_target_cross_model_analog_candidate" + if action_status == "best_next_measurement_action": + return "remeasure_best_next_candidate_same_workload" + if action_status == "parallelism_tradeoff_action": + return "measure_parallelism_tradeoff_candidates" + if action_status == "phase_timing_instrumentation_action": + return "rerun_with_phase_timing" + if action_status == "same_workload_scaling_action": + return "measure_same_workload_scaling_pair" + if action_status == "throughput_efficiency_tradeoff_action": + return "measure_throughput_efficiency_frontier_candidates" + if action_status == "simulator_surface_action": + return "add_simulator_surface_support_or_run_probe" + if action_status == "winner_tracking_action": + return "rerun_current_winner_with_required_instrumentation" + return "inspect_measurement_candidate" + + +def _max_label_for_scores(labels: list[str], scores: dict[str, float]) -> str | None: + return max(labels, key=lambda label: (scores.get(label, float("-inf")), label), default=None) + + +def _unique_in_order(values: list[str]) -> list[str]: + seen: set[str] = set() + ordered: list[str] = [] + for value in values: + if value in seen: + continue + seen.add(value) + ordered.append(value) + return ordered + + +def _scenario_validation_actions( + *, + candidate_reasons: dict[str, list[str]], + candidate_priority_scores: dict[str, float], + candidate_priority_per_gpu: dict[str, float], + candidate_cost_gpus: dict[str, int], + candidate_config_overrides: dict[str, list[str]], +) -> list[ScenarioValidationAction]: + labels_by_category: dict[str, list[str]] = {} + reasons_by_category: dict[str, set[str]] = {} + for label, reasons in candidate_reasons.items(): + for category in sorted({_measurement_reason_category(reason) for reason in reasons}): + labels_by_category.setdefault(category, []).append(label) + reasons_by_category.setdefault(category, set()).update( + reason for reason in reasons if _reason_matches_category(reason, category) + ) + + actions: list[ScenarioValidationAction] = [] + for category, labels in labels_by_category.items(): + action_reasons = reasons_by_category.get(category, set()) + reason_statuses = _reason_suffixes_for_category(action_reasons, category) + action_status = _validation_action_status(category, reason_statuses) + required_measurement = _validation_action_required_measurement(action_status) + score_label = _max_label_for_scores(labels, candidate_priority_scores) + per_gpu_label = _max_label_for_scores(labels, candidate_priority_per_gpu) + config_overrides = _unique_in_order( + [override for label in labels for override in candidate_config_overrides.get(label, [])] + ) + actions.append( + ScenarioValidationAction( + action_status=action_status, + priority=_VALIDATION_ACTION_PRIORITY_BY_CATEGORY.get(category, 10), + required_measurement=required_measurement, + reason_category=category, + candidate_count=len(labels), + candidate_labels=labels, + total_gpu_count=sum(candidate_cost_gpus.get(label, 0) for label in labels), + max_priority_score=candidate_priority_scores.get(score_label) if score_label is not None else None, + max_priority_label=score_label, + max_priority_per_gpu=( + candidate_priority_per_gpu.get(per_gpu_label) if per_gpu_label is not None else None + ), + max_priority_per_gpu_label=per_gpu_label, + parallelism_axis_names=( + _parallelism_axis_names_for_reasons(action_reasons) if category == "parallelism_axis_gap" else [] + ), + reason_statuses=reason_statuses, + config_overrides=config_overrides, + ) + ) + return sorted( + actions, + key=lambda action: ( + -action.priority, + action.reason_category, + action.required_measurement, + action.max_priority_label or "", + ), + ) + + +def _append_unique(items: list[str], values: tuple[str, ...]) -> None: + for value in values: + if value not in items: + items.append(value) + + +def _measurement_candidate_config_overrides(candidate_reasons: dict[str, list[str]]) -> dict[str, list[str]]: + overrides_by_label: dict[str, list[str]] = {} + for label, reasons in candidate_reasons.items(): + overrides: list[str] = [] + if any(reason.startswith("phase_timing_gap:") for reason in reasons): + _append_unique(overrides, _PHASE_TIMING_MEASUREMENT_CONFIG_OVERRIDES) + if "partial_simulator_surface_support" in reasons: + _append_unique(overrides, _PHASE_TIMING_MEASUREMENT_CONFIG_OVERRIDES) + if any(reason.startswith("memory_pressure_probe:") for reason in reasons): + _append_unique(overrides, _MEMORY_PRESSURE_MEASUREMENT_CONFIG_OVERRIDES) + if any( + reason.startswith("parallelism_axis_gap:") and "blocked" in reason.split(":", 2)[-1] for reason in reasons + ): + _append_unique(overrides, _FIT_BOUNDARY_MEASUREMENT_CONFIG_OVERRIDES) + if overrides: + overrides_by_label[label] = overrides + return overrides_by_label + + +def _scenario_needs_phase_timing_design_overrides(report: ScenarioReport) -> bool: + blockers = set(report.decision_summary.scenario_prediction_fidelity_blockers) + if blockers & _PHASE_TIMING_SCENARIO_BLOCKERS: + return True + return any( + any(reason.startswith("phase_timing_gap:") for reason in reasons) + for reasons in report.decision_summary.measurement_candidate_reasons.values() + ) + + +def _measurement_design_config_overrides( + report: ScenarioReport, + required_measurement: str, + *, + base_overrides: tuple[str, ...] = (), +) -> tuple[str, ...]: + overrides = list(base_overrides) + if required_measurement in _PHASE_TIMING_DESIGN_MEASUREMENTS and _scenario_needs_phase_timing_design_overrides( + report + ): + _append_unique(overrides, _PHASE_TIMING_MEASUREMENT_CONFIG_OVERRIDES) + return tuple(overrides) + + +def _parse_config_override_value(raw_value: str) -> Any: + value = yaml.safe_load(raw_value) + if value is None and raw_value.strip().lower() not in {"null", "none", "~"}: + return raw_value + return value + + +def _apply_config_override(raw_config: dict[str, Any], override: str) -> None: + if "=" not in override: + raise ValueError(f"config override must use dotted.path=value syntax: {override!r}") + dotted_path, raw_value = override.split("=", 1) + path = [part.strip() for part in dotted_path.split(".") if part.strip()] + if not path: + raise ValueError(f"config override has no path: {override!r}") + target = raw_config + for part in path[:-1]: + child = target.get(part) + if not isinstance(child, dict): + child = {} + target[part] = child + target = child + target[path[-1]] = _parse_config_override_value(raw_value) + + +def _measurement_config_filename(index: int, label: str) -> str: + stem = re.sub(r"[^A-Za-z0-9_.-]+", "_", label).strip("._-") + if not stem: + stem = "candidate" + return f"{index:02d}_{stem[:180]}.yaml" + + +def _set_sample_packing_sequence_len(raw_config: dict[str, Any], sequence_len: int) -> None: + if requested_simulator_surface(raw_config) == "server_forward_backward": + server = raw_config.get("server") + target = server if isinstance(server, dict) and server else raw_config + else: + target = raw_config.setdefault("data", {}) + if not isinstance(target, dict): + target = {} + raw_config["data"] = target + target["sample_packing_sequence_len"] = sequence_len + + +def _config_from_candidate( + report: ScenarioReport, + base_config: dict[str, Any], + candidate: ScenarioCandidate, + *, + overrides: list[str] | None = None, +) -> dict[str, Any]: + topology = candidate.topology + raw_config = _mutated_config( + base_config, + world_size=topology.world_size, + micro_batch_size=topology.micro_batch_size, + gradient_accumulation_steps=topology.gradient_accumulation_steps, + expert_parallel_size=topology.expert_parallel_size, + tensor_parallel_size=topology.tensor_parallel_size, + pipeline_parallel_size=topology.pipeline_parallel_size, + ulysses_parallel_size=topology.ulysses_parallel_size, + ringattn_parallel_size=topology.ringattn_parallel_size, + data_parallel_replicate_size=topology.data_parallel_replicate_size, + data_parallel_shard_size=topology.data_parallel_shard_size, + ) + if topology.sample_packing_sequence_len is not None: + _set_sample_packing_sequence_len(raw_config, topology.sample_packing_sequence_len) + _set_balanced_routing(raw_config, report.balanced_routing) + for override in overrides or []: + _apply_config_override(raw_config, override) + return raw_config + + +def materialize_measurement_candidate_configs(report: ScenarioReport) -> list[ScenarioMeasurementConfig]: + """Render runnable YAML payloads for the report's bounded measurement portfolio.""" + base_config = load_training_config(report.base_config_path) + candidates_by_label = {candidate.label: candidate for candidate in report.candidates} + rendered: list[ScenarioMeasurementConfig] = [] + for index, label in enumerate(report.decision_summary.measurement_candidate_labels, start=1): + candidate = candidates_by_label[label] + raw_config = _config_from_candidate( + report, + base_config, + candidate, + overrides=report.decision_summary.measurement_candidate_config_overrides.get(label, []), + ) + rendered.append( + ScenarioMeasurementConfig( + label=label, + filename=_measurement_config_filename(index, label), + config=raw_config, + ) + ) + return rendered + + +def _design_anchor_candidate(report: ScenarioReport) -> ScenarioCandidate | None: + return ( + report.best_risk_adjusted + or report.best_raw + or report.best_next_measurement + or report.best_promotable + or (report.candidates[0] if report.candidates else None) + ) + + +def _runtime_variant_anchor_candidate( + candidates: list[ScenarioCandidate], + runtime_mismatch_dimensions: list[str], +) -> ScenarioCandidate | None: + if not runtime_mismatch_dimensions: + return None + mismatch_flags = {f"runtime_mismatch:{dimension}" for dimension in runtime_mismatch_dimensions} + eligible = [candidate for candidate in candidates if mismatch_flags & set(candidate.risk_flags)] + if not eligible: + return None + + calibration_scope_rank = { + "exact_calibrated": 0, + "inside_measured_envelope": 1, + "outside_measured_envelope": 2, + "outside_sequence_calibration_envelope": 3, + "cross_model_analog": 4, + "no_calibration": 5, + } + + def key(candidate: ScenarioCandidate) -> tuple[int, int, bool, float, float, int, str]: + support_rank = 0 if candidate.simulator_support_status == "supported_local_non_pp" else 1 + scope_rank = calibration_scope_rank.get(candidate.calibration_scope, len(calibration_scope_rank)) + headroom = candidate.memory_headroom_gb if candidate.memory_headroom_gb is not None else float("-inf") + peak = candidate.estimated_peak_mem_gb if candidate.estimated_peak_mem_gb is not None else float("inf") + return ( + scope_rank, + support_rank, + _is_memory_blocked(candidate), + -headroom, + peak, + candidate.topology.global_batch_size, + candidate.label, + ) + + return min(eligible, key=key) + + +def _matched_behavior_point_for_candidate( + report: ScenarioReport, candidate: ScenarioCandidate | None +) -> BenchmarkBehaviorPoint | None: + if candidate is None: + return None + matched_labels = {part.strip() for part in (candidate.behavior.matched_label or "").split(",") if part.strip()} + if not matched_labels: + return None + benchmark_dirs = [] + if report.benchmark_dir is not None: + benchmark_dirs.append(Path(report.benchmark_dir)) + benchmark_dirs.extend(Path(path) for path in report.supplemental_benchmark_dirs) + for benchmark_dir in benchmark_dirs: + for point in load_benchmark_behavior_points(benchmark_dir): + if point.label in matched_labels: + return point + return None + + +def _nearby_positive_values(value: int) -> list[int]: + candidates = {value + 1, value * 2} + if value > 1: + candidates.add(value - 1) + if value > 2: + candidates.add(max(1, value // 2)) + return sorted(candidate for candidate in candidates if candidate > 0 and candidate != value) + + +def _nearby_sequence_lengths(sequence_len: int) -> list[int]: + candidates = {min(sequence_len * 2, 131_072)} + if sequence_len > 512: + candidates.add(max(512, sequence_len // 2)) + return sorted(candidate for candidate in candidates if candidate > 0 and candidate != sequence_len) + + +def _workload_design_variants(anchor: Topology) -> list[tuple[str, dict[str, int]]]: + variant_groups = [ + ( + "micro_batch_size", + "mbs", + [ + { + "micro_batch_size": value, + "gradient_accumulation_steps": anchor.gradient_accumulation_steps, + "sample_packing_sequence_len": anchor.sample_packing_sequence_len, + } + for value in _nearby_positive_values(anchor.micro_batch_size) + ], + ), + ( + "gradient_accumulation_steps", + "ga", + [ + { + "micro_batch_size": anchor.micro_batch_size, + "gradient_accumulation_steps": value, + "sample_packing_sequence_len": anchor.sample_packing_sequence_len, + } + for value in _nearby_positive_values(anchor.gradient_accumulation_steps) + ], + ), + ( + "sample_packing_sequence_len", + "seq", + [ + { + "micro_batch_size": anchor.micro_batch_size, + "gradient_accumulation_steps": anchor.gradient_accumulation_steps, + "sample_packing_sequence_len": value, + } + for value in _nearby_sequence_lengths(anchor.sample_packing_sequence_len) + ], + ), + ] + variants: list[tuple[str, dict[str, int]]] = [] + for varied_field, prefix, group in variant_groups: + if group: + variants.append((f"{prefix}{group[0][varied_field]}", group[0])) + for varied_field, prefix, group in variant_groups: + for item in group[1:]: + variants.append((f"{prefix}{item[varied_field]}", item)) + return variants + + +def _nested_config_value(raw_config: dict[str, Any], path: tuple[str, str], default: Any = None) -> Any: + section = raw_config.get(path[0]) + if not isinstance(section, dict): + return default + return section.get(path[1], default) + + +def _runtime_variant_design_value( + *, + base_config: dict[str, Any], + dimension: str, + reference_point: BenchmarkBehaviorPoint | None, +) -> Any: + if reference_point is not None: + reference_value = getattr(reference_point, dimension, None) + if reference_value is not None: + return reference_value + if dimension not in _RUNTIME_VARIANT_BOOL_DEFAULTS: + return None + path = _RUNTIME_VARIANT_CONFIG_PATHS.get(dimension) + if path is None: + return None + current_value = _nested_config_value(base_config, path, _RUNTIME_VARIANT_BOOL_DEFAULTS[dimension]) + return not bool(current_value) + + +def _runtime_variant_label_value(value: Any) -> str: + return re.sub(r"[^A-Za-z0-9_.-]+", "_", str(value).strip().lower()).strip("._-") or "value" + + +def _runtime_design_variants( + base_config: dict[str, Any], + runtime_mismatch_dimensions: list[str], + *, + reference_point: BenchmarkBehaviorPoint | None = None, +) -> list[tuple[str, tuple[str, ...]]]: + variants: list[tuple[str, tuple[str, ...]]] = [] + for dimension in runtime_mismatch_dimensions: + path = _RUNTIME_VARIANT_CONFIG_PATHS.get(dimension) + if path is None: + continue + variant_value = _runtime_variant_design_value( + base_config=base_config, + dimension=dimension, + reference_point=reference_point, + ) + if variant_value is None: + continue + override_value = str(variant_value).lower() if isinstance(variant_value, bool) else str(variant_value) + variants.append( + ( + f"runtime_{dimension}_{_runtime_variant_label_value(override_value)}", + (f"{path[0]}.{path[1]}={override_value}",), + ) + ) + return variants + + +def _design_config_from_topology( + *, + report: ScenarioReport, + base_config: dict[str, Any], + required_measurement: str, + design_kind: str, + index: int, + world_size: int, + local_world_size: int, + micro_batch_size: int, + gradient_accumulation_steps: int, + sample_packing_sequence_len: int, + expert_parallel_size: int, + tensor_parallel_size: int, + pipeline_parallel_size: int, + ulysses_parallel_size: int, + ringattn_parallel_size: int, + data_parallel_replicate_size: int | None = None, + data_parallel_shard_size: int | None = None, + config_overrides: tuple[str, ...] = (), +) -> ScenarioMeasurementConfig | None: + try: + raw_config = _mutated_config( + base_config, + world_size=world_size, + micro_batch_size=micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + expert_parallel_size=expert_parallel_size, + tensor_parallel_size=tensor_parallel_size, + pipeline_parallel_size=pipeline_parallel_size, + ulysses_parallel_size=ulysses_parallel_size, + ringattn_parallel_size=ringattn_parallel_size, + data_parallel_replicate_size=data_parallel_replicate_size, + data_parallel_shard_size=data_parallel_shard_size, + ) + _set_sample_packing_sequence_len(raw_config, sample_packing_sequence_len) + _set_balanced_routing(raw_config, report.balanced_routing) + for override in config_overrides: + _apply_config_override(raw_config, override) + topology = resolve_topology(raw_config, world_size=world_size, local_world_size=local_world_size) + except (TypeError, ValueError): + return None + if topology.ep_fsdp_size is None: + return None + + label = f"design:{required_measurement}:{design_kind}:{_topology_label(topology)}" + return ScenarioMeasurementConfig( + label=label, + filename=_measurement_config_filename(index, label), + config=raw_config, + ) + + +def _memory_pressure_design_variants(anchor: Topology) -> list[tuple[str, dict[str, int]]]: + variants: list[tuple[str, dict[str, int]]] = [] + reduced_mbs = max(1, anchor.micro_batch_size // 2) + if reduced_mbs == anchor.micro_batch_size and anchor.micro_batch_size > 1: + reduced_mbs = anchor.micro_batch_size - 1 + reduced_seq = None + if anchor.sample_packing_sequence_len and anchor.sample_packing_sequence_len > 512: + reduced_seq = max(512, anchor.sample_packing_sequence_len // 2) + + if reduced_mbs != anchor.micro_batch_size: + variants.append( + ( + f"fit_mbs{reduced_mbs}", + { + "micro_batch_size": reduced_mbs, + "gradient_accumulation_steps": anchor.gradient_accumulation_steps, + "sample_packing_sequence_len": anchor.sample_packing_sequence_len, + }, + ) + ) + if reduced_seq is not None and reduced_seq != anchor.sample_packing_sequence_len: + variants.append( + ( + f"fit_seq{reduced_seq}", + { + "micro_batch_size": anchor.micro_batch_size, + "gradient_accumulation_steps": anchor.gradient_accumulation_steps, + "sample_packing_sequence_len": reduced_seq, + }, + ) + ) + if ( + reduced_mbs != anchor.micro_batch_size + and reduced_seq is not None + and reduced_seq != anchor.sample_packing_sequence_len + ): + variants.append( + ( + f"fit_mbs{reduced_mbs}_seq{reduced_seq}", + { + "micro_batch_size": reduced_mbs, + "gradient_accumulation_steps": anchor.gradient_accumulation_steps, + "sample_packing_sequence_len": reduced_seq, + }, + ) + ) + if not variants and anchor.sample_packing_sequence_len is not None: + variants.append( + ( + "profile_current_shape", + { + "micro_batch_size": anchor.micro_batch_size, + "gradient_accumulation_steps": anchor.gradient_accumulation_steps, + "sample_packing_sequence_len": anchor.sample_packing_sequence_len, + }, + ) + ) + return variants + + +def _append_design_config( + rendered: list[ScenarioMeasurementConfig], + seen: set[tuple[str, str]], + design: ScenarioMeasurementConfig | None, +) -> None: + if design is None: + return + required_measurement = design.label.split(":", 3)[1] if ":" in design.label else design.label + key = (required_measurement, yaml.safe_dump(design.config, sort_keys=True)) + if key in seen: + return + seen.add(key) + rendered.append(design) + + +def _design_kind_slug(value: str) -> str: + slug = re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("._-") + return slug[:80] or "candidate" + + +def _fit_boundary_candidate_design_kind(design_kind: str, candidate: ScenarioCandidate) -> str: + if ( + design_kind == "fit_boundary" + and candidate.topology.ulysses_parallel_size > 1 + and re.search(r"(?:^|[-_:])u(?:[2-9]|\d{2,})(?:$|[-_:])", candidate.label) + ): + return "fit_boundary_ulysses" + return design_kind + + +def _action_candidate_axis_family(action: ScenarioValidationAction, candidate: ScenarioCandidate) -> str | None: + action_axes = set(action.parallelism_axis_names) + topology = candidate.topology + if "ulysses" in action_axes and topology.ulysses_parallel_size > 1: + return "ulysses" + if "ringattn" in action_axes and topology.ringattn_parallel_size > 1: + return "ring" + if {"expert_parallel", "ep_fsdp"} & action_axes: + return "ep" + if "tensor_parallel" in action_axes and topology.tensor_parallel_size > 1: + return "tp" + if "pipeline_parallel" in action_axes and topology.pipeline_parallel_size > 1: + return "pp" + if "dp_replicate" in action_axes and topology.data_parallel_replicate_size > 1: + return "dp_replicate" + if {"world_size", "dp_shard"} & action_axes: + return "world" + return None + + +def _action_candidate_design_kind( + design_kind: str, + action: ScenarioValidationAction, + candidate: ScenarioCandidate, +) -> str: + axis_family = _action_candidate_axis_family(action, candidate) + if axis_family is None: + return design_kind + return f"{design_kind}_{axis_family}" + + +_GAP_ACTION_DESIGN_MEASUREMENTS = { + "add_fit_failure_boundary_near_blocked_axes": ( + ("same_workload_same_runtime_axis_pair_with_fit_failure_boundary", "fit_boundary"), + ("launch_memory_pressure_fit_probe_with_memory_profiling", "memory_fit_probe"), + ), + "add_memory_pressure_fit_probe_or_reduce_batch": ( + ("launch_memory_pressure_fit_probe_with_memory_profiling", "memory_fit_probe"), + ), + "add_partial_surface_support_or_direct_probe": (("add_simulator_surface_support_or_run_probe", "surface_probe"),), +} + + +_DIRECT_ACTION_DESIGN_MEASUREMENTS = { + "add_simulator_surface_support_or_run_probe": "action_surface_probe", + "launch_memory_pressure_fit_probe_with_memory_profiling": "action_memory_fit_probe", + "same_workload_same_runtime_axis_pair": "action_axis_pair", + "same_workload_same_runtime_axis_pair_with_fit_failure_boundary": "action_boundary_axis_pair", + "measure_parallelism_tradeoff_candidates": "action_parallelism_tradeoff", + "measure_target_cross_model_interval_top_candidates": "action_cross_model_interval", + "measure_throughput_efficiency_frontier_candidates": "action_efficiency_frontier", + "remeasure_best_next_candidate_same_workload": "action_remeasure_best", + "rerun_current_winner_with_required_instrumentation": "action_winner_instrumented", + "rerun_with_phase_timing": "action_phase_timing", +} + + +def _action_backed_design_config( + *, + report: ScenarioReport, + base_config: dict[str, Any], + required_measurement: str, + design_kind: str, + index: int, + candidate: ScenarioCandidate, +) -> ScenarioMeasurementConfig: + overrides = list(report.decision_summary.measurement_candidate_config_overrides.get(candidate.label, [])) + if required_measurement == "add_fit_failure_boundary_near_blocked_axes": + _append_unique(overrides, _MEMORY_PRESSURE_MEASUREMENT_CONFIG_OVERRIDES) + overrides = list( + _measurement_design_config_overrides( + report, + required_measurement, + base_overrides=tuple(overrides), + ) + ) + raw_config = _config_from_candidate( + report, + base_config, + candidate, + overrides=overrides, + ) + topology = candidate.topology + design_kind = _fit_boundary_candidate_design_kind(design_kind, candidate) + label = ( + f"design:{required_measurement}:{design_kind}_{index:02d}_{_design_kind_slug(candidate.label)}:" + f"{_topology_label(topology)}" + ) + return ScenarioMeasurementConfig( + label=label, + filename=_measurement_config_filename(index, label), + config=raw_config, + ) + + +def _candidate_replay_design_config( + *, + report: ScenarioReport, + base_config: dict[str, Any], + required_measurement: str, + index: int, + candidate: ScenarioCandidate, +) -> ScenarioMeasurementConfig: + overrides = _measurement_design_config_overrides( + report, + required_measurement, + base_overrides=tuple(report.decision_summary.measurement_candidate_config_overrides.get(candidate.label, [])), + ) + raw_config = _config_from_candidate( + report, + base_config, + candidate, + overrides=list(overrides), + ) + topology = candidate.topology + label = ( + f"design:{required_measurement}:scored_replay_{index:02d}_{_design_kind_slug(candidate.label)}:" + f"{_topology_label(topology)}" + ) + return ScenarioMeasurementConfig( + label=label, + filename=_measurement_config_filename(index, label), + config=raw_config, + ) + + +def _same_workload_ga_for_topology( + *, + target_global_batch_size: int, + world_size: int, + micro_batch_size: int, + tensor_parallel_size: int, + pipeline_parallel_size: int, + ulysses_parallel_size: int, + ringattn_parallel_size: int, +) -> int | None: + non_dp = tensor_parallel_size * pipeline_parallel_size * ulysses_parallel_size * ringattn_parallel_size + if non_dp <= 0 or world_size % non_dp: + return None + data_parallel_size = world_size // non_dp + divisor = micro_batch_size * data_parallel_size + if divisor <= 0 or target_global_batch_size % divisor: + return None + return target_global_batch_size // divisor + + +def _valid_dp_split_for_size( + data_parallel_size: int, + *, + preferred_replicate_size: int | None = None, + preferred_shard_size: int | None = None, +) -> tuple[int, int]: + if data_parallel_size <= 0: + raise ValueError("data_parallel_size must be positive") + if ( + preferred_replicate_size is not None + and preferred_shard_size is not None + and preferred_replicate_size > 0 + and preferred_shard_size > 0 + and preferred_replicate_size * preferred_shard_size == data_parallel_size + ): + return preferred_replicate_size, preferred_shard_size + if ( + preferred_replicate_size is not None + and preferred_replicate_size > 0 + and data_parallel_size % preferred_replicate_size == 0 + ): + return preferred_replicate_size, data_parallel_size // preferred_replicate_size + if preferred_shard_size is not None and preferred_shard_size > 0 and data_parallel_size % preferred_shard_size == 0: + return data_parallel_size // preferred_shard_size, preferred_shard_size + return 1, data_parallel_size + + +def _topology_values_with_dp_split( + values: dict[str, int], + *, + preferred_replicate_size: int, + preferred_shard_size: int, +) -> dict[str, int] | None: + non_dp_size = ( + values["tensor_parallel_size"] + * values["pipeline_parallel_size"] + * values["ulysses_parallel_size"] + * values["ringattn_parallel_size"] + ) + if non_dp_size <= 0 or values["world_size"] % non_dp_size: + return None + data_parallel_size = values["world_size"] // non_dp_size + replicate_size, shard_size = _valid_dp_split_for_size( + data_parallel_size, + preferred_replicate_size=preferred_replicate_size, + preferred_shard_size=preferred_shard_size, + ) + return { + **values, + "data_parallel_replicate_size": replicate_size, + "data_parallel_shard_size": shard_size, + } + + +def _dp_replicate_design_values(anchor: Topology, base: dict[str, int]) -> list[dict[str, int]]: + variants: list[dict[str, int]] = [] + candidates = [1 if anchor.data_parallel_replicate_size > 1 else 2] + candidates.extend([2, 4, 8, anchor.node_count, anchor.data_parallel_size]) + for replicate_size in candidates: + if replicate_size <= 0 or replicate_size == anchor.data_parallel_replicate_size: + continue + if anchor.data_parallel_size % replicate_size: + continue + shard_size = anchor.data_parallel_size // replicate_size + updated = { + **base, + "data_parallel_replicate_size": replicate_size, + "data_parallel_shard_size": shard_size, + } + if updated not in variants: + variants.append(updated) + return variants + + +def _parallelism_design_values( + anchor: Topology, + metadata: ModelMetadata, +) -> list[tuple[str, dict[str, int]]]: + base_axes = { + "world_size": anchor.world_size, + "expert_parallel_size": anchor.expert_parallel_size, + "tensor_parallel_size": anchor.tensor_parallel_size, + "pipeline_parallel_size": anchor.pipeline_parallel_size, + "ulysses_parallel_size": anchor.ulysses_parallel_size, + "ringattn_parallel_size": anchor.ringattn_parallel_size, + } + base = _topology_values_with_dp_split( + base_axes, + preferred_replicate_size=anchor.data_parallel_replicate_size, + preferred_shard_size=anchor.data_parallel_shard_size, + ) + if base is None: + return [] + variants: list[tuple[str, dict[str, int]]] = [] + + world_size_values = [anchor.world_size * 2] + if anchor.world_size > anchor.local_world_size and anchor.world_size % 2 == 0: + world_size_values.append(anchor.world_size // 2) + for value in world_size_values: + if value > 0 and value != anchor.world_size: + updated = _topology_values_with_dp_split( + {**base_axes, "world_size": value}, + preferred_replicate_size=anchor.data_parallel_replicate_size, + preferred_shard_size=anchor.data_parallel_shard_size, + ) + if updated is not None: + variants.append(("world", updated)) + + for updated in _dp_replicate_design_values(anchor, base): + variants.append(("dp_replicate", updated)) + + for field, axis, values in ( + ("expert_parallel_size", "ep", _auto_ep_sizes(anchor)), + ("tensor_parallel_size", "tp", _auto_tensor_parallel_sizes(anchor, metadata)), + ("pipeline_parallel_size", "pp", _auto_pipeline_parallel_sizes(anchor, metadata)), + ("ulysses_parallel_size", "ulysses", _auto_ulysses_parallel_sizes(anchor, metadata)), + ("ringattn_parallel_size", "ring", _auto_ringattn_parallel_sizes(anchor)), + ): + for value in values: + if value == base_axes[field]: + continue + updated = _topology_values_with_dp_split( + {**base_axes, field: value}, + preferred_replicate_size=anchor.data_parallel_replicate_size, + preferred_shard_size=anchor.data_parallel_shard_size, + ) + if updated is not None: + variants.append((axis, updated)) + return variants + + +_SCENARIO_AXIS_DESIGN_PREFERRED_VARIANTS = { + "world_size": ("world",), + "dp_replicate": ("dp_replicate", "world"), + "dp_shard": ("world", "ulysses", "ring"), + "tensor_parallel": ("tp",), + "pipeline_parallel": ("pp",), + "expert_parallel": ("ep",), + "ep_fsdp": ("ep", "world"), + "ulysses": ("ulysses",), + "ringattn": ("ring",), +} + +_SCENARIO_PARALLELISM_AXIS_DESIGN_MEASUREMENTS = { + "add_same_workload_same_runtime_parallelism_axis_variants", + "add_same_workload_same_runtime_axis_pairs", +} + +_SCENARIO_PARALLELISM_AXIS_COMPATIBLE_MEASUREMENTS = { + *_SCENARIO_PARALLELISM_AXIS_DESIGN_MEASUREMENTS, + "add_workload_and_parallelism_variants", +} + +_SCENARIO_PARALLELISM_AXIS_CONFIG_BUDGET = 6 +_SCENARIO_FIT_BOUNDARY_AXIS_CONFIG_BUDGET = 8 +_SCENARIO_COMBINED_WORKLOAD_PARALLELISM_CONFIG_BUDGET = 10 + + +def _scenario_gap_parallelism_axes(gap: ScenarioCaptureGap) -> list[str]: + return _unique_in_order( + [ + *gap.missing_parallelism_axis_names, + *gap.blocked_parallelism_axis_names, + *gap.confounded_parallelism_axis_names, + *gap.unscored_parallelism_axis_names, + ] + ) + + +def _scenario_gap_preferred_design_axes(gap: ScenarioCaptureGap) -> list[str]: + axis_order: list[str] = [] + for gap_axis in _scenario_gap_parallelism_axes(gap): + for axis in _SCENARIO_AXIS_DESIGN_PREFERRED_VARIANTS.get(gap_axis, ()): + if axis not in axis_order: + axis_order.append(axis) + return axis_order + + +def _scenario_gap_design_config_budget(gap: ScenarioCaptureGap, default_budget: int) -> int: + axis_family_budget = max(_SCENARIO_PARALLELISM_AXIS_CONFIG_BUDGET, len(_scenario_gap_preferred_design_axes(gap))) + if gap.required_measurement == "add_workload_and_parallelism_variants": + return max( + default_budget, + _SCENARIO_COMBINED_WORKLOAD_PARALLELISM_CONFIG_BUDGET, + default_budget + axis_family_budget, + ) + if gap.required_measurement == "add_fit_failure_boundary_near_blocked_axes": + return max(default_budget, _SCENARIO_FIT_BOUNDARY_AXIS_CONFIG_BUDGET, default_budget + axis_family_budget) + if gap.required_measurement in _SCENARIO_PARALLELISM_AXIS_DESIGN_MEASUREMENTS: + return max(default_budget, axis_family_budget) + return default_budget + + +def _parallelism_design_values_apply_to_workload( + axis: str, + values: dict[str, int], + workload_values: dict[str, int], +) -> bool: + seq_len = workload_values["sample_packing_sequence_len"] or 0 + if axis == "ulysses" and values["ulysses_parallel_size"] > 1 and seq_len < _MIN_ULYSSES_SEQUENCE_LEN: + return False + if axis == "ring" and values["ringattn_parallel_size"] > 1 and seq_len < _MIN_RINGATTN_SEQUENCE_LEN: + return False + return True + + +def _prioritized_parallelism_design_values_for_gap( + anchor: Topology, + metadata: ModelMetadata, + gap: ScenarioCaptureGap, +) -> list[tuple[str, dict[str, int]]]: + values = _parallelism_design_values(anchor, metadata) + grouped: dict[str, list[dict[str, int]]] = {} + for axis, design_values in values: + grouped.setdefault(axis, []).append(design_values) + + axis_order: list[str] = [] + for axis in _scenario_gap_preferred_design_axes(gap): + if axis in grouped and axis not in axis_order: + axis_order.append(axis) + for axis in grouped: + if axis not in axis_order: + axis_order.append(axis) + + ordered: list[tuple[str, dict[str, int]]] = [] + max_group_size = max((len(grouped[axis]) for axis in axis_order), default=0) + for offset in range(max_group_size): + for axis in axis_order: + variants = grouped[axis] + if offset < len(variants): + ordered.append((axis, variants[offset])) + return ordered + + +def materialize_measurement_design_configs( + report: ScenarioReport, + *, + max_configs_per_measurement: int = 4, +) -> list[ScenarioMeasurementConfig]: + """Render bounded YAML design rows for scenario-capture gaps not backed by existing candidates.""" + anchor_candidate = _design_anchor_candidate(report) + if anchor_candidate is None: + return [] + + base_config = load_training_config(report.base_config_path) + anchor = anchor_candidate.topology + metadata = resolve_model_metadata(base_config) + candidates_by_label = {candidate.label: candidate for candidate in report.candidates} + validation_actions_by_measurement: dict[str, list[ScenarioValidationAction]] = {} + for action in report.decision_summary.validation_actions: + validation_actions_by_measurement.setdefault(action.required_measurement, []).append(action) + rendered: list[ScenarioMeasurementConfig] = [] + seen: set[tuple[str, str]] = set() + index = 1 + + def add_design(design: ScenarioMeasurementConfig | None) -> None: + nonlocal index + before = len(rendered) + _append_design_config(rendered, seen, design) + if len(rendered) > before: + index += 1 + + for gap in sorted( + report.decision_summary.scenario_capture_gaps, key=lambda item: (-item.priority, item.gap_status) + ): + count_for_measurement = 0 + design_config_budget = _scenario_gap_design_config_budget(gap, max_configs_per_measurement) + + def add_limited(design: ScenarioMeasurementConfig | None) -> bool: + nonlocal count_for_measurement + if count_for_measurement >= design_config_budget: + return False + before = len(rendered) + add_design(design) + if len(rendered) > before: + count_for_measurement += 1 + return True + return False + + for action_measurement, design_kind in _GAP_ACTION_DESIGN_MEASUREMENTS.get(gap.required_measurement, ()): + for action in validation_actions_by_measurement.get(action_measurement, []): + for label in action.candidate_labels: + if count_for_measurement >= max_configs_per_measurement: + break + candidate = candidates_by_label.get(label) + if candidate is None: + continue + add_limited( + _action_backed_design_config( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + design_kind=design_kind, + index=index, + candidate=candidate, + ) + ) + + if gap.required_measurement == "add_fit_failure_boundary_near_blocked_axes": + added_fit_boundary_axis_families: set[str] = set() + for axis, values in _prioritized_parallelism_design_values_for_gap(anchor, metadata, gap): + if axis in added_fit_boundary_axis_families: + continue + gradient_accumulation_steps = _same_workload_ga_for_topology( + target_global_batch_size=anchor.global_batch_size, + world_size=values["world_size"], + micro_batch_size=anchor.micro_batch_size, + tensor_parallel_size=values["tensor_parallel_size"], + pipeline_parallel_size=values["pipeline_parallel_size"], + ulysses_parallel_size=values["ulysses_parallel_size"], + ringattn_parallel_size=values["ringattn_parallel_size"], + ) + if gradient_accumulation_steps is None: + continue + if add_limited( + _design_config_from_topology( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + design_kind=f"fit_boundary_{axis}", + index=index, + world_size=values["world_size"], + local_world_size=min(anchor.local_world_size, values["world_size"]), + micro_batch_size=anchor.micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + sample_packing_sequence_len=anchor.sample_packing_sequence_len, + expert_parallel_size=values["expert_parallel_size"], + tensor_parallel_size=values["tensor_parallel_size"], + pipeline_parallel_size=values["pipeline_parallel_size"], + ulysses_parallel_size=values["ulysses_parallel_size"], + ringattn_parallel_size=values["ringattn_parallel_size"], + data_parallel_replicate_size=values["data_parallel_replicate_size"], + data_parallel_shard_size=values["data_parallel_shard_size"], + config_overrides=_measurement_design_config_overrides( + report, + gap.required_measurement, + base_overrides=_MEMORY_PRESSURE_MEASUREMENT_CONFIG_OVERRIDES, + ), + ) + ): + added_fit_boundary_axis_families.add(axis) + + if gap.required_measurement == "add_memory_pressure_fit_probe_or_reduce_batch" and count_for_measurement == 0: + for design_kind, values in _memory_pressure_design_variants(anchor): + add_limited( + _design_config_from_topology( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + design_kind=design_kind, + index=index, + world_size=anchor.world_size, + local_world_size=anchor.local_world_size, + micro_batch_size=values["micro_batch_size"], + gradient_accumulation_steps=values["gradient_accumulation_steps"], + sample_packing_sequence_len=values["sample_packing_sequence_len"], + expert_parallel_size=anchor.expert_parallel_size, + tensor_parallel_size=anchor.tensor_parallel_size, + pipeline_parallel_size=anchor.pipeline_parallel_size, + ulysses_parallel_size=anchor.ulysses_parallel_size, + ringattn_parallel_size=anchor.ringattn_parallel_size, + data_parallel_replicate_size=anchor.data_parallel_replicate_size, + data_parallel_shard_size=anchor.data_parallel_shard_size, + config_overrides=_measurement_design_config_overrides( + report, + gap.required_measurement, + base_overrides=_MEMORY_PRESSURE_MEASUREMENT_CONFIG_OVERRIDES, + ), + ) + ) + + if ( + gap.required_measurement + in { + "add_scored_measurements_for_existing_sweep", + "score_unscored_capture_candidates", + } + and count_for_measurement == 0 + ): + for label in report.decision_summary.measurement_candidate_labels: + candidate = candidates_by_label.get(label) + if candidate is None or candidate.simulator_support_status.startswith("unsupported_"): + continue + add_limited( + _candidate_replay_design_config( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + index=index, + candidate=candidate, + ) + ) + + if gap.required_measurement in { + "add_same_parallelism_runtime_workload_variants", + "add_workload_and_parallelism_variants", + }: + workload_only_budget = ( + max_configs_per_measurement + if gap.required_measurement == "add_workload_and_parallelism_variants" + else design_config_budget + ) + for design_kind, values in _workload_design_variants(anchor): + if count_for_measurement >= workload_only_budget: + break + add_limited( + _design_config_from_topology( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + design_kind=design_kind, + index=index, + world_size=anchor.world_size, + local_world_size=anchor.local_world_size, + micro_batch_size=values["micro_batch_size"], + gradient_accumulation_steps=values["gradient_accumulation_steps"], + sample_packing_sequence_len=values["sample_packing_sequence_len"], + expert_parallel_size=anchor.expert_parallel_size, + tensor_parallel_size=anchor.tensor_parallel_size, + pipeline_parallel_size=anchor.pipeline_parallel_size, + ulysses_parallel_size=anchor.ulysses_parallel_size, + ringattn_parallel_size=anchor.ringattn_parallel_size, + data_parallel_replicate_size=anchor.data_parallel_replicate_size, + data_parallel_shard_size=anchor.data_parallel_shard_size, + config_overrides=_measurement_design_config_overrides(report, gap.required_measurement), + ) + ) + + if gap.required_measurement == "add_workload_and_parallelism_variants": + workload_variants = _workload_design_variants(anchor) + added_combined_axis_families: set[str] = set() + for axis, values in _prioritized_parallelism_design_values_for_gap(anchor, metadata, gap): + if axis in added_combined_axis_families: + continue + for workload_design_kind, workload_values in workload_variants: + if not _parallelism_design_values_apply_to_workload(axis, values, workload_values): + continue + if add_limited( + _design_config_from_topology( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + design_kind=f"combined_{workload_design_kind}_{axis}", + index=index, + world_size=values["world_size"], + local_world_size=min(anchor.local_world_size, values["world_size"]), + micro_batch_size=workload_values["micro_batch_size"], + gradient_accumulation_steps=workload_values["gradient_accumulation_steps"], + sample_packing_sequence_len=workload_values["sample_packing_sequence_len"], + expert_parallel_size=values["expert_parallel_size"], + tensor_parallel_size=values["tensor_parallel_size"], + pipeline_parallel_size=values["pipeline_parallel_size"], + ulysses_parallel_size=values["ulysses_parallel_size"], + ringattn_parallel_size=values["ringattn_parallel_size"], + data_parallel_replicate_size=values["data_parallel_replicate_size"], + data_parallel_shard_size=values["data_parallel_shard_size"], + config_overrides=_measurement_design_config_overrides(report, gap.required_measurement), + ) + ): + added_combined_axis_families.add(axis) + break + + if gap.required_measurement == "add_same_parallelism_workload_runtime_variants": + runtime_anchor = _runtime_variant_anchor_candidate(report.candidates, gap.runtime_mismatch_dimensions) + runtime_anchor_topology = runtime_anchor.topology if runtime_anchor is not None else anchor + reference_point = _matched_behavior_point_for_candidate(report, runtime_anchor) + for design_kind, config_overrides in _runtime_design_variants( + base_config, + gap.runtime_mismatch_dimensions, + reference_point=reference_point, + ): + add_limited( + _design_config_from_topology( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + design_kind=design_kind, + index=index, + world_size=runtime_anchor_topology.world_size, + local_world_size=runtime_anchor_topology.local_world_size, + micro_batch_size=runtime_anchor_topology.micro_batch_size, + gradient_accumulation_steps=runtime_anchor_topology.gradient_accumulation_steps, + sample_packing_sequence_len=runtime_anchor_topology.sample_packing_sequence_len, + expert_parallel_size=runtime_anchor_topology.expert_parallel_size, + tensor_parallel_size=runtime_anchor_topology.tensor_parallel_size, + pipeline_parallel_size=runtime_anchor_topology.pipeline_parallel_size, + ulysses_parallel_size=runtime_anchor_topology.ulysses_parallel_size, + ringattn_parallel_size=runtime_anchor_topology.ringattn_parallel_size, + data_parallel_replicate_size=runtime_anchor_topology.data_parallel_replicate_size, + data_parallel_shard_size=runtime_anchor_topology.data_parallel_shard_size, + config_overrides=_measurement_design_config_overrides( + report, + gap.required_measurement, + base_overrides=config_overrides, + ), + ) + ) + + if gap.required_measurement in _SCENARIO_PARALLELISM_AXIS_COMPATIBLE_MEASUREMENTS: + added_axis_families: set[str] = set() + for axis, values in _prioritized_parallelism_design_values_for_gap(anchor, metadata, gap): + if axis in added_axis_families: + continue + gradient_accumulation_steps = _same_workload_ga_for_topology( + target_global_batch_size=anchor.global_batch_size, + world_size=values["world_size"], + micro_batch_size=anchor.micro_batch_size, + tensor_parallel_size=values["tensor_parallel_size"], + pipeline_parallel_size=values["pipeline_parallel_size"], + ulysses_parallel_size=values["ulysses_parallel_size"], + ringattn_parallel_size=values["ringattn_parallel_size"], + ) + if gradient_accumulation_steps is None: + continue + if add_limited( + _design_config_from_topology( + report=report, + base_config=base_config, + required_measurement=gap.required_measurement, + design_kind=axis, + index=index, + world_size=values["world_size"], + local_world_size=min(anchor.local_world_size, values["world_size"]), + micro_batch_size=anchor.micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_steps, + sample_packing_sequence_len=anchor.sample_packing_sequence_len, + expert_parallel_size=values["expert_parallel_size"], + tensor_parallel_size=values["tensor_parallel_size"], + pipeline_parallel_size=values["pipeline_parallel_size"], + ulysses_parallel_size=values["ulysses_parallel_size"], + ringattn_parallel_size=values["ringattn_parallel_size"], + data_parallel_replicate_size=values["data_parallel_replicate_size"], + data_parallel_shard_size=values["data_parallel_shard_size"], + config_overrides=_measurement_design_config_overrides(report, gap.required_measurement), + ) + ): + added_axis_families.add(axis) + + direct_action_counts_by_measurement = Counter( + design.label.split(":", 3)[1] for design in rendered if design.label.startswith("design:") + ) + for action in report.decision_summary.validation_actions: + direct_design_kind = _DIRECT_ACTION_DESIGN_MEASUREMENTS.get(action.required_measurement) + if direct_design_kind is None: + continue + for label in action.candidate_labels: + if direct_action_counts_by_measurement[action.required_measurement] >= max_configs_per_measurement: + break + candidate = candidates_by_label.get(label) + if candidate is None: + continue + added_before = len(rendered) + add_design( + _action_backed_design_config( + report=report, + base_config=base_config, + required_measurement=action.required_measurement, + design_kind=_action_candidate_design_kind(direct_design_kind, action, candidate), + index=index, + candidate=candidate, + ) + ) + if len(rendered) > added_before: + direct_action_counts_by_measurement[action.required_measurement] += 1 + return rendered + + +def materialize_measurement_configs(report: ScenarioReport) -> list[ScenarioMeasurementConfig]: + """Render candidate and scenario-capture design YAML payloads.""" + candidate_configs = materialize_measurement_candidate_configs(report) + design_configs = materialize_measurement_design_configs(report) + return [ + *candidate_configs, + *[ + ScenarioMeasurementConfig( + label=item.label, + filename=f"design_{item.filename}", + config=item.config, + ) + for item in design_configs + ], + ] + + +def write_measurement_candidate_configs( + report: ScenarioReport, output_dir: str | Path +) -> list[ScenarioMeasurementConfig]: + """Write the report's measurement portfolio as YAML configs and return the rendered payloads.""" + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + rendered = materialize_measurement_candidate_configs(report) + for item in rendered: + (output_path / item.filename).write_text( + yaml.safe_dump(runtime_training_config(item.config), sort_keys=False), + encoding="utf-8", + ) + return rendered + + +def write_measurement_configs(report: ScenarioReport, output_dir: str | Path) -> list[ScenarioMeasurementConfig]: + """Write candidate and design measurement configs as YAML files.""" + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + rendered = materialize_measurement_configs(report) + for item in rendered: + (output_path / item.filename).write_text( + yaml.safe_dump(runtime_training_config(item.config), sort_keys=False), + encoding="utf-8", + ) + return rendered + + +def _attach_measurement_design_summary(report: ScenarioReport) -> ScenarioReport: + design_configs = materialize_measurement_design_configs(report) + labels = [item.label for item in design_configs] + filenames = [f"design_{item.filename}" for item in design_configs] + scenario_readiness = replace( + report.decision_summary.scenario_readiness, + measurement_design_config_count=len(design_configs), + measurement_design_config_labels=labels, + measurement_design_config_filenames=filenames, + ) + decision_summary = replace( + report.decision_summary, + measurement_design_config_count=len(design_configs), + measurement_design_config_labels=labels, + measurement_design_config_filenames=filenames, + scenario_readiness=scenario_readiness, + ) + return replace(report, decision_summary=decision_summary) + + +def _cross_model_analog_candidates(candidates: list[ScenarioCandidate]) -> list[ScenarioCandidate]: + return [ + candidate + for candidate in candidates + if candidate.prediction_confidence == "cross_model_extrapolated" or "cross_model_analog" in candidate.risk_flags + ] + + +_CROSS_MODEL_ANALOG_FACTOR_FIELDS = ( + ("active_param_scale", "cross_model_active_param_scale"), + ("sequence_length_factor", "cross_model_sequence_length_factor"), + ("parallelism_factor", "cross_model_parallelism_factor"), + ("memory_factor", "cross_model_memory_factor"), +) + + +def _scored_cross_model_analog_candidates(candidates: list[ScenarioCandidate]) -> list[ScenarioCandidate]: + return [ + candidate + for candidate in _cross_model_analog_candidates(candidates) + if candidate.score_tokens_per_sec is not None + ] + + +def _cross_model_analog_factor_key(candidate: ScenarioCandidate) -> tuple[float | None, ...]: + return tuple(getattr(candidate.behavior, field_name) for _, field_name in _CROSS_MODEL_ANALOG_FACTOR_FIELDS) + + +def _cross_model_analog_unique_factor_count(candidates: list[ScenarioCandidate]) -> int: + return len( + {_cross_model_analog_factor_key(candidate) for candidate in _scored_cross_model_analog_candidates(candidates)} + ) + + +def _cross_model_analog_factor_status(candidates: list[ScenarioCandidate]) -> str: + if not _cross_model_analog_candidates(candidates): + return "not_used" + scored = _scored_cross_model_analog_candidates(candidates) + if not scored: + return "no_scored_cross_model_factors" + if _cross_model_analog_unique_factor_count(scored) == 1: + return "single_cross_model_factor" + return "multiple_cross_model_factors" + + +def _cross_model_analog_factor_ranges(candidates: list[ScenarioCandidate]) -> dict[str, list[float]]: + ranges: dict[str, list[float]] = {} + scored = _scored_cross_model_analog_candidates(candidates) + for range_name, field_name in _CROSS_MODEL_ANALOG_FACTOR_FIELDS: + values = [ + getattr(candidate.behavior, field_name) + for candidate in scored + if getattr(candidate.behavior, field_name) is not None + ] + if values: + ranges[range_name] = [round(min(values), 3), round(max(values), 3)] + return dict(sorted(ranges.items())) + + +def _cross_model_analog_prediction_interval_top(candidates: list[ScenarioCandidate]) -> list[ScenarioCandidate]: + scored = _scored_cross_model_analog_candidates(candidates) + if not scored: + return [] + predicted_order = sorted( + scored, + key=lambda candidate: (candidate.score_tokens_per_sec or float("-inf"), candidate.label), + reverse=True, + ) + predicted_best = predicted_order[0] + best_lower = predicted_best.prediction_interval_lower_tokens_per_sec + best_upper = predicted_best.prediction_interval_upper_tokens_per_sec + if best_lower is None or best_upper is None: + return [] + return [ + candidate + for candidate in predicted_order + if candidate.prediction_interval_lower_tokens_per_sec is not None + and candidate.prediction_interval_upper_tokens_per_sec is not None + and candidate.prediction_interval_lower_tokens_per_sec <= best_upper + and candidate.prediction_interval_upper_tokens_per_sec >= best_lower + ] + + +def _prediction_interval_selectivity_status( + *, + scored_candidate_count: int, + prediction_interval_top_count: int, +) -> str: + if scored_candidate_count == 0: + return "no_scored_cross_model_candidates" + if prediction_interval_top_count == 0: + return "no_prediction_interval_top" + if prediction_interval_top_count == 1: + return "selective_prediction_interval_top" + if prediction_interval_top_count >= scored_candidate_count: + return "nonselective_prediction_interval_top" + return "partial_prediction_interval_top" + + +def _cross_model_analog_support( + candidates: list[ScenarioCandidate], +) -> tuple[str, int, int, int, int, int, int]: + analog_candidates = _cross_model_analog_candidates(candidates) + if not analog_candidates: + return "not_used", 0, 0, 0, 0, 0, 0 + + scored = [candidate for candidate in analog_candidates if candidate.score_tokens_per_sec is not None] + unique_predictions = {round(candidate.score_tokens_per_sec or 0.0, 3) for candidate in scored} + unique_matched_labels = { + candidate.behavior.matched_label for candidate in scored if candidate.behavior.matched_label is not None + } + unique_target_strategies = {_parallelism_strategy_key(candidate) for candidate in scored} + unique_target_runtime_signatures = {candidate.target_runtime_signature for candidate in scored} + if not scored: + status = "no_scored_cross_model_candidates" + elif len(scored) == 1: + status = "single_scored_cross_model_candidate" + elif len(unique_matched_labels) == 1 and len(unique_target_strategies) > 1: + status = "single_reference_cannot_rank_parallelism_variants" + elif len(unique_matched_labels) == 1 and len(unique_target_runtime_signatures) > 1: + status = "single_reference_cannot_rank_runtime_variants" + elif len(unique_predictions) == 1 and len(unique_matched_labels) == 1: + status = "single_reference_tied_tradeoff" + elif len(unique_predictions) == 1: + status = "tied_cross_model_frontier" + else: + status = "scored_cross_model_frontier" + return ( + status, + len(analog_candidates), + len(scored), + len(unique_predictions), + len(unique_matched_labels), + len(unique_target_strategies), + len(unique_target_runtime_signatures), + ) + + +def _model_generalization_support( + candidates: list[ScenarioCandidate], + *, + memory_blocked_count: int, + cross_model_analog_support_status: str, + cross_model_analog_scored_count: int, + cross_model_analog_factor_status: str, + risk_adjusted_interval_overlap_status: str, + cross_model_analog_prediction_interval_selectivity_status: str = "unknown", +) -> tuple[str, list[str]]: + if not candidates: + return "no_candidates", ["no_candidates"] + + cross_model_candidates = _cross_model_analog_candidates(candidates) + if cross_model_candidates: + blockers: list[str] = [] + if cross_model_analog_scored_count == 0: + blockers.append("no_scored_cross_model_candidates") + if cross_model_analog_support_status in { + "single_scored_cross_model_candidate", + "single_reference_cannot_rank_parallelism_variants", + "single_reference_cannot_rank_runtime_variants", + "single_reference_tied_tradeoff", + "tied_cross_model_frontier", + }: + blockers.append(cross_model_analog_support_status) + if cross_model_analog_factor_status == "single_cross_model_factor" and cross_model_analog_scored_count > 1: + blockers.append("single_cross_model_factor") + if ( + cross_model_analog_prediction_interval_selectivity_status == "nonselective_prediction_interval_top" + and cross_model_analog_scored_count > 1 + ): + blockers.append("nonselective_cross_model_prediction_interval_top") + if risk_adjusted_interval_overlap_status == "overlapping_best_interval": + blockers.append("risk_adjusted_interval_overlap") + requires_target_measurement = any( + "requires_remeasurement" in candidate.risk_flags for candidate in cross_model_candidates + ) + blockers = sorted(set(blockers)) + if not blockers and cross_model_analog_support_status == "scored_cross_model_frontier": + if requires_target_measurement: + return ( + "cross_model_generalization_supported_measurement_prior", + ["cross_model_candidates_require_measurement"], + ) + return "cross_model_generalization_supported_frontier", [] + if requires_target_measurement: + blockers = sorted({*blockers, "cross_model_candidates_require_measurement"}) + if "no_scored_cross_model_candidates" in blockers: + return "cross_model_generalization_unscored", blockers + if ( + "risk_adjusted_interval_overlap" in blockers + or "nonselective_cross_model_prediction_interval_top" in blockers + ): + return "cross_model_generalization_interval_uncertain", blockers + return "cross_model_generalization_requires_target_measurement", blockers + + if memory_blocked_count == len(candidates): + return "same_model_unscored_memory_blocked", ["memory_blocked_all_candidates"] + if any(candidate.prediction_confidence != "calibrated" for candidate in candidates): + blockers = ["same_model_extrapolated_predictions"] + if any("requires_remeasurement" in candidate.risk_flags for candidate in candidates): + blockers.append("same_model_candidates_require_measurement") + if any(candidate.memory_coverage_status == "analytic_floor_only" for candidate in candidates): + blockers.append("analytic_floor_only_memory") + return "same_model_extrapolation_requires_measurement", sorted(set(blockers)) + if any(candidate.behavior.model_ref is None for candidate in candidates): + return "same_model_calibrated_unknown_model_ref", ["missing_model_ref"] + return "same_model_calibrated", [] + + +def _scenario_prediction_fidelity_support( + candidates: list[ScenarioCandidate], + *, + memory_blocked_count: int, + phase_bottleneck_candidate_count: int, + routing_regime_status: str, + simulator_support_status_counts: dict[str, int], +) -> tuple[str, list[str]]: + if not candidates: + return "no_scenario_candidates", ["no_candidates"] + + scored = [candidate for candidate in candidates if candidate.score_tokens_per_sec is not None] + if not scored: + blockers = ["no_scored_candidates"] + if memory_blocked_count == len(candidates): + blockers.append("memory_blocked_all_candidates") + return "no_scored_scenario_fidelity", sorted(set(blockers)) + + blockers: set[str] = set() + # Memory-blocked candidates (observed/announced OOM rows) cannot carry throughput fidelity by + # definition; they are measured boundary evidence, not fidelity gaps — mirror the capture + # support's unscored definition and exempt them here too. + unscored_non_blocked = [ + candidate + for candidate in candidates + if candidate.score_tokens_per_sec is None and not _is_memory_blocked(candidate) + ] + if unscored_non_blocked: + blockers.add("unscored_candidates") + if any(candidate.prediction_confidence == "cross_model_extrapolated" for candidate in scored): + blockers.add("cross_model_predictions") + if any(candidate.prediction_confidence not in {"calibrated", "cross_model_extrapolated"} for candidate in scored): + blockers.add("extrapolated_predictions") + if any(candidate.calibration_scope != "exact_calibrated" for candidate in scored): + blockers.add("outside_exact_calibration_scope") + if any(candidate.memory_coverage_status == "analytic_floor_only" for candidate in scored): + blockers.add("analytic_floor_only_memory") + if any( + candidate.memory_coverage_status.startswith("extrapolated_") + or candidate.memory_coverage_status == "calibrated_overhead_peak_with_scaled_residual" + for candidate in scored + ): + blockers.add("scaled_or_extrapolated_memory_peak") + if any(candidate.timing_coverage_status == "no_timing_evidence" for candidate in scored): + blockers.add("no_timing_evidence") + if any( + "extrapolated" in candidate.timing_coverage_status + or candidate.timing_coverage_status.startswith(("reference_", "cross_model_reference_")) + for candidate in scored + ): + blockers.add("reference_or_extrapolated_timing") + if any(candidate.timing_coverage_status.endswith("_total_step_only") for candidate in scored): + blockers.add("missing_phase_timing") + if phase_bottleneck_candidate_count < len(scored): + blockers.add("missing_phase_bottleneck_evidence") + if any("requires_remeasurement" in candidate.risk_flags for candidate in scored): + blockers.add("requires_remeasurement") + if any(status.startswith("unsupported_") for status in simulator_support_status_counts): + blockers.add("unsupported_simulator_surface") + if any( + status != "supported_local_non_pp" and not status.startswith("unsupported_") + for status in simulator_support_status_counts + ): + blockers.add("partial_simulator_surface_support") + if any( + candidate.prediction_uncertainty_fraction is not None and candidate.prediction_uncertainty_fraction >= 0.50 + for candidate in scored + ): + blockers.add("high_prediction_uncertainty") + if routing_regime_status in {"unknown_routing_regime", "mixed_routing_regime"}: + blockers.add(routing_regime_status) + + blocker_list = sorted(blockers) + if "cross_model_predictions" in blockers: + return "cross_model_fidelity_requires_target_measurement", blocker_list + if { + "extrapolated_predictions", + "outside_exact_calibration_scope", + "analytic_floor_only_memory", + "scaled_or_extrapolated_memory_peak", + "no_timing_evidence", + "reference_or_extrapolated_timing", + "requires_remeasurement", + "unsupported_simulator_surface", + "high_prediction_uncertainty", + "unknown_routing_regime", + "mixed_routing_regime", + } & blockers: + return "extrapolated_fidelity_requires_measurement", blocker_list + if "partial_simulator_surface_support" in blockers: + return "partial_surface_fidelity_requires_measurement", blocker_list + total_step_only_blockers = { + "missing_phase_timing", + "missing_phase_bottleneck_evidence", + "unscored_candidates", + } + if blockers and blockers <= total_step_only_blockers: + if "unscored_candidates" in blockers: + return "partial_calibrated_total_step_fidelity", blocker_list + return "calibrated_total_step_fidelity", blocker_list + if "unscored_candidates" in blockers: + return "partial_calibrated_phase_fidelity", blocker_list + return "calibrated_phase_fidelity", [] + + +def _scenario_capture_support( + *, + candidate_count: int, + scored_count: int, + memory_blocked_count: int, + varied_parallelism_dimensions: list[str], + varied_workload_dimensions: list[str], + varied_runtime_dimensions: list[str], + runtime_mismatch_dimensions: list[str], + parallelism_axis_coverage_status_counts: dict[str, int], + simulator_support_status_counts: dict[str, int], +) -> tuple[str, list[str]]: + if candidate_count == 0: + return "no_scenario_candidates", ["no_candidates"] + + has_parallelism_variation = bool(varied_parallelism_dimensions) + has_workload_variation = bool(varied_workload_dimensions) + clean_axis_count = parallelism_axis_coverage_status_counts.get("scored_parallelism_axis", 0) + blockers: set[str] = set() + if candidate_count == 1: + blockers.add("single_candidate") + if scored_count == 0: + blockers.add("no_scored_candidates") + elif scored_count + memory_blocked_count < candidate_count: + # Memory-blocked candidates (observed OOM rows) are measured boundary evidence, not capture + # gaps; only genuinely unscored non-blocked candidates block broad capture. + blockers.add("unscored_candidates") + if memory_blocked_count == candidate_count: + blockers.add("memory_blocked_all_candidates") + if not has_parallelism_variation: + blockers.add("missing_parallelism_variation") + if not has_workload_variation: + blockers.add("missing_workload_variation") + if varied_runtime_dimensions: + blockers.add("runtime_variant_variation") + if runtime_mismatch_dimensions and (scored_count == 0 or memory_blocked_count == candidate_count): + blockers.add("runtime_mismatched_measurement_support") + if has_parallelism_variation and clean_axis_count == 0: + blockers.add("no_clean_parallelism_axis_coverage") + if any("blocked" in status for status in parallelism_axis_coverage_status_counts): + blockers.add("blocked_parallelism_axes") + if any(status.startswith("confounded_") for status in parallelism_axis_coverage_status_counts): + blockers.add("confounded_parallelism_axes") + if any("unscored" in status for status in parallelism_axis_coverage_status_counts): + blockers.add("unscored_parallelism_axes") + if any(status.startswith("unsupported_") for status in simulator_support_status_counts): + blockers.add("unsupported_simulator_surfaces") + if any( + status != "supported_local_non_pp" and not status.startswith("unsupported_") + for status in simulator_support_status_counts + ): + blockers.add("partial_simulator_surfaces") + + blocker_list = sorted(blockers) + if not has_parallelism_variation and not has_workload_variation: + return "single_shape_capture", blocker_list + if has_workload_variation and not has_parallelism_variation: + return "workload_only_capture", blocker_list + if has_parallelism_variation and not has_workload_variation: + if scored_count == 0: + return "unscored_parallelism_capture", blocker_list + if clean_axis_count == 0: + return "parallelism_capture_without_clean_axis", blocker_list + if blockers - {"missing_workload_variation"}: + return "partial_parallelism_capture", blocker_list + return "parallelism_only_capture", blocker_list + + if scored_count == 0: + return "unscored_broad_scenario_capture", blocker_list + if clean_axis_count == 0: + return "coupled_broad_scenario_capture", blocker_list + if blockers: + return "partial_broad_scenario_capture", blocker_list + return "broad_scenario_capture", [] + + +def _scenario_readiness_status( + *, + candidate_count: int, + unique_parallelism_strategy_count: int, + can_capture_scenario: bool, + can_predict_scenario_fidelity: bool, + can_select_parallelism_tradeoff: bool, + can_generalize_model: bool, + scenario_capture_status: str, + scenario_prediction_fidelity_status: str, + model_generalization_status: str, + parallelism_optimality_status: str, +) -> str: + if candidate_count == 0: + return "blocked_by_no_scenario_candidates" + if not can_capture_scenario: + if scenario_capture_status == "single_shape_capture": + return "blocked_by_single_shape_capture" + if scenario_capture_status.startswith("unscored_"): + return "blocked_by_unscored_scenario_capture" + if scenario_capture_status in { + "coupled_broad_scenario_capture", + "parallelism_capture_without_clean_axis", + }: + return "blocked_by_confounded_scenario_capture" + return "blocked_by_incomplete_scenario_capture" + if not can_predict_scenario_fidelity: + if scenario_prediction_fidelity_status == "no_scored_scenario_fidelity": + return "blocked_by_no_scored_scenario_fidelity" + if scenario_prediction_fidelity_status in { + "cross_model_fidelity_requires_target_measurement", + "extrapolated_fidelity_requires_measurement", + }: + return "blocked_by_unvalidated_prediction_fidelity" + if scenario_prediction_fidelity_status in { + "partial_calibrated_total_step_fidelity", + "calibrated_total_step_fidelity", + "partial_calibrated_phase_fidelity", + }: + return "partial_scenario_prediction_fidelity" + return "blocked_by_prediction_fidelity_gaps" + if not can_generalize_model: + if model_generalization_status.endswith("requires_measurement"): + return "blocked_by_model_generalization_measurement" + return "blocked_by_model_generalization_gaps" + if unique_parallelism_strategy_count > 1 and not can_select_parallelism_tradeoff: + if parallelism_optimality_status == "confounded_parallelism_winner": + return "blocked_by_confounded_parallelism_tradeoff" + if parallelism_optimality_status == "requires_measurement_before_parallelism_optimality": + return "blocked_by_parallelism_measurement" + if parallelism_optimality_status == "interval_overlap_parallelism_uncertain": + return "blocked_by_parallelism_interval_overlap" + return "blocked_by_parallelism_tradeoff_gaps" + return "validated_scenario_readiness" + + +def _scenario_readiness( + *, + candidate_count: int, + scored_count: int, + unscored_count: int, + memory_blocked_count: int, + unique_parallelism_strategy_count: int, + scored_parallelism_strategy_count: int, + promotable_parallelism_strategy_count: int, + scenario_capture_status: str, + scenario_capture_blockers: list[str], + scenario_prediction_fidelity_status: str, + scenario_prediction_fidelity_blockers: list[str], + parallelism_optimality_status: str, + parallelism_optimality_blockers: list[str], + model_generalization_status: str, + model_generalization_blockers: list[str], + measurement_readiness_status: str, + measurement_portfolio_coverage_status: str, + measurement_portfolio_coverage_blockers: list[str], + scenario_capture_gaps: list[ScenarioCaptureGap], + scenario_capture_gap_status_counts: dict[str, int], + scenario_capture_gap_required_measurements: list[str], + validation_actions: list[ScenarioValidationAction], + validation_action_status_counts: dict[str, int], + validation_action_required_measurements: list[str], + validation_action_total_gpu_count: int, + measurement_candidate_labels: list[str], + measurement_portfolio_total_gpu_count: int, + measurement_portfolio_parallelism_axis_gap_names: list[str], + measurement_portfolio_cross_model_analog_count: int, + varied_parallelism_dimensions: list[str], + varied_workload_dimensions: list[str], + varied_runtime_dimensions: list[str], + runtime_mismatch_dimensions: list[str], + parallelism_axis_coverage_status_counts: dict[str, int], + scored_parallelism_axis_names: list[str], + blocked_parallelism_axis_names: list[str], + confounded_parallelism_axis_names: list[str], + unscored_parallelism_axis_names: list[str], + missing_parallelism_axis_names: list[str], + simulator_support_status_counts: dict[str, int], + prediction_confidence_counts: dict[str, int], + calibration_scope_counts: dict[str, int], + memory_coverage_status_counts: dict[str, int], + timing_coverage_status_counts: dict[str, int], + cross_model_analog_support_status: str, + cross_model_analog_candidate_count: int, + cross_model_analog_scored_count: int, + benchmark_support: ScenarioBenchmarkSupport, +) -> ScenarioReadiness: + can_capture_scenario = scenario_capture_status == "broad_scenario_capture" + can_predict_scenario_fidelity = scenario_prediction_fidelity_status == "calibrated_phase_fidelity" + can_select_parallelism_tradeoff = parallelism_optimality_status in { + "supported_promotable_parallelism_tradeoff", + "supported_parallelism_winner", + } + can_generalize_model = model_generalization_status in { + "same_model_calibrated", + "cross_model_generalization_supported_frontier", + } + readiness_status = _scenario_readiness_status( + candidate_count=candidate_count, + unique_parallelism_strategy_count=unique_parallelism_strategy_count, + can_capture_scenario=can_capture_scenario, + can_predict_scenario_fidelity=can_predict_scenario_fidelity, + can_select_parallelism_tradeoff=can_select_parallelism_tradeoff, + can_generalize_model=can_generalize_model, + scenario_capture_status=scenario_capture_status, + scenario_prediction_fidelity_status=scenario_prediction_fidelity_status, + model_generalization_status=model_generalization_status, + parallelism_optimality_status=parallelism_optimality_status, + ) + top_capture_gaps = scenario_capture_gaps[:5] + return ScenarioReadiness( + readiness_status=readiness_status, + can_capture_scenario=can_capture_scenario, + can_predict_scenario_fidelity=can_predict_scenario_fidelity, + can_select_parallelism_tradeoff=can_select_parallelism_tradeoff, + can_generalize_model=can_generalize_model, + scenario_capture_status=scenario_capture_status, + scenario_capture_blockers=scenario_capture_blockers, + scenario_prediction_fidelity_status=scenario_prediction_fidelity_status, + scenario_prediction_fidelity_blockers=scenario_prediction_fidelity_blockers, + parallelism_optimality_status=parallelism_optimality_status, + parallelism_optimality_blockers=parallelism_optimality_blockers, + model_generalization_status=model_generalization_status, + model_generalization_blockers=model_generalization_blockers, + measurement_readiness_status=measurement_readiness_status, + measurement_portfolio_coverage_status=measurement_portfolio_coverage_status, + measurement_portfolio_coverage_blockers=measurement_portfolio_coverage_blockers, + required_measurements=_unique_in_order( + [*scenario_capture_gap_required_measurements, *validation_action_required_measurements] + ), + scenario_capture_gap_count=len(scenario_capture_gaps), + scenario_capture_gap_status_counts=scenario_capture_gap_status_counts, + scenario_capture_gap_required_measurements=scenario_capture_gap_required_measurements, + top_scenario_capture_gap_statuses=[gap.gap_status for gap in top_capture_gaps], + validation_action_count=len(validation_actions), + validation_action_status_counts=validation_action_status_counts, + validation_action_required_measurements=validation_action_required_measurements, + validation_action_total_gpu_count=validation_action_total_gpu_count, + measurement_candidate_count=len(measurement_candidate_labels), + measurement_candidate_labels=measurement_candidate_labels, + measurement_portfolio_total_gpu_count=measurement_portfolio_total_gpu_count, + measurement_portfolio_parallelism_axis_gap_names=measurement_portfolio_parallelism_axis_gap_names, + measurement_portfolio_cross_model_analog_count=measurement_portfolio_cross_model_analog_count, + candidate_count=candidate_count, + scored_count=scored_count, + unscored_count=unscored_count, + memory_blocked_count=memory_blocked_count, + unique_parallelism_strategy_count=unique_parallelism_strategy_count, + scored_parallelism_strategy_count=scored_parallelism_strategy_count, + promotable_parallelism_strategy_count=promotable_parallelism_strategy_count, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + runtime_mismatch_dimensions=runtime_mismatch_dimensions, + parallelism_axis_coverage_status_counts=parallelism_axis_coverage_status_counts, + scored_parallelism_axis_names=scored_parallelism_axis_names, + blocked_parallelism_axis_names=blocked_parallelism_axis_names, + confounded_parallelism_axis_names=confounded_parallelism_axis_names, + unscored_parallelism_axis_names=unscored_parallelism_axis_names, + missing_parallelism_axis_names=missing_parallelism_axis_names, + simulator_support_status_counts=simulator_support_status_counts, + prediction_confidence_counts=prediction_confidence_counts, + calibration_scope_counts=calibration_scope_counts, + memory_coverage_status_counts=memory_coverage_status_counts, + timing_coverage_status_counts=timing_coverage_status_counts, + cross_model_analog_support_status=cross_model_analog_support_status, + cross_model_analog_candidate_count=cross_model_analog_candidate_count, + cross_model_analog_scored_count=cross_model_analog_scored_count, + benchmark_support=benchmark_support, + ) + + +def _scenario_capture_gap( + *, + gap_status: str, + priority: int, + required_measurement: str, + reason: str, + blocker_names: list[str], + candidate_count: int, + scored_count: int, + memory_blocked_count: int, + missing_parallelism_axis_names: list[str], + blocked_parallelism_axis_names: list[str], + confounded_parallelism_axis_names: list[str], + unscored_parallelism_axis_names: list[str], + varied_parallelism_dimensions: list[str], + varied_workload_dimensions: list[str], + varied_runtime_dimensions: list[str], + runtime_mismatch_dimensions: list[str], +) -> ScenarioCaptureGap: + return ScenarioCaptureGap( + gap_status=gap_status, + priority=priority, + required_measurement=required_measurement, + reason=reason, + blocker_names=sorted(set(blocker_names)), + candidate_count=candidate_count, + scored_count=scored_count, + unscored_count=max(candidate_count - scored_count, 0), + memory_blocked_count=memory_blocked_count, + missing_parallelism_axis_names=missing_parallelism_axis_names, + blocked_parallelism_axis_names=blocked_parallelism_axis_names, + confounded_parallelism_axis_names=confounded_parallelism_axis_names, + unscored_parallelism_axis_names=unscored_parallelism_axis_names, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + runtime_mismatch_dimensions=runtime_mismatch_dimensions, + ) + + +def _scenario_capture_gap_portfolio( + *, + scenario_capture_blockers: list[str], + candidate_count: int, + scored_count: int, + memory_blocked_count: int, + missing_parallelism_axis_names: list[str], + blocked_parallelism_axis_names: list[str], + confounded_parallelism_axis_names: list[str], + unscored_parallelism_axis_names: list[str], + varied_parallelism_dimensions: list[str], + varied_workload_dimensions: list[str], + varied_runtime_dimensions: list[str], + runtime_mismatch_dimensions: list[str], +) -> list[ScenarioCaptureGap]: + blockers = set(scenario_capture_blockers) + common = { + "candidate_count": candidate_count, + "scored_count": scored_count, + "memory_blocked_count": memory_blocked_count, + "missing_parallelism_axis_names": missing_parallelism_axis_names, + "blocked_parallelism_axis_names": blocked_parallelism_axis_names, + "confounded_parallelism_axis_names": confounded_parallelism_axis_names, + "unscored_parallelism_axis_names": unscored_parallelism_axis_names, + "varied_parallelism_dimensions": varied_parallelism_dimensions, + "varied_workload_dimensions": varied_workload_dimensions, + "varied_runtime_dimensions": varied_runtime_dimensions, + "runtime_mismatch_dimensions": runtime_mismatch_dimensions, + } + gaps: list[ScenarioCaptureGap] = [] + if "no_candidates" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="no_scenario_candidates_need_base_sweep", + priority=130, + required_measurement="add_base_scenario_candidates", + reason="no candidate configurations were generated for this scenario", + blocker_names=["no_candidates"], + **common, + ) + ) + if "single_candidate" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="single_candidate_needs_workload_and_parallelism_variants", + priority=120, + required_measurement="add_workload_and_parallelism_variants", + reason="one candidate cannot prove scenario behavior across workload or topology changes", + blocker_names=["single_candidate"], + **common, + ) + ) + if "missing_parallelism_variation" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="missing_parallelism_variation_needs_axis_sweep", + priority=110, + required_measurement="add_same_workload_same_runtime_parallelism_axis_variants", + reason="scenario capture has no topology variation to test parallelism tradeoffs", + blocker_names=["missing_parallelism_variation"], + **common, + ) + ) + if "confounded_parallelism_axes" in blockers or "no_clean_parallelism_axis_coverage" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="confounded_parallelism_capture_needs_like_for_like_axis_pairs", + priority=105, + required_measurement="add_same_workload_same_runtime_axis_pairs", + reason="parallelism varies, but no clean like-for-like axis comparison proves the scenario", + blocker_names=[ + blocker + for blocker in ["confounded_parallelism_axes", "no_clean_parallelism_axis_coverage"] + if blocker in blockers + ], + **common, + ) + ) + if "blocked_parallelism_axes" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="blocked_parallelism_capture_needs_fit_boundary", + priority=100, + required_measurement="add_fit_failure_boundary_near_blocked_axes", + reason="some topology axes are represented only by memory-blocked candidates", + blocker_names=["blocked_parallelism_axes"], + **common, + ) + ) + if "missing_workload_variation" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="missing_workload_variation_needs_workload_sweep", + priority=95, + required_measurement="add_same_parallelism_runtime_workload_variants", + reason="scenario capture has no workload variation to test scaling over batch or sequence shape", + blocker_names=["missing_workload_variation"], + **common, + ) + ) + if "unsupported_simulator_surfaces" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="unsupported_surface_capture_needs_support_or_probe", + priority=90, + required_measurement="add_simulator_support_or_direct_probe_for_surface", + reason="candidate surface is outside the simulator implementation", + blocker_names=["unsupported_simulator_surfaces"], + **common, + ) + ) + if "memory_blocked_all_candidates" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="memory_blocked_capture_needs_fit_probe", + priority=90, + required_measurement="add_memory_pressure_fit_probe_or_reduce_batch", + reason="every candidate is memory blocked, so no scored scenario behavior is captured", + blocker_names=["memory_blocked_all_candidates"], + **common, + ) + ) + if "runtime_mismatched_measurement_support" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="runtime_mismatch_capture_needs_runtime_variant", + priority=88, + required_measurement="add_same_parallelism_workload_runtime_variants", + reason="supporting measurements differ from target runtime knobs", + blocker_names=["runtime_mismatched_measurement_support"], + **common, + ) + ) + if "no_scored_candidates" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="no_scored_capture_needs_supported_measurement", + priority=85, + required_measurement="add_scored_measurements_for_existing_sweep", + reason="generated candidates do not include any scored throughput measurement", + blocker_names=["no_scored_candidates"], + **common, + ) + ) + if "unscored_candidates" in blockers or "unscored_parallelism_axes" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="unscored_capture_needs_candidate_replay", + priority=80, + required_measurement="score_unscored_capture_candidates", + reason="some generated candidates cannot contribute measured scenario behavior", + blocker_names=[ + blocker for blocker in ["unscored_candidates", "unscored_parallelism_axes"] if blocker in blockers + ], + **common, + ) + ) + if "runtime_variant_variation" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="runtime_variant_capture_needs_runtime_isolation", + priority=75, + required_measurement="replay_with_runtime_dimensions_fixed", + reason="runtime knobs vary with the scenario and confound capture", + blocker_names=["runtime_variant_variation"], + **common, + ) + ) + if "partial_simulator_surfaces" in blockers: + gaps.append( + _scenario_capture_gap( + gap_status="partial_surface_capture_needs_support_or_probe", + priority=65, + required_measurement="add_partial_surface_support_or_direct_probe", + reason="some candidates are only partially covered by the simulator surface", + blocker_names=["partial_simulator_surfaces"], + **common, + ) + ) + return sorted(gaps, key=lambda gap: (-gap.priority, gap.gap_status)) + + +_CROSS_MODEL_ANALOG_SUPPORT_FACTORS = { + "single_scored_cross_model_candidate": 0.85, + "single_reference_cannot_rank_parallelism_variants": 0.60, + "single_reference_cannot_rank_runtime_variants": 0.60, + "single_reference_tied_tradeoff": 0.70, + "tied_cross_model_frontier": 0.80, +} + + +def _replace_decision_factor(factors: list[str], prefix: str, value: str) -> list[str]: + replacement = f"{prefix}{value}" + updated: list[str] = [] + replaced = False + for factor in factors: + if factor.startswith(prefix): + if not replaced: + updated.append(replacement) + replaced = True + continue + updated.append(factor) + if not replaced: + updated.append(replacement) + return updated + + +def _apply_cross_model_analog_support_risk(candidates: list[ScenarioCandidate]) -> list[ScenarioCandidate]: + support_status, *_ = _cross_model_analog_support(candidates) + factor = _CROSS_MODEL_ANALOG_SUPPORT_FACTORS.get(support_status, 1.0) + if factor >= 1.0: + return candidates + + adjusted_candidates: list[ScenarioCandidate] = [] + support_flag = f"cross_model_support:{support_status}" + for candidate in candidates: + is_analog_candidate = ( + candidate.prediction_confidence == "cross_model_extrapolated" + or "cross_model_analog" in candidate.risk_flags + ) + if not is_analog_candidate: + adjusted_candidates.append(candidate) + continue + adjusted_uncertainty = ( + round(min(candidate.prediction_uncertainty_fraction + 0.10, 0.95), 3) + if candidate.prediction_uncertainty_fraction is not None + else None + ) + raw_interval_lower, raw_interval_upper = _prediction_interval( + candidate.score_tokens_per_sec, + adjusted_uncertainty, + ) + if candidate.score_risk_adjusted_tokens_per_sec is None: + decision_factors = _replace_decision_factor( + candidate.decision_factors, + "prediction_uncertainty_fraction=", + _format_factor_float(adjusted_uncertainty), + ) + if raw_interval_lower is not None and raw_interval_upper is not None: + decision_factors = _replace_decision_factor( + decision_factors, + "prediction_interval_tokens_per_sec=", + f"{_format_factor_float(raw_interval_lower)}..{_format_factor_float(raw_interval_upper)}", + ) + decision_factors.extend( + [ + f"cross_model_support={support_status}", + f"cross_model_support_factor={factor:.3f}", + f"prediction_uncertainty_after_cross_model_support={_format_factor_float(adjusted_uncertainty)}", + ] + ) + adjusted_candidates.append( + replace( + candidate, + prediction_uncertainty_fraction=adjusted_uncertainty, + prediction_interval_lower_tokens_per_sec=raw_interval_lower, + prediction_interval_upper_tokens_per_sec=raw_interval_upper, + risk_flags=sorted({*candidate.risk_flags, support_flag}), + decision_factors=decision_factors, + ) + ) + continue + + adjusted_score = round(candidate.score_risk_adjusted_tokens_per_sec * factor, 3) + adjusted_efficiency_score = ( + round(adjusted_score / candidate.topology.world_size, 3) if candidate.topology.world_size else None + ) + risk_interval_lower, risk_interval_upper = _prediction_interval(adjusted_score, adjusted_uncertainty) + decision_factors = _replace_decision_factor( + candidate.decision_factors, + "risk_adjusted_tokens_per_sec=", + _format_factor_float(adjusted_score), + ) + if adjusted_efficiency_score is not None: + decision_factors = _replace_decision_factor( + decision_factors, + "risk_adjusted_tokens_per_gpu_per_sec=", + _format_factor_float(adjusted_efficiency_score), + ) + decision_factors = _replace_decision_factor( + decision_factors, + "prediction_uncertainty_fraction=", + _format_factor_float(adjusted_uncertainty), + ) + if raw_interval_lower is not None and raw_interval_upper is not None: + decision_factors = _replace_decision_factor( + decision_factors, + "prediction_interval_tokens_per_sec=", + f"{_format_factor_float(raw_interval_lower)}..{_format_factor_float(raw_interval_upper)}", + ) + if risk_interval_lower is not None and risk_interval_upper is not None: + decision_factors = _replace_decision_factor( + decision_factors, + "risk_adjusted_prediction_interval_tokens_per_sec=", + f"{_format_factor_float(risk_interval_lower)}..{_format_factor_float(risk_interval_upper)}", + ) + decision_factors.extend( + [ + f"cross_model_support={support_status}", + f"cross_model_support_factor={factor:.3f}", + f"risk_adjusted_after_cross_model_support={_format_factor_float(adjusted_score)}", + f"prediction_uncertainty_after_cross_model_support={_format_factor_float(adjusted_uncertainty)}", + ] + ) + adjusted_candidates.append( + replace( + candidate, + score_risk_adjusted_tokens_per_sec=adjusted_score, + score_risk_adjusted_tokens_per_gpu_per_sec=adjusted_efficiency_score, + prediction_uncertainty_fraction=adjusted_uncertainty, + prediction_interval_lower_tokens_per_sec=raw_interval_lower, + prediction_interval_upper_tokens_per_sec=raw_interval_upper, + risk_adjusted_prediction_interval_lower_tokens_per_sec=risk_interval_lower, + risk_adjusted_prediction_interval_upper_tokens_per_sec=risk_interval_upper, + risk_flags=sorted({*candidate.risk_flags, support_flag}), + decision_factors=decision_factors, + notes=[ + *candidate.notes, + f"cross_model_support_factor={factor:.3f}:{support_status}", + ], + ) + ) + return adjusted_candidates + + +def _format_count_summary(counts: dict[str, int]) -> str: + return ",".join(f"{key}:{value}" for key, value in sorted(counts.items())) + + +def _routing_regime(behavior: BenchmarkBehaviorPrediction) -> str: + if behavior.balanced_routing is True: + return "balanced_synthetic_routing" + if behavior.balanced_routing is False: + return "real_routing" + return "unknown_routing" + + +def _routing_regime_status(counts: dict[str, int]) -> str: + if not counts: + return "no_candidates" + known = [regime for regime in counts if regime != "unknown_routing"] + if not known: + return "unknown_routing_regime" + if len(known) == 1 and ("unknown_routing" not in counts): + return "single_routing_regime" + return "mixed_routing_regime" + + +def _scenario_measurement_guidance( + *, + candidates: list[ScenarioCandidate], + scored_count: int, + memory_blocked_count: int, + parallelism_tradeoff_status: str, + parallelism_optimality_status: str, + parallelism_optimality_blockers: list[str], + parallelism_boundary_status: str, + parallelism_boundary_prediction_status: str, + parallelism_boundary_prediction_blockers: list[str], + throughput_efficiency_tradeoff_status: str, + throughput_efficiency_frontier_count: int, + risk_adjusted_efficiency_frontier_count: int, + raw_dominated_candidate_count: int, + risk_adjusted_dominated_candidate_count: int, + parallelism_axis_comparison_count: int, + isolated_parallelism_axis_comparison_count: int, + coupled_parallelism_axis_comparison_count: int, + parallelism_axis_interval_overlap_count: int, + blocked_parallelism_axis_names: list[str], + confounded_parallelism_axis_names: list[str], + same_workload_scaling_status: str, + same_workload_scaling_candidate_count: int, + min_scaling_efficiency: float | None, + memory_coverage_status_counts: dict[str, int], + timing_coverage_status_counts: dict[str, int], + max_memory_residual_gb: float | None, + phase_bottleneck_bucket_counts: dict[str, int], + max_phase_bottleneck_share: float | None, + max_phase_bottleneck_half_speedup_delta_pct: float | None, + memory_bottleneck_bucket_counts: dict[str, int], + max_memory_bottleneck_fraction_of_peak: float | None, + high_uncertainty_candidate_count: int, + max_prediction_uncertainty_fraction: float | None, + risk_adjusted_interval_overlap_status: str, + risk_adjusted_interval_overlap_contender_count: int, + risk_adjusted_interval_best_vs_next_margin_tokens_per_sec: float | None, + routing_regime_status: str, + routing_regime_counts: dict[str, int], + cross_model_analog_support_status: str, + cross_model_analog_factor_status: str, + cross_model_analog_unique_factor_count: int, + cross_model_analog_unique_target_runtime_signature_count: int, + cross_model_analog_scored_varied_parallelism_dimensions: list[str], + cross_model_analog_scored_varied_workload_dimensions: list[str], + model_generalization_status: str, + model_generalization_blockers: list[str], + promotion_readiness_status: str, + promotable_raw_gap_tokens_per_sec: float | None, + promotable_risk_adjusted_gap_tokens_per_sec: float | None, + best_risk_adjusted: ScenarioCandidate | None, + best_next_measurement: ScenarioCandidate | None, + best_promotable: ScenarioCandidate | None, +) -> tuple[str, list[str]]: + focus = best_next_measurement or best_promotable or best_risk_adjusted + rationale = [f"parallelism_tradeoff={parallelism_tradeoff_status}"] + rationale.append(f"parallelism_optimality={parallelism_optimality_status}") + if parallelism_optimality_blockers: + rationale.append(f"parallelism_optimality_blockers={','.join(parallelism_optimality_blockers)}") + rationale.append(f"parallelism_boundary={parallelism_boundary_status}") + rationale.append(f"parallelism_boundary_prediction={parallelism_boundary_prediction_status}") + if parallelism_boundary_prediction_blockers: + rationale.append( + f"parallelism_boundary_prediction_blockers={','.join(parallelism_boundary_prediction_blockers)}" + ) + rationale.append(f"throughput_efficiency_tradeoff={throughput_efficiency_tradeoff_status}") + rationale.append(f"throughput_efficiency_frontier_count={throughput_efficiency_frontier_count}") + rationale.append(f"risk_adjusted_efficiency_frontier_count={risk_adjusted_efficiency_frontier_count}") + if raw_dominated_candidate_count: + rationale.append(f"raw_dominated_candidate_count={raw_dominated_candidate_count}") + if risk_adjusted_dominated_candidate_count: + rationale.append(f"risk_adjusted_dominated_candidate_count={risk_adjusted_dominated_candidate_count}") + if parallelism_axis_comparison_count: + rationale.append(f"parallelism_axis_comparison_count={parallelism_axis_comparison_count}") + rationale.append(f"isolated_parallelism_axis_comparison_count={isolated_parallelism_axis_comparison_count}") + rationale.append(f"coupled_parallelism_axis_comparison_count={coupled_parallelism_axis_comparison_count}") + if parallelism_axis_interval_overlap_count: + rationale.append(f"parallelism_axis_interval_overlap_count={parallelism_axis_interval_overlap_count}") + if blocked_parallelism_axis_names: + rationale.append(f"blocked_parallelism_axes={','.join(blocked_parallelism_axis_names)}") + if confounded_parallelism_axis_names: + rationale.append(f"confounded_parallelism_axes={','.join(confounded_parallelism_axis_names)}") + rationale.append(f"same_workload_scaling={same_workload_scaling_status}") + if same_workload_scaling_candidate_count: + rationale.append(f"same_workload_scaling_candidate_count={same_workload_scaling_candidate_count}") + if min_scaling_efficiency is not None: + rationale.append(f"min_scaling_efficiency={min_scaling_efficiency:.3f}") + if memory_coverage_status_counts: + rationale.append(f"memory_coverage={_format_count_summary(memory_coverage_status_counts)}") + if timing_coverage_status_counts: + rationale.append(f"timing_coverage={_format_count_summary(timing_coverage_status_counts)}") + if max_memory_residual_gb is not None: + rationale.append(f"max_memory_residual_gb={max_memory_residual_gb:.3f}") + if phase_bottleneck_bucket_counts: + rationale.append(f"phase_bottlenecks={_format_count_summary(phase_bottleneck_bucket_counts)}") + if max_phase_bottleneck_share is not None: + rationale.append(f"max_phase_bottleneck_share={max_phase_bottleneck_share:.3f}") + if max_phase_bottleneck_half_speedup_delta_pct is not None: + rationale.append( + f"max_phase_bottleneck_half_speedup_delta_pct={max_phase_bottleneck_half_speedup_delta_pct:.3f}" + ) + if memory_bottleneck_bucket_counts: + rationale.append(f"memory_bottlenecks={_format_count_summary(memory_bottleneck_bucket_counts)}") + if max_memory_bottleneck_fraction_of_peak is not None: + rationale.append(f"max_memory_bottleneck_fraction_of_peak={max_memory_bottleneck_fraction_of_peak:.3f}") + if high_uncertainty_candidate_count: + rationale.append(f"high_uncertainty_candidate_count={high_uncertainty_candidate_count}") + if max_prediction_uncertainty_fraction is not None: + rationale.append(f"max_prediction_uncertainty_fraction={max_prediction_uncertainty_fraction:.3f}") + if risk_adjusted_interval_overlap_status != "unknown": + rationale.append(f"risk_adjusted_interval_overlap={risk_adjusted_interval_overlap_status}") + if risk_adjusted_interval_overlap_contender_count: + rationale.append( + f"risk_adjusted_interval_overlap_contender_count={risk_adjusted_interval_overlap_contender_count}" + ) + if risk_adjusted_interval_best_vs_next_margin_tokens_per_sec is not None: + rationale.append( + "risk_adjusted_interval_best_vs_next_margin_tokens_per_sec=" + f"{risk_adjusted_interval_best_vs_next_margin_tokens_per_sec:.3f}" + ) + if routing_regime_status == "mixed_routing_regime": + rationale.append(f"routing_regime={routing_regime_status}") + rationale.append(f"routing_regimes={_format_count_summary(routing_regime_counts)}") + if cross_model_analog_support_status != "not_used": + rationale.append(f"cross_model_analog_support={cross_model_analog_support_status}") + if cross_model_analog_factor_status not in {"not_used", "no_scored_cross_model_factors"}: + rationale.append(f"cross_model_analog_factor_status={cross_model_analog_factor_status}") + rationale.append(f"cross_model_analog_unique_factor_count={cross_model_analog_unique_factor_count}") + if cross_model_analog_unique_target_runtime_signature_count > 1: + rationale.append( + "cross_model_analog_unique_target_runtime_signature_count=" + f"{cross_model_analog_unique_target_runtime_signature_count}" + ) + if cross_model_analog_scored_varied_parallelism_dimensions: + rationale.append( + f"cross_model_analog_varied_parallelism={','.join(cross_model_analog_scored_varied_parallelism_dimensions)}" + ) + if cross_model_analog_scored_varied_workload_dimensions: + rationale.append( + f"cross_model_analog_varied_workload={','.join(cross_model_analog_scored_varied_workload_dimensions)}" + ) + rationale.append(f"model_generalization={model_generalization_status}") + if model_generalization_blockers: + rationale.append(f"model_generalization_blockers={','.join(model_generalization_blockers)}") + if best_risk_adjusted is not None: + rationale.append(f"best_risk_adjusted={best_risk_adjusted.label}") + if best_promotable is not None: + rationale.append(f"best_promotable={best_promotable.label}") + rationale.append(f"promotion_readiness={promotion_readiness_status}") + if promotable_raw_gap_tokens_per_sec is not None: + rationale.append(f"promotable_raw_gap_tokens_per_sec={promotable_raw_gap_tokens_per_sec:.3f}") + if promotable_risk_adjusted_gap_tokens_per_sec is not None: + rationale.append( + f"promotable_risk_adjusted_gap_tokens_per_sec={promotable_risk_adjusted_gap_tokens_per_sec:.3f}" + ) + if best_next_measurement is not None: + rationale.append(f"best_next_measurement={best_next_measurement.label}") + if focus is not None: + rationale.append(f"focused_recommendation={focus.recommendation}") + if focus.risk_flags: + rationale.append(f"focused_risks={','.join(focus.risk_flags)}") + + if not candidates: + return "no_candidates", rationale + if best_next_measurement is not None: + if parallelism_tradeoff_status == "scored_parallelism_tradeoff_requires_remeasurement": + return "measure_parallelism_tradeoff_candidate", rationale + if "cross_model_analog" in best_next_measurement.risk_flags: + return "measure_cross_model_analog_candidate", rationale + if best_next_measurement.memory_coverage_status == "analytic_floor_only": + return "measure_analytic_floor_candidate", rationale + return "measure_best_next_candidate", rationale + if best_promotable is not None: + if parallelism_tradeoff_status == "promotable_parallelism_tradeoff": + return "promote_parallelism_tradeoff_winner", rationale + return "promote_best_promotable_candidate", rationale + if best_risk_adjusted is not None: + if best_risk_adjusted.recommendation == "debug_runtime_failure": + return "debug_best_risk_adjusted_candidate", rationale + if best_risk_adjusted.recommendation == "correctness_gate_required": + return "correctness_gate_best_risk_adjusted_candidate", rationale + if best_risk_adjusted.recommendation.startswith("remeasure"): + return "measure_best_risk_adjusted_candidate", rationale + return "review_best_risk_adjusted_candidate", rationale + if memory_blocked_count == len(candidates): + if any(candidate.feasibility_status == "memory_floor_exceeds_safety_margin" for candidate in candidates): + return "measure_memory_safety_margin_candidate", rationale + if memory_coverage_status_counts == {"analytic_floor_only": len(candidates)}: + return "blocked_by_analytic_memory_floor", rationale + return "blocked_by_memory_model", rationale + if scored_count == 0: + return "no_scored_candidate", rationale + return "review_scenario", rationale + + +def _count_values(values: list[str]) -> dict[str, int]: + return dict(sorted(Counter(values).items())) + + +def _count_int_values(values: list[int | None]) -> dict[int, int]: + return dict(sorted(Counter(value for value in values if value is not None).items())) + + +def _count_candidate_model_refs(candidates: list[ScenarioCandidate]) -> dict[str, int]: + return _count_values([candidate.behavior.model_ref or "unknown_model_ref" for candidate in candidates]) + + +def _count_candidate_runtime_signatures(candidates: list[ScenarioCandidate]) -> dict[str, int]: + return _count_values([candidate.target_runtime_signature for candidate in candidates]) + + +def _count_candidate_topology_values(candidates: list[ScenarioCandidate], dimension: str) -> dict[int, int]: + return _count_int_values([_candidate_dimension_value(candidate, dimension) for candidate in candidates]) + + +def _cross_node_dimension_counts(candidates: list[ScenarioCandidate]) -> dict[str, int]: + dimensions = [ + dimension + for candidate in candidates + if candidate.communication is not None + for dimension in candidate.communication.cross_node_dimensions + ] + return _count_values(dimensions) + + +def _is_memory_blocked(candidate: ScenarioCandidate) -> bool: + status = candidate.feasibility_status + return ( + status == "observed_oom" or status.endswith("_exceeds_limit") or status == "memory_floor_exceeds_safety_margin" + ) + + +def _risk_adjusted_interval_overlap_summary( + feasible: list[ScenarioCandidate], + best_risk_adjusted: ScenarioCandidate | None, +) -> tuple[str, int, list[str], float | None]: + if best_risk_adjusted is None: + return "no_scored_interval", 0, [], None + best_lower = best_risk_adjusted.risk_adjusted_prediction_interval_lower_tokens_per_sec + best_upper = best_risk_adjusted.risk_adjusted_prediction_interval_upper_tokens_per_sec + if best_lower is None or best_upper is None: + return "no_scored_interval", 0, [], None + + overlapping_labels: list[str] = [] + contender_upper_bounds: list[float] = [] + best_strategy = _parallelism_strategy_key(best_risk_adjusted) + for candidate in feasible: + if candidate.label == best_risk_adjusted.label: + continue + # Same-strategy candidates are re-measurements of the SAME parallelism choice (winner reruns, + # instrumented reruns, historical twins); their interval overlap cannot make the strategy + # SELECTION uncertain. Only overlap with a genuinely different strategy blocks the choice. + if _parallelism_strategy_key(candidate) == best_strategy: + continue + contender_lower = candidate.risk_adjusted_prediction_interval_lower_tokens_per_sec + contender_upper = candidate.risk_adjusted_prediction_interval_upper_tokens_per_sec + if contender_lower is None or contender_upper is None: + continue + contender_upper_bounds.append(contender_upper) + if contender_lower <= best_upper and contender_upper >= best_lower: + overlapping_labels.append(candidate.label) + + if not contender_upper_bounds: + return "single_scored_interval", 0, [], None + + margin = round(best_lower - max(contender_upper_bounds), 3) + if overlapping_labels: + return "overlapping_best_interval", len(overlapping_labels), sorted(overlapping_labels), margin + return "clear_best_interval", 0, [], margin + + +def _scenario_decision_summary( + candidates: list[ScenarioCandidate], + feasible: list[ScenarioCandidate], + best_raw: ScenarioCandidate | None, + best_risk_adjusted: ScenarioCandidate | None, + best_efficiency: ScenarioCandidate | None, + best_risk_adjusted_efficiency: ScenarioCandidate | None, + best_next_measurement: ScenarioCandidate | None, + best_promotable: ScenarioCandidate | None, + benchmark_support: ScenarioBenchmarkSupport, +) -> ScenarioDecisionSummary: + risk_flags = [flag for candidate in candidates for flag in candidate.risk_flags] + distances = [ + candidate.calibration_distance for candidate in candidates if candidate.calibration_distance is not None + ] + scored_distances = [ + candidate.calibration_distance for candidate in feasible if candidate.calibration_distance is not None + ] + uncertainty_fractions = [ + candidate.prediction_uncertainty_fraction + for candidate in candidates + if candidate.prediction_uncertainty_fraction is not None + ] + scored_uncertainty_fractions = [ + candidate.prediction_uncertainty_fraction + for candidate in feasible + if candidate.prediction_uncertainty_fraction is not None + ] + high_uncertainty_candidate_count = sum( + 1 + for candidate in candidates + if candidate.prediction_uncertainty_fraction is not None and candidate.prediction_uncertainty_fraction >= 0.50 + ) + max_prediction_uncertainty_fraction = round(max(uncertainty_fractions), 3) if uncertainty_fractions else None + ( + risk_adjusted_interval_overlap_status, + risk_adjusted_interval_overlap_contender_count, + risk_adjusted_interval_overlap_contender_labels, + risk_adjusted_interval_best_vs_next_margin, + ) = _risk_adjusted_interval_overlap_summary(feasible, best_risk_adjusted) + memory_residuals = [ + candidate.estimated_memory_residual_gb + for candidate in candidates + if candidate.estimated_memory_residual_gb is not None + ] + memory_residual_fractions = [ + candidate.estimated_memory_residual_fraction_of_peak + for candidate in candidates + if candidate.estimated_memory_residual_fraction_of_peak is not None + ] + phase_bottleneck_candidates = [ + candidate for candidate in candidates if candidate.phase_bottleneck_bucket is not None + ] + phase_bottleneck_shares = [ + candidate.phase_bottleneck_share + for candidate in phase_bottleneck_candidates + if candidate.phase_bottleneck_share is not None + ] + phase_bottleneck_bucket_counts = _count_values( + [ + candidate.phase_bottleneck_bucket + for candidate in phase_bottleneck_candidates + if candidate.phase_bottleneck_bucket + ] + ) + phase_bottleneck_phase_counts = _count_values( + [ + candidate.phase_bottleneck_phase + for candidate in phase_bottleneck_candidates + if candidate.phase_bottleneck_phase + ] + ) + max_phase_bottleneck_share = round(max(phase_bottleneck_shares), 3) if phase_bottleneck_shares else None + max_phase_bottleneck_candidate = _max_positive_candidate_by_field(candidates, "phase_bottleneck_share") + phase_bottleneck_half_speedup_candidates = [ + candidate for candidate in candidates if candidate.phase_bottleneck_half_speedup_delta_pct is not None + ] + phase_bottleneck_half_speedup_deltas = [ + candidate.phase_bottleneck_half_speedup_delta_pct + for candidate in phase_bottleneck_half_speedup_candidates + if candidate.phase_bottleneck_half_speedup_delta_pct is not None + ] + max_phase_bottleneck_half_speedup_delta_pct = ( + round(max(phase_bottleneck_half_speedup_deltas), 3) if phase_bottleneck_half_speedup_deltas else None + ) + max_phase_bottleneck_half_speedup_candidate = _max_positive_candidate_by_field( + candidates, + "phase_bottleneck_half_speedup_delta_pct", + ) + memory_bottleneck_candidates = [ + candidate for candidate in candidates if candidate.memory_bottleneck_bucket is not None + ] + memory_bottleneck_fractions = [ + candidate.memory_bottleneck_fraction_of_peak + for candidate in memory_bottleneck_candidates + if candidate.memory_bottleneck_fraction_of_peak is not None + ] + memory_bottleneck_bucket_counts = _count_values( + [ + candidate.memory_bottleneck_bucket + for candidate in memory_bottleneck_candidates + if candidate.memory_bottleneck_bucket + ] + ) + memory_bottleneck_phase_counts = _count_values( + [ + candidate.memory_bottleneck_phase + for candidate in memory_bottleneck_candidates + if candidate.memory_bottleneck_phase + ] + ) + max_memory_bottleneck_fraction = round(max(memory_bottleneck_fractions), 3) if memory_bottleneck_fractions else None + max_memory_bottleneck_candidate = _max_positive_candidate_by_field( + candidates, + "memory_bottleneck_fraction_of_peak", + ) + ( + unique_strategy_count, + scored_strategy_count, + promotable_strategy_count, + requires_remeasurement_strategy_count, + ) = _parallelism_strategy_counts(candidates) + scored_count = len(feasible) + memory_blocked_count = sum(1 for candidate in candidates if _is_memory_blocked(candidate)) + memory_coverage_status_counts = _count_values([candidate.memory_coverage_status for candidate in candidates]) + timing_coverage_status_counts = _count_values([candidate.timing_coverage_status for candidate in candidates]) + simulator_support_status_counts = _count_values([candidate.simulator_support_status for candidate in candidates]) + simulator_support_blocker_counts = _count_values( + [blocker for candidate in candidates for blocker in candidate.simulator_support_blockers] + ) + max_memory_residual_gb = round(max(memory_residuals), 3) if memory_residuals else None + routing_regime_counts = _count_values([_routing_regime(candidate.behavior) for candidate in candidates]) + routing_regime_status = _routing_regime_status(routing_regime_counts) + parallelism_tradeoff_status = _parallelism_tradeoff_status( + unique_strategy_count=unique_strategy_count, + scored_strategy_count=scored_strategy_count, + promotable_strategy_count=promotable_strategy_count, + requires_remeasurement_strategy_count=requires_remeasurement_strategy_count, + ) + ( + cross_model_analog_support_status, + cross_model_analog_candidate_count, + cross_model_analog_scored_count, + cross_model_analog_unique_prediction_count, + cross_model_analog_unique_matched_label_count, + cross_model_analog_unique_target_strategy_count, + cross_model_analog_unique_target_runtime_signature_count, + ) = _cross_model_analog_support(candidates) + cross_model_analog_scored_candidates = [ + candidate + for candidate in _cross_model_analog_candidates(candidates) + if candidate.score_tokens_per_sec is not None + ] + cross_model_analog_scored_varied_parallelism_dimensions = _varied_candidate_dimensions( + cross_model_analog_scored_candidates, _SCENARIO_PARALLELISM_DIMENSIONS + ) + cross_model_analog_scored_varied_workload_dimensions = _varied_candidate_dimensions( + cross_model_analog_scored_candidates, _SCENARIO_WORKLOAD_DIMENSIONS + ) + cross_model_analog_factor_status = _cross_model_analog_factor_status(candidates) + cross_model_analog_unique_factor_count = _cross_model_analog_unique_factor_count(candidates) + cross_model_analog_factor_ranges = _cross_model_analog_factor_ranges(candidates) + cross_model_analog_prediction_interval_top = _cross_model_analog_prediction_interval_top(candidates) + cross_model_analog_prediction_interval_top_count = len(cross_model_analog_prediction_interval_top) + cross_model_analog_prediction_interval_top_fraction = ( + round(cross_model_analog_prediction_interval_top_count / cross_model_analog_scored_count, 3) + if cross_model_analog_scored_count + else None + ) + if cross_model_analog_candidate_count == 0: + cross_model_analog_prediction_interval_selectivity_status = "not_used" + else: + cross_model_analog_prediction_interval_selectivity_status = _prediction_interval_selectivity_status( + scored_candidate_count=cross_model_analog_scored_count, + prediction_interval_top_count=cross_model_analog_prediction_interval_top_count, + ) + cross_model_analog_prediction_interval_top_labels = [ + candidate.label for candidate in cross_model_analog_prediction_interval_top + ] + model_generalization_status, model_generalization_blockers = _model_generalization_support( + candidates, + memory_blocked_count=memory_blocked_count, + cross_model_analog_support_status=cross_model_analog_support_status, + cross_model_analog_scored_count=cross_model_analog_scored_count, + cross_model_analog_factor_status=cross_model_analog_factor_status, + risk_adjusted_interval_overlap_status=risk_adjusted_interval_overlap_status, + cross_model_analog_prediction_interval_selectivity_status=( + cross_model_analog_prediction_interval_selectivity_status + ), + ) + scenario_prediction_fidelity_status, scenario_prediction_fidelity_blockers = _scenario_prediction_fidelity_support( + candidates, + memory_blocked_count=memory_blocked_count, + phase_bottleneck_candidate_count=len(phase_bottleneck_candidates), + routing_regime_status=routing_regime_status, + simulator_support_status_counts=simulator_support_status_counts, + ) + throughput_efficiency_tradeoff_status = _throughput_efficiency_tradeoff_status( + best_raw=best_raw, + best_risk_adjusted=best_risk_adjusted, + best_efficiency=best_efficiency, + best_risk_adjusted_efficiency=best_risk_adjusted_efficiency, + ) + throughput_efficiency_frontier_labels = _throughput_efficiency_frontier_labels( + feasible, + throughput_attr="score_tokens_per_sec", + efficiency_attr="score_tokens_per_gpu_per_sec", + ) + risk_adjusted_efficiency_frontier_labels = _throughput_efficiency_frontier_labels( + feasible, + throughput_attr="score_risk_adjusted_tokens_per_sec", + efficiency_attr="score_risk_adjusted_tokens_per_gpu_per_sec", + ) + raw_dominated_candidate_count = sum(1 for candidate in feasible if candidate.raw_dominated_by_label is not None) + risk_adjusted_dominated_candidate_count = sum( + 1 for candidate in feasible if candidate.risk_adjusted_dominated_by_label is not None + ) + parallelism_axis_comparisons = _parallelism_axis_comparisons(feasible) + parallelism_axis_coverage = _scenario_parallelism_axis_coverage(candidates) + scored_parallelism_axis_names = [ + coverage.axis for coverage in parallelism_axis_coverage if coverage.status == "scored_parallelism_axis" + ] + blocked_parallelism_axis_names = [ + coverage.axis for coverage in parallelism_axis_coverage if "blocked" in coverage.status + ] + confounded_parallelism_axis_names = [ + coverage.axis for coverage in parallelism_axis_coverage if coverage.status.startswith("confounded_") + ] + unscored_parallelism_axis_names = [ + coverage.axis for coverage in parallelism_axis_coverage if coverage.status == "unscored_parallelism_axis" + ] + missing_parallelism_axis_names = [ + coverage.axis for coverage in parallelism_axis_coverage if coverage.status == "missing_parallelism_axis" + ] + parallelism_axis_coverage_status_counts = _count_values([coverage.status for coverage in parallelism_axis_coverage]) + varied_parallelism_dimensions = _varied_candidate_dimensions(candidates, _SCENARIO_PARALLELISM_DIMENSIONS) + varied_workload_dimensions = _varied_candidate_dimensions(candidates, _SCENARIO_WORKLOAD_DIMENSIONS) + varied_runtime_dimensions = _varied_candidate_runtime_dimensions(candidates) + runtime_mismatch_dimensions = _candidate_runtime_mismatch_dimensions(candidates) + promotable_labels = {candidate.label for candidate in feasible if candidate.promotable} + interval_overlap_only_promotable_tie = ( + risk_adjusted_interval_overlap_status == "overlapping_best_interval" + and best_risk_adjusted is not None + and best_risk_adjusted.promotable + and bool(risk_adjusted_interval_overlap_contender_labels) + and set(risk_adjusted_interval_overlap_contender_labels) <= promotable_labels + ) + parallelism_optimality_status, parallelism_optimality_blockers = _parallelism_optimality_support( + unique_strategy_count=unique_strategy_count, + scored_strategy_count=scored_strategy_count, + memory_blocked_count=memory_blocked_count, + best_risk_adjusted=best_risk_adjusted, + best_promotable=best_promotable, + risk_adjusted_interval_overlap_status=risk_adjusted_interval_overlap_status, + interval_overlap_only_promotable_tie=interval_overlap_only_promotable_tie, + parallelism_tradeoff_status=parallelism_tradeoff_status, + parallelism_axis_coverage_status_counts=parallelism_axis_coverage_status_counts, + varied_runtime_dimensions=varied_runtime_dimensions, + simulator_support_status_counts=simulator_support_status_counts, + ) + scenario_capture_status, scenario_capture_blockers = _scenario_capture_support( + candidate_count=len(candidates), + scored_count=scored_count, + memory_blocked_count=memory_blocked_count, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + runtime_mismatch_dimensions=runtime_mismatch_dimensions, + parallelism_axis_coverage_status_counts=parallelism_axis_coverage_status_counts, + simulator_support_status_counts=simulator_support_status_counts, + ) + scenario_capture_gaps = _scenario_capture_gap_portfolio( + scenario_capture_blockers=scenario_capture_blockers, + candidate_count=len(candidates), + scored_count=scored_count, + memory_blocked_count=memory_blocked_count, + missing_parallelism_axis_names=missing_parallelism_axis_names, + blocked_parallelism_axis_names=blocked_parallelism_axis_names, + confounded_parallelism_axis_names=confounded_parallelism_axis_names, + unscored_parallelism_axis_names=unscored_parallelism_axis_names, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + runtime_mismatch_dimensions=runtime_mismatch_dimensions, + ) + scenario_capture_gap_status_counts = _count_values([gap.gap_status for gap in scenario_capture_gaps]) + scenario_capture_gap_required_measurements = _unique_in_order( + [gap.required_measurement for gap in scenario_capture_gaps] + ) + parallelism_boundary_groups = _scenario_parallelism_boundary_groups(candidates) + parallelism_boundary_status = _scenario_parallelism_boundary_status(parallelism_boundary_groups, candidates) + parallelism_boundary_axis_coverage = _scenario_parallelism_boundary_axis_coverage(parallelism_boundary_groups) + parallelism_boundary_measured_axis_names = [ + coverage.axis + for coverage in parallelism_boundary_axis_coverage + if coverage.status == "measured_parallelism_boundary_axis" + ] + parallelism_boundary_confounded_axis_names = [ + coverage.axis + for coverage in parallelism_boundary_axis_coverage + if coverage.status == "confounded_parallelism_boundary_axis" + ] + parallelism_boundary_missing_axis_names = [ + coverage.axis + for coverage in parallelism_boundary_axis_coverage + if coverage.status == "missing_parallelism_boundary_axis" + ] + parallelism_boundary_axis_coverage_status_counts = _count_values( + [coverage.status for coverage in parallelism_boundary_axis_coverage] + ) + parallelism_boundary_candidates = [ + candidate for candidate in candidates if _scenario_boundary_outcome(candidate) is not None + ] + parallelism_boundary_fit_count = sum( + 1 for candidate in parallelism_boundary_candidates if _scenario_boundary_outcome(candidate) == "fit" + ) + parallelism_boundary_failure_count = sum( + 1 for candidate in parallelism_boundary_candidates if _scenario_boundary_outcome(candidate) == "failure" + ) + parallelism_boundary_best_fit = max( + (group for group in parallelism_boundary_groups if group.best_fit_tokens_per_sec is not None), + key=lambda group: (group.best_fit_tokens_per_sec or float("-inf"), group.best_fit_label or ""), + default=None, + ) + parallelism_boundary_confounded_dimensions = sorted( + { + dimension + for group in parallelism_boundary_groups + for dimension in [*group.confounded_workload_dimensions, *group.confounded_runtime_dimensions] + } + ) + ( + parallelism_boundary_prediction_status, + parallelism_boundary_prediction_blockers, + ) = _scenario_parallelism_boundary_prediction_support( + parallelism_boundary_status=parallelism_boundary_status, + parallelism_boundary_fit_count=parallelism_boundary_fit_count, + parallelism_boundary_failure_count=parallelism_boundary_failure_count, + parallelism_boundary_measured_axis_names=parallelism_boundary_measured_axis_names, + parallelism_boundary_confounded_axis_names=parallelism_boundary_confounded_axis_names, + parallelism_boundary_missing_axis_names=parallelism_boundary_missing_axis_names, + parallelism_boundary_confounded_dimensions=parallelism_boundary_confounded_dimensions, + ) + isolated_parallelism_axis_comparison_count = sum( + 1 for comparison in parallelism_axis_comparisons if comparison.coupling_status == "isolated_axis_comparison" + ) + coupled_parallelism_axis_comparison_count = sum( + 1 for comparison in parallelism_axis_comparisons if comparison.coupling_status == "coupled_axis_comparison" + ) + parallelism_axis_interval_overlap_count = sum( + 1 + for comparison in parallelism_axis_comparisons + if comparison.risk_adjusted_interval_overlap_status == "overlapping_best_interval" + ) + scaling_candidates = [ + candidate + for candidate in feasible + if candidate.scaling_efficiency is not None + and candidate.scaling_gpu_ratio is not None + and candidate.scaling_gpu_ratio > 1.0 + ] + risk_adjusted_scaling_candidates = [ + candidate for candidate in scaling_candidates if candidate.risk_adjusted_scaling_efficiency is not None + ] + best_scaling_efficiency = ( + max(scaling_candidates, key=lambda candidate: (candidate.scaling_efficiency or float("-inf"), candidate.label)) + if scaling_candidates + else None + ) + best_risk_adjusted_scaling_efficiency = ( + max( + risk_adjusted_scaling_candidates, + key=lambda candidate: (candidate.risk_adjusted_scaling_efficiency or float("-inf"), candidate.label), + ) + if risk_adjusted_scaling_candidates + else None + ) + scaling_efficiencies = [ + candidate.scaling_efficiency for candidate in scaling_candidates if candidate.scaling_efficiency is not None + ] + risk_adjusted_scaling_efficiencies = [ + candidate.risk_adjusted_scaling_efficiency + for candidate in risk_adjusted_scaling_candidates + if candidate.risk_adjusted_scaling_efficiency is not None + ] + same_workload_scaling_status = _same_workload_scaling_status(scaling_candidates) + promotion_readiness_status = _scenario_promotion_readiness_status( + best_raw=best_raw, + best_risk_adjusted=best_risk_adjusted, + best_promotable=best_promotable, + ) + promotable_raw_gap_tokens_per_sec = _candidate_score_gap( + best_raw, + best_promotable, + "score_tokens_per_sec", + ) + promotable_raw_gap_percentage = _candidate_score_gap_percentage( + best_raw, + "score_tokens_per_sec", + promotable_raw_gap_tokens_per_sec, + ) + promotable_risk_adjusted_gap_tokens_per_sec = _candidate_score_gap( + best_risk_adjusted, + best_promotable, + "score_risk_adjusted_tokens_per_sec", + ) + promotable_risk_adjusted_gap_percentage = _candidate_score_gap_percentage( + best_risk_adjusted, + "score_risk_adjusted_tokens_per_sec", + promotable_risk_adjusted_gap_tokens_per_sec, + ) + measurement_readiness_status, measurement_rationale = _scenario_measurement_guidance( + candidates=candidates, + scored_count=scored_count, + memory_blocked_count=memory_blocked_count, + parallelism_tradeoff_status=parallelism_tradeoff_status, + parallelism_optimality_status=parallelism_optimality_status, + parallelism_optimality_blockers=parallelism_optimality_blockers, + parallelism_boundary_status=parallelism_boundary_status, + parallelism_boundary_prediction_status=parallelism_boundary_prediction_status, + parallelism_boundary_prediction_blockers=parallelism_boundary_prediction_blockers, + throughput_efficiency_tradeoff_status=throughput_efficiency_tradeoff_status, + throughput_efficiency_frontier_count=len(throughput_efficiency_frontier_labels), + risk_adjusted_efficiency_frontier_count=len(risk_adjusted_efficiency_frontier_labels), + raw_dominated_candidate_count=raw_dominated_candidate_count, + risk_adjusted_dominated_candidate_count=risk_adjusted_dominated_candidate_count, + parallelism_axis_comparison_count=len(parallelism_axis_comparisons), + isolated_parallelism_axis_comparison_count=isolated_parallelism_axis_comparison_count, + coupled_parallelism_axis_comparison_count=coupled_parallelism_axis_comparison_count, + parallelism_axis_interval_overlap_count=parallelism_axis_interval_overlap_count, + blocked_parallelism_axis_names=blocked_parallelism_axis_names, + confounded_parallelism_axis_names=confounded_parallelism_axis_names, + same_workload_scaling_status=same_workload_scaling_status, + same_workload_scaling_candidate_count=len(scaling_candidates), + min_scaling_efficiency=round(min(scaling_efficiencies), 3) if scaling_efficiencies else None, + memory_coverage_status_counts=memory_coverage_status_counts, + timing_coverage_status_counts=timing_coverage_status_counts, + max_memory_residual_gb=max_memory_residual_gb, + phase_bottleneck_bucket_counts=phase_bottleneck_bucket_counts, + max_phase_bottleneck_share=max_phase_bottleneck_share, + max_phase_bottleneck_half_speedup_delta_pct=max_phase_bottleneck_half_speedup_delta_pct, + memory_bottleneck_bucket_counts=memory_bottleneck_bucket_counts, + max_memory_bottleneck_fraction_of_peak=max_memory_bottleneck_fraction, + high_uncertainty_candidate_count=high_uncertainty_candidate_count, + max_prediction_uncertainty_fraction=max_prediction_uncertainty_fraction, + risk_adjusted_interval_overlap_status=risk_adjusted_interval_overlap_status, + risk_adjusted_interval_overlap_contender_count=risk_adjusted_interval_overlap_contender_count, + risk_adjusted_interval_best_vs_next_margin_tokens_per_sec=risk_adjusted_interval_best_vs_next_margin, + routing_regime_status=routing_regime_status, + routing_regime_counts=routing_regime_counts, + cross_model_analog_support_status=cross_model_analog_support_status, + cross_model_analog_factor_status=cross_model_analog_factor_status, + cross_model_analog_unique_factor_count=cross_model_analog_unique_factor_count, + cross_model_analog_unique_target_runtime_signature_count=( + cross_model_analog_unique_target_runtime_signature_count + ), + cross_model_analog_scored_varied_parallelism_dimensions=( + cross_model_analog_scored_varied_parallelism_dimensions + ), + cross_model_analog_scored_varied_workload_dimensions=cross_model_analog_scored_varied_workload_dimensions, + model_generalization_status=model_generalization_status, + model_generalization_blockers=model_generalization_blockers, + promotion_readiness_status=promotion_readiness_status, + promotable_raw_gap_tokens_per_sec=promotable_raw_gap_tokens_per_sec, + promotable_risk_adjusted_gap_tokens_per_sec=promotable_risk_adjusted_gap_tokens_per_sec, + best_risk_adjusted=best_risk_adjusted, + best_next_measurement=best_next_measurement, + best_promotable=best_promotable, + ) + ( + measurement_candidate_labels, + measurement_candidate_reasons, + measurement_candidate_priority_scores, + measurement_candidate_priority_per_gpu, + measurement_candidate_cost_gpus, + measurement_candidate_priority_factors, + ) = _measurement_portfolio( + candidates=candidates, + throughput_efficiency_frontier_labels=throughput_efficiency_frontier_labels, + risk_adjusted_efficiency_frontier_labels=risk_adjusted_efficiency_frontier_labels, + parallelism_axis_coverage=parallelism_axis_coverage, + parallelism_tradeoff_status=parallelism_tradeoff_status, + cross_model_analog_support_status=cross_model_analog_support_status, + cross_model_analog_prediction_interval_selectivity_status=( + cross_model_analog_prediction_interval_selectivity_status + ), + cross_model_analog_prediction_interval_top_labels=cross_model_analog_prediction_interval_top_labels, + same_workload_scaling_status=same_workload_scaling_status, + best_raw=best_raw, + best_risk_adjusted=best_risk_adjusted, + best_efficiency=best_efficiency, + best_risk_adjusted_efficiency=best_risk_adjusted_efficiency, + best_next_measurement=best_next_measurement, + best_promotable=best_promotable, + ) + measurement_candidate_config_overrides = _measurement_candidate_config_overrides(measurement_candidate_reasons) + ( + measurement_portfolio_coverage_status, + measurement_portfolio_coverage_blockers, + measurement_portfolio_reason_category_counts, + measurement_portfolio_parallelism_axis_gap_names, + measurement_portfolio_cross_model_analog_count, + ) = _measurement_portfolio_coverage_status( + candidate_reasons=measurement_candidate_reasons, + best_next_measurement=best_next_measurement, + parallelism_tradeoff_status=parallelism_tradeoff_status, + throughput_efficiency_tradeoff_status=throughput_efficiency_tradeoff_status, + same_workload_scaling_status=same_workload_scaling_status, + cross_model_analog_support_status=cross_model_analog_support_status, + ) + validation_actions = _scenario_validation_actions( + candidate_reasons=measurement_candidate_reasons, + candidate_priority_scores=measurement_candidate_priority_scores, + candidate_priority_per_gpu=measurement_candidate_priority_per_gpu, + candidate_cost_gpus=measurement_candidate_cost_gpus, + candidate_config_overrides=measurement_candidate_config_overrides, + ) + validation_action_status_counts = _count_values([action.action_status for action in validation_actions]) + validation_action_required_measurements = _unique_in_order( + [action.required_measurement for action in validation_actions] + ) + measurement_rationale = [ + *measurement_rationale, + f"scenario_capture={scenario_capture_status}", + f"benchmark_support={benchmark_support.support_status}", + f"scenario_prediction_fidelity={scenario_prediction_fidelity_status}", + f"measurement_portfolio_coverage={measurement_portfolio_coverage_status}", + "simulator_support_status_counts=" + + ",".join(f"{status}:{count}" for status, count in simulator_support_status_counts.items()), + ] + if benchmark_support.support_blockers: + measurement_rationale.append(f"benchmark_support_blockers={','.join(benchmark_support.support_blockers)}") + if scenario_capture_blockers: + measurement_rationale.append(f"scenario_capture_blockers={','.join(scenario_capture_blockers)}") + if scenario_capture_gaps: + measurement_rationale.append( + "scenario_capture_gaps=" + + ",".join(f"{status}:{count}" for status, count in scenario_capture_gap_status_counts.items()) + ) + measurement_rationale.append( + "scenario_capture_required_measurements=" + ",".join(scenario_capture_gap_required_measurements) + ) + if varied_runtime_dimensions: + measurement_rationale.append(f"varied_runtime_dimensions={','.join(varied_runtime_dimensions)}") + if runtime_mismatch_dimensions: + measurement_rationale.append(f"runtime_mismatch_dimensions={','.join(runtime_mismatch_dimensions)}") + if cross_model_analog_prediction_interval_selectivity_status not in { + "not_used", + "no_scored_cross_model_candidates", + }: + measurement_rationale.append( + "cross_model_analog_prediction_interval_selectivity=" + f"{cross_model_analog_prediction_interval_selectivity_status}" + ) + if scenario_prediction_fidelity_blockers: + measurement_rationale.append( + f"scenario_prediction_fidelity_blockers={','.join(scenario_prediction_fidelity_blockers)}" + ) + if measurement_portfolio_coverage_blockers: + measurement_rationale.append( + f"measurement_portfolio_coverage_blockers={','.join(measurement_portfolio_coverage_blockers)}" + ) + if measurement_portfolio_parallelism_axis_gap_names: + measurement_rationale.append( + f"measurement_portfolio_parallelism_axis_gaps={','.join(measurement_portfolio_parallelism_axis_gap_names)}" + ) + if validation_actions: + measurement_rationale.append( + "validation_actions=" + + ",".join(f"{status}:{count}" for status, count in validation_action_status_counts.items()) + ) + measurement_rationale.append( + "validation_action_required_measurements=" + ",".join(validation_action_required_measurements) + ) + measurement_portfolio_total_gpu_count = sum(measurement_candidate_cost_gpus.values()) + prediction_confidence_counts = _count_values([candidate.prediction_confidence for candidate in candidates]) + calibration_scope_counts = _count_values([candidate.calibration_scope for candidate in candidates]) + scenario_readiness = _scenario_readiness( + candidate_count=len(candidates), + scored_count=scored_count, + unscored_count=len(candidates) - scored_count, + memory_blocked_count=memory_blocked_count, + unique_parallelism_strategy_count=unique_strategy_count, + scored_parallelism_strategy_count=scored_strategy_count, + promotable_parallelism_strategy_count=promotable_strategy_count, + scenario_capture_status=scenario_capture_status, + scenario_capture_blockers=scenario_capture_blockers, + scenario_prediction_fidelity_status=scenario_prediction_fidelity_status, + scenario_prediction_fidelity_blockers=scenario_prediction_fidelity_blockers, + parallelism_optimality_status=parallelism_optimality_status, + parallelism_optimality_blockers=parallelism_optimality_blockers, + model_generalization_status=model_generalization_status, + model_generalization_blockers=model_generalization_blockers, + measurement_readiness_status=measurement_readiness_status, + measurement_portfolio_coverage_status=measurement_portfolio_coverage_status, + measurement_portfolio_coverage_blockers=measurement_portfolio_coverage_blockers, + scenario_capture_gaps=scenario_capture_gaps, + scenario_capture_gap_status_counts=scenario_capture_gap_status_counts, + scenario_capture_gap_required_measurements=scenario_capture_gap_required_measurements, + validation_actions=validation_actions, + validation_action_status_counts=validation_action_status_counts, + validation_action_required_measurements=validation_action_required_measurements, + validation_action_total_gpu_count=measurement_portfolio_total_gpu_count, + measurement_candidate_labels=measurement_candidate_labels, + measurement_portfolio_total_gpu_count=measurement_portfolio_total_gpu_count, + measurement_portfolio_parallelism_axis_gap_names=measurement_portfolio_parallelism_axis_gap_names, + measurement_portfolio_cross_model_analog_count=measurement_portfolio_cross_model_analog_count, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + runtime_mismatch_dimensions=runtime_mismatch_dimensions, + parallelism_axis_coverage_status_counts=parallelism_axis_coverage_status_counts, + scored_parallelism_axis_names=scored_parallelism_axis_names, + blocked_parallelism_axis_names=blocked_parallelism_axis_names, + confounded_parallelism_axis_names=confounded_parallelism_axis_names, + unscored_parallelism_axis_names=unscored_parallelism_axis_names, + missing_parallelism_axis_names=missing_parallelism_axis_names, + simulator_support_status_counts=simulator_support_status_counts, + prediction_confidence_counts=prediction_confidence_counts, + calibration_scope_counts=calibration_scope_counts, + memory_coverage_status_counts=memory_coverage_status_counts, + timing_coverage_status_counts=timing_coverage_status_counts, + cross_model_analog_support_status=cross_model_analog_support_status, + cross_model_analog_candidate_count=cross_model_analog_candidate_count, + cross_model_analog_scored_count=cross_model_analog_scored_count, + benchmark_support=benchmark_support, + ) + return ScenarioDecisionSummary( + candidate_count=len(candidates), + scored_count=scored_count, + unscored_count=len(candidates) - scored_count, + feasible_count=scored_count, + promotable_count=sum(1 for candidate in feasible if candidate.promotable), + requires_remeasurement_count=sum( + 1 for candidate in candidates if "requires_remeasurement" in candidate.risk_flags + ), + memory_blocked_count=memory_blocked_count, + unique_parallelism_strategy_count=unique_strategy_count, + best_raw_label=best_raw.label if best_raw is not None else None, + best_raw_score_tokens_per_sec=best_raw.score_tokens_per_sec if best_raw is not None else None, + best_risk_adjusted_label=best_risk_adjusted.label if best_risk_adjusted is not None else None, + best_risk_adjusted_score_tokens_per_sec=( + best_risk_adjusted.score_risk_adjusted_tokens_per_sec if best_risk_adjusted is not None else None + ), + best_efficiency_label=best_efficiency.label if best_efficiency is not None else None, + best_efficiency_score_tokens_per_gpu_per_sec=( + best_efficiency.score_tokens_per_gpu_per_sec if best_efficiency is not None else None + ), + best_risk_adjusted_efficiency_label=( + best_risk_adjusted_efficiency.label if best_risk_adjusted_efficiency is not None else None + ), + best_risk_adjusted_efficiency_score_tokens_per_gpu_per_sec=( + best_risk_adjusted_efficiency.score_risk_adjusted_tokens_per_gpu_per_sec + if best_risk_adjusted_efficiency is not None + else None + ), + best_next_measurement_label=best_next_measurement.label if best_next_measurement is not None else None, + best_next_measurement_score_tokens_per_sec=( + best_next_measurement.score_tokens_per_sec if best_next_measurement is not None else None + ), + best_promotable_label=best_promotable.label if best_promotable is not None else None, + candidate_model_ref_counts=_count_candidate_model_refs(candidates), + scored_model_ref_counts=_count_candidate_model_refs(feasible), + candidate_world_size_counts=_count_candidate_topology_values(candidates, "world_size"), + scored_world_size_counts=_count_candidate_topology_values(feasible, "world_size"), + candidate_sequence_length_counts=_count_candidate_topology_values(candidates, "sample_packing_sequence_len"), + scored_sequence_length_counts=_count_candidate_topology_values(feasible, "sample_packing_sequence_len"), + candidate_global_batch_size_counts=_count_candidate_topology_values(candidates, "global_batch_size"), + scored_global_batch_size_counts=_count_candidate_topology_values(feasible, "global_batch_size"), + candidate_runtime_signature_counts=_count_candidate_runtime_signatures(candidates), + scored_runtime_signature_counts=_count_candidate_runtime_signatures(feasible), + best_promotable_score_tokens_per_sec=( + best_promotable.score_tokens_per_sec if best_promotable is not None else None + ), + best_promotable_score_risk_adjusted_tokens_per_sec=( + best_promotable.score_risk_adjusted_tokens_per_sec if best_promotable is not None else None + ), + promotable_raw_gap_tokens_per_sec=promotable_raw_gap_tokens_per_sec, + promotable_raw_gap_percentage=promotable_raw_gap_percentage, + promotable_risk_adjusted_gap_tokens_per_sec=promotable_risk_adjusted_gap_tokens_per_sec, + promotable_risk_adjusted_gap_percentage=promotable_risk_adjusted_gap_percentage, + promotion_readiness_status=promotion_readiness_status, + throughput_efficiency_frontier_labels=throughput_efficiency_frontier_labels, + risk_adjusted_efficiency_frontier_labels=risk_adjusted_efficiency_frontier_labels, + throughput_efficiency_frontier_count=len(throughput_efficiency_frontier_labels), + risk_adjusted_efficiency_frontier_count=len(risk_adjusted_efficiency_frontier_labels), + raw_dominated_candidate_count=raw_dominated_candidate_count, + risk_adjusted_dominated_candidate_count=risk_adjusted_dominated_candidate_count, + throughput_efficiency_tradeoff_status=throughput_efficiency_tradeoff_status, + same_workload_scaling_status=same_workload_scaling_status, + same_workload_scaling_group_count=len( + {candidate.scaling_baseline_label for candidate in scaling_candidates if candidate.scaling_baseline_label} + ), + same_workload_scaling_candidate_count=len(scaling_candidates), + best_scaling_efficiency_label=(best_scaling_efficiency.label if best_scaling_efficiency is not None else None), + best_scaling_efficiency=( + best_scaling_efficiency.scaling_efficiency if best_scaling_efficiency is not None else None + ), + mean_scaling_efficiency=( + round(sum(scaling_efficiencies) / len(scaling_efficiencies), 3) if scaling_efficiencies else None + ), + min_scaling_efficiency=round(min(scaling_efficiencies), 3) if scaling_efficiencies else None, + best_risk_adjusted_scaling_efficiency_label=( + best_risk_adjusted_scaling_efficiency.label if best_risk_adjusted_scaling_efficiency is not None else None + ), + best_risk_adjusted_scaling_efficiency=( + best_risk_adjusted_scaling_efficiency.risk_adjusted_scaling_efficiency + if best_risk_adjusted_scaling_efficiency is not None + else None + ), + mean_risk_adjusted_scaling_efficiency=( + round(sum(risk_adjusted_scaling_efficiencies) / len(risk_adjusted_scaling_efficiencies), 3) + if risk_adjusted_scaling_efficiencies + else None + ), + min_risk_adjusted_scaling_efficiency=( + round(min(risk_adjusted_scaling_efficiencies), 3) if risk_adjusted_scaling_efficiencies else None + ), + measurement_readiness_status=measurement_readiness_status, + measurement_rationale=measurement_rationale, + measurement_candidate_count=len(measurement_candidate_labels), + measurement_candidate_labels=measurement_candidate_labels, + measurement_candidate_reasons=measurement_candidate_reasons, + measurement_candidate_priority_scores=measurement_candidate_priority_scores, + measurement_candidate_priority_per_gpu=measurement_candidate_priority_per_gpu, + measurement_candidate_cost_gpus=measurement_candidate_cost_gpus, + measurement_candidate_priority_factors=measurement_candidate_priority_factors, + measurement_candidate_config_overrides=measurement_candidate_config_overrides, + measurement_portfolio_total_gpu_count=measurement_portfolio_total_gpu_count, + measurement_portfolio_max_priority_score=( + max(measurement_candidate_priority_scores.values()) if measurement_candidate_priority_scores else None + ), + measurement_portfolio_max_priority_label=_max_score_label(measurement_candidate_priority_scores), + measurement_portfolio_max_priority_per_gpu=( + max(measurement_candidate_priority_per_gpu.values()) if measurement_candidate_priority_per_gpu else None + ), + measurement_portfolio_max_priority_per_gpu_label=_max_score_label(measurement_candidate_priority_per_gpu), + measurement_portfolio_coverage_status=measurement_portfolio_coverage_status, + measurement_portfolio_coverage_blockers=measurement_portfolio_coverage_blockers, + measurement_portfolio_reason_category_counts=measurement_portfolio_reason_category_counts, + measurement_portfolio_parallelism_axis_gap_names=measurement_portfolio_parallelism_axis_gap_names, + measurement_portfolio_cross_model_analog_count=measurement_portfolio_cross_model_analog_count, + validation_action_count=len(validation_actions), + validation_action_status_counts=validation_action_status_counts, + validation_action_required_measurements=validation_action_required_measurements, + validation_action_total_gpu_count=measurement_portfolio_total_gpu_count, + validation_actions=validation_actions, + max_calibration_distance=round(max(distances), 3) if distances else None, + max_calibration_distance_label=_max_positive_candidate_label(candidates, "calibration_distance"), + mean_scored_calibration_distance=( + round(sum(scored_distances) / len(scored_distances), 3) if scored_distances else None + ), + high_uncertainty_candidate_count=high_uncertainty_candidate_count, + max_prediction_uncertainty_fraction=max_prediction_uncertainty_fraction, + max_prediction_uncertainty_fraction_label=_max_positive_candidate_label( + candidates, + "prediction_uncertainty_fraction", + ), + mean_scored_prediction_uncertainty_fraction=( + round(sum(scored_uncertainty_fractions) / len(scored_uncertainty_fractions), 3) + if scored_uncertainty_fractions + else None + ), + risk_adjusted_interval_overlap_status=risk_adjusted_interval_overlap_status, + risk_adjusted_interval_overlap_contender_count=risk_adjusted_interval_overlap_contender_count, + risk_adjusted_interval_overlap_contender_labels=risk_adjusted_interval_overlap_contender_labels, + risk_adjusted_interval_best_vs_next_margin_tokens_per_sec=risk_adjusted_interval_best_vs_next_margin, + parallelism_tradeoff_status=parallelism_tradeoff_status, + parallelism_optimality_status=parallelism_optimality_status, + parallelism_optimality_blockers=parallelism_optimality_blockers, + scored_parallelism_strategy_count=scored_strategy_count, + promotable_parallelism_strategy_count=promotable_strategy_count, + requires_remeasurement_parallelism_strategy_count=requires_remeasurement_strategy_count, + parallelism_axis_comparison_count=len(parallelism_axis_comparisons), + isolated_parallelism_axis_comparison_count=isolated_parallelism_axis_comparison_count, + coupled_parallelism_axis_comparison_count=coupled_parallelism_axis_comparison_count, + parallelism_axis_interval_overlap_count=parallelism_axis_interval_overlap_count, + parallelism_axis_comparisons=parallelism_axis_comparisons, + scored_parallelism_axis_names=scored_parallelism_axis_names, + blocked_parallelism_axis_names=blocked_parallelism_axis_names, + confounded_parallelism_axis_names=confounded_parallelism_axis_names, + unscored_parallelism_axis_names=unscored_parallelism_axis_names, + missing_parallelism_axis_names=missing_parallelism_axis_names, + parallelism_axis_coverage_status_counts=parallelism_axis_coverage_status_counts, + parallelism_axis_coverage=parallelism_axis_coverage, + parallelism_boundary_status=parallelism_boundary_status, + parallelism_boundary_prediction_status=parallelism_boundary_prediction_status, + parallelism_boundary_prediction_blockers=parallelism_boundary_prediction_blockers, + parallelism_boundary_group_count=len(parallelism_boundary_groups), + parallelism_boundary_candidate_count=len(parallelism_boundary_candidates), + parallelism_boundary_fit_count=parallelism_boundary_fit_count, + parallelism_boundary_failure_count=parallelism_boundary_failure_count, + parallelism_boundary_best_fit_label=( + parallelism_boundary_best_fit.best_fit_label if parallelism_boundary_best_fit is not None else None + ), + parallelism_boundary_confounded_dimensions=parallelism_boundary_confounded_dimensions, + parallelism_boundary_measured_axis_names=parallelism_boundary_measured_axis_names, + parallelism_boundary_confounded_axis_names=parallelism_boundary_confounded_axis_names, + parallelism_boundary_missing_axis_names=parallelism_boundary_missing_axis_names, + parallelism_boundary_axis_coverage_status_counts=parallelism_boundary_axis_coverage_status_counts, + parallelism_boundary_axis_coverage=parallelism_boundary_axis_coverage, + parallelism_boundary_groups=parallelism_boundary_groups, + cross_model_analog_support_status=cross_model_analog_support_status, + cross_model_analog_candidate_count=cross_model_analog_candidate_count, + cross_model_analog_scored_count=cross_model_analog_scored_count, + cross_model_analog_unique_prediction_count=cross_model_analog_unique_prediction_count, + cross_model_analog_unique_matched_label_count=cross_model_analog_unique_matched_label_count, + cross_model_analog_unique_target_strategy_count=cross_model_analog_unique_target_strategy_count, + cross_model_analog_unique_target_runtime_signature_count=( + cross_model_analog_unique_target_runtime_signature_count + ), + cross_model_analog_scored_varied_parallelism_dimensions=( + cross_model_analog_scored_varied_parallelism_dimensions + ), + cross_model_analog_scored_varied_workload_dimensions=cross_model_analog_scored_varied_workload_dimensions, + cross_model_analog_factor_status=cross_model_analog_factor_status, + cross_model_analog_unique_factor_count=cross_model_analog_unique_factor_count, + cross_model_analog_factor_ranges=cross_model_analog_factor_ranges, + cross_model_analog_prediction_interval_top_count=cross_model_analog_prediction_interval_top_count, + cross_model_analog_prediction_interval_top_fraction=cross_model_analog_prediction_interval_top_fraction, + cross_model_analog_prediction_interval_top_labels=cross_model_analog_prediction_interval_top_labels, + cross_model_analog_prediction_interval_selectivity_status=( + cross_model_analog_prediction_interval_selectivity_status + ), + model_generalization_status=model_generalization_status, + model_generalization_blockers=model_generalization_blockers, + scenario_capture_status=scenario_capture_status, + scenario_capture_blockers=scenario_capture_blockers, + scenario_capture_gap_count=len(scenario_capture_gaps), + scenario_capture_gap_status_counts=scenario_capture_gap_status_counts, + scenario_capture_gap_required_measurements=scenario_capture_gap_required_measurements, + scenario_capture_gaps=scenario_capture_gaps, + benchmark_support=benchmark_support, + scenario_prediction_fidelity_status=scenario_prediction_fidelity_status, + scenario_prediction_fidelity_blockers=scenario_prediction_fidelity_blockers, + varied_parallelism_dimensions=varied_parallelism_dimensions, + varied_workload_dimensions=varied_workload_dimensions, + varied_runtime_dimensions=varied_runtime_dimensions, + runtime_mismatch_dimensions=runtime_mismatch_dimensions, + prediction_confidence_counts=prediction_confidence_counts, + calibration_scope_counts=calibration_scope_counts, + memory_basis_counts=_count_values([candidate.memory_basis for candidate in candidates]), + memory_coverage_status_counts=memory_coverage_status_counts, + simulator_support_status_counts=simulator_support_status_counts, + simulator_support_blocker_counts=simulator_support_blocker_counts, + timing_coverage_status_counts=timing_coverage_status_counts, + max_estimated_memory_residual_gb=max_memory_residual_gb, + max_estimated_memory_residual_gb_label=_max_positive_candidate_label( + candidates, + "estimated_memory_residual_gb", + ), + max_estimated_memory_residual_fraction_of_peak=( + round(max(memory_residual_fractions), 3) if memory_residual_fractions else None + ), + max_estimated_memory_residual_fraction_of_peak_label=_max_positive_candidate_label( + candidates, + "estimated_memory_residual_fraction_of_peak", + ), + phase_bottleneck_candidate_count=len(phase_bottleneck_candidates), + phase_bottleneck_bucket_counts=phase_bottleneck_bucket_counts, + phase_bottleneck_phase_counts=phase_bottleneck_phase_counts, + max_phase_bottleneck_share=max_phase_bottleneck_share, + max_phase_bottleneck_share_label=( + max_phase_bottleneck_candidate.label if max_phase_bottleneck_candidate is not None else None + ), + max_phase_bottleneck_phase=( + max_phase_bottleneck_candidate.phase_bottleneck_phase + if max_phase_bottleneck_candidate is not None + else None + ), + max_phase_bottleneck_bucket=( + max_phase_bottleneck_candidate.phase_bottleneck_bucket + if max_phase_bottleneck_candidate is not None + else None + ), + phase_bottleneck_half_speedup_candidate_count=len(phase_bottleneck_half_speedup_candidates), + max_phase_bottleneck_half_speedup_delta_pct=max_phase_bottleneck_half_speedup_delta_pct, + max_phase_bottleneck_half_speedup_delta_label=( + max_phase_bottleneck_half_speedup_candidate.label + if max_phase_bottleneck_half_speedup_candidate is not None + else None + ), + max_phase_bottleneck_half_speedup_phase=( + max_phase_bottleneck_half_speedup_candidate.phase_bottleneck_phase + if max_phase_bottleneck_half_speedup_candidate is not None + else None + ), + max_phase_bottleneck_half_speedup_bucket=( + max_phase_bottleneck_half_speedup_candidate.phase_bottleneck_bucket + if max_phase_bottleneck_half_speedup_candidate is not None + else None + ), + memory_bottleneck_candidate_count=len(memory_bottleneck_candidates), + memory_bottleneck_bucket_counts=memory_bottleneck_bucket_counts, + memory_bottleneck_phase_counts=memory_bottleneck_phase_counts, + max_memory_bottleneck_fraction_of_peak=max_memory_bottleneck_fraction, + max_memory_bottleneck_fraction_label=( + max_memory_bottleneck_candidate.label if max_memory_bottleneck_candidate is not None else None + ), + max_memory_bottleneck_phase=( + max_memory_bottleneck_candidate.memory_bottleneck_phase + if max_memory_bottleneck_candidate is not None + else None + ), + max_memory_bottleneck_bucket=( + max_memory_bottleneck_candidate.memory_bottleneck_bucket + if max_memory_bottleneck_candidate is not None + else None + ), + cross_node_dimension_counts=_cross_node_dimension_counts(candidates), + feasibility_status_counts=_count_values([candidate.feasibility_status for candidate in candidates]), + routing_regime_status=routing_regime_status, + routing_regime_counts=routing_regime_counts, + recommendation_counts=_count_values([candidate.recommendation for candidate in candidates]), + risk_flag_counts=_count_values(risk_flags), + scenario_readiness=scenario_readiness, + ) + + +def _max_positive_candidate_by_field( + candidates: list[ScenarioCandidate], + field_name: str, +) -> ScenarioCandidate | None: + max_row: tuple[float, str] | None = None + max_candidate: ScenarioCandidate | None = None + for candidate in candidates: + value = getattr(candidate, field_name) + if not isinstance(value, int | float) or value <= 0: + continue + row = (float(value), candidate.label) + if max_row is None or row > max_row: + max_row = row + max_candidate = candidate + return max_candidate + + +def _max_score_label(scores: dict[str, float]) -> str | None: + return max(scores.items(), key=lambda item: (item[1], item[0]), default=(None, 0.0))[0] + + +def _max_positive_candidate_label(candidates: list[ScenarioCandidate], field_name: str) -> str | None: + max_candidate = _max_positive_candidate_by_field(candidates, field_name) + return max_candidate.label if max_candidate is not None else None + + +def _csv(values: list[int] | None) -> str | None: + if values is None: + return None + return ",".join(str(value) for value in values) + + +def _append_cli_option(args: list[str], option: str, value: object | None) -> None: + if value is None: + return + args.extend([option, str(value)]) + + +def _scenario_planner_context( + *, + base_path: Path, + benchmark_dir: str | Path | None, + supplemental_benchmark_dirs: list[str | Path] | None, + analog_benchmark_dirs: list[str | Path] | None, + requested_world_size: int | None, + candidate_world_sizes: list[int], + resolved_local_world_size: int, + micro_batch_values: list[int], + gradient_accumulation_values: list[int], + sample_packing_sequence_values: list[int], + expert_parallel_sizes: list[int] | None, + tensor_parallel_sizes: list[int] | None, + pipeline_parallel_sizes: list[int] | None, + ulysses_parallel_sizes: list[int] | None, + ringattn_parallel_sizes: list[int] | None, + topology_sweep: str, + balanced_routing: bool, + device_memory_limit_gb: float, + memory_safety_factor: float, +) -> dict[str, Any]: + cli_args = [ + "xorl-sim-plan", + "--config", + str(base_path), + ] + _append_cli_option(cli_args, "--benchmark-dir", benchmark_dir) + for supplemental_dir in supplemental_benchmark_dirs or []: + _append_cli_option(cli_args, "--supplemental-benchmark-dir", supplemental_dir) + for analog_dir in analog_benchmark_dirs or []: + _append_cli_option(cli_args, "--analog-benchmark-dir", analog_dir) + _append_cli_option(cli_args, "--world-size", requested_world_size) + _append_cli_option(cli_args, "--world-sizes", _csv(candidate_world_sizes)) + _append_cli_option(cli_args, "--local-world-size", resolved_local_world_size) + _append_cli_option(cli_args, "--micro-batch-sizes", _csv(micro_batch_values)) + _append_cli_option(cli_args, "--gradient-accumulation-steps", _csv(gradient_accumulation_values)) + _append_cli_option(cli_args, "--sample-packing-sequence-lengths", _csv(sample_packing_sequence_values)) + _append_cli_option(cli_args, "--expert-parallel-sizes", _csv(expert_parallel_sizes)) + _append_cli_option(cli_args, "--tensor-parallel-sizes", _csv(tensor_parallel_sizes)) + _append_cli_option(cli_args, "--pipeline-parallel-sizes", _csv(pipeline_parallel_sizes)) + _append_cli_option(cli_args, "--ulysses-parallel-sizes", _csv(ulysses_parallel_sizes)) + _append_cli_option(cli_args, "--ringattn-parallel-sizes", _csv(ringattn_parallel_sizes)) + _append_cli_option(cli_args, "--topology-sweep", topology_sweep) + if balanced_routing: + cli_args.append("--balanced-routing") + _append_cli_option(cli_args, "--device-memory-limit-gb", device_memory_limit_gb) + _append_cli_option(cli_args, "--memory-safety-factor", memory_safety_factor) + write_args = [*cli_args, "--write-measurement-configs", ""] + return { + "requested_args": { + "benchmark_dir": str(benchmark_dir) if benchmark_dir is not None else None, + "supplemental_benchmark_dirs": [ + str(supplemental_dir) for supplemental_dir in supplemental_benchmark_dirs or [] + ], + "analog_benchmark_dirs": [str(analog_dir) for analog_dir in analog_benchmark_dirs or []], + "world_size": requested_world_size, + "world_sizes": candidate_world_sizes, + "local_world_size": resolved_local_world_size, + "micro_batch_sizes": micro_batch_values, + "gradient_accumulation_steps": gradient_accumulation_values, + "sample_packing_sequence_lengths": sample_packing_sequence_values, + "expert_parallel_sizes": expert_parallel_sizes, + "tensor_parallel_sizes": tensor_parallel_sizes, + "pipeline_parallel_sizes": pipeline_parallel_sizes, + "ulysses_parallel_sizes": ulysses_parallel_sizes, + "ringattn_parallel_sizes": ringattn_parallel_sizes, + "topology_sweep": topology_sweep, + "balanced_routing": balanced_routing, + "device_memory_limit_gb": device_memory_limit_gb, + "memory_safety_factor": memory_safety_factor, + }, + "measurement_config_command": write_args, + } + + +def plan_scenario( + base_config_path: str | Path, + *, + benchmark_dir: str | Path | None = None, + supplemental_benchmark_dirs: list[str | Path] | None = None, + analog_benchmark_dirs: list[str | Path] | None = None, + world_size: int | None = None, + world_sizes: list[int] | None = None, + local_world_size: int | None = None, + micro_batch_sizes: list[int] | None = None, + gradient_accumulation_steps: list[int] | None = None, + sample_packing_sequence_lengths: list[int] | None = None, + expert_parallel_sizes: list[int] | None = None, + tensor_parallel_sizes: list[int] | None = None, + pipeline_parallel_sizes: list[int] | None = None, + ulysses_parallel_sizes: list[int] | None = None, + ringattn_parallel_sizes: list[int] | None = None, + topology_sweep: str = "base", + balanced_routing: bool = False, + device_memory_limit_gb: float = 80.0, + memory_safety_factor: float = 1.15, +) -> ScenarioReport: + if topology_sweep not in {"base", "auto"}: + raise ValueError("topology_sweep must be 'base' or 'auto'") + base_path = Path(base_config_path) + base_config = load_training_config(base_path) + base_topology = resolve_topology(base_config, world_size=world_size, local_world_size=local_world_size) + resolved_world_size = world_size or base_topology.world_size + resolved_local_world_size = local_world_size or base_topology.local_world_size + candidate_world_sizes = _dedupe_sorted(world_sizes or [resolved_world_size]) + primary_behavior_points = load_benchmark_behavior_points(benchmark_dir) if benchmark_dir is not None else [] + raw_supplemental_behavior_points: list[BenchmarkBehaviorPoint] = [] + for supplemental_dir in supplemental_benchmark_dirs or []: + raw_supplemental_behavior_points.extend(load_benchmark_behavior_points(supplemental_dir)) + supplemental_behavior_points = [ + point for point in raw_supplemental_behavior_points if not behavior_point_model_mismatches(point, base_config) + ] + supplemental_model_mismatch_count = len(raw_supplemental_behavior_points) - len(supplemental_behavior_points) + analog_behavior_points: list[BenchmarkBehaviorPoint] = [] + for analog_dir in analog_benchmark_dirs or []: + analog_behavior_points.extend(load_benchmark_behavior_points(analog_dir)) + same_model_behavior_points = primary_behavior_points + supplemental_behavior_points + behavior_points = same_model_behavior_points + analog_behavior_points + metadata = resolve_model_metadata(base_config) + benchmark_support = _scenario_benchmark_support( + same_model_behavior_points, + base_config=base_config, + base_topology=base_topology, + ) + + default_behavior_points = same_model_behavior_points or behavior_points + micro_batch_values = micro_batch_sizes or _default_micro_batch_sizes(base_topology, default_behavior_points) + gradient_accumulation_values = gradient_accumulation_steps or [base_topology.gradient_accumulation_steps] + sample_packing_sequence_values = sample_packing_sequence_lengths or ( + [base_topology.sample_packing_sequence_len] if base_topology.sample_packing_sequence_len is not None else [] + ) + + candidates: list[ScenarioCandidate] = [] + warnings: list[str] = [] + if not sample_packing_sequence_values: + warnings.append("skipped all scenarios: data.sample_packing_sequence_len is not set") + include_sequence_len_in_label = len(sample_packing_sequence_values) > 1 + seen: set[tuple[str, str]] = set() + for candidate_world_size in candidate_world_sizes: + candidate_local_world_size = min(resolved_local_world_size, candidate_world_size) + for sample_packing_sequence_len in sample_packing_sequence_values: + try: + candidate_base_values = _topology_values_with_dp_split( + { + "world_size": candidate_world_size, + "expert_parallel_size": base_topology.expert_parallel_size, + "tensor_parallel_size": base_topology.tensor_parallel_size, + "pipeline_parallel_size": base_topology.pipeline_parallel_size, + "ulysses_parallel_size": base_topology.ulysses_parallel_size, + "ringattn_parallel_size": base_topology.ringattn_parallel_size, + }, + preferred_replicate_size=base_topology.data_parallel_replicate_size, + preferred_shard_size=base_topology.data_parallel_shard_size, + ) + if candidate_base_values is None: + raise ValueError("candidate base topology has invalid DP split") + candidate_base_config = _mutated_config( + base_config, + world_size=candidate_world_size, + micro_batch_size=base_topology.micro_batch_size, + gradient_accumulation_steps=base_topology.gradient_accumulation_steps, + expert_parallel_size=base_topology.expert_parallel_size, + tensor_parallel_size=base_topology.tensor_parallel_size, + pipeline_parallel_size=base_topology.pipeline_parallel_size, + ulysses_parallel_size=base_topology.ulysses_parallel_size, + ringattn_parallel_size=base_topology.ringattn_parallel_size, + data_parallel_replicate_size=candidate_base_values["data_parallel_replicate_size"], + data_parallel_shard_size=candidate_base_values["data_parallel_shard_size"], + ) + _set_sample_packing_sequence_len(candidate_base_config, sample_packing_sequence_len) + candidate_base_topology = resolve_topology( + candidate_base_config, + world_size=candidate_world_size, + local_world_size=candidate_local_world_size, + ) + except ValueError as exc: + warnings.append( + f"skipped world_size={candidate_world_size}, " + f"sample_packing_sequence_len={sample_packing_sequence_len}: {exc}" + ) + continue + + if topology_sweep == "auto": + ep_values = expert_parallel_sizes or _auto_ep_sizes(candidate_base_topology) + tp_values = tensor_parallel_sizes or _auto_tensor_parallel_sizes(candidate_base_topology, metadata) + pp_values = pipeline_parallel_sizes or _auto_pipeline_parallel_sizes(candidate_base_topology, metadata) + ulysses_values = ulysses_parallel_sizes or _auto_ulysses_parallel_sizes( + candidate_base_topology, metadata + ) + ring_values = ringattn_parallel_sizes or _auto_ringattn_parallel_sizes(candidate_base_topology) + else: + ep_values = expert_parallel_sizes or [candidate_base_topology.expert_parallel_size] + tp_values = tensor_parallel_sizes or [candidate_base_topology.tensor_parallel_size] + pp_values = pipeline_parallel_sizes or [candidate_base_topology.pipeline_parallel_size] + ulysses_values = ulysses_parallel_sizes or [candidate_base_topology.ulysses_parallel_size] + ring_values = ringattn_parallel_sizes or [candidate_base_topology.ringattn_parallel_size] + + for pp in pp_values: + for tp in tp_values: + for ulysses in ulysses_values: + for ringattn in ring_values: + for ep in ep_values: + for micro_batch_size in micro_batch_values: + for gradient_accumulation_step in gradient_accumulation_values: + try: + candidate_values = _topology_values_with_dp_split( + { + "world_size": candidate_world_size, + "expert_parallel_size": ep, + "tensor_parallel_size": tp, + "pipeline_parallel_size": pp, + "ulysses_parallel_size": ulysses, + "ringattn_parallel_size": ringattn, + }, + preferred_replicate_size=base_topology.data_parallel_replicate_size, + preferred_shard_size=base_topology.data_parallel_shard_size, + ) + if candidate_values is None: + raise ValueError("candidate topology has invalid DP split") + raw_config = _mutated_config( + base_config, + world_size=candidate_world_size, + micro_batch_size=micro_batch_size, + gradient_accumulation_steps=gradient_accumulation_step, + expert_parallel_size=ep, + tensor_parallel_size=tp, + pipeline_parallel_size=pp, + ulysses_parallel_size=ulysses, + ringattn_parallel_size=ringattn, + data_parallel_replicate_size=candidate_values[ + "data_parallel_replicate_size" + ], + data_parallel_shard_size=candidate_values["data_parallel_shard_size"], + ) + _set_sample_packing_sequence_len(raw_config, sample_packing_sequence_len) + _set_balanced_routing(raw_config, balanced_routing) + topology = resolve_topology( + raw_config, + world_size=candidate_world_size, + local_world_size=candidate_local_world_size, + ) + except ValueError as exc: + warnings.append( + f"skipped world={candidate_world_size}, " + f"seq={sample_packing_sequence_len}, mbs={micro_batch_size}, " + f"ga={gradient_accumulation_step}, ep={ep}, tp={tp}, pp={pp}, " + f"u={ulysses}, r={ringattn}: {exc}" + ) + continue + if topology.ep_fsdp_size is None: + warnings.append( + f"skipped {_topology_label(topology)}: ep_fsdp is not integral" + ) + continue + if ( + topology.num_experts is not None + and topology.num_experts % topology.expert_parallel_size + ): + warnings.append( + f"skipped {_topology_label(topology)}: EP does not divide num_experts" + ) + continue + + shape = build_shape_ledger(topology, balanced_routing=True) + memory = build_memory_ledger( + raw_config, + topology=topology, + model_metadata=metadata, + ) + communication = _communication_ledger(topology) + memory_peak_estimate = _calibrated_memory_peak_estimate( + behavior_points, + base_config, + raw_config, + topology, + shape, + metadata, + default_world_size=candidate_world_size, + default_local_world_size=candidate_local_world_size, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + ) + exact_points = [ + point + for point in behavior_points + if behavior_point_matches_topology(point, topology) + and behavior_point_matches_workload(point, raw_config) + ] + label_topology = _candidate_topology_label( + topology, + include_sequence_len=include_sequence_len_in_label, + ) + if exact_points: + for point in exact_points: + behavior = predict_benchmark_behavior( + [point], topology, shape, raw_config + ) + label = f"{label_topology}:{point.label}" + key = (label, point.source) + if key in seen: + continue + seen.add(key) + candidates.append( + _candidate_from_prediction( + label=label, + config_path=str(base_path), + topology=topology, + shape=shape, + behavior=behavior, + prediction_confidence="calibrated", + promotable=point.correctness_status == "k3_pass", + behavior_points=behavior_points, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + memory_ownership_notes=_memory_ownership_notes(memory), + communication=communication, + notes=list(point.notes), + ) + ) + continue + + behavior, extrapolation_notes = _extrapolate_behavior( + behavior_points, + topology, + shape, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + ) + label_suffix = ( + "cross_model_extrapolated" + if behavior.status == "cross_model_extrapolated" + else "extrapolated" + ) + label = f"{label_topology}:{label_suffix}" + key = (label, behavior.source or "") + if key in seen: + continue + seen.add(key) + candidates.append( + _candidate_from_prediction( + label=label, + config_path=None, + topology=topology, + shape=shape, + behavior=behavior, + prediction_confidence=behavior.status, + promotable=False, + behavior_points=behavior_points, + raw_config=raw_config, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + analytic_peak_floor_gb=memory.analytic_peak_floor_gb, + memory_peak_estimate=memory_peak_estimate, + memory_ownership_notes=_memory_ownership_notes(memory), + communication=communication, + notes=extrapolation_notes, + ) + ) + + candidates = _apply_cross_model_analog_support_risk(candidates) + candidates = _apply_same_workload_scaling_metrics(candidates) + candidates = _apply_frontier_dominance(candidates) + candidates = sorted(candidates, key=_candidate_sort_key, reverse=True) + feasible = [candidate for candidate in candidates if candidate.score_tokens_per_sec is not None] + best_raw = feasible[0] if feasible else None + risk_adjusted = [candidate for candidate in feasible if candidate.score_risk_adjusted_tokens_per_sec is not None] + best_risk_adjusted = max(risk_adjusted, key=_risk_adjusted_sort_key) if risk_adjusted else None + efficiency = [candidate for candidate in feasible if candidate.score_tokens_per_gpu_per_sec is not None] + best_efficiency = max(efficiency, key=_efficiency_sort_key) if efficiency else None + risk_adjusted_efficiency = [ + candidate for candidate in feasible if candidate.score_risk_adjusted_tokens_per_gpu_per_sec is not None + ] + best_risk_adjusted_efficiency = ( + max(risk_adjusted_efficiency, key=_risk_adjusted_efficiency_sort_key) if risk_adjusted_efficiency else None + ) + next_measurement = [candidate for candidate in risk_adjusted if "requires_remeasurement" in candidate.risk_flags] + best_next_measurement = max(next_measurement, key=_risk_adjusted_sort_key) if next_measurement else None + promotable = [candidate for candidate in feasible if candidate.promotable] + best_promotable = promotable[0] if promotable else None + decision_summary = _scenario_decision_summary( + candidates, + feasible, + best_raw, + best_risk_adjusted, + best_efficiency, + best_risk_adjusted_efficiency, + best_next_measurement, + best_promotable, + benchmark_support, + ) + if best_raw is not None and not best_raw.promotable: + warnings.append(f"best raw scenario {best_raw.label} is not correctness-promotable") + if best_raw is not None and best_risk_adjusted is not None and best_raw.label != best_risk_adjusted.label: + warnings.append( + f"best raw scenario {best_raw.label} differs from risk-adjusted choice {best_risk_adjusted.label}" + ) + if best_promotable is None: + warnings.append("no correctness-promotable scenario found") + if supplemental_behavior_points: + warnings.append(f"loaded {len(supplemental_behavior_points)} supplemental benchmark behavior points") + if supplemental_model_mismatch_count: + warnings.append( + f"ignored {supplemental_model_mismatch_count} supplemental benchmark behavior points with model mismatch" + ) + if analog_behavior_points: + warnings.append(f"loaded {len(analog_behavior_points)} analog benchmark behavior points") + + report = ScenarioReport( + base_config_path=str(base_path), + benchmark_dir=str(benchmark_dir) if benchmark_dir is not None else None, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + topology_sweep=topology_sweep, + balanced_routing=balanced_routing, + world_sizes=candidate_world_sizes, + candidate_count=len(candidates), + feasible_count=len(feasible), + best_raw=best_raw, + best_risk_adjusted=best_risk_adjusted, + best_efficiency=best_efficiency, + best_risk_adjusted_efficiency=best_risk_adjusted_efficiency, + best_next_measurement=best_next_measurement, + best_promotable=best_promotable, + decision_summary=decision_summary, + candidates=candidates, + planner_context=_scenario_planner_context( + base_path=base_path, + benchmark_dir=benchmark_dir, + supplemental_benchmark_dirs=supplemental_benchmark_dirs, + analog_benchmark_dirs=analog_benchmark_dirs, + requested_world_size=world_size, + candidate_world_sizes=candidate_world_sizes, + resolved_local_world_size=resolved_local_world_size, + micro_batch_values=micro_batch_values, + gradient_accumulation_values=gradient_accumulation_values, + sample_packing_sequence_values=sample_packing_sequence_values, + expert_parallel_sizes=expert_parallel_sizes, + tensor_parallel_sizes=tensor_parallel_sizes, + pipeline_parallel_sizes=pipeline_parallel_sizes, + ulysses_parallel_sizes=ulysses_parallel_sizes, + ringattn_parallel_sizes=ringattn_parallel_sizes, + topology_sweep=topology_sweep, + balanced_routing=balanced_routing, + device_memory_limit_gb=device_memory_limit_gb, + memory_safety_factor=memory_safety_factor, + ), + warnings=warnings, + supplemental_benchmark_dirs=[str(path) for path in supplemental_benchmark_dirs or []], + analog_benchmark_dirs=[str(path) for path in analog_benchmark_dirs or []], + primary_behavior_point_count=len(primary_behavior_points), + supplemental_behavior_point_count=len(supplemental_behavior_points), + analog_behavior_point_count=len(analog_behavior_points), + total_behavior_point_count=len(behavior_points), + ) + return _attach_measurement_design_summary(report) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pack", help="Built-in calibration-pack name") + parser.add_argument("--config", type=Path, default=None) + parser.add_argument("--benchmark-dir", type=Path, default=None) + parser.add_argument( + "--supplemental-benchmark-dir", + action="append", + type=Path, + default=[], + help="Additional same-model benchmark directory included as scenario calibration and support evidence", + ) + parser.add_argument( + "--analog-benchmark-dir", + action="append", + type=Path, + default=[], + help="Additional benchmark directory used only as low-confidence analog evidence when same-model data is absent", + ) + parser.add_argument("--world-size", type=int, default=None) + parser.add_argument("--world-sizes", default=None, help="Comma list of world sizes to compare") + parser.add_argument("--local-world-size", type=int, default=None) + parser.add_argument("--micro-batch-sizes", default=None, help="Comma list, or auto when omitted") + parser.add_argument( + "--gradient-accumulation-steps", default=None, help="Comma list, or base config GA when omitted" + ) + parser.add_argument( + "--sample-packing-sequence-lengths", + default=None, + help="Comma list, or base config data.sample_packing_sequence_len when omitted", + ) + parser.add_argument("--expert-parallel-sizes", default=None, help="Comma list, or base config EP when omitted") + parser.add_argument("--tensor-parallel-sizes", default=None, help="Comma list, or base config TP when omitted") + parser.add_argument("--pipeline-parallel-sizes", default=None, help="Comma list, or base config PP when omitted") + parser.add_argument( + "--ulysses-parallel-sizes", default=None, help="Comma list, or base config Ulysses when omitted" + ) + parser.add_argument("--ringattn-parallel-sizes", default=None, help="Comma list, or base config Ring when omitted") + parser.add_argument( + "--topology-sweep", + choices=("base", "auto"), + default="base", + help="Use base topology dimensions, or derive an automatic TP/PP/CP/EP sweep", + ) + parser.add_argument( + "--balanced-routing", + action="store_true", + help="Match measured scenarios that used XORL_MOE_SYNTHETIC_ROUTING=balanced", + ) + parser.add_argument("--device-memory-limit-gb", type=float, default=80.0) + parser.add_argument("--memory-safety-factor", type=float, default=1.15) + parser.add_argument("--output", type=Path, default=None) + parser.add_argument( + "--write-measurement-configs", + type=Path, + default=None, + help="Write the bounded measurement portfolio as runnable YAML configs in this directory", + ) + args = parser.parse_args() + + args.config, args.benchmark_dir = resolve_pack_inputs(args.pack, args.config, args.benchmark_dir) + if args.config is None: + parser.error("provide --pack or --config") + + report = plan_scenario( + args.config, + benchmark_dir=args.benchmark_dir, + supplemental_benchmark_dirs=args.supplemental_benchmark_dir, + analog_benchmark_dirs=args.analog_benchmark_dir, + world_size=args.world_size, + world_sizes=_parse_int_list(args.world_sizes), + local_world_size=args.local_world_size, + micro_batch_sizes=_parse_int_list(args.micro_batch_sizes), + gradient_accumulation_steps=_parse_int_list(args.gradient_accumulation_steps), + sample_packing_sequence_lengths=_parse_int_list(args.sample_packing_sequence_lengths), + expert_parallel_sizes=_parse_int_list(args.expert_parallel_sizes), + tensor_parallel_sizes=_parse_int_list(args.tensor_parallel_sizes), + pipeline_parallel_sizes=_parse_int_list(args.pipeline_parallel_sizes), + ulysses_parallel_sizes=_parse_int_list(args.ulysses_parallel_sizes), + ringattn_parallel_sizes=_parse_int_list(args.ringattn_parallel_sizes), + topology_sweep=args.topology_sweep, + balanced_routing=args.balanced_routing, + device_memory_limit_gb=args.device_memory_limit_gb, + memory_safety_factor=args.memory_safety_factor, + ) + rendered = json.dumps(to_jsonable(report), indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + if args.write_measurement_configs is not None: + write_measurement_configs(report, args.write_measurement_configs) + + +if __name__ == "__main__": + main() diff --git a/src/xorl/sim/schemas.py b/src/xorl/sim/schemas.py new file mode 100644 index 00000000..0d011b05 --- /dev/null +++ b/src/xorl/sim/schemas.py @@ -0,0 +1,1912 @@ +"""Dataclasses shared by the local training-engine simulator.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field, is_dataclass +from typing import Any + + +def to_jsonable(value: Any) -> Any: + """Convert simulator dataclasses into plain JSON-compatible containers.""" + if is_dataclass(value): + return {key: to_jsonable(item) for key, item in asdict(value).items()} + if isinstance(value, dict): + return {str(key): to_jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [to_jsonable(item) for item in value] + return value + + +@dataclass(frozen=True) +class ModelMetadata: + model_path: str | None + config_path: str | None + source: str + num_experts: int | None = None + top_k: int | None = None + num_hidden_layers: int | None = None + hidden_size: int | None = None + intermediate_size: int | None = None + moe_intermediate_size: int | None = None + shared_expert_intermediate_size: int | None = None + num_attention_heads: int | None = None + num_key_value_heads: int | None = None + head_dim: int | None = None + vocab_size: int | None = None + tie_word_embeddings: bool | None = None + # Hybrid GatedDeltaNet models (Qwen3.5/3.6): every full_attention_interval-th layer is full + # attention, the rest are GatedDeltaNet linear-attention. None => all layers are full attention. + full_attention_interval: int | None = None + # Gated attention (Qwen3.5/3.6): full-attention q_proj is 2x width (query + sigmoid output gate), + # and the GatedDeltaNet layers carry a g_proj. None/False => ungated. + attn_output_gate: bool | None = None + # GatedDeltaNet linear-attention dimensions (present only for hybrid models). + linear_num_key_heads: int | None = None + linear_num_value_heads: int | None = None + linear_key_head_dim: int | None = None + linear_value_head_dim: int | None = None + linear_conv_kernel_dim: int | None = None + + +@dataclass(frozen=True) +class Topology: + world_size: int + local_world_size: int + node_count: int + data_parallel_size: int + data_parallel_replicate_size: int + data_parallel_shard_size: int + tensor_parallel_size: int + pipeline_parallel_size: int + expert_parallel_size: int + ep_fsdp_size: int | None + ulysses_parallel_size: int + ringattn_parallel_size: int + micro_batch_size: int + gradient_accumulation_steps: int + global_batch_size: int + sample_packing_sequence_len: int | None + num_experts: int | None = None + top_k: int | None = None + + @property + def sequence_parallel_size(self) -> int: + return self.ulysses_parallel_size * self.ringattn_parallel_size + + +@dataclass(frozen=True) +class RunFingerprint: + config_path: str + config_sha256: str + config_name: str + repo_commit: str | None + balanced_routing: bool + topology: Topology + model_metadata: ModelMetadata + + +@dataclass(frozen=True) +class BalancedRoutingLedger: + total_slots: int + num_experts: int + counts_by_expert: list[int] + max_slots_per_expert: int + min_slots_per_expert: int + imbalance_slots: int + + +@dataclass(frozen=True) +class ShapeLedger: + microbatch_tokens_per_dp_rank: int | None + global_tokens_per_microbatch: int | None + global_tokens_per_train_step: int | None + tokens_per_gpu_per_train_step: float | None + sequence_parallel_size: int + tokens_per_model_rank_per_microbatch: int | None + routed_slots_per_model_rank_microbatch: int | None + routed_slots_per_train_step_model_rank: int | None + balanced_routing: BalancedRoutingLedger | None + ep_rank_slots_per_microbatch: list[int] | None = None + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class StepObservation: + source: str + step: int + max_steps: str + loss: float | None = None + grad_norm: float | None = None + lr: float | None = None + tflops_per_gpu: float | None = None + mfu: float | None = None + tokens_per_sec: float | None = None + step_time_s: float | None = None + peak_mem_gb: float | None = None + phase_memory_gb: dict[str, float] = field(default_factory=dict) + extra: dict[str, float] = field(default_factory=dict) + + +@dataclass(frozen=True) +class PhaseObservation: + source: str + prefix: str + step: int + max_steps: str + metrics: dict[str, float] + + +@dataclass(frozen=True) +class MemoryPhaseObservation: + source: str + prefix: str + step: int + max_steps: str + metrics: dict[str, float] + + +@dataclass(frozen=True) +class ObservedRun: + sources: list[str] + steps: list[StepObservation] + phases: list[PhaseObservation] = field(default_factory=list) + memory_phases: list[MemoryPhaseObservation] = field(default_factory=list) + + +@dataclass(frozen=True) +class MemoryBucket: + name: str + gb: float + source: str + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class MemoryLedger: + deepep_buffer_size_gb: float | None + observed_peak_mem_gb_max: float | None + calibrated_peak_mem_gb: float | None + calibrated_peak_source: str | None + observed_phase_peak_gb: dict[str, float] + estimated_total_params_b: float | None = None + estimated_local_params_b: float | None = None + persistent_model_state_gb: float | None = None + gradient_state_gb: float | None = None + optimizer_state_gb: float | None = None + analytic_peak_floor_gb: float | None = None + analytic_floor_fraction_of_calibrated_peak: float | None = None + calibrated_residual_peak_gb: float | None = None + calibrated_residual_fraction_of_peak: float | None = None + memory_coverage_status: str = "unresolved_analytic_floor" + top_memory_buckets: list[MemoryBucket] = field(default_factory=list) + unsupported_buckets: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class TimingLedger: + source: str | None + timing_coverage_status: str + forward_backward_s: float | None + forward_s: float | None + loss_s: float | None + backward_s: float | None + optimizer_s: float | None + input_s: float | None + step_time_s: float | None + phase_time_sec: dict[str, float] = field(default_factory=dict) + phase_time_share: dict[str, float] = field(default_factory=dict) + phase_bottleneck_phase: str | None = None + phase_bottleneck_bucket: str | None = None + phase_bottleneck_share: float | None = None + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class SimulatorSupportLedger: + requested_surface: str + support_status: str + support_blockers: list[str] + supported_outputs: list[str] + unsupported_outputs: list[str] + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class CommLedger: + tensor_parallel_cross_node: bool + pipeline_parallel_cross_node: bool + expert_parallel_cross_node: bool + context_parallel_cross_node: bool + fsdp_cross_node: bool + cross_node_dimensions: list[str] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class BenchmarkBehaviorPoint: + label: str + source: str + micro_batch_size: int | None + global_batch_size: int | None + tokens_per_sec: float | None + step_time_sec: float | None + gradient_accumulation_steps: int | None = None + tokens_per_sec_std: float | None = None + tokens_per_sec_cv: float | None = None + step_time_sec_std: float | None = None + step_time_sec_cv: float | None = None + # Median over post-warmup steps of (tokens_per_sec x step_time) computed PER STEP — the realized + # per-step token load, free of the mean(tps) x mean/median(step) cross-aggregate biases. + tokens_per_step: float | None = None + phase_time_sec: dict[str, float] = field(default_factory=dict) + # Cross-rank MEAN companion of phase_time_sec (which is the cross-rank MAX convention): lets + # balanced-rank term comparisons separate rank asymmetry from term error. + phase_time_rank_mean_sec: dict[str, float] = field(default_factory=dict) + phase_time_share: dict[str, float] = field(default_factory=dict) + phase_memory_peak_gb: dict[str, float] = field(default_factory=dict) + mfu_percent: float | None = None + tflops_per_gpu: float | None = None + peak_mem_gb: float | None = None + allocator_retries: int | None = None + measured_steps: int | None = None + warmup_steps: int | None = None + gpu_count: int | None = None + model_ref: str | None = None + sample_packing_sequence_len: int | None = None + data_parallel_replicate_size: int | None = None + data_parallel_shard_size: int | None = None + tensor_parallel_size: int | None = None + pipeline_parallel_size: int | None = None + ulysses_parallel_size: int | None = None + ringattn_parallel_size: int | None = None + expert_parallel_size: int | None = None + ep_fsdp_size: int | None = None + deepep_async_combine: bool | None = None + deepep_num_sms: int | None = None + deepep_buffer_size_gb: float | None = None + enable_compile: bool | None = None + gradient_checkpointing_method: str | None = None + enable_activation_offload: bool | None = None + activation_offload_prefetch_count: int | None = None + skip_param_upcast: bool | None = None + fsdp_reduce_dtype: str | None = None + ce_mode: str | None = None + moe_implementation: str | None = None + moe_checkpoint_method: str | None = None + muon_momentum: float | None = None + muon_update_dtype: str | None = None + attention_backend: str | None = None + balanced_routing: bool | None = None + status: str = "observed" + correctness_status: str | None = None + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class BenchmarkBehaviorPrediction: + status: str + matched_label: str | None + source: str | None + tokens_per_sec: float | None + tokens_per_sec_per_gpu: float | None + step_time_sec: float | None + mfu_percent: float | None + tflops_per_gpu: float | None + promised_tflops_per_gpu: float | None + peak_mem_gb: float | None + allocator_retries: int | None + derived_global_tokens_per_step: int | None + tokens_per_sec_std: float | None = None + tokens_per_sec_cv: float | None = None + step_time_sec_std: float | None = None + step_time_sec_cv: float | None = None + phase_time_sec: dict[str, float] = field(default_factory=dict) + phase_time_share: dict[str, float] = field(default_factory=dict) + phase_memory_peak_gb: dict[str, float] = field(default_factory=dict) + measured_steps: int | None = None + warmup_steps: int | None = None + model_ref: str | None = None + balanced_routing: bool | None = None + correctness_status: str | None = None + cross_model_active_param_ratio: float | None = None + cross_model_active_param_scale: float | None = None + cross_model_reference_active_params_b: float | None = None + cross_model_target_active_params_b: float | None = None + cross_model_sequence_length_factor: float | None = None + cross_model_parallelism_factor: float | None = None + cross_model_memory_factor: float | None = None + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class PredictionReport: + fingerprint: RunFingerprint + shape: ShapeLedger + memory: MemoryLedger + timing: TimingLedger + support: SimulatorSupportLedger + benchmark_behavior: BenchmarkBehaviorPrediction | None = None + observed_summary: dict[str, Any] | None = None + calibration_sources: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class TradeoffCandidate: + label: str + config_path: str | None + behavior_source: str + topology: Topology | None + behavior: BenchmarkBehaviorPrediction + promotable: bool + score_tokens_per_sec: float | None + score_tflops_per_gpu: float | None + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class TradeoffReport: + benchmark_dir: str + status: str + candidate_count: int + best_raw: TradeoffCandidate | None + best_promotable: TradeoffCandidate | None + candidates: list[TradeoffCandidate] + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ScenarioCandidate: + label: str + config_path: str | None + topology: Topology + behavior: BenchmarkBehaviorPrediction + prediction_confidence: str + promotable: bool + feasibility_status: str + score_tokens_per_sec: float | None + score_tokens_per_gpu_per_sec: float | None + score_risk_adjusted_tokens_per_sec: float | None + score_risk_adjusted_tokens_per_gpu_per_sec: float | None + prediction_uncertainty_fraction: float | None + prediction_interval_lower_tokens_per_sec: float | None + prediction_interval_upper_tokens_per_sec: float | None + risk_adjusted_prediction_interval_lower_tokens_per_sec: float | None + risk_adjusted_prediction_interval_upper_tokens_per_sec: float | None + analytic_peak_floor_gb: float | None + estimated_peak_mem_gb: float | None + memory_basis: str + memory_coverage_status: str + memory_headroom_gb: float | None + estimated_memory_residual_gb: float | None + estimated_memory_residual_fraction_of_peak: float | None + max_ep_rank_slots_per_microbatch: int | None + phase_bottleneck_phase: str | None + phase_bottleneck_bucket: str | None + phase_bottleneck_share: float | None + phase_bottleneck_time_sec: float | None + memory_bottleneck_phase: str | None + memory_bottleneck_bucket: str | None + memory_bottleneck_peak_gb: float | None + memory_bottleneck_fraction_of_peak: float | None + timing_coverage_status: str + timing_source_label: str | None + timing_step_time_s: float | None + timing_forward_backward_s: float | None + calibration_scope: str + recommendation: str + phase_bottleneck_half_speedup_scale: float | None = None + phase_bottleneck_half_speedup_tokens_per_sec: float | None = None + phase_bottleneck_half_speedup_delta_pct: float | None = None + phase_bottleneck_half_speedup_risk_adjusted_tokens_per_sec: float | None = None + phase_bottleneck_half_speedup_risk_adjusted_delta_pct: float | None = None + simulator_surface: str = "unknown_config_surface" + simulator_support_status: str = "unknown" + simulator_support_blockers: list[str] = field(default_factory=list) + target_runtime_signature: str = "unknown" + calibration_distance: float | None = None + scaling_baseline_label: str | None = None + scaling_baseline_world_size: int | None = None + scaling_gpu_ratio: float | None = None + scaling_speedup: float | None = None + scaling_efficiency: float | None = None + risk_adjusted_scaling_speedup: float | None = None + risk_adjusted_scaling_efficiency: float | None = None + raw_frontier_member: bool | None = None + raw_dominated_by_label: str | None = None + raw_dominance_margin_tokens_per_sec: float | None = None + raw_dominance_margin_tokens_per_gpu_per_sec: float | None = None + risk_adjusted_frontier_member: bool | None = None + risk_adjusted_dominated_by_label: str | None = None + risk_adjusted_dominance_margin_tokens_per_sec: float | None = None + risk_adjusted_dominance_margin_tokens_per_gpu_per_sec: float | None = None + calibration_distance_factors: list[str] = field(default_factory=list) + memory_calibration_source: str | None = None + memory_calibration_notes: list[str] = field(default_factory=list) + memory_ownership_notes: list[str] = field(default_factory=list) + communication: CommLedger | None = None + decision_factors: list[str] = field(default_factory=list) + risk_flags: list[str] = field(default_factory=list) + notes: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ParallelismAxisComparison: + axis: str + varied_dimensions: list[str] + primary_varied_dimensions: list[str] + co_varied_axis_dimensions: list[str] + coupling_status: str + group_key: str + candidate_count: int + raw_best_label: str | None + raw_best_axis_value: str | None + raw_best_score_tokens_per_sec: float | None + raw_worst_label: str | None + raw_worst_axis_value: str | None + raw_worst_score_tokens_per_sec: float | None + raw_spread_tokens_per_sec: float | None + raw_spread_ratio: float | None + risk_adjusted_best_label: str | None + risk_adjusted_best_axis_value: str | None + risk_adjusted_best_score_tokens_per_sec: float | None + risk_adjusted_worst_label: str | None + risk_adjusted_worst_axis_value: str | None + risk_adjusted_worst_score_tokens_per_sec: float | None + risk_adjusted_spread_tokens_per_sec: float | None + risk_adjusted_spread_ratio: float | None + risk_adjusted_winner_matches_raw: bool | None + comparison_status: str + risk_adjusted_best_interval_lower_tokens_per_sec: float | None = None + risk_adjusted_best_interval_upper_tokens_per_sec: float | None = None + risk_adjusted_worst_interval_lower_tokens_per_sec: float | None = None + risk_adjusted_worst_interval_upper_tokens_per_sec: float | None = None + risk_adjusted_interval_overlap_status: str = "unknown" + risk_adjusted_interval_overlap_candidate_count: int = 0 + risk_adjusted_interval_overlap_candidate_labels: list[str] = field(default_factory=list) + risk_adjusted_interval_margin_tokens_per_sec: float | None = None + + +@dataclass(frozen=True) +class ScenarioParallelismAxisCoverage: + axis: str + status: str + candidate_group_count: int + candidate_count: int + scored_count: int + blocked_count: int + unscored_count: int + varied_dimensions: list[str] = field(default_factory=list) + primary_varied_dimensions: list[str] = field(default_factory=list) + co_varied_axis_dimensions: list[str] = field(default_factory=list) + confounded_runtime_dimensions: list[str] = field(default_factory=list) + feasibility_status_counts: dict[str, int] = field(default_factory=dict) + + +@dataclass(frozen=True) +class ScenarioBenchmarkSupport: + support_status: str = "no_benchmark_support" + support_blockers: list[str] = field(default_factory=list) + point_count: int = 0 + scored_count: int = 0 + memory_blocked_count: int = 0 + varied_parallelism_dimensions: list[str] = field(default_factory=list) + varied_workload_dimensions: list[str] = field(default_factory=list) + varied_runtime_dimensions: list[str] = field(default_factory=list) + parallelism_axis_coverage_status_counts: dict[str, int] = field(default_factory=dict) + scored_parallelism_axis_names: list[str] = field(default_factory=list) + blocked_parallelism_axis_names: list[str] = field(default_factory=list) + confounded_parallelism_axis_names: list[str] = field(default_factory=list) + unscored_parallelism_axis_names: list[str] = field(default_factory=list) + missing_parallelism_axis_names: list[str] = field(default_factory=list) + point_labels: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ScenarioValidationAction: + action_status: str + priority: int + required_measurement: str + reason_category: str + candidate_count: int + candidate_labels: list[str] + total_gpu_count: int + max_priority_score: float | None + max_priority_label: str | None + max_priority_per_gpu: float | None + max_priority_per_gpu_label: str | None + parallelism_axis_names: list[str] = field(default_factory=list) + reason_statuses: list[str] = field(default_factory=list) + config_overrides: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ScenarioCaptureGap: + gap_status: str + priority: int + required_measurement: str + reason: str + blocker_names: list[str] + candidate_count: int + scored_count: int + unscored_count: int + memory_blocked_count: int + missing_parallelism_axis_names: list[str] = field(default_factory=list) + blocked_parallelism_axis_names: list[str] = field(default_factory=list) + confounded_parallelism_axis_names: list[str] = field(default_factory=list) + unscored_parallelism_axis_names: list[str] = field(default_factory=list) + varied_parallelism_dimensions: list[str] = field(default_factory=list) + varied_workload_dimensions: list[str] = field(default_factory=list) + varied_runtime_dimensions: list[str] = field(default_factory=list) + runtime_mismatch_dimensions: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ScenarioReadiness: + readiness_status: str = "unknown_scenario_readiness" + can_capture_scenario: bool = False + can_predict_scenario_fidelity: bool = False + can_select_parallelism_tradeoff: bool = False + can_generalize_model: bool = False + scenario_capture_status: str = "unknown" + scenario_capture_blockers: list[str] = field(default_factory=list) + scenario_prediction_fidelity_status: str = "unknown" + scenario_prediction_fidelity_blockers: list[str] = field(default_factory=list) + parallelism_optimality_status: str = "unknown" + parallelism_optimality_blockers: list[str] = field(default_factory=list) + model_generalization_status: str = "unknown" + model_generalization_blockers: list[str] = field(default_factory=list) + measurement_readiness_status: str = "unknown" + measurement_portfolio_coverage_status: str = "unknown" + measurement_portfolio_coverage_blockers: list[str] = field(default_factory=list) + required_measurements: list[str] = field(default_factory=list) + scenario_capture_gap_count: int = 0 + scenario_capture_gap_status_counts: dict[str, int] = field(default_factory=dict) + scenario_capture_gap_required_measurements: list[str] = field(default_factory=list) + top_scenario_capture_gap_statuses: list[str] = field(default_factory=list) + validation_action_count: int = 0 + validation_action_status_counts: dict[str, int] = field(default_factory=dict) + validation_action_required_measurements: list[str] = field(default_factory=list) + validation_action_total_gpu_count: int = 0 + measurement_candidate_count: int = 0 + measurement_candidate_labels: list[str] = field(default_factory=list) + measurement_design_config_count: int = 0 + measurement_design_config_labels: list[str] = field(default_factory=list) + measurement_design_config_filenames: list[str] = field(default_factory=list) + measurement_portfolio_total_gpu_count: int = 0 + measurement_portfolio_parallelism_axis_gap_names: list[str] = field(default_factory=list) + measurement_portfolio_cross_model_analog_count: int = 0 + candidate_count: int = 0 + scored_count: int = 0 + unscored_count: int = 0 + memory_blocked_count: int = 0 + unique_parallelism_strategy_count: int = 0 + scored_parallelism_strategy_count: int = 0 + promotable_parallelism_strategy_count: int = 0 + varied_parallelism_dimensions: list[str] = field(default_factory=list) + varied_workload_dimensions: list[str] = field(default_factory=list) + varied_runtime_dimensions: list[str] = field(default_factory=list) + runtime_mismatch_dimensions: list[str] = field(default_factory=list) + parallelism_axis_coverage_status_counts: dict[str, int] = field(default_factory=dict) + scored_parallelism_axis_names: list[str] = field(default_factory=list) + blocked_parallelism_axis_names: list[str] = field(default_factory=list) + confounded_parallelism_axis_names: list[str] = field(default_factory=list) + unscored_parallelism_axis_names: list[str] = field(default_factory=list) + missing_parallelism_axis_names: list[str] = field(default_factory=list) + simulator_support_status_counts: dict[str, int] = field(default_factory=dict) + prediction_confidence_counts: dict[str, int] = field(default_factory=dict) + calibration_scope_counts: dict[str, int] = field(default_factory=dict) + memory_coverage_status_counts: dict[str, int] = field(default_factory=dict) + timing_coverage_status_counts: dict[str, int] = field(default_factory=dict) + cross_model_analog_support_status: str = "not_used" + cross_model_analog_candidate_count: int = 0 + cross_model_analog_scored_count: int = 0 + benchmark_support: ScenarioBenchmarkSupport = field(default_factory=ScenarioBenchmarkSupport) + + +@dataclass(frozen=True) +class ScenarioDecisionSummary: + candidate_count: int + scored_count: int + unscored_count: int + feasible_count: int + promotable_count: int + requires_remeasurement_count: int + memory_blocked_count: int + unique_parallelism_strategy_count: int + best_raw_label: str | None + best_raw_score_tokens_per_sec: float | None + best_risk_adjusted_label: str | None + best_risk_adjusted_score_tokens_per_sec: float | None + best_efficiency_label: str | None + best_efficiency_score_tokens_per_gpu_per_sec: float | None + best_risk_adjusted_efficiency_label: str | None + best_risk_adjusted_efficiency_score_tokens_per_gpu_per_sec: float | None + best_next_measurement_label: str | None + best_next_measurement_score_tokens_per_sec: float | None + best_promotable_label: str | None + candidate_model_ref_counts: dict[str, int] = field(default_factory=dict) + scored_model_ref_counts: dict[str, int] = field(default_factory=dict) + candidate_world_size_counts: dict[int, int] = field(default_factory=dict) + scored_world_size_counts: dict[int, int] = field(default_factory=dict) + candidate_sequence_length_counts: dict[int, int] = field(default_factory=dict) + scored_sequence_length_counts: dict[int, int] = field(default_factory=dict) + candidate_global_batch_size_counts: dict[int, int] = field(default_factory=dict) + scored_global_batch_size_counts: dict[int, int] = field(default_factory=dict) + best_promotable_score_tokens_per_sec: float | None = None + best_promotable_score_risk_adjusted_tokens_per_sec: float | None = None + promotable_raw_gap_tokens_per_sec: float | None = None + promotable_raw_gap_percentage: float | None = None + promotable_risk_adjusted_gap_tokens_per_sec: float | None = None + promotable_risk_adjusted_gap_percentage: float | None = None + promotion_readiness_status: str = "unknown" + throughput_efficiency_frontier_labels: list[str] = field(default_factory=list) + risk_adjusted_efficiency_frontier_labels: list[str] = field(default_factory=list) + throughput_efficiency_frontier_count: int = 0 + risk_adjusted_efficiency_frontier_count: int = 0 + raw_dominated_candidate_count: int = 0 + risk_adjusted_dominated_candidate_count: int = 0 + throughput_efficiency_tradeoff_status: str = "unknown" + same_workload_scaling_status: str = "unknown" + same_workload_scaling_group_count: int = 0 + same_workload_scaling_candidate_count: int = 0 + best_scaling_efficiency_label: str | None = None + best_scaling_efficiency: float | None = None + mean_scaling_efficiency: float | None = None + min_scaling_efficiency: float | None = None + best_risk_adjusted_scaling_efficiency_label: str | None = None + best_risk_adjusted_scaling_efficiency: float | None = None + mean_risk_adjusted_scaling_efficiency: float | None = None + min_risk_adjusted_scaling_efficiency: float | None = None + measurement_readiness_status: str = "unknown" + measurement_rationale: list[str] = field(default_factory=list) + measurement_candidate_count: int = 0 + measurement_candidate_labels: list[str] = field(default_factory=list) + measurement_candidate_reasons: dict[str, list[str]] = field(default_factory=dict) + measurement_candidate_priority_scores: dict[str, float] = field(default_factory=dict) + measurement_candidate_priority_per_gpu: dict[str, float] = field(default_factory=dict) + measurement_candidate_cost_gpus: dict[str, int] = field(default_factory=dict) + measurement_candidate_priority_factors: dict[str, list[str]] = field(default_factory=dict) + measurement_candidate_config_overrides: dict[str, list[str]] = field(default_factory=dict) + measurement_design_config_count: int = 0 + measurement_design_config_labels: list[str] = field(default_factory=list) + measurement_design_config_filenames: list[str] = field(default_factory=list) + measurement_portfolio_total_gpu_count: int = 0 + measurement_portfolio_max_priority_score: float | None = None + measurement_portfolio_max_priority_label: str | None = None + measurement_portfolio_max_priority_per_gpu: float | None = None + measurement_portfolio_max_priority_per_gpu_label: str | None = None + measurement_portfolio_coverage_status: str = "unknown" + measurement_portfolio_coverage_blockers: list[str] = field(default_factory=list) + measurement_portfolio_reason_category_counts: dict[str, int] = field(default_factory=dict) + measurement_portfolio_parallelism_axis_gap_names: list[str] = field(default_factory=list) + measurement_portfolio_cross_model_analog_count: int = 0 + validation_action_count: int = 0 + validation_action_status_counts: dict[str, int] = field(default_factory=dict) + validation_action_required_measurements: list[str] = field(default_factory=list) + validation_action_total_gpu_count: int = 0 + validation_actions: list[ScenarioValidationAction] = field(default_factory=list) + max_calibration_distance: float | None = None + max_calibration_distance_label: str | None = None + mean_scored_calibration_distance: float | None = None + high_uncertainty_candidate_count: int = 0 + max_prediction_uncertainty_fraction: float | None = None + max_prediction_uncertainty_fraction_label: str | None = None + mean_scored_prediction_uncertainty_fraction: float | None = None + risk_adjusted_interval_overlap_status: str = "unknown" + risk_adjusted_interval_overlap_contender_count: int = 0 + risk_adjusted_interval_overlap_contender_labels: list[str] = field(default_factory=list) + risk_adjusted_interval_best_vs_next_margin_tokens_per_sec: float | None = None + parallelism_tradeoff_status: str = "unknown" + parallelism_optimality_status: str = "unknown" + parallelism_optimality_blockers: list[str] = field(default_factory=list) + scored_parallelism_strategy_count: int = 0 + promotable_parallelism_strategy_count: int = 0 + requires_remeasurement_parallelism_strategy_count: int = 0 + parallelism_axis_comparison_count: int = 0 + isolated_parallelism_axis_comparison_count: int = 0 + coupled_parallelism_axis_comparison_count: int = 0 + parallelism_axis_interval_overlap_count: int = 0 + parallelism_axis_comparisons: list[ParallelismAxisComparison] = field(default_factory=list) + scored_parallelism_axis_names: list[str] = field(default_factory=list) + blocked_parallelism_axis_names: list[str] = field(default_factory=list) + confounded_parallelism_axis_names: list[str] = field(default_factory=list) + unscored_parallelism_axis_names: list[str] = field(default_factory=list) + missing_parallelism_axis_names: list[str] = field(default_factory=list) + parallelism_axis_coverage_status_counts: dict[str, int] = field(default_factory=dict) + parallelism_axis_coverage: list[ScenarioParallelismAxisCoverage] = field(default_factory=list) + parallelism_boundary_status: str = "unknown" + parallelism_boundary_prediction_status: str = "unknown" + parallelism_boundary_prediction_blockers: list[str] = field(default_factory=list) + parallelism_boundary_group_count: int = 0 + parallelism_boundary_candidate_count: int = 0 + parallelism_boundary_fit_count: int = 0 + parallelism_boundary_failure_count: int = 0 + parallelism_boundary_best_fit_label: str | None = None + parallelism_boundary_confounded_dimensions: list[str] = field(default_factory=list) + parallelism_boundary_measured_axis_names: list[str] = field(default_factory=list) + parallelism_boundary_confounded_axis_names: list[str] = field(default_factory=list) + parallelism_boundary_missing_axis_names: list[str] = field(default_factory=list) + parallelism_boundary_axis_coverage_status_counts: dict[str, int] = field(default_factory=dict) + parallelism_boundary_axis_coverage: list[ParallelismBoundaryAxisCoverage] = field(default_factory=list) + parallelism_boundary_groups: list[ParallelismBoundaryGroup] = field(default_factory=list) + cross_model_analog_support_status: str = "not_used" + cross_model_analog_candidate_count: int = 0 + cross_model_analog_scored_count: int = 0 + cross_model_analog_unique_prediction_count: int = 0 + cross_model_analog_unique_matched_label_count: int = 0 + cross_model_analog_unique_target_strategy_count: int = 0 + cross_model_analog_unique_target_runtime_signature_count: int = 0 + cross_model_analog_scored_varied_parallelism_dimensions: list[str] = field(default_factory=list) + cross_model_analog_scored_varied_workload_dimensions: list[str] = field(default_factory=list) + cross_model_analog_factor_status: str = "not_used" + cross_model_analog_unique_factor_count: int = 0 + cross_model_analog_factor_ranges: dict[str, list[float]] = field(default_factory=dict) + cross_model_analog_prediction_interval_top_count: int = 0 + cross_model_analog_prediction_interval_top_fraction: float | None = None + cross_model_analog_prediction_interval_top_labels: list[str] = field(default_factory=list) + cross_model_analog_prediction_interval_selectivity_status: str = "not_used" + model_generalization_status: str = "unknown" + model_generalization_blockers: list[str] = field(default_factory=list) + scenario_capture_status: str = "unknown" + scenario_capture_blockers: list[str] = field(default_factory=list) + scenario_capture_gap_count: int = 0 + scenario_capture_gap_status_counts: dict[str, int] = field(default_factory=dict) + scenario_capture_gap_required_measurements: list[str] = field(default_factory=list) + scenario_capture_gaps: list[ScenarioCaptureGap] = field(default_factory=list) + benchmark_support: ScenarioBenchmarkSupport = field(default_factory=ScenarioBenchmarkSupport) + scenario_prediction_fidelity_status: str = "unknown" + scenario_prediction_fidelity_blockers: list[str] = field(default_factory=list) + varied_parallelism_dimensions: list[str] = field(default_factory=list) + varied_workload_dimensions: list[str] = field(default_factory=list) + varied_runtime_dimensions: list[str] = field(default_factory=list) + runtime_mismatch_dimensions: list[str] = field(default_factory=list) + candidate_runtime_signature_counts: dict[str, int] = field(default_factory=dict) + scored_runtime_signature_counts: dict[str, int] = field(default_factory=dict) + prediction_confidence_counts: dict[str, int] = field(default_factory=dict) + calibration_scope_counts: dict[str, int] = field(default_factory=dict) + memory_basis_counts: dict[str, int] = field(default_factory=dict) + memory_coverage_status_counts: dict[str, int] = field(default_factory=dict) + simulator_support_status_counts: dict[str, int] = field(default_factory=dict) + simulator_support_blocker_counts: dict[str, int] = field(default_factory=dict) + max_estimated_memory_residual_gb: float | None = None + max_estimated_memory_residual_gb_label: str | None = None + max_estimated_memory_residual_fraction_of_peak: float | None = None + max_estimated_memory_residual_fraction_of_peak_label: str | None = None + phase_bottleneck_candidate_count: int = 0 + phase_bottleneck_bucket_counts: dict[str, int] = field(default_factory=dict) + phase_bottleneck_phase_counts: dict[str, int] = field(default_factory=dict) + max_phase_bottleneck_share: float | None = None + max_phase_bottleneck_share_label: str | None = None + max_phase_bottleneck_phase: str | None = None + max_phase_bottleneck_bucket: str | None = None + phase_bottleneck_half_speedup_candidate_count: int = 0 + max_phase_bottleneck_half_speedup_delta_pct: float | None = None + max_phase_bottleneck_half_speedup_delta_label: str | None = None + max_phase_bottleneck_half_speedup_phase: str | None = None + max_phase_bottleneck_half_speedup_bucket: str | None = None + memory_bottleneck_candidate_count: int = 0 + memory_bottleneck_bucket_counts: dict[str, int] = field(default_factory=dict) + memory_bottleneck_phase_counts: dict[str, int] = field(default_factory=dict) + max_memory_bottleneck_fraction_of_peak: float | None = None + max_memory_bottleneck_fraction_label: str | None = None + max_memory_bottleneck_phase: str | None = None + max_memory_bottleneck_bucket: str | None = None + timing_coverage_status_counts: dict[str, int] = field(default_factory=dict) + cross_node_dimension_counts: dict[str, int] = field(default_factory=dict) + feasibility_status_counts: dict[str, int] = field(default_factory=dict) + routing_regime_status: str = "unknown" + routing_regime_counts: dict[str, int] = field(default_factory=dict) + recommendation_counts: dict[str, int] = field(default_factory=dict) + risk_flag_counts: dict[str, int] = field(default_factory=dict) + scenario_readiness: ScenarioReadiness = field(default_factory=ScenarioReadiness) + + +@dataclass(frozen=True) +class ScenarioMeasurementConfig: + label: str + filename: str + config: dict[str, Any] + + +@dataclass(frozen=True) +class ScenarioReport: + base_config_path: str + benchmark_dir: str | None + device_memory_limit_gb: float + memory_safety_factor: float + topology_sweep: str + balanced_routing: bool + world_sizes: list[int] + candidate_count: int + feasible_count: int + best_raw: ScenarioCandidate | None + best_risk_adjusted: ScenarioCandidate | None + best_efficiency: ScenarioCandidate | None + best_risk_adjusted_efficiency: ScenarioCandidate | None + best_next_measurement: ScenarioCandidate | None + best_promotable: ScenarioCandidate | None + decision_summary: ScenarioDecisionSummary + candidates: list[ScenarioCandidate] + planner_context: dict[str, Any] = field(default_factory=dict) + warnings: list[str] = field(default_factory=list) + supplemental_benchmark_dirs: list[str] = field(default_factory=list) + analog_benchmark_dirs: list[str] = field(default_factory=list) + primary_behavior_point_count: int = 0 + supplemental_behavior_point_count: int = 0 + analog_behavior_point_count: int = 0 + total_behavior_point_count: int = 0 + + +@dataclass(frozen=True) +class CalibrationHoldout: + label: str + source: str + topology_label: str + actual_tokens_per_sec: float | None + predicted_tokens_per_sec: float | None + prediction_uncertainty_fraction: float | None + prediction_interval_lower_tokens_per_sec: float | None + prediction_interval_upper_tokens_per_sec: float | None + actual_tokens_in_prediction_interval: bool | None + actual_peak_mem_gb: float | None + predicted_peak_mem_gb: float | None + prediction_status: str + matched_label: str | None + absolute_error_tokens_per_sec: float | None + absolute_percentage_error: float | None + analytic_peak_floor_gb: float | None + memory_prediction_basis: str + memory_coverage_status: str + memory_feasibility_status: str + predicted_memory_residual_gb: float | None + predicted_memory_residual_fraction_of_peak: float | None + actual_memory_residual_gb: float | None + actual_memory_residual_fraction_of_peak: float | None + memory_absolute_error_gb: float | None + memory_absolute_percentage_error: float | None + actual_memory_bottleneck_phase: str | None + actual_memory_bottleneck_bucket: str | None + actual_memory_bottleneck_peak_gb: float | None + actual_memory_bottleneck_fraction_of_peak: float | None + predicted_memory_bottleneck_phase: str | None + predicted_memory_bottleneck_bucket: str | None + predicted_memory_bottleneck_peak_gb: float | None + predicted_memory_bottleneck_fraction_of_peak: float | None + memory_bottleneck_phase_match: bool | None + memory_bottleneck_bucket_match: bool | None + memory_bottleneck_peak_absolute_error_gb: float | None + memory_bottleneck_fraction_absolute_error: float | None + actual_phase_bottleneck_phase: str | None + actual_phase_bottleneck_bucket: str | None + actual_phase_bottleneck_share: float | None + predicted_phase_bottleneck_phase: str | None + predicted_phase_bottleneck_bucket: str | None + predicted_phase_bottleneck_share: float | None + phase_bottleneck_phase_match: bool | None + phase_bottleneck_bucket_match: bool | None + phase_bottleneck_share_absolute_error: float | None + actual_phase_top3: list[str] + predicted_phase_top3: list[str] + actual_phase_bucket_top3: list[str] + predicted_phase_bucket_top3: list[str] + phase_top3_overlap_count: int | None + phase_top3_overlap_rate: float | None + phase_bucket_top3_overlap_count: int | None + phase_bucket_top3_overlap_rate: float | None + memory_calibration_source: str | None + calibrated_from_count: int + memory_calibration_notes: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class CalibrationValidationGap: + gap_status: str + priority: int + required_measurement: str + reason: str + affected_holdout_count: int + affected_holdout_labels: list[str] + blocker_names: list[str] + max_absolute_percentage_error: float | None = None + max_absolute_percentage_error_label: str | None = None + max_memory_absolute_error_gb: float | None = None + max_memory_absolute_error_label: str | None = None + max_phase_bottleneck_share_absolute_error: float | None = None + max_phase_bottleneck_share_absolute_error_label: str | None = None + missing_memory_count: int = 0 + missing_phase_bottleneck_count: int = 0 + + +@dataclass(frozen=True) +class CalibrationReport: + base_config_path: str + benchmark_dir: str + status: str + measured_point_count: int + evaluated_count: int + skipped_count: int + mean_absolute_percentage_error: float | None + median_absolute_percentage_error: float | None + max_absolute_percentage_error: float | None + max_absolute_percentage_error_label: str | None + max_absolute_percentage_error_prediction_status: str | None + max_absolute_percentage_error_in_prediction_interval: bool | None + prediction_interval_coverage_count: int + prediction_interval_coverage_rate: float | None + mean_prediction_uncertainty_fraction: float | None + max_prediction_uncertainty_fraction: float | None + max_prediction_uncertainty_label: str | None + memory_evaluated_count: int + mean_memory_absolute_error_gb: float | None + max_memory_absolute_error_gb: float | None + max_memory_absolute_error_label: str | None + mean_memory_absolute_percentage_error: float | None + max_memory_absolute_percentage_error: float | None + max_memory_absolute_percentage_error_label: str | None + memory_prediction_basis_counts: dict[str, int] + memory_coverage_status_counts: dict[str, int] + memory_feasibility_status_counts: dict[str, int] + max_predicted_memory_residual_gb: float | None + max_predicted_memory_residual_gb_label: str | None + max_predicted_memory_residual_fraction_of_peak: float | None + max_predicted_memory_residual_fraction_of_peak_label: str | None + max_actual_memory_residual_gb: float | None + max_actual_memory_residual_gb_label: str | None + max_actual_memory_residual_fraction_of_peak: float | None + max_actual_memory_residual_fraction_of_peak_label: str | None + memory_bottleneck_evaluated_count: int + memory_bottleneck_phase_match_count: int + memory_bottleneck_phase_match_rate: float | None + memory_bottleneck_bucket_match_count: int + memory_bottleneck_bucket_match_rate: float | None + mean_memory_bottleneck_peak_absolute_error_gb: float | None + max_memory_bottleneck_peak_absolute_error_gb: float | None + max_memory_bottleneck_peak_absolute_error_label: str | None + mean_memory_bottleneck_fraction_absolute_error: float | None + max_memory_bottleneck_fraction_absolute_error: float | None + max_memory_bottleneck_fraction_absolute_error_label: str | None + memory_bottleneck_phase_mismatch_labels: list[str] + memory_bottleneck_bucket_mismatch_labels: list[str] + phase_bottleneck_evaluated_count: int + phase_bottleneck_phase_match_count: int + phase_bottleneck_phase_match_rate: float | None + phase_bottleneck_bucket_match_count: int + phase_bottleneck_bucket_match_rate: float | None + mean_phase_bottleneck_share_absolute_error: float | None + max_phase_bottleneck_share_absolute_error: float | None + max_phase_bottleneck_share_absolute_error_label: str | None + phase_bottleneck_phase_mismatch_labels: list[str] + phase_bottleneck_bucket_mismatch_labels: list[str] + phase_top3_evaluated_count: int + mean_phase_top3_overlap_rate: float | None + min_phase_top3_overlap_rate: float | None + min_phase_top3_overlap_rate_label: str | None + mean_phase_bucket_top3_overlap_rate: float | None + min_phase_bucket_top3_overlap_rate: float | None + min_phase_bucket_top3_overlap_rate_label: str | None + calibration_fidelity_status: str + calibration_fidelity_blockers: list[str] + calibration_validation_gap_count: int + calibration_validation_gap_status_counts: dict[str, int] + calibration_validation_gap_required_measurements: list[str] + calibration_validation_gaps: list[CalibrationValidationGap] + prediction_status_counts: dict[str, int] + holdouts: list[CalibrationHoldout] + warnings: list[str] = field(default_factory=list) + prediction_uncertainty_calibration_status: str = "not_evaluated" + mean_empirical_required_uncertainty_fraction: float | None = None + max_empirical_required_uncertainty_fraction: float | None = None + max_empirical_required_uncertainty_label: str | None = None + measurement_design_config_count: int = 0 + measurement_design_config_labels: list[str] = field(default_factory=list) + measurement_design_config_filenames: list[str] = field(default_factory=list) + calibration_support_benchmark_dirs: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class FeasibilityHoldout: + label: str + source: str + topology_label: str + actual_outcome: str + predicted_outcome: str + actual_tokens_per_sec: float | None + actual_peak_mem_gb: float | None + predicted_tokens_per_sec: float | None + predicted_peak_mem_gb: float | None + prediction_status: str + matched_label: str | None + memory_prediction_basis: str + analytic_peak_floor_gb: float | None + memory_coverage_status: str + predicted_memory_residual_gb: float | None + predicted_memory_residual_fraction_of_peak: float | None + actual_memory_residual_gb: float | None + actual_memory_residual_fraction_of_peak: float | None + memory_calibration_source: str | None + predicted_feasibility_status: str + classified_correctly: bool + calibrated_from_count: int + memory_calibration_notes: list[str] = field(default_factory=list) + risk_flags: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class FeasibilityReport: + base_config_path: str + benchmark_dir: str + status: str + observed_point_count: int + evaluated_count: int + skipped_count: int + actual_fit_count: int + actual_oom_count: int + predicted_fit_count: int + predicted_blocked_count: int + predicted_unknown_count: int + correct_count: int + false_fit_count: int + false_blocked_count: int + accuracy: float | None + fit_recall: float | None + oom_recall: float | None + prediction_status_counts: dict[str, int] + memory_prediction_basis_counts: dict[str, int] + memory_coverage_status_counts: dict[str, int] + feasibility_status_counts: dict[str, int] + max_predicted_memory_residual_gb: float | None + max_predicted_memory_residual_gb_label: str | None + max_predicted_memory_residual_fraction_of_peak: float | None + max_predicted_memory_residual_fraction_of_peak_label: str | None + max_actual_memory_residual_gb: float | None + max_actual_memory_residual_gb_label: str | None + max_actual_memory_residual_fraction_of_peak: float | None + max_actual_memory_residual_fraction_of_peak_label: str | None + risk_flag_counts: dict[str, int] + holdouts: list[FeasibilityHoldout] + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class AnalogHoldout: + label: str + source: str + target_model_ref: str | None + target_sequence_len: int | None + target_runtime_signature: str + topology_label: str + target_world_size: int + actual_tokens_per_sec: float + actual_tokens_per_gpu_per_sec: float + predicted_tokens_per_sec: float | None + predicted_tokens_per_gpu_per_sec: float | None + prediction_uncertainty_fraction: float | None + prediction_interval_lower_tokens_per_sec: float | None + prediction_interval_upper_tokens_per_sec: float | None + actual_tokens_in_prediction_interval: bool | None + prediction_status: str + matched_label: str | None + matched_model_ref: str | None + matched_sequence_len: int | None + matched_topology_label: str | None + calibration_distance: float | None + calibration_distance_factors: list[str] + absolute_error_tokens_per_sec: float | None + absolute_percentage_error: float | None + analog_point_count: int + cross_model_active_param_ratio: float | None = None + cross_model_active_param_scale: float | None = None + cross_model_reference_active_params_b: float | None = None + cross_model_target_active_params_b: float | None = None + cross_model_sequence_length_factor: float | None = None + cross_model_parallelism_factor: float | None = None + cross_model_memory_factor: float | None = None + actual_tflops_per_gpu: float | None = None + matched_analog_tflops_per_gpu: float | None = None + # measured target tflops / measured analog tflops: the MFU-regime transfer ratio (None unless + # both rows logged tflops). Far from 1.0 => the equal-MFU analog assumption is measured-false. + mfu_regime_ratio: float | None = None + nearest_analog_label: str | None = None + nearest_analog_model_ref: str | None = None + nearest_analog_sequence_len: int | None = None + nearest_analog_topology_label: str | None = None + nearest_analog_sequence_length_factor: float | None = None + nearest_analog_runtime_mismatch_count: int | None = None + nearest_analog_runtime_mismatches: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + target_local_world_size: int | None = None + target_micro_batch_size: int | None = None + target_gradient_accumulation_steps: int | None = None + target_global_batch_size: int | None = None + target_data_parallel_replicate_size: int | None = None + target_data_parallel_shard_size: int | None = None + target_tensor_parallel_size: int | None = None + target_pipeline_parallel_size: int | None = None + target_expert_parallel_size: int | None = None + target_ep_fsdp_size: int | None = None + target_ulysses_parallel_size: int | None = None + target_ringattn_parallel_size: int | None = None + + +@dataclass(frozen=True) +class CrossModelFactorCoverage: + factor: str + status: str + scored_candidate_count: int + observed_value_count: int + missing_value_count: int + unique_value_count: int + min_value: float | None = None + max_value: float | None = None + values: list[float] = field(default_factory=list) + + +@dataclass(frozen=True) +class AnalogDecisionSummary: + status: str + scored_candidate_count: int + actual_best_label: str | None + predicted_best_label: str | None + actual_best_tokens_per_sec: float | None + selected_actual_tokens_per_sec: float | None + selected_regret_tokens_per_sec: float | None + selected_regret_percentage: float | None + selected_is_actual_best: bool | None + actual_best_in_predicted_top_tie: bool | None + predicted_top_tie_count: int + top_tie_best_actual_tokens_per_sec: float | None + top_tie_regret_tokens_per_sec: float | None + top_tie_regret_percentage: float | None + actual_best_in_prediction_interval_top: bool | None + prediction_interval_top_labels: list[str] + prediction_interval_top_count: int + prediction_interval_top_fraction: float | None + prediction_interval_selectivity_status: str + interval_top_best_actual_tokens_per_sec: float | None + interval_top_regret_tokens_per_sec: float | None + interval_top_regret_percentage: float | None + pairwise_ordering_pair_count: int + pairwise_ordering_correct_count: int + pairwise_ordering_accuracy: float | None + mean_absolute_rank_error: float | None + max_absolute_rank_error: int | None + max_absolute_rank_error_label: str | None + actual_efficiency_best_label: str | None + predicted_efficiency_best_label: str | None + actual_efficiency_best_tokens_per_gpu_per_sec: float | None + selected_efficiency_actual_tokens_per_gpu_per_sec: float | None + efficiency_regret_tokens_per_gpu_per_sec: float | None + efficiency_regret_percentage: float | None + efficiency_selected_is_actual_best: bool | None + efficiency_pairwise_ordering_pair_count: int + efficiency_pairwise_ordering_correct_count: int + efficiency_pairwise_ordering_accuracy: float | None + efficiency_mean_absolute_rank_error: float | None + efficiency_max_absolute_rank_error: int | None + efficiency_max_absolute_rank_error_label: str | None + analog_support_status: str = "unknown" + cross_model_prediction_status: str = "unknown" + cross_model_prediction_blockers: list[str] = field(default_factory=list) + analog_generalization_scope_status: str = "unknown" + analog_generalization_scope_blockers: list[str] = field(default_factory=list) + larger_model_generalization_status: str = "unknown" + larger_model_generalization_blockers: list[str] = field(default_factory=list) + scored_unique_prediction_count: int = 0 + scored_unique_matched_label_count: int = 0 + scored_unique_matched_topology_count: int = 0 + scored_unique_target_topology_count: int = 0 + scored_unique_target_runtime_signature_count: int = 0 + scored_cross_model_factor_status: str = "unknown" + scored_unique_cross_model_factor_count: int = 0 + scored_cross_model_factor_ranges: dict[str, list[float]] = field(default_factory=dict) + scored_cross_model_factor_coverage_status_counts: dict[str, int] = field(default_factory=dict) + scored_cross_model_factor_coverage: list[CrossModelFactorCoverage] = field(default_factory=list) + scored_varied_cross_model_factor_names: list[str] = field(default_factory=list) + scored_degenerate_cross_model_factor_names: list[str] = field(default_factory=list) + scored_missing_cross_model_factor_names: list[str] = field(default_factory=list) + larger_model_generalization_readiness: LargerModelGeneralizationReadiness = field( + default_factory=lambda: LargerModelGeneralizationReadiness() + ) + + +@dataclass(frozen=True) +class AnalogValidationGap: + gap_status: str + priority: int + required_measurement: str + reason: str + affected_target_count: int + affected_target_labels: list[str] + matched_analog_label_count: int + matched_analog_labels: list[str] + matched_analog_model_refs: list[str] + matched_analog_sequence_lengths: list[int] + nearest_analog_label_count: int + nearest_analog_labels: list[str] + nearest_analog_model_refs: list[str] + nearest_analog_sequence_lengths: list[int] + nearest_analog_sequence_length_factors: list[float] + nearest_analog_runtime_mismatches: list[str] + target_sequence_lengths: list[int] + analog_sequence_lengths: list[int] + target_workload_dimensions: list[str] + target_runtime_dimensions: list[str] + target_parallelism_dimensions: list[str] + cross_model_factor_ranges: dict[str, list[float]] + max_calibration_distance: float | None + max_calibration_distance_label: str | None + analog_support_status: str + cross_model_prediction_status: str + analog_generalization_scope_status: str + blocker_names: list[str] + mfu_regime_ratios: list[float] = field(default_factory=list) + + +@dataclass(frozen=True) +class LargerModelGeneralizationReadiness: + readiness_status: str = "unknown_larger_model_generalization_readiness" + can_generalize_to_larger_model: bool = False + support_status: str = "unknown_larger_model_generalization_support" + support_blockers: list[str] = field(default_factory=list) + analog_support_status: str = "unknown" + cross_model_prediction_status: str = "unknown" + cross_model_prediction_blockers: list[str] = field(default_factory=list) + analog_generalization_scope_status: str = "unknown" + analog_generalization_scope_blockers: list[str] = field(default_factory=list) + required_measurements: list[str] = field(default_factory=list) + validation_gap_count: int = 0 + validation_gap_status_counts: dict[str, int] = field(default_factory=dict) + top_validation_gap_statuses: list[str] = field(default_factory=list) + top_validation_gap_required_measurements: list[str] = field(default_factory=list) + top_validation_gap_target_labels: list[str] = field(default_factory=list) + top_validation_gap_nearest_analog_labels: list[str] = field(default_factory=list) + top_validation_gap_nearest_analog_model_refs: list[str] = field(default_factory=list) + top_validation_gap_nearest_analog_sequence_lengths: list[int] = field(default_factory=list) + top_validation_gap_nearest_analog_sequence_length_factors: list[float] = field(default_factory=list) + top_validation_gap_nearest_analog_runtime_mismatches: list[str] = field(default_factory=list) + unscored_target_nearest_analog_labels: list[str] = field(default_factory=list) + unscored_target_nearest_analog_model_refs: list[str] = field(default_factory=list) + unscored_target_nearest_analog_sequence_lengths: list[int] = field(default_factory=list) + unscored_target_nearest_analog_sequence_length_factors: list[float] = field(default_factory=list) + unscored_target_nearest_analog_runtime_mismatches: list[str] = field(default_factory=list) + unscored_target_context_gap_summaries: list[dict[str, Any]] = field(default_factory=list) + measurement_design_config_count: int = 0 + measurement_design_config_labels: list[str] = field(default_factory=list) + measurement_design_config_filenames: list[str] = field(default_factory=list) + evaluated_count: int = 0 + scored_candidate_count: int = 0 + unscored_count: int = 0 + analog_point_count: int = 0 + scored_unique_prediction_count: int = 0 + scored_unique_matched_label_count: int = 0 + scored_unique_matched_topology_count: int = 0 + scored_unique_target_topology_count: int = 0 + scored_unique_target_runtime_signature_count: int = 0 + scored_cross_model_factor_status: str = "unknown" + scored_unique_cross_model_factor_count: int = 0 + scored_cross_model_factor_ranges: dict[str, list[float]] = field(default_factory=dict) + scored_cross_model_factor_coverage_status_counts: dict[str, int] = field(default_factory=dict) + scored_cross_model_factor_coverage: list[CrossModelFactorCoverage] = field(default_factory=list) + scored_varied_cross_model_factor_names: list[str] = field(default_factory=list) + scored_degenerate_cross_model_factor_names: list[str] = field(default_factory=list) + scored_missing_cross_model_factor_names: list[str] = field(default_factory=list) + scored_varied_target_workload_dimensions: list[str] = field(default_factory=list) + scored_varied_target_runtime_dimensions: list[str] = field(default_factory=list) + scored_varied_target_parallelism_dimensions: list[str] = field(default_factory=list) + target_model_refs: list[str] = field(default_factory=list) + analog_model_refs: list[str] = field(default_factory=list) + target_sequence_lengths: list[int] = field(default_factory=list) + analog_sequence_lengths: list[int] = field(default_factory=list) + selected_is_actual_best: bool | None = None + actual_best_in_predicted_top_tie: bool | None = None + prediction_interval_top_fraction: float | None = None + prediction_interval_selectivity_status: str = "unknown" + pairwise_ordering_accuracy: float | None = None + efficiency_selected_is_actual_best: bool | None = None + efficiency_pairwise_ordering_accuracy: float | None = None + max_absolute_percentage_error: float | None = None + prediction_interval_coverage_rate: float | None = None + + +@dataclass(frozen=True) +class AnalogReport: + base_config_path: str + benchmark_dir: str + analog_benchmark_dirs: list[str] + status: str + coverage_status: str + measured_point_count: int + evaluated_count: int + unscored_count: int + skipped_count: int + target_coverage_fraction: float | None + analog_point_count: int + analog_support_status: str + cross_model_prediction_status: str + cross_model_prediction_blockers: list[str] + analog_generalization_scope_status: str + analog_generalization_scope_blockers: list[str] + larger_model_generalization_status: str + larger_model_generalization_blockers: list[str] + analog_validation_gap_count: int + analog_validation_gap_status_counts: dict[str, int] + analog_validation_gap_required_measurements: list[str] + analog_validation_gaps: list[AnalogValidationGap] + scored_unique_prediction_count: int + scored_unique_matched_label_count: int + scored_unique_matched_topology_count: int + scored_unique_target_topology_count: int + scored_unique_target_runtime_signature_count: int + scored_cross_model_factor_status: str + scored_unique_cross_model_factor_count: int + scored_cross_model_factor_ranges: dict[str, list[float]] + scored_cross_model_factor_coverage_status_counts: dict[str, int] + scored_cross_model_factor_coverage: list[CrossModelFactorCoverage] + scored_varied_cross_model_factor_names: list[str] + scored_degenerate_cross_model_factor_names: list[str] + scored_missing_cross_model_factor_names: list[str] + scored_varied_target_workload_dimensions: list[str] + scored_varied_target_runtime_dimensions: list[str] + scored_varied_target_parallelism_dimensions: list[str] + target_model_refs: list[str] + analog_model_refs: list[str] + target_sequence_lengths: list[int] + evaluated_target_sequence_lengths: list[int] + unscored_target_sequence_lengths: list[int] + unscored_target_labels: list[str] + unscored_target_sequence_length_counts: dict[int, int] + unscored_target_reason_counts: dict[str, int] + unscored_target_nearest_analog_labels: list[str] + unscored_target_nearest_analog_model_refs: list[str] + unscored_target_nearest_analog_sequence_lengths: list[int] + unscored_target_nearest_analog_sequence_length_factors: list[float] + unscored_target_nearest_analog_runtime_mismatches: list[str] + unscored_target_context_gap_summaries: list[dict[str, Any]] + analog_sequence_lengths: list[int] + mean_absolute_percentage_error: float | None + median_absolute_percentage_error: float | None + max_absolute_percentage_error: float | None + max_absolute_percentage_error_label: str | None + max_absolute_percentage_error_prediction_status: str | None + max_absolute_percentage_error_in_prediction_interval: bool | None + mean_scored_calibration_distance: float | None + max_calibration_distance: float | None + max_calibration_distance_label: str | None + prediction_interval_coverage_count: int + prediction_interval_coverage_rate: float | None + prediction_interval_top_fraction: float | None + prediction_interval_selectivity_status: str + mean_prediction_uncertainty_fraction: float | None + max_prediction_uncertainty_fraction: float | None + max_prediction_uncertainty_label: str | None + prediction_status_counts: dict[str, int] + decision_summary: AnalogDecisionSummary + holdouts: list[AnalogHoldout] + larger_model_generalization_readiness: LargerModelGeneralizationReadiness = field( + default_factory=LargerModelGeneralizationReadiness + ) + warnings: list[str] = field(default_factory=list) + supplemental_benchmark_dirs: list[str] = field(default_factory=list) + primary_target_point_count: int = 0 + supplemental_target_point_count: int = 0 + primary_measured_point_count: int = 0 + supplemental_measured_point_count: int = 0 + + +@dataclass(frozen=True) +class DecisionCandidatePrediction: + label: str + source: str + topology_label: str + actual_tokens_per_sec: float + actual_tokens_per_gpu_per_sec: float | None + predicted_tokens_per_sec: float | None + predicted_tokens_per_gpu_per_sec: float | None + predicted_risk_adjusted_tokens_per_sec: float | None + predicted_risk_adjusted_tokens_per_gpu_per_sec: float | None + prediction_uncertainty_fraction: float | None + predicted_interval_lower_tokens_per_sec: float | None + predicted_interval_upper_tokens_per_sec: float | None + predicted_risk_adjusted_interval_lower_tokens_per_sec: float | None + predicted_risk_adjusted_interval_upper_tokens_per_sec: float | None + prediction_status: str + matched_label: str | None + calibration_distance: float | None + feasibility_status: str + actual_rank: int + actual_efficiency_rank: int + predicted_rank: int | None + risk_adjusted_rank: int | None + predicted_efficiency_rank: int | None + risk_adjusted_efficiency_rank: int | None + actual_frontier_member: bool + predicted_frontier_member: bool + risk_adjusted_frontier_member: bool + actual_scaling_baseline_label: str | None + actual_scaling_gpu_ratio: float | None + actual_scaling_speedup: float | None + actual_scaling_efficiency: float | None + predicted_scaling_baseline_label: str | None + predicted_scaling_speedup: float | None + predicted_scaling_efficiency: float | None + risk_adjusted_scaling_baseline_label: str | None + risk_adjusted_scaling_speedup: float | None + risk_adjusted_scaling_efficiency: float | None + risk_flags: list[str] = field(default_factory=list) + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class DecisionHoldout: + heldout_label: str + heldout_source: str + actual_best_label: str + predicted_best_label: str | None + actual_best_tokens_per_sec: float + selected_actual_tokens_per_sec: float | None + regret_tokens_per_sec: float | None + regret_percentage: float | None + risk_adjusted_predicted_best_label: str | None + risk_adjusted_selected_actual_tokens_per_sec: float | None + risk_adjusted_regret_tokens_per_sec: float | None + risk_adjusted_regret_percentage: float | None + actual_efficiency_best_label: str + actual_efficiency_best_tokens_per_gpu_per_sec: float | None + predicted_efficiency_best_label: str | None + selected_actual_efficiency_tokens_per_gpu_per_sec: float | None + efficiency_regret_tokens_per_gpu_per_sec: float | None + efficiency_regret_percentage: float | None + risk_adjusted_efficiency_predicted_best_label: str | None + risk_adjusted_efficiency_selected_actual_tokens_per_gpu_per_sec: float | None + risk_adjusted_efficiency_regret_tokens_per_gpu_per_sec: float | None + risk_adjusted_efficiency_regret_percentage: float | None + selected_is_actual_best: bool + actual_best_in_predicted_top_tie: bool + risk_adjusted_selected_is_actual_best: bool + actual_best_in_risk_adjusted_top_tie: bool + actual_best_in_risk_adjusted_interval_top: bool + risk_adjusted_interval_top_labels: list[str] + risk_adjusted_interval_top_count: int + risk_adjusted_interval_top_best_actual_tokens_per_sec: float | None + risk_adjusted_interval_top_regret_tokens_per_sec: float | None + risk_adjusted_interval_top_regret_percentage: float | None + pairwise_ordering_pair_count: int + pairwise_ordering_correct_count: int + pairwise_ordering_accuracy: float | None + risk_adjusted_pairwise_ordering_pair_count: int + risk_adjusted_pairwise_ordering_correct_count: int + risk_adjusted_pairwise_ordering_accuracy: float | None + mean_absolute_rank_error: float | None + max_absolute_rank_error: int | None + risk_adjusted_mean_absolute_rank_error: float | None + risk_adjusted_max_absolute_rank_error: int | None + efficiency_pairwise_ordering_pair_count: int + efficiency_pairwise_ordering_correct_count: int + efficiency_pairwise_ordering_accuracy: float | None + risk_adjusted_efficiency_pairwise_ordering_pair_count: int + risk_adjusted_efficiency_pairwise_ordering_correct_count: int + risk_adjusted_efficiency_pairwise_ordering_accuracy: float | None + efficiency_mean_absolute_rank_error: float | None + efficiency_max_absolute_rank_error: int | None + risk_adjusted_efficiency_mean_absolute_rank_error: float | None + risk_adjusted_efficiency_max_absolute_rank_error: int | None + efficiency_selected_is_actual_best: bool + actual_efficiency_best_in_predicted_top_tie: bool + risk_adjusted_efficiency_selected_is_actual_best: bool + actual_efficiency_best_in_risk_adjusted_top_tie: bool + actual_frontier_labels: list[str] + predicted_frontier_labels: list[str] + risk_adjusted_frontier_labels: list[str] + actual_frontier_count: int + predicted_frontier_count: int + risk_adjusted_frontier_count: int + actual_frontier_in_predicted_count: int + actual_frontier_in_risk_adjusted_count: int + actual_frontier_predicted_coverage_fraction: float | None + actual_frontier_risk_adjusted_coverage_fraction: float | None + predicted_frontier_extra_count: int + risk_adjusted_frontier_extra_count: int + actual_frontier_missed_labels: list[str] + actual_frontier_risk_adjusted_missed_labels: list[str] + scaling_candidate_count: int + predicted_scaling_candidate_count: int + risk_adjusted_scaling_candidate_count: int + mean_scaling_efficiency_absolute_error: float | None + max_scaling_efficiency_absolute_error: float | None + max_scaling_efficiency_absolute_error_label: str | None + mean_risk_adjusted_scaling_efficiency_absolute_error: float | None + max_risk_adjusted_scaling_efficiency_absolute_error: float | None + max_risk_adjusted_scaling_efficiency_absolute_error_label: str | None + predicted_top_tie_count: int + risk_adjusted_top_tie_count: int + predicted_efficiency_top_tie_count: int + risk_adjusted_efficiency_top_tie_count: int + candidate_count: int + predicted_unscored_count: int + risk_adjusted_unscored_count: int + predicted_efficiency_unscored_count: int + risk_adjusted_efficiency_unscored_count: int + candidates: list[DecisionCandidatePrediction] + warnings: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ParallelismBoundaryGroup: + signature: str + candidate_count: int + fit_count: int + failure_count: int + best_fit_label: str | None + best_fit_tokens_per_sec: float | None + failure_labels: list[str] + varied_parallelism_dimensions: list[str] + confounded_workload_dimensions: list[str] + confounded_runtime_dimensions: list[str] + + +@dataclass(frozen=True) +class ParallelismBoundaryAxisCoverage: + axis: str + status: str + group_count: int + candidate_count: int + fit_count: int + failure_count: int + varied_parallelism_dimensions: list[str] + co_varied_parallelism_dimensions: list[str] + confounded_workload_dimensions: list[str] + confounded_runtime_dimensions: list[str] + boundary_candidate_labels: list[str] = field(default_factory=list) + best_fit_labels: list[str] = field(default_factory=list) + failure_labels: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ParallelismAxisCoverage: + axis: str + status: str + candidate_count: int + varied_parallelism_dimensions: list[str] + co_varied_parallelism_dimensions: list[str] + like_for_like_group_count: int + evaluated_count: int + confounded_workload_dimensions: list[str] + confounded_runtime_dimensions: list[str] + candidate_labels: list[str] = field(default_factory=list) + evaluated_labels: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class ParallelismValidationGap: + axis: str + gap_status: str + priority: int + required_measurement: str + reason: str + throughput_axis_status: str + boundary_axis_status: str + candidate_count: int + boundary_candidate_count: int + fit_count: int + failure_count: int + co_varied_parallelism_dimensions: list[str] + confounded_workload_dimensions: list[str] + confounded_runtime_dimensions: list[str] + throughput_candidate_labels: list[str] = field(default_factory=list) + throughput_evaluated_labels: list[str] = field(default_factory=list) + boundary_candidate_labels: list[str] = field(default_factory=list) + boundary_best_fit_labels: list[str] = field(default_factory=list) + boundary_failure_labels: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class OptimalParallelismReadiness: + readiness_status: str = "unknown_optimal_parallelism_readiness" + can_predict_optimal_tradeoff: bool = False + support_status: str = "unknown_optimal_parallelism_support" + support_blockers: list[str] = field(default_factory=list) + prediction_status: str = "unknown_parallelism_prediction" + validation_scope_status: str = "unknown_parallelism_validation_scope" + boundary_prediction_status: str = "unknown_parallelism_boundary_prediction" + required_measurements: list[str] = field(default_factory=list) + validation_gap_count: int = 0 + validation_gap_axis_names: list[str] = field(default_factory=list) + validation_gap_status_counts: dict[str, int] = field(default_factory=dict) + top_validation_gap_statuses: list[str] = field(default_factory=list) + top_validation_gap_required_measurements: list[str] = field(default_factory=list) + top_validation_gap_axis_names: list[str] = field(default_factory=list) + measured_parallelism_axis_names: list[str] = field(default_factory=list) + isolated_measured_parallelism_axis_names: list[str] = field(default_factory=list) + coupled_measured_parallelism_axis_names: list[str] = field(default_factory=list) + confounded_parallelism_axis_names: list[str] = field(default_factory=list) + missing_parallelism_axis_names: list[str] = field(default_factory=list) + parallelism_evaluated_count: int = 0 + parallelism_axis_evaluated_count: int = 0 + like_for_like_parallelism_group_count: int = 0 + parallelism_boundary_group_count: int = 0 + parallelism_boundary_fit_count: int = 0 + parallelism_boundary_failure_count: int = 0 + parallelism_top1_selection_hit_rate: float | None = None + risk_adjusted_parallelism_top1_selection_hit_rate: float | None = None + risk_adjusted_parallelism_interval_selection_hit_rate: float | None = None + parallelism_pairwise_ordering_accuracy: float | None = None + risk_adjusted_parallelism_pairwise_ordering_accuracy: float | None = None + efficiency_parallelism_pairwise_ordering_accuracy: float | None = None + risk_adjusted_efficiency_parallelism_pairwise_ordering_accuracy: float | None = None + parallelism_frontier_coverage_hit_rate: float | None = None + mean_parallelism_frontier_coverage_fraction: float | None = None + risk_adjusted_parallelism_frontier_coverage_hit_rate: float | None = None + mean_risk_adjusted_parallelism_frontier_coverage_fraction: float | None = None + parallelism_frontier_missed_labels: list[str] = field(default_factory=list) + risk_adjusted_parallelism_frontier_missed_labels: list[str] = field(default_factory=list) + measurement_design_config_count: int = 0 + measurement_design_config_labels: list[str] = field(default_factory=list) + measurement_design_config_filenames: list[str] = field(default_factory=list) + + +@dataclass(frozen=True) +class DecisionReport: + base_config_path: str + benchmark_dir: str + status: str + measured_point_count: int + promotable_point_count: int + best_actual_label: str | None + best_actual_tokens_per_sec: float | None + best_actual_promotable_label: str | None + best_actual_promotable_tokens_per_sec: float | None + promotable_actual_gap_tokens_per_sec: float | None + promotable_actual_gap_percentage: float | None + promotion_readiness_status: str + evaluated_count: int + skipped_count: int + hit_count: int + selection_hit_rate: float | None + selected_best_count: int + top1_selection_hit_rate: float | None + mean_regret_percentage: float | None + max_regret_percentage: float | None + max_regret_percentage_label: str | None + risk_adjusted_hit_count: int + risk_adjusted_selection_hit_rate: float | None + risk_adjusted_selected_best_count: int + risk_adjusted_top1_selection_hit_rate: float | None + mean_risk_adjusted_regret_percentage: float | None + max_risk_adjusted_regret_percentage: float | None + max_risk_adjusted_regret_percentage_label: str | None + risk_adjusted_interval_hit_count: int + risk_adjusted_interval_selection_hit_rate: float | None + mean_risk_adjusted_interval_top_count: float | None + max_risk_adjusted_interval_top_count: int | None + mean_risk_adjusted_interval_regret_percentage: float | None + max_risk_adjusted_interval_regret_percentage: float | None + pairwise_ordering_pair_count: int + pairwise_ordering_correct_count: int + pairwise_ordering_accuracy: float | None + mean_pairwise_ordering_accuracy: float | None + risk_adjusted_pairwise_ordering_pair_count: int + risk_adjusted_pairwise_ordering_correct_count: int + risk_adjusted_pairwise_ordering_accuracy: float | None + mean_risk_adjusted_pairwise_ordering_accuracy: float | None + mean_absolute_rank_error: float | None + max_absolute_rank_error: int | None + max_absolute_rank_error_label: str | None + mean_holdout_absolute_rank_error: float | None + risk_adjusted_mean_absolute_rank_error: float | None + risk_adjusted_max_absolute_rank_error: int | None + risk_adjusted_max_absolute_rank_error_label: str | None + mean_risk_adjusted_holdout_absolute_rank_error: float | None + efficiency_hit_count: int + efficiency_selection_hit_rate: float | None + efficiency_selected_best_count: int + efficiency_top1_selection_hit_rate: float | None + mean_efficiency_regret_percentage: float | None + max_efficiency_regret_percentage: float | None + max_efficiency_regret_percentage_label: str | None + risk_adjusted_efficiency_hit_count: int + risk_adjusted_efficiency_selection_hit_rate: float | None + risk_adjusted_efficiency_selected_best_count: int + risk_adjusted_efficiency_top1_selection_hit_rate: float | None + mean_risk_adjusted_efficiency_regret_percentage: float | None + max_risk_adjusted_efficiency_regret_percentage: float | None + max_risk_adjusted_efficiency_regret_percentage_label: str | None + efficiency_pairwise_ordering_pair_count: int + efficiency_pairwise_ordering_correct_count: int + efficiency_pairwise_ordering_accuracy: float | None + mean_efficiency_pairwise_ordering_accuracy: float | None + risk_adjusted_efficiency_pairwise_ordering_pair_count: int + risk_adjusted_efficiency_pairwise_ordering_correct_count: int + risk_adjusted_efficiency_pairwise_ordering_accuracy: float | None + mean_risk_adjusted_efficiency_pairwise_ordering_accuracy: float | None + efficiency_mean_absolute_rank_error: float | None + efficiency_max_absolute_rank_error: int | None + efficiency_max_absolute_rank_error_label: str | None + mean_holdout_efficiency_absolute_rank_error: float | None + risk_adjusted_efficiency_mean_absolute_rank_error: float | None + risk_adjusted_efficiency_max_absolute_rank_error: int | None + risk_adjusted_efficiency_max_absolute_rank_error_label: str | None + mean_risk_adjusted_efficiency_holdout_absolute_rank_error: float | None + frontier_hit_count: int + frontier_coverage_hit_rate: float | None + mean_frontier_coverage_fraction: float | None + min_frontier_coverage_fraction: float | None + mean_predicted_frontier_extra_count: float | None + max_predicted_frontier_extra_count: int | None + frontier_missed_labels: list[str] + risk_adjusted_frontier_hit_count: int + risk_adjusted_frontier_coverage_hit_rate: float | None + mean_risk_adjusted_frontier_coverage_fraction: float | None + min_risk_adjusted_frontier_coverage_fraction: float | None + mean_risk_adjusted_frontier_extra_count: float | None + max_risk_adjusted_frontier_extra_count: int | None + risk_adjusted_frontier_missed_labels: list[str] + scaling_evaluated_count: int + mean_scaling_efficiency_absolute_error: float | None + max_scaling_efficiency_absolute_error: float | None + max_scaling_efficiency_absolute_error_label: str | None + risk_adjusted_scaling_evaluated_count: int + mean_risk_adjusted_scaling_efficiency_absolute_error: float | None + max_risk_adjusted_scaling_efficiency_absolute_error: float | None + max_risk_adjusted_scaling_efficiency_absolute_error_label: str | None + varied_parallelism_dimensions: list[str] + varied_workload_dimensions: list[str] + varied_runtime_dimensions: list[str] + unique_parallelism_strategy_count: int + parallelism_coverage_status: str + parallelism_prediction_status: str + parallelism_prediction_blockers: list[str] + parallelism_validation_scope_status: str + parallelism_validation_scope_blockers: list[str] + optimal_parallelism_support_status: str + optimal_parallelism_support_blockers: list[str] + measured_parallelism_axis_names: list[str] + isolated_measured_parallelism_axis_names: list[str] + coupled_measured_parallelism_axis_names: list[str] + confounded_parallelism_axis_names: list[str] + missing_parallelism_axis_names: list[str] + parallelism_axis_coverage_status_counts: dict[str, int] + parallelism_axis_coverage: list[ParallelismAxisCoverage] + parallelism_validation_gap_count: int + parallelism_validation_gap_axis_names: list[str] + parallelism_validation_gap_status_counts: dict[str, int] + parallelism_validation_gaps: list[ParallelismValidationGap] + like_for_like_parallelism_group_count: int + parallelism_evaluated_count: int + parallelism_axis_group_count: int + parallelism_axis_evaluated_count: int + parallelism_axis_hit_count: int + parallelism_axis_selection_hit_rate: float | None + parallelism_axis_selected_best_count: int + parallelism_axis_top1_selection_hit_rate: float | None + risk_adjusted_parallelism_axis_hit_count: int + risk_adjusted_parallelism_axis_selection_hit_rate: float | None + risk_adjusted_parallelism_axis_selected_best_count: int + risk_adjusted_parallelism_axis_top1_selection_hit_rate: float | None + parallelism_hit_count: int + parallelism_selection_hit_rate: float | None + parallelism_selected_best_count: int + parallelism_top1_selection_hit_rate: float | None + risk_adjusted_parallelism_hit_count: int + risk_adjusted_parallelism_selection_hit_rate: float | None + risk_adjusted_parallelism_selected_best_count: int + risk_adjusted_parallelism_top1_selection_hit_rate: float | None + risk_adjusted_parallelism_interval_hit_count: int + risk_adjusted_parallelism_interval_selection_hit_rate: float | None + mean_risk_adjusted_parallelism_interval_top_count: float | None + max_risk_adjusted_parallelism_interval_top_count: int | None + parallelism_pairwise_ordering_pair_count: int + parallelism_pairwise_ordering_correct_count: int + parallelism_pairwise_ordering_accuracy: float | None + mean_parallelism_pairwise_ordering_accuracy: float | None + risk_adjusted_parallelism_pairwise_ordering_pair_count: int + risk_adjusted_parallelism_pairwise_ordering_correct_count: int + risk_adjusted_parallelism_pairwise_ordering_accuracy: float | None + mean_risk_adjusted_parallelism_pairwise_ordering_accuracy: float | None + parallelism_mean_absolute_rank_error: float | None + parallelism_max_absolute_rank_error: int | None + parallelism_max_absolute_rank_error_label: str | None + mean_parallelism_holdout_absolute_rank_error: float | None + risk_adjusted_parallelism_mean_absolute_rank_error: float | None + risk_adjusted_parallelism_max_absolute_rank_error: int | None + risk_adjusted_parallelism_max_absolute_rank_error_label: str | None + mean_risk_adjusted_parallelism_holdout_absolute_rank_error: float | None + parallelism_axis_pairwise_ordering_pair_count: int + parallelism_axis_pairwise_ordering_correct_count: int + parallelism_axis_pairwise_ordering_accuracy: float | None + mean_parallelism_axis_pairwise_ordering_accuracy: float | None + risk_adjusted_parallelism_axis_pairwise_ordering_pair_count: int + risk_adjusted_parallelism_axis_pairwise_ordering_correct_count: int + risk_adjusted_parallelism_axis_pairwise_ordering_accuracy: float | None + mean_risk_adjusted_parallelism_axis_pairwise_ordering_accuracy: float | None + parallelism_axis_mean_absolute_rank_error: float | None + parallelism_axis_max_absolute_rank_error: int | None + parallelism_axis_max_absolute_rank_error_label: str | None + mean_parallelism_axis_holdout_absolute_rank_error: float | None + risk_adjusted_parallelism_axis_mean_absolute_rank_error: float | None + risk_adjusted_parallelism_axis_max_absolute_rank_error: int | None + risk_adjusted_parallelism_axis_max_absolute_rank_error_label: str | None + mean_risk_adjusted_parallelism_axis_holdout_absolute_rank_error: float | None + efficiency_parallelism_hit_count: int + efficiency_parallelism_selection_hit_rate: float | None + efficiency_parallelism_selected_best_count: int + efficiency_parallelism_top1_selection_hit_rate: float | None + risk_adjusted_efficiency_parallelism_hit_count: int + risk_adjusted_efficiency_parallelism_selection_hit_rate: float | None + risk_adjusted_efficiency_parallelism_selected_best_count: int + risk_adjusted_efficiency_parallelism_top1_selection_hit_rate: float | None + efficiency_parallelism_pairwise_ordering_pair_count: int + efficiency_parallelism_pairwise_ordering_correct_count: int + efficiency_parallelism_pairwise_ordering_accuracy: float | None + mean_efficiency_parallelism_pairwise_ordering_accuracy: float | None + risk_adjusted_efficiency_parallelism_pairwise_ordering_pair_count: int + risk_adjusted_efficiency_parallelism_pairwise_ordering_correct_count: int + risk_adjusted_efficiency_parallelism_pairwise_ordering_accuracy: float | None + mean_risk_adjusted_efficiency_parallelism_pairwise_ordering_accuracy: float | None + efficiency_parallelism_mean_absolute_rank_error: float | None + efficiency_parallelism_max_absolute_rank_error: int | None + efficiency_parallelism_max_absolute_rank_error_label: str | None + mean_efficiency_parallelism_holdout_absolute_rank_error: float | None + risk_adjusted_efficiency_parallelism_mean_absolute_rank_error: float | None + risk_adjusted_efficiency_parallelism_max_absolute_rank_error: int | None + risk_adjusted_efficiency_parallelism_max_absolute_rank_error_label: str | None + mean_risk_adjusted_efficiency_parallelism_holdout_absolute_rank_error: float | None + efficiency_parallelism_axis_pairwise_ordering_pair_count: int + efficiency_parallelism_axis_pairwise_ordering_correct_count: int + efficiency_parallelism_axis_pairwise_ordering_accuracy: float | None + mean_efficiency_parallelism_axis_pairwise_ordering_accuracy: float | None + risk_adjusted_efficiency_parallelism_axis_pairwise_ordering_pair_count: int + risk_adjusted_efficiency_parallelism_axis_pairwise_ordering_correct_count: int + risk_adjusted_efficiency_parallelism_axis_pairwise_ordering_accuracy: float | None + mean_risk_adjusted_efficiency_parallelism_axis_pairwise_ordering_accuracy: float | None + efficiency_parallelism_axis_mean_absolute_rank_error: float | None + efficiency_parallelism_axis_max_absolute_rank_error: int | None + efficiency_parallelism_axis_max_absolute_rank_error_label: str | None + mean_efficiency_parallelism_axis_holdout_absolute_rank_error: float | None + risk_adjusted_efficiency_parallelism_axis_mean_absolute_rank_error: float | None + risk_adjusted_efficiency_parallelism_axis_max_absolute_rank_error: int | None + risk_adjusted_efficiency_parallelism_axis_max_absolute_rank_error_label: str | None + mean_risk_adjusted_efficiency_parallelism_axis_holdout_absolute_rank_error: float | None + parallelism_frontier_hit_count: int + parallelism_frontier_coverage_hit_rate: float | None + mean_parallelism_frontier_coverage_fraction: float | None + min_parallelism_frontier_coverage_fraction: float | None + parallelism_frontier_missed_labels: list[str] + risk_adjusted_parallelism_frontier_hit_count: int + risk_adjusted_parallelism_frontier_coverage_hit_rate: float | None + mean_risk_adjusted_parallelism_frontier_coverage_fraction: float | None + min_risk_adjusted_parallelism_frontier_coverage_fraction: float | None + risk_adjusted_parallelism_frontier_missed_labels: list[str] + parallelism_scaling_evaluated_count: int + mean_parallelism_scaling_efficiency_absolute_error: float | None + max_parallelism_scaling_efficiency_absolute_error: float | None + max_parallelism_scaling_efficiency_absolute_error_label: str | None + risk_adjusted_parallelism_scaling_evaluated_count: int + mean_risk_adjusted_parallelism_scaling_efficiency_absolute_error: float | None + max_risk_adjusted_parallelism_scaling_efficiency_absolute_error: float | None + max_risk_adjusted_parallelism_scaling_efficiency_absolute_error_label: str | None + parallelism_boundary_status: str + parallelism_boundary_prediction_status: str + parallelism_boundary_prediction_blockers: list[str] + parallelism_boundary_group_count: int + parallelism_boundary_candidate_count: int + parallelism_boundary_fit_count: int + parallelism_boundary_failure_count: int + parallelism_boundary_best_fit_label: str | None + parallelism_boundary_confounded_dimensions: list[str] + parallelism_boundary_measured_axis_names: list[str] + parallelism_boundary_confounded_axis_names: list[str] + parallelism_boundary_missing_axis_names: list[str] + parallelism_boundary_axis_coverage_status_counts: dict[str, int] + parallelism_boundary_axis_coverage: list[ParallelismBoundaryAxisCoverage] + parallelism_boundary_groups: list[ParallelismBoundaryGroup] + holdouts: list[DecisionHoldout] + optimal_parallelism_readiness: OptimalParallelismReadiness = field(default_factory=OptimalParallelismReadiness) + warnings: list[str] = field(default_factory=list) + supplemental_benchmark_dirs: list[str] = field(default_factory=list) + primary_behavior_point_count: int = 0 + supplemental_behavior_point_count: int = 0 + primary_measured_point_count: int = 0 + supplemental_measured_point_count: int = 0 diff --git a/experiments/local_benchmark/training_sim/shape_engine.py b/src/xorl/sim/shape_engine.py similarity index 100% rename from experiments/local_benchmark/training_sim/shape_engine.py rename to src/xorl/sim/shape_engine.py diff --git a/src/xorl/sim/simulator_support.py b/src/xorl/sim/simulator_support.py new file mode 100644 index 00000000..c49b3923 --- /dev/null +++ b/src/xorl/sim/simulator_support.py @@ -0,0 +1,211 @@ +"""Declare which training-engine surface a static simulator report covers.""" + +from __future__ import annotations + +from typing import Any + + +try: + from .schemas import MemoryLedger, SimulatorSupportLedger, TimingLedger, Topology +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import MemoryLedger, SimulatorSupportLedger, TimingLedger, Topology + + +_SERVER_FLAT_KEYS = { + "ce_mode", + "data_parallel_mode", + "enable_packing", + "sample_packing_sequence_len", + "skip_initial_checkpoint", + "sync_inference_method", + "worker_connection_timeout", + "worker_max_retries", +} + + +def _section(raw_config: dict[str, Any], name: str) -> dict[str, Any]: + value = raw_config.get(name, {}) + return value if isinstance(value, dict) else {} + + +def _looks_like_flat_server_config(raw_config: dict[str, Any]) -> bool: + if _section(raw_config, "train"): + return False + if "model_path" not in raw_config and "config_path" not in raw_config: + return False + return any(key in raw_config for key in _SERVER_FLAT_KEYS) + + +def requested_simulator_surface(raw_config: dict[str, Any]) -> str: + if _section(raw_config, "server") or _looks_like_flat_server_config(raw_config): + return "server_forward_backward" + if _section(raw_config, "train"): + return "local_training" + return "unknown_config_surface" + + +def _configured_int( + raw_config: dict[str, Any], + key: str, + *, + surface: str, + topology_value: int | None = None, + default: int = 1, +) -> int: + sections: list[dict[str, Any]] + if surface == "server_forward_backward": + sections = [_section(raw_config, "server"), raw_config, _section(raw_config, "train")] + else: + sections = [_section(raw_config, "train"), raw_config, _section(raw_config, "server")] + + for section in sections: + value = section.get(key) + if value is not None: + return int(value) + if topology_value is not None: + return int(topology_value) + return default + + +def _runtime_output_support(memory: MemoryLedger | None, timing: TimingLedger | None) -> tuple[list[str], list[str]]: + supported: list[str] = [] + unsupported: list[str] = [] + + if memory is None: + unsupported.append("memory_ledger_not_built") + else: + if memory.analytic_peak_floor_gb is not None: + supported.append("analytic_memory_floor") + else: + unsupported.append("analytic_memory_floor") + if memory.calibrated_peak_mem_gb is not None or memory.observed_peak_mem_gb_max is not None: + supported.append("calibrated_peak_memory") + else: + unsupported.append("calibrated_peak_memory_without_observed_or_benchmark_anchor") + + if timing is None: + unsupported.append("timing_ledger_not_built") + else: + if timing.step_time_s is not None: + supported.append("step_time_prediction") + else: + unsupported.append("step_time_prediction_without_observed_or_benchmark_anchor") + if timing.phase_time_sec: + supported.append("phase_timing") + else: + unsupported.append("phase_timing_without_observed_or_benchmark_anchor") + + return supported, unsupported + + +def resolve_simulator_support( + raw_config: dict[str, Any], + *, + topology: Topology | None = None, + memory: MemoryLedger | None = None, + timing: TimingLedger | None = None, +) -> SimulatorSupportLedger: + surface = requested_simulator_surface(raw_config) + pipeline_parallel_size = _configured_int( + raw_config, + "pipeline_parallel_size", + surface=surface, + topology_value=topology.pipeline_parallel_size if topology is not None else None, + ) + runtime_supported, runtime_unsupported = _runtime_output_support(memory, timing) + + if surface == "server_forward_backward": + supported = ["config_surface_detection", "model_metadata_resolution"] + if topology is not None: + supported.append("server_parallelism_topology") + blockers = [ + "server_forward_backward_backend_missing", + "opd_loss_path_missing", + "hidden_cache_prefetch_timing_missing", + "server_phase_memory_attribution_missing", + ] + unsupported = [ + "server_forward_backward_timing", + "opd_loss_path_timing", + "hidden_cache_fetch_prefetch_timing", + "server_phase_memory_attribution", + *runtime_unsupported, + ] + notes = [ + "server configs use a flat YAML surface and need a server-specific forward/backward backend", + "server and local forward/backward surfaces must be calibrated separately", + ] + if pipeline_parallel_size > 1: + blockers.extend( + [ + "pp_schedule_event_model_missing", + "pp_forward_backward_peak_not_separated", + ] + ) + unsupported.extend(["pp_schedule_timing", "separate_forward_backward_peak"]) + return SimulatorSupportLedger( + requested_surface=surface, + support_status="unsupported_server_forward_backward", + support_blockers=sorted(set(blockers)), + supported_outputs=sorted(set(supported)), + unsupported_outputs=sorted(set(unsupported)), + notes=notes, + ) + + if surface != "local_training": + return SimulatorSupportLedger( + requested_surface=surface, + support_status="unsupported_unknown_config_surface", + support_blockers=["config_surface_not_recognized"], + supported_outputs=["config_surface_detection"], + unsupported_outputs=sorted( + { + "local_training_topology", + "server_forward_backward_timing", + *runtime_unsupported, + } + ), + notes=["local training configs must use nested train/data/model sections"], + ) + + supported = [ + "config_surface_detection", + "model_metadata_resolution", + "local_training_topology", + "shape_ledger", + *runtime_supported, + ] + unsupported = list(runtime_unsupported) + if pipeline_parallel_size > 1: + blockers = [ + "pp_schedule_event_model_missing", + "pp_forward_backward_peak_not_separated", + "pp_phase_timing_not_separated", + ] + unsupported.extend( + [ + "pp_schedule_timing", + "separate_forward_backward_peak", + "pp_activation_liveness_peak", + ] + ) + return SimulatorSupportLedger( + requested_surface=surface, + support_status="partial_local_pp_memory_only", + support_blockers=blockers, + supported_outputs=sorted({*supported, "pp_stage_parameter_ownership"}), + unsupported_outputs=sorted(set(unsupported)), + notes=[ + "PP reports include topology and analytic ownership, but not a schedule-level event model", + "XoRL reports a combined fwd+bwd peak for PP until schedule attribution is added", + ], + ) + + return SimulatorSupportLedger( + requested_surface=surface, + support_status="supported_local_non_pp", + support_blockers=[], + supported_outputs=sorted(set(supported)), + unsupported_outputs=sorted(set(unsupported)), + notes=[], + ) diff --git a/src/xorl/sim/timing_ledger.py b/src/xorl/sim/timing_ledger.py new file mode 100644 index 00000000..766806ff --- /dev/null +++ b/src/xorl/sim/timing_ledger.py @@ -0,0 +1,300 @@ +"""Timing ledger construction for prediction reports.""" + +from __future__ import annotations + +import math +from typing import Any + + +try: + from .schemas import BenchmarkBehaviorPrediction, TimingLedger +except ImportError: # pragma: no cover - exercised by direct script execution + from schemas import BenchmarkBehaviorPrediction, TimingLedger + + +_FORWARD_PHASES = ("model_forward", "forward", "fwd") +_LOSS_PHASES = ("loss_compute", "loss", "cross_entropy", "ce_loss") +_BACKWARD_PHASES = ("backward", "model_backward", "bwd") +_OPTIMIZER_PHASES = ("optimizer_step", "optimizer", "optim", "clip_and_step") +_INPUT_PHASES = ("dataloader", "get_batch", "input", "microbatch_to_device", "collator", "tokenize") +_FORWARD_BACKWARD_PHASES = ("forward_backward", "forward_backward_total", "fwd_bwd", "fwd_bwd_total") + + +def _float_dict(value: Any) -> dict[str, float]: + if not isinstance(value, dict): + return {} + output: dict[str, float] = {} + for key, item in value.items(): + try: + numeric = float(item) + except (TypeError, ValueError): + continue + if math.isfinite(numeric): + output[str(key)] = round(numeric, 6) + return output + + +def _float_or_none(value: Any) -> float | None: + try: + numeric = float(value) + except (TypeError, ValueError): + return None + if not math.isfinite(numeric): + return None + return round(numeric, 6) + + +def _source_for_observed(calibration_sources: list[str] | None) -> str: + if not calibration_sources: + return "observed_logs" + if len(calibration_sources) == 1: + return f"observed_logs:{calibration_sources[0]}" + return f"observed_logs:{len(calibration_sources)} sources" + + +def _source_for_benchmark(benchmark_behavior: BenchmarkBehaviorPrediction) -> str: + if benchmark_behavior.matched_label: + return f"benchmark_behavior:{benchmark_behavior.matched_label}" + if benchmark_behavior.source: + return f"benchmark_behavior:{benchmark_behavior.source}" + return "benchmark_behavior" + + +def _phase_value(phase_time_sec: dict[str, float], names: tuple[str, ...]) -> float | None: + lowered = {phase.lower(): value for phase, value in phase_time_sec.items()} + for name in names: + value = lowered.get(name) + if value is not None: + return value + return None + + +def _phase_bucket(phase: str) -> str: + lowered = phase.lower() + if any(part in lowered for part in ("dataloader", "get_batch", "input", "tokenize", "collator")): + return "input" + if any(part in lowered for part in ("optimizer", "optim", "clip_and_step", "lr_scheduler", "zero_grad")): + return "optimizer" + if any( + part in lowered + for part in ( + "sync", + "all_reduce", + "reduce_scatter", + "all_gather", + "nccl", + "deepep", + "dispatch", + "combine", + "communication", + "data_movement", + "a2a", + "fsdp", + ) + ): + return "communication" + if lowered in {"fwd", "bwd", "fwd+bwd"} or any( + part in lowered for part in ("forward", "backward", "loss", "recompute", "moe", "attention") + ): + return "model_compute" + return "other" + + +def _is_composite_phase_for_bottleneck(phase: str, phases: set[str]) -> bool: + lowered = phase.lower() + lowered_phases = {item.lower() for item in phases} + if lowered == "train_step_total": + return True + if lowered in set(_FORWARD_BACKWARD_PHASES): + return bool( + lowered_phases + & { + *_FORWARD_PHASES, + *_LOSS_PHASES, + *_BACKWARD_PHASES, + } + ) + if lowered == "clip_and_step_total": + return bool( + lowered_phases + & { + "clip_gradients", + "optimizer_step", + "optimizer", + "optim", + "lr_scheduler_step", + } + ) + return False + + +def _derive_phase_share( + phase_time_sec: dict[str, float], phase_time_share: dict[str, float] +) -> tuple[dict[str, float], str | None]: + if phase_time_share: + return phase_time_share, None + denominator = phase_time_sec.get("train_step_total") + if denominator is None: + denominator = sum(value for phase, value in phase_time_sec.items() if phase != "train_step_total") + if not denominator: + return {}, None + derived = { + phase: round(value / denominator, 6) for phase, value in phase_time_sec.items() if phase != "train_step_total" + } + return derived, "phase_time_share=derived_from_phase_time_sec" + + +def _derive_phase_time_sec( + phase_time_share: dict[str, float], + step_time_s: float | None, +) -> tuple[dict[str, float], str | None]: + if not phase_time_share or step_time_s is None or step_time_s <= 0: + return {}, None + derived = { + phase: round(share * step_time_s, 6) for phase, share in phase_time_share.items() if phase != "train_step_total" + } + if derived: + derived["train_step_total"] = round(step_time_s, 6) + return derived, "phase_time_sec=derived_from_phase_time_share_and_step_time" + return {}, None + + +def _phase_bottleneck(phase_time_share: dict[str, float]) -> tuple[str, float] | None: + visible = { + phase: share + for phase, share in phase_time_share.items() + if not _is_composite_phase_for_bottleneck(phase, set(phase_time_share)) + } + if not visible: + visible = {phase: share for phase, share in phase_time_share.items() if phase != "train_step_total"} + if not visible: + return None + phase = max(visible, key=visible.get) + return phase, visible[phase] + + +def _forward_backward_time( + *, + phase_time_sec: dict[str, float], + forward_s: float | None, + loss_s: float | None, + backward_s: float | None, +) -> tuple[float | None, str | None]: + direct = _phase_value(phase_time_sec, _FORWARD_BACKWARD_PHASES) + if direct is not None: + return direct, None + if forward_s is None or backward_s is None: + return None, None + total = forward_s + backward_s + if loss_s is not None: + total += loss_s + return round(total, 6), "forward_backward_s=summed_phase_components" + + +def build_timing_ledger( + observed_summary: dict[str, Any] | None, + benchmark_behavior: BenchmarkBehaviorPrediction | None, + *, + calibration_sources: list[str] | None = None, +) -> TimingLedger: + """Build a typed timing ledger from observed logs or matched benchmark behavior.""" + notes: list[str] = [] + source = None + timing_coverage_status = "no_timing_calibration" + phase_time_sec: dict[str, float] = {} + phase_time_share: dict[str, float] = {} + step_time_s: float | None = None + + observed_phase_time_sec = _float_dict((observed_summary or {}).get("phase_time_sec")) + observed_phase_time_share = _float_dict((observed_summary or {}).get("phase_time_share")) + observed_step_time = _float_or_none((observed_summary or {}).get("step_time_s_mean")) + benchmark_phase_time_sec = _float_dict(benchmark_behavior.phase_time_sec if benchmark_behavior else None) + benchmark_phase_time_share = _float_dict(benchmark_behavior.phase_time_share if benchmark_behavior else None) + benchmark_step_time = _float_or_none(benchmark_behavior.step_time_sec if benchmark_behavior else None) + if observed_phase_time_sec: + source = _source_for_observed(calibration_sources) + timing_coverage_status = "observed_phase_timing" + phase_time_sec = observed_phase_time_sec + phase_time_share = observed_phase_time_share + step_time_s = observed_step_time + elif observed_phase_time_share and observed_step_time is not None: + source = _source_for_observed(calibration_sources) + timing_coverage_status = "observed_phase_timing" + phase_time_share = observed_phase_time_share + step_time_s = observed_step_time + phase_time_sec, phase_sec_note = _derive_phase_time_sec(phase_time_share, step_time_s) + if phase_sec_note is not None: + notes.append(phase_sec_note) + elif benchmark_phase_time_sec: + assert benchmark_behavior is not None + source = _source_for_benchmark(benchmark_behavior) + timing_coverage_status = "benchmark_phase_timing" + phase_time_sec = benchmark_phase_time_sec + phase_time_share = benchmark_phase_time_share + step_time_s = benchmark_step_time + elif benchmark_phase_time_share and benchmark_step_time is not None and benchmark_behavior is not None: + source = _source_for_benchmark(benchmark_behavior) + timing_coverage_status = "benchmark_phase_timing" + phase_time_share = benchmark_phase_time_share + step_time_s = benchmark_step_time + phase_time_sec, phase_sec_note = _derive_phase_time_sec(phase_time_share, step_time_s) + if phase_sec_note is not None: + notes.append(phase_sec_note) + else: + if observed_step_time is not None: + source = _source_for_observed(calibration_sources) + timing_coverage_status = "observed_total_step_only" + step_time_s = observed_step_time + elif benchmark_step_time is not None and benchmark_behavior is not None: + source = _source_for_benchmark(benchmark_behavior) + timing_coverage_status = "benchmark_total_step_only" + step_time_s = benchmark_step_time + + if step_time_s is None: + step_time_s = phase_time_sec.get("train_step_total") + share_note = None + phase_time_share, share_note = _derive_phase_share(phase_time_sec, phase_time_share) + if share_note is not None: + notes.append(share_note) + if timing_coverage_status.endswith("_total_step_only"): + notes.append("phase_breakdown_unavailable") + + forward_s = _phase_value(phase_time_sec, _FORWARD_PHASES) + loss_s = _phase_value(phase_time_sec, _LOSS_PHASES) + backward_s = _phase_value(phase_time_sec, _BACKWARD_PHASES) + optimizer_s = _phase_value(phase_time_sec, _OPTIMIZER_PHASES) + input_s = _phase_value(phase_time_sec, _INPUT_PHASES) + forward_backward_s, forward_backward_note = _forward_backward_time( + phase_time_sec=phase_time_sec, + forward_s=forward_s, + loss_s=loss_s, + backward_s=backward_s, + ) + if forward_backward_note is not None: + notes.append(forward_backward_note) + + bottleneck = _phase_bottleneck(phase_time_share) + phase_bottleneck_phase = None + phase_bottleneck_bucket = None + phase_bottleneck_share = None + if bottleneck is not None: + phase_bottleneck_phase, phase_bottleneck_share = bottleneck + phase_bottleneck_bucket = _phase_bucket(phase_bottleneck_phase) + + return TimingLedger( + source=source, + timing_coverage_status=timing_coverage_status, + forward_backward_s=forward_backward_s, + forward_s=forward_s, + loss_s=loss_s, + backward_s=backward_s, + optimizer_s=optimizer_s, + input_s=input_s, + step_time_s=step_time_s, + phase_time_sec=phase_time_sec, + phase_time_share=phase_time_share, + phase_bottleneck_phase=phase_bottleneck_phase, + phase_bottleneck_bucket=phase_bottleneck_bucket, + phase_bottleneck_share=round(phase_bottleneck_share, 6) if phase_bottleneck_share is not None else None, + notes=notes, + ) diff --git a/experiments/local_benchmark/training_sim/tradeoff_ranker.py b/src/xorl/sim/tradeoff_ranker.py similarity index 96% rename from experiments/local_benchmark/training_sim/tradeoff_ranker.py rename to src/xorl/sim/tradeoff_ranker.py index e0807cc7..3b1081a1 100644 --- a/experiments/local_benchmark/training_sim/tradeoff_ranker.py +++ b/src/xorl/sim/tradeoff_ranker.py @@ -9,11 +9,13 @@ try: from .benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from .calibration_packs import resolve_calibration_pack from .config_fingerprint import load_training_config, resolve_topology from .schemas import BenchmarkBehaviorPoint, Topology, TradeoffCandidate, TradeoffReport, to_jsonable from .shape_engine import build_shape_ledger except ImportError: # pragma: no cover - exercised by direct script execution from benchmark_behavior import load_benchmark_behavior_points, predict_benchmark_behavior + from calibration_packs import resolve_calibration_pack from config_fingerprint import load_training_config, resolve_topology from schemas import BenchmarkBehaviorPoint, Topology, TradeoffCandidate, TradeoffReport, to_jsonable from shape_engine import build_shape_ledger @@ -111,7 +113,7 @@ def rank_benchmark_tradeoffs( world_size: int | None = None, local_world_size: int | None = None, ) -> TradeoffReport: - benchmark_path = Path(benchmark_dir) + benchmark_path = resolve_calibration_pack(benchmark_dir) configs = sorted((benchmark_path / "configs").glob("*.yaml")) behavior_points = load_benchmark_behavior_points(benchmark_path) warnings: list[str] = [] @@ -179,7 +181,7 @@ def rank_benchmark_tradeoffs( def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("benchmark_dir", type=Path) + parser.add_argument("benchmark_dir", help="Path or built-in calibration-pack name") parser.add_argument("--world-size", type=int, default=None) parser.add_argument("--local-world-size", type=int, default=None) parser.add_argument("--output", type=Path, default=None) diff --git a/src/xorl/sim/validate.py b/src/xorl/sim/validate.py new file mode 100644 index 00000000..d79684a0 --- /dev/null +++ b/src/xorl/sim/validate.py @@ -0,0 +1,143 @@ +"""Run the portable simulator, calibration-pack, and analytical golden gates.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from .analytical_ledgers import activation_ledger, communication_ledger, flops_ledger +from .benchmark_behavior import load_benchmark_behavior_points +from .calibration_packs import ( + list_calibration_packs, + load_calibration_pack, + validate_calibration_pack, +) +from .config_fingerprint import load_training_config +from .model_metadata import resolve_model_metadata +from .predict import build_report +from .schemas import to_jsonable +from .tradeoff_ranker import rank_benchmark_tradeoffs + + +def _check(name: str, expected: Any, actual: Any, *, tolerance: float | None = None) -> dict[str, Any]: + if tolerance is None: + passed = expected == actual + else: + passed = expected is not None and actual is not None and abs(float(expected) - float(actual)) <= tolerance + return { + "name": name, + "status": "pass" if passed else "fail", + "expected": expected, + "actual": actual, + } + + +def validate_simulator_pack(name: str) -> dict[str, Any]: + pack = load_calibration_pack(name) + manifest = pack.manifest + golden = manifest.get("golden", {}) + balanced_routing = bool(manifest.get("balanced_routing", False)) + sanitation = validate_calibration_pack(pack.path) + points = load_benchmark_behavior_points(pack.path) + raw_config = load_training_config(pack.default_config) + report = build_report( + pack.default_config, + world_size=None, + local_world_size=None, + balanced_routing=balanced_routing, + num_experts=None, + top_k=None, + benchmark_dir=pack.path, + ) + topology = report.fingerprint.topology + metadata = resolve_model_metadata(raw_config) + train = raw_config.get("train", {}) + flops = flops_ledger(metadata, topology) + activations = activation_ledger(metadata, topology, train) + communication = communication_ledger(metadata, topology, train) + tradeoffs = rank_benchmark_tradeoffs(pack.path) + + default_tps = report.benchmark_behavior.tokens_per_sec if report.benchmark_behavior is not None else None + best_raw_tps = tradeoffs.best_raw.score_tokens_per_sec if tradeoffs.best_raw is not None else None + best_promotable_tps = ( + tradeoffs.best_promotable.score_tokens_per_sec if tradeoffs.best_promotable is not None else None + ) + checks = [ + _check("pack_sanitation", "pass", sanitation["status"]), + _check("behavior_point_count", golden.get("behavior_point_count"), len(points)), + _check("world_size", golden.get("world_size"), topology.world_size), + _check("global_batch_size", golden.get("global_batch_size"), topology.global_batch_size), + _check("default_tokens_per_sec", golden.get("default_tokens_per_sec"), default_tps, tolerance=0.001), + _check("best_raw_tokens_per_sec", golden.get("best_raw_tokens_per_sec"), best_raw_tps, tolerance=0.001), + _check( + "best_promotable_tokens_per_sec", + golden.get("best_promotable_tokens_per_sec"), + best_promotable_tps, + tolerance=0.001 if golden.get("best_promotable_tokens_per_sec") is not None else None, + ), + _check( + "analytic_peak_floor_gb", + golden.get("analytic_peak_floor_gb"), + report.memory.analytic_peak_floor_gb, + tolerance=0.001, + ), + _check("flops_ledger", "exact_analytic", flops.get("status")), + _check("activation_ledger", "exact_analytic_lower_bound", activations.get("status")), + _check("communication_ledger", "exact_analytic_bytes", communication.get("status")), + _check("simulator_support", True, report.support.support_status.startswith("supported_")), + ] + return { + "name": pack.name, + "status": "pass" if all(check["status"] == "pass" for check in checks) else "fail", + "manifest": manifest, + "checks": checks, + "sanitation": sanitation, + "prediction_report": to_jsonable(report), + "tradeoff_report": to_jsonable(tradeoffs), + "analytical": { + "flops": flops, + "activations": activations, + "communication": communication, + }, + } + + +def validate_simulator(pack_names: list[str] | None = None) -> dict[str, Any]: + names = pack_names or list_calibration_packs() + packs = [validate_simulator_pack(name) for name in names] + check_count = sum(len(pack["checks"]) + len(pack["sanitation"]["checks"]) for pack in packs) + failed_check_count = sum( + check["status"] != "pass" for pack in packs for check in [*pack["checks"], *pack["sanitation"]["checks"]] + ) + return { + "schema_version": 1, + "status": "pass" if packs and all(pack["status"] == "pass" for pack in packs) else "fail", + "pack_count": len(packs), + "check_count": check_count, + "failed_check_count": failed_check_count, + "packs": packs, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--pack", action="append", default=None, help="Validate only this built-in pack; repeatable") + parser.add_argument("--output", type=Path, default=None) + parser.add_argument("--no-fail-on-error", action="store_true") + args = parser.parse_args() + + payload = validate_simulator(args.pack) + rendered = json.dumps(payload, indent=2, sort_keys=True) + "\n" + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(rendered, encoding="utf-8") + else: + print(rendered, end="") + if payload["status"] != "pass" and not args.no_fail_on_error: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tests/experiments/test_training_sim.py b/tests/experiments/test_training_sim.py index dcaf145a..87c41f63 100644 --- a/tests/experiments/test_training_sim.py +++ b/tests/experiments/test_training_sim.py @@ -2,13 +2,24 @@ import yaml -from experiments.local_benchmark.training_sim.benchmark_behavior import load_benchmark_behavior_points -from experiments.local_benchmark.training_sim.calibration_evaluator import evaluate_calibration -from experiments.local_benchmark.training_sim.collect_calibration import parse_log_text, summarize_observed_run -from experiments.local_benchmark.training_sim.config_fingerprint import build_fingerprint, resolve_topology -from experiments.local_benchmark.training_sim.model_metadata import resolve_model_metadata -from experiments.local_benchmark.training_sim.scenario_planner import plan_scenario -from experiments.local_benchmark.training_sim.shape_engine import balanced_counts, build_shape_ledger +from xorl.sim.analytical_ledgers import activation_ledger, communication_ledger, flops_ledger +from xorl.sim.benchmark_behavior import load_benchmark_behavior_points +from xorl.sim.calibration_evaluator import evaluate_calibration +from xorl.sim.calibration_packs import ( + list_calibration_packs, + load_calibration_pack, + validate_calibration_pack, +) +from xorl.sim.collect_calibration import parse_log_text, summarize_observed_run +from xorl.sim.config_fingerprint import build_fingerprint, load_training_config, resolve_topology +from xorl.sim.feasibility_evaluator import evaluate_feasibility +from xorl.sim.kernel_variants import compare_kernel_variants, rank_kernel_variants +from xorl.sim.model_metadata import resolve_model_metadata +from xorl.sim.predict import build_report +from xorl.sim.scenario_planner import plan_scenario +from xorl.sim.shape_engine import balanced_counts, build_shape_ledger +from xorl.sim.tradeoff_ranker import rank_benchmark_tradeoffs +from xorl.sim.validate import validate_simulator def test_balanced_counts_round_robin_distribution() -> None: @@ -227,7 +238,7 @@ def test_scenario_planner_keeps_observed_fit_feasible_when_safety_margin_is_tigh assert report.best_raw.score_tokens_per_sec == 1_100.0 assert report.best_raw.feasibility_status == "feasible_calibrated_peak_high_pressure" assert report.best_raw.memory_headroom_gb == -0.5 - assert report.best_raw.recommendation == "correctness_gate_required" + assert report.best_raw.recommendation == "remeasure_for_stability" def test_build_fingerprint_reads_config_file(tmp_path: Path) -> None: @@ -370,7 +381,10 @@ def _write_q235_config_fixture(config_path: Path) -> None: "model": { "model_path": "Qwen/Qwen3-235B-A22B", "ep_dispatch": "deepep", + "moe_implementation": "quack", "deepep_buffer_size_gb": 2.0, + "deepep_num_sms": 24, + "deepep_async_combine": False, }, "data": { "sample_packing_sequence_len": 4096, @@ -388,8 +402,12 @@ def _write_q235_config_fixture(config_path: Path) -> None: "gradient_accumulation_steps": 1, "optimizer": "muon", "optimizer_dtype": "bf16", - "muon_momentum": 0.0, + "muon_momentum": 0.95, "enable_mixed_precision": True, + "skip_param_upcast": True, + "fsdp_reduce_dtype": "fp32", + "ce_mode": "quack_linear", + "gradient_checkpointing_method": "recompute_before_dispatch", }, } config_path.write_text(yaml.safe_dump(config), encoding="utf-8") @@ -417,7 +435,7 @@ def test_qwen235_scenario_planner_uses_markdown_calibration_for_ga_tradeoff(tmp_ assert report.best_raw.prediction_confidence == "calibrated" assert report.best_raw.score_tokens_per_sec == 8_400.0 assert report.best_raw.behavior.tokens_per_sec_per_gpu == 262.5 - assert report.best_raw.analytic_peak_floor_gb == 29.363 + assert report.best_raw.analytic_peak_floor_gb == 56.812 assert report.best_raw.estimated_peak_mem_gb == 68.3 assert report.best_raw.memory_basis == "calibrated_peak" assert report.best_raw.feasibility_status == "feasible_calibrated_peak_high_pressure" @@ -444,13 +462,13 @@ def test_qwen235_scenario_planner_extrapolates_ga_asymptote_from_step_time_fit(t assert report.best_raw is not None assert report.best_raw.label == "mbs1-gb256-ep8-efsdp4-tp1-pp1-u1-r1:extrapolated" assert report.best_raw.prediction_confidence == "extrapolated_step_time_fit" - assert report.best_raw.score_tokens_per_sec == 9646.513 - assert report.best_raw.behavior.step_time_sec == 108.7 + assert report.best_raw.score_tokens_per_sec == 10_200.0 + assert report.best_raw.behavior.step_time_sec == 102.801569 assert report.best_raw.estimated_peak_mem_gb == 68.3 - assert report.best_raw.memory_basis == "extrapolated_peak" - assert report.best_raw.feasibility_status == "feasible_extrapolated_peak_high_pressure" + assert report.best_raw.memory_basis == "calibrated_overhead_peak" + assert report.best_raw.feasibility_status == "feasible_calibrated_overhead_peak_high_pressure" assert report.best_raw.calibration_scope == "outside_measured_envelope" - assert report.best_raw.score_risk_adjusted_tokens_per_sec == 4796.729 + assert report.best_raw.score_risk_adjusted_tokens_per_sec == 4666.194 assert report.best_raw.recommendation == "remeasure_before_ranking" assert "outside_measured_envelope" in report.best_raw.risk_flags assert "requires_remeasurement" in report.best_raw.risk_flags @@ -463,7 +481,7 @@ def test_qwen235_scenario_planner_extrapolates_ga_asymptote_from_step_time_fit(t by_label = {candidate.label: candidate for candidate in report.candidates} ga4 = by_label["mbs1-gb128-ep8-efsdp4-tp1-pp1-u1-r1:extrapolated"] assert ga4.prediction_confidence == "extrapolated_step_time_fit" - assert ga4.score_tokens_per_sec == 9181.926 + assert ga4.score_tokens_per_sec == 9_520.0 def test_qwen235_calibration_evaluator_reports_leave_one_out_ga_error(tmp_path: Path) -> None: @@ -484,15 +502,15 @@ def test_qwen235_calibration_evaluator_reports_leave_one_out_ga_error(tmp_path: assert report.evaluated_count == 2 assert report.skipped_count == 0 assert report.prediction_status_counts == {"extrapolated": 2} - assert report.mean_absolute_percentage_error == 21.288 + assert report.mean_absolute_percentage_error == 19.16 by_label = {holdout.label: holdout for holdout in report.holdouts} ga1 = by_label["q235_markdown:n4_ep8_bd_pk4096"] ga2 = by_label["q235_markdown:n4_ep8_bd_pk4096_ga2"] assert ga1.topology_label == "mbs1-gb32-ep8-efsdp4-tp1-pp1-u1-r1" - assert ga1.predicted_tokens_per_sec == 8_400.0 - assert ga1.absolute_percentage_error == 23.529 - assert ga2.predicted_tokens_per_sec == 6_800.0 - assert ga2.absolute_percentage_error == 19.048 + assert ga1.predicted_tokens_per_sec == 7_560.0 + assert ga1.absolute_percentage_error == 11.176 + assert ga2.predicted_tokens_per_sec == 6_120.0 + assert ga2.absolute_percentage_error == 27.143 def test_qwen235_scenario_planner_does_not_exact_match_observed_row_to_tp_what_if(tmp_path: Path) -> None: @@ -520,7 +538,7 @@ def test_qwen235_scenario_planner_does_not_exact_match_observed_row_to_tp_what_i assert tp2.prediction_confidence == "extrapolated" assert tp2.behavior.matched_label == "q235_markdown:n4_ep8_bd_pk4096_ga2" assert "TP extrapolation uses conservative communication penalty" in tp2.behavior.warnings - assert tp2.score_tokens_per_sec == 7_560.0 + assert tp2.score_tokens_per_sec == 6_804.0 def test_qwen235_scenario_planner_auto_sweeps_parallelism_strategy_space(tmp_path: Path) -> None: @@ -622,3 +640,116 @@ def test_qwen235_scenario_planner_marks_matching_oom_pack_infeasible(tmp_path: P assert candidate.score_tokens_per_sec is None assert candidate.calibration_scope == "exact_calibrated" assert "observed_oom_boundary:q235_markdown:n4_ep8_bd_pk16k" in candidate.risk_flags + + +def test_builtin_calibration_packs_are_sanitized_and_versioned() -> None: + assert list_calibration_packs() == ["qwen3_235b_a22b", "qwen3_5_397b_a17b", "qwen3_6_35b_a3b"] + for name in list_calibration_packs(): + pack = load_calibration_pack(name) + validation = validate_calibration_pack(pack.path) + assert pack.manifest["schema_version"] == 1 + assert pack.default_config.is_file() + assert validation["status"] == "pass" + + +def test_builtin_qwen35_pack_preserves_raw_and_promotable_winners() -> None: + pack = load_calibration_pack("qwen3_5_397b_a17b") + points = load_benchmark_behavior_points(pack.path) + report = rank_benchmark_tradeoffs(pack.path) + + assert len(points) == 6 + assert report.best_raw is not None + assert report.best_raw.score_tokens_per_sec == 59_217.0 + assert report.best_raw.promotable is False + assert report.best_promotable is not None + assert report.best_promotable.score_tokens_per_sec == 59_188.0 + assert report.best_promotable.promotable is True + + +def test_builtin_qwen36_pack_matches_default_config_but_remains_ungated() -> None: + pack = load_calibration_pack("qwen3_6_35b_a3b") + report = build_report( + pack.default_config, + world_size=None, + local_world_size=None, + balanced_routing=True, + num_experts=None, + top_k=None, + benchmark_dir=pack.path, + ) + + assert report.benchmark_behavior is not None + assert report.benchmark_behavior.matched_label == "readme_reference_mbs8" + assert report.benchmark_behavior.tokens_per_sec == 261_000.0 + assert report.benchmark_behavior.correctness_status == "raw_speed_not_promoted_without_matching_k3_pass" + assert report.support.support_status == "supported_local_non_pp" + assert report.timing.timing_coverage_status == "benchmark_total_step_only" + + +def test_builtin_qwen235_pack_replays_fit_and_oom_boundaries() -> None: + pack = load_calibration_pack("qwen3_235b_a22b") + report = evaluate_feasibility(pack.default_config, benchmark_dir=pack.path) + + assert report.status == "ok" + assert report.evaluated_count == 3 + assert report.accuracy == 1.0 + assert report.fit_recall == 1.0 + assert report.oom_recall == 1.0 + assert {holdout.actual_outcome for holdout in report.holdouts} == {"fit", "oom"} + + +def test_portable_analytical_core_covers_flops_activations_and_communication() -> None: + pack = load_calibration_pack("qwen3_5_397b_a17b") + raw_config = load_training_config(pack.default_config) + topology = resolve_topology(raw_config) + metadata = resolve_model_metadata(raw_config) + + flops = flops_ledger(metadata, topology) + activations = activation_ledger(metadata, topology, raw_config["train"]) + communication = communication_ledger(metadata, topology, raw_config["train"]) + + assert metadata.full_attention_interval == 4 + assert metadata.linear_num_value_heads == 64 + assert flops["status"] == "exact_analytic" + assert flops["total_flops"] > 0 + assert activations["status"] == "exact_analytic_lower_bound" + assert activations["analytic_activation_lower_bound_gb"] > 0 + assert communication["status"] == "exact_analytic_bytes" + assert communication["total_per_rank_gb"] > 0 + + +def test_kernel_variant_ranking_requires_a_correctness_gate() -> None: + rows = [ + { + "family": "attention", + "variant": "fast-ungated", + "workload": "qwen35-seq4096", + "latency_ms": 8.0, + "correctness_status": "not_promoted", + }, + { + "family": "attention", + "variant": "validated", + "workload": "qwen35-seq4096", + "latency_ms": 10.0, + "correctness_status": "pass", + }, + ] + + report = rank_kernel_variants(rows) + comparison = compare_kernel_variants(rows[1], rows[0]) + + assert report["status"] == "ok" + assert report["best"]["variant"] == "validated" + assert report["measurements"][0]["variant"] == "fast-ungated" + assert comparison["speedup"] == 1.25 + assert comparison["candidate_promotable"] is False + + +def test_consolidated_validator_covers_all_builtin_packs() -> None: + report = validate_simulator() + + assert report["status"] == "pass" + assert report["pack_count"] == 3 + assert report["check_count"] >= 200 + assert report["failed_check_count"] == 0 From 0daf6284b480ab0b64fd84e3232560db17d42921 Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Jul 2026 03:10:26 +0000 Subject: [PATCH 3/4] Fix simulator spelling checks --- src/xorl/sim/benchmark_behavior.py | 2 +- src/xorl/sim/calibration_evaluator.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/xorl/sim/benchmark_behavior.py b/src/xorl/sim/benchmark_behavior.py index 8ff0173d..f9af483f 100644 --- a/src/xorl/sim/benchmark_behavior.py +++ b/src/xorl/sim/benchmark_behavior.py @@ -563,7 +563,7 @@ def _readme_topology_defaults(readme_text: str) -> dict[str, int | float | bool field_patterns = { "data_parallel_replicate_size": ( r"\bdata_parallel_replicate_size[=: ]+(?P\d+)\b", - r"\bdp[_-]?rep(?:licate)?[=_ ](?P\d+)\b", + r"\bdp[_-]?rep(?:lica(?:te)?)?[=_ ](?P\d+)\b", ), "data_parallel_shard_size": ( r"\bdata_parallel_shard_size[=: ]+(?P\d+)\b", diff --git a/src/xorl/sim/calibration_evaluator.py b/src/xorl/sim/calibration_evaluator.py index b3c48682..3c18fb8f 100644 --- a/src/xorl/sim/calibration_evaluator.py +++ b/src/xorl/sim/calibration_evaluator.py @@ -856,7 +856,7 @@ def _calibration_validation_gap_portfolio( gap_status="throughput_error_needs_attribution", priority=120, required_measurement="replay_high_error_holdouts_with_component_timing", - reason="at least one scored holdout exceeds the throughput MAPE target", + reason="at least one scored holdout exceeds the throughput mean absolute percentage error target", affected_holdouts=affected, blocker_names=["max_throughput_mape_exceeds_8_percent"], ) From 435178c0f89c6cca46cf9886da56c21e7c905bda Mon Sep 17 00:00:00 2001 From: Ashwinee Panda Date: Fri, 10 Jul 2026 03:12:16 +0000 Subject: [PATCH 4/4] Restrict simulator metadata paths --- src/xorl/sim/calibration_packs.py | 10 +++++----- src/xorl/sim/model_metadata.py | 19 ++++++++++++++++--- tests/experiments/test_training_sim.py | 23 +++++++++++++++++++++++ 3 files changed, 44 insertions(+), 8 deletions(-) diff --git a/src/xorl/sim/calibration_packs.py b/src/xorl/sim/calibration_packs.py index d2309175..112f377c 100644 --- a/src/xorl/sim/calibration_packs.py +++ b/src/xorl/sim/calibration_packs.py @@ -55,12 +55,12 @@ def resolve_calibration_pack(value: str | Path) -> Path: raw = str(value) name = raw.removeprefix("builtin:") - builtin = _PACK_ROOT / name - if raw.startswith("builtin:") or (not Path(raw).exists() and (builtin / "manifest.json").is_file()): - if not (builtin / "manifest.json").is_file(): - available = ", ".join(list_calibration_packs()) or "none" + available_names = list_calibration_packs() + if raw.startswith("builtin:") or (not Path(raw).exists() and name in available_names): + if name not in available_names: + available = ", ".join(available_names) or "none" raise ValueError(f"unknown built-in calibration pack {name!r}; available: {available}") - return builtin + return _PACK_ROOT / name return Path(raw) diff --git a/src/xorl/sim/model_metadata.py b/src/xorl/sim/model_metadata.py index 4dedd41d..456f2f27 100644 --- a/src/xorl/sim/model_metadata.py +++ b/src/xorl/sim/model_metadata.py @@ -4,6 +4,7 @@ import json import os +import re from pathlib import Path from typing import Any @@ -233,13 +234,25 @@ def default_hf_cache_roots() -> list[Path]: def _candidate_config_paths(model_ref: str, hf_cache_roots: list[Path]) -> list[Path]: ref_path = Path(model_ref).expanduser() + allowed_roots = [root.expanduser().resolve() for root in hf_cache_roots] + + def is_allowed(path: Path) -> bool: + resolved = path.resolve() + return any(resolved.is_relative_to(root) for root in allowed_roots) + candidates: list[Path] = [] - if ref_path.is_file(): + if ref_path.name == "config.json" and ref_path.is_file() and is_allowed(ref_path): candidates.append(ref_path) - elif ref_path.is_dir(): + elif ref_path.is_dir() and is_allowed(ref_path): candidates.append(ref_path / "config.json") - if "/" in model_ref and not ref_path.exists(): + repo_parts = model_ref.split("/") + is_hf_repo_id = ( + len(repo_parts) == 2 + and all(part not in {"", ".", ".."} for part in repo_parts) + and all(re.fullmatch(r"[A-Za-z0-9_.-]+", part) is not None for part in repo_parts) + ) + if is_hf_repo_id and not ref_path.exists(): cache_name = "models--" + model_ref.replace("/", "--") for root in hf_cache_roots: snapshots_dir = root / cache_name / "snapshots" diff --git a/tests/experiments/test_training_sim.py b/tests/experiments/test_training_sim.py index 87c41f63..46e39f74 100644 --- a/tests/experiments/test_training_sim.py +++ b/tests/experiments/test_training_sim.py @@ -1,5 +1,6 @@ from pathlib import Path +import pytest import yaml from xorl.sim.analytical_ledgers import activation_ledger, communication_ledger, flops_ledger @@ -8,6 +9,7 @@ from xorl.sim.calibration_packs import ( list_calibration_packs, load_calibration_pack, + resolve_calibration_pack, validate_calibration_pack, ) from xorl.sim.collect_calibration import parse_log_text, summarize_observed_run @@ -652,6 +654,27 @@ def test_builtin_calibration_packs_are_sanitized_and_versioned() -> None: assert validation["status"] == "pass" +def test_builtin_pack_prefix_rejects_path_traversal() -> None: + with pytest.raises(ValueError, match="unknown built-in calibration pack"): + resolve_calibration_pack("builtin:../qwen3_6_35b_a3b") + + +def test_model_metadata_restricts_local_config_reads_to_approved_roots(tmp_path: Path) -> None: + allowed_root = tmp_path / "allowed" + blocked_config = tmp_path / "blocked" / "config.json" + allowed_root.mkdir() + blocked_config.parent.mkdir() + blocked_config.write_text('{"num_experts": 999}', encoding="utf-8") + + metadata = resolve_model_metadata( + {"model": {"model_path": str(blocked_config)}}, + hf_cache_roots=[allowed_root], + ) + + assert metadata.config_path is None + assert metadata.num_experts is None + + def test_builtin_qwen35_pack_preserves_raw_and_promotable_winners() -> None: pack = load_calibration_pack("qwen3_5_397b_a17b") points = load_benchmark_behavior_points(pack.path)