diff --git a/lib/python/examples/_metadata/baselines.yaml b/lib/python/examples/_metadata/baselines.yaml index 8ee6e05df..435405455 100644 --- a/lib/python/examples/_metadata/baselines.yaml +++ b/lib/python/examples/_metadata/baselines.yaml @@ -331,8 +331,17 @@ baselines: # latency variance is much higher than the n=10 smoke default. stalenessPolicy: exact inc_model_version_per_data_id: true - max_iterations_per_data_id: 15 + # Iteration-per-data_id cap left UNBOUNDED (None) for the sync baselines: + # the variance-plateau force-commit is a fluxtune optimization (charter + # Opt-2) that did not exist in vanilla FwdLLM, so fwdllm/fwdllm_plus run + # the original unbounded variance gate (grind until variance passes). var_threshold: 0.3 + # Opt-1 comm-redundancy fix (charter §5c): learning-neutral, enabled on + # ALL baselines so E4 bytes are fair. fwdllm re-sent the identical model + # to the same K trainers every iteration (~90% redundant, ~22 GB) because + # the return-driven staleness guard never marked actively-training clients + # current. Set false to revert (byte-on-wire identical off). + suppress_redundant_weights: true reselect_each_iteration: false trainer: hyperparameters: @@ -378,8 +387,14 @@ baselines: # into a model_version-53 cycle) happen in the first place. stalenessPolicy: round_data_id inc_model_version_per_data_id: true - max_iterations_per_data_id: 15 + # Unbounded (None) — the iteration cap is a fluxtune optimization + # (charter Opt-2), not part of vanilla FwdLLM_Plus; sync baselines run + # the original unbounded variance gate. See fwdllm block above. var_threshold: 0.3 + # Opt-1 comm-redundancy fix (charter §5c), learning-neutral. Already near + # clean here (per-iteration reselect keeps the guard current, ~3.5%), but + # enabled for consistency so the fix covers all baselines. Set false to revert. + suppress_redundant_weights: true reselect_each_iteration: true trainer: hyperparameters: @@ -454,9 +469,30 @@ baselines: # has no fwdllm/agnews entry. Starting point; retune against # fluxtune's live smoke test. learning_rate: 0.075 + # Aggregation rate (per-update weight). DEFAULT for fluxtune is now C3 + # `grad_aware` (Opt-3, charter §5c) -- fluxtune's OWN aggregation. `type: new` + # is the borrowed FeLiX scalar staleness×utility (magnitude only, charter N2) + # and is NOT a fluxtune contribution; set it only for the FeLiX-baseline + # ablation arm. grad_aware weights by DIRECTION/reliability, bounded ≤ base so + # the effective LR never inflates (no LR re-tune): + # base: new # keep staleness×utility as the base (mobiperf async); + # # `neutral` = 1.0 (pure gradient-aware, no staleness) + # align_gate: true # S2 (primary): down-weight updates whose cos vs the + # # running aggregate < align_floor -> 0 at cos=-1 + # # (fixes H1: averaging anti-aligned JVP estimates) + # align_floor: 0.0 # gate below this cosine (0 = gate only anti-aligned) + # inverse_var: false # S1 (optional stronger C3): *min(1, var_ref/(var_i+eps)) + # var_ref: 0.3 # reliability reference (defaults to var_threshold) + # `grad_aware_gated_total` + `agg_rate_type` emitted on agg_round. + # ABLATION: to revert this arm to FeLiX aggregation set `type: new`. agg_rate_conf: - type: new - scale: 0.4 + type: grad_aware + base: new + align_gate: true + align_floor: 0.0 + inverse_var: false + var_ref: 0.3 + scale: 0.4 # base=new (staleness×utility) params a_exp: 0.25 b_exp: 0.1 hyperparameters: @@ -473,6 +509,42 @@ baselines: stalenessPolicy: none inc_model_version_per_data_id: true var_threshold: 0.3 + # --- fluxtune optimizations (charter §5c; flag gated) --- + # Opt-1: suppress byte-identical intra-databin weight re-sends (shared + # sync+async fix, enabled on all baselines). Within a data-bin the + # model_version is constant and the full WEIGHTS+GRAD_POOL payload is + # identical across iterations; a trainer that already received it this + # cycle is downgraded to the tiny VAR=bad "keep training" message (it + # caches its weights). Learning-neutral; here it was 71% redundant / ~68 GB. + # Set false to revert (byte-on-wire identical off). + suppress_redundant_weights: true + # Opt-2 — variance-plateau force-commit (charter §5c/§5e; DEFAULT ON for + # fluxtune). The plain var<=threshold gate is a FIXED ABSOLUTE bar: as the + # model converges, gradients shrink but the perturbation-noise floor doesn't, + # so a bin needs ever-more JVP samples to average down to 0.30 (measured + # 15->34 iters/bin early->late at α=1) -> 95.7% of aggregation iterations + # commit nothing. This policy adds TWO commit triggers on top of var<=thr: + # * plateau (PRIMARY, helps early): commit once the per-bin variance curve + # FLATTENS -- relative var drop over the last `var_plateau_patience` (N) + # iters < `var_plateau_rel_delta` (ε) while var is still above threshold. + # Scale-invariant, so it commits the DENOISED estimate at whatever iter it + # flattens (early bins ~12-17). This is the "more denoising buys nothing" cut. + # * cap (LATE-STAGE BACKSTOP): `max_iterations_per_data_id`. Late in training + # the plateau onset drifts out (~33 iters) -- the model is already stable so + # waiting that long is wasteful; the cap force-commits at 20 first. Sits + # ABOVE the early plateau fire (so plateau stays primary) and BELOW the late + # grind (so it caps the tail). Whichever trigger fires first wins; the + # `commit_reason` field on agg_round telemetry records which (natural/cap/plateau). + # Learning is compared against the earlier var<=0.3-only fluxtune run (no A/B). + # TO DISABLE (revert to var<=thr only, byte-identical): var_stopping_policy: off. + # TO TUNE: lower ε -> plateau fires later/less; raise ε -> earlier/more (ε=0.25 + # fires ~94% at iter ~11). Lower the cap -> more late-stage forcing (and it + # engages on shorter bins, e.g. the smoke). Re-characterize at the more-IID + # ablations (α∈{10,100}) where the floor drops and the grind may soften. + var_stopping_policy: plateau + var_plateau_patience: 3 # N: look-back window for the flatten test + var_plateau_rel_delta: 0.15 # ε: commit when rel var-drop over last N < ε + max_iterations_per_data_id: 20 # late-stage force-commit ceiling trainer: hyperparameters: select_perturbation_using_jvp: true diff --git a/lib/python/examples/async_cifar10/PARITY.md b/lib/python/examples/async_cifar10/PARITY.md index bb07ce2a7..cb86f51e9 100644 --- a/lib/python/examples/async_cifar10/PARITY.md +++ b/lib/python/examples/async_cifar10/PARITY.md @@ -44,6 +44,13 @@ Last green (Jun 24): selector+mode+sim+parity **351 pass / 7 skip** (incl. `Test A fix that could perturb another baseline → serialize (one baseline per run round). 4. **Shortest run that exhibits the issue** (table below). Reserve long runs for C1/C2. 5. **Crisp comments (≤1 sentence); context-free names** (Naming discipline below). +6. **Telemetry-FIRST; a cluster run is the LAST resort, never the debugger.** Before launching anything to + validate/refute a hypothesis, name the exact stored field/log-line that would confirm it and go read the + banked telemetry — most roots are already visible there. Ship every new mechanism WITH its telemetry + + plot + pytest in the same change (over-instrument: cheap to log, expensive to re-run for), so the next + root is catchable from disk. Launch ONLY to observe an emergent quantity no stored telemetry can yield + (fresh convergence / concurrency-after-a-change / a longer trajectory), then the shortest length that + shows it. (fwdllm restates this as its principle #11.) ### Run-length budget (state min duration up front, keyed here; never default to 3–4h) | validating | min run | why | diff --git a/lib/python/examples/async_cifar10/real-sim_parity_checker_plan.md b/lib/python/examples/async_cifar10/real-sim_parity_checker_plan.md index 2290af0fa..4412ddf08 100644 --- a/lib/python/examples/async_cifar10/real-sim_parity_checker_plan.md +++ b/lib/python/examples/async_cifar10/real-sim_parity_checker_plan.md @@ -338,8 +338,10 @@ python -m scripts.parity.cli \ (real vs sim), overlap-factor bars, GPU-vs-budget, convergence-by-round. - Exit code per §3 verdict rule. -`compare_overnight.sh` is rewritten to call this one command per (baseline) pair -and to drop its separate `compare_parity.py` invocation. +Post-run parity is now invoked via `debug_run.sh --after parity` (the old +standalone `compare_overnight.sh` was deleted — its per-baseline sim-vs-real +parity + cross-baseline plots are absorbed into `--after parity,plot`, on the +maintained `scripts.parity.cli` engine). ### 4.2b Multi-baseline invocation (test parity per selector as we converge) @@ -433,7 +435,8 @@ test import path. The `analyze_*.py` scripts are **kept** but demoted to 4. **Build `report.py` + `cli.py`** single output; fix the convergence self-compare bug (C3). Add `--batch` multi-baseline mode (§4.2b). 5. **Add per-selector test layer** (`tests/mode/test_selector_invariants.py`, §4.2c). -6. **Rewrite `compare_overnight.sh`** to the single command; demote `analyze_*` +6. **`compare_overnight.sh` deleted** — folded into `debug_run.sh --after parity,plot` + (shared `examples/scripts/expt_runner.sh` dispatcher); demote `analyze_*` behind `--diagnostics`. 7. **Delete** the superseded files (§4.4); update any docs/HANDOFF references. 8. **Run** the new checker (`--batch`) on felix + refl pairs; paste the FAILing diff --git a/lib/python/examples/async_cifar10/scripts/compare_overnight.sh b/lib/python/examples/async_cifar10/scripts/compare_overnight.sh deleted file mode 100755 index 48a56dda3..000000000 --- a/lib/python/examples/async_cifar10/scripts/compare_overnight.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/bin/bash -# Post-run comparison for the overnight n300 runs. Two views: -# 1. PER-BASELINE sim-vs-real tracking (the only apples-to-apples pairing): -# parity_check.py for each baseline's sim run dir vs its real run dir. -# 2. CROSS-BASELINE, separately for sim and for real (one plot set each): -# analyze_run.py --compare-streaming over the 4 runs of a mode. -# -# compare_overnight.sh # auto-discovers latest run dir per (baseline,mode) -# Output: /tmp/overnight_compare/{parity_.txt, parity_.json, sim_cross/, real_cross/} -set -u - -# --- robust conda activation (handles non-default install locations) --- -ENVNAME="${FLAME_CONDA_ENV:-dg_flame}" -CB="" -if command -v conda >/dev/null 2>&1; then - CB="$(conda info --base 2>/dev/null)" -elif [ -n "${CONDA_EXE:-}" ]; then - CB="$(dirname "$(dirname "$CONDA_EXE")")" -fi -if [ -z "$CB" ] || [ ! -f "$CB/etc/profile.d/conda.sh" ]; then - for c in "$HOME/miniconda3" "/coc/scratch/${USER%??}/miniconda3" \ - "/coc/scratch/$USER/miniconda3" "$HOME/anaconda3" /opt/conda; do - [ -f "$c/etc/profile.d/conda.sh" ] && CB="$c" && break - done -fi -if [ -z "$CB" ] || [ ! -f "$CB/etc/profile.d/conda.sh" ]; then - echo "ERROR: conda not found. Activate '$ENVNAME' yourself or set CONDA_EXE." >&2; exit 1 -fi -source "$CB/etc/profile.d/conda.sh" -conda activate "$ENVNAME" || { echo "ERROR: 'conda activate $ENVNAME' failed" >&2; exit 1; } - -# paths derived from this script's location (portable across nodes) -EX="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" -cd "$EX" || exit 1 -ROOT="$(cd "$EX/../../../.." && pwd)" -OUT=/tmp/overnight_compare; mkdir -p "$OUT" -BASELINES="felix oort refl feddance" - -# latest run dir for a (baseline, mode) — names end in _stream_ -rundir() { ls -dt experiments/run_*_${1}_n300_*_stream_${2}* 2>/dev/null | head -1; } - -echo "### 1. PER-BASELINE sim-vs-real parity ###" -for b in $BASELINES; do - dr=$(rundir "$b" real); ds=$(rundir "$b" sim) - if [ -z "$dr" ] || [ -z "$ds" ]; then - echo " $b: missing run dir (real='$dr' sim='$ds')"; continue - fi - echo " $b: real=$(basename "$dr") sim=$(basename "$ds")" - python scripts/parity_check.py \ - --real "$dr" --sim "$ds" \ - --json-out "$OUT/parity_${b}.json" \ - > "$OUT/parity_${b}.txt" 2>&1 - grep -E "\[OK\]|\[!!\]|\[XX\]|CHECKS" "$OUT/parity_${b}.txt" | sed 's/^/ /' -done - -echo "### 2. CROSS-BASELINE (sim plots, then real plots) ###" -for mode in sim real; do - dirs=(); labels=() - for b in $BASELINES; do - d=$(rundir "$b" "$mode"); [ -z "$d" ] && continue - dirs+=("$d/telemetry"); labels+=("$b") - done - if [ ${#dirs[@]} -ge 2 ]; then - echo " $mode: ${labels[*]}" - python "$ROOT"/scripts/analysis/analyze_run.py \ - --compare-streaming "${dirs[@]}" --labels "${labels[@]}" \ - --out "$OUT/${mode}_cross" > "$OUT/${mode}_cross.log" 2>&1 - echo " -> $OUT/${mode}_cross/ ($(find "$OUT/${mode}_cross" -name '*.pdf' 2>/dev/null | wc -l) plots)" - else - echo " $mode: <2 run dirs, skipping cross-baseline" - fi -done -echo "DONE -> $OUT" diff --git a/lib/python/examples/async_cifar10/scripts/debug_run.sh b/lib/python/examples/async_cifar10/scripts/debug_run.sh index fe24b6326..e5d20a6dd 100755 --- a/lib/python/examples/async_cifar10/scripts/debug_run.sh +++ b/lib/python/examples/async_cifar10/scripts/debug_run.sh @@ -24,29 +24,17 @@ # time for sync baselines like Refl). set -u -# --- robust conda activation --- -ENVNAME="${FLAME_CONDA_ENV:-dg_flame}" -CB="" -if command -v conda >/dev/null 2>&1; then - CB="$(conda info --base 2>/dev/null)" -elif [ -n "${CONDA_EXE:-}" ]; then - CB="$(dirname "$(dirname "$CONDA_EXE")")" -fi -if [ -z "$CB" ] || [ ! -f "$CB/etc/profile.d/conda.sh" ]; then - for c in "$HOME/miniconda3" "/coc/scratch/${USER%??}/miniconda3" \ - "/coc/scratch/$USER/miniconda3" "$HOME/anaconda3" /opt/conda; do - [ -f "$c/etc/profile.d/conda.sh" ] && CB="$c" && break - done -fi -if [ -z "$CB" ] || [ ! -f "$CB/etc/profile.d/conda.sh" ]; then - echo "ERROR: conda not found. Activate '$ENVNAME' yourself or set CONDA_EXE." >&2; exit 1 -fi -source "$CB/etc/profile.d/conda.sh" -conda activate "$ENVNAME" || { echo "ERROR: 'conda activate $ENVNAME' failed" >&2; exit 1; } -echo "conda: base=$CB env=$ENVNAME python=$(which python)" - # repo example dir (portable across nodes) -EX="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +EX="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # .../examples/async_cifar10 +REPO_ROOT="$(cd "$EX/../../../.." && pwd)" # flame/ + +# shared harness: conda activation, PYTHONPATH pin, launch/progress loop, log +# asserts, preflight bridge -- moved here from inline (shared with fwdllm). +# shellcheck source=../../scripts/expt_runner.sh +source "$REPO_ROOT/lib/python/examples/scripts/expt_runner.sh" + +expt_activate_conda dg_flame # default env dg_flame (FLAME_CONDA_ENV overrides) +expt_pin_pythonpath "$REPO_ROOT" cd "$EX" || exit 1 SCR=expt_scripts_2026 LOGDIR="${FLAME_LOGDIR:-/tmp/debug_run_logs}"; mkdir -p "$LOGDIR" @@ -59,6 +47,11 @@ SIM_WALL_CEILING_S="" # empty = max_experiment_runtime_s (1×, tight guard; sim MODE="both" # sim | real | both — which time_mode variant(s) of each baseline to run NUM_TRAINERS="" # empty = use whatever's in the parity config (300); non-smoke override only ALPHA="" # empty = use the parity config's dirichlet_alpha (0.1); e.g. 100 for homogeneous +DRY_RUN=0 # --dry-run: show the pre-flight table + checks, generate cfg, DON'T launch +SHOW_ALL=0 # --show-all: expand tier ③ + list passing checks +STRICT=0 # --strict: a BLOCKING pre-flight check aborts (default: warn + continue, so + # smoke_suite.sh's non-interactive timeout-wrapped runs never hang/abort) +AFTER="" # --after: comma list of post-launch hooks (parity,plot) -- see after_* below usage() { echo "usage: $0 [--baselines 'felix refl'] [--runtime-s 3600] [--mode sim|real|both] [--sim-wall-ceiling-s 2700] [--trace syn_20]" @@ -84,6 +77,14 @@ usage() { echo " --alpha Dirichlet alpha override (default: parity config's 0.1). Supported" echo " values have an n300 split: 0.1 / 1.0 / 10.0 / 100.0 (100=homogeneous)." echo " When set, the split lookup uses the n300 partition for that alpha." + echo " --dry-run show the pre-flight hyperparameter table + feasibility checks and the" + echo " generated cfg, then exit WITHOUT launching." + echo " --show-all expand tier ③ (config-baked rows) + list the passing checks too." + echo " --strict abort if a pre-flight check is BLOCKING (default: warn + continue, so" + echo " smoke_suite.sh's non-interactive runs never hang/abort)." + echo " --after HOOKS comma-separated post-launch hooks: parity (scripts.parity.cli --batch" + echo " sim-vs-real per baseline) and/or plot (analyze_run cross-baseline" + echo " streaming figs). Absorbs the old compare_overnight.sh. e.g. --after parity,plot" exit 2 } @@ -98,6 +99,10 @@ if [ "${1:-}" = "smoke" ]; then --mode) MODE="$2"; shift 2 ;; --trace) TRACE="$2"; shift 2 ;; --alpha) ALPHA="$2"; shift 2 ;; + --dry-run) DRY_RUN=1; shift ;; + --show-all) SHOW_ALL=1; shift ;; + --strict) STRICT=1; shift ;; + --after) AFTER="$2"; shift 2 ;; *) shift ;; esac done @@ -113,6 +118,10 @@ else --trace) TRACE="$2"; shift 2 ;; --num-trainers) NUM_TRAINERS="$2"; shift 2 ;; --alpha) ALPHA="$2"; shift 2 ;; + --dry-run) DRY_RUN=1; shift ;; + --show-all) SHOW_ALL=1; shift ;; + --strict) STRICT=1; shift ;; + --after) AFTER="$2"; shift 2 ;; # --node is DEPRECATED (node1/node2 split removed): baselines are filtered # from a single node-agnostic parity config, so the node is irrelevant. # Accept+ignore so existing wrappers don't hard-error. @@ -309,46 +318,113 @@ print(len(d.get('experiments', []))) PY } +# Thin wrapper over the shared harness's expt_launch, kept as a named function so +# the two call sites below are unchanged. run_node() { local label="$1" cfg="$2" budget_s="${3:-0}" n_exps="${4:-1}" - local start_ts; start_ts=$(date +%s) - # Baseline run-dir count — new dirs that appear are newly-started experiments. - local initial_runs; initial_runs=$(find experiments -maxdepth 1 -name "run_*" -type d 2>/dev/null | wc -l) + expt_launch "$label" "$cfg" "$EX" "$budget_s" "$n_exps" "$LOGDIR" +} - echo "[$(date '+%F %T')] START $label ($n_exps exp(s), ~${budget_s}s budget)" | tee -a "$LOGDIR/debug_run.log" +# cifar_preflight -- render the tiered hyperparameter table + feasibility +# checks (shared expt_runner.py) for the generated cfg. Returns 2 if a check is +# BLOCKING; by default that only warns (so smoke_suite.sh's non-interactive runs +# never hang), --strict makes it fatal. +GPUS_VISIBLE="$( (command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi --query-gpu=index --format=csv,noheader 2>/dev/null | wc -l) || echo 0)" +cifar_preflight() { + local cfg="$1" + EXPT_RUNNER_DIR="$EXPT_RUNNER_DIR" CFG="$cfg" \ + BASELINES="$BASELINES" MODE="$MODE" RUNTIME_S="$RUNTIME_S" TRACE="$TRACE" \ + ALPHA="$ALPHA" NUM_TRAINERS="$NUM_TRAINERS" SIM_WALL_CEILING_S="$SIM_WALL_CEILING_S" \ + GPUS_VISIBLE="$GPUS_VISIBLE" DRY_RUN="$DRY_RUN" SHOW_ALL="$SHOW_ALL" \ + EX="$EX" LOGDIR="$LOGDIR" \ + python - <<'PY' +import os, sys, yaml +sys.path.insert(0, os.environ["EXPT_RUNNER_DIR"]) +import expt_runner +env = os.environ.get +cfg = yaml.safe_load(open(env("CFG"), encoding="utf-8")) +exps = cfg.get("experiments", []) +gpus_vis = int(env("GPUS_VISIBLE") or "0") - # Background progress ticker: fires every 30s, prints elapsed/remaining/percent - # and how many experiments have started (each start creates a new run_* dir). - ( - while true; do - sleep 30 - local now; now=$(date +%s) - local elapsed=$(( now - start_ts )) - local pct=0 remaining=0 - if [ "$budget_s" -gt 0 ]; then - pct=$(( elapsed * 100 / budget_s )) - remaining=$(( budget_s - elapsed )) - [ "$pct" -gt 100 ] && pct=100 - [ "$remaining" -lt 0 ] && remaining=0 - fi - local curr; curr=$(find experiments -maxdepth 1 -name "run_*" -type d 2>/dev/null | wc -l) - local started=$(( curr - initial_runs )) - [ "$started" -lt 0 ] && started=0 - printf " [%s] %s | %ds elapsed / ~%ds (%d%%) | exp started: %d/%d\n" \ - "$(date '+%T')" "$label" "$elapsed" "$budget_s" "$pct" "$started" "$n_exps" - done - ) & - local ticker_pid=$! +rows2, checks = [], [] +for e in exps: + h = e["aggregator"]["config_overrides"]["hyperparameters"] + n = e.get("trainer", {}).get("num_trainers") + ng = e.get("execution", {}).get("num_gpus") + mts = h.get("min_trainers_to_start") + a = (e.get("trainer", {}).get("dataset", {}) or {}).get("dirichlet_alpha") + rows2.append({"label": e.get("name", "?")[:26], + "value": f"n_trainers={n} n_gpus={ng} min_start={mts} rounds={h.get('rounds')} alpha={a}"}) + if isinstance(n, int) and isinstance(mts, int) and n < mts: + checks.append({"name": f"num_trainers >= min_trainers_to_start ({e.get('name')})", + "level": "error", "detail": f"{n} < {mts} — join barrier never clears"}) + if isinstance(ng, int) and gpus_vis and ng > gpus_vis: + checks.append({"name": f"num_gpus <= gpus_visible ({e.get('name')})", + "level": "error", "detail": f"num_gpus={ng} > visible={gpus_vis}"}) - python -m flame.launch.run_experiment "$cfg" --example-dir "$EX" \ - < /dev/null >> "$LOGDIR/${label}.out" 2>&1 - local rc=$? +tiers = [ + {"name": "① REVIEW EVERY RUN", "rows": [ + {"label": "baselines", "value": env("BASELINES")}, + {"label": "mode", "value": env("MODE"), + **({"level": "warn", "note": "single-sided: parity needs both"} if env("MODE") != "both" else {})}, + {"label": "runtime_s", "value": env("RUNTIME_S")}, + {"label": "trace", "value": env("TRACE") or "", + **({"level": "warn", "note": "not syn_0"} if (env("TRACE") and env("TRACE") != "syn_0") else {})}, + {"label": "sim_wall_ceiling", "value": env("SIM_WALL_CEILING_S") or "= runtime_s (auto)"}, + {"label": "alpha", "value": env("ALPHA") or ""}, + ]}, + {"name": "② PER-EXPERIMENT (moderate)", "rows": rows2}, + {"name": "③ RARELY CHANGED", "collapsed": True, "rows": [ + {"label": "env", "value": os.environ.get("CONDA_DEFAULT_ENV", "?")}, + {"label": "gpus_visible", "value": str(gpus_vis)}, + {"label": "example_dir", "value": env("EX")}, + {"label": "logdir", "value": env("LOGDIR")}, + ]}, +] +checks.append({"name": "run names carry _real/_sim tags for parity glob", "level": "ok", + "detail": "make_debug_yaml keeps the parity config's _sim/_real suffixes"}) +spec = {"title": "CIFAR DEBUG RUN", "subtitle": f"{len(exps)} experiment(s)", + "dry_run": env("DRY_RUN") == "1", "tiers": tiers, "checks": checks} +sys.exit(expt_runner.render_and_gate(spec, show_all=(env("SHOW_ALL") == "1"))) +PY +} - kill "$ticker_pid" 2>/dev/null - wait "$ticker_pid" 2>/dev/null +# gate_or_continue -- shared post-preflight decision for both paths. +gate_or_continue() { + local rc="$1" + if [ "$DRY_RUN" = "1" ]; then + echo "--dry-run: generated cfg in $LOGDIR. Nothing launched."; exit 0 + fi + if [ "$rc" -eq 2 ]; then + if [ "$STRICT" = "1" ]; then + echo "Pre-flight BLOCKED (exit 2) and --strict set. Nothing launched." >&2; exit 2 + fi + echo "WARNING: pre-flight flagged a BLOCKING check (continuing; pass --strict to abort)." >&2 + fi +} - local elapsed=$(( $(date +%s) - start_ts )) - echo "[$(date '+%F %T')] DONE $label exit=$rc (took ${elapsed}s / ~${budget_s}s budget)" | tee -a "$LOGDIR/debug_run.log" +# ---- post-launch hooks (--after ...), dispatched by expt_dispatch_after ---- +# after_parity / after_plot: per-baseline sim-vs-real parity + cross-baseline +# streaming plots, off the maintained scripts.parity.cli engine. +after_parity() { + python -m scripts.parity.cli --batch --experiments-dir experiments \ + --baselines $BASELINES --json-out "$LOGDIR/parity_.json" +} +after_plot() { + local ar="$REPO_ROOT/scripts/analysis/analyze_run.py" mode b d + [ -f "$ar" ] || { echo " [after:plot] $ar not found — skipping" >&2; return 0; } + for mode in sim real; do + local dirs=() labels=() + for b in $BASELINES; do + d=$(ls -dt experiments/run_*dbg_*"${b}"*_"${mode}"* 2>/dev/null | head -1) + [ -n "$d" ] && [ -d "$d/telemetry" ] && { dirs+=("$d/telemetry"); labels+=("$b"); } + done + if [ "${#dirs[@]}" -ge 2 ]; then + echo " [after:plot] $mode cross-baseline: ${labels[*]}" + python "$ar" --compare-streaming "${dirs[@]}" --labels "${labels[@]}" \ + --out "$LOGDIR/${mode}_cross" || true + fi + done } # ---- smoke mode ---- @@ -361,7 +437,12 @@ if [ "$SMOKE" = "1" ]; then make_debug_yaml "$BASELINES" 240 "$cfg" 1 "$SIM_WALL_CEILING_S" "$MODE" "$TRACE" "" "$ALPHA" if [ -f "$cfg" ]; then _n=$(_count_exps "$cfg") + cifar_preflight "$cfg"; gate_or_continue $? run_node "dbg_smoke" "$cfg" $(( _n * 240 )) "$_n" + expt_assert_run "$EX" "$EXPT_LAST_MARKER" "dbg_smoke" + [ -n "$AFTER" ] && expt_dispatch_after "$AFTER" + elif [ "$DRY_RUN" = "1" ]; then + echo "--dry-run: no experiments matched baselines='$BASELINES' mode=$MODE. Nothing to show."; exit 0 fi echo "=== SMOKE RESULTS ===" for dd in experiments/run_*dbg_smoke_*; do @@ -389,6 +470,9 @@ fi _n_exps=$(_count_exps "$cfg") _budget=$(( _n_exps * RUNTIME_S )) echo " queued: $_n_exps exp(s), estimated budget ~${_budget}s (sim finishes faster than real)" +cifar_preflight "$cfg"; gate_or_continue $? run_node "debug_run" "$cfg" "$_budget" "$_n_exps" +expt_assert_run "$EX" "$EXPT_LAST_MARKER" "debug_run" +[ -n "$AFTER" ] && expt_dispatch_after "$AFTER" echo "Logs: $LOGDIR/debug_run.out" echo "Run dirs: experiments/run_*dbg_*" diff --git a/lib/python/examples/async_cifar10/scripts/parity/checks.py b/lib/python/examples/async_cifar10/scripts/parity/checks.py index 984939709..9720864a1 100644 --- a/lib/python/examples/async_cifar10/scripts/parity/checks.py +++ b/lib/python/examples/async_cifar10/scripts/parity/checks.py @@ -351,6 +351,38 @@ def _per_round_last_event(agg_rounds: list) -> dict: return by_round +def _progress_axis(agg_rounds: list) -> str: + """The run's true progress axis. Normal FL advances FL `round`; fwdllm holds + `round` static (one model, grads aggregated in place) and advances committed + `data_id`, so a round-keyed clock rung divides by a counter stuck at 1. Use + `round` when the run advances it (>1 distinct), else `data_id` when the fwdllm + cadence field `cycle_data_id` is present, else `round` -- async_cifar10 stays + round-keyed (byte-identical); only fwdllm re-keys.""" + rounds = {e.get("round") for e in agg_rounds if e.get("round") is not None} + if len(rounds) > 1: + return "round" + if any(e.get("cycle_data_id") is not None for e in agg_rounds): + return "data_id" + return "round" + + +def _per_progress_last_event(agg_rounds: list, axis: str) -> dict: + """{progress_unit -> last event on that unit (by ts)} on the given axis. + Mirrors _per_round_last_event but keyed on the run's true progress axis + (`round` or fwdllm's `cycle_data_id`), so the clock family measures + progress-per-time on the axis the run actually advances.""" + if axis == "round": + return _per_round_last_event(agg_rounds) + out: dict = {} + for e in agg_rounds: + k = e.get("cycle_data_id") + if k is None: + continue + if k not in out or e.get("ts", 0) > out[k].get("ts", 0): + out[k] = e + return out + + def _per_round_max_speed(agg_rounds: list) -> dict: """Per FL round: max trainer_speed_s across all commits in that round.""" out: dict = {} @@ -364,17 +396,46 @@ def _per_round_max_speed(agg_rounds: list) -> dict: return out +def _real_intrinsic_clock(agg_rounds: list) -> Optional[dict]: + """Real's genuine-time coordinate for the clock-rate rungs, or None. + + When the aggregator emits ``intrinsic_span_s`` (fwdllm), returns + {id(agg_round_event): cumulative_intrinsic_s} -- the running sum of per-cycle + algorithmic spans (barrier + fedavg + eval), the real analog of the sim's + vclock. Real's raw wall Δts bundles a ~constant inter-round transport artifact + (mqtt re-fetch / redistribute / drain-tail / sleeps) the sim omits by design, + so anchoring on this intrinsic clock compares real-genuine vs sim-vclock like + for like. None when ``intrinsic_span_s`` is absent (async_cifar10 -> callers + fall back to ``ts``, byte-identical).""" + evs = [e for e in agg_rounds if e.get("event") == "agg_round"] + if not any(e.get("intrinsic_span_s") is not None for e in evs): + return None + coord, run = {}, 0.0 + for e in evs: + run += (e.get("intrinsic_span_s") or 0.0) + coord[id(e)] = run + return coord + + def _per_round_advances(agg_rounds: list, use_vclock: bool) -> list: - """Compute per-FL-round time advances. + """Per-progress-unit time advances (positive only). - use_vclock=True: Δvclock_now between consecutive rounds (sim mode). - use_vclock=False: Δts (wall) between consecutive rounds (real mode). - Returns list of positive advances. + use_vclock=True: Δvclock_now between consecutive units (sim mode). + use_vclock=False: real mode -- Δ(intrinsic algorithmic clock) when the + aggregator emits ``intrinsic_span_s`` (see _real_intrinsic_clock), else Δts + (wall) -> async_cifar10 byte-identical. + + Keyed on the run's true progress axis (_progress_axis), not raw `round`: + fwdllm holds `round` static and advances `data_id`, so a round key yields <2 + units. async_cifar10 advances `round` -> byte-identical; only fwdllm re-keys. """ - by_round = _per_round_last_event(agg_rounds) + axis = _progress_axis(agg_rounds) + by_round = _per_progress_last_event(agg_rounds, axis) rounds_sorted = sorted(by_round.keys()) if len(rounds_sorted) < 2: return [] + # Real: prefer real's intrinsic algorithmic clock over raw wall ts. + real_coord = None if use_vclock else _real_intrinsic_clock(agg_rounds) advances = [] for i in range(1, len(rounds_sorted)): e_prev = by_round[rounds_sorted[i - 1]] @@ -382,15 +443,15 @@ def _per_round_advances(agg_rounds: list, use_vclock: bool) -> list: if use_vclock: v_prev = e_prev.get("vclock_now") v_curr = e_curr.get("vclock_now") - if v_prev is None or v_curr is None: - continue - adv = v_curr - v_prev + elif real_coord is not None: + v_prev = real_coord.get(id(e_prev)) + v_curr = real_coord.get(id(e_curr)) else: - t_prev = e_prev.get("ts") - t_curr = e_curr.get("ts") - if t_prev is None or t_curr is None: - continue - adv = t_curr - t_prev + v_prev = e_prev.get("ts") + v_curr = e_curr.get("ts") + if v_prev is None or v_curr is None: + continue + adv = v_curr - v_prev if adv > 0: advances.append(adv) return advances @@ -803,11 +864,14 @@ def _by_round_selection(selection_train: list) -> dict: return out -# Selectors whose per-round SET selection is a deterministic function of -# (candidate set, seed). Currently empty — every shipped selector samples -# from a join-order-dependent candidate list, making exact per-round set -# identity unattainable across real/sim. participation_parity is the -# enforced selection invariant for stochastic selectors. +# Selectors whose per-round SET selection is deterministic independent of the +# draw. Currently empty — every shipped selector samples a join-order-dependent +# candidate list, so a stochastic SUBSET draw is not set-identical across +# real/sim. Set-identity is also attainable when the run is FULL-COHORT +# (K >= candidate pool), detected data-drivenly by `_full_cohort_selection` / +# `_selection_is_deterministic`, which un-gates the set/sequence rungs for +# full-cohort runs while leaving subset selectors gated. participation_parity +# is the enforced invariant for the stochastic-subset case. DETERMINISTIC_SELECTORS: set = set() @@ -821,12 +885,69 @@ def _selector_name(*loaded: dict) -> str: return "" +def _has_cohort_counts(loaded: dict) -> bool: + """True iff selection telemetry carries the num_chosen/num_candidates fields + the full-cohort gate needs (real runs always do; synthetic/legacy may not).""" + for e in loaded.get("selection_train", []): + if e.get("num_candidates") is not None or e.get("num_chosen") is not None: + return True + return False + + +def _full_cohort_selection(loaded: dict) -> bool: + """True iff EVERY selection round chose the whole candidate pool + (num_chosen == num_candidates, pool > 0). + + When K >= the candidate pool the selected SET is the entire pool — + deterministic regardless of the draw — so the set/sequence selection rungs + become exact-enforceable. Under scarcity (K < pool) or asymmetric eligibility + across modes some round is not full-cohort → False → the rung stays gated + rather than false-failing a genuinely stochastic/divergent selection. Missing + telemetry also returns False — never assert determinism we cannot see. Data- + driven, so it self-disables under Phase-2 unavailability with no config change.""" + train = loaded.get("selection_train") or [] + saw = False + for e in train: + nchosen, ncand = e.get("num_chosen"), e.get("num_candidates") + if nchosen is None or ncand is None or ncand <= 0: + return False + saw = True + if nchosen != ncand: + return False + return saw + + +def _selection_is_deterministic(real: dict, sim: dict) -> bool: + """Whether the per-round selected SET is a deterministic function of the + candidate pool, so the set/sequence selection rungs should ENFORCE rather + than gate to a trivial pass. + + True when the selector is declared deterministic (DETERMINISTIC_SELECTORS) + OR the run is full-cohort in BOTH modes (_full_cohort_selection). NOTE: the + variance-cadence rung `cohort_sequence` is ungated-EXACT for ALL fwdllm + baselines (receive-order is deterministic by design) and is intentionally NOT + gated here. + + Fallback: when neither mode carries num_chosen/num_candidates telemetry + (synthetic/legacy runs) revert to the selector-name rule — enforce on an + unknown/deterministic selector, gate a known stochastic one — to avoid + silently weakening a check on telemetry that predates the count fields.""" + selector = _selector_name(real, sim) + if selector and selector in DETERMINISTIC_SELECTORS: + return True + if _has_cohort_counts(real) and _has_cohort_counts(sim): + return _full_cohort_selection(real) and _full_cohort_selection(sim) + return not bool(selector) # legacy fallback = old `not gated` semantics + + def selection_parity(real: dict, sim: dict, max_rounds: Optional[int] = None, warn_jaccard: float = 0.7) -> dict: """S1/S2: Per-round selection overlap (Jaccard). - Enforced only for DETERMINISTIC_SELECTORS; gated to WARN for stochastic - selectors (participation_parity is the enforced invariant for those). + Enforced when selection is deterministic (`_selection_is_deterministic`: + a DETERMINISTIC_SELECTORS selector OR a full-cohort run); gated to a + trivial pass otherwise (participation_parity is the enforced invariant for + the stochastic-subset case). """ r = _by_round_selection(real["selection_train"]) s = _by_round_selection(sim["selection_train"]) @@ -841,7 +962,7 @@ def selection_parity(real: dict, sim: dict, max_rounds: Optional[int] = None, exact += 1 mean_j = sum(js) / len(js) if js else float("nan") selector = _selector_name(real, sim) - gated = bool(selector) and selector not in DETERMINISTIC_SELECTORS + gated = not _selection_is_deterministic(real, sim) enforced_ok = (not js) or mean_j >= warn_jaccard return { "ok": True if gated else enforced_ok, @@ -1001,9 +1122,11 @@ def aggregation_sequence_parity(real: dict, sim: dict, max_rounds: Optional[int] = None) -> dict: """P1 / U1: Per-round set of contributing trainers matches across modes. - Enforced only for DETERMINISTIC_SELECTORS; gated to WARN for stochastic - selectors. Exact per-round contributing-set identity is unattainable for a - stochastic, streaming, path-dependent selector (a trainer is chosen in + Enforced when selection is deterministic (`_selection_is_deterministic`: + a DETERMINISTIC_SELECTORS selector OR a full-cohort run — fwdllm syn_0); + gated to a trivial pass otherwise. Exact per-round contributing-set + identity is unattainable for a stochastic-SUBSET, streaming, path-dependent + selector (a trainer is chosen in *different* rounds across modes), the same reason S1 (selection_parity) is gated. The enforced selection invariants for stochastic selectors are participation_parity (S2) + the pooled distributions; this check stays as a @@ -1021,7 +1144,7 @@ def by_round(agg_rounds): rounds = [x for x in rounds if x <= max_rounds] matches = sum(1 for rd in rounds if r[rd] == s[rd]) selector = _selector_name(real, sim) - gated = bool(selector) and selector not in DETERMINISTIC_SELECTORS + gated = not _selection_is_deterministic(real, sim) enforced_ok = (not rounds) or matches == len(rounds) return { "ok": True if gated else enforced_ok, @@ -1385,6 +1508,11 @@ def class_share(counter): speed_class_tvd = 0.5 * sum(abs(rcs.get(b, 0) - scs.get(b, 0)) for b in buckets) selector = _selector_name(real, sim) + # NOT un-gated by the full-cohort rule: participation keys on `round`, which is + # CONSTANT for fwdllm, so the matched-round window degenerates to nmatch=1 and a + # mechanical KS=1.0. fwdllm's per-cycle cohort enforcement is cohort_sequence; + # here the stochastic speed-class TVD branch is the right call. Set-based round + # rungs are full-cohort-safe (all-K union both sides); count rungs aren't. gated = bool(selector) and selector not in DETERMINISTIC_SELECTORS tvd_tol = 0.15 if gated and speed_class_tvd is not None: @@ -1606,7 +1734,7 @@ def per_trainer_utils(agg_rounds): avg_mean_diff = sum(mean_diffs) / len(mean_diffs) if mean_diffs else float("nan") selector = _selector_name(real, sim) - gated = bool(selector) and selector not in DETERMINISTIC_SELECTORS + gated = not _selection_is_deterministic(real, sim) pooled_ok = math.isnan(pooled_ks) or pooled_ks <= max_ks per_trainer_ok = math.isnan(max_ks_val) or max_ks_val <= max_ks # Stochastic: enforce the pooled distribution only. Deterministic: also @@ -1766,13 +1894,23 @@ def sim_rate_ok(sim: dict, min_rate: float = 0.01, max_rate: float = 100.0) -> d def failsafe_ok(sim: dict, budget_s: Optional[float] = None, - max_overshoot: float = 0.20) -> dict: - """K5 [INV]: sim wall must not overshoot sim_wall_ceiling_s by > 20%.""" + max_overshoot: float = 0.20, + real_compute_sim: Optional[bool] = None) -> dict: + """K5 [INV]: sim wall must not overshoot the budget by > 20%. + + For a REAL-COMPUTE sim (fwdllm runs the real forward-grad GPU pass in sim + mode) sim wall ≫ vclock by construction, so the vclock is the wrong budget -- + compare sim wall against the RUN WALL budget (`max_runtime_s`, passed as + budget_s); if none is available, SKIP rather than falling back to the vclock. + Auto-detected via the progress axis when not passed. A cheap-compute sim + (async_cifar10, wall≈vclock) keeps the vclock fallback -> byte-identical.""" rounds = [e for e in sim["agg_rounds"] if e.get("event") == "agg_round"] all_evs = sim.get("_all_events", sim["agg_rounds"]) # agg_rounds used as proxy if len(rounds) < 2: return {"ok": True, "tier": "INV", "status": "SKIP", "note": "fewer than 2 agg_round events"} + if real_compute_sim is None: + real_compute_sim = _progress_axis(sim["agg_rounds"]) == "data_id" wall_elapsed = rounds[-1]["ts"] - rounds[0]["ts"] failsafe_fired = any( "SIM_WALL_CEILING" in str(e.get("stop_reason", "")) or @@ -1780,6 +1918,11 @@ def failsafe_ok(sim: dict, budget_s: Optional[float] = None, for e in all_evs ) if budget_s is None: + if real_compute_sim: + return {"ok": True, "tier": "INV", "status": "SKIP", + "note": "real-compute sim (sim wall ≫ vclock by construction); " + "no run wall budget to compare against — SKIP not " + "wall-vs-vclock (§H #9)"} vclock_final = rounds[-1].get("vclock_now") if not vclock_final: return {"ok": True, "tier": "INV", "status": "SKIP", @@ -1795,6 +1938,7 @@ def failsafe_ok(sim: dict, budget_s: Optional[float] = None, "overshoot_frac": round(overshoot, 3), "failsafe_fired": failsafe_fired, "max_overshoot": max_overshoot, + "real_compute_sim": real_compute_sim, } @@ -1818,16 +1962,28 @@ def throughput_parity(real: dict, sim: dict, tol_rel: float = 0.05) -> dict: return {"ok": False, "tier": "EXACT", "note": "K10: no vclock_now in sim agg_round events — cannot compute throughput"} final_vclock = max(sim_vclock_vals) - sim_by_round = _per_round_last_event(sim["agg_rounds"]) + # Progress-axis re-key: count units on the axis the run advances -- `round` + # for normal FL, committed `data_id` for fwdllm (else n_rounds==1). + axis = "data_id" if "data_id" in (_progress_axis(sim["agg_rounds"]), + _progress_axis(real["agg_rounds"])) else "round" + sim_by_round = _per_progress_last_event(sim["agg_rounds"], axis) n_sim_rounds = len(sim_by_round) sim_throughput = n_sim_rounds / final_vclock if final_vclock > 0 else 0.0 - real_ts = [e["ts"] for e in real["agg_rounds"] if e.get("ts") is not None] - if not real_ts or len(real_ts) < 2: - return {"ok": True, "tier": "EXACT", "status": "SKIP", - "note": "insufficient real ts data (< 2 agg_round events)"} - wall_elapsed = max(real_ts) - min(real_ts) - real_by_round = _per_round_last_event(real["agg_rounds"]) + # REAL denominator is real's genuine algorithmic time -- total intrinsic span + # (barrier+fedavg+eval) when emitted, else wall ts span (async byte-identical). + # Excludes real's inter-round transport artifact so rounds-per-genuine-second + # compares like-for-like vs the sim's rounds-per-vclock-second. + real_coord = _real_intrinsic_clock(real["agg_rounds"]) + if real_coord is not None: + wall_elapsed = max(real_coord.values()) if real_coord else 0.0 + else: + real_ts = [e["ts"] for e in real["agg_rounds"] if e.get("ts") is not None] + if not real_ts or len(real_ts) < 2: + return {"ok": True, "tier": "EXACT", "status": "SKIP", + "note": "insufficient real ts data (< 2 agg_round events)"} + wall_elapsed = max(real_ts) - min(real_ts) + real_by_round = _per_progress_last_event(real["agg_rounds"], axis) n_real_rounds = len(real_by_round) real_throughput = n_real_rounds / wall_elapsed if wall_elapsed > 0 else 0.0 @@ -1902,6 +2058,102 @@ def per_round_advance_parity(real: dict, sim: dict, } +def wall_disparity(real: dict, sim: dict) -> dict: + """wall_disparity [DIAG]: |real_genuine − sim_vclock| per matched progress + unit -- the sanity metric to drive to ~0, surfaced every run without gating. + Real's coordinate is the cumulative intrinsic span (barrier+fedavg+eval) when + ``intrinsic_span_s`` is emitted -- NOT raw wall, which bundles the inter-round + transport artifact the sim omits (chasing real's full wall would over-charge + the vclock). Falls back to wall ts for async (byte-identical). Both clocks + cumulative from the first matched unit. Keyed on the progress axis. Never fails.""" + axis = "data_id" if "data_id" in (_progress_axis(sim["agg_rounds"]), + _progress_axis(real["agg_rounds"])) else "round" + real_by = _per_progress_last_event(real["agg_rounds"], axis) + sim_by = _per_progress_last_event(sim["agg_rounds"], axis) + # Real: genuine algorithmic clock when emitted, else raw wall ts. + real_coord = _real_intrinsic_clock(real["agg_rounds"]) + _real_t = ((lambda e: real_coord.get(id(e))) if real_coord is not None + else (lambda e: e.get("ts"))) + matched = [k for k in sorted(set(real_by) & set(sim_by)) + if _real_t(real_by[k]) is not None + and sim_by[k].get("vclock_now") is not None] + if len(matched) < 2: + return {"ok": True, "tier": "DIAG", "status": "SKIP", + "note": "fewer than 2 matched units with both real time and " + "sim vclock_now"} + ts0 = _real_t(real_by[matched[0]]) + v0 = sim_by[matched[0]]["vclock_now"] + residuals, per_unit = [], {} + for k in matched: + real_genuine = _real_t(real_by[k]) - ts0 + sim_vclock = sim_by[k]["vclock_now"] - v0 + resid = abs(real_genuine - sim_vclock) + residuals.append(resid) + per_unit[k] = round(resid, 2) + mean_resid = sum(residuals) / len(residuals) + return { + "ok": True, # DIAG: informational, never gates the ladder + "tier": "DIAG", + "axis": axis, + "anchor": "intrinsic_span" if real_coord is not None else "wall_ts", + "mean_abs_disparity_s": round(mean_resid, 2), + "max_abs_disparity_s": round(max(residuals), 2), + "n_matched_units": len(residuals), + "per_unit_abs_disparity_s": per_unit, + } + + +def _wall_span_s(agg: dict) -> Optional[float]: + """Physical wall seconds spanned by a run's agg_round events. Prefers the + emitted `wall_elapsed_s` (measured from the re-anchored agg start, excludes + the join wait) when present; falls back to the ts epoch span.""" + evs = [e for e in agg["agg_rounds"] if e.get("event") == "agg_round"] + we = [e.get("wall_elapsed_s") for e in evs if e.get("wall_elapsed_s") is not None] + if we: + return max(we) + ts = [e["ts"] for e in evs if e.get("ts") is not None] + if len(ts) >= 2: + return max(ts) - min(ts) + return None + + +def sim_speedup(real: dict, sim: dict, min_rate: float = 0.98) -> dict: + """sim_speedup [DIAG]: the sim must be a SPEEDUP, not a slowdown. Two numbers: + - sim_rate = final_vclock / sim_wall (virtual-s per wall-s). Invariant: + sim_rate >= 1 (vclock advances at least as fast as wall). K7 `sim_rate` + only checks the range [0.01,100], so a slowdown passes it silently -- + this rung is the invariant. + - wall_speedup = real_wall / sim_wall (how many times faster the sim + finishes the same work than the real run; > 1 is the whole point). + DIAG: surfaces every run, does not gate the ladder. For a real-compute sim + (fwdllm) the GPU pass is irreducible wall, so once the transport waits are + skipped sim_rate -> (gpu+D)/gpu >= 1.""" + vclock_vals = [e.get("vclock_now") for e in sim["agg_rounds"] + if e.get("vclock_now") is not None] + sim_wall = _wall_span_s(sim) + real_wall = _wall_span_s(real) + if not vclock_vals or not sim_wall or sim_wall <= 0: + return {"ok": True, "tier": "DIAG", "status": "SKIP", + "note": "no sim vclock_now or zero sim wall span"} + final_vclock = max(vclock_vals) + sim_rate = final_vclock / sim_wall + wall_speedup = (real_wall / sim_wall) if (real_wall and sim_wall > 0) else None + ok = sim_rate >= min_rate + return { + "ok": ok, # DIAG but ok reflects the #13 invariant so it shows red + "tier": "DIAG", + "sim_rate": round(sim_rate, 4), + "is_speedup": sim_rate >= min_rate, + "wall_speedup": round(wall_speedup, 3) if wall_speedup is not None else None, + "final_vclock_s": round(final_vclock, 1), + "sim_wall_s": round(sim_wall, 1), + "real_wall_s": round(real_wall, 1) if real_wall else None, + "min_rate": min_rate, + "note": ("SLOWDOWN — sim_rate < 1, the sim is broken (root #13)" + if not ok else "speedup healthy"), + } + + def overlap_factor(real: dict, sim: dict, tol: float = 0.3) -> dict: """K4 [DIAG]: async overlap factor diagnostic. @@ -1977,15 +2229,37 @@ def total_commits_parity(real: dict, sim: dict, tol_rel: float = 0.05) -> dict: real_ts = [e["ts"] for e in real["agg_rounds"] if e.get("ts") is not None] if not real_ts: return {"ok": False, "tier": "EXACT", "note": "no ts in real events"} + # REAL: matched-budget window on real's genuine algorithmic clock (cumulative + # intrinsic span) when emitted, else raw wall ts (async byte-identical). + # _real_time(e) is 0-based cumulative-intrinsic OR ts-real_t0. + real_coord = _real_intrinsic_clock(real["agg_rounds"]) real_t0 = min(real_ts) - final_real_wall = max(real_ts) - real_t0 + if real_coord is not None: + _real_time = lambda e: real_coord.get(id(e)) + final_real_wall = max(real_coord.values()) if real_coord else 0.0 + else: + _real_time = lambda e: (e["ts"] - real_t0) if e.get("ts") is not None else None + final_real_wall = max(real_ts) - real_t0 V = min(final_sim_vclock, final_real_wall) if V <= 0: return {"ok": True, "tier": "EXACT", "status": "SKIP", "note": "matched virtual budget V ≤ 0 — run too short to measure"} - n_sim = sum(1 for e in sim["agg_rounds"] if (e.get("vclock_now") or 0) <= V + 1e-9) - n_real = sum(1 for e in real["agg_rounds"] - if e.get("ts") is not None and (e["ts"] - real_t0) <= V + 1e-9) + # Progress-axis re-key: count commits on the axis the run advances. Normal FL + # commits once per agg_round; fwdllm's committed unit is the `data_id` (a + # variance-FAIL cycle rolls back, so raw cycle events overcount), so count + # DISTINCT committed data_ids within V. + axis = "data_id" if "data_id" in (_progress_axis(sim["agg_rounds"]), + _progress_axis(real["agg_rounds"])) else "round" + if axis == "round": + n_sim = sum(1 for e in sim["agg_rounds"] if (e.get("vclock_now") or 0) <= V + 1e-9) + n_real = sum(1 for e in real["agg_rounds"] + if _real_time(e) is not None and _real_time(e) <= V + 1e-9) + else: + sim_units = _per_progress_last_event(sim["agg_rounds"], axis) + real_units = _per_progress_last_event(real["agg_rounds"], axis) + n_sim = sum(1 for e in sim_units.values() if (e.get("vclock_now") or 0) <= V + 1e-9) + n_real = sum(1 for e in real_units.values() + if _real_time(e) is not None and _real_time(e) <= V + 1e-9) if max(n_sim, n_real, 1) == 0: return {"ok": True, "tier": "EXACT", "note": "no commits in V window"} rel_diff = abs(n_sim - n_real) / max(n_sim, n_real) @@ -2017,25 +2291,38 @@ def terminal_state_parity(real: dict, sim: dict, real_ts_all = [e["ts"] for e in real["agg_rounds"] if e.get("ts") is not None] if not real_ts_all: return {"ok": False, "tier": "EXACT", "note": "no ts in real events"} + # REAL: matched-budget window on real's genuine algorithmic clock (see + # total_commits_parity); falls back to wall ts for async (byte-identical). + real_coord = _real_intrinsic_clock(real["agg_rounds"]) real_t0 = min(real_ts_all) - final_real_wall = max(real_ts_all) - real_t0 + if real_coord is not None: + _real_time = lambda e: real_coord.get(id(e)) + final_real_wall = max(real_coord.values()) if real_coord else 0.0 + else: + _real_time = lambda e: (e["ts"] - real_t0) if e.get("ts") is not None else None + final_real_wall = max(real_ts_all) - real_t0 V = min(final_sim_vclock, final_real_wall) if V <= 0: return {"ok": True, "tier": "EXACT", "status": "SKIP", "note": "matched virtual budget V ≤ 0 — run too short to measure"} - sim_by_round = _per_round_last_event(sim["agg_rounds"]) - real_by_round = _per_round_last_event(real["agg_rounds"]) + # Progress-axis re-key: "rounds at V" is really "progress units at V" -- FL + # rounds for normal FL, committed data_ids for fwdllm (round static). + axis = "data_id" if "data_id" in (_progress_axis(sim["agg_rounds"]), + _progress_axis(real["agg_rounds"])) else "round" + unit_key = "round" if axis == "round" else "cycle_data_id" + sim_by_round = _per_progress_last_event(sim["agg_rounds"], axis) + real_by_round = _per_progress_last_event(real["agg_rounds"], axis) sim_rounds_at_V = {r for r, e in sim_by_round.items() if (e.get("vclock_now") or 0) <= V + 1e-9} real_rounds_at_V = {r for r, e in real_by_round.items() - if e.get("ts") is not None and (e["ts"] - real_t0) <= V + 1e-9} + if _real_time(e) is not None and _real_time(e) <= V + 1e-9} - def _trainers(agg_rounds, round_set): + def _trainers(agg_rounds, unit_set): ts = set() for e in agg_rounds: - if e.get("round") in round_set: + if e.get(unit_key) in unit_set: ts.update(e.get("contributing_trainers", [])) return ts @@ -2153,6 +2440,80 @@ def gpu_budget_ok(trainers: dict, warn_overrun_frac: float = 0.25) -> dict: } +def _overrun_stats(trainers: dict) -> tuple: + """(#rounds, #overran, earliest (data_id, iter) overran) from trainer_round + telemetry. An event counts only when `training_overran` is present.""" + n = over = 0 + first = None + for _tid, d in trainers.items(): + for e in d.get("trainer_round", []): + if "training_overran" not in e: + continue + n += 1 + if e.get("training_overran"): + over += 1 + did = e.get("data_id") + if did is not None: + it = e.get("iteration_per_data_id") + cand = (did, it if it is not None else 0) + if first is None or cand < first: + first = cand + return n, over, first + + +def timing_overrun(real_trainers: dict, sim_trainers: dict, + warn_frac: float = 0.05) -> dict: + """Ovr [DIAG]: fraction of trainer rounds where the real GPU pass OVERRAN the + modeled mobile-device delay budget. + + Precondition-for-parity localizer, NOT a real↔sim diff. A trainer's arrival + order is deterministic only while gpu_time_s <= the modeled delay D (the GPU + hides inside the device wall). When gpu > D the update completes after the + vclock passed its sct → the sim can commit it out of order → the deterministic + arrival order that cohort_sequence/v2_var_trajectory require is broken. So a + cohort/var divergence with a HIGH overrun fraction is a timing-MODEL limitation + (GPU too slow for the budget), fixable by raising `delay_factor` or lowering + `perturbation_count`, NOT a sim ordering bug. A near-zero fraction is the + precondition for exact cohort/var parity. + + DIAG (never fails); reports per-mode overrun fraction + the earliest + (data_id, iteration) an overrun occurs (cross-check against cohort_sequence + first_divergence). SKIPs when `training_overran` is absent (delays disabled). + """ + rn, ro, rf = _overrun_stats(real_trainers) + sn, so, sf = _overrun_stats(sim_trainers) + if rn == 0 and sn == 0: + return {"ok": True, "tier": "DIAG", "status": "SKIP", + "note": "no training_overran telemetry (delays disabled or " + "pre-P2-6 run)"} + r_frac = ro / rn if rn else 0.0 + s_frac = so / sn if sn else 0.0 + worst = max(r_frac, s_frac) + if worst == 0.0: + verdict = ("no overrun — GPU within the modeled delay budget in both " + "modes; per-trainer arrival order is deterministic → exact " + "cohort/var parity is attainable") + elif worst < warn_frac: + verdict = (f"marginal overrun ({worst:.1%} < {warn_frac:.0%}) — order " + "determinism mostly holds; watch the cohort_sequence tail") + else: + verdict = (f"OVERRUN {worst:.1%} — GPU exceeds the modeled delay budget; " + "update order can flip → expect cohort_sequence/v2 breaks. " + "Raise delay_factor or lower perturbation_count (fluxtune).") + return { + "ok": True, + "tier": "DIAG", + "real_overrun_frac": round(r_frac, 4), + "sim_overrun_frac": round(s_frac, 4), + "real_overran": ro, "real_rounds": rn, + "sim_overran": so, "sim_rounds": sn, + "real_first_overrun": list(rf) if rf else None, + "sim_first_overrun": list(sf) if sf else None, + "warn_frac": warn_frac, + "verdict": verdict, + } + + _PHASE_FIELDS = ("pre_train_s", "gpu_compute_s", "mqtt_fetch_s", "weights_to_gpu_s", "weights_to_ram_s", "post_train_s") @@ -2215,8 +2576,12 @@ def collect_phases(trainers: dict) -> dict: ("selection.num_eligible", "sel", "num_eligible", "both"), ("selection.avail_composition", "sel", "avail_composition", "both"), ("selection.num_chosen", "sel", "num_chosen", "both"), - ("trainer_round.gpu_compute_s", "trainer_round", "gpu_compute_s", "both"), - ("trainer_round.training_budget_s", "trainer_round", "training_budget_s", "both"), + # fwdllm's trainer emits the same data under different names: its forward-grad + # "compute" is real_gpu_time_s and its budget is + # sim_round_duration_s (gpu + modeled delay). Accept either spelling so the + # coverage matrix agrees on both examples instead of false-FAILing fwdllm. + ("trainer_round.gpu_compute_s", "trainer_round", ("gpu_compute_s", "real_gpu_time_s"), "both"), + ("trainer_round.training_budget_s", "trainer_round", ("training_budget_s", "sim_round_duration_s"), "both"), ("task_recv.sim_send_ts", "task_recv", "sim_send_ts", "sim"), ] @@ -2230,10 +2595,15 @@ def field_coverage(real_agg: dict, sim_agg: dict, into "these fields are absent in sim". FAIL-LOUD when an expected field has zero density in a mode that requires it. """ - def _density(events: list, field: str) -> Optional[float]: + def _density(events: list, field) -> Optional[float]: + # `field` may be a single name or a tuple of accepted aliases (an event + # counts as covered if ANY alias is present) -- lets one canonical spec + # row match a different-but-equivalent field name per example. if not events: return None - n = sum(1 for e in events if e.get(field) not in (None, [], {})) + fields = field if isinstance(field, tuple) else (field,) + n = sum(1 for e in events + if any(e.get(f) not in (None, [], {}) for f in fields)) return n / len(events) def _agg_evs(agg, src): @@ -3424,6 +3794,584 @@ def _curve(evs): "eval_rounds_compared": len(diffs), "loss_tol": loss_tol}, budget_s) +# ═══════════════════════════════════════════════════════════════════ +# §F FwdLLM variance-cadence layer (PARITY.md §F.4) — V/DK/G rungs +# ═══════════════════════════════════════════════════════════════════ +# +# fwdllm's commit cadence is ENDOGENOUS: at each agg-goal boundary aggregate() +# computes a gradient-pool `var`; var<=var_threshold commits (advance data_id, +# eval, clear cached_v) else rolls back and retries the same data_id. So +# updates-per-data_id is a random variable of the variance trajectory. These +# rungs verify that trajectory is mode-invariant given matched inputs — the +# fwdllm-specific layer the async_cifar10 ladder does not model. Never fix an +# EMERGENT rung directly: walk to the lowest rung whose *inputs* are matched. +# All read the per-cycle agg_round series (fwdllm_aggregator emits one event +# per agg-goal boundary carrying cycle_data_id / cycle_iteration / var / +# var_threshold / var_good_enough / force_commit_planned / grad_pool_size / +# cached_v_size). Inputs source: fwdllm_aggregator.py _process_aggregation_goal_met. + + +def _fwd_cadence_cycles(agg: dict, max_bin: Optional[int] = None) -> list: + """Ordered per-cycle agg_round events carrying fwdllm cadence fields. + + A cycle is one agg-goal boundary (a variance gate). Non-fwdllm runs (no + `var_good_enough`/`cycle_data_id`) yield [] -> the V/DK/G rungs SKIP. + + `max_bin` restricts to cycles whose `cycle_data_id` <= max_bin (the + first-data-bin logical-parity window, simulate_fwdllm.md §A); None = all. + A cycle with no `cycle_data_id` is kept only when unwindowed. + """ + out = [e for e in agg.get("agg_rounds", []) + if "cycle_data_id" in e or "var_good_enough" in e] + if max_bin is not None: + out = [e for e in out + if e.get("cycle_data_id") is not None + and e["cycle_data_id"] <= max_bin] + return out + + +def _iters_per_data_id(cycles: list) -> dict: + """{cycle_data_id -> #cycles spent on it} = realized dynamic-K per data_id. + + Exact for both natural-pass and force-commit paths because each agg-goal + boundary emits exactly one cadence event tagged with the data_id it worked + on (cycle_data_id), pre-advance.""" + out: dict = {} + for e in cycles: + d = e.get("cycle_data_id") + if d is None: + continue + out[d] = out.get(d, 0) + 1 + return out + + +def iters_per_data_id_parity(real: dict, sim: dict, ks_tol: float = 0.2, + mean_tol_rel: float = 0.15, + max_bin: Optional[int] = None) -> dict: + """V1 [DIST]: iterations-per-data_id distribution (realized dynamic K). + + The number of accumulation cycles a data_id needs to pass the variance gate. + A divergence means the contributing set/order (U5/U4 upstream) differs, so + the accumulated grad-pool composition — and thus the variance trajectory — + differs. KS on the per-data_id iteration counts + a mean guard. + """ + rc, sc = _fwd_cadence_cycles(real, max_bin), _fwd_cadence_cycles(sim, max_bin) + if not rc or not sc: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no fwdllm cadence events (non-fwdllm run or telemetry absent)"} + r_iters = list(_iters_per_data_id(rc).values()) + s_iters = list(_iters_per_data_id(sc).values()) + if not r_iters or not s_iters: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no cycle_data_id in cadence events — cannot bin by data_id"} + ks = ks_stat(s_iters, r_iters) + r_mean, _ = mean_std(r_iters) + s_mean, _ = mean_std(s_iters) + mean_rel = (abs(r_mean - s_mean) / max(r_mean, s_mean) + if max(r_mean, s_mean) > 0 else 0.0) + return { + "ok": ks <= ks_tol and mean_rel <= mean_tol_rel, + "tier": "DIST", + "real_mean_iters": round(r_mean, 3), + "sim_mean_iters": round(s_mean, 3), + "mean_rel_diff": round(mean_rel, 3), + "ks_stat": round(ks, 3), + "ks_tol": ks_tol, + "mean_tol_rel": mean_tol_rel, + "n_real_data_ids": len(r_iters), + "n_sim_data_ids": len(s_iters), + } + + +def var_trajectory_parity(real: dict, sim: dict, ks_tol: float = 0.2, + mean_tol_rel: float = 0.02, + max_bin: Optional[int] = None) -> dict: + """V2 [DIST]: per-cycle `var` trajectory distribution. + + With V1's inputs matched, the variance *signal* itself must match; a + divergence with matched iterations points at a grad-pool accumulation-order + bug (a true sim bug, not an input divergence). KS over the per-cycle var + values (None dropped — a cycle before the first gate has no var), PLUS a + relative-mean guard: KS alone is blind to a systematic offset (a uniform + ~1% shift barely moves the empirical CDFs → KS≈0), which is exactly the + grad-desync signature (simulate_fwdllm.md §A). The mean guard fails it. + """ + rc, sc = _fwd_cadence_cycles(real, max_bin), _fwd_cadence_cycles(sim, max_bin) + r_var = [e["var"] for e in rc if e.get("var") is not None] + s_var = [e["var"] for e in sc if e.get("var") is not None] + if not r_var or not s_var: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no non-null `var` in cadence events"} + ks = ks_stat(s_var, r_var) + r_mean, _ = mean_std(r_var) + s_mean, _ = mean_std(s_var) + mean_rel = (abs(r_mean - s_mean) / max(abs(r_mean), abs(s_mean)) + if max(abs(r_mean), abs(s_mean)) > 0 else 0.0) + return { + "ok": ks <= ks_tol and mean_rel <= mean_tol_rel, + "tier": "DIST", + "real_mean_var": round(r_mean, 6), + "sim_mean_var": round(s_mean, 6), + "mean_rel_diff": round(mean_rel, 4), + "mean_tol_rel": mean_tol_rel, + "ks_stat": round(ks, 3), + "ks_tol": ks_tol, + "n_real_cycles": len(r_var), + "n_sim_cycles": len(s_var), + } + + +def cohort_sequence_parity(real: dict, sim: dict, max_bin: Optional[int] = None, + var_rel_tol: float = 1e-3) -> dict: + """L1 [EXACT]: the ordered per-aggregation logical sequence is IDENTICAL. + + The strongest logical-parity rung (simulate_fwdllm.md §A). Unlike + `aggregation_sequence_parity` (per-`round`, gated to + WARN for stochastic selectors), this keys on the fwdllm variance-cadence + *cycle* stream (`_fwd_cadence_cycles`, ordered) and asserts, cycle-by-cycle: + + - COHORT SET : same trainers committed together (which is fluxtune's bug) + - COHORT ORDER : same receive/commit order — NOT benign; the fwdllm + variance is a split-half statistic over the commit-ORDERED grad list, so + a reshuffle changes `var` even for an identical set + - CADENCE : (cycle_data_id, iteration_per_data_id, agg_goal_count, + var_good_enough, force_commit_planned) identical + - VAR VALUE : per-cycle `var` matches within var_rel_tol (grads are + deterministic GIVEN matched order → var must match; a ~1% gap is the + RNG-desync tell) + + Real receive-order is deterministic in both modes by design, so exact match + is the correct target (NOT gated). Enforced EXACT: any divergence fails. + SKIPs cleanly on non-fwdllm runs (no cadence fields → async_cifar10 etc.). + `max_bin` restricts to the first-data-bin window. + """ + rc = _fwd_cadence_cycles(real, max_bin) + sc = _fwd_cadence_cycles(sim, max_bin) + if not rc or not sc: + return {"ok": True, "tier": "EXACT", "status": "SKIP", + "note": "no fwdllm cadence events (non-fwdllm run or telemetry absent)"} + + def cohort(e): # receive/commit-ordered contributing trainers + return list(e.get("contributing_trainers") or []) + + def cadence(e): + return (e.get("cycle_data_id"), e.get("iteration_per_data_id"), + e.get("agg_goal_count"), e.get("var_good_enough"), + e.get("force_commit_planned")) + + n = min(len(rc), len(sc)) + set_m = order_m = cad_m = var_m = 0 + first_div = None + for i in range(n): + r, s = rc[i], sc[i] + rc_ord, sc_ord = cohort(r), cohort(s) + set_ok = sorted(rc_ord) == sorted(sc_ord) + order_ok = rc_ord == sc_ord + cad_ok = cadence(r) == cadence(s) + rv, sv = r.get("var"), s.get("var") + var_ok = (rv is None and sv is None) or ( + rv is not None and sv is not None + and abs(rv - sv) <= var_rel_tol * max(abs(rv), abs(sv), 1e-9)) + set_m += set_ok; order_m += order_ok; cad_m += cad_ok; var_m += var_ok + if first_div is None and not (set_ok and order_ok and cad_ok and var_ok): + first_div = { + "cycle_index": i, + "real": {"data_id": r.get("cycle_data_id"), + "iter": r.get("iteration_per_data_id"), + "cohort": rc_ord, "var": rv, + "var_good": r.get("var_good_enough")}, + "sim": {"data_id": s.get("cycle_data_id"), + "iter": s.get("iteration_per_data_id"), + "cohort": sc_ord, "var": sv, + "var_good": s.get("var_good_enough")}, + "set_ok": set_ok, "order_ok": order_ok, + "cadence_ok": cad_ok, "var_ok": var_ok, + } + ok = (len(rc) == len(sc) and set_m == n and order_m == n + and cad_m == n and var_m == n) + return { + "ok": ok, + "tier": "EXACT", + "cycles_compared": n, + "n_real_cycles": len(rc), + "n_sim_cycles": len(sc), + "set_match_frac": round(set_m / n, 3) if n else None, + "order_match_frac": round(order_m / n, 3) if n else None, + "cadence_match_frac": round(cad_m / n, 3) if n else None, + "var_match_frac": round(var_m / n, 3) if n else None, + "max_bin": max_bin, + "first_divergence": first_div, + } + + +def cached_v_pool_parity(real: dict, sim: dict, ks_tol: float = 0.25, + max_bin: Optional[int] = None) -> dict: + """V3 [DIAG]: `cached_v` (carried aggregated grad-pool) size over time. + + A rollback/cache bookkeeping divergence looks like a variance bug but is + stateful accounting — this DIAG localizes it. cached_v carries across + variance-FAIL rollbacks and clears on a commit; its size trajectory should + match once V1 matches. SKIP if the field was not emitted. + """ + rc, sc = _fwd_cadence_cycles(real, max_bin), _fwd_cadence_cycles(sim, max_bin) + r_sz = [e["cached_v_size"] for e in rc if e.get("cached_v_size") is not None] + s_sz = [e["cached_v_size"] for e in sc if e.get("cached_v_size") is not None] + if not r_sz or not s_sz: + return {"ok": True, "tier": "DIAG", "status": "SKIP", + "note": "no cached_v_size in cadence events (field not emitted)"} + ks = ks_stat(s_sz, r_sz) + return { + "ok": ks <= ks_tol, + "tier": "DIAG", + "real_mean_cached_v": round(sum(r_sz) / len(r_sz), 3), + "sim_mean_cached_v": round(sum(s_sz) / len(s_sz), 3), + "ks_stat": round(ks, 3), + "ks_tol": ks_tol, + } + + +def force_commit_rate_parity(real: dict, sim: dict, tol: float = 0.05, + max_bin: Optional[int] = None) -> dict: + """V4 [DIST]: force-commit frequency (max_iterations_per_data_id bypass rate). + + The fraction of cycles that hit the iteration cap and force-commit despite a + failed variance gate. A rate divergence = chronic variance divergence (the + cap fires at a different frequency), not a separate bug — walk to V1. + var_threshold / max_iterations_per_data_id are baseline-defining config + knobs, NOT parity levers. + """ + rc, sc = _fwd_cadence_cycles(real, max_bin), _fwd_cadence_cycles(sim, max_bin) + if not rc or not sc: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no fwdllm cadence events"} + r_rate = sum(1 for e in rc if e.get("force_commit_planned")) / len(rc) + s_rate = sum(1 for e in sc if e.get("force_commit_planned")) / len(sc) + return { + "ok": abs(r_rate - s_rate) <= tol, + "tier": "DIST", + "real_force_commit_rate": round(r_rate, 4), + "sim_force_commit_rate": round(s_rate, 4), + "abs_diff": round(abs(r_rate - s_rate), 4), + "tol": tol, + "n_real_cycles": len(rc), + "n_sim_cycles": len(sc), + } + + +def variance_pass_ratio_parity(real: dict, sim: dict, tol: float = 0.05, + max_bin: Optional[int] = None) -> dict: + """V5 [DIST]: genuine variance-pass ratio (the rollup feeding DynamicKC). + + A *genuine* pass is var<=var_threshold — a force-commit (var>threshold, + committed only because the iteration cap fired) is NOT a variance pass and is + excluded, so V5 tracks the true gate-pass rate the DynamicKC policy consumes. + EMERGENT rollup of V1/V2; localize down, do not tune it. + """ + def _ratio(cycles: list): + n, passes = 0, 0 + for e in cycles: + var, thr = e.get("var"), e.get("var_threshold") + if var is None or thr is None: + continue + n += 1 + if var <= thr: + passes += 1 + return (passes / n if n else None), n + + rc, sc = _fwd_cadence_cycles(real, max_bin), _fwd_cadence_cycles(sim, max_bin) + r_ratio, r_n = _ratio(rc) + s_ratio, s_n = _ratio(sc) + if r_ratio is None or s_ratio is None: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no var/var_threshold pairs in cadence events"} + return { + "ok": abs(r_ratio - s_ratio) <= tol, + "tier": "DIST", + "real_pass_ratio": round(r_ratio, 4), + "sim_pass_ratio": round(s_ratio, 4), + "abs_diff": round(abs(r_ratio - s_ratio), 4), + "tol": tol, + "n_real_gated": r_n, + "n_sim_gated": s_n, + } + + +def agg_goal_trajectory_parity(real: dict, sim: dict, ks_tol: float = 0.2) -> dict: + """DK1 [DIST]: K (`_agg_goal`) trajectory. + + When DynamicKC is enabled it moves K from observed metrics; a divergence + feeds back into cadence. INERT for the fixed-K baselines: if K is constant + and equal across modes there is no dynamic behavior to check -> SKIP so a + fixed-K run does not spuriously PASS/FAIL a mechanism it never exercises. + """ + rc, sc = _fwd_cadence_cycles(real), _fwd_cadence_cycles(sim) + r_k = [e["agg_goal"] for e in rc if e.get("agg_goal") is not None] + s_k = [e["agg_goal"] for e in sc if e.get("agg_goal") is not None] + if not r_k or not s_k: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no agg_goal in cadence events"} + if len(set(r_k)) <= 1 and len(set(s_k)) <= 1: + ok = set(r_k) == set(s_k) + return {"ok": ok, "tier": "DIST", "status": "SKIP", + "note": f"DynamicKC disabled — constant K (real={r_k[0]}, sim={s_k[0]})", + "real_k": r_k[0], "sim_k": s_k[0]} + ks = ks_stat(s_k, r_k) + return { + "ok": ks <= ks_tol, "tier": "DIST", + "real_mean_k": round(sum(r_k) / len(r_k), 3), + "sim_mean_k": round(sum(s_k) / len(s_k), 3), + "ks_stat": round(ks, 3), "ks_tol": ks_tol, + } + + +def dynamic_c_trajectory_parity(real: dict, sim: dict, ks_tol: float = 0.2) -> dict: + """DK2 [DIST]: C (`dynamic_c`) concurrency-target trajectory. + + SKIP unless a `dynamic_c` field is emitted (DynamicKC enabled). Do not fork + the shared controller per baseline (PARITY.md §F.4 locked principle 5). + """ + rc, sc = _fwd_cadence_cycles(real), _fwd_cadence_cycles(sim) + r_c = [e["dynamic_c"] for e in rc if e.get("dynamic_c") is not None] + s_c = [e["dynamic_c"] for e in sc if e.get("dynamic_c") is not None] + if not r_c or not s_c: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no dynamic_c in cadence events (DynamicKC disabled)"} + ks = ks_stat(s_c, r_c) + return { + "ok": ks <= ks_tol, "tier": "DIST", + "real_mean_c": round(sum(r_c) / len(r_c), 3), + "sim_mean_c": round(sum(s_c) / len(s_c), 3), + "ks_stat": round(ks, 3), "ks_tol": ks_tol, + } + + +def eligible_ends_metric_parity(real: dict, sim: dict, ks_tol: float = 0.2) -> dict: + """DK3 [DIST/CONTROL]: eligible-ends-count metric fed to the DynamicKC policy. + + Validate the policy *input* before the policy (CONTROL before MECHANISM): a + diverging input means fix the metric, not the policy. SKIP unless the + n_eligible_train/n_eligible_eval counts are emitted — deferred while no + baseline enables DynamicKC, not silently dropped. + """ + rc, sc = _fwd_cadence_cycles(real), _fwd_cadence_cycles(sim) + r_e = [e["n_eligible_train"] for e in rc if e.get("n_eligible_train") is not None] + s_e = [e["n_eligible_train"] for e in sc if e.get("n_eligible_train") is not None] + if not r_e or not s_e: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no n_eligible_train in cadence events (DK3 emit deferred, §K-D10)"} + ks = ks_stat(s_e, r_e) + return { + "ok": ks <= ks_tol, "tier": "DIST", + "real_mean_eligible": round(sum(r_e) / len(r_e), 3), + "sim_mean_eligible": round(sum(s_e) / len(s_e), 3), + "ks_stat": round(ks, 3), "ks_tol": ks_tol, + } + + +def grad_norm_parity(real: dict, sim: dict, ks_tol: float = 0.2) -> dict: + """G1 [DIST]: per-update grad/JVP norm distribution. + + Gradient values are mode-invariant given identical input + perturbation seed, + so G1 should be ~0; a FAIL means a perturbation seed/order leaked across + modes. SKIP unless a per-update `grad_norm` field is emitted — trainer-side + per-update emit deferred, logged not silently dropped. + """ + def _norms(agg): + out = [] + for e in _fwd_cadence_cycles(agg): + gn = e.get("grad_norm") + if isinstance(gn, list): + out.extend(x for x in gn if x is not None) + elif gn is not None: + out.append(gn) + return out + + r_n, s_n = _norms(real), _norms(sim) + if not r_n or not s_n: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no grad_norm in cadence events (G1 emit deferred, §K-D10)"} + ks = ks_stat(s_n, r_n) + return { + "ok": ks <= ks_tol, "tier": "DIST", + "real_mean_grad_norm": round(sum(r_n) / len(r_n), 6), + "sim_mean_grad_norm": round(sum(s_n) / len(s_n), 6), + "ks_stat": round(ks, 3), "ks_tol": ks_tol, + } + + +def grad_pool_size_parity(real: dict, sim: dict, ks_tol: float = 0.2, + mean_tol_rel: float = 0.15) -> dict: + """G2 [DIST]: grad_pool size at commit (realized contributions per data_id). + + The number of gradients accumulated into the pool when a data_id commits — + an emergent rollup of V1 x K. Measured only on committed cycles + (var_good_enough True). SKIP if grad_pool_size was not emitted. + """ + def _sizes(agg): + return [e["grad_pool_size"] for e in _fwd_cadence_cycles(agg) + if e.get("var_good_enough") and e.get("grad_pool_size") is not None] + + r_sz, s_sz = _sizes(real), _sizes(sim) + if not r_sz or not s_sz: + return {"ok": True, "tier": "DIST", "status": "SKIP", + "note": "no grad_pool_size on committed cycles"} + ks = ks_stat(s_sz, r_sz) + r_mean, _ = mean_std(r_sz) + s_mean, _ = mean_std(s_sz) + mean_rel = (abs(r_mean - s_mean) / max(r_mean, s_mean) + if max(r_mean, s_mean) > 0 else 0.0) + return { + "ok": ks <= ks_tol and mean_rel <= mean_tol_rel, + "tier": "DIST", + "real_mean_pool": round(r_mean, 3), + "sim_mean_pool": round(s_mean, 3), + "mean_rel_diff": round(mean_rel, 3), + "ks_stat": round(ks, 3), "ks_tol": ks_tol, "mean_tol_rel": mean_tol_rel, + "n_real_commits": len(r_sz), "n_sim_commits": len(s_sz), + } + + +# ═══════════════════════════════════════════════════════════════════ +# §3.5 FwdLLM async residence rungs (R1 / W1) — simulate_fwdllm.md §L.3 +# ═══════════════════════════════════════════════════════════════════ +# +# Catch a residence violation on the async grad path (surplus grads dropped + +# re-dispatched every agg-goal cycle → sim doing ~2x real's forward passes). +# R1 is the finest check (per-trainer interval overlap); W1 is the coarse +# compute-conservation tell that first flags the wasted recompute. + +_R1_EPS = 1e-6 + + +def _overlap_fraction(cycles: list) -> tuple: + """(overlap_frac, n_pairs, n_intervals) over per-trainer dispatch->commit + intervals reconstructed from the agg_round `contributor_intervals` field. + + A trainer's contribution i "overlaps" if its dispatch_ts precedes the + latest commit_ts among that trainer's earlier contributions -- i.e. it was + re-dispatched while a prior update was still outstanding, the one-in-flight + residence violation. 0.0 = strict residence; both modes must be ~0. + """ + by_end: dict = {} + for e in cycles: + for iv in (e.get("contributor_intervals") or []): + d, c = iv.get("dispatch_ts"), iv.get("commit_ts") + if d is None or c is None: + continue + by_end.setdefault(iv.get("end"), []).append((float(d), float(c))) + n_pairs = 0 + n_overlap = 0 + n_intervals = 0 + for ivs in by_end.values(): + ivs.sort() + n_intervals += len(ivs) + running_max_commit = float("-inf") + for i, (d, c) in enumerate(ivs): + if i > 0: + n_pairs += 1 + if d < running_max_commit - _R1_EPS: + n_overlap += 1 + running_max_commit = max(running_max_commit, c) + frac = (n_overlap / n_pairs) if n_pairs else 0.0 + return frac, n_pairs, n_intervals + + +def inflight_overlap_parity(real: dict, sim: dict, tol_frac: float = 0.02) -> dict: + """R1 [INV]: one-in-flight-per-trainer residence on the grad path. + + Per-trainer dispatch->commit intervals must NOT overlap (a trainer is + re-pickable only after its update commits). Real satisfies this by channel + construction (~0%); sim must model it (commit-then-carry + slot hold). A + non-zero sim fraction with real ~0 is the residence bug. Checked per mode -- + both must sit under tol_frac. + """ + rc, sc = _fwd_cadence_cycles(real), _fwd_cadence_cycles(sim) + r_frac, r_pairs, r_n = _overlap_fraction(rc) + s_frac, s_pairs, s_n = _overlap_fraction(sc) + if r_pairs == 0 and s_pairs == 0: + return {"ok": True, "tier": "INV", "status": "SKIP", + "note": "no contributor_intervals with >=2 contributions per " + "trainer (field absent, non-fwdllm run, or no re-selection)"} + ok = r_frac <= tol_frac and s_frac <= tol_frac + return { + "ok": ok, + "tier": "INV", + "real_overlap_frac": round(r_frac, 4), + "sim_overlap_frac": round(s_frac, 4), + "tol_frac": tol_frac, + "n_real_pairs": r_pairs, + "n_sim_pairs": s_pairs, + "n_real_intervals": r_n, + "n_sim_intervals": s_n, + "interpretation": ( + f"real {r_frac:.1%} / sim {s_frac:.1%} of same-trainer intervals " + f"overlap a prior one; >0 = re-dispatched while still in flight " + f"(residence violation)." + ), + } + + +def _forward_passes(trainers: dict) -> int: + """Total forward passes = trainer_round events across all trainers.""" + return sum(len(t.get("trainer_round", []) or []) for t in trainers.values()) + + +def _committed_grads(agg: dict) -> int: + """Total committed grads = contributors summed over committed cycles.""" + total = 0 + for e in _fwd_cadence_cycles(agg): + ivs = e.get("contributor_intervals") + if ivs is not None: + total += len(ivs) + else: + total += len(e.get("contributing_trainers") or []) + return total + + +def compute_conservation_parity(real: dict, sim: dict, + real_trainers: dict, sim_trainers: dict, + ratio_tol: float = 0.25) -> dict: + """W1 [DIAG]: compute-conservation — sim must not WASTE forward passes. + + `trainer_round` counts forward-pass STARTS, so forward/commit ~= 1 plus an + in-flight-at-stop + stale-reject tail. W1 catches the residence bug where sim + RE-DISPATCHES dropped grads and thus does far MORE forward passes per commit + than real (the 2x-recompute). ASYMMETRIC: only a sim EXCESS over real is a + violation -- a live async real system accrues a larger in-flight START tail + than the clock-gated sim, so real > sim is expected. Localizes to R1; R1==0 + + received==committed is the precise residence signal, W1 the coarse tell. + """ + r_fwd, s_fwd = _forward_passes(real_trainers), _forward_passes(sim_trainers) + r_com, s_com = _committed_grads(real), _committed_grads(sim) + if r_com == 0 or s_com == 0 or r_fwd == 0 or s_fwd == 0: + return {"ok": True, "tier": "DIAG", "status": "SKIP", + "note": "no forward passes or committed grads (non-fwdllm run " + "or telemetry absent)"} + r_ratio = r_fwd / r_com + s_ratio = s_fwd / s_com + # signed excess of sim over real (positive = sim wastes more compute) + excess = (s_ratio - r_ratio) / max(r_ratio, s_ratio) + direction = ("sim OVER-computes (recompute waste — check R1)" if excess > ratio_tol + else "sim under-computes (async start-tail; benign for compute waste)" + if excess < -ratio_tol else "matched") + return { + "ok": excess <= ratio_tol, + "tier": "DIAG", + "real_forward_passes": r_fwd, + "sim_forward_passes": s_fwd, + "real_committed": r_com, + "sim_committed": s_com, + "real_fwd_per_commit": round(r_ratio, 3), + "sim_fwd_per_commit": round(s_ratio, 3), + "sim_excess_rel": round(excess, 3), + "ratio_tol": ratio_tol, + "interpretation": ( + f"real {r_ratio:.2f} vs sim {s_ratio:.2f} forward passes per commit; " + f"{direction}." + ), + } + + # ═══════════════════════════════════════════════════════════════════ # §4 Consolidated run_all_parity (extended) # ═══════════════════════════════════════════════════════════════════ @@ -3435,7 +4383,8 @@ def run_all_parity(real_agg: dict, sim_agg: dict, rounds_cap: Optional[int] = None, budget_s: Optional[float] = None, real_ground_truth: Optional[dict] = None, - sim_ground_truth: Optional[dict] = None) -> dict: + sim_ground_truth: Optional[dict] = None, + max_bin: Optional[int] = None) -> dict: """Run the full parity + invariant battery; returns {name: result_dict}. Ordered HIGH → MID → LOW so coarse failures surface first: @@ -3466,6 +4415,8 @@ def run_all_parity(real_agg: dict, sim_agg: dict, results["overlap_factor"] = overlap_factor(real_agg, sim_agg) results["per_round_advance"] = per_round_advance_parity(real_agg, sim_agg) results["throughput"] = throughput_parity(real_agg, sim_agg) + results["wall_disparity"] = wall_disparity(real_agg, sim_agg) + results["sim_speedup"] = sim_speedup(real_agg, sim_agg) # ── Stage 2 Availability ── results["avail_composition"] = avail_composition_parity(real_agg, sim_agg) @@ -3508,6 +4459,7 @@ def run_all_parity(real_agg: dict, sim_agg: dict, results["trainer_phase"] = trainer_phase_parity(real_trainers, sim_trainers) results["gpu_budget_real"] = gpu_budget_ok(real_trainers) results["gpu_budget_sim"] = gpu_budget_ok(sim_trainers) + results["timing_overrun"] = timing_overrun(real_trainers, sim_trainers) results["sim_send_ts"] = sim_send_ts_ok(real_trainers, sim_trainers) # ── Stage 5 Update return & ordering ── @@ -3525,6 +4477,30 @@ def run_all_parity(real_agg: dict, sim_agg: dict, results["aggregation_sequence"] = aggregation_sequence_parity( real_agg, sim_agg, max_rounds) + # ── Stage 6'/3'/7' FwdLLM variance-cadence layer (PARITY.md §F.4) ── + # Pure functions over the per-cycle agg_round series; SKIP cleanly on + # non-fwdllm runs (no cadence fields emitted). V/G rungs feed off Stage-5 + # ordering + Stage-1 clock; DK rungs are inert unless DynamicKC is enabled. + results["cohort_sequence"] = cohort_sequence_parity(real_agg, sim_agg, max_bin=max_bin) + results["v1_iter_per_data_id"] = iters_per_data_id_parity(real_agg, sim_agg, max_bin=max_bin) + results["v2_var_trajectory"] = var_trajectory_parity(real_agg, sim_agg, max_bin=max_bin) + results["v3_cached_v_pool"] = cached_v_pool_parity(real_agg, sim_agg, max_bin=max_bin) + results["v4_force_commit_rate"] = force_commit_rate_parity(real_agg, sim_agg, max_bin=max_bin) + results["v5_variance_pass_ratio"] = variance_pass_ratio_parity(real_agg, sim_agg, max_bin=max_bin) + results["dk1_agg_goal_trajectory"] = agg_goal_trajectory_parity(real_agg, sim_agg) + results["dk2_dynamic_c"] = dynamic_c_trajectory_parity(real_agg, sim_agg) + results["dk3_eligible_ends_metric"] = eligible_ends_metric_parity(real_agg, sim_agg) + results["g1_grad_norm"] = grad_norm_parity(real_agg, sim_agg) + results["g2_grad_pool_size"] = grad_pool_size_parity(real_agg, sim_agg) + + # ── Stage 3' FwdLLM async residence (R1/W1, §L.3) ── + # R1 is the finest residence check (per-trainer interval overlap); W1 is the + # coarse compute-conservation tell that feeds V1/K2. Both SKIP cleanly when + # contributor_intervals is absent (sync baselines / non-fwdllm runs). + results["r1_inflight_overlap"] = inflight_overlap_parity(real_agg, sim_agg) + results["w1_compute_conservation"] = compute_conservation_parity( + real_agg, sim_agg, real_trainers, sim_trainers) + # ── Stage 7 Statistical utility ── results["utility"] = utility_parity(real_agg, sim_agg) @@ -3567,6 +4543,8 @@ def run_all_parity(real_agg: dict, sim_agg: dict, "overlap_factor": {"stage": 1, "role": "DIAG", "deps": ("trainer_speed", "sim_commit_monotone")}, "per_round_advance": {"stage": 1, "role": "EMERGENT", "deps": ("overhead_residual",)}, "throughput": {"stage": 1, "role": "EMERGENT", "deps": ("per_round_advance",)}, + "wall_disparity": {"stage": 1, "role": "DIAG", "deps": ("throughput",)}, + "sim_speedup": {"stage": 1, "role": "DIAG", "deps": ("sim_rate",)}, # ── Stage 2 Availability ── "avail_composition": {"stage": 2, "role": "MECHANISM", "deps": ()}, "eligibility": {"stage": 2, "role": "MECHANISM", "deps": ("avail_composition",)}, @@ -3605,6 +4583,7 @@ def run_all_parity(real_agg: dict, sim_agg: dict, "trainer_phase": {"stage": 4, "role": "DIAG", "deps": ()}, "gpu_budget_real": {"stage": 4, "role": "MECHANISM", "deps": ("training_budget",)}, "gpu_budget_sim": {"stage": 4, "role": "MECHANISM", "deps": ("training_budget",)}, + "timing_overrun": {"stage": 4, "role": "DIAG", "deps": ("gpu_budget_real", "gpu_budget_sim")}, "sim_send_ts": {"stage": 4, "role": "CONTROL", "deps": ("vclock_telemetry",)}, # ── Stage 5 Update return & ordering ── "inter_arrival_order": {"stage": 5, "role": "MECHANISM", "deps": ("per_round_advance", "selection_detail")}, @@ -3617,7 +4596,27 @@ def run_all_parity(real_agg: dict, sim_agg: dict, "withheld_delivery": {"stage": 6, "role": "DIAG", "deps": ("staleness", "abandon_timeout")}, "commit_promptness": {"stage": 6, "role": "CONTROL", "deps": ("withheld_delivery",)}, "aggregation_sequence": {"stage": 6, "role": "EMERGENT", "deps": ("participation", "inter_arrival_order")}, + "cohort_sequence": {"stage": 6, "role": "EMERGENT", "deps": ("participation", "inter_arrival_order", "r1_inflight_overlap")}, "first_divergence_summary": {"stage": 6, "role": "DIAG", "deps": ()}, + # ── Stage 3' FwdLLM async residence (R1/W1, simulate_fwdllm.md §L.3) ── + # R1 is the residence INV; W1 the compute-conservation DIAG that first flags a + # violation and localizes to R1. V1's cadence divergence is DOWNSTREAM of R1 + # (a residence violation changes the contributing set/order), so V1 deps on it. + "r1_inflight_overlap": {"stage": 3, "role": "MECHANISM", "deps": ("participation",)}, + "w1_compute_conservation": {"stage": 3, "role": "DIAG", "deps": ("r1_inflight_overlap",)}, + # ── Stage 6' FwdLLM variance-gated aggregation cadence (PARITY.md §F.4) ── + "v1_iter_per_data_id": {"stage": 6, "role": "MECHANISM", "deps": ("inter_arrival_order", "r1_inflight_overlap")}, + "v2_var_trajectory": {"stage": 6, "role": "MECHANISM", "deps": ("v1_iter_per_data_id",)}, + "v3_cached_v_pool": {"stage": 6, "role": "DIAG", "deps": ("v1_iter_per_data_id",)}, + "v4_force_commit_rate": {"stage": 6, "role": "MECHANISM", "deps": ("v1_iter_per_data_id",)}, + "v5_variance_pass_ratio": {"stage": 6, "role": "EMERGENT", "deps": ("v1_iter_per_data_id", "v2_var_trajectory")}, + # ── Stage 3' Dynamic K/C trajectory (inert unless DynamicKC enabled) ── + "dk1_agg_goal_trajectory": {"stage": 3, "role": "MECHANISM", "deps": ("v5_variance_pass_ratio",)}, + "dk2_dynamic_c": {"stage": 3, "role": "MECHANISM", "deps": ("v5_variance_pass_ratio",)}, + "dk3_eligible_ends_metric": {"stage": 3, "role": "CONTROL", "deps": ("avail_composition",)}, + # ── Stage 7' Forward-gradient quality ── + "g1_grad_norm": {"stage": 7, "role": "EMERGENT", "deps": ("selection_detail",)}, + "g2_grad_pool_size": {"stage": 7, "role": "EMERGENT", "deps": ("v1_iter_per_data_id", "dk1_agg_goal_trajectory")}, # ── Stage 7 Statistical utility ── "utility": {"stage": 7, "role": "EMERGENT", "deps": ("participation", "phase_gpu_compute", "staleness")}, # ── Stage 8 Emergent outcomes ── diff --git a/lib/python/examples/async_cifar10/scripts/parity/cli.py b/lib/python/examples/async_cifar10/scripts/parity/cli.py index 7a4b2bbc6..0f78e7a90 100644 --- a/lib/python/examples/async_cifar10/scripts/parity/cli.py +++ b/lib/python/examples/async_cifar10/scripts/parity/cli.py @@ -45,7 +45,8 @@ def _run_pair(real_dir: str, sim_dir: str, agg_goal: int, rounds_cap, budget_s, strict: bool, lenient: bool, json_out, plot_out, - real_label: str = "", sim_label: str = "") -> bool: + real_label: str = "", sim_label: str = "", + max_bin=None) -> bool: """Load, check, and report one real/sim pair. Returns True if passed.""" # Import here to avoid circular import issues when run as __main__ import sys as _sys @@ -87,6 +88,7 @@ def _run_pair(real_dir: str, sim_dir: str, budget_s=budget_s, real_ground_truth=real_ground_truth, sim_ground_truth=sim_ground_truth, + max_bin=max_bin, ) # Add first_divergence as a diagnostic summary entry (always ok — index=0 is expected for async) @@ -127,6 +129,9 @@ def main() -> None: help="rounds cap from config (enables K9 truncation check)") parser.add_argument("--budget-s", type=float, default=None, help="max_experiment_runtime_s / sim_wall_ceiling_s (enables K5/K9)") + parser.add_argument("--max-bin", type=int, default=None, + help="restrict fwdllm cadence rungs (cohort_sequence/V*) to " + "cycle_data_id <= MAX_BIN (first-data-bin logical parity)") parser.add_argument("--strict", action="store_true", help="Treat WARN as FAIL") parser.add_argument("--lenient", action="store_true", @@ -194,6 +199,7 @@ def main() -> None: lenient=args.lenient, json_out=json_out, plot_out=plot_out, + max_bin=args.max_bin, ) # Collect for roll-up (re-load results from JSON if written) import json as _json @@ -224,6 +230,7 @@ def main() -> None: lenient=args.lenient, json_out=args.json_out, plot_out=args.plot_out, + max_bin=args.max_bin, ) sys.exit(0 if ok else 1) diff --git a/lib/python/examples/async_cifar10/scripts/parity/validate_real.py b/lib/python/examples/async_cifar10/scripts/parity/validate_real.py index 48d636ab3..630cce2d3 100644 --- a/lib/python/examples/async_cifar10/scripts/parity/validate_real.py +++ b/lib/python/examples/async_cifar10/scripts/parity/validate_real.py @@ -45,7 +45,14 @@ if str(_HERE.parent) not in sys.path: sys.path.insert(0, str(_HERE.parent)) -from parity.checks import load_run_dir, agg_goal_cycles_ok # noqa: E402 +from parity.checks import ( # noqa: E402 + load_run_dir, + agg_goal_cycles_ok, + _fwd_cadence_cycles, + _overlap_fraction, + _forward_passes, + _committed_grads, +) # Allow a tiny fraction of telemetry edge cases (interleaving/late events) before @@ -250,6 +257,35 @@ def check_aggregation(agg: dict) -> dict: } +def check_grad_residence(agg: dict, trainer: dict) -> dict: + """FwdLLM async residence (R1) + compute-conservation (W1) on the REAL side — + the gate blocking the sim mechanism change: if real itself overlaps (R1 > ~0), + fix real FIRST, never tune sim to it. W1 is reported for context. SKIPs on a + non-fwdllm real run (no intervals). + """ + cycles = _fwd_cadence_cycles(agg) + have_intervals = any(e.get("contributor_intervals") for e in cycles) + if not have_intervals: + return {"ok": True, "status": "SKIP", + "note": "no contributor_intervals (non-fwdllm real run)"} + frac, n_pairs, n_intervals = _overlap_fraction(cycles) + fwd = _forward_passes(trainer) + com = _committed_grads(agg) + ratio = (fwd / com) if com else None + # R1 is the hard invariant; W1 ratio is reported (a huge real ratio would + # itself be suspicious, so bound it generously — real overcommit is small). + ok = frac <= _VIOLATION_FRAC_TOL + return { + "ok": ok, + "r1_overlap_frac": round(frac, 4), + "r1_overlap_pairs": n_pairs, + "n_intervals": n_intervals, + "w1_forward_passes": fwd, + "w1_committed": com, + "w1_fwd_per_commit": round(ratio, 3) if ratio is not None else None, + } + + def validate_real(run_dir: str) -> bool: agg, trainer = load_run_dir(run_dir) tiv = _train_intervals(trainer) @@ -264,9 +300,11 @@ def validate_real(run_dir: str) -> bool: "concurrency": check_concurrency(agg, tiv), "selection": check_selection(agg), "aggregation": check_aggregation(agg), + "grad_residence": check_grad_residence(agg, trainer), } all_ok = True for name, r in results.items(): + r.pop("status", None) # SKIP marker — informational, not printed as a field ok = r.pop("ok") all_ok = all_ok and ok tag = "[OK] " if ok else "[FAIL]" diff --git a/lib/python/examples/async_cifar10/scripts/parity_checks.py b/lib/python/examples/async_cifar10/scripts/parity_checks.py index a82adaa30..258f6494d 100644 --- a/lib/python/examples/async_cifar10/scripts/parity_checks.py +++ b/lib/python/examples/async_cifar10/scripts/parity_checks.py @@ -33,6 +33,9 @@ # selection constants DETERMINISTIC_SELECTORS, _selector_name, + _has_cohort_counts, + _full_cohort_selection, + _selection_is_deterministic, # §3.B selection selection_parity, # §3.D updates @@ -52,6 +55,7 @@ # §3.C sim invariants sim_send_ts_ok, gpu_budget_ok, + timing_overrun, trainer_phase_parity, # Stage 0 / 1 / 2 / 4 / 8 additions (causal ladder) field_coverage, @@ -78,9 +82,30 @@ throughput_parity, per_round_advance_parity, overlap_factor, + wall_disparity, + sim_speedup, total_commits_parity, terminal_state_parity, budget_not_cap, + # §F fwdllm variance-cadence layer (V/DK/G rungs) + _iters_per_data_id, + cohort_sequence_parity, + iters_per_data_id_parity, + var_trajectory_parity, + cached_v_pool_parity, + force_commit_rate_parity, + variance_pass_ratio_parity, + agg_goal_trajectory_parity, + dynamic_c_trajectory_parity, + eligible_ends_metric_parity, + grad_norm_parity, + grad_pool_size_parity, + # fwdllm async residence (R1/W1, simulate_fwdllm.md §L.3) + _overlap_fraction, + _forward_passes, + _committed_grads, + inflight_overlap_parity, + compute_conservation_parity, # overall run_all_parity, overall_verdict, diff --git a/lib/python/examples/async_cifar10/scripts/run_felix_streaming.sh b/lib/python/examples/async_cifar10/scripts/run_felix_streaming.sh index d639ed5a8..a34dbfc12 100755 --- a/lib/python/examples/async_cifar10/scripts/run_felix_streaming.sh +++ b/lib/python/examples/async_cifar10/scripts/run_felix_streaming.sh @@ -14,26 +14,14 @@ # See docs/EXPERIMENT_felix_streaming.md (Distributed execution & pooling). set -u -# --- robust conda activation (mirrors debug_run.sh) --- -ENVNAME="${FLAME_CONDA_ENV:-dg_flame}" -CB="" -if command -v conda >/dev/null 2>&1; then - CB="$(conda info --base 2>/dev/null)" -elif [ -n "${CONDA_EXE:-}" ]; then - CB="$(dirname "$(dirname "$CONDA_EXE")")" -fi -if [ -z "$CB" ] || [ ! -f "$CB/etc/profile.d/conda.sh" ]; then - for c in "$HOME/miniconda3" "$HOME/anaconda3" /opt/conda; do - [ -f "$c/etc/profile.d/conda.sh" ] && CB="$c" && break - done -fi -[ -f "$CB/etc/profile.d/conda.sh" ] || { echo "ERROR: conda not found; activate '$ENVNAME' yourself." >&2; exit 1; } -source "$CB/etc/profile.d/conda.sh" -conda activate "$ENVNAME" || { echo "ERROR: conda activate $ENVNAME failed" >&2; exit 1; } -echo "conda: base=$CB env=$ENVNAME python=$(which python)" - EX="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # async_cifar10 example dir REPO="$(cd "$EX/../../../.." && pwd)" # repo root + +# shared harness: robust conda activation (same block debug_run.sh uses). +# shellcheck source=../../scripts/expt_runner.sh +source "$REPO/lib/python/examples/scripts/expt_runner.sh" +expt_activate_conda dg_flame + cd "$EX" || exit 1 SCR=expt_scripts_2026 CFG="$SCR/n50_alpha0.1_syn0_stream_unif_sim.yaml" diff --git a/lib/python/examples/async_cifar10/scripts/smoke_suite.sh b/lib/python/examples/async_cifar10/scripts/smoke_suite.sh index a1fe02753..dd0c27538 100755 --- a/lib/python/examples/async_cifar10/scripts/smoke_suite.sh +++ b/lib/python/examples/async_cifar10/scripts/smoke_suite.sh @@ -55,6 +55,12 @@ EX_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" # async_cifar10/ LIB_DIR="$(cd "$EX_DIR/../.." && pwd)" # lib/python/ DEBUG_RUN="$SCRIPT_DIR/debug_run.sh" +# Shared harness: the timeout + process-group-kill primitive (expt_timed_run), +# moved here from inline in _run_baseline. No conda activation here -- each +# debug_run.sh child activates its own env. +# shellcheck source=../../scripts/expt_runner.sh +source "$LIB_DIR/examples/scripts/expt_runner.sh" + # ── Defaults ───────────────────────────────────────────────────────────────── RUNTIME_SYN0_S=900 RUNTIME_SYN20_S=1800 @@ -202,64 +208,14 @@ _run_baseline() { local ts_marker="$run_dir/.ts_start" touch "$ts_marker" - # set -m (job control) forces bash to assign PGID = runner_pid to the - # background job regardless of whether the suite is running interactively - # or not. All descendants inherit that PGID (flame's spawner.py uses plain - # subprocess.Popen with no start_new_session/os.setsid), so - # kill -TERM/-KILL on -$runner_pid reliably reaches the whole process tree. - # Rationale: setsid forks when the calling process is already a pg-leader - # (interactive terminals do this), making runner_pid point to a dead parent - # instead of the actual session leader → kill misses the tree entirely. - set -m - env FLAME_LOGDIR="$run_dir" bash "$DEBUG_RUN" "$@" \ - >"$run_dir/shell.log" 2>&1 & - local runner_pid=$! - set +m - - local deadline=$(( ts_start + wall_timeout )) - local timed_out=0 - local _ela _kill_in _prun _pdone # ticker temporaries - while kill -0 "$runner_pid" 2>/dev/null; do - sleep 5 - _ela=$(( $(date +%s) - ts_start )) - _kill_in=$(( deadline - $(date +%s) )); [[ "$_kill_in" -lt 0 ]] && _kill_in=0 - _prun=0; [[ "$runtime_s" -gt 0 ]] && _prun=$(( _ela * 100 / runtime_s )) - [[ "$_prun" -gt 100 ]] && _prun=100 - _pdone=0; [[ "$TOTAL_RUNS" -gt 0 ]] && _pdone=$(( COMPLETED_RUNS * 100 / TOTAL_RUNS )) - printf '\r %-52s %4ds/%-4ds(%3d%%) kill in %4ds | %d/%d done(%d%%) ' \ - "[$label]" "$_ela" "$runtime_s" "$_prun" "$_kill_in" \ - "$COMPLETED_RUNS" "$TOTAL_RUNS" "$_pdone" >&2 - if [[ "$(date +%s)" -ge "$deadline" ]]; then - printf '\n' >&2 - _log " [$label] TIMEOUT after ${wall_timeout}s — killing process group $runner_pid" - # SIGTERM first: lets ExperimentRunner._signal_handler call _cleanup() - # (terminate_all trainers + terminate aggregator). 20s grace lets Python - # flush open files and release MQTT connections before the hard kill. - kill -TERM -"$runner_pid" 2>/dev/null || true - sleep 20 - # SIGKILL for anything that survived (hung GPU op, stuck MQTT recv). - kill -KILL -"$runner_pid" 2>/dev/null || true - wait "$runner_pid" 2>/dev/null - timed_out=1 - break - fi - done - printf '\n' >&2 - [[ "$timed_out" == "0" ]] && wait "$runner_pid" 2>/dev/null + # Launch under the shared timeout+process-group-kill primitive: runs the child in + # its own PGID, SIGTERM->SIGKILL the whole tree on overrun, sweeps stragglers, then + # waits KILL_SETTLE_S for GPU memory to drain. Returns 124 on timeout. + EXPT_KILL_SETTLE_S="$KILL_SETTLE_S" \ + expt_timed_run "$label" "$runtime_s" "$TIMEOUT_BUFFER_S" "$run_dir/shell.log" -- \ + env FLAME_LOGDIR="$run_dir" bash "$DEBUG_RUN" "$@" local run_rc=$? - - # ── Post-kill cleanup ───────────────────────────────────────────────────── - # When SIGKILL fires, run_experiment_batch's finally block (_sweep_stragglers) - # is in the killed group and may not complete. Replicate it here: hard-kill - # any surviving trainer/aggregator processes and wait for GPU memory to drain - # so the next run starts from a clean slate. - if [[ "$timed_out" == "1" ]]; then - _log " [$label] post-kill sweep: clearing straggler trainer/aggregator processes" - pkill -9 -f "trainer/pytorch/main.py" 2>/dev/null || true - pkill -9 -f "aggregator/pytorch/main_" 2>/dev/null || true - _log " [$label] waiting ${KILL_SETTLE_S}s for GPU memory to drain before next run" - sleep "$KILL_SETTLE_S" - fi + local timed_out=0; [[ "$run_rc" == "124" ]] && timed_out=1 local elapsed=$(( $(date +%s) - ts_start )) diff --git a/lib/python/examples/fwdllm/EXPERIMENTS.md b/lib/python/examples/fwdllm/EXPERIMENTS.md new file mode 100644 index 000000000..38094a710 --- /dev/null +++ b/lib/python/examples/fwdllm/EXPERIMENTS.md @@ -0,0 +1,537 @@ +# FLUXTUNE vs FWDLLM / FWDLLM_PLUS — experiment design (living doc) + +**Status:** N=100 α=1 runs landed (2×2 opt ablation + baseline comparison — charter). E1 headline holds on +**peak** accuracy; runs do not yet *hold* the minimum (Issue I-1, **root-caused** → `fluxtune_contributions.md` +§8, next = S1 server optimizer). Tooling validated end-to-end; `main` condition gated. +**Owner:** dgarg39 · **Branch:** `dg/fwdllm_sim_unavail` + +Human design doc. Its machine-readable twin [`experiments.yaml`](experiments.yaml) is what the tooling +**consumes** — `run_sequential.sh` launches the run-set from it, `compare_baselines.py` reads which runs feed +which metric. Keep them in sync. On disagreement, `experiments.yaml` is source of truth for **what ran**; this +doc for **what we intend and why**. + +**Paper ⇄ code reconciliation:** [`EXPTS_CHARTER.md`](EXPTS_CHARTER.md) reconciles this doc with the paper +draft [`05-evaluation.tex`](05-evaluation.tex). The **run ledger** (which log on which node feeds which +result) is §10. + +Related: [`simulate_fwdllm.md`](simulate_fwdllm.md) (principles), [`PARITY_LOGICAL_TASKS.md`](PARITY_LOGICAL_TASKS.md) +(real↔sim parity), [`fluxtune_contributions.md`](fluxtune_contributions.md) (systems/ML contributions, incl. +the memory/inference-only-NPU thesis — a motivation/design claim, not an eval experiment). + +--- + +## 0. Architecture: runs vs. analyses (the anti-redundancy backbone) + +The organizing rule that makes cross-baseline numbers reusable AND trustworthy: + +- **Run-set** — the expensive GPU work. Produced **once** per `(baseline × condition)`. Each run dir + self-describes via `snapshot.yaml` + `telemetry/*.jsonl`. Governed by the **convergence stop** (§2) with + `max_runtime_s` / `max_data_id_progress` as safety caps. +- **Analyses** — cheap pure functions over `telemetry/`. **All five experiments are views over the same + run-set.** Expt 1 *defines* the runs; Expts 2–5 add **zero** GPU runs — they are reducers. +- **Mix-guard** — `compare_baselines.py` discovers the latest run per baseline (exact-token regex, so + `fwdllm` never captures `fwdllm_plus`) and **warns if the baselines' shared axes (N/partition/trace) + disagree** — the post-hoc twin of the launch-time `condition_fp`. The fingerprint proves the runs launched + identically; the mix-guard proves the compared dirs still agree. + +Consequence: "Experiment 6" later = a reducer, not new hardware. Re-deriving a metric after a telemetry fix = +re-run the analysis, not the experiment. + +``` +experiments.yaml ──> run_sequential.sh ──> experiments/run_*/telemetry/*.jsonl + │ │ + └────────────> compare_baselines.py <──────────┘ (reducers → one table/CSV + overlay plots) +``` + +--- + +## 1. Baselines (what distinguishes them) + +Substance lives in `_metadata/baselines.yaml`; run YAMLs only pick `baseline:` + a few overrides. + +| Knob | **fwdllm** | **fwdllm_plus** | **fluxtune** | +|---|---|---|---| +| sync / async | sync | sync | **async** (`fedbuff`) | +| selector | `random` | `random` | `async_oort` | +| availability tracking | unaware | **ORACULAR** | client_notify (self-report) | +| reselect | per-round | per-iteration | continuous (async) | +| staleness | `exact` | `round_data_id` | `none` (down-weight) | +| JVP perturbation scoring | off | off | **on** (`jvp_perf_opt`) | +| native `agg_goal` | 10 | 10 | 3 | + +**Signed-off experiment override:** the `main` condition matches **`agg_goal=10` across all three** (fluxtune +uses `agg_goal=10` with `C=30`, NOT its native fedbuff 3) — same aggregation batch for a fair head-to-head. A +pre-flight check enforces the match. See §7.0. + +**Design risk (Phase 2, mobiperf):** fwdllm_plus's sync barrier (`agg_goal ≥ available trainers`) **cannot +assemble** under unavailability — the pre-flight gate **blocks** this. A fair Expt-1 comparison under a +`mobiperf_*` trace needs a deliberate answer for what fwdllm_plus does. Phase 1 (the signed-off `main` +condition) uses `syn_0` (100% available) where all three complete; the mobiperf policy is a Phase-2 decision. + +--- + +## 2. Convergence stop (run termination) — runner-side watcher + +**A run terminates when the aggregator has completed `W` consecutive data bins all remaining ≥ target +accuracy `τ`.** Defaults `W=20`, `τ` per-condition in `experiments.yaml`. + +- **Signal:** `agg_eval` telemetry events. Verified cadence: **exactly one `agg_eval` per data bin**, emitted + at bin completion (`data_id` k → the eval at its final `iteration_per_data_id`). If a bin emits multiple, + the watcher takes the **last**. +- **Mechanism:** a poller in the `expt_launch` ticker loop tails the aggregator telemetry, maintains the + per-`data_id` representative accuracy, and fires when the last `W` distinct completed bins are all ≥ `τ` → + clean process-group termination. +- **Health verdicts:** `CONVERGED` (window satisfied), `STALLED` (early-terminated, not learning), + `DID_NOT_CONVERGE` (hit wall ceiling still learning), plus `COMPLETED` / `CRASH` / `WALL_CEILING`. +- **Output:** per-run `converge.json` = time-to-converge in **wall + vclock (sim) + data_id + round** (or + `stall.json` on a stall). *This is the Expt-1 metric captured at the source*, not reconstructed. +- **Wall ceiling = 48h.** Convergence runs are accuracy-governed, so `max_runtime_s` defaults to **172800s + (48h)** whenever `--target-acc` is set. `max_data_id_progress` also bounds it. +- **Stall guard (early-out).** Terminate BEFORE the 48h ceiling if clearly not learning: **no PROGRESS within + `stall_window_s` (default 7200s = 2h)** → verdict `STALLED`. Any progress resets the clock; convergence is + checked first, so a just-converged run is never called stalled. `stall_window_s=0` disables it. **Progress + is set by `--stall-on` (default `either`):** + - `acc` — best accuracy gained ≥ `stall_min_delta` (default 0.01 = **1% absolute**; accuracy ∈ [0,1]). + - `loss` — best (running-min) test-loss dropped ≥ `loss_min_rel_delta` (default 0.01 = **1% relative**; loss + is unbounded/scale-dependent, so the bar is *fractional vs the running-best*, not absolute). + - `either` — reset if **either** fired. Motivating case: a run can plateau in accuracy while test-loss keeps + falling (still learning) — `either`/`loss` keeps it alive; `acc` would kill it. + + Both signals use the running-best (max acc / min loss), so a single noisy eval can neither reset the clock + nor fake progress. *Validated: flat acc + flat loss → `STALLED [either]`; flat acc + steadily-falling loss → + clock resets on the loss signal.* + +New flags on `run_sequential.sh`: `--target-acc τ`, `--converge-window W`, `--stall-window-s S` (or the hours +alias `--stall-window-h H`, e.g. `6` ⇒ 21600s), `--stall-min-delta D`, `--stall-on acc|loss|either`, +`--loss-min-rel-delta R` (all also settable from the registry via `--run-set`). **The CLI value overrides the +registry** — a run that stalls too eagerly (flat accuracy but loss still falling) can be re-launched with a +wider idle window, e.g. `--stall-window-h 6`, without editing `experiments.yaml`. The stall window is part of +`condition_fp` but is a *termination-policy* knob, not a scientific-condition one: the cross-baseline +mix-guard compares only N/partition/trace, so widening it for one baseline's re-run does not taint the Expt-1 +comparison. + +--- + +## 3. Telemetry provenance + +Most metrics are **already emitted** (map below). Two instrumentation additions were made, both "route an +already-known quantity to `telemetry.emit`", off when `FLAME_TELEMETRY_DIR` unset (zero cost elsewhere). +**Both implemented and validated** (2026-07-06): + +- **(WS3-a) Network bytes + message counts** — a dedicated **`comm` event** + (`flame/telemetry/events.py:build_comm`), emitted **both directions** with `direction` + (`agg_to_trainer`/`trainer_to_agg`), `size_bytes`, `payload_kind` (`weights`/`var_bad`/`gradients`), + `n_tensors`, tagged with round/data_id. Sizes reuse the debug-log values: + - trainer upload at [`fwdllm_trainer.py`](../../flame/mode/horizontal/syncfl/fwdllm_trainer.py) `_send_weights`; + - aggregator dispatch at **both** send sites — sync `_distribute_weights` **and** async + `_distribute_weights_async` (the async site is fluxtune's; instrumenting only sync silently produced zero + agg→trainer events until fixed). +- **(WS3-b) Perturbations / forward-passes per client** — module-global counters in + [`fwdgrad_utils.py`](trainer/forward_training/fwdgrad_utils.py) (`_FWD_PASSES` / `_JVP_EVALS`, incremented + in `calculate_jvp*`; each trainer is its own process ⇒ per-client), read into `trainer_round` as + `forward_passes_iter/total` + `perturbations_iter/total`. Gives Expt 3 a **hardware-independent** compute + denominator immune to the 8-GPU contention confound. (Validated: 10 perturbations = 20 forward passes/iter.) + +**Aggregator "active GPU time" = aggregator compute wall-time** — measured from the NON-OVERLAPPING `agg_round` +phase fields (`aggregate_fedavg_s` + `eval_s`), **not** summed `step_timing` (double-counts — `timer_decorator` +calls nest), **not** CUDA-event-isolated GPU time. CUDA-event timing only if a reviewer challenges the number. + +--- + +## 4. The five experiments (analyses over the one run-set) + +Metric logic lives in `expt_scripts/plotlib/reducers.py` (`load_run` → `RunResult`), consumed by +`compare_baselines.py` (the `expt1..5_*` table wrappers) and `plot_run.py`. Legend — **provenance**: `EMIT` +already in telemetry · `DERIVE` reducer over existing telemetry · `WS3` the instrumentation add (now emitted) +· `WS2` from the convergence watcher. + +> ⚠ **Reducer-audit findings (2026-07-07, tracked in [`EXPTS_CHARTER.md`](EXPTS_CHARTER.md) §2b):** +> **N3** — E1 time-to-τ is **reconstructed** from `agg_eval` (streak-over-window scan, over the loss-truncated +> series), it does **not** read `converge.json`; can silently diverge from the watcher's verdict. **N5** — E2 +> idle is only `1−busy_frac`; `mqtt_fetch_s` is emitted but unused, and `barrier_wait_s`/`drain_tail_s` read +> ≈0 in sim. **N6** — E5 sync sessions use `contributor_intervals` (dispatch→commit) for *all* baselines (the +> one-round-span method is unimplemented); `agg_round.contributing_trainers` is emitted but never consumed. +> Fix or re-scope before the claims land. + +### Experiment 1 — Time to target accuracy +> **Takeaway:** Fluxtune reaches target accuracy faster than FwdLLM and FwdLLM_Plus. +- **Config:** the `main` run-set. Baselines: all three. **Reuse:** none (this *defines* the runs). +- **Metrics reported:** + - Time to reach `τ` (the convergence event) — **wall, #rounds, #data_bins, #iterations**. *(WS2 intent; + currently DERIVE — see N3.)* **Virtual-clock is deferred** (real-mode runs emit no vclock; sim vclock is + unvalidated, `sim_rate≈0.50`) — add later if a validated sim lands. + - Maximum accuracy attained. *(DERIVE: max `agg_eval.test-accuracy`)* + +> ✅ **Observed (N=100 α=1, 2026-07-08) — supports the takeaway on PEAK accuracy + speed.** Peak test accuracy: +> **FluxTune (R4 full) 84.1% @ 3.9h** reaches target · FwdLLM++ 80.9% @ 7.4h (−3.1 from target, ~2× slower) · +> FwdLLM 30.3% (never learned). Within the fluxtune 2×2, peak rises with the opt ladder (R1 FluxTune-base 83.0 +> → R2 82.2 → R3 grad-aware 83.9 → R4 full 84.1; R1 = C1 guided-JVP base, **not FeLiX** — forward-mode LLM +> fine-tuning, only borrowing FeLiX's scalar agg rate). Plots (`e1_acc_vs_time.pdf`, both sets) mark each run's +> peak with a ★ + legend value against the 84% target line. +> ⚠ **Caveat (Issue I-1):** peak is **transient** — runs oscillate and collapse after the round-1 peak, +> **root-caused** as an undamped high-variance optimizer (NOT the once-suspected epoch bug; see charter I-1 + +> `fluxtune_contributions.md` §8, fix = server optimizer S1 / M2 below). E1 uses the round-1 peak, plots +> clipped there (`--cutoff-mode peak_acc`); the time-to-τ streak never fires (accuracy only grazes 84% amid +> oscillation). + +### Experiment 2 — Resource utilization (wait-time reduction) +> **Takeaway:** Fluxtune improves utilization by cutting wait times at trainers (primary, thousands) and the +> aggregator (secondary, single). +- **Reuse:** **Expt-1 run-set** (no new runs). +- **Approach:** **derived busy/idle time-fraction** (operator call — no hardware sampler). + - Trainer busy = Σ`gpu_compute_s` / (last−first ts); idle = `mqtt_fetch_s` + inter-round gaps. + - Aggregator busy = Σ(`aggregate_fedavg_s`+`eval_s`) / `agg_round.wall_elapsed_s` (non-overlapping phases, + NOT summed `step_timing`); wait fractions = Σ`barrier_wait_s` / Σ`drain_tail_s` over wall. +- **Metrics reported:** P50 / P90 / P99 busy-fraction **across trainers** over runtime; aggregator busy vs. + barrier-wait vs. drain fraction. *(DERIVE)* + +### Experiment 3 — Compute productivity (learning per unit compute) +> **Takeaway:** At resource-constrained clients, Fluxtune yields more learning per unit compute. +- **Reuse:** **Expt-1 run-set**. +- **Metrics reported:** `Δloss / cumulative compute`, against **two** compute denominators: + - **Forward passes** (perturbations): Σ per-client perturbation count — **PRIMARY / clean**. *(WS3-b — + hardware-independent: counts actual passes regardless of contention.)* + - GPU-seconds: Σ(trainer `gpu_compute_s`) + aggregator compute wall-time — **SECONDARY, confounded**. + *(DERIVE.)* ⚠ `gpu_compute_s` is **wall-time** GPU work; 100 trainers time-share **8 GPUs**, so contention + inflates it (~8–10 ms/pass clean → ~0.21 s/pass under load, ≈20×), and Fluxtune's ~30 concurrent clients + contend differently than the sync baselines' K=10 bursts — so this denominator measures scheduling + contention, not algorithmic compute. Report with the caveat; lean on the forward-pass denominator. + - `Δloss` = first `agg_eval.test-loss` − final `agg_eval.test-loss`, **in time order** (never keyed by + `data_id`, which cycles per round). *(EMIT)* + +> ⚠ **Observed (N=100 smoke, 2026-07-07) — does NOT yet support the takeaway.** Learning-per-compute ranks +> **FwdLLM++ > Fluxtune > FwdLLM** on *both* denominators (Δloss/GPU-h 0.055 vs 0.016 vs −0.014; Δloss/M-fwd +> **5.62 vs 1.27** vs −1.26). Fluxtune attains the most *total* Δloss (0.82) and the fastest wall-clock +> convergence (E1) but spends **~5–6× the forward-pass compute** to get there, so per-unit it is *less* +> efficient. GPU-h is confounded by 8-GPU contention (§7.1); the clean forward-pass denominator still favors +> FwdLLM++. +> **Why:** Fluxtune's async high-concurrency design does more *unproductive* compute — (i) fedbuff +> **staleness**: concurrent clients train on stale models and their updates are down-weighted, so part of the +> forward-pass compute yields little Δloss; (ii) it keeps ~30 clients busy speculatively. FwdLLM++ is +> synchronous with **oracular** availability, so every forward pass feeds a fresh, fully-weighted update on a +> client that will contribute. Concurrency buys wall-clock speed, not compute efficiency. +> **Optimize:** staleness-aware admission / adaptive concurrency + cutting JVP forward-pass overhead. (The +> per-resource-constrained-client framing still needs a per-client-normalized metric — a metric decision, not +> a plot bug.) + +### Experiment 4 — Data transmitted over the network +> **Takeaway:** Fluxtune incurs lower total data overhead despite more messages per round/iteration. +- **Reuse:** **Expt-1 run-set** (requires WS3-a telemetry present at run time). +- **Metrics reported:** total messages sent (each side); total bytes transmitted (each side); per-message size + distribution. *(WS3-a)* + +> ⚠ **Observed (N=100 smoke, 2026-07-07) — contradicts the takeaway.** Fluxtune transmits **~1.9× more total +> bytes** than FwdLLM++ (**146 vs 79 GB**) and ~1.6× more messages — it does *not* incur lower overhead here. +> Per-message size is identical (1.8 MB/upload), so the gap is message **count**, not payload. +> **Why:** (i) **Model distribution dominates** — FwdLLM++ (sync) sends the model once per round to the K +> selected clients (weights **9 GB**); Fluxtune (async) has ~30 clients *continuously re-pull* the latest +> global model as they finish and re-enlist → weights balloon to **88 GB**. (ii) **Uploads scale with +> iterations** — one gradient upload (1.8 MB) per trainer-iteration; Fluxtune runs ~1.6× more iterations +> (concurrency) → 58 vs 36 GB up. FwdLLM++'s synchronous rounds amortize both, and oracular selection avoids +> dispatching to non-contributors (it *does* pay a 33 GB method-specific `var_bad` payload, yet still totals +> less). **Same root cause as E3:** async concurrency does more total work. +> **Optimize:** delta/compressed model distribution on re-pull + staleness-aware throttling would cut +> Fluxtune's dominant weight-download term. + +> 🔧 **Single-contribution results — efficiency levers not yet on.** Contributions (charter B2): **C1** guided +> perturbations · **C2** dynamic K/C · **C3** gradient-aware aggregation. In `main`: **C1 active**; **C2 OFF** +> (static `agg_goal=10`/`C=30` for the agg_goal-matched head-to-head, §7.0; `dynamic_kc.enabled=false`); **C3** +> runs the fedbuff "new" staleness×utility *scalar* rate + `var≤0.3` gate — a borrowed FeLiX placeholder, not +> the intended gradient-aware rule (charter N2; `grad_aware`=Opt-3). C2/C3 target exactly the E2-E4 +> inefficiencies. **So E2/E3/E4 efficiency claims need a separate full-system run (C2 + real C3 on)** — +> enabling dynamic K/C breaks the agg_goal match. **E1 (speed + final accuracy) holds on C1 alone.** + +### Experiment 3-adjacent — C3 intelligent (gradient-aware) aggregation: design investigation +> **Status: NOT the intended contribution yet.** The active weighting (above) is a borrowed scalar rate. +> Fluxtune needs a **gradient-aware** rule. **Full design starter:** +> [`docs/aggregation_design.md`](docs/aggregation_design.md) (FedBuff→FeLiX→FluxTune regime, hypothesis, +> 5-axis design space, candidate schemes S0–S5). This section seeds that design (implement later, then move +> the feature doc to [`fluxtune_contributions.md`](fluxtune_contributions.md) and delete from here). +> **Dimensions to evaluate:** +> 1. **Staleness under iteration-based progression.** Staleness = `agg_model_version − trainer_version`; +> `_model_version` advances **per data-bin completion**, and a data-bin completes only when the **variance +> threshold is met** ([`fwdllm_aggregator.py:1422`](../../flame/mode/horizontal/syncfl/fwdllm_aggregator.py#L1422), +> [`fedbuff.py:194`](../../flame/optimizer/fedbuff.py#L194)) → staleness accrues at the *variance-gated +> data-bin rate, not wall-clock*. Hypothesis: staleness grows **slower** in Fluxtune than round-based FL. +> Quantify the effective staleness distribution vs a round-based baseline. +> 2. **Scalar rate vs gradient-aware combination.** Scalar-multiplying a *gradient* update (forward-mode JVP +> estimate) may be the wrong operator vs down-weighting *weights*. Explore direction-/variance-aware +> combination (weight by JVP magnitude / SNR / agreement with the running aggregate), not just staleness×loss. +> 3. **Interaction with C1 + the variance gate.** Updates already passed `var ≤ var_threshold`; does +> re-weighting by loss (`stat_utility`) double-count what the gate filtered? +> 4. **Interaction with C2 (dynamic K/C).** Concurrency C sets how many stale/in-flight updates coexist; +> aggregation rule and concurrency controller co-determine the wasted work E3/E4 measure. + +### Experiment 5 — Client training-session durations & participation +> **Takeaway:** Fluxtune's client sessions are much shorter than FwdLLM's. +- **Reuse:** **Expt-1 run-set**. +- **Active-session definition (precise, per operator):** + - **async (fluxtune):** selected → next reselection ≈ `dispatch_ts → commit_ts` + (`agg_round.contributor_intervals`). + - **sync (fwdllm / fwdllm_plus, round-based reselect):** one-round span from `trainer_round`/`agg_round` + timestamps. +- **Metrics reported:** + - P50 / P90 / P99 of session duration across clients + a histogram. *(DERIVE)* + - **Per-client participation counts at three granularities across baselines** — #rounds, #data_bins, + #iterations each client participated in. *(DERIVE: `trainer_round` + `agg_round.contributing_trainers`)* + +### Experiments M1–M2 — training-stability track (motivation / ablation) +> Context: the α=1 N=100 runs oscillate and single-class-collapse (charter I-1). The H0 diagnostic **refuted** +> data-class-bias as the cause (`fluxtune_contributions.md` §8, F11-F15) → the driver is the **undamped, +> high-variance forward-gradient optimizer**. These two experiments motivate and validate the fix. + +**M1 — Data-bin size × heterogeneity (the forward-mode bias/variance tradeoff).** +- **Hypothesis (bidirectional):** per-commit **variance** falls as bin/cohort size grows (favors *large* bins), + BUT a large, class-mixed bin **averages conflicting per-sample gradients → small mean-gradient magnitude → + the scalar JVP signal sinks below its noise floor and learning stalls** (favors *small* bins). ⇒ an + **α-dependent optimum**: too-small = oscillation (observed), too-large = no learning (recalled). +- **Design:** sweep `train_batch_size` (bin) ∈ {2,4,8,16,32,…} × α ∈ {0.1,1,100}, N=100, fluxtune. Report + peak/final acc, Δloss, and **per-commit JVP SNR / variance** (WS3-b + var telemetry). Cross-baseline (bin + size is a shared knob), flag-gated. +- **Reads on:** whether small bins are *required* for a usable forward-gradient signal — i.e. we **cannot** + just "enlarge bins to kill variance" (the H2 caveat) → motivates fixing variance at the **aggregator** (M2), + not the data. + +**M2 — Aggregator optimizer (descent to a minimum) vs. random walk. [FLUXTUNE CONTRIBUTION]** +- **Claim:** FwdLLM applies each committed JVP estimate as a **raw, undamped SGD step** + (`FedSgdAggregator.py:322-324`) → under small-bin noise the global model **random-walks** (charter I-1; §8 + F8/F13). Fluxtune's contribution is a **server-side optimizer that integrates the noisy-but-informative + small-bin updates into a smooth descent and holds the minimum** — momentum / EMA-of-weights / adaptive step + (§8 S1). **Distinct from and composable with C3**: C3 reweights contributions *within* a commit; the + optimizer damps *across* commits. +- **Metric:** sustained peak (no post-peak divergence) · monotone test-loss envelope · and the real prize — + the run **fills the W=20 convergence window** (never fires today, I-1) → turns E1's *transient* peak into a + genuine time-to-τ. Compare undamped baseline vs optimizer, α=1 N=100. +- **Opportunity experiments:** (a) momentum/EMA coefficient sweep; (b) does the optimizer let the run + *converge & sustain* 84%?; (c) optimizer × grad-aware (Opt-3) — does damping remove the "R4 diverges worst" + effect?; (d) optimizer × bin size (M1) — does a real optimizer widen the usable bin range? +- **Validate / fix before finalizing:** (i) optimizer-not-data established (✅ H0); (ii) confirm momentum does + not **double-damp** with the fedbuff staleness/utility rate and the variance gate; (iii) **real↔sim parity** + preserved (server optimizer state must be deterministic under the frozen update order); (iv) gain holds + across α, not just α=1. Flag-gated, default off; A/B before permanent (contributions §8 rule 6). + +--- + +## 5. Consolidated metric map (incl. the operator's second list) + +| # | Metric | Provenance | Feeds | +|---|--------|-----------|-------| +| 1 | Experiment wall-time before exit | WS2 `converge.json` + `agg_round.wall_elapsed_s` | 1 | +| 2 | Aggregator active GPU time (compute wall-time) | `step_timing`/`agg_round.aggregate_fedavg_s`,`eval_s` | 2,3 | +| 3 | Each client active GPU time | `trainer_round.gpu_compute_s` | 2,3 | +| 4 | Loss at first iteration | first `agg_eval.test-loss` | 3 | +| 5 | Loss after training complete | last `agg_eval.test-loss` | 3 | +| 6 | Updates used per iteration | `agg_round.agg_goal_count`/`updates_in_queue` | — | +| 7 | Total updates over training | Σ `agg_round.agg_goal_count` | — | +| 8 | **Forward passes / perturbations per client** | **WS3-b** `trainer_round.forward_passes_total`/`perturbations_total` | 3 | +| 9 | Data sent from aggregator | **WS3-a** `comm{direction=agg_to_trainer}.size_bytes` | 4 | +| 10 | Data sent from clients | **WS3-a** `comm{direction=trainer_to_agg}.size_bytes` | 4 | +| 11 | Client active: selected→next reselection (async) | `agg_round.contributor_intervals` | 5 | +| 12 | Client active per round (sync reselect) | `trainer_round`/`agg_round` ts | 5 | +| 13 | **Participation counts: rounds / data_bins / iterations per client** | `trainer_round` + `contributing_trainers` | 5 | + +--- + +## 6. Implementation workstreams & status + +| WS | Item | State | +|----|------|-------| +| WS1 | `EXPERIMENTS.md` + `experiments.yaml` (backbone) | ✅ | +| WS2 | Convergence-stop watcher (`--target-acc`,`--converge-window`, `CONVERGED`/`DID_NOT_CONVERGE`, `converge.json`) | ✅ validated N=10 (fired at bin 0) | +| WS3-a | Network bytes + message-count telemetry (`comm` event, both directions) | ✅ validated (agg→trainer + trainer→agg) | +| WS3-b | Per-client perturbation / forward-pass counter (`fwdgrad_utils` counters → `trainer_round`) | ✅ validated (20 passes / 10 perturbations per iter) | +| WS4 | `compare_baselines.py` cross-baseline reducer → table/CSV + overlay plots | ✅ (mix-guard + degrades on missing telemetry) | +| — | Enhanced gate: `condition_fp` (incl. delay_factor), tier ② baseline internals, agg_goal-match check, `--run-set` registry launch | ✅ | +| — | Phase 2: fwdllm_plus-under-mobiperf policy; N=100/8-GPU utilization confound | ⬜ deferred | + +**Validated end-to-end at N=10** (run `run_20260706_152025_fluxtune_n10_smoke_syn_0_real`): `CONVERGED` fired, +`comm` bytes (both directions) + `forward_passes`/`perturbations` landed in JSONL, comparison table produced. +**Sign-off applied** (2026-07-06): `main` is now N=100, delay_factor=2, agg_goal=10 matched, target 0.82. +**Next: the N=100 real convergence runs** (§7). + +--- + +## 7. Run workflow (multi-node, misconfig-proof) + +`experiments.yaml` + the enhanced gate define a comparison run **once**; every node verifies it launched the +same thing. + +### 7.0 Signed-off condition (2026-07-06) — `run_set: main` +N=100 · K=10 · C=10 (sync) / 30 (fluxtune) · **agg_goal=10 matched across all three** · partition alpha=1 · +syn_0 · delays ON at **delay_factor=2** · target 0.84 / window 20 · **wall ceiling 48h · stall-out after 2h +with no ≥1% gain**. A pre-flight check enforces the agg_goal match; the `condition_fp` (includes delay_factor) +must match across nodes. **Manual pre-run check:** verify the partition group `niid_label_clients=100_alpha=1` +exists in `agnews_partition.h5` (the gate only warns — it can't read the H5 from the launch env); a missing +group crashes all 100 trainers. + +### 7.1 Baseline defaults — REVIEW before the first real run +Resolved from `_metadata/baselines.yaml` (the gate's tier ② shows these live per run): + +| knob | fwdllm | fwdllm_plus | fluxtune | +|---|---|---|---| +| mode | sync | sync | **async** | +| selector | random | random | **async_oort** | +| optimizer | fedavg | fedavg | **fedbuff** (lr 0.075) | +| native agg_goal | 10 | 10 | **3** | +| native k / c | 5 / 15 | 5 / 15 | 5 / 15 | +| availability tracking | unaware | **ORACULAR** | client_notify | + +**agg_goal / k coupling (decided 2026-07-06):** `--c N` fans `agg_goal=N` + `minInit` to **every** baseline +and `--k N` sets k on all — so `--run-set main` (C.sync=10, K=10) drives fluxtune to `agg_goal=10, k=10`. +**Intended:** the experiment matches `agg_goal=10` across all three, and fluxtune's async concurrency is set by +`C.async=30` (its `c` becomes 30, `agg_goal` stays 10). The agg_goal-match pre-flight check enforces it. A +future experiment wanting fluxtune's native `agg_goal=3` needs a per-baseline agg_goal override (not currently +wired). + +### 7.2 Single-node run (all three baselines) +```bash +cd lib/python/examples/fwdllm/expt_scripts +bash run_sequential.sh --run-set main --mode real --yes +``` + +### 7.3 Multi-node run (split baselines, ONE source of truth) +All nodes read the SAME condition from `experiments.yaml` via `--run-set main`; only `--only` differs. **The +gate prints `condition_fp` — it MUST be identical on every node.** If fingerprints differ, a knob was +mistyped: stop and fix. + +**Three-node run (one baseline per node) — pass `--clean` so each node auto-clears any stray workers from a +prior run before launching:** +```bash +# node A +bash run_sequential.sh --run-set main --only fwdllm --mode real --clean --yes +# node B +bash run_sequential.sh --run-set main --only fwdllm_plus --mode real --clean --yes +# node C (shared filesystem: run dirs land in the same experiments/) +bash run_sequential.sh --run-set main --only fluxtune --mode real --clean --yes +``` +Pre-run checklist (the gate does most of this — eyeball, don't skip): +1. `condition_fp` identical across nodes (expect `04d64814` for the current `main`). +2. tier ② `mode/selector/optim` match the baseline table above (right algorithm per baseline). +3. `target_acc`, `trace`, `part` are the intended values (🟢 = from flag/registry); `part` = `niid_label_clients=100_alpha=1`. +4. no ✗ pre-flight checks (a ⚠ on the niid partition group just says "verify it exists"). +5. `[] clean slate verified` printed before launch (the clean-slate guard, §7.6). + +### 7.4 After the runs — compare +```bash +python compare_baselines.py --variant real --target-acc 0.84 --window 20 --plots +``` +Reducers pull from the same run dirs; the **mix-guard** warns if the baselines' shared axes (N/partition/trace) +disagree — the post-hoc twin of `condition_fp`. + +### 7.5 Convergence-stop knobs (recap) +`--target-acc 0.84 --converge-window 20` arm the watcher (WS2). A run ends `CONVERGED` (window met, +`converge.json` written) or `DID_NOT_CONVERGE` (hit a safety cap). `--run-set main` sets these from the +registry, so you rarely pass them. + +### 7.6 Stopping a run & the clean-slate guard +The run's trainers/aggregator run in their **own process group** (so the watcher can signal the whole tree), +which means a bare terminal **Ctrl+C would not reach them**. `run_sequential.sh`/`expt_runner.sh` install a +**SIGINT/SIGTERM trap**: one Ctrl+C tears down the run's process group + the convergence watcher +(`converge_watch.py`) + the progress ticker, escalates SIGTERM→SIGKILL after a short grace +(`EXPT_INT_GRACE_S`, default 5s), sweeps stragglers, and frees GPU/RAM. + +**Manual teardown** (if a run was killed the wrong way and left orphans): +```bash +pkill -TERM -f 'flame.launch.run_experiment'; sleep 3 +pkill -9 -f 'trainer/forward_training'; pkill -9 -f 'trainer/pytorch/main.py' +pkill -9 -f 'aggregator/pytorch/main_'; pkill -9 -f converge_watch.py +pkill -9 -f run_sequential.sh +# verify clean (want: nothing, GPU ~0 MiB) +pgrep -af -u "$USER" -f 'run_experiment|forward_training|trainer/pytorch/main.py|aggregator/pytorch/main_|converge_watch.py' || echo clean +nvidia-smi --query-gpu=index,memory.used --format=csv +``` + +**Clean-slate guard (`expt_assert_clean_slate`).** Before every launch the runner checks for stray FL workers +(own procs only) and residual GPU memory: +- **default:** if the node is dirty it **ABORTS** and prints the kill command (never nukes a process you didn't + sign off on — safe on shared boxes); +- **`--clean`** (or `EXPT_AUTOCLEAN=1`): kills the stragglers, re-verifies, and only aborts if still dirty; +- `EXPT_GPU_FREE_MB` (default 500) warns on residual GPU memory; `EXPT_GPU_STRICT=1` turns that warning into an + abort. + +--- + +## 8. Files & entry points + +| File | Role | +|---|---| +| [`experiments.yaml`](experiments.yaml) | machine registry — run-sets, conditions, analyses (source of truth for *what runs*) | +| [`expt_scripts/run_sequential.sh`](expt_scripts/run_sequential.sh) | launcher — `--run-set`, condition_fp gate, tier ② internals, agg_goal-match check, convergence flags | +| [`../scripts/expt_runner.sh`](../scripts/expt_runner.sh) | shared harness — `expt_launch` (arms watcher + SIGINT/SIGTERM teardown), `expt_assert_clean_slate` (pre-launch guard), `expt_assert_run` (`CONVERGED`/`DID_NOT_CONVERGE`) | +| [`../scripts/converge_watch.py`](../scripts/converge_watch.py) | WS2 side-car — polls `agg_eval`, writes `converge.json`, kills the run on convergence | +| [`expt_scripts/compare_baselines.py`](expt_scripts/compare_baselines.py) | WS4 reducer — 5-experiment table/CSV + overlay plots + mix-guard (cross-baseline) | +| [`expt_scripts/plot_run.py`](expt_scripts/plot_run.py) | per-run twin — streams ONE run's telemetry (handles the >1 GB agg JSONL) → full 5-experiment plot set + `summary.json` | +| `flame/telemetry/events.py` | `build_comm` (WS3-a) | +| `flame/.../fwdllm_aggregator.py`, `fwdllm_trainer.py` | `comm` emit sites (WS3-a, both dispatch paths + upload) | +| `trainer/forward_training/fwdgrad_utils.py`, `FedSgdTrainer.py` | forward-pass counters (WS3-b) | + +Session artifacts: `expt_scripts/smoke_logs//` (`converge_.json`, manifest, gate spec); per-run +`experiments/run_*/telemetry/*.jsonl`; comparison output `experiments/_compare/`. + +--- + +## 9. Changelog +- **2026-07-08 — stability track (M1–M2).** Added M1 (data-bin size × α bias/variance tradeoff) + M2 + (aggregator optimizer vs. random walk — fluxtune contribution). Motivated by charter I-1 + the H0 diagnostic + refuting data-class-bias (`fluxtune_contributions.md` §8, F11-F15). Flag-gated; ledger row added (§10). +- **2026-07-07 (g) — paper ⇄ code reconciliation ([`EXPTS_CHARTER.md`](EXPTS_CHARTER.md)).** Corrections: (i) + contribution taxonomy → C1/C2/C3 (async is structural); (ii) C3 is a borrowed fedbuff scalar-rate placeholder, + not gradient-aware (+ C3 design investigation, §4); (iii) E1 units → wall/rounds/data-bins/iterations, vclock + deferred; (iv) E3 forward-pass denominator primary, GPU-seconds secondary; (v) reducer gaps N3/N5/N6 logged; + (vi) run ledger added (§10). Open: α (config 1, runs ran 0.1). Paper-side rewrites in `05-evaluation.tex`. +- **2026-07-07 (f) — E3/E4 observed (Fluxtune trades efficiency for speed).** N=100 smoke contradicts E3/E4: + FwdLLM++ more compute-efficient (Δloss/M-fwd 5.62 vs 1.27) + communication-efficient (79 vs 146 GB). Both: + async concurrency does more total work (staleness + speculative compute; weight re-pulls 88 vs 9 GB). Fluxtune's + win is E1 speed. §4 has optimization directions. Caveat: only C1 on; C2/C3 OFF → E2/E3/E4 need a full-system run. +- **2026-07-07 (e) — target 0.82 → 0.84 + paper-figure pipeline.** Raised `main` target to **0.84** in + `experiments.yaml` + plot defaults. Added `expt_scripts/plotlib/` + `make_paper_figs.py`; + `plot_run.py`/`compare_baselines.py` on the shared reducer. Fixed the data_id-keyed eval bug (cycles per round, + corrupted Δloss) → time-ordered series. Per-run cutoff at last significant test-loss improvement; EMA `--smooth`. +- **2026-07-07 (c) — loss-aware stall guard (`--stall-on`).** Idle clock resets on **acc / loss / either** + (default `either`). Loss = relative drop vs running-best (`--loss-min-rel-delta` 1%); acc = absolute 1%; both + running-best. Fixes killing a run whose acc plateaus while loss still falls. Wired through `converge_watch.py`, + `expt_runner.sh`, `run_sequential.sh` (+ `condition_fp`), `experiments.yaml`. Validated end-to-end. +- **2026-07-07 (b) — CLI stall window (hours alias).** `--stall-window-h H` on `run_sequential.sh` overrides the + registry `stall_window_s`. Motivated by a fwdllm N=100 run the 2h guard killed at the start (loss still + falling). Re-launch: `run_sequential.sh --run-set main --only fwdllm --mode real --clean --yes --stall-window-h 6`. +- **2026-07-07 — per-run plots + ticker-orphan fix.** Added `plot_run.py` (single-run twin, streams the >1 GB agg + JSONL → 5-experiment plot set + `summary.json`; first N=100 fluxtune run `run_20260706_185045…` STALLED at + **84.08%**, 139 bins). Fixed a harness hang: the backgrounded ticker wasn't in its own process group, so the + process-group `kill` missed it and `wait` blocked forever; the ticker now self-terminates by PID when the run dies. +- **2026-07-06 (c) — teardown, clean-slate guard, alpha=1.** Ctrl+C/SIGTERM tears down the run (own process + group) + watcher + ticker + frees GPU/RAM. Added `expt_assert_clean_slate` (`--clean`/`EXPT_AUTOCLEAN`), §7.6. + **`main` partition alpha=0.1 → alpha=1** (group present in `agnews_partition.h5`); `condition_fp` now `04d64814`. +- **2026-07-06 (b) — termination policy.** **48h wall ceiling** (was 1h) when `--target-acc` set + **stall guard** + (`STALLED` if best-acc gain < `stall_min_delta` 1% within `stall_window_s` 2h). Wired through watcher, registry, + gate, `--run-set`. Unit + integration tested. +- **2026-07-06 — implementation landed & N=10-validated.** WS2 watcher; WS3-a `comm` telemetry both directions; + WS3-b forward-pass/perturbation counters; WS4 `compare_baselines.py`; enhanced gate (`condition_fp` incl. + delay_factor, tier ② internals, agg_goal-match, `--run-set`). Sign-off: N=100, delay_factor=2, agg_goal=10, target 0.82. +- _(init)_ Doc + registry scaffold. Backbone, convergence-stop spec, metric map (incl. #8 forward passes, #13 + participation granularities). + +--- + +## 10. Run ledger (which log feeds which result) + +Update as runs land — how we know which log file on which node backs each figure/claim. `Status`: SMOKE +(validation, not for paper) · FINAL (paper number) · STALLED/CONVERGED/DNC (verdict). ⚠ The three N=100 runs +below ran at **α=0.1** (log filenames say `alpha0p1`). **α=0.1 is now excluded** — learning was too slow across +all baselines to complete convergence runs — so the paper's primary condition is **α=1** +(`experiments.yaml main`), and these runs will be **re-run at α=1** for final numbers. Treat the α=0.1 runs as +smoke / evidence-that-0.1-is-too-slow, not paper numbers. + +| Run dir | Baseline | Node | Condition | Verdict | Feeds | Notes | +|---|---|---|---|---|---|---| +| `run_20260707_015846_fwdllm_n100_smoke_syn_0_real` | fwdllm | shepherd | N=100, syn_0, α0.1, df=2, agg_goal=10 | SMOKE | E1–E5 (baseline) | log `07_07_26_01_59_random_n100_default_alpha0p1_syn0_*` | +| `run_20260706_185023_fwdllm_plus_n100_smoke_syn_0_real` | fwdllm_plus | kaylee | N=100, syn_0, α0.1, df=2, agg_goal=10 | SMOKE | E1–E5 (baseline) | log `06_07_26_18_50_random_n100_oracular_alpha0p1_syn0_*`; oracular **inert** at syn_0 | +| `run_20260706_185045_fluxtune_n100_smoke_syn_0_real` | fluxtune | shepherd | N=100, syn_0, α0.1, df=2, agg_goal=10, C=30 | STALLED @84.08% (139 bins) | E1–E5 (C1-only, superseded) | log `06_07_26_18_51_…`; C2 off, C3=placeholder. **Superseded by R4 (`…025716…`) for the baseline comparison.** | +| `run_20260708_025543_fluxtune_n100_smoke_syn_0_real` | fluxtune **R1** base | — | N=100, α=1, agg_goal=10, C=30, var-stop=off, agg-rate=new | max 83.00% | ablation (Opt-2/3 off) | 2×2 opt ablation; `figs_ablation.yaml`. **FluxTune-base (C1 guided-JVP + Opt-1), NOT FeLiX.** α=1 (dir mislabeled `alpha0p1`) | +| `run_20260708_025616_fluxtune_n100_smoke_syn_0_real` | fluxtune **R2** +var-stop | — | …var-stop=plateau, agg-rate=new | max 82.25% | ablation (Opt-2 only) | isolates Opt-2 | +| `run_20260708_025636_fluxtune_n100_smoke_syn_0_real` | fluxtune **R3** +grad-aware | — | …var-stop=off, agg-rate=grad_aware | max 83.91% | ablation (Opt-3 only) | isolates Opt-3 | +| `run_20260708_025716_fluxtune_n100_smoke_syn_0_real` | fluxtune **R4** full (=default) | — | …var-stop=plateau, agg-rate=grad_aware | max 84.08% | ablation (full) **+ E1–E5 baseline comparison** | full-stack default; feeds `figs.yaml fluxtune` | +| _pending_ | fluxtune (full-system) | — | C2 dynamic K/C **ON** + real C3 ON | — | E2/E3/E4 efficiency | breaks agg_goal match → separate run-set | +| _pending_ | all three | — | `mobiperf_*` (real-world availability) | — | E1 headline | needs fwdllm_plus-under-scarcity policy | +| _pending_ | fwdllm (or port) | — | fidelity: accuracy vs `xu2024fwdllm` | — | Setup (D3) | locate **old** run data; accuracy parity, not time | +| _pending_ | ablations | — | JVP-sens / K-C-sens / α∈{0.1,0.5} | — | Ablation §§ | tooling TBD (charter §2e) | +| _pending_ | fluxtune | — | stability: bin-size×α sweep (M1) + server-optimizer (M2) | — | I-1 fix / E1 convergence | flag-gated; see `fluxtune_contributions.md` §8 | diff --git a/lib/python/examples/fwdllm/EXPTS_CHARTER.md b/lib/python/examples/fwdllm/EXPTS_CHARTER.md new file mode 100644 index 000000000..684908a49 --- /dev/null +++ b/lib/python/examples/fwdllm/EXPTS_CHARTER.md @@ -0,0 +1,174 @@ +# FLUXTUNE evaluation charter (living doc) + +**Purpose.** Current state + resolved decisions for the fluxtune eval, reconciling the paper +(`05-evaluation.tex`) with the code (`EXPERIMENTS.md`, telemetry, plots). History lives in git; this doc +holds only what's *current* plus rationale a diff wouldn't explain. Run ledger: `EXPERIMENTS.md` §10. +**Owner:** dgarg39 · **Branch:** `dg/fwdllm_sim_unavail` + +## Status + +Paper reconciliation ✅ delivered. Bottleneck optimizations built — all **flag-gated, byte-identical +off, fluxtune-only, unit-tested**. Measurements are at the paper's primary **α=1** point (experiments +never go below α=1; ablations go UP to α∈{10,100} — [[fwdllm-alpha-convention]]; the `…185045…alpha0p1…` +dir name mislabels an α=1 run — **verify α from the loaded-partition log line, not the dir name**). + +| Opt | Feature | Flag(s) | Status | +|---|---|---|---| +| 1 | Suppress redundant intra-databin weight re-sends | `suppress_redundant_weights` (all baselines) | ✅ validated: 0% redundant, −79% down-bytes | +| 2 | Variance-plateau force-commit (plateau early + max-iter cap late) | `var_stopping_policy=plateau`, `var_plateau_patience`, `var_plateau_rel_delta`, `max_iterations_per_data_id` | ✅ default ON fluxtune (N=3, ε=0.15, cap=20) | +| 3 | Gradient-aware aggregation / C3 (align gate + inverse-var) | `agg_rate_conf.type=grad_aware` | ✅ default ON fluxtune (`type=new` = FeLiX ablation only) | +| 4 | Dynamic C | `dynamic_kc.enabled` | ⬚ wired, not enabled — remaining | + +**Default fluxtune = full stack** (C1 JVP + Opt-1 + Opt-2 + Opt-3). Code: +`flame/mode/horizontal/syncfl/fwdllm_aggregator.py`, `examples/fwdllm/aggregator/FedSgdAggregator.py`; +config `examples/_metadata/baselines.yaml`; tools `expt_scripts/{characterize_variance_curve, +audit_weight_redundancy}.py`; tests `tests/mode/test_fwdllm_*`. + +**Remaining opts:** Opt-4 dynamic C (wired, unenabled) · Opt-5 finer staleness clock (bundle w/ Opt-3) · +Opt-1 cross-databin delta-cache (larger comm piece). **Parked** (measured dead-ends): dynamic-K to lower +the floor (floor is structural non-IID) and staleness-based-C (staleness too low at syn_0). + +## Four planned full runs + +N=100, `target_acc=84%`, parallel — a **2×2 over the two learning-affecting opts** (C1+Opt-1 constant). +Toggled via `run_sequential.sh` CLI flags (they patch per-run config_overrides, which WIN over the +baselines.yaml catalog at launch — dry-run validated); no per-run config files needed. + +Common: `--only fluxtune --mode real --num-trainers 100 --partition-method niid_label_clients=100_alpha=1 +--agg-goal 10 --c 30 --target-acc 0.84 --yes` (agg_goal/C match the n100 reference run). + +**Naming:** R1 is the **FluxTune-base** (C1 guided JVP perturbations + Opt-1, both ON in all four runs), +NOT FeLiX — it does forward-mode LLM perturbation fine-tuning and merely *borrows* FeLiX's scalar +aggregation rate (`agg_rate_conf.type=new`, decision N2). FeLiX targets the CNN/speech family with no +perturbation-based fine-tuning (`baselines.yaml` fluxtune spec). `select_perturbation_using_jvp=true` is a +**trainer-side** flag (the aggregator config's copy reads false and is unused). + +| Run | Opt-2 var-stop | Opt-3 grad-aware | Extra CLI flags | Isolates | +|---|:---:|:---:|---|---| +| **R1** FluxTune-base | ✗ | ✗ | `--var-stopping-policy off --agg-rate-type new` | C1-only reference (≈ last night) | +| **R2** + var-stop | ✓ | ✗ | `--agg-rate-type new` | Opt-2 alone | +| **R3** + grad-aware | ✗ | ✓ | `--var-stopping-policy off` | Opt-3 alone | +| **R4** all (= default) | ✓ | ✓ | *(none)* | full fluxtune | + +R2−R1 & R4−R3 = Opt-2 effect; R3−R1 & R4−R2 = Opt-3 effect; the 2×2 also gives the interaction. After: +read `commit_reason` (natural/cap/plateau) + `grad_aware_gated_total` telemetry; tune ε/cap + +align_floor/inverse_var; re-characterize at α∈{10,100}. + +**LANDED 2026-07-08** (N=100, α=1 — dir names say `alpha0p1`, MISLABELED; loaded-partition log reads +`niid_label_clients=100_alpha=1`). Runs identified by `aggregator_config.json` (`var_stopping_policy` × +`agg_rate_conf.type`); figures via `expt_scripts/figs_ablation.yaml` + blue-ramp styles in +`plotlib/baselines.py`. R4 = full default, also feeds the baseline comparison (`figs.yaml fluxtune`). + +**Results — peak test accuracy (the paper number; every run diverges after → see Issue I-1).** +Peak is at/near target and rises with the opt ladder (R4 full ≈ target); no run *sustains* it. + +| Run | var_stopping_policy | agg-rate type | Run dir | **peak acc @ round-1** | gap to 84% | +|---|---|---|---|---|---| +| **R1** FluxTune-base | off | new | `run_20260708_025543_fluxtune_n100_smoke_syn_0_real` | 83.00% @ 2.91h | −1.0 | +| **R2** +var-stop | plateau | new | `run_20260708_025616_fluxtune_n100_smoke_syn_0_real` | 82.25% @ 4.07h | −1.8 | +| **R3** +grad-aware | off | grad_aware | `run_20260708_025636_fluxtune_n100_smoke_syn_0_real` | 83.91% @ 2.44h | −0.1 | +| **R4** full (=default) | plateau | grad_aware | `run_20260708_025716_fluxtune_n100_smoke_syn_0_real` | **84.08% @ 3.90h** | **+0.1** | + +**Baseline comparison (`figs.yaml`, R4 = fluxtune):** FluxTune **84.1% @ 3.9h** · FwdLLM++ 80.9% @ 7.4h +(−3.1 from target, ~2× slower) · FwdLLM 30.3% (never learned, shown full). E1 headline (speed + reaching +target) holds; the peak-accuracy ★ + legend value on every E1 acc plot shows each baseline's gap to target. + +## Bottleneck (measured, α=1 run `run_20260706_185045_fluxtune_n100…`) + +**The variance gate is the hub:** the fixed absolute `var≤0.3` gate sits *below* the achievable variance +floor (~0.45), so a data-bin commits only on a noise dip and grinds 15→34 iters. The resulting slow +model-version makes staleness degenerate and drives long non-commit comm stretches. + +| # | Metric | Measured | Meaning | +|---|--------|----------|---------| +| M-1 | agg iters that commit | 4.3% | 95.7% of compute advances nothing | +| M-3 | variance floor vs threshold | ~0.45 med / 0.30 min vs thr 0.30 | gate crossed only by noise dips | +| M-5 | iters/data-bin (max) | 18 (61) | bins grind past diminishing returns | +| M-6 | force-commit firings | 0.0% | escape valve inert → Opt-2 | +| M-8 | staleness (median/p90) | 1 / 3 | model-version crawls → rate can't differentiate | +| M-10 | redundant weight bytes | ~88 GB, ≈99% re-sends | unchanged model re-pulled → Opt-1 | +| M-12 | acc 50/75/100% wall | 0.825 / 0.743 / 0.806 | mid-run regression → aggregation instability → Opt-3 | + +Floor is **structural** (grad pool grows but `var` asymptotes) → more samples won't lower it; "weight +smarter" (C3) or "accept the floor" (plateau) will → dynamic-K de-prioritized. + +### ⚠ Issue I-1 — runs don't hold the minimum (ROOT-CAUSED; fix = `fluxtune_contributions.md` §8) + +**Status: root-caused, fix in progress (§8 track, next = S1).** All 4 ablation runs learn in round 1 +(peak 82-84%, min loss 0.55-0.66) then oscillate and collapse to ~25% (mcc 0) with loss blow-up (R4 +loss→4.9); force-stopped 2026-07-08. Paper uses each run's round-1 peak; E1 plots clipped at peak. Not a +blocker for the E1 headline. + +**Root cause (H0 diagnostic — `fluxtune_contributions.md` §8, F1-F15):** an **undamped, high-variance +forward-gradient optimizer** — each noisy JVP commit applied raw (`FedSgdAggregator.py:322`, no momentum/EMA) +→ the model random-walks. Severity tracks aggregation aggressiveness (R4 grad-aware worst; R1 least → Opt-3 +as tuned *amplifies*, `align_floor=0`/`inverse_var` off). **Refuted:** the epoch-boundary reset bug (F10 — +`_model_version` monotonic across the boundary) and data class-bias (F11 — K=10 cohort near class-balanced +at α=1). The frozen data schedule only *freezes* the noise → position-locked collapse (F13). + +**Fix:** server-side optimizer (momentum/weight-EMA) = descent, not random walk (§8 **S1** / EXPERIMENTS.md M2); +interim tempering = cross-round LR decay + grad-aware `align_floor>0`/`inverse_var` (§8 S3). Also make the stall +guard **divergence-aware** (fire on rising loss) — today it can't: W=20 needs consecutive bins ≥0.84 (accuracy +only grazed it) and the `either` guard stays alive while round-1 loss still improves, so it can't fire before +round 2 destroys the model. + +**Opt tuning (kept, `characterize_variance_curve.py`):** var@commit≈0.29 (noise dip), plateau ~0.45; cap-12 ≈ +−48% iters → chose plateau **ε=0.15** + **cap=20**. Opt-1 audit: fwdllm 90%/22 GB, fluxtune 71%/68 GB → 0% post-fix. + +## Latest figures (for paper embedding) + +Two PDF sets, same 7 basenames, from `make_paper_figs.py` (cutoff `--cutoff-mode peak_acc` default → +every run clipped at its peak, Issue I-1 tail excluded; E1 acc plots carry a ★ + legend "peak X%" per +run). Rebuild: `cd expt_scripts && python make_paper_figs.py --manifest [--out-root ]`. `latest` +symlinks the newest timestamped dir; copy PDFs into Overleaf by basename. + +| Set | Manifest | Dir (`latest` symlink) | +|---|---|---| +| **Baseline comparison** (FwdLLM / FwdLLM++ / FluxTune=R4) | `expt_scripts/figs.yaml` | `expt_scripts/paper_figs/latest/` | +| **FluxTune opt-ablation** (R1–R4, blue ramp) | `expt_scripts/figs_ablation.yaml` | `expt_scripts/paper_figs_ablation/latest/` | + +Figure basenames (both sets): `e1_acc_vs_time.pdf` (time-to-acc, peak ★) · `e1_loss_vs_time.pdf` · +`e2_trainer_busy_cdf.pdf` · `e3_dloss_per_gpu_hour.pdf` · `e3_dloss_per_mfwd.pdf` · +`e4_network_bytes.pdf` · `e5_session_cdf.pdf`. Each dir also has `manifest.json` (run dirs + cutoff h). + +## Resolved decisions + +| # | Decision | +|---|---| +| **A1** Model | DistilBERT-base-uncased (66.4M, frozen) + AdapterHub adapters (~1.5% trainable). Drop 7B claims. | +| **A2** Task | AG News 4-class topic classification. | +| **A3** Hardware | A40, 8-GPU box + a *modeled* mobile delay (`delay_factor=2`); mobile figures argued from structure, not measured. | +| **A4** Trace | `mobiperf_*` is the condition of record; REFL (`third_party/REFL/`) is a scoped future item. | +| **A5** α | Primary **α=1**, never below 1; non-IID ablation → **α∈{10,100}** (more IID). [[fwdllm-alpha-convention]] | +| **B1** FwdLLM_Plus | FwdLLM + per-iteration reselection [defining] + oracular availability (inert at syn_0 → scope to unavailability) + relaxed staleness. Sync, random, agg_goal=10. | +| **B2** Contributions | **C1** guided (JVP-magnitude) perturbation · **C2** dynamic K/C · **C3** gradient-aware aggregation. Async/iteration-level is the *substrate*, not numbered. | +| **C1/C2** E3/E4 | Retain E3/E4 as hypotheses (assume C2+C3 deliver); E1 stands on C1 alone. | +| **N2** C3 | `type=new` fedbuff rate is a borrowed **FeLiX** placeholder (scalar staleness×utility), NOT fluxtune's C3 → replaced by `grad_aware` (Opt-3, default on). | +| **N4** Clock | Report wall + rounds + data-bins + iterations; drop virtual-clock (note as future). | +| **N7** Memory | Motivation/design claim argued from structure (peak mem bounded by inference); no memory experiment. | +| **D1** Seeds | Single seed now; paper TODO for ≥3-seed medians on headline E1. | +| **D2** GPU-sec | Forward-pass count is the primary denominator; GPU-seconds secondary w/ an 8-GPU-contention caveat. | +| **D3** Fidelity | Accuracy parity (not time-to-accuracy) vs `xu2024fwdllm`; do **not** compare wall-clock. | +| **D4** SPRY | Exclude split/personalized-per-device FL; 1-line Baselines pointer + Background *why* (bib: spry, split-learning, HeteroFL, personalized-FL survey). | +| **D6** Surcharge | P=10 ⇒ 2P=20 fwd/iter ⇒ ~10× sync compute at P=10 → parity at P=1; fidelity opt −37% GPU (sync −68%), zero gradient change. | + +## Remaining work + +- **Paper repo (operator-owned):** D4 SPRY/excluded-baseline justification + `\tbd` bib keys; fidelity + accuracy value; confirm the `~1.5%` / `delay_factor=2` / `~3.6s` numerics. +- **Telemetry/reducers:** N3 (E1 time-to-τ reconstructs convergence instead of reading `converge.json`) · + N5 (E2 idle `mqtt_fetch_s` emitted-but-unused) · N6 (E5 sync-session should use one-round-span, not + `contributor_intervals`) · `real_gpu_time_s`↔`gpu_compute_s` schema footgun. +- **Figures (`plotlib/figures.py`):** E1 iterations-to-target metric · E2 aggregator-breakdown bar · + E3 Δloss-vs-cumulative-compute trajectory · E4 msg-size inset · E5 participation-count bars · + unify `FwdLLM++`/`FwdLLM_Plus` label. +- **Runs:** full-system E2/E3/E4 (C2+C3 on) · mobiperf E1 (needs the fwdllm_plus-under-scarcity barrier + decision) · fidelity accuracy-parity (locate old run) · optional ≥3-seed medians for E1. +- **Ablations:** C1 JVP sensitivity (threshold, refresh) · C2 K/C sensitivity · α∈{10,100}. + +## Principles + +**P1** substrate honesty — describe what ran, or run the claimed substrate before submission. **P2** +isolation vs full-system are two conditions — efficiency claims (E2–E4) need C2+C3 on. **P3** no takeaway +ships ahead of its evidence. **P4** directional source-of-truth — implementation facts flow code→paper, +narrative paper→code. **P5** one name per concept. diff --git a/lib/python/examples/fwdllm/PARITY_LOGICAL_TASKS.md b/lib/python/examples/fwdllm/PARITY_LOGICAL_TASKS.md new file mode 100644 index 000000000..42d0160c9 --- /dev/null +++ b/lib/python/examples/fwdllm/PARITY_LOGICAL_TASKS.md @@ -0,0 +1,249 @@ +# TEMP task tracker — logical real↔sim parity (delete when folded into simulate_fwdllm.md §G) + +**Goal.** Prove the sim takes the SAME logical steps in the SAME order as real (same cohorts, receive order, +variance cadence, grads) up to data bin 1 across fwdllm / fwdllm_plus / fluxtune, and make the parity CHECKS + +PYTESTS exhibit these. Correctness before speed; no hacks (simulate_fwdllm.md principles #14/#16). + +--- + +## ⏸⏸ SIM DEBUG PAUSED (2026-07-06) — switching to REAL-experiment telemetry/impl on this branch + +**PAUSING fluxtune sim debugging.** Next work on `dg/fwdllm_sim_unavail` is a set of **real-experiment-run +telemetry + implementation changes** to launch first (separate from the sim loop). Resume from "RESUME HERE" +once those land. No sim code is mid-edit — the tree is clean; everything needed to resume is here + in #15. + +### ✅ P3 was RUN — phantom stall CONFIRMED FIXED, but `sim_rate` still <1 for TWO NEW (non-drain-gate) reasons +Pair: sim `run_20260706_112114_fluxtune_n10_smoke_syn_0_sim` / real `run_20260706_110555_..._real` (both +`--max-data-id 2`, `--delay-factor 1`, `sim_compute_truthful_gate=True`, `jvp_perf_opt=True` on trainers — +verified symmetric). Sim ran clean to `data_id=2` (an earlier run was mistakenly Ctrl-C'd ~48s into steady state). + +- **Phantom fix WORKS + is load-bearing:** `SIM_GRAD_STUCK_EVICT=0`; `phantom_skip` rises **~1 per commit** (→65); + 30s failsafe stalls GONE (max gap **12.8s once**, rest ≤5s; baseline had 31s failsafe + 24 gaps ≥10s). **Original + #15 root (phantom `_sim_inflight_expected` 30s failsafe) is RESOLVED.** +- **But `sim_rate` did NOT cross 1** (steady-state ~0.49 == pre-fix 0.494; overall 0.31 startup-dominated on this + short scope). **Per-commit `Δvclock/Δwall = 0.465`** — recurring ~4–5s wall gaps advance vclock only 0–1s. The + phantom fix was **necessary, not sufficient**; the ceiling is now elsewhere: + - **(a) vclock OMITS the aggregator's serialized compute.** `aggregate()` (variance pass, O(pool); grows with the + 126-entry `cached_v`) runs 1–3s **between every commit** = **32s in sim**, **never credited to vclock**. + [`fwdllm_aggregator.py:1783`](../../flame/mode/horizontal/syncfl/fwdllm_aggregator.py#L1783) folds **eval** + (`_vclock.advance(now+_eval_s)`) but has **no fold for variance-compute time**. Genuine, symmetric compute: + **real `aggregate()`=29s (mean 1.25s), sim=32s (mean 1.34s)** — the K-D25 "fold genuine unmodeled compute" + class. This ~32s + ~29s drain/distribute/MQTT = the 61s steady-state deficit (145s wall − 84s vclock). + - **(b) No GPU-vs-D skip headroom (#1d).** `sct=send+max(gpu,D)`; real JVP GPU ≈4–5s (contended, 10 trainers/8 + GPUs; ~2.4s uncontended floor) ≈ min modeled `D=4.0` → `max(gpu,D)≈gpu` → sim BLOCKS in `drain_ready` on the + real GPU pass for ~0 virtual headroom. `buf_depth` pinned at **6** (grads ready) while the strict-sct drain + waits for the next-sct grad's GPU. Net: **sim wall (145s) > real (91s)** — the sct-ordered drain + over-serializes behind genuinely-computing trainers. + +### ▶ RESUME HERE (after the real-experiment work) — resolve the residual, then re-validate P3 +1. **Fold the non-overlapped `aggregate()` variance-compute time into the vclock** (mirror the eval fold at + `:1783`). Biggest, most principled lever. **CAVEAT:** in real async that compute partially OVERLAPS other + trainers' GPU passes, so folding the *full* time over-counts — fold only the *non-overlapped* portion, after + measuring real's `aggregate()`-GPU overlap. Baseline-affecting vclock change → operator sign-off (principles + #1/#12/#16); land behind a flag, byte-identical off. +2. **Restore GPU-vs-D headroom (#1d):** pin **1 trainer/GPU** (GPU → ~2.4s floor, under min D=4.0) and/or raise + `--delay-factor` so `max(gpu,D)=D` gives real skip-able headroom. Lifts #15-residual-(b) AND #1d cohort overrun. +3. **Then re-validate:** a LONGER flag-on run (past the ~125s startup that dominated this short scope) PLUS a + flag-OFF A/B at matched scope to prove parity UNCHANGED (this run's 11 parity fails all look pre-existing, but + not provable without the A/B). Re-run command below. + +```bash +cd lib/python/examples/fwdllm/expt_scripts +bash run_sequential.sh --only fluxtune --mode both --delays on --delay-factor 1 --max-data-id 2 --yes +python run_parity.py --yes --max-bin 1 --baselines fluxtune +# P3 gate greps (new sim agg log): +# FS=$(ls -td ../experiments/run_*_fluxtune_n10_smoke_syn_0_sim | head -1); AGG=$(ls "$FS"/*aggregator.log|head -1) +# grep -c SIM_GRAD_STUCK_EVICT "$AGG" # expect ~0 (WAS 0 ✓) +# grep SIM_GRAD_RECV "$AGG" | grep -oE 'phantom_skip=[0-9]+' | tail -1 # expect >0 (WAS 65 ✓) +# grep VCLOCK_PROGRESS "$AGG" | tail -1 # the sim_rate verdict (WAS ~0.49 steady ✗) +``` +**EXPECT (PASS):** `STUCK_EVICT~0` ✓ + `phantom_skip` rising ✓ (both hold) AND **`sim_rate` → >1** (the open part) +AND `cohort_sequence`/`var`/`staleness` parity **UNCHANGED** vs the flag-off run. **If parity MOVES:** compute cap +too tight → raise `sim_gate_compute_cap_s` (yaml, default 10.0); the fix must only remove dead wall, not change +WHICH grad commits WHEN. + +**Telemetry to add when resuming:** per-commit `Δvclock` vs `Δwall` line (ratio<1 tell); `aggregate()` +overlap-with-GPU fraction (decides fold (1)); explicit `[VCLOCK_DEFICIT] aggregate_s=… drain_s=… uncredited_s=…` +at each `VCLOCK_PROGRESS`. Repro one-liners in the #15 "Repro" block below. + +### Other OPEN fronts (after P3), priority order +1. **bin-7 nondeterminism → relax the parity target (SHARED, correctness-of-CHECK).** SYNC full-run cadence breaks + at bin 7 even with order 41/41 matched (K-D31): ~1e-3 GPU fp16 grad jitter, amplified by the split-half variance + ratio, flips `var<0.3` at (7,2) — NOT a sim bug. **Task:** keep `cohort_sequence` EXACT scoped to `--max-bin 1`; + ADD a DISTRIBUTIONAL cadence/var rung (mean-band + KS + `var_good` fraction) for the full run. Confirm with a + 2-real-run diff first (= **P0-2**, below). +2. **P0-2 — grad-determinism-given-order confirmation.** Force identical commit order in a sim + a real short run + (or two real runs), diff per-iteration `var` + a grad norm. Decides exact (sim-bug) vs distributional + (nondeterminism) target beyond bin 1. Also unblocks the P1-6 LIVE sim==real grad-determinism pytest (TODO). +3. **#1d / P2-5 — fluxtune cohort SET diverges (thin-margin overrun), fluxtune only.** `set_match=3/272`. + `jvp_perf_opt` (K-D32) cut GPU mean 7.57→3.61s (<4.0s budget) but the TAIL (4.1–5.4s) overruns on the two + doubled GPUs (10/8) → order flips → wrong 3-of-K commit. **Task:** once #15's faster commits land, re-check; if + a residual tail remains, tune `perturbation_count`↓ (P2-5, wired, default 10; LOWERING changes the baseline + algorithm → operator call) so GPU < min cohort D, and/or `--delay-factor` up. +4. **#7 — fwdllm_plus real ~4× slower/round; at syn_0 real sees ~4.9 eligible vs sim ~9.6.** Not a sim bug. + Profile per-iteration reselection + oracular-read cost from the banked per-phase log; explain the eligible gap. +5. **#11 — real-mode critical-path waste** (`sleep(0.1)` MQTT-settle, one-grad-per-poll drain tail), real-only, + ZERO parity impact (sim already skips). Deferred to a validated pass (removing it needs a real run; #8/#11c). +6. **Phase 3 — extend beyond bin 1** once bin-1 parity holds; re-check where it breaks; iterate. Then C1/C2 + convergence (distributional target) at matched `data_id` → gate to Phase 2 (unavailability). + +### D3 (optional, secondary) — N=20/C=10/K=3 run +Adds a 10-trainer idle pool; confirms real fills C from fresh idle trainers while a returner waits for its commit +(validates F5) and the sim reproduces the selection sequence. `run_sequential.sh --only fluxtune --mode both`, +`num_trainers=20`, selector `c: 10`, `agg_goal: 3`. + +--- + +## 📚 FELIX GROUNDING (async_cifar10) — the reference model for #15 (durable; keep) + +*How async_cifar10's felix drain ACTUALLY behaves (from code + `async_cifar10/PARITY.md`). Anchors: **asyncfl** = +`flame/mode/horizontal/asyncfl/top_aggregator.py`; **fwdllm** = `flame/mode/horizontal/syncfl/fwdllm_aggregator.py`.* + +- **F1 — Felix's expected-sct arrival gate is INERT (gate_holds=0), not load-bearing.** It stamps + `_sim_inflight_expected[end]=dispatch_vclock+budget` (asyncfl:1668-71, a LOWER BOUND) and CAN block + (`earlier_stuck`, asyncfl:430-446), but over a cifar run `gate_holds=0` (asyncfl:157-162; + PARITY.md:294-296/321-323): cifar GPU ≈ **0.4s**, so every in-flight update is arrived+buffered when the drain + runs. Felix's EFFECTIVE behavior = sort arrived updates by sct and commit. The gate is a dormant safety net. +- **F2 — Felix re-dispatches on COMMIT; the felix trainer idles in recv — sub-second only because compute≈0.** + `_sim_hold_busy_slots` (asyncfl:1453-1500) at the agg-goal boundary holds each busy trainer in + `selected_ends`/`all_selected`/`_sim_pending_commit` until its update commits (released asyncfl:618-627); the + trainer blocks in `_fetch_weights`→`channel.recv` (`syncfl/trainer.py:187`) meanwhile. So felix IS + hold-to-commit — **fwdllm's K-D17b is a FAITHFUL port.** Precondition: return→commit sub-second in cifar → ≈0 + idle wall. +- **F3 — Two ledgers already SEPARATE in felix.** PHYSICAL (GPU/MQTT, `_sim_buffer`, `cleanup_recvd_ends`) vs + VIRTUAL (selection eligibility `_sim_pending_commit`/`all_selected`/`selected_ends`, staleness gate + `_sim_inflight_expected`, vclock). "Busy ≠ UN_AVL": a busy trainer holds a SLOT (`extra = c − + len(selected_ends)`), released on commit (asyncfl:1459-62; PARITY.md:324-328). +- **F4 — Regime table (why the same mechanism transfers to sync fwdllm but breaks fluxtune):** + + | baseline | GPU | agg_goal vs c | return→commit idle | gate | outcome | + |---|---|---|---|---|---| + | cifar felix | ~0.4s | (varies) | sub-second | inert | sim_rate ≫ 1 | + | fwdllm/plus (sync) | ~1.0s | K=c (barrier) | ~1 GPU pass, ALL commit | inert | `sim_rate` 2.9-3.0 ✓ | + | **fluxtune (async)** | **~4s** | **3 ≪ 10** | **multi-sec → 30s** | **load-bearing → stall** | **`sim_rate` 0.50 ⛔** | + +- **F5 — Hold-to-commit is a CORRECTNESS CHECK, not over-restriction.** A trainer is freed (re-selectable) ONLY + once its returned update is PROCESSED/COMMITTED. Three guards: **(a)** never dispatch a version it already + computed (same-shard×same-version = wasted grad); **(b)** never dispatch while still computing; **(c)** never + dispatch while its returned update is uncommitted. Do NOT weaken it. Concurrency: with **N>C** a fast returner + is held while fresh idle trainers keep C busy; with **N=C** (fluxtune 10/10/3, NO idle pool) concurrency is + agg-cadence-limited and real's **3.37 is a duty cycle** (gpu/max(gpu,D) ≈ 4/12 ≈ 0.34 → N×0.34), NOT + under-utilization. +- **F6 — The bug is purely the COMMIT PATH STALLING; commit RATE is the throughput lever.** Trainers freed only on + commit ⇒ commit rate = free rate = throughput. cifar's commit path is instant → held trainers barely idle; + fluxtune's STALLS (#15) → held trainers idle 30s. Fix = fast/non-stalling commit path; the gate must wait only + on a genuinely-COMPUTING trainer, never a phantom. Hold-to-commit untouched. +- **F7 — Anchors.** Felix gate/HOLD asyncfl:359-448; clamp 461-476; `_sim_hold_busy_slots` 1453-1500; commit-release + 618-627; expected-sct seed 128-137/1668-1671; sct formula (trainer) `async_cifar10/trainer/pytorch/main.py:838-846`; + PARITY.md gate-inert 294-296/321-323, residence §3.resid, past-dating 537-563. + +--- + +## ⭐ #15 fluxtune `sim_rate=0.50` — COMMIT-PATH STALL (phantom fix LANDED+VALIDATED; residual = vclock-omits-aggregate + GPU≈D) + +**Superseded framings (do not revisit):** (1) "GPU-pipelining loss / decouple dispatch" — felix's gate is INERT +(F1). (2) "re-dispatch on return / hold-to-commit is over-restrictive" — hold-to-commit is a CORRECTNESS check +(F5). Final root = a commit-path stall. + +**ROOT (D1/D2 CONFIRMED on `run_20260705_204619` sim / `_202448` real):** hold-to-commit frees a trainer only on +commit, so commit RATE sets throughput (F5/F6). fluxtune's `_sim_recv_min_grad` `earlier_stuck` gate blocks real +wall on a PHANTOM `_sim_inflight_expected` entry — a trainer stamped expected-at-DISPATCH that isn't computing +(idle-in-recv, waiting for weights the gate-blocked single-threaded aggregator can't send) → 30s failsafe → +correctly-held trainers idle ~30s instead of ~one GPU pass → sim_rate 0.50. + +**EVIDENCE:** +- Gate holds ALREADY-ARRIVED grads: every top stall commits a grad that arrived (`[MSG_ARRIVAL]`) 20-23s before + the stall; `buf_depth=7` on 99.7% of commits. Awaited trainers provably `recv_wrapper`-IDLE (0370 idle 13s, + 0371 44.6s spanning the stall). +- Dominant tax: inter-commit gap mean 2.20s; 496 commits pay ~2s + 24 gaps ≥10s + 1×31s failsafe = **1974s = 82% + of the 2398s wall**. `sim_rate` 0.494. +- Closed self-throttle: arrival rate 0.459/s == commit rate 0.456/s. Concurrent compute **sim 1.65 vs real 7.69**. + `recv_wrapper` idle: sim mean 17.9s vs real median 0.01s. Compute mode-invariant (~3.6s). +- Mechanism = **(c1)** optimistic `exp` stamped at DISPATCH (`:2819`, contention-free lower-bound budget) + + **(c3)** `sim_staggered_redispatch=False` → `_sst=_round_now` (`:2813`) collapses fresh-cohort exp onto the + round frontier so it always looks earlier than buffered later-sct grads → `earlier_stuck` stays armed + **(c2)** + single-threaded agg BLOCKS in the `for _pass` loop (`:772`), can't run `_distribute_weights_async` → can't send + weights to the trainer it waits on → deadlock to grace/30s failsafe. `SIM_R1_DISPATCH=0` — hold-to-commit clean. +- Also: fluxtune's self-throttle is the variance-gating keep-training loop — after the initial `Sent 10 WEIGHTS`, + distributes are `0 WEIGHTS + 1–2 VAR=bad` (only 1–2 trainers kept training/cycle). Real frees on grad PROCESS + (recv 0.01s, non-residence path) → ~8 compute; sim defers to commit → slow commit starves re-dispatch. + +**FIX — ✅ P1/P2 LANDED (`sim_compute_truthful_gate`, flag-gated, default off = byte-identical; +fwdllm_aggregator-only → async_cifar10 untouched):** +- **P1:** `_sim_dispatch_wall[end]` stamped at the real `channel.send` (`fwdllm_aggregator.py:2836`). The + `earlier_stuck` gate (`:856`) skips any `_sim_inflight_expected` entry whose last dispatch is older than + `sim_gate_compute_cap_s` (default **10.0s** > ~5.6s max JVP compute) OR never dispatched → a stamped-but-idle + phantom no longer blocks a ready commit; a genuine in-window straggler is STILL held (sct order preserved). + Hold-to-commit, sct-ordered drain, K-D12 carried surplus, K-D27 `_sim_pending_commit` — all UNTOUCHED. + `[SIM_GRAD_RECV]` now emits `phantom_skip=`. Enabled in `expt_scripts/fluxtune_n10_smoke_sim.yaml` + (`sim_compute_truthful_gate: true`, `sim_gate_compute_cap_s: 10.0`). +- **P2:** 4 new `TestComputeTruthfulGate` (phantom skipped; stale-dispatch=phantom; in-window straggler STILL held; + flag-off byte-identical) in `tests/mode/test_fwdllm_sim_grad_loop.py`; +189 fwdllm mode tests green. +- **P3 = RAN 2026-07-06 (pair `run_20260706_112114` sim / `_110555` real).** Phantom fix **VALIDATED**: + `STUCK_EVICT=0`, `phantom_skip`→65 (~1/commit, load-bearing), 30s failsafes gone (max gap 12.8s). **But + `sim_rate` still <1** (steady ~0.49; per-commit `Δvclock/Δwall=0.465`). **Residual root (NEW, non-drain-gate):** + (a) vclock omits `aggregate()` variance compute (sim 32s / real 29s, uncredited — `:1783` folds eval only) + (b) + GPU≈D no-headroom (#1d). Full evidence + resume plan at the PAUSED checkpoint (top). **P4:** §G/§K one-liner + after `sim_rate>1` + parity A/B. + +**Repro (read-only, from `lib/python/examples/fwdllm`):** +```bash +FS=$(ls -td experiments/run_*_fluxtune_n10_smoke_syn_0_sim | head -1); AGG=$(ls "$FS"/*aggregator.log|head -1) +grep SIM_GRAD_RECV "$AGG" | python3 -c "import sys;from datetime import datetime as D;p=None;b=w=t=0 +for l in sys.stdin: + s=D.strptime(l.split(' | ')[0],'%Y-%m-%d %H:%M:%S,%f').timestamp() + if p is not None: + d=s-p;t+=1 + if d>2:b+=1;w+=d + p=s +print(f'commits={t+1} gaps>2s={b} wall_in_waits={w:.0f}s')" # baseline sim: 48% / ~1974s -> expect near-0 +grep SIM_GRAD_STUCK_EVICT "$AGG"; grep SIM_GRAD_RECV "$AGG" | grep -oE 'phantom_skip=[0-9]+' | tail -1 +``` + +--- + +## DONE (landed + tested — terse; do not redo) +- **#14/#1c/#13/#12c fixed:** MQTT join-notify race (#14); R1 two-ledger bridge (K-D27, `SIM_R1_DISPATCH` 238→0, #1c); + drain-stall felix port (K-D28, `sim_rate` 0.06→0.30, #13); `--delay-factor 1` → sync `sim_rate` 0.94→2.9-3.0 (#12c). +- **K-D29 remainder-wait delay model:** real sleeps `max(0,D−gpu)`, sct = `send + max(gpu,D)`, `[TIMING_OVERRUN]` + telemetry (P2-1/6). crc32 straggler offset disabled (P2-3). `perturbation_count` knob, default 10 (P2-5, wired). +- **K-D30 full-cohort determinism gate + `timing_overrun` DIAG rung** (P1-5): un-gates + selection/aggregation_sequence/utility for fwdllm (K=all); fluxtune/fwdllm_plus stay gated. P1-4 assessed redundant. +- **K-D31 canonical `(D, trainer_id)` cohort commit order — VALIDATED (P2-7/P2-7a).** Databin1 `--max-bin 1`: sync + `cohort_sequence` ok=true, set/order/var/cadence=1.0 (the benign delay-tie closed; var already bit-identical). 428 + mode + 9 canon + 115 async parity green; async byte-identical. +- **K-D32 fluxtune JVP perf-opt (`jvp_perf_opt`, config-gated, bit-identical):** trainable-only FD + skip 3 diagnostic + passes + reuse winner JVP → fluxtune −37% (7.57→3.61s mean). NOT retained: vmap (fp32 FD cancellation), fwd-AD. +- **K-D33 pinning:** trainer `[PIN]` self-report, `[LOAD_BALANCE]` check, aggregator GPU pin (least-loaded/idle). 8 + GPUs balanced round-robin (earlier "under-provisioned" read was a misread). 120 launch tests green. +- **Checks/pytests (P1-1/2/3/8):** `cohort_sequence` rung (EXACT, ungated), V2 mean-guard, `--max-bin` window. Banked + ENFORCED ref (full run): fwdllm 41/13/21, fwdllm_plus 36/17/21, fluxtune 35/19/19. +- **#15 P1/P2** (this session) — see the #15 section above. + +## KEY LEARNINGS (durable) +- **fluxtune #15 = commit-path stall, not residence/pipelining/GPU.** Hold-to-commit is correct (F5); commit RATE is + the throughput lever (F6). Felix's gate is inert (F1) — cifar hides the cost via sub-second compute. +- **SYNC divergence root = commit ORDER via split-half var, but exact cadence has a bin-7 float-nondeterminism + wall.** fwdllm var is a split-half stat over the commit-ordered grad list → wrong order → wrong var → `var<0.3` + flips → per-trainer RNG (seeded once, never reset) desyncs → grads diverge ~1%. Grads ARE deterministic given + matched order ⇒ exact parity achievable ≤bin 1; beyond ~bin 6, ~1e-3 GPU fp16 jitter amplified by the variance + ratio breaks it → target must go DISTRIBUTIONAL (open task 1). +- **Order made deterministic** by remainder-wait `max(gpu,D)` + per-trainer D ≫ GPU (K-D29) + canonical tie-break + (K-D31). fwdllm is forward-ONLY → the GPU lever is `perturbation_count`, not backprop. + +## ANCHORS +- Split-half var `aggregator/.../fwdgrad_utils.py:133-158`; commit-order append `fwdllm_aggregator.py:718-721`. +- RNG-once seed `.../tc_transformer_trainer_distribute.py:222-225`. +- Delay/sct `FedSgdTrainer.py:510-538`(sleep)/`:623-629`(sct). Sim drain `fwdllm_aggregator.py:_sim_recv_min_grad` + (~752-970); dispatch stamp `:2819`; send `:2836`; compute-truthful gate `:856`. +- Checker rungs `async_cifar10/scripts/parity/checks.py` — `aggregation_sequence`, `inter_arrival_order`, + `v2 var_trajectory`, gating `DETERMINISTIC_SELECTORS`, `cohort_sequence_parity`. Standalone diff + `expt_scripts/logical_parity.py`. + +## RESOLVED DECISIONS (operator, 2026-07-05/06) +1. Delay model → remainder-wait (`wall = max(gpu,D)`), aligning to async_cifar10. +2. Delay magnitude → configurable; full registry (`--delay-factor 1`) for the current sync runs. +3. GPU overhead → optimize now (K-D32 landed); forward < backprop held with margin. +4. Checker rungs → hard-FAIL (EXACT + V2 mean-guard, ungated). +5. #15 fix → **compute-truthful gate** (NOT re-dispatch-on-return / NOT weaken hold-to-commit); **flag-gated + P3 + parity gate**; raise `sim_gate_compute_cap_s` if parity moves. diff --git a/lib/python/examples/fwdllm/aggregator/FedSgdAggregator.py b/lib/python/examples/fwdllm/aggregator/FedSgdAggregator.py index 20bbd2d42..65897188c 100755 --- a/lib/python/examples/fwdllm/aggregator/FedSgdAggregator.py +++ b/lib/python/examples/fwdllm/aggregator/FedSgdAggregator.py @@ -201,6 +201,24 @@ def aggregate(self, current_round): logger.info(f"self.var = {self.var}") logger.info(f"snr of jvps = {self.snr}") logger.info(f"coefficient of variation = {c_of_variation}") + + # Opt-2 (charter §5c/§5e): variance-plateau force-commit. Under the + # 'plateau' policy, additionally force a commit once the per-bin variance + # curve has flattened (relative drop over the last N cycles < rel_delta) + # while var is still above threshold -- more denoising buys nothing, so + # commit the denoised estimate. var_prev_iter_list is the per-bin var + # history (reset on commit), already including this cycle's var. Policy + # off/absent => never arms a commit => byte-identical. + self._force_commit_reason = None + self._plateau_fired_this_cycle = self._should_force_commit_on_plateau() + if self._plateau_fired_this_cycle: + self._force_commit_this_cycle = True + logger.info( + f"[VarPlateau] curve flattened over N=" + f"{getattr(self, '_var_plateau_patience', 3)} " + f"(var={self.var_prev_iter_list[-1]:.4f} > " + f"thr={self.var_threshold}); force-committing at plateau." + ) logger.debug( f"self.grad_for_var_check_list size: {len(self.grad_for_var_check_list)}" ) @@ -314,6 +332,7 @@ def aggregate(self, current_round): f"keeping weight update, clearing cached_v." ) self.var_good_enough = True + self._force_commit_reason = "natural" # 方差满足要求 self.cached_v = [] elif _force_commit: @@ -352,9 +371,14 @@ def aggregate(self, current_round): self.last_round_update = [ p.clone().detach() for p in weighted_gradient_sum ] + self._force_commit_reason = ( + "plateau" if getattr(self, "_plateau_fired_this_cycle", False) + else "cap" + ) logger.info( f"[MaxIterBypass] Variance FAILED (var={self.var} > " - f"thr={self.var_threshold}) but force-commit is set; " + f"thr={self.var_threshold}) but force-commit is set " + f"(reason={self._force_commit_reason}); " f"committing weights anyway, clearing cached_v." ) self.var_good_enough = True diff --git a/lib/python/examples/fwdllm/configs/trainer_base.yaml b/lib/python/examples/fwdllm/configs/trainer_base.yaml index 16b433cd2..95f9824cf 100644 --- a/lib/python/examples/fwdllm/configs/trainer_base.yaml +++ b/lib/python/examples/fwdllm/configs/trainer_base.yaml @@ -81,6 +81,12 @@ hyperparameters: var_control: true perturbation_sampling: true select_perturbation_using_jvp: false + # Fluxtune JVP perf optimizations (simulate_fwdllm.md §L) — trainable-only + # finite difference + skip diagnostic-only forward passes + reuse the selected + # perturbation's JVP. All BIT-IDENTICAL to the grads (validated by + # scripts/profile_jvp_opt.py). Default false = byte-identical; the fluxtune + # yamls set it true. Must match real<->sim (both read this config). + jvp_perf_opt: false evaluate_during_training_steps: 100 frequency_of_the_test: 1 is_debug_mode: 0 diff --git a/lib/python/examples/fwdllm/docs/aggregation_design.md b/lib/python/examples/fwdllm/docs/aggregation_design.md new file mode 100644 index 000000000..70cc85fd1 --- /dev/null +++ b/lib/python/examples/fwdllm/docs/aggregation_design.md @@ -0,0 +1,176 @@ +# FluxTune aggregation (C3) — design starter + +**What this is.** A starter design doc for FluxTune's **intelligent (gradient-aware) aggregation** +contribution (C3): (1) the async-FL aggregation regime *today* — FedBuff → FeLiX → current FluxTune, with +code anchors; (2) the **hypothesis** for why the current rule is probably wrong for FluxTune; (3) the +**design space** for weighing updates. Seed for the C3 investigation in +[`EXPERIMENTS.md`](../EXPERIMENTS.md) §4 and [`EXPTS_CHARTER.md`](../EXPTS_CHARTER.md) §2d; companion +efficiency contribution is [`docs/dynamic_kc_design.md`](dynamic_kc_design.md) (C2). + +**Status:** design only — nothing here is implemented; current FluxTune uses the FeLiX rule (below). +Any change here is **baseline-affecting** → operator sign-off, flag-gated, byte-identical off, and must +preserve real↔sim parity ([`simulate_fwdllm.md`](../simulate_fwdllm.md)). + +--- + +## 1. The regime today + +All three run through the **FedBuff optimizer** ([`flame/optimizer/fedbuff.py`](../../../flame/optimizer/fedbuff.py)). +An update is a **delta** buffered until `agg_goal` arrive; each is multiplied by a **scalar rate**, summed, then +scaled by the server LR and added to base weights (`do()` → `aggregate_fn(tres, rate)` → `scale_add_agg_weights` +at `fedbuff.py:181-228`). The three differ only in **how `rate` is computed.** + +| | rate formula | signals used | lineage | +|---|---|---|---| +| **FedBuff** | `rate = 1 / √(1 + Δv)`, `Δv = version − trainer_version` | **staleness only** (round-based) | `agg_rate_conf.type = "old"` (`fedbuff.py:180`) | +| **FeLiX** | `rate = scale·α(Δv) + (1−scale)·β(u)` | staleness **+** statistical utility | `type = "new"` (`fedbuff.py:183-200`); async_cifar10 / REFL | +| **FluxTune (now)** | **same as FeLiX** + a **variance commit-gate** | staleness + utility, gated on grad variance | FeLiX rate + `var < var_threshold` gate | + +Where, for the FeLiX/FluxTune "new" rate (`fedbuff.py:91-139`): +- `α(Δv) = 1 / (1 + Δv)^a_exp` — **staleness decay** (polynomial; `alpha_exponential` also available). +- `β(u) = 1 − 1/(1+u)^b_exp + 0.5` — **utility upshift**, increasing in `u`. +- `u = tres.stat_utility` = the **Oort statistical utility** (`I_m` ≈ the client's average training loss + in its last engaged round — `flame/selector/feddance.py:7`). So β rewards high-loss (informative) clients. +- FluxTune config: `scale=0.4, a_exp=0.25, b_exp=0.1` (`_metadata/baselines.yaml:457-461`); + `stalenessPolicy=none` (down-weight, never *reject*, `baselines.yaml:473`). + +**The variance gate (FluxTune only).** Before an update is eligible, `FedSgdAggregator.aggregate()` computes +a batch of gradient statistics and commits the data-bin only when variance is under threshold +(`aggregator/FedSgdAggregator.py:194-202, 132`): +- `calculate_var(grads)`, `calculate_real_var(jvps)`, `calculate_snr(jvps)`, + `calculate_snr_gradients(grads)`, `calculate_cv(grads)`. +- commit iff `self.var < self.var_threshold` (`var_threshold`: distilbert default **0.1**, overridden to + **0.3** for the baseline; `FedSgdAggregator.py:73-89`, `baselines.yaml:475`). + +**Key structural fact:** these rich per-update **gradient statistics are computed and then discarded for +weighting** — they gate whether the bin commits, but the *rate* that weights each update ignores them. + +--- + +## 2. Why this is probably the wrong rule for FluxTune (the hypothesis) + +Observed in N=100 smoke runs (E3/E4): FluxTune spends ~5–6× the forward-pass compute of the oracular sync +baseline for *less* Δloss per unit — its async concurrency does more *unproductive* work. The aggregation rule +is one suspect. Three reasons it likely under-differentiates FluxTune's updates: + +**H1 — Weights ≠ gradients.** FedBuff/FeLiX target **weight-averaging** async FL: the buffered delta is a +*weight update*, so scalar down-weighting is sensible. FluxTune's update is a **forward-mode gradient estimate** +(a JVP-scaled direction). A **scalar** rate changes *magnitude*, not *direction quality*: two opposite-pointing +estimates are still **averaged**, not filtered — a wasteful/harmful combination scalar weighting can't express. +Gradients want a *direction-aware* combiner, not a magnitude knob. + +**H2 — Staleness doesn't move fast enough to differentiate.** `Δv = agg_model_version − trainer_version`, and +`_model_version` advances **per data-bin completion**, itself **gated on the variance threshold** +(`fwdllm_aggregator.py:1422`, `fedbuff.py:194`). So version advances **slowly and in lockstep**, concurrent +in-flight updates carry **near-identical `Δv`** → `α(Δv) ≈ constant` across the buffer → the **staleness axis +adds almost no differentiation**. The rate collapses onto `β(u)` (Oort-loss) alone — one coarse, quantized +signal. Net: the "tradeoff" is mostly one-dimensional, and that dimension is weak. + +**H3 — The best signals are computed but unused.** Each update's gradient/JVP variance/SNR/CV already exist +(§1) but only drive the binary gate. They are exactly *how trustworthy* an estimate is — the natural weighting +signal for noisy gradient estimates — and they're thrown away. + +--- + +## 3. Design space — how to actually weigh gradient updates + +Think along five axes; a concrete scheme picks one option per axis. + +**Axis A — what signal(s) to weight by.** +- **Staleness** — but on *which clock*? (see Axis D). Current: `Δmodel_version` (coarse). +- **Statistical utility** — Oort `I_m` (now), or test-loss delta, or per-update loss reduction. +- **Gradient magnitude** — `|JVP|` (FluxTune already selects the max-|JVP| direction in C1; magnitude is a + proxy for informative descent). +- **Gradient trust — variance / SNR / CV** (already computed): inverse-variance ≈ precision, the + statistically optimal weight for combining noisy unbiased estimates. +- **Direction agreement** — cosine of the update against the *running aggregate* (or against the buffer's + mean). Anti-aligned updates (negative cosine) are candidates to reject or down-weight. +- **Representativeness** — client sample count / class coverage (non-IID aware). + +**Axis B — how to combine (the core "weights≠gradients" choice).** +- **B0 scalar rate** (current): one scalar per update, magnitude-only. +- **B1 scalar but gradient-aware**: keep scalar form, feed it the trust/alignment signals (cheap upgrade). +- **B2 inverse-variance / precision weighting**: `w_i ∝ 1/var_i` (or `SNR_i`) — principled for noisy JVP + estimates; uses signals already on hand. +- **B3 direction-aware / alignment**: project updates onto the running descent direction; weight by (or + gate on) cosine agreement; optionally drop anti-aligned components rather than average them. +- **B4 coordinate/block-wise**: per-parameter-block trust instead of one global scalar (adapters are + small — feasible). + +**Axis C — reject vs down-weight vs hybrid.** Current: a **variance gate** (hard reject at bin level) + +soft down-weight (rate). Options: loosen the gate and push the discrimination into the *weight*; or keep a +hard gate on *direction* (anti-aligned) and soft-weight on *trust*. + +**Axis D — the staleness clock (fixes H2).** Replace/augment `Δmodel_version` with a clock that differentiates +concurrent updates: +- `#commits since dispatch` (aggregations while this update was in flight), +- wall-age or vclock-age of the update, +- `#iterations` since dispatch, +- or **normalize** `Δv` by the current data-bin progression rate so a slow, variance-gated clock still yields + spread. Quantify the staleness *distribution* first (§4) before picking. + +**Axis E — normalization & LR coupling.** Do the weights sum to 1 / to `agg_goal`? A change in weight scale +silently rescales the effective server LR (`scale_add`, `fedbuff.py:230+`), so re-normalize or re-tune LR together. + +--- + +## 4. Candidate schemes to prototype + +All flag-gated, A/B against S0, measured on E1 (time-to-acc) and E3 (Δloss per forward-pass). + +| id | scheme | axes | rationale | +|---|---|---|---| +| **S0** | current FeLiX scalar rate | B0 | control | +| **S1** | **inverse-variance** weighting `w∝1/var` (or `SNR`) | A:trust, B2 | principled for noisy JVP estimates; reuses §1 signals; directly attacks H3 | +| **S2** | **alignment-gated**: reject/down-weight anti-aligned (cosine<0 vs running aggregate), weight by \|cos\| | A:agreement, B3, C | attacks H1 — stops averaging opposing gradients | +| **S3** | **magnitude×SNR** `w∝\|JVP\|·SNR`, staleness-free | A:magnitude+trust, B1 | isolates informativeness from the dead staleness axis (H2) | +| **S4** | **finer staleness clock** (Δcommits or wall-age) in the FeLiX form | A:staleness, D | tests whether H2 alone explains the collapse | +| **S5** | **hybrid**: inverse-variance × alignment × finer-staleness | B2+B3+D | the "everything" combiner if S1–S4 each help | + +--- + +## 5. How to know if it's working (evaluate the *aggregation*, not just the run) + +Before E1/E3 outcomes, instrument the aggregation itself: +- **Staleness spread** — distribution/entropy of `Δv` across the buffer at each commit. If ~degenerate, + H2 confirmed → a finer clock (Axis D) is mandatory. +- **Weight spread** — distribution/entropy of the assigned `rate`s. If FeLiX assigns near-uniform weights, + it's not differentiating regardless of formula. +- **Weight ↔ usefulness correlation** — does a higher assigned weight predict a larger realized Δloss + contribution? (post-hoc, per committed update). The real test of a good rule. +- **Wasted-work fraction** — share of committed forward-pass compute in low-weight / anti-aligned / dropped + updates (ties back to the E3/E4 root cause). + +--- + +## 6. Interactions & guardrails +- **C1 variance gate** — already filters high-variance updates; a variance-based *weight* (S1) risks + **double-counting**. Decide: loosen the gate and move discrimination into the weight, or keep the gate for + commit-timing and weight on an orthogonal signal (alignment/magnitude). +- **C2 dynamic K/C** ([`dynamic_kc_design.md`](dynamic_kc_design.md)) — concurrency `C` sets how many + stale/in-flight updates coexist, i.e. how much there is to differentiate. High C makes a good rule matter + more; the two together set the wasted-work E3/E4 measure. +- **Fidelity / sim parity** — any rule change is baseline-affecting: flag (byte-identical off), operator + sign-off, re-validate real↔sim parity (variance cadence / commit order are parity-checked — a new weight that + reorders commits moves parity). +- **Server LR coupling** — re-normalize weights or re-tune LR together (Axis E). + +--- + +## 7. Open questions +- Is the right object a **scalar weight** at all, or a **direction-aware combine** (B3/B4)? H1 says maybe not. +- Should stale/anti-aligned updates be **rejected** (waste the compute but protect the model) or **kept and + corrected**? Interacts with C2 (whether to have dispatched them at all). +- Does the Oort `stat_utility` (training-loss proxy) even correlate with *gradient* usefulness here, or + should utility be redefined in gradient terms (|JVP|, SNR)? +- What's the cheapest signal that recovers most of the differentiation — so C3 doesn't add aggregator cost + that eats FluxTune's speed win? + +--- + +## 8. Anchors +- FedBuff rates + apply: [`flame/optimizer/fedbuff.py`](../../../flame/optimizer/fedbuff.py) — `alpha_*`/`beta_*` `:91-139`, `weight_factor` `:110-139`, `do()` rate select `:178-207`, `scale_add`/LR `:215-270`. +- Variance/SNR/CV signals + gate: [`aggregator/FedSgdAggregator.py`](../aggregator/FedSgdAggregator.py) `:194-202` (compute), `:132` (`var launches a run_set's (baseline x condition) matrix +# * compare_baselines.py -> maps each analysis to the run dirs + metric ids it needs +# +# Human twin: EXPERIMENTS.md (intent & why). This file is the source of truth for +# WHAT RUNS and HOW METRICS ARE WIRED. Metric *implementations* live in +# compare_baselines.py, keyed by the `id`s below -- this file is the contract, not +# the code. Keep in sync with EXPERIMENTS.md (pre-flight-gate discipline). +# ============================================================================ + +version: 1 + +# ---- run-time defaults (per-condition overridable) ------------------------- +defaults: + target_accuracy: 0.84 # tau: convergence-stop accuracy threshold + converge_window: 20 # W : consecutive data bins all >= tau -> stop + stop: converge # converge | data_id | runtime + # safety caps (a non-converging run is bounded by these -> DID_NOT_CONVERGE) + max_runtime_s: 172800 # 48h wall ceiling — a convergence run needs room to reach tau + max_data_id_progress: 9999 + # stall guard: terminate EARLY (before 48h) if the run is clearly not learning — + # no PROGRESS within stall_window_s on the armed signal (stall_on). + # -> verdict STALLED (distinct from DID_NOT_CONVERGE). 0 window disables it. + stall_window_s: 7200 # 2h with no progress ⇒ stalled + stall_min_delta: 0.01 # acc: 1% ABSOLUTE test-accuracy gain counts as progress + stall_on: either # acc | loss | either — 'either' keeps a run alive while EITHER + # accuracy OR test-loss is still improving (a run can plateau in + # accuracy while loss keeps falling ⇒ still learning, don't kill it) + loss_min_rel_delta: 0.01 # loss: 1% RELATIVE drop vs running-best counts as progress + # (loss is unbounded, so progress is relative, not absolute) + +# ---- baselines (descriptive; substance lives in _metadata/baselines.yaml) --- +baselines: + fwdllm: { sync: true, selector: random, tracking: default, agg_goal_default: 10 } + fwdllm_plus: { sync: true, selector: random, tracking: oracular, agg_goal_default: 10 } + fluxtune: { sync: false, selector: async_oort, tracking: client_notify, agg_goal_default: 3 } + +# ---- run-sets: the expensive GPU work, produced ONCE, reused by all analyses - +run_sets: + main: + description: > + Canonical comparison run-set. Experiment 1 defines it; Experiments 2-5 are + pure reducers over these run dirs (no new GPU runs). One run per baseline. + baselines: [fwdllm, fwdllm_plus, fluxtune] + # Active condition = the real comparison (operator-signed-off 2026-07-06): + # N=100, 84% target, syn_0 (Phase 1), data heterogeneity alpha=1, modeled + # mobile-latency delays ON at delay_factor=2. agg_goal MATCHES across all three + # baselines (=10); fluxtune's concurrency C=30 (async) while sync baselines C=10. + condition: + N: 100 # trainer.num_trainers + K: 10 # selector.kwargs.k (all baselines) + C: { sync: 10, async: 30 } # selector.kwargs.c — fluxtune uses async=30, sync baselines=10 + # agg_goal is fanned from C.sync -> 10 for EVERY baseline (matched across + # baselines, per sign-off). fluxtune therefore uses agg_goal=10 (NOT its + # native fedbuff 3) with C=30. A pre-flight check enforces the match. + partition_method: niid_label_clients=100_alpha=1 # data heterogeneity alpha=1 (was 0.1) + avail_trace: syn_0 # 100% available (Phase 1); -> mobiperf_* at Phase 2 + delays: on # enable_training_delays, matched both sides (K-D8) + delay_factor: 2 # divides the yaml base compute-latency delay (full/2) — mobile-latency model + target_accuracy: 0.84 # convergence stop threshold + converge_window: 20 # W consecutive bins >= target + scale_n100: # target config, DEFERRED until pipeline proven + N: 100 + K: 10 + C: { sync: 10, async: 30 } + partition_method: niid_label_clients=100_alpha=0.1 + avail_trace: mobiperf_2st + # NOTE: fwdllm_plus sync barrier cannot assemble under mobiperf unavailability + # (agg_goal >= available trainers). Needs a policy decision before this runs; + # the pre-flight gate will BLOCK it otherwise. See EXPERIMENTS.md sec 1. + _blocked_pending: fwdllm_plus_mobiperf_policy + +# ---- analyses: every experiment is a view over a run_set -------------------- +# reuse: run_set analyses share the SAME run dirs. `metrics[].id` are implemented +# in compare_baselines.py. `source` documents telemetry provenance: +# emit = already in telemetry derive = reducer over existing telemetry +# ws2 = convergence watcher ws3a = network telemetry (new) +# ws3b = perturbation counter (new) +analyses: + - id: 1_time_to_target + takeaway: Fluxtune reaches target accuracy faster than FwdLLM and FwdLLM_Plus. + run_set: main + reuses: null # this analysis DEFINES the run-set + metrics: + - { id: time_to_target, source: ws2, report: [wall_s, vclock_s, rounds, data_bins] } + - { id: max_accuracy, source: derive, report: [accuracy] } + + - id: 2_resource_utilization + takeaway: Fluxtune cuts wait time at trainers (primary) and the aggregator (secondary). + run_set: main + reuses: main # NO new runs -- reuse Expt-1 telemetry + approach: derived_time_fraction # operator call: no hardware sampler + metrics: + - { id: trainer_busy_fraction, source: derive, report: [p50, p90, p99] } + - { id: agg_busy_vs_wait, source: derive, report: [busy_frac, barrier_wait_frac, drain_frac] } + + - id: 3_compute_productivity + takeaway: At constrained clients, Fluxtune yields more learning per unit compute. + run_set: main + reuses: main + metrics: + - { id: delta_loss, source: emit, report: [first_loss, final_loss, delta] } + - { id: productivity_per_gpu_s, source: derive, report: [delta_loss_per_gpu_s] } + - { id: productivity_per_fwd_pass, source: ws3b, report: [delta_loss_per_forward_pass] } + + - id: 4_network_data + takeaway: Fluxtune's total data overhead is lower despite more messages per round. + run_set: main + reuses: main # requires ws3a telemetry present AT RUN TIME + requires_telemetry: [ws3a] + metrics: + - { id: total_messages, source: ws3a, report: [agg_side, client_side] } + - { id: total_bytes, source: ws3a, report: [agg_side, client_side] } + - { id: msg_size_dist, source: ws3a, report: [p50, p90, p99, histogram] } + + - id: 5_client_sessions + takeaway: Fluxtune's client sessions are much shorter than FwdLLM's. + run_set: main + reuses: main + session_def: # active-session semantics (operator-precise) + async: dispatch_to_commit # selected -> next reselection (contributor_intervals) + sync: one_round_span # round-based reselect: one-round duration + metrics: + - { id: session_duration, source: derive, report: [p50, p90, p99, histogram] } + - { id: participation_counts, source: derive, report: [rounds, data_bins, iterations] } + +# ---- outputs ---------------------------------------------------------------- +outputs: + table_csv: compare_baselines.csv # one row per (analysis, baseline, metric) + overlay_plots: true # N-way accuracy/comm/util overlays (extends compare_streaming) + run_manifest: run_manifest.json # analysis -> run dirs + config fingerprint (mix-guard) diff --git a/lib/python/examples/fwdllm/expt_scripts/.gitignore b/lib/python/examples/fwdllm/expt_scripts/.gitignore new file mode 100644 index 000000000..5ab7371d4 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/.gitignore @@ -0,0 +1,4 @@ +__pycache__/ +*.pyc +paper_figs/ +paper_figs_ablation/ diff --git a/lib/python/examples/fwdllm/expt_scripts/audit_weight_redundancy.py b/lib/python/examples/fwdllm/expt_scripts/audit_weight_redundancy.py new file mode 100644 index 000000000..ca075c10b --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/audit_weight_redundancy.py @@ -0,0 +1,118 @@ +#!/usr/bin/env python3 +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Audit intra-databin weight-download redundancy from a run's telemetry. + +Charter §5c Opt-1: within a data-bin the model_version is constant and the full +WEIGHTS payload is byte-identical across iterations, so a trainer should receive +it AT MOST ONCE per bin (later dispatches get the tiny VAR=bad "keep training" +message). Measures, per data-bin, how many full-weight sends each trainer got. + +Robust to the `data_id` cycling footgun (data_id repeats within a round): the +true data-bin id is reconstructed by counting COMMITS (`agg_round` with +`var_good_enough=True`) in stream order; weight-sends for data-bin N precede +commit N. + +Usage: + python audit_weight_redundancy.py [--max-redundant-frac F] +Exit 1 if the redundant fraction exceeds --max-redundant-frac (default off), so +this doubles as a CI regression check over a smoke run. +""" +import argparse +import glob +import json +import os +import statistics as st +import sys +from collections import defaultdict + + +def _agg_jsonl(path): + if os.path.isdir(path): + hits = glob.glob(os.path.join(path, "telemetry", "aggregator_*.jsonl")) + if not hits: + sys.exit(f"no aggregator_*.jsonl under {path}/telemetry/") + return hits[0] + return path + + +def audit(path): + f = _agg_jsonl(path) + databin = 0 # running commit count == true data-bin id + # databin -> peer -> [n_weights, n_varbad] + sends = defaultdict(lambda: defaultdict(lambda: [0, 0])) + wbytes = defaultdict(int) + with open(f) as fh: + for line in fh: + # cheap prefilter before json.loads (files run >1 GB) + is_comm = '"comm"' in line and "agg_to_trainer" in line + is_round = '"agg_round"' in line + if not (is_comm or is_round): + continue + try: + e = json.loads(line) + except Exception: + continue + ev = e.get("event") + if ev == "agg_round": + if e.get("var_good_enough"): + databin += 1 # commit -> next data-bin starts + elif ev == "comm" and e.get("direction") == "agg_to_trainer": + pk = e.get("payload_kind") + peer = e.get("peer_id") + if pk == "weights": + sends[databin][peer][0] += 1 + wbytes[databin] += e.get("size_bytes", 0) + elif pk == "var_bad": + sends[databin][peer][1] += 1 + return f, sends, wbytes + + +def report(path): + f, sends, wbytes = audit(path) + if not sends: + print(f"{path}: no agg_to_trainer comm events found (WS3-a telemetry missing?)") + return None + uniq, wsent, redun_sends, redun_bytes = [], [], 0, 0 + tot_w = tot_vb = 0 + n_bins_with_repeat = 0 + for db, peers in sends.items(): + u = len(peers) + w = sum(p[0] for p in peers.values()) + vb = sum(p[1] for p in peers.values()) + r = sum(max(0, p[0] - 1) for p in peers.values()) # sends beyond first/peer + uniq.append(u); wsent.append(w); tot_w += w; tot_vb += vb + redun_sends += r + if r > 0: + n_bins_with_repeat += 1 + if w: + redun_bytes += (wbytes[db] / w) * r + n = len(sends) + frac = redun_sends / tot_w if tot_w else 0.0 + print(f"===== {os.path.basename(os.path.dirname(os.path.dirname(f))) or f} =====") + print(f"data-bins (reconstructed via commit count): {n}") + print(f"unique trainers/bin: median={st.median(uniq):.0f} max={max(uniq)}") + print(f"weight-sends/bin: median={st.median(wsent):.0f} max={max(wsent)}") + print(f"total weight-sends={tot_w} var_bad-sends={tot_vb}") + print(f"data-bins with a trainer sent weights 2+ times: {n_bins_with_repeat}/{n} " + f"({100*n_bins_with_repeat/n:.1f}%)") + print(f"REDUNDANT weight-sends (same trainer 2+ in a bin): {redun_sends} " + f"({100*frac:.1f}% of weight-sends)") + print(f"REDUNDANT weight-bytes: {redun_bytes/1e9:.2f} GB") + return frac + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("path", help="run dir or aggregator_*.jsonl") + ap.add_argument("--max-redundant-frac", type=float, default=None, + help="exit 1 if redundant weight-send fraction exceeds this") + a = ap.parse_args() + frac = report(a.path) + if a.max_redundant_frac is not None and frac is not None and frac > a.max_redundant_frac: + print(f"FAIL: redundant fraction {frac:.3f} > {a.max_redundant_frac}") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/lib/python/examples/fwdllm/expt_scripts/characterize_variance_curve.py b/lib/python/examples/fwdllm/expt_scripts/characterize_variance_curve.py new file mode 100644 index 000000000..93bcd93fe --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/characterize_variance_curve.py @@ -0,0 +1,355 @@ +#!/usr/bin/env python3 +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Characterize the per-databin variance-decay curve from a run's telemetry. + +Charter §5c Opt-2: before designing the variance-gate `stopping_policy` +(fixed_cap / plateau / adaptive), measure the actual decay curve -- where var +starts, how fast it falls, where it plateaus, and whether the plateau drops as +training progresses. + +Streams the aggregator telemetry, reconstructs each data-bin (robust to the +`data_id` cycling footgun via COMMIT counting -- see audit_weight_redundancy.py), +extracts each bin's (iteration -> var) curve, and runs two counterfactual sweeps: + + * fixed_cap sweep -- for a cap K: iterations/forward-passes saved, bins + affected, and how much worse the committed var is vs the natural commit var + (the accuracy risk of committing early). + * plateau sweep -- for (patience N, tol eps): where a diminishing-returns + rule fires and at what var, vs the fixed cap. + +Usage: + python characterize_variance_curve.py [--caps 7,8,10,12,15] + [--plateau-N 3 --plateau-eps 0.05,0.10,0.15] [--json OUT.json] + +Pure read; streaming/prefilter style so it runs over multi-GB logs. +""" +import argparse +import glob +import json +import os +import statistics as st +import sys +from collections import defaultdict + + +def _agg_jsonl(path): + if os.path.isdir(path): + hits = glob.glob(os.path.join(path, "telemetry", "aggregator_*.jsonl")) + if not hits: + sys.exit(f"no aggregator_*.jsonl under {path}/telemetry/") + return hits[0] + return path + + +def _pct(xs, q): + """Nearest-rank percentile (q in [0,100]); None on empty.""" + if not xs: + return None + s = sorted(xs) + if q <= 0: + return s[0] + if q >= 100: + return s[-1] + k = max(0, min(len(s) - 1, int(round((q / 100.0) * (len(s) - 1))))) + return s[k] + + +def load_curves(path, max_databins=None): + """Stream agg_round events -> {databin_id: [(iteration, var, committed), ...]}. + + Data-bin id is the running commit count: a bin ends when an agg_round commits + (`var_good_enough=True`). Uses `cycle_iteration` (unambiguous, unlike the + post-mutation `iteration_per_data_id`) as the intra-bin axis. `max_databins` + caps to the first N committed bins for like-for-like truncation across runs. + + Also returns each committed bin's `commit_reason` (Opt-2 telemetry: + natural / cap / plateau; None on legacy runs). + """ + f = _agg_jsonl(path) + databin = 0 + curves = defaultdict(list) # databin -> [(iter, var, committed)] + reasons = {} # databin -> commit_reason on the committing cycle + thr_seen = set() + with open(f) as fh: + for line in fh: + if '"agg_round"' not in line: + continue + try: + e = json.loads(line) + except Exception: + continue + if e.get("event") != "agg_round": + continue + var = e.get("var") + it = e.get("cycle_iteration") + if it is None: + it = e.get("iteration_per_data_id") + committed = bool(e.get("var_good_enough")) + thr = e.get("var_threshold") + if thr is not None: + thr_seen.add(thr) + if var is not None: + curves[databin].append((it, float(var), committed)) + if committed: + reasons[databin] = e.get("commit_reason") + databin += 1 + if max_databins is not None and databin >= max_databins: + break + # sort each bin by iteration; drop a trailing open (never-committed) bin + out = {} + for b, pts in curves.items(): + if max_databins is not None and b >= max_databins: + continue + out[b] = sorted(pts, key=lambda p: (p[0] if p[0] is not None else 0)) + threshold = min(thr_seen) if thr_seen else None + return out, threshold, reasons + + +def _plateau_onset(vars_, N, eps): + """First iteration index i (>=N) where the relative var drop over the last N + steps is < eps -- i.e. diminishing returns. Returns (onset_idx, var_at_onset) + or (None, None) if it never plateaus within the bin.""" + for i in range(N, len(vars_)): + prev = vars_[i - N] + cur = vars_[i] + if prev <= 0: + continue + rel_drop = (prev - cur) / prev + if rel_drop < eps: + return i, cur + return None, None + + +def characterize(curves, threshold, caps, plateau_N, plateau_eps_list, reasons=None): + committed_bins = {b: pts for b, pts in curves.items() if any(c for _, _, c in pts)} + n_bins = len(committed_bins) + + # Opt-2 commit-reason split (natural gate / max-iter cap / plateau); None on + # legacy runs (policy off) -> reported as "natural(legacy)". + reason_counts = defaultdict(int) + if reasons: + for b in committed_bins: + r = reasons.get(b) + reason_counts["natural(legacy)" if r is None else r] += 1 + + per_bin = [] # list of dicts, in databin order + for b in sorted(committed_bins): + pts = committed_bins[b] + vars_ = [v for _, v, _ in pts] + n_iters = len(pts) + commit_idx = next((i for i, (_, _, c) in enumerate(pts) if c), n_iters - 1) + v0 = vars_[0] + v_commit = vars_[commit_idx] + v_min = min(vars_) + # plateau onset at the *reference* (N, eps[0]) for characterization + onset, v_onset = _plateau_onset(vars_, plateau_N, plateau_eps_list[0]) + per_bin.append({ + "databin": b, + "n_iters": n_iters, + "commit_idx": commit_idx, # 0-based iteration at which it committed + "var_initial": v0, + "var_commit": v_commit, + "var_min": v_min, + "plateau_onset": onset, + "var_at_plateau": v_onset, + }) + + def dist(key): + xs = [d[key] for d in per_bin if d[key] is not None] + return {"p10": _pct(xs, 10), "p50": _pct(xs, 50), + "p90": _pct(xs, 90), "mean": (st.mean(xs) if xs else None), "n": len(xs)} + + summary = {k: dist(k) for k in + ["n_iters", "commit_idx", "var_initial", "var_commit", + "var_min", "plateau_onset", "var_at_plateau"]} + + # --- evolution over training: split databins into thirds by order --- + thirds = {} + order = sorted(committed_bins) + if order: + third = max(1, len(order) // 3) + buckets = {"early": order[:third], "mid": order[third:2 * third], "late": order[2 * third:]} + idx = {d["databin"]: d for d in per_bin} + for name, bs in buckets.items(): + iters = [idx[b]["n_iters"] for b in bs] + vcommit = [idx[b]["var_commit"] for b in bs] + vmin = [idx[b]["var_min"] for b in bs] + onset = [idx[b]["plateau_onset"] for b in bs if idx[b]["plateau_onset"] is not None] + thirds[name] = { + "n_bins": len(bs), + "median_n_iters": _pct(iters, 50), + "median_var_commit": _pct(vcommit, 50), + "median_var_min": _pct(vmin, 50), + "median_plateau_onset": _pct(onset, 50), + } + + # --- counterfactual fixed_cap sweep --- + # A cap K forces commit at iteration index min(K-1, natural_commit_idx). We only + # SAVE work on bins whose natural commit happened later than K-1. + cap_sweep = [] + total_natural_iters = sum(d["n_iters"] for d in per_bin) + for K in caps: + saved = 0 + affected = 0 + for d in per_bin: + nat_idx = d["commit_idx"] + if nat_idx > K - 1: # bin grinds past the cap -> we cut it + affected += 1 + saved += nat_idx - (K - 1) + cap_sweep.append({ + "cap": K, + "bins_affected": affected, + "frac_bins_affected": (affected / n_bins) if n_bins else None, + "iters_saved": saved, + "frac_iters_saved": (saved / total_natural_iters) if total_natural_iters else None, + }) + + # need the raw vars to report forced-commit var quality; recompute cleanly + cap_quality = [] + for K in caps: + forced = [] + for b in sorted(committed_bins): + pts = committed_bins[b] + vars_ = [v for _, v, _ in pts] + commit_idx = next((i for i, (_, _, c) in enumerate(pts) if c), len(pts) - 1) + if commit_idx > K - 1 and len(vars_) >= K: + forced.append(vars_[K - 1]) + cap_quality.append({ + "cap": K, + "var_at_cap_p50": _pct(forced, 50), + "var_at_cap_p90": _pct(forced, 90), + "var_at_cap_max": _pct(forced, 100), + "n": len(forced), + }) + + # --- counterfactual plateau sweep --- + plateau_sweep = [] + for eps in plateau_eps_list: + fired = 0 + fire_idx = [] + fire_var = [] + for b in sorted(committed_bins): + vars_ = [v for _, v, _ in committed_bins[b]] + onset, v = _plateau_onset(vars_, plateau_N, eps) + if onset is not None: + fired += 1 + fire_idx.append(onset) + fire_var.append(v) + plateau_sweep.append({ + "patience_N": plateau_N, + "eps": eps, + "bins_fired": fired, + "frac_bins_fired": (fired / n_bins) if n_bins else None, + "fire_iter_p50": _pct(fire_idx, 50), + "fire_iter_p90": _pct(fire_idx, 90), + "var_at_fire_p50": _pct(fire_var, 50), + "var_at_fire_p90": _pct(fire_var, 90), + }) + + return { + "n_committed_bins": n_bins, + "var_threshold": threshold, + "total_natural_iters": total_natural_iters, + "commit_reasons": dict(reason_counts), + "summary": summary, + "evolution_thirds": thirds, + "cap_sweep": cap_sweep, + "cap_quality": cap_quality, + "plateau_sweep": plateau_sweep, + } + + +def _fmt(x, nd=3): + return "None" if x is None else (f"{x:.{nd}f}" if isinstance(x, float) else str(x)) + + +def print_report(r): + print(f"\n=== Variance-decay curve characterization ===") + print(f"committed data-bins: {r['n_committed_bins']} " + f"var_threshold: {_fmt(r['var_threshold'])} " + f"total iters (compute): {r['total_natural_iters']}") + if r.get("commit_reasons"): + split = " ".join(f"{k}={v}" for k, v in sorted(r["commit_reasons"].items())) + print(f"commit reasons: {split}") + + print(f"\n-- per-bin distribution (p10 / p50 / p90 / mean) --") + s = r["summary"] + rows = [ + ("iters per bin", "n_iters", 1), + ("commit iter idx", "commit_idx", 1), + ("var @ iter 0", "var_initial", 3), + ("var @ commit", "var_commit", 3), + ("var min in bin", "var_min", 3), + ("plateau onset iter", "plateau_onset", 1), + ("var @ plateau", "var_at_plateau", 3), + ] + for label, key, nd in rows: + d = s[key] + print(f" {label:20s}: {_fmt(d['p10'],nd):>8} / {_fmt(d['p50'],nd):>8} / " + f"{_fmt(d['p90'],nd):>8} / {_fmt(d['mean'],nd):>8} (n={d['n']})") + + print(f"\n-- evolution over training (databins split in thirds) --") + ev = r["evolution_thirds"] + print(f" {'phase':6s} {'bins':>5} {'med_iters':>10} {'med_var_commit':>15} " + f"{'med_var_min':>12} {'med_onset':>10}") + for name in ["early", "mid", "late"]: + if name in ev: + d = ev[name] + print(f" {name:6s} {d['n_bins']:>5} {_fmt(d['median_n_iters'],1):>10} " + f"{_fmt(d['median_var_commit'],3):>15} {_fmt(d['median_var_min'],3):>12} " + f"{_fmt(d['median_plateau_onset'],1):>10}") + + print(f"\n-- fixed_cap sweep (iters/forward-passes saved) --") + print(f" {'cap':>4} {'bins_hit':>9} {'%bins':>7} {'iters_saved':>12} {'%iters':>8}") + for d in r["cap_sweep"]: + print(f" {d['cap']:>4} {d['bins_affected']:>9} " + f"{_fmt(100*(d['frac_bins_affected'] or 0),1):>7} {d['iters_saved']:>12} " + f"{_fmt(100*(d['frac_iters_saved'] or 0),1):>8}") + print(f"\n-- fixed_cap quality (var we'd COMMIT at on capped bins; lower=safer) --") + print(f" {'cap':>4} {'var@cap p50':>12} {'p90':>8} {'max':>8} (natural commit var p50~thr)") + for d in r["cap_quality"]: + print(f" {d['cap']:>4} {_fmt(d['var_at_cap_p50'],3):>12} " + f"{_fmt(d['var_at_cap_p90'],3):>8} {_fmt(d['var_at_cap_max'],3):>8} (n={d['n']})") + + print(f"\n-- plateau rule sweep (patience N, tol eps) --") + print(f" {'N':>3} {'eps':>6} {'bins_fired':>11} {'%':>6} " + f"{'fire_iter p50/p90':>18} {'var@fire p50/p90':>18}") + for d in r["plateau_sweep"]: + fi = f"{_fmt(d['fire_iter_p50'],0)}/{_fmt(d['fire_iter_p90'],0)}" + vf = f"{_fmt(d['var_at_fire_p50'],3)}/{_fmt(d['var_at_fire_p90'],3)}" + print(f" {d['patience_N']:>3} {_fmt(d['eps'],2):>6} {d['bins_fired']:>11} " + f"{_fmt(100*(d['frac_bins_fired'] or 0),1):>6} {fi:>18} {vf:>18}") + print() + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("run", help="run dir or aggregator_*.jsonl") + ap.add_argument("--caps", default="7,8,10,12,15", + help="comma list of fixed_cap values to sweep") + ap.add_argument("--plateau-N", type=int, default=3, help="plateau patience window") + ap.add_argument("--plateau-eps", default="0.05,0.10,0.15", + help="comma list of relative-drop tolerances to sweep") + ap.add_argument("--max-databins", type=int, default=None, + help="only analyze the first N committed data-bins (like-for-like truncation)") + ap.add_argument("--json", default=None, help="write full result as JSON here") + a = ap.parse_args() + + caps = [int(x) for x in a.caps.split(",") if x.strip()] + eps_list = [float(x) for x in a.plateau_eps.split(",") if x.strip()] + + curves, threshold, reasons = load_curves(a.run, max_databins=a.max_databins) + if not curves: + sys.exit("no agg_round events with var found") + r = characterize(curves, threshold, caps, a.plateau_N, eps_list, reasons=reasons) + print_report(r) + if a.json: + with open(a.json, "w") as fh: + json.dump(r, fh, indent=2) + print(f"wrote {a.json}") + + +if __name__ == "__main__": + main() diff --git a/lib/python/examples/fwdllm/expt_scripts/compare_baselines.py b/lib/python/examples/fwdllm/expt_scripts/compare_baselines.py new file mode 100644 index 000000000..8fd902794 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/compare_baselines.py @@ -0,0 +1,313 @@ +#!/usr/bin/env python3 +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Cross-baseline metric reducer (EXPERIMENTS.md WS4). + +Discovers the latest run per baseline (exact-token regex, so `fwdllm` never captures +`fwdllm_plus`) and emits one comparison table + CSV (+ optional overlay plots). The +five experiments are computed by `plotlib.reducers` (shared with plot_run.py and +make_paper_figs.py); this module owns only discovery, the mix-guard and the table. +`--plots` renders via `plotlib.figures` so the look matches the paper. Missing +baselines / metrics degrade to skipped / None, never a crash. + + python compare_baselines.py [--variant real|sim] [--target-acc 0.84] [--plots] +""" + +from __future__ import annotations + +import argparse +import csv +import glob +import json +import os +import re +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +if HERE not in sys.path: + sys.path.insert(0, HERE) +from plotlib import reducers as R # noqa: E402 (single source of the metrics) + +# `run____n_smoke[_]_` +_RUN_RE = re.compile( + r"^run_(?P\d{8}_\d{6})_(?P.+)_n(?P\d+)_smoke" + r"(?:_(?P.+))?_(?Preal|sim)$" +) +_DEFAULT_BASELINES = ["fwdllm", "fwdllm_plus", "fluxtune"] + + +# --------------------------------------------------------------------------- # +# discovery +# --------------------------------------------------------------------------- # +def discover(experiments_dir: str, variant: str) -> dict: + """{baseline: (n, trace, path)} — latest run per baseline for the variant. + + Exact-token match on baseline avoids the fwdllm ⊃ fwdllm_plus glob trap.""" + latest: dict = {} + for path in glob.glob(os.path.join(experiments_dir, "run_*")): + m = _RUN_RE.match(os.path.basename(path)) + if not m or m["variant"] != variant: + continue + b = m["baseline"] + prev = latest.get(b) + if prev is None or m["ts"] > prev[0]: + latest[b] = (m["ts"], int(m["n"]), m["trace"] or "", path) + return {b: (n, tr, p) for b, (ts, n, tr, p) in latest.items()} + + +# --------------------------------------------------------------------------- # +# small stats helpers +# --------------------------------------------------------------------------- # +def _pct(values, q: float): + xs = sorted(v for v in values if v is not None) + if not xs: + return None + if len(xs) == 1: + return xs[0] + pos = (q / 100.0) * (len(xs) - 1) + lo = int(pos) + hi = min(lo + 1, len(xs) - 1) + frac = pos - lo + return xs[lo] * (1 - frac) + xs[hi] * frac + + +def _p(values): + return {"p50": _pct(values, 50), "p90": _pct(values, 90), "p99": _pct(values, 99)} + + +# --------------------------------------------------------------------------- # +# the five experiments — PURE over a RunResult (from plotlib.reducers) +# --------------------------------------------------------------------------- # +def expt1_time_to_target(rr: R.RunResult, target: float, window: int) -> dict: + max_acc, final_acc = rr.max_accuracy(), rr.final_accuracy() + e = rr.target_event(target, window) + if e is not None: + return { + "reached": True, "target": target, "window": window, + "wall_s": round(e["ts"] - rr.t0, 2) if rr.t0 is not None else None, + "vclock_s": rr.vclock_at(e["ts"]), "rounds": e.get("round"), + # cumulative bin-evals to convergence (monotonic; data_id cycles per round) + "data_bins": rr.evals.index(e) + 1, + "max_accuracy": max_acc, "final_accuracy": final_acc, + } + return { + "reached": False, "target": target, "window": window, + "wall_s": None, "vclock_s": None, "rounds": None, "data_bins": None, + "max_accuracy": max_acc, "final_accuracy": final_acc, + } + + +def expt2_utilization(rr: R.RunResult) -> dict: + tp = _p(rr.busy_frac) + w = rr.agg_wall_s + frac = lambda x: (round(x / w, 4) if (w and w > 0 and x is not None) else None) + return { + "trainer_busy_frac_p50": tp["p50"], "trainer_busy_frac_p90": tp["p90"], + "trainer_busy_frac_p99": tp["p99"], "n_trainers_measured": len(rr.busy_frac), + "agg_busy_frac": frac(rr.agg_compute_s), "agg_barrier_wait_frac": frac(rr.agg_barrier_s), + "agg_drain_frac": frac(rr.agg_drain_s), + "agg_wall_s": round(w, 2) if w else None, + } + + +def expt3_productivity(rr: R.RunResult) -> dict: + losses = [e["loss"] for e in rr.evals if e["loss"] is not None] + first_loss = losses[0] if losses else None + final_loss = losses[-1] if losses else None + delta = rr.delta_loss() + gpu_total = rr.gpu_s_total() + return { + "first_loss": first_loss, "final_loss": final_loss, "delta_loss": delta, + "trainer_gpu_s": round(rr.trainer_gpu_s, 2), "agg_compute_s": round(rr.agg_compute_s, 2), + "gpu_s_total": round(gpu_total, 2), + "delta_loss_per_gpu_s": (round(delta / gpu_total, 6) if (delta and gpu_total) else None), + "forward_passes_total": rr.fwd_total if rr.have_fwd else None, + "perturbations_total": rr.pert_total if rr.have_fwd else None, + "delta_loss_per_forward_pass": ( + round(delta / rr.fwd_total, 9) if (delta and rr.have_fwd and rr.fwd_total) else None), + } + + +def expt4_network(rr: R.RunResult) -> dict: + if not rr.have_comm and not rr.up_sizes and not rr.down_sizes: + return {"comm_telemetry": False} + a_sizes, c_sizes = rr.down_sizes, rr.up_sizes + ap, cp = _p(a_sizes), _p(c_sizes) + return { + "comm_telemetry": True, + "msgs_agg_to_trainer": len(a_sizes), "msgs_trainer_to_agg": len(c_sizes), + "bytes_agg_to_trainer": sum(a_sizes), "bytes_trainer_to_agg": sum(c_sizes), + "bytes_total": sum(a_sizes) + sum(c_sizes), + "msg_size_agg_p50": ap["p50"], "msg_size_agg_p90": ap["p90"], "msg_size_agg_p99": ap["p99"], + "msg_size_client_p50": cp["p50"], "msg_size_client_p90": cp["p90"], "msg_size_client_p99": cp["p99"], + } + + +def expt5_sessions(rr: R.RunResult, is_async: bool) -> dict: + sp = _p(rr.session_durs) + return { + "session_def": "dispatch_to_commit" if is_async else "one_round_span", + "session_s_p50": sp["p50"], "session_s_p90": sp["p90"], "session_s_p99": sp["p99"], + "n_sessions": len(rr.session_durs), + "part_rounds_p50": _pct(rr.part_rounds, 50), "part_rounds_p90": _pct(rr.part_rounds, 90), + "part_databins_p50": _pct(rr.part_bins, 50), "part_databins_p90": _pct(rr.part_bins, 90), + "part_iters_p50": _pct(rr.part_iters, 50), "part_iters_p90": _pct(rr.part_iters, 90), + "part_iters_total": sum(rr.part_iters), + } + + +# --------------------------------------------------------------------------- # +# config fingerprint mix-guard +# --------------------------------------------------------------------------- # +def _fingerprint(run_dir: str, n: int, trace: str) -> dict: + """Shared condition axes that MUST match across baselines (N/partition/trace). + agg_goal/c/k/selector legitimately differ per baseline, so they're excluded.""" + part = None + cfg = os.path.join(run_dir, "aggregator_config.json") + if os.path.exists(cfg): + try: + h = json.load(open(cfg)).get("hyperparameters", {}) + part = h.get("partition_method") + except (ValueError, OSError): + pass + return {"N": n, "trace": trace, "partition": part} + + +# --------------------------------------------------------------------------- # +# driver +# --------------------------------------------------------------------------- # +def _fmt(v): + if v is None: + return "–" + if isinstance(v, float): + return f"{v:.4g}" + if isinstance(v, int) and abs(v) >= 100000: + return f"{v/1e6:.2f}M" + return str(v) + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--experiments-dir", default=os.path.join(HERE, "..", "experiments")) + ap.add_argument("--variant", choices=["real", "sim"], default="real") + ap.add_argument("--baselines", default=",".join(_DEFAULT_BASELINES)) + ap.add_argument("--target-acc", type=float, default=0.84) + ap.add_argument("--window", type=int, default=20) + ap.add_argument("--out", default=None, help="output dir for CSV (+plots); default = experiments/_compare") + ap.add_argument("--plots", action="store_true", help="also write paper overlay plots (via plotlib)") + ap.add_argument("--smooth", type=float, default=0.7, + help="EMA smoothing factor for overlay learning curves (0 = raw)") + ap.add_argument("--loss-plateau-rel", type=float, default=0.01, + help="cut each run at its last cumulative test-loss drop of this " + "relative size (default 0.01 = 1%%; end of productive learning)") + ap.add_argument("--post-peak-grace-min", type=float, default=0.0, + help="minutes to keep past the loss-plateau point (default 0)") + ap.add_argument("--no-cutoff", action="store_true", help="use full telemetry (no cutoff)") + args = ap.parse_args() + grace_s = None if args.no_cutoff else args.post_peak_grace_min * 60.0 + + exp_dir = os.path.abspath(args.experiments_dir) + baselines = [b.strip() for b in args.baselines.split(",") if b.strip()] + out_dir = os.path.abspath(args.out) if args.out else os.path.join(exp_dir, "_compare") + os.makedirs(out_dir, exist_ok=True) + + found = discover(exp_dir, args.variant) + rows, fps, results = {}, {}, {} + for b in baselines: + if b not in found: + print(f" [compare] WARN no {args.variant} run found for baseline '{b}' — skipping", file=sys.stderr) + continue + n, trace, path = found[b] + rr = R.load_run(path, key=b, post_peak_grace_s=grace_s, + loss_plateau_rel=args.loss_plateau_rel) + if rr is None or not rr.evals: + print(f" [compare] WARN empty aggregator telemetry for '{b}' ({path}) — skipping", file=sys.stderr) + continue + results[b] = rr + is_async = (b == "fluxtune") + rows[b] = { + "run_dir": os.path.basename(path), "N": n, "trace": trace or "syn_0", + **{f"e1_{k}": v for k, v in expt1_time_to_target(rr, args.target_acc, args.window).items()}, + **{f"e2_{k}": v for k, v in expt2_utilization(rr).items()}, + **{f"e3_{k}": v for k, v in expt3_productivity(rr).items()}, + **{f"e4_{k}": v for k, v in expt4_network(rr).items()}, + **{f"e5_{k}": v for k, v in expt5_sessions(rr, is_async).items()}, + } + fps[b] = _fingerprint(path, n, trace) + + if not rows: + print(" [compare] no baselines with telemetry found — nothing to compare.", file=sys.stderr) + return 1 + + # mix-guard: shared condition axes must agree across baselines + uniq = {tuple(sorted(fp.items())) for fp in fps.values()} + if len(uniq) > 1: + print(" [compare] ⚠ MIX-GUARD: baselines differ on a SHARED condition axis " + "(N/partition/trace) — comparison may be apples-to-oranges:", file=sys.stderr) + for b, fp in fps.items(): + print(f" {b}: {fp}", file=sys.stderr) + + cols = [] + for b in baselines: + for k in rows.get(b, {}): + if k not in cols: + cols.append(k) + + csv_path = os.path.join(out_dir, "compare_baselines.csv") + with open(csv_path, "w", newline="") as fh: + w = csv.writer(fh) + w.writerow(["baseline"] + cols) + for b in baselines: + if b in rows: + w.writerow([b] + [rows[b].get(c, "") for c in cols]) + + headline = [ + ("e1_reached", "conv?"), ("e1_wall_s", "t→acc wall_s"), ("e1_data_bins", "t→acc bins"), + ("e1_max_accuracy", "max_acc"), ("e2_trainer_busy_frac_p50", "trn_busy_p50"), + ("e2_agg_busy_frac", "agg_busy"), ("e3_delta_loss", "Δloss"), + ("e3_delta_loss_per_gpu_s", "Δloss/gpu_s"), ("e3_delta_loss_per_forward_pass", "Δloss/fwd"), + ("e4_bytes_total", "net_bytes"), ("e4_msgs_trainer_to_agg", "msgs↑"), + ("e5_session_s_p50", "sess_p50"), ("e5_part_iters_total", "iters_tot"), + ] + w1 = max(len(b) for b in rows) + print(f"\n=== CROSS-BASELINE COMPARISON ({args.variant}, τ={args.target_acc}, W={args.window}) ===") + hdr = f"{'baseline':<{w1}} " + " ".join(f"{lbl:>13}" for _, lbl in headline) + print(hdr) + print("-" * len(hdr)) + for b in baselines: + if b not in rows: + continue + cells = " ".join(f"{_fmt(rows[b].get(k)):>13}" for k, _ in headline) + print(f"{b:<{w1}} {cells}") + print(f"\nCSV: {csv_path}") + + if args.plots: + _plots(results, out_dir, args.target_acc, args.smooth) + print(f"Out: {out_dir}") + return 0 + + +def _plots(results: dict, out_dir: str, target, smooth: float = 0.0): + """Overlay plots via the shared plotlib (same look as the paper figures).""" + try: + from plotlib import baselines as B + from plotlib import figures as F + from plotlib import style as S + except Exception as e: # noqa: BLE001 + print(f" [compare] plots skipped (plotlib/matplotlib unavailable: {e})", file=sys.stderr) + return + S.use_paper_style() + ordered = [results[k] for k in B.ordered(results.keys())] + n = 0 + for name, builder in F.FIG_BUILDERS.items(): + fig = builder(ordered, target=target, smooth=smooth) + if fig is None: + continue + S.save_pdf(fig, out_dir, name) + n += 1 + print(f" [compare] {n} overlay plot(s) written to {out_dir}") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/lib/python/examples/fwdllm/expt_scripts/diagnose_partition_binning.py b/lib/python/examples/fwdllm/expt_scripts/diagnose_partition_binning.py new file mode 100644 index 000000000..72330318b --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/diagnose_partition_binning.py @@ -0,0 +1,395 @@ +#!/usr/bin/env python +"""H0 — partition & data-bin diagnostic (fluxtune stability track). + +READ-ONLY. Checks whether the data a trainer sees at each `data_id` focuses on one +class, and whether that predicts the single-class collapses seen in run telemetry. + +Measures, per Dirichlet alpha (100-client distribution): + + Layer A per-CLIENT class skew (from the partition h5 + labels): dominant-class + fraction, #classes present, normalized entropy -- "how non-IID is each + client" across alpha=0.1 / 1 / 100. + + Layer B per-BIN narrowness (a data_id = one `train_batch_size`-sample bin): + per-bin dominant frac / entropy; per-client single-class-bin fraction. + Order source is FAITHFUL (the frozen per-client pickle cache the run + trained on) when present, else raw partition index order (flagged); for + alpha=1 both are reported to expose the read_instance_from_h5 shuffle. + + Layer C per-DATA_ID pooled bias (proxy for what the aggregator sees at commit): + pool bin #k across clients -> pooled dominant frac / entropy. Overlaid + (alpha=1) against the OBSERVED collapse data_ids from agg_eval telemetry + to test whether low-entropy data_ids predict collapse. + +Usage: + python diagnose_partition_binning.py # all defaults + python diagnose_partition_binning.py --alphas 0.1,1,100 --plots +""" +import argparse, json, math, os, pickle, sys +from collections import Counter + +import numpy as np + +# ---- defaults resolved from the run's aggregator_config.json (the /coc/scratch h5s) ---- +DEF_PART = "/coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5" +DEF_DATA = "/coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5" +DEF_CACHE = "/home/dgarg39/gaurav/flame/lib/python/examples/fwdllm/expts/run_tc_expts/cache_dir" +DEF_RUN = ("/home/dgarg39/flame/lib/python/examples/fwdllm/experiments/" + "run_20260708_025543_fluxtune_n100_smoke_syn_0_real") +CACHE_TMPL = ("distilbert_distilbert-base-uncased_cached_192_ClassificationModel_" + "agnews_niid_label_clients=100_alpha={alpha}_{cid}") +N_CLASSES = 4 # AG News + + +def norm_entropy(counts): + """Shannon entropy normalized to [0,1] over N_CLASSES (1 = uniform, 0 = single class).""" + tot = sum(counts.values()) + if tot == 0: + return float("nan") + h = 0.0 + for c in counts.values(): + if c > 0: + p = c / tot + h -= p * math.log(p) + return h / math.log(N_CLASSES) + + +def dom_frac(counts): + tot = sum(counts.values()) + return (max(counts.values()) / tot) if tot else float("nan") + + +def load_label_map(data_h5): + """Full {index:int label} from the data h5 Y group. Built once, reused for index-order alphas.""" + import h5py + print(f" reading labels from {data_h5} ...", flush=True) + with h5py.File(data_h5, "r") as f: + Y = f["Y"] + m = {} + for k in Y.keys(): + m[int(k)] = int(Y[k][()].decode("utf-8")) + print(f" loaded {len(m)} labels; classes present: {sorted(set(m.values()))}", flush=True) + return m + + +def client_labels_faithful(cache_dir, alpha, cid): + """Labels in the frozen order the run trained on (from the pickle cache), or None.""" + path = os.path.join(cache_dir, CACHE_TMPL.format(alpha=alpha, cid=cid)) + if not os.path.exists(path): + return None + with open(path, "rb") as h: + tup = pickle.load(h) + train_examples = tup[0] + return [int(e.label) for e in train_examples] + + +def client_labels_index(partition_h5_group, cid, label_map): + """Labels in raw partition-index order (fallback when no cache).""" + idx = partition_h5_group["partition_data"][str(cid)]["train"][()] + return [label_map[int(i)] for i in idx] + + +def analyze_alpha(alpha, part_h5, label_map, cache_dir, batch_size, n_clients): + import h5py + group_name = f"niid_label_clients={n_clients}_alpha={alpha}" + per_client = [] # dicts of client-level metrics + per_client_bins = [] # list of list-of-bin-label-Counters (for Layer C pooling) + order_src = Counter() + with h5py.File(part_h5, "r") as f: + if group_name not in f: + raise SystemExit(f"partition group '{group_name}' not in {part_h5}") + g = f[group_name] + cids = sorted((int(c) for c in g["partition_data"].keys())) + for cid in cids: + labels = client_labels_faithful(cache_dir, alpha, cid) + if labels is not None: + order_src["faithful(cache)"] += 1 + else: + labels = client_labels_index(g, cid, label_map) + order_src["index(partition)"] += 1 + cc = Counter(labels) + # bins + bins = [Counter(labels[i:i + batch_size]) for i in range(0, len(labels), batch_size)] + single = sum(1 for b in bins if len([1 for v in b.values() if v > 0]) == 1) + per_client.append(dict( + cid=cid, n=len(labels), n_classes=len([1 for v in cc.values() if v > 0]), + dom_frac=dom_frac(cc), entropy=norm_entropy(cc), + n_bins=len(bins), single_class_bin_frac=single / len(bins) if bins else float("nan"), + mean_bin_entropy=float(np.nanmean([norm_entropy(b) for b in bins])) if bins else float("nan"), + class_counts={int(k): int(v) for k, v in cc.items()}, + )) + per_client_bins.append(bins) + return group_name, per_client, per_client_bins, order_src + + +def pooled_per_dataid(per_client_bins, max_k=None): + """Layer C: pool bin #k across all clients -> pooled Counter per data_id k.""" + if max_k is None: + max_k = max(len(b) for b in per_client_bins) + pooled = [] + for k in range(max_k): + agg = Counter() + n_contrib = 0 + for bins in per_client_bins: + if k < len(bins): + agg += bins[k] + n_contrib += 1 + pooled.append(dict(data_id=k, n_contrib=n_contrib, + dom_frac=dom_frac(agg), entropy=norm_entropy(agg))) + return pooled + + +def cohort_sampled(per_client_bins, K, R, seed=0): + """Realistic aggregation batch: for each data_id, draw R random K-client cohorts and + pool their bin #k. Returns arrays of pooled entropy and dominant-fraction over all + (data_id, cohort) draws — the class bias an agg_goal=K commit actually experiences.""" + rng = np.random.default_rng(seed) + n_clients = len(per_client_bins) + max_k = max(len(b) for b in per_client_bins) + ent, dfr = [], [] + for k in range(max_k): + avail = [ci for ci in range(n_clients) if k < len(per_client_bins[ci])] + if len(avail) < K: + continue + for _ in range(R): + coh = rng.choice(avail, size=K, replace=False) + agg = Counter() + for ci in coh: + agg += per_client_bins[ci][k] + ent.append(norm_entropy(agg)) + dfr.append(dom_frac(agg)) + return np.array(ent), np.array(dfr) + + +def observed_collapse(run_dir): + """Round-1 agg_eval: (data_id -> accuracy) and the set of single-class-collapse data_ids.""" + tele = os.path.join(run_dir, "telemetry") + agg = [f for f in os.listdir(tele) if f.startswith("aggregator")] + if not agg: + return None, None + acc_by_did, collapse = {}, set() + with open(os.path.join(tele, agg[0])) as fh: + for line in fh: + if '"agg_eval"' not in line: + continue + d = json.loads(line) + if d.get("round") != 1: + continue + did, a, m = d.get("data_id"), d["test-accuracy"], d.get("mcc", 0.0) + acc_by_did[did] = a + if abs(a - 0.25) < 0.006 and abs(m) < 0.02: + collapse.add(did) + return acc_by_did, collapse + + +def pct(vals, p): + v = [x for x in vals if not math.isnan(x)] + return float(np.percentile(v, p)) if v else float("nan") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--partition", default=DEF_PART) + ap.add_argument("--data", default=DEF_DATA) + ap.add_argument("--cache-dir", default=DEF_CACHE) + ap.add_argument("--run-dir", default=DEF_RUN, help="run for the alpha=1 collapse overlay (Layer C)") + ap.add_argument("--alphas", default="0.1,1,100") + ap.add_argument("--n-clients", type=int, default=100) + ap.add_argument("--batch-size", type=int, default=8) + ap.add_argument("--agg-goal", type=int, default=10, help="clients pooled per commit (Layer C')") + ap.add_argument("--cohort-draws", type=int, default=200, help="random cohorts sampled per data_id") + ap.add_argument("--overlay-alpha", default="1", help="which alpha to overlay against run telemetry") + ap.add_argument("--out", default=None, help="dir for csv/json/plots (default: ./_diag_partition)") + ap.add_argument("--plots", action="store_true") + ap.add_argument("--dist", action="store_true", + help="sample→trainer→bin distribution report + heatmaps (heterogeneity study)") + args = ap.parse_args() + + alphas = [a.strip() for a in args.alphas.split(",")] + out = args.out or os.path.join(os.path.dirname(os.path.abspath(__file__)), "_diag_partition") + os.makedirs(out, exist_ok=True) + + # labels only needed if some alpha lacks a cache (index-order fallback) + need_labels = any( + not os.path.exists(os.path.join(args.cache_dir, CACHE_TMPL.format(alpha=a, cid=0))) + for a in alphas) + label_map = load_label_map(args.data) if need_labels else {} + + summary = {} + pooled_all = {} + bins_all = {} + for a in alphas: + print(f"\n===== alpha = {a} (group niid_label_clients={args.n_clients}_alpha={a}) =====") + gname, pc, pcb, osrc = analyze_alpha(a, args.partition, label_map, args.cache_dir, + args.batch_size, args.n_clients) + print(f" order source: {dict(osrc)}") + df = [c["dom_frac"] for c in pc] + ncl = [c["n_classes"] for c in pc] + ent = [c["entropy"] for c in pc] + scb = [c["single_class_bin_frac"] for c in pc] + mbe = [c["mean_bin_entropy"] for c in pc] + nsz = [c["n"] for c in pc] + print(f" Layer A per-client (n={len(pc)} clients, samples/client p10/50/90=" + f"{pct(nsz,10):.0f}/{pct(nsz,50):.0f}/{pct(nsz,90):.0f}):") + print(f" dominant-class fraction p10/50/90 = {pct(df,10):.3f} / {pct(df,50):.3f} / {pct(df,90):.3f}" + f" (mean {np.mean(df):.3f})") + print(f" #classes present p10/50/90 = {pct(ncl,10):.1f} / {pct(ncl,50):.1f} / {pct(ncl,90):.1f}") + print(f" normalized entropy p10/50/90 = {pct(ent,10):.3f} / {pct(ent,50):.3f} / {pct(ent,90):.3f}") + print(f" Layer B per-bin (batch_size={args.batch_size}):") + print(f" single-class-bin frac p10/50/90 = {pct(scb,10):.3f} / {pct(scb,50):.3f} / {pct(scb,90):.3f}" + f" (mean {np.mean(scb):.3f})") + print(f" mean bin entropy p10/50/90 = {pct(mbe,10):.3f} / {pct(mbe,50):.3f} / {pct(mbe,90):.3f}") + pooled = pooled_per_dataid(pcb) + pooled_all[a] = pooled + bins_all[a] = pcb + pe = [p["entropy"] for p in pooled] + pdf = [p["dom_frac"] for p in pooled] + print(f" Layer C per-data_id pooled over cohort ({len(pooled)} data_ids):") + print(f" pooled entropy p10/50/90 = {pct(pe,10):.3f} / {pct(pe,50):.3f} / {pct(pe,90):.3f}") + print(f" pooled dominant frac p10/50/90 = {pct(pdf,10):.3f} / {pct(pdf,50):.3f} / {pct(pdf,90):.3f}") + # realistic aggregation batch: agg_goal=K random cohorts + ce, cd = cohort_sampled(pcb, args.agg_goal, args.cohort_draws) + frac_biased = float(np.mean(cd > 0.5)) + print(f" Layer C' realistic cohort (K={args.agg_goal} clients/commit, {args.cohort_draws} draws/data_id):") + print(f" cohort pooled entropy p10/50/90 = {pct(list(ce),10):.3f} / {pct(list(ce),50):.3f} / {pct(list(ce),90):.3f}") + print(f" cohort dominant frac p10/50/90 = {pct(list(cd),10):.3f} / {pct(list(cd),50):.3f} / {pct(list(cd),90):.3f}") + print(f" P(cohort dominant frac > 0.5) = {frac_biased:.3f} (fraction of commits that are majority-one-class)") + summary[a] = dict(group=gname, order=dict(osrc), per_client=pc, pooled=pooled, + cohort=dict(K=args.agg_goal, entropy_p50=pct(list(ce),50), + dom_p50=pct(list(cd),50), frac_biased=frac_biased)) + + # ---- Layer C overlay: predicted-biased data_ids vs observed collapses (overlay alpha) ---- + oa = args.overlay_alpha + if oa in pooled_all and os.path.isdir(args.run_dir): + print(f"\n===== Layer C overlay: alpha={oa} predicted bias vs OBSERVED collapse ({os.path.basename(args.run_dir)}) =====") + acc_by_did, collapse = observed_collapse(args.run_dir) + if acc_by_did: + pooled = pooled_all[oa] + common = [p["data_id"] for p in pooled if p["data_id"] in acc_by_did] + pe = np.array([pooled[d]["entropy"] for d in common]) + ac = np.array([acc_by_did[d] for d in common]) + # rank correlation (low pooled entropy should co-occur with low accuracy) + from scipy.stats import spearmanr + rho, pval = spearmanr(pe, ac) + print(f" observed round-1 collapse data_ids (acc~0.25 & mcc~0): {sorted(collapse)}") + # predicted-biased = lowest-entropy quartile + thr = np.percentile(pe, 25) + predicted = set(d for d in common if pooled[d]["entropy"] <= thr) + hit = len(predicted & collapse) + print(f" predicted-biased data_ids (pooled entropy <= p25={thr:.3f}): n={len(predicted)}") + print(f" overlap predicted ∩ collapsed = {hit} / {len(collapse)} collapses " + f"({100*hit/max(1,len(collapse)):.0f}% of collapses are in the biased quartile)") + print(f" Spearman(pooled_entropy, accuracy) over {len(common)} data_ids = {rho:+.3f} (p={pval:.1e})") + print(" -> positive rho ⇒ higher-entropy (class-mixed) data_ids have higher accuracy = data-bias hypothesis supported") + summary["_overlay"] = dict(alpha=oa, collapse=sorted(collapse), spearman=rho, p=pval, + predicted_biased=sorted(predicted), overlap=hit) + else: + print(" no agg_eval telemetry found in run-dir; overlay skipped") + + with open(os.path.join(out, "diag_summary.json"), "w") as fh: + json.dump(summary, fh, indent=2, default=float) + print(f"\nwrote {os.path.join(out,'diag_summary.json')}") + + if args.plots: + make_plots(summary, pooled_all, args, out) + if args.dist: + dist_report(summary, bins_all, alphas, out) + + +def _canon_matrix(per_client): + """100×4 counts matrix with class ids canonicalized to 0..3 by rank over the α's label domain.""" + domain = sorted({k for c in per_client for k in c["class_counts"].keys()}) + cmap = {lbl: i for i, lbl in enumerate(domain)} + K = len(domain) + M = np.zeros((len(per_client), K), dtype=int) + for r, c in enumerate(per_client): + for lbl, n in c["class_counts"].items(): + M[r, cmap[lbl]] = n + return M, domain + + +def dist_report(summary, bins_all, alphas, out): + """How samples distribute across trainers and their data-bins, as heterogeneity varies.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + print("\n===== DISTRIBUTION STUDY (sample → trainer → bin; N=100) =====") + fig, axes = plt.subplots(1, len(alphas), figsize=(4.2 * len(alphas), 4.6), squeeze=False) + for j, a in enumerate(alphas): + pc = summary[a]["per_client"] + M, domain = _canon_matrix(pc) # 100×4 counts + tot = M.sum() + per_class_tot = M.sum(0) + # global structural facts + sizes = M.sum(1) + print(f"\n α={a}: total samples {tot} | samples/trainer min/med/max = " + f"{sizes.min()}/{int(np.median(sizes))}/{sizes.max()} | per-class totals {per_class_tot.tolist()}") + # concentration: for each class, share held by its top-10 trainers (of 100) + top10 = [] + for c in range(M.shape[1]): + col = np.sort(M[:, c])[::-1] + top10.append(col[:10].sum() / max(1, col.sum())) + print(f" class spread — top-10-trainer share per class: " + f"{[round(x,2) for x in top10]} (0.10 = perfectly even; 1.0 = all in 10 trainers)") + # #classes present per trainer histogram + ncl = (M > 0).sum(1) + hist = {k: int((ncl == k).sum()) for k in range(1, M.shape[1] + 1)} + print(f" #classes/trainer histogram (1..4 classes): {hist}") + # per-bin composition across ALL bins of all trainers + bindom = [dom_frac(b) for bins in bins_all[a] for b in bins] + print(f" per-bin dominant-class frac: mean {np.mean(bindom):.3f} " + f"(fraction of bins that are ≥50% one class: {np.mean(np.array(bindom)>0.5):.3f})") + # heatmap: trainers (rows sorted by dominant class then dom frac) × class, row-normalized + rown = M / M.sum(1, keepdims=True) + order = sorted(range(len(rown)), key=lambda r: (int(np.argmax(rown[r])), -rown[r].max())) + ax = axes[0][j] + im = ax.imshow(rown[order], aspect="auto", cmap="viridis", vmin=0, vmax=1, + interpolation="nearest") + ax.set_title(f"α={a}\n(mean dom-frac {np.mean([c['dom_frac'] for c in pc]):.2f})") + ax.set_xlabel("class"); ax.set_xticks(range(len(domain))) + if j == 0: + ax.set_ylabel("trainer (sorted by dominant class)") + fig.colorbar(im, ax=axes[0].tolist(), fraction=0.025, label="within-trainer class proportion") + fig.suptitle("Class distribution across N=100 trainers as heterogeneity decreases (α↑)") + p = os.path.join(out, "trainer_class_heatmap.pdf") + fig.savefig(p, bbox_inches="tight"); plt.close(fig) + print(f"\n wrote {p}") + + +def make_plots(summary, pooled_all, args, out): + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + alphas = list(pooled_all.keys()) + # (1) CDF of per-client dominant-class fraction + fig, ax = plt.subplots(figsize=(5, 3.2)) + for a in alphas: + df = np.sort([c["dom_frac"] for c in summary[a]["per_client"]]) + ax.plot(df, np.linspace(0, 1, len(df)), label=f"α={a}") + ax.set_xlabel("per-client dominant-class fraction"); ax.set_ylabel("CDF") + ax.set_title("Client class skew by α (100 clients)"); ax.legend(); ax.grid(alpha=.3) + fig.tight_layout(); fig.savefig(os.path.join(out, "clientskew_cdf.pdf")); plt.close(fig) + # (2) alpha=1 pooled entropy vs observed accuracy per data_id + oa = args.overlay_alpha + ov = summary.get("_overlay") + if oa in pooled_all and ov: + acc, collapse = observed_collapse(args.run_dir) + pooled = pooled_all[oa] + ks = [p["data_id"] for p in pooled if p["data_id"] in acc] + fig, ax1 = plt.subplots(figsize=(7, 3.2)) + ax1.plot(ks, [pooled[k]["entropy"] for k in ks], color="tab:blue", label="pooled bin entropy") + ax1.set_xlabel("data_id"); ax1.set_ylabel("pooled entropy", color="tab:blue") + ax2 = ax1.twinx() + ax2.plot(ks, [acc[k] for k in ks], color="tab:red", alpha=.6, label="observed accuracy") + ax2.axhline(0.25, ls="--", color="grey", lw=.8) + for k in collapse: + ax2.axvline(k, color="tab:red", alpha=.08) + ax2.set_ylabel("round-1 accuracy", color="tab:red") + ax1.set_title(f"α={oa}: pooled data-bin entropy vs observed accuracy (collapses shaded)") + fig.tight_layout(); fig.savefig(os.path.join(out, "dataid_entropy_vs_acc.pdf")); plt.close(fig) + print(f"wrote plots to {out}") + + +if __name__ == "__main__": + main() diff --git a/lib/python/examples/fwdllm/expt_scripts/figs.yaml b/lib/python/examples/fwdllm/expt_scripts/figs.yaml new file mode 100644 index 000000000..c8a27ce1d --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/figs.yaml @@ -0,0 +1,21 @@ +# SOCC-2026 paper-figure manifest — baseline → run dir (the reproducible mapping). +# +# Paths resolve relative to the fwdllm example dir (…/examples/fwdllm) or absolute. +# Fill entries as logs land; missing baselines are skipped, not fatal. Override a +# single entry without editing this file via: make_paper_figs.py --run KEY=PATH +# +# Keys are INTERNAL baseline names (telemetry keys); paper display names/colors +# live in plotlib/baselines.py (e.g. our system's codename is switchable there). + +runs: + # fwdllm: terminated at ~6h (usable now; refresh path when the full run lands) + fwdllm: experiments/run_20260707_015846_fwdllm_n100_smoke_syn_0_real + fwdllm_plus: experiments/run_20260706_185023_fwdllm_plus_n100_smoke_syn_0_real + # fluxtune = full-stack default (C1 + Opt-1 + Opt-2 var-stop + Opt-3 grad-aware), + # N=100 α=1 launched 2026-07-08 (== R4 of the opt-ablation, figs_ablation.yaml). + # Was run_20260706_185045 (FeLiX-era, C3=placeholder); refreshed to the full stack. + fluxtune: experiments/run_20260708_025716_fluxtune_n100_smoke_syn_0_real + # fluxtune_nojvp: # also register it in plotlib/baselines.py + +target_acc: 0.84 +window: 20 diff --git a/lib/python/examples/fwdllm/expt_scripts/figs_ablation.yaml b/lib/python/examples/fwdllm/expt_scripts/figs_ablation.yaml new file mode 100644 index 000000000..ddb947bc7 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/figs_ablation.yaml @@ -0,0 +1,25 @@ +# SOCC-2026 FluxTune opt-ablation figure manifest — the 2×2 over the two +# learning-affecting opts (Opt-2 var-stop × Opt-3 grad-aware), C1+Opt-1 held +# constant. See EXPTS_CHARTER.md "Four planned full runs" and EXPERIMENTS.md §10. +# +# Landed N=100, α=1 (dir names say alpha0p1 — MISLABELED; the loaded-partition log +# line reads niid_label_clients=100_alpha=1), target 0.84, launched 2026-07-08. +# Identified by config (aggregator_config.json): var_stopping_policy × agg_rate type. +# +# Render: python make_paper_figs.py --manifest figs_ablation.yaml +# Keys are INTERNAL; blue-shade styles/labels live in plotlib/baselines.py. + +runs: + # R1 FluxTune-base — var-stop OFF, agg-rate=new (neither Opt-2 nor Opt-3). + # C1 guided JVP perturbations + Opt-1 ON (as in all four). NOT FeLiX: it does + # forward-mode LLM perturbation fine-tuning, only borrowing FeLiX's scalar agg rate. + fluxtune_r1_base: experiments/run_20260708_025543_fluxtune_n100_smoke_syn_0_real + # R2 + var-stop (Opt-2) — var-stop=plateau, agg-rate=new + fluxtune_r2_varstop: experiments/run_20260708_025616_fluxtune_n100_smoke_syn_0_real + # R3 + grad-aware (Opt-3) — var-stop OFF, agg-rate=grad_aware + fluxtune_r3_gradaware: experiments/run_20260708_025636_fluxtune_n100_smoke_syn_0_real + # R4 all (= full FluxTune default) — var-stop=plateau, agg-rate=grad_aware + fluxtune_r4_full: experiments/run_20260708_025716_fluxtune_n100_smoke_syn_0_real + +target_acc: 0.84 +window: 20 diff --git a/lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke.yaml b/lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke.yaml index 22bc34477..e14c2b319 100644 --- a/lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke.yaml +++ b/lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke.yaml @@ -32,10 +32,11 @@ experiments: # file exists (or ever will) for agnews; skips that lookup entirely. path_style: true availability: - # Cosmetic here -- fluxtune's 3-tier signal comes from - # trainer.client_notify (set in the baseline), not this mode. Kept - # consistent with the baseline's mobiperf_3st_50 trace. - mode: mobiperf_3st_50 + # Phase 1 is syn_0 (100% availability): client_notify then reports all + # trainers available. Switch to mobiperf_3st_50 (via --avail-trace) only + # for the Phase-2 unavailability runs. run_sequential.sh's --avail-trace + # patches trainer mode + client_notify + trackTrainerAvail together. + mode: syn_0 # FedFwd has no simulated-clock support; always real. time_mode: real # Skip the registry's 4-18s artificial delay for a fast smoke run. @@ -45,6 +46,12 @@ experiments: client_idx_modulo: 100 config_overrides: hyperparameters: + # Fluxtune JVP perf optimizations (simulate_fwdllm.md §L) — all + # BIT-IDENTICAL to the grads (scripts/profile_jvp_opt.py). Cuts + # fluxtune's per-batch forward passes ~37% (25->20) + trainable-only + # FD, so GPU drops under the modeled mobile delay. Set true on BOTH + # real+sim (must match). Set false to revert to the un-optimized path. + jvp_perf_opt: true data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 # Uniform/IID partition for smoke tests -- avoids the niid_label @@ -62,10 +69,15 @@ experiments: id: fluxtune_n10_smoke hyperparameters: rounds: 50 # short smoke run, not the full 300 + # Async fedbuff: accept stale grads + down-weight by (V'-V), not reject + # (K-D13). Set identically real+sim; never a parity lever. camelCase + # MUST match the base key to override it — snake_case is a dead key + # that loses at pydantic resolution (K-D15). + stalenessPolicy: fedbuff # Smoke-test stop bar: whichever of these fires first ends the run # (`rounds` alone is much too coarse, since one round is 150 # data-id completions). - max_data_id_progress: 10 # stop once data_id reaches 10 + max_data_id_progress: 9999 # effectively unbounded; runs governed by max_runtime_s (pass --max-data-id for a short run) max_runtime_s: 1800 # ...or after 30 minutes wall time data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 @@ -79,7 +91,7 @@ experiments: # the runner fans it into both this kwarg and # hyperparameters.aggGoal so they can't drift apart; don't # duplicate it here. - c: 8 + c: 10 minInitialTrainers: 10 execution: diff --git a/lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke_sim.yaml b/lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke_sim.yaml new file mode 100644 index 000000000..9dddc453e --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke_sim.yaml @@ -0,0 +1,152 @@ +# FedFwd SIMULATED-clock smoke test: 10 trainers, agnews H5 partitions +# (client_idx 0-9 of the 100 available niid_label_clients=100_alpha=1 +# partitions). Sim sibling of fluxtune_n10_smoke.yaml -- IDENTICAL except +# time_mode: simulated and the job id (so run dirs don't collide). +# +# fluxtune is the ASYNC baseline (async_oort selector, fedbuff), so this +# variant exercises the Batch-1 async grad-loop sim drive: _sim_recv_min_grad +# (sct-ordered reorder buffer + in-flight gate + virtual-clock advance) with +# agg-goal-boundary slot release. Trainers still do real forward-grad compute +# (gradient VALUES mode-invariant); only the clock/ordering differ. +# +# enable_training_delays stays false (as in the real smoke) for direct +# comparability; modeled-delay-on (D>0) is for later parity runs, enabled in +# BOTH real reference and sim run together (see simulate_fwdllm.md §K-D8). +# +# Run: +# python -m flame.launch.run_experiment \ +# lib/python/examples/fwdllm/expt_scripts/fluxtune_n10_smoke_sim.yaml + +experiments: + - name: fluxtune_n10_smoke_sim + description: "FedFwd SIM smoke (10 trainers, fluxtune baseline: async async_oort+fedbuff+JVP+3-tier, time_mode=simulated). Validates the Batch-1 async grad-loop sim drive." + + baseline: fluxtune + + example: + # fwdllm's trainer entrypoint has no pytorch/ subdirectory, unlike the + # ExampleConfig default (trainer/pytorch/main.py). + trainer_main: trainer/main.py + + trainer: + num_trainers: 10 + start_id: 1 + dataset: + name: agnews + # H5 path-style dataset -- no _metadata/dataset_splits/ index-list + # file exists (or ever will) for agnews; skips that lookup entirely. + path_style: true + availability: + # Phase 1 is syn_0 (100% availability): client_notify then reports all + # trainers available. Switch to mobiperf_3st_50 (via --avail-trace) only + # for the Phase-2 unavailability runs. run_sequential.sh's --avail-trace + # patches trainer mode + client_notify + trackTrainerAvail together. + mode: syn_0 + # Batch 1: drive the simulated virtual clock instead of real sleeps. + time_mode: simulated + # Skip the registry's 4-18s artificial delay for a fast smoke run. + enable_training_delays: false + # client_idx = (trainer_id - 1) % 100 -- wraps 10 trainers onto a + # distinct, non-colliding subset (0-9) of the 100 H5 partitions. + client_idx_modulo: 100 + config_overrides: + hyperparameters: + # Fluxtune JVP perf optimizations (simulate_fwdllm.md §L) — BIT-IDENTICAL + # grads; MUST match the real yaml (both sides read this config, so the + # optimization keeps real<->sim parity). Set false to revert. + jvp_perf_opt: true + data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 + partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 + # Uniform/IID partition for smoke tests -- avoids the niid_label + # baseline default's data-skew confound while validating mechanics. + partition_method: uniform + # Phase 4c (B2, #6): sim barrier straggler spread to match real's + # per-trainer completion dispersion (uniform[0,spread), spread ≈ + # real-compute-std·√12). Active with --delays on; calibrate to + # wall_disparity→~0 at the sign-off run. + sim_straggler_spread_s: 0.0 # P2-3/K-D29: per-trainer registry delays now supply the completion spread (remainder-wait max(gpu,D)); the crc32 offset would re-noise the deterministic order + + aggregator: + config_template: configs/aggregator_base.json + selector: async_oort + tracking_mode: client_notify + agg_goal: 3 + log_to_wandb: false + config_overrides: + job: + id: fluxtune_n10_smoke_sim + hyperparameters: + rounds: 50 # short smoke run, not the full 300 + # SIM-only: with c=10 >> agg_goal=3, commit-then-carry the surplus + # grads + hold busy trainers across the boundary, instead of dropping + # + re-dispatching them (the 2x-forward-pass residence bug, K-D12). + # Real satisfies this by channel construction, so it's inert there. + sim_inflight_residence: true + # #13 step 3: drain each in-flight end's rxq DIRECTLY (channel.drain_ready) + # instead of the recv_fifo streamer. The streamer blocks the full recv + # grace PER not-ready end (the 4-10s inter-burst gaps that dominated the + # post-step-2 wall) and can strand a delivered grad out of the buffer's + # view (past-dated commit). drain_ready sweeps ready rxqs non-blocking and + # returns on the first arrival. SIM-only (real never enters _sim_recv_min_grad). + sim_sct_ordered_drain: true + # #13 step 4: freed-slot staggered re-dispatch. INVESTIGATED + DISABLED + # (K-D28d, run_..._034550): neutral-to-slightly-worse for fwdllm (sim_rate + # 0.289->0.274, the 11-12s holds persisted). The holds are the drain + # correctly waiting (strict sct order) for the earliest-sct in-flight + # straggler, NOT a dispatch-bunching artifact; staggering feeds an OLD + # freed-slot vclock as send_ts, making the gate expect trainers even + # earlier -> more holds. Code kept (flag-gated) for a possible reworked + # attempt; see §H. Leave OFF -> validated steps 1-3 drain. + sim_staggered_redispatch: false + # #15: compute-truthful commit gate. The earlier_stuck gate blocks real + # wall on `_sim_inflight_expected` entries stamped at DISPATCH; a trainer + # whose payload was sent long ago but hasn't returned is idle-in-recv + # (not computing) behind the single-threaded drain, so waiting for it + # burns the grace floor / 30s failsafe and throttles commits -> via + # hold-to-commit, starves re-dispatch (sim concurrency 1.65 vs real 7.69). + # When on, the gate only blocks on a trainer still within its modeled + # compute window (sim_gate_compute_cap_s). Hold-to-commit UNTOUCHED; this + # only unblocks the commit path so held trainers are freed promptly. + # SIM-only; must not move cohort/var/staleness parity (PARITY_LOGICAL_TASKS.md #15). + sim_compute_truthful_gate: true + sim_gate_compute_cap_s: 10.0 + # Async fedbuff: accept stale carried grads + down-weight by (V'-V), + # not reject (K-D13). Set identically in the real sibling; never a + # parity lever. camelCase MUST match the base key to override it — + # snake_case is a dead key that loses at pydantic resolution (K-D15). + stalenessPolicy: fedbuff + # Smoke-test stop bar: whichever of these fires first ends the run + # (`rounds` alone is much too coarse, since one round is 150 + # data-id completions). + max_data_id_progress: 9999 # effectively unbounded; runs governed by max_runtime_s (pass --max-data-id for a short run) + max_runtime_s: 1800 # ...or after 30 minutes wall time + # Phase 4c (B1, #6): charge the measured server eval_s onto the vclock + # so it models real wall (self-calibrating). Delay-independent. + sim_model_eval_time: true + data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 + partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 + partition_method: uniform + selector: + kwargs: + # Scaled down from the baseline's n=150-ish defaults + # (c=15, aggGoal=5, minInitialTrainers=15) to fit 10 trainers. + # minInitialTrainers=10 = wait for all 10 smoke trainers to join. + # aggGoal is set once via the top-level `agg_goal: 3` above -- + # the runner fans it into both this kwarg and + # hyperparameters.aggGoal so they can't drift apart; don't + # duplicate it here. + c: 10 + minInitialTrainers: 10 + + execution: + # 8, not 1, even though this YAML's own default is 10 trainers -- + # run_sequential.sh's --num-trainers commonly scales this up to 100+ + # without a matching GPU bump (it has no num_gpus override of its + # own), and cramming 100 trainer processes onto 1 GPU OOMs almost + # immediately (see fluxtune_n10_smoke run 20260629_221116). Spreading + # across 8 GPUs even at the 10-trainer default is harmless headroom. + num_gpus: 8 + sleep_between_spawns: 2.0 + aggregator_warmup_time: 20 + monitoring: + enabled: false diff --git a/lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke.yaml b/lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke.yaml index 348eba5de..07871f491 100644 --- a/lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke.yaml +++ b/lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke.yaml @@ -64,7 +64,7 @@ experiments: # Smoke-test stop bar: whichever of these fires first ends the run # (`rounds` alone is much too coarse, since one round is 150 # data-id completions). - max_data_id_progress: 10 # stop once data_id reaches 10 + max_data_id_progress: 9999 # effectively unbounded; runs governed by max_runtime_s (pass --max-data-id for a short run) max_runtime_s: 1800 # ...or after 30 minutes wall time data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 diff --git a/lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke_sim.yaml b/lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke_sim.yaml new file mode 100644 index 000000000..ee6255704 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke_sim.yaml @@ -0,0 +1,109 @@ +# FedFwd SIMULATED-clock smoke test: 10 trainers, agnews H5 partitions +# (client_idx 0-9 of the 100 available niid_label_clients=100_alpha=1 +# partitions). Sim sibling of fwdllm_n10_smoke.yaml -- IDENTICAL except +# time_mode: simulated and the job id (so run dirs don't collide). +# +# time_mode: simulated exercises the Batch-1 sim path: the trainer skips the +# real emulated-delay sleep and stamps a modeled sim_completion_ts; the +# aggregator orders commits by that sct on a virtual clock (sync barrier here, +# via _sync_sim_recv_first_k). Trainers still do real forward-grad compute, so +# gradient VALUES are mode-invariant -- only the clock/ordering differ. +# +# enable_training_delays stays false (as in the real smoke) so this variant is +# DIRECTLY comparable to fwdllm_n10_smoke.yaml: with D=0 the sct = dispatch + +# real GPU time, which still drives the reorder buffer + vclock advance and +# validates the mechanics. Modeled-delay-on (D>0, the full sim behavior) is for +# the later convergence/parity runs, enabled in BOTH the real reference and the +# sim run together (see simulate_fwdllm.md §K-D8). +# +# Run: +# python -m flame.launch.run_experiment \ +# lib/python/examples/fwdllm/expt_scripts/fwdllm_n10_smoke_sim.yaml + +experiments: + - name: fwdllm_n10_smoke_sim + description: "FedFwd SIM smoke (10 trainers, fwdllm baseline: sync random+fedavg, time_mode=simulated). Validates the Batch-1 sim clock/ordering port." + + baseline: fwdllm + + example: + # fwdllm's trainer entrypoint has no pytorch/ subdirectory, unlike the + # ExampleConfig default (trainer/pytorch/main.py). + trainer_main: trainer/main.py + + trainer: + num_trainers: 10 + start_id: 1 + dataset: + name: agnews + # H5 path-style dataset -- no _metadata/dataset_splits/ index-list + # file exists (or ever will) for agnews; skips that lookup entirely. + path_style: true + availability: + mode: syn_0 + # Batch 1: drive the simulated virtual clock instead of real sleeps. + time_mode: simulated + # Skip the registry's 4-18s artificial delay for a fast smoke run. + enable_training_delays: false + # client_idx = (trainer_id - 1) % 100 -- wraps 10 trainers onto a + # distinct, non-colliding subset (0-9) of the 100 H5 partitions. + client_idx_modulo: 100 + config_overrides: + hyperparameters: + data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 + partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 + # Uniform/IID partition for smoke tests -- the niid_label baseline + # default adds a data-skew confound we don't want while validating + # the launcher mechanics. + partition_method: uniform + # Phase 4c (B2, #6): widen the sim barrier's k-th-smallest sct to + # match real's per-trainer completion dispersion (banked real compute + # std ~0.26s; uniform[0,spread) offset ⇒ spread ≈ 0.26·√12). Only + # active with --delays on (convergence run); inert at D=0. Calibrate + # to wall_disparity→~0 at the sign-off run. + sim_straggler_spread_s: 0.0 # P2-3/K-D29: per-trainer registry delays now supply the completion spread (remainder-wait max(gpu,D)); the crc32 offset would re-noise the deterministic order + client_notify: + trace: syn_0 + + aggregator: + config_template: configs/aggregator_base.json + selector: random + tracking_mode: default + # Matches selector.kwargs.c below -- every selected trainer is required. + agg_goal: 10 + log_to_wandb: false + config_overrides: + job: + id: fwdllm_n10_smoke_sim + hyperparameters: + rounds: 50 # short smoke run, not the full 300 + # Smoke-test stop bar: whichever of these fires first ends the run + # (`rounds` alone is much too coarse, since one round is 150 + # data-id completions). + max_data_id_progress: 9999 # effectively unbounded; runs governed by max_runtime_s (pass --max-data-id for a short run) + max_runtime_s: 1800 # ...or after 30 minutes wall time + # Phase 4c (B1, #6): charge the MEASURED server eval_s (~13s/data_id -- + # the syn_0 wall dominator) onto the vclock so it models real wall and + # sim_rate approaches/exceeds 1 (self-calibrating: uses the actual + # eval wall, not a constant). Delay-independent; the point of Phase 2/4. + sim_model_eval_time: true + data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 + partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 + partition_method: uniform + selector: + kwargs: + # Scaled down from the baseline's n=150-ish defaults (k=5, c=15) + # to fit 10 trainers. + k: 5 + c: 10 + minInitialTrainers: 10 + + execution: + # 8, not 1 -- see fluxtune_n10_smoke.yaml's execution.num_gpus comment; + # run_sequential.sh's --num-trainers scales trainer count without a + # matching GPU bump, and 100 trainers on 1 GPU OOMs almost immediately. + num_gpus: 8 + sleep_between_spawns: 2.0 + aggregator_warmup_time: 20 + monitoring: + enabled: false diff --git a/lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke.yaml b/lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke.yaml index 119238eff..42dcdb686 100644 --- a/lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke.yaml +++ b/lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke.yaml @@ -31,10 +31,12 @@ experiments: # file exists (or ever will) for agnews; skips that lookup entirely. path_style: true availability: - # ORACULAR reads trainer 1-10's mobiperf_2st trace from _metadata -- - # mode here just keeps the trainer-side avl_events_* fields - # consistent with what the aggregator reads. - mode: mobiperf_2st + # Phase 1 is syn_0 (100% availability); the ORACULAR aggregator then + # reads all-available, so the sync barrier assembles the full agg_goal=10 + # cohort. Switch to a mobiperf_* trace (via --avail-trace) only for the + # Phase-2 unavailability runs. Keep trainer + aggregator on the SAME + # trace (run_sequential.sh's --avail-trace patches both). + mode: syn_0 # FedFwd has no simulated-clock support; always real. time_mode: real # Skip the registry's 4-18s artificial delay for a fast smoke run. @@ -54,7 +56,7 @@ experiments: config_template: configs/aggregator_base.json selector: random tracking_mode: oracular - agg_goal: 2 + agg_goal: 10 log_to_wandb: false config_overrides: job: @@ -64,7 +66,7 @@ experiments: # Smoke-test stop bar: whichever of these fires first ends the run # (`rounds` alone is much too coarse, since one round is 150 # data-id completions). - max_data_id_progress: 10 # stop once data_id reaches 10 + max_data_id_progress: 9999 # effectively unbounded; runs governed by max_runtime_s (pass --max-data-id for a short run) max_runtime_s: 1800 # ...or after 30 minutes wall time data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 diff --git a/lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke_sim.yaml b/lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke_sim.yaml new file mode 100644 index 000000000..10f496558 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke_sim.yaml @@ -0,0 +1,104 @@ +# FedFwd SIMULATED-clock smoke test: 10 trainers, agnews H5 partitions +# (client_idx 0-9 of the 100 available niid_label_clients=100_alpha=1 +# partitions). Sim sibling of fwdllm_plus_n10_smoke.yaml -- IDENTICAL except +# time_mode: simulated and the job id (so run dirs don't collide). +# +# fwdllm_plus is SYNC (random selector, fedavg, ORACULAR availability read, +# per-iteration reselection), so this variant exercises the Batch-1 sync +# barrier sim drive (_sync_sim_recv_first_k: commit k-smallest-sct, advance the +# virtual clock to the k-th). Trainers still do real forward-grad compute +# (gradient VALUES mode-invariant); only the clock/ordering differ. +# +# enable_training_delays stays false (as in the real smoke) for direct +# comparability; modeled-delay-on (D>0) is for later parity runs, enabled in +# BOTH real reference and sim run together (see simulate_fwdllm.md §K-D8). +# +# Run: +# python -m flame.launch.run_experiment \ +# lib/python/examples/fwdllm/expt_scripts/fwdllm_plus_n10_smoke_sim.yaml + +experiments: + - name: fwdllm_plus_n10_smoke_sim + description: "FedFwd SIM smoke (10 trainers, fwdllm_plus baseline: sync random+fedavg+ORACULAR, time_mode=simulated). Validates the Batch-1 sync barrier sim drive." + + baseline: fwdllm_plus + + example: + # fwdllm's trainer entrypoint has no pytorch/ subdirectory, unlike the + # ExampleConfig default (trainer/pytorch/main.py). + trainer_main: trainer/main.py + + trainer: + num_trainers: 10 + start_id: 1 + dataset: + name: agnews + # H5 path-style dataset -- no _metadata/dataset_splits/ index-list + # file exists (or ever will) for agnews; skips that lookup entirely. + path_style: true + availability: + # Phase 1 is syn_0 (100% availability); the ORACULAR aggregator then + # reads all-available, so the sync barrier assembles the full agg_goal=10 + # cohort. Switch to a mobiperf_* trace (via --avail-trace) only for the + # Phase-2 unavailability runs. Keep trainer + aggregator on the SAME + # trace (run_sequential.sh's --avail-trace patches both). + mode: syn_0 + # Batch 1: drive the simulated virtual clock instead of real sleeps. + time_mode: simulated + # Skip the registry's 4-18s artificial delay for a fast smoke run. + enable_training_delays: false + # client_idx = (trainer_id - 1) % 100 -- wraps 10 trainers onto a + # distinct, non-colliding subset (0-9) of the 100 H5 partitions. + client_idx_modulo: 100 + config_overrides: + hyperparameters: + data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 + partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 + # Uniform/IID partition for smoke tests -- avoids the niid_label + # baseline default's data-skew confound while validating mechanics. + partition_method: uniform + # Phase 4c (B2, #6): sim barrier straggler spread to match real's + # per-trainer completion dispersion (uniform[0,spread), spread ≈ + # real-compute-std·√12). Active with --delays on; calibrate to + # wall_disparity→~0 at the sign-off run. + sim_straggler_spread_s: 0.0 # P2-3/K-D29: per-trainer registry delays now supply the completion spread (remainder-wait max(gpu,D)); the crc32 offset would re-noise the deterministic order + + aggregator: + config_template: configs/aggregator_base.json + selector: random + tracking_mode: oracular + agg_goal: 10 + log_to_wandb: false + config_overrides: + job: + id: fwdllm_plus_n10_smoke_sim + hyperparameters: + rounds: 50 # short smoke run, not the full 300 + # Smoke-test stop bar: whichever of these fires first ends the run + # (`rounds` alone is much too coarse, since one round is 150 + # data-id completions). + max_data_id_progress: 9999 # effectively unbounded; runs governed by max_runtime_s (pass --max-data-id for a short run) + max_runtime_s: 1800 # ...or after 30 minutes wall time + # Phase 4c (B1, #6): charge the measured server eval_s onto the vclock + # so it models real wall (self-calibrating). Delay-independent. + sim_model_eval_time: true + data_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/data_files/agnews_data.h5 + partition_file_path: /coc/scratch/dgarg/fl_datasets/fwdllm/fednlp_data/partition_files/agnews_partition.h5 + partition_method: uniform + selector: + kwargs: + # Scaled down from the baseline's n=150-ish defaults (k=5, c=15) + # to fit 10 trainers. + k: 5 + c: 10 + minInitialTrainers: 10 + + execution: + # 8, not 1 -- see fluxtune_n10_smoke.yaml's execution.num_gpus comment; + # run_sequential.sh's --num-trainers scales trainer count without a + # matching GPU bump, and 100 trainers on 1 GPU OOMs almost immediately. + num_gpus: 8 + sleep_between_spawns: 2.0 + aggregator_warmup_time: 20 + monitoring: + enabled: false diff --git a/lib/python/examples/fwdllm/expt_scripts/logical_parity.py b/lib/python/examples/fwdllm/expt_scripts/logical_parity.py new file mode 100644 index 000000000..cd6456fef --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/logical_parity.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""Time-STRIPPED logical parity diff, real vs sim, up to a data-bin cap. + +Does the sim take the SAME steps in the SAME order as real, IGNORING every +wall/vclock timestamp? Logical parity is checked before the time dimension +(sim_rate). Reads already-banked aggregator telemetry — no re-run. + +The authoritative record is the `agg_round` event stream (data_id, +iteration_per_data_id, receive-ordered contributing_trainers, variance decision). +Restricts to data_id <= --max-bin (default 1) and compares, per aggregation cycle: + - receive SET : which trainers committed together (async cohort composition) + - receive ORDER : arrival order (only meaningful for async; a sync barrier + commits the whole cohort, so order within a set is irrelevant) + - cadence tuple : (data_id, iter, agg_goal_count, var_good_enough, force_commit) +plus iterations-to-clear-bin-0 and the per-cycle variance trajectory. + +Usage: python logical_parity.py [--max-bin 1] [--baselines fluxtune ...] +""" +from __future__ import annotations +import argparse, glob, json, os +from pathlib import Path + +_EXP = Path(__file__).resolve().parent.parent / "experiments" + + +def _short(t): return t[-3:] if len(t) >= 3 else t + + +def _agg_trace(run_dir: Path, max_bin: int): + f = next(iter(glob.glob(str(run_dir / "telemetry" / "aggregator_*.jsonl"))), None) + out = [] + for line in open(f): + try: + e = json.loads(line) + except Exception: + continue + if e.get("event") != "agg_round": + continue + did = e.get("data_id") + if did is None or did > max_bin: + continue + out.append({ + "data_id": did, + "iter": e.get("iteration_per_data_id"), + "goal": e.get("agg_goal_count"), + "order": [_short(t) for t in (e.get("contributing_trainers") or [])], + "var_good": e.get("var_good_enough"), + "force": e.get("force_commit_planned"), + "var": e.get("var"), + }) + return out + + +def _find_pair(baseline: str): + def pick(variant): + cands = [p for p in glob.glob(str(_EXP / f"run_*_{baseline}_n*_smoke_*_{variant}")) + if f"_{baseline}_n" in os.path.basename(p)] # fwdllm !=> fwdllm_plus + return Path(sorted(cands)[-1]) if cands else None + return pick("real"), pick("sim") + + +def _iters_to_clear_bin0(trace): + return sum(1 for x in trace if x["data_id"] == 0) + + +def _cmp(baseline: str, max_bin: int): + rdir, sdir = _find_pair(baseline) + print(f"\n==================== {baseline} (data bin <= {max_bin}) ====================") + print(f" real: {rdir.name if rdir else None}") + print(f" sim : {sdir.name if sdir else None}") + if not (rdir and sdir): + print(" MISSING pair"); return + r, s = _agg_trace(rdir, max_bin), _agg_trace(sdir, max_bin) + + n = min(len(r), len(s)) + set_match = sum(1 for i in range(n) if sorted(r[i]["order"]) == sorted(s[i]["order"])) + ord_match = sum(1 for i in range(n) if r[i]["order"] == s[i]["order"]) + cad = lambda x: (x["data_id"], x["iter"], x["goal"], x["var_good"], x["force"]) + cad_match = sum(1 for i in range(n) if cad(r[i]) == cad(s[i])) + + print(f" aggregations compared: {n} (real={len(r)}, sim={len(s)})") + print(f" receive-SET identical: {set_match}/{n}") + print(f" receive-ORDER identical: {ord_match}/{n} (sync barrier: order-within-set is benign)") + print(f" cadence identical: {cad_match}/{n} (data_id,iter,goal,var_good,force)") + print(f" iters to clear data bin 0: real={_iters_to_clear_bin0(r)} sim={_iters_to_clear_bin0(s)}") + + verdict = ("LOGICAL PARITY" if cad_match == n and set_match == n else + "CADENCE PARITY (cohorts differ)" if cad_match == n else + "LOGICAL DIVERGENCE") + print(f" => {verdict}") + + if cad_match != n or set_match != n: + print(f" {'REAL':<44} | SIM") + for i in range(max(len(r), len(s))): + rr = (r[i]["data_id"], r[i]["iter"], tuple(sorted(r[i]["order"])), + r[i]["var_good"], round(r[i]["var"] or 0, 3)) if i < len(r) else "" + ss = (s[i]["data_id"], s[i]["iter"], tuple(sorted(s[i]["order"])), + s[i]["var_good"], round(s[i]["var"] or 0, 3)) if i < len(s) else "" + flag = "" if (rr and ss and rr[2] == ss[2]) else " <- cohort differs" + print(f" {str(rr):<44} | {ss}{flag}") + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--max-bin", type=int, default=1) + ap.add_argument("--baselines", nargs="+", default=["fwdllm", "fwdllm_plus", "fluxtune"]) + a = ap.parse_args() + for b in a.baselines: + _cmp(b, a.max_bin) diff --git a/lib/python/examples/fwdllm/expt_scripts/make_paper_figs.py b/lib/python/examples/fwdllm/expt_scripts/make_paper_figs.py new file mode 100644 index 000000000..26067590d --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/make_paper_figs.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Render the SOCC-2026 paper figures from a baseline→run-dir mapping. + +The mapping is data: `figs.yaml` maps each baseline to a run dir, overridable with +`--run KEY=PATH`. Missing baselines are skipped. Output goes to a timestamped +`paper_figs//` (+ `latest` symlink) with stable basenames + a manifest.json. + + python make_paper_figs.py [--run fluxtune=/data/run_...] [--figures e1_acc_vs_time,...] +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +if HERE not in sys.path: + sys.path.insert(0, HERE) + +from plotlib import baselines as B # noqa: E402 +from plotlib import figures as F # noqa: E402 +from plotlib import reducers as R # noqa: E402 +from plotlib import style as S # noqa: E402 + +# run dirs in the manifest are resolved relative to the fwdllm example dir +EXAMPLE_DIR = os.path.abspath(os.path.join(HERE, "..")) + + +def _resolve(path: str) -> str: + if os.path.isabs(path): + return path + for base in (os.getcwd(), EXAMPLE_DIR): + cand = os.path.join(base, path) + if os.path.exists(cand): + return os.path.abspath(cand) + return os.path.abspath(os.path.join(EXAMPLE_DIR, path)) # best effort + + +def _load_manifest(path: str) -> dict: + if not path or not os.path.exists(path): + return {} + try: + import yaml + with open(path) as fh: + return yaml.safe_load(fh) or {} + except Exception as e: # noqa: BLE001 + print(f" [paper-figs] could not read manifest {path}: {e}", file=sys.stderr) + return {} + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--manifest", default=os.path.join(HERE, "figs.yaml"), + help="baseline→run-dir manifest (default: expt_scripts/figs.yaml)") + ap.add_argument("--run", action="append", default=[], metavar="KEY=PATH", + help="add/override one baseline's run dir (repeatable)") + ap.add_argument("--out-root", default=os.path.join(HERE, "paper_figs"), + help="root for timestamped output dirs") + ap.add_argument("--figures", default=None, + help="comma list of figure names (default: all)") + ap.add_argument("--target-acc", type=float, default=None) + ap.add_argument("--smooth", type=float, default=0.7, + help="EMA line-smoothing factor ∈ [0,1) for learning curves " + "(visual only; 0 = raw). Default 0.7") + ap.add_argument("--loss-plateau-rel", type=float, default=0.01, + help="cut each run at its last cumulative test-loss drop of this " + "relative size — the deterministic end of productive " + "learning (default 0.01 = 1%%; the stall-guard rule)") + ap.add_argument("--post-peak-grace-min", type=float, default=0.0, + help="minutes of tail to keep past the loss-plateau point " + "(default 0 — the plateau is deterministic, no buffer needed)") + ap.add_argument("--no-cutoff", action="store_true", + help="disable the cutoff (use full telemetry)") + ap.add_argument("--cutoff-mode", choices=["peak_acc", "loss_plateau"], + default="peak_acc", + help="where to clip each run: 'peak_acc' (default) at its " + "highest-accuracy eval — every baseline stops very close " + "to its best accuracy, clipping post-peak drift/divergence; " + "'loss_plateau' at the last significant test-loss drop. " + "A run that never learned (flat loss) is never clipped.") + args = ap.parse_args() + grace_s = None if args.no_cutoff else args.post_peak_grace_min * 60.0 + + manifest = _load_manifest(args.manifest) + runs_map = dict(manifest.get("runs", {}) or {}) + target = args.target_acc if args.target_acc is not None else manifest.get("target_acc") + for spec in args.run: # CLI overrides the manifest + if "=" not in spec: + print(f" [paper-figs] ignoring malformed --run '{spec}' (want KEY=PATH)", + file=sys.stderr) + continue + k, v = spec.split("=", 1) + runs_map[k.strip()] = v.strip() + if not runs_map: + print(" [paper-figs] no baselines mapped — populate figs.yaml or pass --run", + file=sys.stderr) + return 1 + + # load each mapped run (streaming); skip + warn on missing/empty telemetry + loaded = {} + for key in B.ordered(runs_map.keys()): + rr = R.load_run(_resolve(runs_map[key]), key=key, post_peak_grace_s=grace_s, + loss_plateau_rel=args.loss_plateau_rel, + cutoff_mode=args.cutoff_mode) + if rr is None or not rr.evals: + print(f" [paper-figs] WARN no usable telemetry for '{key}' " + f"({runs_map[key]}) — skipping", file=sys.stderr) + continue + loaded[key] = rr + cut = " (no cutoff — never learned)" if rr.cutoff_ts is None else ( + f", cut @{args.cutoff_mode} {(rr.cutoff_ts - rr.t0)/3600:.2f}h") + print(f" [paper-figs] loaded {key}: {len(rr.evals)} evals, " + f"{rr.n_trainers} trainers, max_acc={100*(rr.max_accuracy() or 0):.2f}%{cut}") + if not loaded: + print(" [paper-figs] nothing loaded — nothing to plot.", file=sys.stderr) + return 1 + + ordered_runs = [loaded[k] for k in B.ordered(loaded.keys())] + which = ([f.strip() for f in args.figures.split(",")] + if args.figures else list(F.FIG_BUILDERS)) + + S.use_paper_style() + out_dir = S.timestamped_outdir(args.out_root) + written = [] + for name in which: + builder = F.FIG_BUILDERS.get(name) + if builder is None: + print(f" [paper-figs] unknown figure '{name}' — skipping", file=sys.stderr) + continue + fig = builder(ordered_runs, target=target, smooth=args.smooth) + if fig is None: + print(f" [paper-figs] {name}: no data across baselines — skipped") + continue + S.save_pdf(fig, out_dir, name) + written.append(name) + + with open(os.path.join(out_dir, "manifest.json"), "w") as fh: + json.dump({ + "runs": {k: loaded[k].run_dir for k in loaded}, + "baselines_order": B.ordered(loaded.keys()), + "target_acc": target, + "figures": written, + "system_label": B.SYSTEM_LABEL, + "smooth": args.smooth, + "cutoff_mode": None if args.no_cutoff else args.cutoff_mode, + "loss_plateau_rel": None if args.no_cutoff else args.loss_plateau_rel, + "post_peak_grace_min": None if args.no_cutoff else args.post_peak_grace_min, + "cutoff_h": {k: (None if loaded[k].cutoff_ts is None + else round((loaded[k].cutoff_ts - loaded[k].t0) / 3600, 3)) + for k in loaded}, + }, fh, indent=2) + + print(f"\n wrote {len(written)} figure(s) → {out_dir}") + print(f" latest → {os.path.join(os.path.abspath(args.out_root), 'latest')}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/lib/python/examples/fwdllm/expt_scripts/plot_run.py b/lib/python/examples/fwdllm/expt_scripts/plot_run.py new file mode 100644 index 000000000..d2324e377 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/plot_run.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Single-run diagnostics — rich multi-panel PDFs for ONE run dir (per-run twin of +compare_baselines.py). Metrics come from `plotlib.reducers` (shared, streaming). + + python plot_run.py --run-dir experiments/run_..._fluxtune_..._real [--target-acc 0.84] + +Writes vector-PDF composite panels + summary.json. Learning curves are over +wall-clock / eval order (not data_id, which cycles per round), round boundaries marked. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +HERE = os.path.dirname(os.path.abspath(__file__)) +if HERE not in sys.path: + sys.path.insert(0, HERE) +from plotlib import reducers as R # noqa: E402 (single source of the metrics) +from plotlib import style as S # noqa: E402 + + +def _pct(xs, q): + import numpy as np + xs = [x for x in xs if x is not None] + return float(np.percentile(xs, q)) if xs else None + + +def _pctiles(xs): + return {"p50": _pct(xs, 50), "p90": _pct(xs, 90), "p99": _pct(xs, 99)} + + +# --------------------------------------------------------------------------- # +# plotting — composite per-run panels +# --------------------------------------------------------------------------- # +def make_plots(rr: R.RunResult, out_dir, target, smooth=0.0): + import numpy as np + import matplotlib.pyplot as plt + S.use_paper_style() + + def save(fig, name): + fig.tight_layout() + fig.savefig(os.path.join(out_dir, f"{name}.pdf")) + plt.close(fig) + + hrs, acc_raw, loss_raw, _rnd = rr.learning_curve() + acc = S.ema(acc_raw, smooth) # visual smoothing (scalars use raw) + loss = S.ema(loss_raw, smooth) + idx = list(range(len(rr.evals))) # cumulative eval index + trans = rr.round_transition_indices() + + # ---- Experiment 1: learning curves (time / eval-index / loss) ---------- + fig, ax = plt.subplots(1, 3, figsize=(15, 4.0)) + if hrs: + ax[0].plot(hrs, acc, lw=1.6, color="#2b6cb0") + ax[0].axhline(target * 100, ls="--", color="#c53030", lw=1, + label=f"target {target*100:.0f}%") + for i in trans: + ax[0].axvline(hrs[i], color="#dd6b20", ls=":", lw=0.9, alpha=0.7) + pa = rr.max_accuracy() * 100 + pi = max(idx, key=lambda i: (acc_raw[i] if acc_raw[i] is not None else -1)) + ax[0].scatter([hrs[pi]], [acc_raw[pi]], color="#dd6b20", zorder=5, + label=f"peak {pa:.2f}%") + ax[0].set(xlabel="wall-clock time (h)", ylabel="test accuracy (%)", + title="E1 — accuracy vs wall time") + ax[0].legend(fontsize=8); ax[0].grid(alpha=.3) + ax[1].plot(idx, acc, lw=1.4, color="#2b6cb0") + for i in trans: + ax[1].axvline(i, color="#dd6b20", ls=":", lw=0.9, alpha=0.7, + label="round start" if i == trans[0] else None) + ax[1].set(xlabel="cumulative eval (data bin)", ylabel="test accuracy (%)", + title="E1 — accuracy vs eval index") + if trans: + ax[1].legend(fontsize=8) + ax[1].grid(alpha=.3) + ax[2].plot(idx, loss, lw=1.4, color="#805ad5") + for i in trans: + ax[2].axvline(i, color="#dd6b20", ls=":", lw=0.9, alpha=0.7) + ax[2].set(xlabel="cumulative eval (data bin)", ylabel="test loss", + title="E1 — loss vs eval index") + ax[2].grid(alpha=.3) + save(fig, "e1_learning_curves") + + # ---- Experiment 2: utilization ----------------------------------------- + fig, ax = plt.subplots(1, 2, figsize=(11, 4.0)) + bf = rr.busy_frac + if bf: + ax[0].hist(np.array(bf) * 100, bins=25, color="#38a169", alpha=.75) + for q, c in [(50, "#2b6cb0"), (90, "#dd6b20"), (99, "#c53030")]: + v = _pct(bf, q) * 100 + ax[0].axvline(v, ls="--", color=c, lw=1, label=f"P{q}={v:.1f}%") + ax[0].set(xlabel="trainer busy fraction (%)", ylabel="# trainers", + title=f"E2 — trainer busy fraction (n={len(bf)})") + ax[0].legend(fontsize=8) + w = rr.agg_wall_s or 1.0 + comp, barr, drain = rr.agg_compute_s, rr.agg_barrier_s, rr.agg_drain_s + other = max(0.0, w - comp - barr - drain) + parts = [("compute", comp), ("barrier-wait", barr), ("drain", drain), ("idle/other", other)] + cols = ["#38a169", "#dd6b20", "#805ad5", "#a0aec0"] + ax[1].bar([p[0] for p in parts], [p[1] / w * 100 for p in parts], color=cols) + ax[1].set(ylabel="% of aggregator wall", title="E2 — aggregator time decomposition") + for i, p in enumerate(parts): + ax[1].text(i, p[1] / w * 100, f"{p[1]/w*100:.1f}%", ha="center", va="bottom", fontsize=8) + save(fig, "e2_utilization") + + # ---- Experiment 3: compute productivity -------------------------------- + fig, ax = plt.subplots(1, 2, figsize=(11, 4.0)) + if loss: + dloss = [loss[0] - v for v in loss] + ax[0].plot(idx, dloss, lw=1.6, color="#2b6cb0") + for i in trans: + ax[0].axvline(i, color="#dd6b20", ls=":", lw=0.9, alpha=0.7) + ax[0].set(xlabel="cumulative eval (data bin)", ylabel="Δloss (from first)", + title="E3 — cumulative loss reduction") + ax[0].grid(alpha=.3) + gpu = rr.gpu_s_total() + dl = rr.delta_loss() + labels = ["Δloss / GPU-hour", "Δloss / Mfwd-pass"] + vals = [ + (dl / (gpu / 3600)) if (dl and gpu) else 0, + (dl / (rr.fwd_total / 1e6)) if (dl and rr.fwd_total) else 0, + ] + ax[1].bar(labels, vals, color=["#38a169", "#805ad5"]) + for i, v in enumerate(vals): + ax[1].text(i, v, f"{v:.4g}", ha="center", va="bottom", fontsize=9) + ax[1].set(title="E3 — learning per unit compute") + save(fig, "e3_productivity") + + # ---- Experiment 4: network --------------------------------------------- + fig, ax = plt.subplots(1, 3, figsize=(15, 4.0)) + down = sum(rr.down_sizes) + up = sum(rr.up_sizes) + ax[0].bar(["agg→trainer\n(down)", "trainer→agg\n(up)"], [down / 1e9, up / 1e9], + color=["#2b6cb0", "#dd6b20"]) + for i, v in enumerate([down, up]): + ax[0].text(i, v / 1e9, f"{v/1e9:.2f} GB", ha="center", va="bottom", fontsize=9) + ax[0].set(ylabel="total bytes (GB)", title="E4 — data transmitted") + if rr.down_sizes: + ax[1].hist(np.array(rr.down_sizes) / 1e6, bins=40, color="#2b6cb0", alpha=.7) + ax[1].set(xlabel="message size (MB)", ylabel="# messages", + title="E4 — agg→trainer msg sizes") + if rr.up_sizes: + ax[2].hist(np.array(rr.up_sizes) / 1e6, bins=40, color="#dd6b20", alpha=.7) + ax[2].set(xlabel="message size (MB)", ylabel="# messages", + title="E4 — trainer→agg msg sizes") + save(fig, "e4_network") + + # ---- Experiment 5: sessions & participation ---------------------------- + fig, ax = plt.subplots(1, 3, figsize=(15, 4.0)) + sd = rr.session_durs + if sd: + ax[0].hist(sd, bins=40, color="#805ad5", alpha=.75) + for q, c in [(50, "#2b6cb0"), (90, "#dd6b20"), (99, "#c53030")]: + v = _pct(sd, q) + ax[0].axvline(v, ls="--", color=c, lw=1, label=f"P{q}={v:.1f}s") + ax[0].set(xlabel="session duration (s)", ylabel="# sessions", + title=f"E5 — session durations (n={len(sd)})") + ax[0].legend(fontsize=8) + xs = np.sort(sd) + ax[1].plot(xs, np.arange(1, len(xs) + 1) / len(xs) * 100, color="#805ad5") + ax[1].set(xlabel="session duration (s)", ylabel="CDF (%)", + title="E5 — session-duration CDF") + ax[1].grid(alpha=.3) + if rr.part_iters: + ax[2].hist(rr.part_iters, bins=25, color="#38a169", alpha=.6, label="iterations") + ax[2].hist(rr.part_bins, bins=25, color="#dd6b20", alpha=.6, label="data bins") + ax[2].set(xlabel="count per client", ylabel="# clients", + title="E5 — participation per client") + ax[2].legend(fontsize=8) + save(fig, "e5_sessions") + + print(f" [plot_run] wrote 5 PDF panels to {out_dir}") + + +# --------------------------------------------------------------------------- # +# driver +# --------------------------------------------------------------------------- # +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--run-dir", required=True) + ap.add_argument("--target-acc", type=float, default=0.84) + ap.add_argument("--out", default=None, help="default: /plots") + ap.add_argument("--smooth", type=float, default=0.7, + help="EMA line-smoothing factor ∈ [0,1) (visual only; 0 = raw)") + ap.add_argument("--loss-plateau-rel", type=float, default=0.01, + help="cut at last cumulative test-loss drop of this relative size " + "(default 0.01 = 1%%; end of productive learning)") + ap.add_argument("--post-peak-grace-min", type=float, default=0.0, + help="minutes to keep past the loss-plateau point (default 0)") + ap.add_argument("--no-cutoff", action="store_true", help="use full telemetry") + args = ap.parse_args() + grace_s = None if args.no_cutoff else args.post_peak_grace_min * 60.0 + + run_dir = os.path.abspath(args.run_dir) + print(f" [plot_run] streaming telemetry from {os.path.basename(run_dir)}…") + rr = R.load_run(run_dir, key=os.path.basename(run_dir), post_peak_grace_s=grace_s, + loss_plateau_rel=args.loss_plateau_rel) + if rr is None: + print(f" [plot_run] no aggregator telemetry under {run_dir}/telemetry", file=sys.stderr) + return 1 + out_dir = os.path.abspath(args.out) if args.out else os.path.join(run_dir, "plots") + os.makedirs(out_dir, exist_ok=True) + + losses = [e["loss"] for e in rr.evals if e["loss"] is not None] + dl = rr.delta_loss() + gpu = rr.gpu_s_total() + down, up = sum(rr.down_sizes), sum(rr.up_sizes) + w = rr.agg_wall_s or None + wall_h = round((rr.evals[-1]["ts"] - rr.t0) / 3600, 3) if (rr.evals and rr.t0) else None + summary = { + "run_dir": os.path.basename(run_dir), "key": rr.key, "target_acc": args.target_acc, + "e1_max_accuracy": rr.max_accuracy(), "e1_final_accuracy": rr.final_accuracy(), + "e1_n_evals": len(rr.evals), "e1_wall_h": wall_h, + "e2_trainer_busy_frac": _pctiles(rr.busy_frac), + "e2_agg_compute_frac": round(rr.agg_compute_s / w, 4) if w else None, + "e2_agg_barrier_frac": round(rr.agg_barrier_s / w, 4) if w else None, + "e2_agg_drain_frac": round(rr.agg_drain_s / w, 4) if w else None, + "e3_delta_loss": dl, "e3_gpu_s_total": round(gpu, 1), + "e3_forward_passes_total": rr.fwd_total if rr.have_fwd else None, + "e3_perturbations_total": rr.pert_total if rr.have_fwd else None, + "e3_delta_loss_per_gpu_hour": round(dl / (gpu / 3600), 6) if (dl and gpu) else None, + "e3_delta_loss_per_Mfwd": round(dl / (rr.fwd_total / 1e6), 6) + if (dl and rr.have_fwd and rr.fwd_total) else None, + "e4_bytes_down": down, "e4_bytes_up": up, "e4_bytes_total": down + up, + "e4_msgs_down": len(rr.down_sizes), "e4_msgs_up": len(rr.up_sizes), + "e4_bytes_by_kind": {f"{d}/{k}": v for (d, k), v in rr.comm_by_kind.items()}, + "e5_session_s": _pctiles(rr.session_durs), "e5_n_sessions": len(rr.session_durs), + "e5_part_iters": _pctiles(rr.part_iters), "e5_part_databins": _pctiles(rr.part_bins), + "n_trainers": rr.n_trainers, + } + summary["cutoff_h"] = (None if rr.cutoff_ts is None + else round((rr.cutoff_ts - rr.t0) / 3600, 3)) + summary["loss_plateau_rel"] = None if args.no_cutoff else args.loss_plateau_rel + summary["post_peak_grace_min"] = None if args.no_cutoff else args.post_peak_grace_min + with open(os.path.join(out_dir, "summary.json"), "w") as fh: + json.dump(summary, fh, indent=2) + + make_plots(rr, out_dir, args.target_acc, args.smooth) + + print(f"\n=== {os.path.basename(run_dir)} ===") + print(f" E1 max acc : {(summary['e1_max_accuracy'] or 0)*100:.2f}% " + f"over {summary['e1_n_evals']} evals / {summary['e1_wall_h']} h") + bf = summary["e2_trainer_busy_frac"] + if bf["p50"] is not None: + print(f" E2 trainer busy : P50={bf['p50']*100:.1f}% P90={bf['p90']*100:.1f}% " + f"P99={bf['p99']*100:.1f}% agg compute={summary['e2_agg_compute_frac']}") + print(f" E3 Δloss : {summary['e3_delta_loss']} " + f"| {summary['e3_gpu_s_total']} GPU-s | {summary['e3_forward_passes_total']} fwd-passes") + print(f" E4 network : {summary['e4_bytes_total']/1e9:.2f} GB total " + f"({summary['e4_msgs_down']:,}↓ / {summary['e4_msgs_up']:,}↑ msgs)") + ss = summary["e5_session_s"] + if ss["p50"] is not None: + print(f" E5 session dur : P50={ss['p50']:.1f}s P90={ss['p90']:.1f}s P99={ss['p99']:.1f}s " + f"({summary['e5_n_sessions']:,} sessions)") + print(f"\n summary.json + plots → {out_dir}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/lib/python/examples/fwdllm/expt_scripts/plot_step_timing.py b/lib/python/examples/fwdllm/expt_scripts/plot_step_timing.py new file mode 100644 index 000000000..a3d57af0c --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/plot_step_timing.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Decompose fwdllm trainer wall time by compute step, for GPU optimization. + +Two granularities, both read from a run's ``telemetry/trainer_*.jsonl``: + + * FINE — the ``step_timing`` events emitted by ``@timer_decorator``: + per-function wall (``_make_model_functional``, ``_setup_training_state`` / + perturbation selection, ``_train_one_batch`` / JVP forward passes, + ``_emulate_training_delay``) — the "where does the GPU time go" view that + tells us WHAT to optimize. + * COARSE — the per-round phase fields already on ``trainer_round`` + (``gpu_compute_s``, ``pre_train_s``, ``weights_to_gpu_s``, + ``weights_to_ram_s``, ``mqtt_fetch_s``, ``post_train_s``). Always present; + a fallback when a run predates the step_timing telemetry. + +Prints a breakdown table and (unless ``--no-plot``) writes PNGs next to the run. + +Usage: + python plot_step_timing.py # latest sim run per baseline + python plot_step_timing.py --run-dir DIR ... # explicit run dir(s) + python plot_step_timing.py --baselines fluxtune +""" +from __future__ import annotations + +import argparse +import collections +import json +import os +import pathlib +import re +import sys + +_HERE = pathlib.Path(__file__).resolve().parent +_DEFAULT_EXPERIMENTS = _HERE.parent / "experiments" + +# Same run-dir grammar as run_parity._RUN_RE — the `_n\d+_smoke` anchor makes the +# baseline token EXACT, so "fwdllm" does not swallow "fwdllm_plus". +_RUN_RE = re.compile( + r"^run_(?P\d{8}_\d{6})_(?P.+)_n(?P\d+)_smoke" + r"(?:_(?P.+))?_(?Preal|sim)$" +) + +# Coarse trainer_round phase fields, in pipeline order. +_PHASES = ("mqtt_fetch_s", "weights_to_ram_s", "weights_to_gpu_s", + "pre_train_s", "gpu_compute_s", "post_train_s") + + +def _iter_trainer_events(run_dir: str): + """Yield (short_id, event_dict) for every trainer telemetry record.""" + tdir = pathlib.Path(run_dir) / "telemetry" + for f in sorted(tdir.glob("trainer_*.jsonl")): + sid = f.stem[-4:] + with open(f) as fp: + for line in fp: + line = line.strip() + if not line: + continue + try: + yield sid, json.loads(line) + except json.JSONDecodeError: + continue + + +def _collect(run_dir: str) -> dict: + """Aggregate step_timing (fine) + trainer_round phases (coarse) for a run.""" + step_tot: dict = collections.defaultdict(float) + step_n: dict = collections.defaultdict(int) + phase_tot: dict = collections.defaultdict(float) + phase_n: dict = collections.defaultdict(int) + n_rounds = 0 + for _sid, e in _iter_trainer_events(run_dir): + ev = e.get("event") + if ev == "step_timing": + fn, d = e.get("func"), e.get("duration_s") + if fn is not None and d is not None: + step_tot[fn] += float(d) + step_n[fn] += 1 + elif ev == "trainer_round": + n_rounds += 1 + for p in _PHASES: + v = e.get(p) + if v is not None: + phase_tot[p] += float(v) + phase_n[p] += 1 + return {"step_tot": dict(step_tot), "step_n": dict(step_n), + "phase_tot": dict(phase_tot), "phase_n": dict(phase_n), + "n_rounds": n_rounds} + + +def _fmt_table(label: str, agg: dict) -> str: + lines = [f"\n=== {label} ({agg['n_rounds']} trainer rounds) ==="] + # coarse phases (always present) + lines.append(" COARSE trainer_round phases (mean s/round | total s):") + ptot = agg["phase_tot"] + for p in _PHASES: + if p in ptot: + mean = ptot[p] / max(agg["phase_n"].get(p, 1), 1) + lines.append(f" {p:<18} {mean:8.3f} | {ptot[p]:9.1f}") + # fine steps (present only on runs with step_timing telemetry) + st = agg["step_tot"] + if st: + lines.append(" FINE step_timing (mean s/call | calls | total s), " + "by total desc:") + for fn in sorted(st, key=st.get, reverse=True): + n = agg["step_n"].get(fn, 0) + mean = st[fn] / max(n, 1) + lines.append(f" {fn:<30} {mean:7.3f} | {n:6d} | {st[fn]:8.1f}") + else: + lines.append(" FINE step_timing: NONE — run predates step_timing " + "telemetry (re-run to populate the per-step GPU breakdown).") + return "\n".join(lines) + + +def _plot(label: str, agg: dict, out_path: str) -> bool: + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except Exception as exc: # pragma: no cover + print(f" [plot skipped] matplotlib unavailable: {exc}") + return False + + st = agg["step_tot"] + ptot = agg["phase_tot"] + fig, axes = plt.subplots(1, 2, figsize=(13, 5)) + + # left: coarse phase mean s/round + ph = [(p, ptot[p] / max(agg["phase_n"].get(p, 1), 1)) + for p in _PHASES if p in ptot] + if ph: + names, vals = zip(*ph) + axes[0].barh(range(len(names)), vals, color="#4C72B0") + axes[0].set_yticks(range(len(names))) + axes[0].set_yticklabels(names) + axes[0].invert_yaxis() + axes[0].set_xlabel("mean seconds / round") + axes[0].set_title(f"{label}: coarse phase (trainer_round)") + + # right: fine step_timing mean s/call, top 12 by total + ax = axes[1] + if st: + top = sorted(st, key=st.get, reverse=True)[:12] + means = [st[fn] / max(agg["step_n"].get(fn, 1), 1) for fn in top] + ax.barh(range(len(top)), means, color="#C44E52") + ax.set_yticks(range(len(top))) + ax.set_yticklabels(top, fontsize=8) + ax.invert_yaxis() + ax.set_xlabel("mean seconds / call") + ax.set_title(f"{label}: fine step (step_timing)") + else: + ax.text(0.5, 0.5, "no step_timing telemetry\n(re-run to populate)", + ha="center", va="center", transform=ax.transAxes) + ax.set_axis_off() + + fig.tight_layout() + fig.savefig(out_path) + plt.close(fig) + return True + + +def _discover_latest(experiments_dir: str, baselines, side: str) -> dict: + """{baseline -> newest run_dir} for the requested side, EXACT baseline match + (via _RUN_RE, so fwdllm != fwdllm_plus). ts is fixed-width → lexical == time.""" + want = set(baselines) + best: dict = {} + root = pathlib.Path(experiments_dir) + if not root.is_dir(): + return {} + for d in root.iterdir(): + if not d.is_dir(): + continue + m = _RUN_RE.match(d.name) + if not m or m["variant"] != side or m["baseline"] not in want: + continue + base = m["baseline"] + if base not in best or m["ts"] > best[base][0]: + best[base] = (m["ts"], str(d)) + return {b: p for b, (_ts, p) in best.items()} + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--run-dir", nargs="+", default=None, + help="explicit run dir(s); default = latest per baseline") + ap.add_argument("--experiments-dir", default=str(_DEFAULT_EXPERIMENTS)) + ap.add_argument("--baselines", nargs="+", + default=["fwdllm", "fwdllm_plus", "fluxtune"]) + ap.add_argument("--side", default="sim", choices=["sim", "real"], + help="which side to profile when auto-discovering (default sim)") + ap.add_argument("--out-dir", default=None, + help="where to write PNGs (default /_timing_plots)") + ap.add_argument("--no-plot", action="store_true", help="table only, no PNGs") + args = ap.parse_args(argv) + + if args.run_dir: + runs = {os.path.basename(d.rstrip("/")): d for d in args.run_dir} + else: + runs = _discover_latest(args.experiments_dir, args.baselines, args.side) + if not runs: + print("No run dirs found. Pass --run-dir or check --experiments-dir.") + return 1 + + out_dir = args.out_dir or os.path.join(args.experiments_dir, "_timing_plots") + if not args.no_plot: + os.makedirs(out_dir, exist_ok=True) + + for label, rdir in runs.items(): + agg = _collect(rdir) + print(_fmt_table(label, agg)) + if not args.no_plot: + pdf = os.path.join(out_dir, f"step_timing_{label}.pdf") + if _plot(label, agg, pdf): + print(f" -> {pdf}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/lib/python/examples/fwdllm/expt_scripts/plotlib/__init__.py b/lib/python/examples/fwdllm/expt_scripts/plotlib/__init__.py new file mode 100644 index 000000000..fa524b230 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/plotlib/__init__.py @@ -0,0 +1,11 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""plotlib — shared plotting core for the FWDLLM/FLUXTUNE experiment suite. + + style SOCC-2026 rcParams + PDF/timestamp output convention (one "look") + baselines display registry: names, colors, emphasis (one place to re-brand) + reducers streaming telemetry -> RunResult (one source of the 5 metrics) + figures single-panel paper-figure builders over RunResult lists + +`make_paper_figs.py`, `plot_run.py` and `compare_baselines.py` all build on these. +""" diff --git a/lib/python/examples/fwdllm/expt_scripts/plotlib/baselines.py b/lib/python/examples/fwdllm/expt_scripts/plotlib/baselines.py new file mode 100644 index 000000000..9b11a779d --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/plotlib/baselines.py @@ -0,0 +1,136 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Baseline display registry — the single place paper names, colors, markers and +legend emphasis live (telemetry keys stay internal: fwdllm/fwdllm_plus/fluxtune/…). + +Families carry identity by hue (FwdLLM=orange, our system=blue — the colorblind-safe +pair); members within a family differ by a lighter validated tint + linestyle + +marker (grayscale/CVD-safe, not hue alone). Head-to-head systems are legend-bold, +FwdLLM++ (a derived config) regular-italic. Add an ablation = one BASELINES row; +unknown keys fall back by family prefix. Rename our system via SYSTEM_LABEL. +""" + +from __future__ import annotations + +from dataclasses import dataclass + + +# ---- the switchable codename for OUR system ------------------------------- # +# logs still say "fluxtune"; only the printed/plotted label changes here. +SYSTEM_LABEL = "FluxTune" + + +# family root hue + one validated lighter tint (index 0 = root). Roots+tint pass +# dataviz/validate_palette.js (lightness/chroma/CVD>=12/contrast); a 3rd goes gray. +_FAMILY_TINTS = { + "fwdllm": ["#D55E00", "#E8853A"], # orange + "fluxtune": ["#0072B2", "#4A97CC"], # blue (our system + ablations) +} + +# emphasis -> (fontweight, fontstyle) for the legend text +_EMPHASIS = { + "bold": ("bold", "normal"), + "regular": ("normal", "normal"), + "italic": ("normal", "italic"), +} + + +@dataclass(frozen=True) +class Style: + key: str + label: str + color: str + linestyle: str + marker: str + emphasis: str # key into _EMPHASIS + order: int # legend / z-order (lower plots first, listed first) + + def legend_font(self): + return _EMPHASIS.get(self.emphasis, _EMPHASIS["regular"]) + + +def _tint(family: str, level: int) -> str: + ramp = _FAMILY_TINTS[family] + return ramp[min(level, len(ramp) - 1)] + + +def _mk(key, family, tint, linestyle, emphasis, order, label, marker="o", + color=None) -> Style: + # `color` overrides the family tint ramp — used by the fluxtune ablation, whose + # four variants need a dedicated sequential-blue ramp (light→dark = more opts on), + # not the two categorical family tints. + return Style(key=key, label=label, color=(color or _tint(family, tint)), + linestyle=linestyle, marker=marker, emphasis=emphasis, order=order) + + +# Sequential CVD-safe blue ramp (ColorBrewer "Blues"), light→dark encodes the +# ablation ladder: more learning-affecting opts enabled ⇒ darker. Ordered, not +# categorical — the four fluxtune configs form a 2×2 over Opt-2/Opt-3. +_ABLATION_BLUES = ["#9ecae1", "#4292c6", "#2171b5", "#08306b"] + + +# --------------------------------------------------------------------------- # +# THE registry — one row per baseline/ablation +# --------------------------------------------------------------------------- # +# Grayscale/CVD-safe: orange & blue sit at near-identical grayscale luminance, so +# each baseline also gets a distinct linestyle + marker (identity survives B&W). +BASELINES: dict[str, Style] = { + "fwdllm": _mk("fwdllm", "fwdllm", 0, "-", "bold", 0, "FwdLLM", marker="o"), + "fwdllm_plus": _mk("fwdllm_plus", "fwdllm", 1, "--", "italic", 1, "FwdLLM++", marker="s"), + "fluxtune": _mk("fluxtune", "fluxtune", 0, "-", "bold", 2, SYSTEM_LABEL, marker="^"), + # ablations (uncomment when the runs exist): + # "fluxtune_nojvp": _mk("fluxtune_nojvp", "fluxtune", 1, "--", "regular", 3, + # f"{SYSTEM_LABEL}−JVP", marker="v"), + + # ---- 2×2 opt ablation (N=100, α=1) --------------------------------------- + # Constant across all four: C1 guided JVP perturbations (trainer-side + # select_perturbation_using_jvp=true) + Opt-1 weight-suppression. Varies: + # Opt-2 var-stop, Opt-3 grad-aware. R1 is the FluxTune base (forward-mode LLM + # perturbation fine-tuning; borrows only FeLiX's scalar aggregation rate + # `type=new`); R4 == the full FluxTune default ("fluxtune" above). Blue ramp + + # distinct marker/linestyle so identity survives grayscale/CVD. + "fluxtune_r1_base": _mk("fluxtune_r1_base", "fluxtune", 0, ":", "regular", 10, + f"{SYSTEM_LABEL}-base", marker="o", color=_ABLATION_BLUES[0]), + "fluxtune_r2_varstop": _mk("fluxtune_r2_varstop", "fluxtune", 0, "-.", "regular", 11, + "+ var-stop", marker="s", color=_ABLATION_BLUES[1]), + "fluxtune_r3_gradaware": _mk("fluxtune_r3_gradaware", "fluxtune", 0, "--", "regular", 12, + "+ grad-aware", marker="^", color=_ABLATION_BLUES[2]), + "fluxtune_r4_full": _mk("fluxtune_r4_full", "fluxtune", 0, "-", "bold", 13, + f"{SYSTEM_LABEL} (full)", marker="D", color=_ABLATION_BLUES[3]), +} + +# linestyles + markers cycled for un-registered ablations of the same family +_FALLBACK_LS = ["--", ":", "-."] +_FALLBACK_MARKERS = ["v", "D", "P", "X"] + + +def style_for(key: str) -> Style: + """Resolve a baseline key to a Style, with a safe fallback for unknown keys.""" + if key in BASELINES: + return BASELINES[key] + # fallback: attach to a family by name prefix, lighter tint + cycled linestyle + family = "fluxtune" if key.startswith("fluxtune") else \ + "fwdllm" if key.startswith("fwdllm") else "fluxtune" + n_sibling = sum(1 for s in BASELINES.values() + if s.color in _FAMILY_TINTS[family]) + label = SYSTEM_LABEL + key[len("fluxtune"):].replace("_", " ") \ + if key.startswith("fluxtune") else key + return _mk(key, family, 1, _FALLBACK_LS[n_sibling % len(_FALLBACK_LS)], + "regular", 100 + len(BASELINES), label, + marker=_FALLBACK_MARKERS[n_sibling % len(_FALLBACK_MARKERS)]) + + +def ordered(keys) -> list[str]: + """Keys sorted by registry order (stable legend/plot order across figures).""" + return sorted(keys, key=lambda k: style_for(k).order) + + +def apply_legend_emphasis(legend): + """Set per-entry fontweight/style from the registry (bold vs regular-italic).""" + for txt in legend.get_texts(): + label = txt.get_text() + st = next((s for s in BASELINES.values() if s.label == label), None) + if st is not None: + weight, style = st.legend_font() + txt.set_fontweight(weight) + txt.set_fontstyle(style) diff --git a/lib/python/examples/fwdllm/expt_scripts/plotlib/figures.py b/lib/python/examples/fwdllm/expt_scripts/plotlib/figures.py new file mode 100644 index 000000000..bb3af2bb9 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/plotlib/figures.py @@ -0,0 +1,247 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Paper figure builders — one single-panel figure per experiment. + +Each builder takes an ordered list of RunResult and returns a column-sized Figure; +color/label/linestyle/marker/emphasis come from `baselines.style_for()` (figures +make no independent style decisions). Baselines with no data are skipped, so partial +run-sets just plot what exists. FIG_BUILDERS maps a stable name -> builder; the name +is also the PDF basename, so adding a figure = one entry. +""" + +from __future__ import annotations + +import numpy as np + +from . import baselines as B +from . import style as S + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # +def _new_ax(fraction=1.0, aspect=S.DEFAULT_ASPECT): + import matplotlib.pyplot as plt + fig, ax = plt.subplots(figsize=S.column_figsize(fraction, aspect)) + return fig, ax + + +def _finish(fig, ax, order_keys): + """Legend (registry order + emphasis) + light frame cleanup.""" + handles, labels = ax.get_legend_handles_labels() + if handles: + leg = ax.legend(loc="best") + B.apply_legend_emphasis(leg) + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + return fig + + +def _cdf_xy(values): + xs = np.sort(np.asarray([v for v in values if v is not None], dtype=float)) + if xs.size == 0: + return None, None + ys = np.arange(1, xs.size + 1) / xs.size * 100.0 + return xs, ys + + +def _plot_curve(ax, x, y_raw, st, smooth, label): + """Smoothed line + a faint RAW envelope underneath — the envelope preserves the + true peaks (e.g. the accuracy actually reached) that EMA smoothing pulls down.""" + if smooth and smooth > 0: + ax.plot(x, y_raw, color=st.color, linestyle=st.linestyle, lw=0.7, + alpha=0.22, zorder=2 + st.order) + y = S.ema(y_raw, smooth) + else: + y = y_raw + ax.plot(x, y, color=st.color, linestyle=st.linestyle, marker=st.marker, + markevery=0.1, markersize=4.5, markeredgecolor="white", + markeredgewidth=0.5, label=label, zorder=3 + st.order) + + +def _round_boundaries(ax, rr, hrs): + for i in rr.round_transition_indices(): + ax.axvline(hrs[i], color=B.style_for(rr.key).color, ls=(0, (1, 2)), + lw=0.9, alpha=0.55, zorder=2) + + +def _mark_peak(ax, hrs, acc, st): + """Star at a run's highest accuracy (over the RAW series, so it survives EMA + smoothing) + return the peak value so the caller can print it in the legend — + lets the reader read each baseline's best acc and its gap to the target line.""" + pts = [(h, a) for h, a in zip(hrs, acc) if a is not None] + if not pts: + return None + hp, ap = max(pts, key=lambda p: p[1]) + ax.scatter([hp], [ap], marker="*", s=95, color=st.color, edgecolor="white", + linewidth=0.6, zorder=8 + st.order) + return ap + + +# --------------------------------------------------------------------------- # +# Experiment 1 — accuracy / loss vs wall-clock time (adjacent paper figures) +# --------------------------------------------------------------------------- # +def fig_e1_acc_vs_time(runs, target=None, smooth=0.0, **_): + fig, ax = _new_ax(aspect=0.68) + for rr in runs: + hrs, acc, _loss, _rnd = rr.learning_curve() + if not hrs: + continue + st = B.style_for(rr.key) + peak = _mark_peak(ax, hrs, acc, st) # ★ + peak value in the legend + label = st.label if peak is None else f"{st.label} — peak {peak:.1f}%" + _plot_curve(ax, hrs, acc, st, smooth, label) + _round_boundaries(ax, rr, hrs) # data_id cycles per round → mark epochs + if target is not None: + ax.axhline(target * 100, ls=":", color="#555555", lw=1.0, zorder=2) + ax.text(ax.get_xlim()[1], target * 100, f" target {target*100:.0f}%", + va="center", ha="left", fontsize=7, color="#555555") + ax.set_xlabel("wall-clock time (h)") + ax.set_ylabel("test accuracy (%)") + ax.plot([], [], color="#888888", ls=(0, (1, 2)), lw=0.9, label="round boundary") + ax.scatter([], [], marker="*", s=95, color="#555555", edgecolor="white", + linewidth=0.6, label="peak accuracy") + return _finish(fig, ax, [r.key for r in runs]) + + +def fig_e1_loss_vs_time(runs, smooth=0.0, **_): + """Test-loss vs wall time — the grounded companion to the accuracy curve.""" + fig, ax = _new_ax(aspect=0.68) + any_data = False + for rr in runs: + hrs, _acc, loss, _rnd = rr.learning_curve() + if not hrs or all(v is None for v in loss): + continue + _plot_curve(ax, hrs, loss, B.style_for(rr.key), smooth, B.style_for(rr.key).label) + _round_boundaries(ax, rr, hrs) + any_data = True + if not any_data: + return None + ax.set_xlabel("wall-clock time (h)") + ax.set_ylabel("test loss") + ax.plot([], [], color="#888888", ls=(0, (1, 2)), lw=0.9, label="round boundary") + return _finish(fig, ax, [r.key for r in runs]) + + +# --------------------------------------------------------------------------- # +# Experiment 2 — trainer busy-fraction CDF across clients +# --------------------------------------------------------------------------- # +def fig_e2_trainer_busy_cdf(runs, **_): + fig, ax = _new_ax() + any_data = False + for rr in runs: + st = B.style_for(rr.key) + xs, ys = _cdf_xy([b * 100 for b in rr.busy_frac]) + if xs is None: + continue + ax.plot(xs, ys, color=st.color, linestyle=st.linestyle, + marker=st.marker, markevery=0.12, markersize=4.5, + markeredgecolor="white", markeredgewidth=0.5, + label=st.label, zorder=3 + st.order) + any_data = True + if not any_data: + return None + ax.set_xlabel("trainer busy fraction (%)") + ax.set_ylabel("CDF of clients (%)") + ax.set_ylim(0, 100) + return _finish(fig, ax, [r.key for r in runs]) + + +# --------------------------------------------------------------------------- # +# Experiment 3 — learning per unit compute (ONE figure per denominator: the two +# metrics differ ~100x in scale, so a shared axis makes one set of bars invisible) +# --------------------------------------------------------------------------- # +def _bar_by_baseline(runs, valfn, ylabel, fmt="{:.3g}"): + fig, ax = _new_ax(aspect=0.72) + keys = [r.key for r in runs] + if not keys: + return None + xs = np.arange(len(keys)) + vals = [(valfn(rr) or 0.0) for rr in runs] + ax.bar(xs, vals, 0.62, color=[B.style_for(k).color for k in keys], zorder=3) + ax.axhline(0, color="#888888", lw=0.6) + ax.set_xticks(xs) + ax.set_xticklabels([B.style_for(k).label for k in keys]) + for t, k in zip(ax.get_xticklabels(), keys): # bold/italic per registry + w, s = B.style_for(k).legend_font() + t.set_fontweight(w); t.set_fontstyle(s) + for xi, v in zip(xs, vals): + ax.text(xi, v, fmt.format(v), ha="center", + va="bottom" if v >= 0 else "top", fontsize=7) + ax.set_ylabel(ylabel) + ax.spines["top"].set_visible(False) + ax.spines["right"].set_visible(False) + return fig + + +def fig_e3_dloss_per_gpu_hour(runs, **_): + return _bar_by_baseline( + runs, lambda rr: (rr.delta_loss() / (rr.gpu_s_total() / 3600.0)) + if (rr.delta_loss() and rr.gpu_s_total()) else 0.0, "Δloss per GPU-hour") + + +def fig_e3_dloss_per_mfwd(runs, **_): + return _bar_by_baseline( + runs, lambda rr: (rr.delta_loss() / (rr.fwd_total / 1e6)) + if (rr.delta_loss() and rr.have_fwd and rr.fwd_total) else 0.0, + "Δloss per M forward-pass") + + +# --------------------------------------------------------------------------- # +# Experiment 4 — data transmitted (grouped bar, up vs down) +# --------------------------------------------------------------------------- # +def fig_e4_network_bytes(runs, **_): + fig, ax = _new_ax(aspect=0.7) + runs = [r for r in runs if r.have_comm or r.up_sizes or r.down_sizes] + if not runs: + return None + labels = ["agg→trainer\n(down)", "trainer→agg\n(up)"] + x = np.arange(len(labels)) + n = len(runs) + width = 0.8 / n + for j, rr in enumerate(runs): + st = B.style_for(rr.key) + down = sum(rr.down_sizes) / 1e9 + up = sum(rr.up_sizes) / 1e9 + ax.bar(x + (j - (n - 1) / 2) * width, [down, up], width * 0.92, + color=st.color, label=st.label, zorder=3) + ax.set_xticks(x) + ax.set_xticklabels(labels) + ax.set_ylabel("data transmitted (GB)") + return _finish(fig, ax, [r.key for r in runs]) + + +# --------------------------------------------------------------------------- # +# Experiment 5 — client session-duration CDF (the E5 money plot) +# --------------------------------------------------------------------------- # +def fig_e5_session_cdf(runs, **_): + fig, ax = _new_ax() + any_data = False + for rr in runs: + st = B.style_for(rr.key) + xs, ys = _cdf_xy(rr.session_durs) + if xs is None: + continue + ax.plot(xs, ys, color=st.color, linestyle=st.linestyle, + marker=st.marker, markevery=0.12, markersize=4.5, + markeredgecolor="white", markeredgewidth=0.5, + label=st.label, zorder=3 + st.order) + any_data = True + if not any_data: + return None + ax.set_xlabel("client session duration (s)") + ax.set_ylabel("CDF of sessions (%)") + ax.set_ylim(0, 100) + return _finish(fig, ax, [r.key for r in runs]) + + +# stable name -> builder (name is also the PDF basename) +FIG_BUILDERS = { + "e1_acc_vs_time": fig_e1_acc_vs_time, + "e1_loss_vs_time": fig_e1_loss_vs_time, + "e2_trainer_busy_cdf": fig_e2_trainer_busy_cdf, + "e3_dloss_per_gpu_hour": fig_e3_dloss_per_gpu_hour, + "e3_dloss_per_mfwd": fig_e3_dloss_per_mfwd, + "e4_network_bytes": fig_e4_network_bytes, + "e5_session_cdf": fig_e5_session_cdf, +} diff --git a/lib/python/examples/fwdllm/expt_scripts/plotlib/reducers.py b/lib/python/examples/fwdllm/expt_scripts/plotlib/reducers.py new file mode 100644 index 000000000..8e69096d4 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/plotlib/reducers.py @@ -0,0 +1,337 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Unified telemetry reducer — the single source of the five experiments' metrics, +shared by plot_run / compare_baselines / make_paper_figs via load_run(). + +Streams the aggregator JSONL in one pass (can be multi-GB) and keeps evals as an +ORDERED series, not a data_id-keyed dict: data_id cycles per round (round 2 reuses +data_id 0), so keying by it silently overwrites earlier rounds and corrupts Δloss. +Field names per EXPERIMENTS.md §5. +""" + +from __future__ import annotations + +import glob +import json +import os +from dataclasses import dataclass, field + + +# --------------------------------------------------------------------------- # +# streaming JSONL iterator (bounded memory) +# --------------------------------------------------------------------------- # +def _iter_jsonl(path): + try: + with open(path, "r", encoding="utf-8") as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + yield json.loads(line) + except json.JSONDecodeError: + continue # partial trailing line + except OSError: + return + + +# --------------------------------------------------------------------------- # +# the per-run result — scalars + BOUNDED plot series (never the raw JSONL) +# --------------------------------------------------------------------------- # +@dataclass +class RunResult: + key: str + run_dir: str + # E1/E3 — ordered evals, emission order preserved. Each: dict(ts,round,data_id, + # iter,acc,loss). Bounded by #bins * #rounds (~hundreds), safe to keep. + evals: list = field(default_factory=list) + t0: float | None = None + # E2 — aggregator wall decomposition + agg_compute_s: float = 0.0 + agg_barrier_s: float = 0.0 + agg_drain_s: float = 0.0 + agg_wall_s: float = 0.0 + # E2/E3 — trainer aggregates + busy_frac: list = field(default_factory=list) + trainer_gpu_s: float = 0.0 + fwd_total: int = 0 + pert_total: int = 0 + have_fwd: bool = False + part_rounds: list = field(default_factory=list) + part_bins: list = field(default_factory=list) + part_iters: list = field(default_factory=list) + n_trainers: int = 0 + # E4 — communication + down_sizes: list = field(default_factory=list) + up_sizes: list = field(default_factory=list) + comm_by_kind: dict = field(default_factory=dict) + have_comm: bool = False + # E5 — sessions + session_durs: list = field(default_factory=list) + # E1 (sim) — (ts, vclock_now) from agg_round, for vclock-at-convergence + vclock_track: list = field(default_factory=list) + # productive-learning cutoff (see load_run): telemetry beyond `cutoff_ts` is + # IGNORED so a non-productive stalled tail can't skew the accumulated metrics. + cutoff_ts: float | None = None # None = no cutoff applied (full telemetry) + plateau_ts: float | None = None # ts of the last significant loss improvement + + # ---- derived views (pure over the fields above) ----------------------- # + def learning_curve(self): + """Ordered (hours, acc%, loss, round) arrays over the eval series.""" + if not self.evals or self.t0 is None: + return [], [], [], [] + hrs = [(e["ts"] - self.t0) / 3600 for e in self.evals] + acc = [e["acc"] * 100 if e["acc"] is not None else None for e in self.evals] + loss = [e["loss"] for e in self.evals] + rnd = [e["round"] for e in self.evals] + return hrs, acc, loss, rnd + + def round_transition_indices(self): + """Indices into the eval series where a new round begins (round increments + or data_id wraps to a lower value) — the epoch boundaries to mark on E1.""" + idx = [] + for i in range(1, len(self.evals)): + prev, cur = self.evals[i - 1], self.evals[i] + r0, r1 = prev.get("round"), cur.get("round") + if r0 is not None and r1 is not None and r1 > r0: + idx.append(i) + elif (cur.get("data_id") is not None and prev.get("data_id") is not None + and cur["data_id"] < prev["data_id"]): + idx.append(i) + return idx + + def max_accuracy(self): + accs = [e["acc"] for e in self.evals if e["acc"] is not None] + return max(accs) if accs else None + + def final_accuracy(self): + accs = [e["acc"] for e in self.evals if e["acc"] is not None] + return accs[-1] if accs else None + + def delta_loss(self): + """first eval loss − last eval loss, in TIME order (not data_id order).""" + losses = [e["loss"] for e in self.evals if e["loss"] is not None] + return (losses[0] - losses[-1]) if len(losses) >= 1 else None + + def target_event(self, target: float, window: int): + """The eval dict at which the trailing `window` consecutive evals (time + order) first stay all >= target. None if never reached.""" + streak = 0 + for e in self.evals: + if e["acc"] is None: + continue + streak = streak + 1 if e["acc"] >= target else 0 + if streak >= window: + return e + return None + + def time_to_target(self, target: float, window: int): + """Wall-seconds to the convergence event (see target_event). None if never.""" + e = self.target_event(target, window) + if e is None or self.t0 is None: + return None + return e["ts"] - self.t0 + + def vclock_at(self, ts): + """Last agg_round vclock_now at or before `ts` (sim runs). None if absent.""" + v = None + for t, vc in self.vclock_track: + if ts is not None and t is not None and t <= ts: + v = vc + return v + + def gpu_s_total(self): + return self.trainer_gpu_s + self.agg_compute_s + + +# --------------------------------------------------------------------------- # +# streaming loaders (cutoff-aware — see load_run) +# --------------------------------------------------------------------------- # +def _read_aggregator(path: str) -> dict: + """One streaming pass → the aggregator's BOUNDED raw structures (each carries a + ts so a post-peak cutoff can be applied AFTER the peak time is known).""" + evals, agg_rounds, comm_down, sessions, vclock = [], [], [], [], [] + t0 = None + for e in _iter_jsonl(path): + ev = e.get("event") + ts = e.get("ts") + if ts is not None: + t0 = ts if t0 is None else min(t0, ts) + if ev == "agg_eval": + if e.get("data_id") is None: + continue + evals.append({ + "ts": ts, "round": e.get("round"), + "data_id": int(e["data_id"]), + "iter": e.get("iteration_per_data_id"), + "acc": (float(e["test-accuracy"]) if e.get("test-accuracy") is not None else None), + "loss": (float(e["test-loss"]) if e.get("test-loss") is not None else None), + }) + elif ev == "agg_round": + agg_rounds.append({ + "ts": ts, + "compute": (e.get("aggregate_fedavg_s") or 0.0) + (e.get("eval_s") or 0.0), + "barrier": e.get("barrier_wait_s") or 0.0, + "drain": e.get("drain_tail_s") or 0.0, + "wall": e.get("wall_elapsed_s"), + }) + if e.get("vclock_now") is not None: + vclock.append((ts, float(e["vclock_now"]))) + for iv in (e.get("contributor_intervals") or []): + d, c = iv.get("dispatch_ts"), iv.get("commit_ts") + if d is not None and c is not None and c >= d: + sessions.append((c, c - d)) # (commit_ts, duration) + elif ev == "comm" and e.get("direction") == "agg_to_trainer": + sz = e.get("size_bytes") + if sz is not None: + comm_down.append((ts, sz, e.get("payload_kind"))) + return {"evals": evals, "agg_rounds": agg_rounds, "comm_down": comm_down, + "sessions": sessions, "vclock": vclock, "t0": t0} + + +def _compute_cutoff(evals, grace_s, loss_rel, mode="peak_acc"): + """(cutoff_ts, anchor_ts): where the productive window ends. + + Two anchors, same "did it actually learn?" gate: + * ``loss_plateau`` — the last eval where running-best test-loss dropped + cumulatively >= ``loss_rel`` since its previous milestone. Loss is the + grounded, un-quantized learning signal. + * ``peak_acc`` (default) — the first eval attaining the run's global-max + accuracy, so a post-peak accuracy drift or a round-boundary divergence is + clipped off (stop each baseline near its highest accuracy). + + The loss scan is ALSO the learned-at-all detector: a run whose loss never + dropped ``loss_rel`` (flat at chance) returns inf (no cutoff → shown in full + so the flat non-learning line stays visible), regardless of mode. + Deterministic, so grace_s defaults to 0; grace_s=None disables cutoff entirely. + """ + if grace_s is None or not evals: + return float("inf"), None + # loss-improvement scan — doubles as the "did it learn?" gate + rmin = milestone = plateau_ts = None + advanced = False + for e in evals: + L, ts = e["loss"], e["ts"] + if L is None or ts is None: + continue + if rmin is None: + rmin = milestone = L + plateau_ts = ts + continue + rmin = min(rmin, L) + if rmin <= milestone * (1 - loss_rel): # cumulative >= loss_rel drop + milestone = rmin + plateau_ts = ts + advanced = True + if plateau_ts is None or not advanced: + return float("inf"), None # never learned → show full run + if mode == "loss_plateau": + return plateau_ts + grace_s, plateau_ts + # peak_acc: first occurrence of the global-max accuracy (strict > keeps the + # earliest peak, so a later noisy re-touch of the same max doesn't extend the tail) + best = peak_ts = None + for e in evals: + a, ts = e["acc"], e["ts"] + if a is None or ts is None: + continue + if best is None or a > best: + best, peak_ts = a, ts + if peak_ts is None: + return plateau_ts + grace_s, plateau_ts # no accuracy series → fall back + return peak_ts + grace_s, peak_ts + + +def _read_trainers(tdir: str, cutoff: float, rr: RunResult): + """Per-trainer aggregates, counting only trainer telemetry at ts <= cutoff.""" + for f in glob.glob(os.path.join(tdir, "trainer_*.jsonl")): + rr.n_trainers += 1 + rounds_seen, bins_seen, n_iters = set(), set(), 0 + gpu = 0.0 + ts_lo = ts_hi = None + fp_max = pt_max = 0 + saw_fp = False + for e in _iter_jsonl(f): + ev = e.get("event") + ts = e.get("ts") + if ts is not None and ts > cutoff: + continue # past the cutoff — ignore + if ev == "trainer_round": + gpu += e.get("gpu_compute_s") or 0.0 + n_iters += 1 + if e.get("round") is not None: + rounds_seen.add(e["round"]) + if e.get("data_id") is not None: + bins_seen.add(e["data_id"]) + if e.get("forward_passes_total") is not None: + fp_max = max(fp_max, e["forward_passes_total"]); saw_fp = True + if e.get("perturbations_total") is not None: + pt_max = max(pt_max, e["perturbations_total"]) + if ts is not None: + ts_lo = ts if ts_lo is None else min(ts_lo, ts) + ts_hi = ts if ts_hi is None else max(ts_hi, ts) + elif ev == "comm" and e.get("direction") == "trainer_to_agg": + # trainer_to_agg (upload) is emitted TRAINER-side only. + sz = e.get("size_bytes") + if sz is not None: + rr.up_sizes.append(sz) + rr.have_comm = True + key = ("trainer_to_agg", e.get("payload_kind")) + rr.comm_by_kind[key] = rr.comm_by_kind.get(key, 0) + sz + rr.trainer_gpu_s += gpu + if saw_fp: + rr.fwd_total += fp_max + rr.pert_total += pt_max + rr.have_fwd = True + if n_iters: + rr.part_rounds.append(len(rounds_seen)) + rr.part_bins.append(len(bins_seen)) + rr.part_iters.append(n_iters) + if ts_lo is not None and ts_hi is not None and ts_hi > ts_lo: + rr.busy_frac.append(min(1.0, gpu / (ts_hi - ts_lo))) + + +def load_run(run_dir: str, key: str | None = None, + post_peak_grace_s: float | None = 0.0, + loss_plateau_rel: float = 0.01, + cutoff_mode: str = "peak_acc") -> RunResult | None: + """Stream one run dir's telemetry into a RunResult (None if no aggregator file). + Truncates the productive-learning tail (`cutoff_mode`: 'peak_acc' cuts at the + highest-accuracy eval, 'loss_plateau' at the last significant loss improvement) + so a post-peak drift/divergence can't skew the metrics; `post_peak_grace_s`=None + disables it. The learned-at-all gate is loss-based either way. See _compute_cutoff.""" + run_dir = os.path.abspath(run_dir) + tdir = os.path.join(run_dir, "telemetry") + agg_files = sorted(glob.glob(os.path.join(tdir, "aggregator_*.jsonl"))) + if not agg_files: + return None + raw = _read_aggregator(agg_files[0]) + cutoff, plateau_ts = _compute_cutoff(raw["evals"], post_peak_grace_s, + loss_plateau_rel, mode=cutoff_mode) + + rr = RunResult(key=key or os.path.basename(run_dir), run_dir=run_dir) + rr.t0 = raw["t0"] + rr.cutoff_ts = None if cutoff == float("inf") else cutoff + rr.plateau_ts = plateau_ts + + within = lambda t: t is not None and t <= cutoff + rr.evals = [e for e in raw["evals"] if within(e["ts"])] + used_rounds = [r for r in raw["agg_rounds"] if within(r["ts"])] + rr.agg_compute_s = sum(r["compute"] for r in used_rounds) + rr.agg_barrier_s = sum(r["barrier"] for r in used_rounds) + rr.agg_drain_s = sum(r["drain"] for r in used_rounds) + walls = [r["wall"] for r in used_rounds if r["wall"] is not None] + if walls: + rr.agg_wall_s = max(walls) + elif rr.evals and rr.t0 is not None: + rr.agg_wall_s = max(0.0, rr.evals[-1]["ts"] - rr.t0) + rr.vclock_track = [(t, v) for (t, v) in raw["vclock"] if within(t)] + for (t, sz, kind) in raw["comm_down"]: + if within(t): + rr.down_sizes.append(sz) + key_k = ("agg_to_trainer", kind) + rr.comm_by_kind[key_k] = rr.comm_by_kind.get(key_k, 0) + sz + rr.have_comm = bool(rr.down_sizes) + rr.session_durs = [d for (c, d) in raw["sessions"] if within(c)] + + _read_trainers(tdir, cutoff, rr) + return rr diff --git a/lib/python/examples/fwdllm/expt_scripts/plotlib/style.py b/lib/python/examples/fwdllm/expt_scripts/plotlib/style.py new file mode 100644 index 000000000..2ded37d83 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/plotlib/style.py @@ -0,0 +1,127 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""SOCC-2026 figure style — the single source of every "look" decision: rcParams, +column geometry, PDF export, EMA smoothing and the timestamped-output convention. + +`pdf.fonttype=42` embeds TrueType glyphs (camera-ready checkers reject Type-3). +`timestamped_outdir()` writes to `//` with stable basenames + a `latest` +symlink, so renders never clobber each other yet copying `latest/` into Overleaf +overwrites the draft's figures by name. +""" + +from __future__ import annotations + +import os +from datetime import datetime + + +# single-column width for a 2-column ACM/IEEE-style layout (inches) +COL_WIDTH_IN = 3.35 +# default aspect (height / width) for a line/CDF panel +DEFAULT_ASPECT = 0.66 + + +# --------------------------------------------------------------------------- # +# rcParams — applied once via use_paper_style() +# --------------------------------------------------------------------------- # +_RC = { + # fonts: embed real glyphs; serif body to match a LaTeX column + "pdf.fonttype": 42, + "ps.fonttype": 42, + "font.family": "serif", + # DejaVu Serif is always present; Times/STIX used if the system has them + "font.serif": ["Times New Roman", "Times", "STIXGeneral", "DejaVu Serif"], + "mathtext.fontset": "stix", + "font.size": 9, + "axes.titlesize": 9, + "axes.labelsize": 9, + "xtick.labelsize": 8, + "ytick.labelsize": 8, + "legend.fontsize": 8, + # thin, recessive frame + grid (dataviz: recessive grid/axes) + "axes.linewidth": 0.8, + "axes.grid": True, + "grid.alpha": 0.3, + "grid.linewidth": 0.5, + "axes.axisbelow": True, + "xtick.direction": "out", + "ytick.direction": "out", + "lines.linewidth": 1.8, + "lines.markersize": 5, + "legend.frameon": False, + "legend.handlelength": 1.8, + "legend.columnspacing": 1.2, + "legend.labelspacing": 0.35, + "figure.dpi": 150, + "savefig.dpi": 300, + "savefig.bbox": "tight", + "savefig.pad_inches": 0.02, +} + + +def use_paper_style(): + """Idempotently apply the SOCC-2026 rcParams to the active matplotlib.""" + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + plt.rcParams.update(_RC) + + +def column_figsize(fraction: float = 1.0, aspect: float = DEFAULT_ASPECT): + """(w, h) in inches for `fraction` of a column width at the given aspect.""" + w = COL_WIDTH_IN * fraction + return (w, w * aspect) + + +# --------------------------------------------------------------------------- # +# output convention — timestamped dir + stable basenames + `latest` symlink +# --------------------------------------------------------------------------- # +def timestamped_outdir(root: str, stamp: str | None = None) -> str: + """Create `//`, refresh a `/latest` symlink, return the dir. + + `stamp` defaults to now (YYYYMMDD_HHMMSS). Basenames written inside are stable, + so old renders are preserved while `latest/` always holds the newest set to + copy into the paper. + """ + stamp = stamp or datetime.now().strftime("%Y%m%d_%H%M%S") + out = os.path.join(os.path.abspath(root), stamp) + os.makedirs(out, exist_ok=True) + link = os.path.join(os.path.abspath(root), "latest") + try: + if os.path.islink(link) or os.path.exists(link): + os.remove(link) + os.symlink(stamp, link) # relative target -> portable if root moves + except OSError: + pass # symlink is a convenience, never fatal + return out + + +# --------------------------------------------------------------------------- # +# line smoothing — VISUAL ONLY (never feeds a scalar metric) +# --------------------------------------------------------------------------- # +def ema(values, factor: float): + """TensorBoard-style exponential-moving-average smoothing of a noisy curve. + + `factor` ∈ [0, 1): 0 = no smoothing (returns input), higher = smoother. None + entries pass through untouched. Applied ONLY to plotted lines — the underlying + accuracy/loss scalars (max-acc, Δloss) always use the raw series. + """ + if not factor or factor <= 0: + return list(values) + out, last = [], None + for v in values: + if v is None: + out.append(None) + continue + last = v if last is None else last * factor + v * (1 - factor) + out.append(last) + return out + + +def save_pdf(fig, out_dir: str, name: str): + """Write `/.pdf` — the only artifact. Figures are vector PDF + (resolution-independent); `savefig.dpi=300` covers any rasterized element.""" + fig.tight_layout(pad=0.3) + fig.savefig(os.path.join(out_dir, f"{name}.pdf")) + import matplotlib.pyplot as plt + plt.close(fig) diff --git a/lib/python/examples/fwdllm/expt_scripts/run_parity.py b/lib/python/examples/fwdllm/expt_scripts/run_parity.py new file mode 100644 index 000000000..cc4c45b70 --- /dev/null +++ b/lib/python/examples/fwdllm/expt_scripts/run_parity.py @@ -0,0 +1,229 @@ +#!/usr/bin/env python3 +"""Repeatable real<->sim parity checker for the fwdllm baselines. + +Discovers the most-recent real/sim run pair per baseline, shows them for +confirmation, runs the shared parity battery, and prints one compact cross- +baseline summary. + +Removes two footguns vs calling scripts.parity.cli by hand: + - baseline is parsed as an exact token from the run-dir name + (`run___n_smoke[_]_`), so the `*fwdllm*` + glob never captures `fwdllm_plus`. + - always takes the latest pair and reads agg_goal from each run's own config, + so no hand-passed --agg-goal drifts. + +Usage: + python run_parity.py # all 3 baselines, latest pairs, confirm + python run_parity.py --baselines fluxtune # one baseline + python run_parity.py --yes # skip the confirm prompt + python run_parity.py --validate # + live-run checks (staleness/vclock_now) +""" +from __future__ import annotations + +import argparse +import glob +import json +import os +import re +import sys +from pathlib import Path + +_HERE = Path(__file__).resolve().parent +_LIB_PYTHON = _HERE.parents[2] # lib/python +_PARITY_SCRIPTS = _LIB_PYTHON / "examples" / "async_cifar10" / "scripts" +_DEFAULT_EXPERIMENTS = _HERE.parent / "experiments" # examples/fwdllm/experiments + +for _p in (str(_LIB_PYTHON), str(_PARITY_SCRIPTS)): + if _p not in sys.path: + sys.path.insert(0, _p) + +# `run____n_smoke[_]_` +_RUN_RE = re.compile( + r"^run_(?P\d{8}_\d{6})_(?P.+)_n(?P\d+)_smoke" + r"(?:_(?P.+))?_(?Preal|sim)$" +) + + +def _discover(experiments_dir: str) -> dict: + """{(baseline, trace): {'real': (ts, path), 'sim': (ts, path)}} keeping the + latest ts per (baseline, trace, variant).""" + out: dict = {} + for path in glob.glob(os.path.join(experiments_dir, "run_*")): + m = _RUN_RE.match(os.path.basename(path)) + if not m: + continue + key = (m["baseline"], m["trace"] or "") + slot = out.setdefault(key, {}) + prev = slot.get(m["variant"]) + if prev is None or m["ts"] > prev[0]: # ts is fixed-width -> lexical == chronological + slot[m["variant"]] = (m["ts"], path) + return out + + +def _agg_goal(run_dir: str) -> int | None: + cfg = os.path.join(run_dir, "aggregator_config.json") + if not os.path.exists(cfg): + return None + try: + h = json.load(open(cfg)).get("hyperparameters", {}) + g = h.get("aggGoal", h.get("agg_goal")) + return int(g) if g is not None else None + except (ValueError, OSError, TypeError): + return None + + +# rungs shown in the summary (name -> short label), in ladder-ish order +_HEADLINE = [ + ("cohort_sequence", "cohort"), ("vclock_telemetry", "vclock"), + ("throughput", "thru"), ("total_commits", "commits"), + ("terminal_state", "terminal"), ("r1_inflight_overlap", "R1"), + ("v1_iter_per_data_id", "V1"), ("v2_var_trajectory", "V2"), + ("staleness", "U3"), ("participation", "S2"), + ("convergence", "conv"), ("convergence_loss", "conv_loss"), +] + + +def _status(res: dict) -> str: + if not isinstance(res, dict): + return "–" + if res.get("status") == "SKIP": + return "–" + if res.get("ok"): + return "✓" + return "~" if res.get("tier") == "DIAG" else "✗" + + +def _live_checks(baseline: str, real_dir: str, sim_dir: str) -> list[str]: + """--validate extras: confirm the config/telemetry a fresh run must carry.""" + notes = [] + for tag, d in (("real", real_dir), ("sim", sim_dir)): + logs = glob.glob(os.path.join(d, "*aggregator*.log")) + pol = "?" + if logs: + for line in open(logs[0], errors="replace"): + if "staleness_policy =" in line: + pol = line.split("staleness_policy =")[1].strip() + notes.append(f"{tag} staleness_policy={pol}") + # vclock_now present in sim agg_round telemetry? + tel = glob.glob(os.path.join(sim_dir, "telemetry", "aggregator_*.jsonl")) + n_vclock = 0 + if tel: + for line in open(tel[0], errors="replace"): + if '"event": "agg_round"' in line or '"event":"agg_round"' in line: + e = json.loads(line) + if e.get("vclock_now") is not None: + n_vclock += 1 + notes.append(f"sim vclock_now agg_rounds={n_vclock}") + return notes + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--experiments-dir", default=str(_DEFAULT_EXPERIMENTS)) + ap.add_argument("--baselines", nargs="+", + default=["fwdllm", "fwdllm_plus", "fluxtune"]) + ap.add_argument("--json-dir", default=None, + help="where to write per-pair JSON (default /_parity_reports)") + ap.add_argument("--yes", action="store_true", help="skip the confirm prompt") + ap.add_argument("--validate", action="store_true", + help="also report staleness_policy + vclock_now presence") + ap.add_argument("--max-bin", type=int, default=None, + help="restrict fwdllm cadence rungs (cohort_sequence/V*) to " + "cycle_data_id <= MAX_BIN (first-data-bin logical parity)") + args = ap.parse_args(argv) + + from parity.checks import load_run_dir, run_all_parity # noqa: E402 + try: + from parity.checks import _WARN_ONLY_CHECKS # noqa: E402 + except ImportError: + _WARN_ONLY_CHECKS = set() + + found = _discover(args.experiments_dir) + # Build the ordered work list: one pair per requested (baseline, trace). + pairs = [] + for base in args.baselines: + matches = sorted(k for k in found if k[0] == base) + if not matches: + print(f" [SKIP] {base}: no run dirs found in {args.experiments_dir}") + continue + for key in matches: + slot = found[key] + if "real" not in slot or "sim" not in slot: + have = ",".join(slot) or "none" + print(f" [SKIP] {'/'.join(k for k in key if k)}: " + f"missing a side (have: {have})") + continue + pairs.append((key, slot["real"], slot["sim"])) + + if not pairs: + print("No complete real/sim pairs to check.") + return 1 + + # ── confirmation: show exactly which dirs will be compared ── + label_w = max(len("/".join(k for k in key if k)) for key, _, _ in pairs) + print("\n Real<->sim pairs to check (latest per baseline):") + for key, (rts, rdir), (sts, sdir) in pairs: + goal_r, goal_s = _agg_goal(rdir), _agg_goal(sdir) + goal = f"agg_goal={goal_r}" + (f"!={goal_s}⚠" if goal_s != goal_r else "") + label = "/".join(k for k in key if k) + print(f" {label.ljust(label_w)} {goal}") + print(f" real {rts} {os.path.basename(rdir)}") + print(f" sim {sts} {os.path.basename(sdir)}") + if not args.yes and sys.stdin.isatty(): + if input("\n Proceed with these pairs? [y/N] ").strip().lower() not in ("y", "yes"): + print(" Aborted.") + return 0 + + json_dir = args.json_dir or os.path.join(args.experiments_dir, "_parity_reports") + os.makedirs(json_dir, exist_ok=True) + + # ── run + collect ── + summary = [] # (label, {name: result}, json_path, tally, live_notes) + for key, (rts, rdir), (sts, sdir) in pairs: + label = "/".join(k for k in key if k) + goal = _agg_goal(rdir) or 0 + real_agg, real_tr = load_run_dir(rdir) + sim_agg, sim_tr = load_run_dir(sdir) + res = run_all_parity(real_agg, sim_agg, real_tr, sim_tr, agg_goal=goal, + max_bin=args.max_bin) + jpath = os.path.join(json_dir, f"parity_{label.replace('/', '_')}_{sts}.json") + json.dump(res, open(jpath, "w"), indent=2, default=str) + n_pass = sum(1 for v in res.values() + if isinstance(v, dict) and v.get("ok") and not v.get("status")) + # a real FAIL = not-ok, not a DIAG warn, not in the warn-only allowlist + fails = [k for k, v in res.items() + if isinstance(v, dict) and v.get("ok") is False + and v.get("tier") != "DIAG" and k not in _WARN_ONLY_CHECKS] + n_skip = sum(1 for v in res.values() + if isinstance(v, dict) and v.get("status") == "SKIP") + tally = (n_pass, len(fails), n_skip) + live = _live_checks(label, rdir, sdir) if args.validate else [] + summary.append((label, res, jpath, tally, fails, live)) + + # ── compact cross-baseline table ── + print("\n" + "=" * 78) + lab_w = max(len(s[0]) for s in summary) + hdr = " " + "baseline".ljust(lab_w) + " " + " ".join(lbl for _, lbl in _HEADLINE) + print(hdr) + for label, res, _jp, _tally, _fails, _live in summary: + cells = [] + for name, lbl in _HEADLINE: + cells.append(_status(res.get(name, {})).center(max(len(lbl), 1))) + print(" " + label.ljust(lab_w) + " " + " ".join(cells)) + print(" legend: ✓ pass ✗ fail ~ warn(diag) – skip") + print("=" * 78) + for label, _res, jpath, tally, fails, live in summary: + p, f, s = tally + print(f" {label}: {p} pass / {f} fail / {s} skip" + + (f" FAILS={fails}" if fails else "")) + for note in live: + print(f" · {note}") + print(f" json: {jpath}") + + any_fail = any(t[1] for _, _, _, t, _, _ in summary) + return 1 if any_fail else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/lib/python/examples/fwdllm/expt_scripts/run_sequential.sh b/lib/python/examples/fwdllm/expt_scripts/run_sequential.sh index 98f5a0465..5b888ad86 100755 --- a/lib/python/examples/fwdllm/expt_scripts/run_sequential.sh +++ b/lib/python/examples/fwdllm/expt_scripts/run_sequential.sh @@ -1,121 +1,90 @@ #!/bin/bash -# Run multiple fwdllm YAMLs (fwdllm, fwdllm_plus, fluxtune) one after -# another, in a single conda env, logging each run separately. +# Drive the fwdllm real<->sim launcher pairs (fwdllm, fwdllm_plus, fluxtune) for +# the parity sign-off runs. Thin driver over the shared harness +# examples/scripts/expt_runner.sh (conda activation, launch+progress loop, +# log-health asserts); this file owns the fwdllm specifics: the +# baseline->(real yaml, sim yaml) map, the knob patching, and the pre-flight spec. # -# Each YAML auto-terminates once data_id reaches a threshold or after a -# wall-time cap, whichever comes first. This script overrides those caps -# per invocation via --max-runtime-s/--max-data-id, generating a patched -# copy of each YAML rather than editing the originals. +# --mode {sim|real|both} pairs each baseline's real+sim runs (default both = the +# parity run); run names carry a _real/_sim tag so scripts.parity.cli globs the +# pair. --delays {on|off} sets enable_training_delays IDENTICALLY on both sides +# (mismatched D would be a false divergence). A pre-flight gate prints the +# hyperparameters in three tiers and refuses infeasible configs; --dry-run shows +# it without launching, --yes skips the confirm, --force overrides a block. # # Usage (from anywhere): -# run_sequential.sh [--max-runtime-s 600] [--max-data-id 10] -# [--num-trainers N] [--num-gpus N] [--c C] [--k K] [--stop-on-fail] -# [--partition-method NAME] [--only name1,name2] +# run_sequential.sh [--mode sim|real|both] [--delays on|off] +# [--max-runtime-s 600] [--max-data-id 10] [--num-trainers N] [--num-gpus N] +# [--c C] [--c-async C] [--k K] [--agg-goal N] [--min-initial-trainers N] +# [--partition-method NAME] [--avail-trace NAME | --avail-traces N1,N2] +# [--only name1,name2] [--stop-on-fail] [--dry-run] [--yes] [--force] +# [--show-all] # -# --max-runtime-s wall-clock cap in seconds for each run (default: 600 = 10 min) -# --max-data-id stop a run once data_id reaches this value (default: 10) -# --num-trainers override trainer.num_trainers (default: each YAML's own, 10) -# --num-gpus override execution.num_gpus (default: each YAML's own, 1). -# Scale this with --num-trainers -- each YAML's default of -# 1 GPU is sized for its own default 10-trainer count. -# --c override selector.kwargs.c + minInitialTrainers + agg_goal -# (agg_goal matches c so no selected trainer goes stranded) -# -- superseded per-field by --agg-goal/--min-initial-trainers -# when those are also passed (see below). -# --c-async override selector.kwargs.c only for the async baseline -# (fluxtune) -- lets sync baselines run concurrency==agg_goal -# via --c/--agg-goal while fluxtune overcommits concurrency -# independent of agg_goal (matches async_oort's design: c -# ends in flight, agg_goal of them counted per round). -# --agg-goal override aggregator.agg_goal directly (fans into -# hyperparameters.aggGoal + selector.kwargs.aggGoal/aggr_num -# per runner.py) independent of --c/--c-async. When --c is -# also given without this, legacy behavior (agg_goal==c) -# still applies. -# --min-initial-trainers override selector.kwargs.minInitialTrainers -# directly, independent of --num-trainers/--c. -# --avail-trace override the availability trace used by ALL baselines: -# trainer.availability.mode (cosmetic/consistency), -# trainer hyperparameters.client_notify.trace (fluxtune's -# real signal), and aggregator -# hyperparameters.trackTrainerAvail.trace (fwdllm_plus's -# real ORACULAR signal). Use e.g. "syn_0" (always -# available) to isolate selection/aggregation bugs from -# trace-driven scarcity/churn. -# --avail-traces comma-separated list of traces, e.g. "syn_0,syn_20" -- -# runs the ENTIRE --only baseline sequence once per trace, -# back to back, in this one invocation/process (for an -# unattended overnight multi-trace comparison; no need to -# babysit and launch the next trace by hand). Takes -# precedence over --avail-trace if both are given. Each -# (baseline, trace) run's name/log/results are -# disambiguated by trace -- see the run-name note below. -# --k override selector.kwargs.k -# --stop-on-fail abort the remaining runs as soon as one exits non-zero -# (default: run all three regardless, report at the end) -# --partition-method override hyperparameters.partition_method on both the -# trainer and aggregator sides (default: each YAML's own, -# "uniform" -- IID, chosen for smoke tests to isolate -# launcher-mechanics validation from data-skew effects). -# Must be one of agnews_partition.h5's own group names, -# e.g. "niid_label_clients=100_alpha=0.1" for the most -# heterogeneous split available in the 100-client group -# (smaller alpha = more skewed/non-IID). -# --only comma-separated subset of baselines to execute, e.g. -# --only fwdllm_plus,fluxtune -# (default: all three -- fwdllm, fwdllm_plus, fluxtune) -# These are plain baseline names, independent of -# --num-trainers -- the "n10" in each source YAML's -# filename is just that file's own default trainer -# count, not part of the run's identity. +# --mode which time_mode variant(s) to run per baseline (default both). +# --delays enable_training_delays for BOTH sides of a pair (default off=D=0). +# --max-runtime-s wall/vclock cap per run (default 600 = 10 min). +# --max-data-id stop a run once data_id reaches this (default 9999 = +# effectively unbounded, so --max-runtime-s governs). Pass a small +# value only for a deliberately data-id-capped run. +# --num-trainers override trainer.num_trainers (default: each YAML's own, 10). +# --num-gpus override execution.num_gpus (default: each YAML's own). +# --c / --c-async / --k / --agg-goal / --min-initial-trainers +# selector/aggregator knobs (see the per-flag notes below). +# --partition-method override hyperparameters.partition_method both sides. +# --var-threshold set the variance-pass gate threshold (hyperparameters.var_threshold) +# on both sides. It VARIES with data heterogeneity, so it's a +# review-every-run knob (shown in tier ①), not a fixed default. +# --max-iter-per-data-id set the force-commit cap (max_iterations_per_data_id) +# on both sides (review-every-run, tier ①). +# --avail-trace / --avail-traces availability trace(s); Phase 1 uses syn_0. +# --only comma-separated baseline subset (default all three). +# --after comma-separated post-launch hooks to run once all launches +# finish: parity (scripts.parity.cli --batch on the real/sim +# pairs), sanity (extract_sanity_checks.py per run dir), plot +# (analyze_run.py over the produced telemetry). e.g. --after parity,sanity +# --stop-on-fail abort remaining runs on first non-zero exit. +# --dry-run show the pre-flight table + checks, generate cfgs, DON'T launch. +# --yes don't prompt to confirm a real (GPU) run. +# --force launch even if a pre-flight check is BLOCKING (error). +# --show-all expand tier ③ (config-baked) + list passing checks. +# +# Per-flag knob notes: +# --c sets selector.kwargs.c (+ minInitialTrainers + agg_goal unless +# --agg-goal/--min-initial-trainers override those per-field). +# --c-async sets selector.kwargs.c only for the async baseline (fluxtune). +# --agg-goal sets aggregator.agg_goal directly (fans into hyperparameters.aggGoal +# + selector aggGoal/aggr_num per runner.py) independent of --c. +# --partition-method must be a group name in agnews_partition.h5 +# (e.g. niid_label_clients=100_alpha=0.1); default "uniform" (IID). set -u -# --- robust conda activation (same pattern as scripts/debug_run.sh) --- -# Env choice: FLAME_CONDA_ENV overrides; otherwise use whatever conda env is -# already active in the launching shell (CONDA_DEFAULT_ENV). No hardcoded -# fallback -- activate an env before calling this script, or set -# FLAME_CONDA_ENV explicitly. -ENVNAME="${FLAME_CONDA_ENV:-${CONDA_DEFAULT_ENV:-}}" -if [ -z "$ENVNAME" ]; then - echo "ERROR: no conda env active in this shell and FLAME_CONDA_ENV not set." >&2 - echo " Activate an env first (conda activate ) or pass FLAME_CONDA_ENV=." >&2 - exit 1 -fi -CB="" -if command -v conda >/dev/null 2>&1; then - CB="$(conda info --base 2>/dev/null)" -elif [ -n "${CONDA_EXE:-}" ]; then - CB="$(dirname "$(dirname "$CONDA_EXE")")" -fi -if [ -z "$CB" ] || [ ! -f "$CB/etc/profile.d/conda.sh" ]; then - for c in "$HOME/miniconda3" "/coc/scratch/${USER%??}/miniconda3" \ - "/coc/scratch/$USER/miniconda3" "$HOME/anaconda3" /opt/conda; do - [ -f "$c/etc/profile.d/conda.sh" ] && CB="$c" && break - done -fi -if [ -z "$CB" ] || [ ! -f "$CB/etc/profile.d/conda.sh" ]; then - echo "ERROR: conda not found. Activate '$ENVNAME' yourself or set CONDA_EXE." >&2; exit 1 -fi -source "$CB/etc/profile.d/conda.sh" -conda activate "$ENVNAME" || { echo "ERROR: 'conda activate $ENVNAME' failed" >&2; exit 1; } -echo "conda: base=$CB env=$ENVNAME python=$(which python)" - # repo paths (portable across nodes/checkouts) SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" EXAMPLE_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" # .../examples/fwdllm REPO_ROOT="$(cd "$EXAMPLE_DIR/../../../.." && pwd)" # flame/ +AC10_DIR="$REPO_ROOT/lib/python/examples/async_cifar10" # hosts scripts.parity.cli + +# shared harness: conda activation, launch+ticker, log asserts, preflight bridge +# shellcheck source=../../scripts/expt_runner.sh +source "$REPO_ROOT/lib/python/examples/scripts/expt_runner.sh" -# Force this checkout's flame package ahead of anything already on -# sys.path (e.g. a stale `pip install -e` editable pointing at a -# different clone) so the code that actually runs matches this repo. -export PYTHONPATH="$REPO_ROOT/lib/python${PYTHONPATH:+:$PYTHONPATH}" +expt_activate_conda # no default env: require an active env / FLAME_CONDA_ENV +expt_pin_pythonpath "$REPO_ROOT" # defaults -MAX_RUNTIME_S=600 # 10 minutes -MAX_DATA_ID=10 +MODE="both" +DELAYS="off" +MAX_RUNTIME_S=600 +# Default high so --max-runtime-s governs, not a silent low data-id cap (cost of a +# high default is zero -- --max-runtime-s still bounds the run). +MAX_DATA_ID=9999 +# "Was this passed on the command line?" companions -- MODE/DELAYS/MAX_RUNTIME_S/ +# MAX_DATA_ID have non-empty defaults, so their value alone can't distinguish +# "passed (override)" from "defaulted". Empty-default knobs don't need this. +MODE_SET=0; DELAYS_SET=0; MAX_RUNTIME_S_SET=0; MAX_DATA_ID_SET=0 STOP_ON_FAIL=0 -NUM_TRAINERS="" # empty = leave each YAML's own value -NUM_GPUS="" # empty = leave each YAML's own value +NUM_TRAINERS="" +NUM_GPUS="" SEL_C="" SEL_C_ASYNC="" SEL_K="" @@ -123,121 +92,172 @@ AGG_GOAL="" MIN_INIT_TRAINERS="" AVAIL_TRACE="" AVAIL_TRACES="" -PARTITION_METHOD="" # empty = leave each YAML's own value ("uniform") -ONLY="" # empty = run all three +PARTITION_METHOD="" +VAR_THRESHOLD="" # variance-pass gate threshold; varies with data heterogeneity -> review every run +MAX_ITER_PER_DATA_ID="" # force-commit cap (max_iterations_per_data_id); review every run +VAR_STOPPING_POLICY="" # Opt-2 stopping policy: off|fixed_cap|plateau. Empty => baselines.yaml + # default (fluxtune=plateau); `off` reverts to var<=threshold only. +AGG_RATE_TYPE="" # Opt-3 aggregation rate: grad_aware|new. Empty => baselines.yaml + # default (fluxtune=grad_aware); `new` = the FeLiX scalar rate. +TARGET_ACC="" # convergence stop (EXPERIMENTS.md WS2): terminate when the last + # --converge-window data bins are ALL >= this test accuracy. +CONVERGE_WINDOW="" # W consecutive-bin window for the convergence stop (default 20 when --target-acc set) +STALL_WINDOW_S="" # stall guard: terminate EARLY if best acc hasn't gained --stall-min-delta + # within this many wall s (empty/0 = off unless registry/--run-set sets it). + # Set via --stall-window-s S or the hours alias --stall-window-h H. +STALL_MIN_DELTA="" # accuracy gain that counts as progress (default 0.01 = 1%) +STALL_ON="" # signal that resets the idle clock: acc | loss | either (default either). + # 'either' keeps a run alive if accuracy gains >=stall_min_delta OR test-loss + # drops >=loss_min_rel_delta (empty => converge_watch default: either). +LOSS_MIN_REL_DELTA="" # RELATIVE test-loss drop vs running-best that counts as progress (default 0.01 = 1%) +DELAY_FACTOR="" # training_delay_factor: divides the registry 4-18s delay. Default (trainer_base) + # 10 (=> 0.4-1.8s); pass 1 for the FULL modeled delay. Fans to both roles via runner.py. +RUN_SET="" # load the SHARED condition from experiments.yaml run_sets[NAME] + # (single source of truth for multi-node runs; CLI flags override) +ONLY="" +AFTER="" # comma list of post-launch hooks: parity,sanity,plot (see after_* below) +DRY_RUN=0 +ASSUME_YES=0 +FORCE=0 +SHOW_ALL=0 +CLEAN=0 # --clean: auto-kill stray workers from a prior run (default: abort if dirty) + +usage() { + echo "usage: $0 [--mode sim|real|both] [--delays on|off] [--max-runtime-s S] [--max-data-id N]" >&2 + echo " [--num-trainers N] [--num-gpus N] [--c C] [--c-async C] [--k K] [--agg-goal N]" >&2 + echo " [--min-initial-trainers N] [--partition-method NAME]" >&2 + echo " [--var-threshold F] [--max-iter-per-data-id N] [--delay-factor F]" >&2 + echo " [--target-acc A] [--converge-window W] [--stall-window-s S | --stall-window-h H] [--stall-min-delta D]" >&2 + echo " [--stall-on acc|loss|either] [--loss-min-rel-delta R]" >&2 + echo " [--run-set NAME] [--avail-trace NAME | --avail-traces N1,N2] [--only n1,n2] [--stop-on-fail]" >&2 + echo " [--dry-run] [--yes] [--force] [--show-all] [--clean]" >&2 + echo " --clean auto-kill stray FL workers from a prior/crashed run before each" >&2 + echo " launch (default: verify clean & ABORT if the node is dirty)." >&2 + exit 2 +} while [[ $# -gt 0 ]]; do case "$1" in - --max-runtime-s) MAX_RUNTIME_S="$2"; shift 2 ;; - --max-data-id) MAX_DATA_ID="$2"; shift 2 ;; - --num-trainers) NUM_TRAINERS="$2"; shift 2 ;; - --num-gpus) NUM_GPUS="$2"; shift 2 ;; - --c) SEL_C="$2"; shift 2 ;; - --c-async) SEL_C_ASYNC="$2"; shift 2 ;; - --k) SEL_K="$2"; shift 2 ;; - --agg-goal) AGG_GOAL="$2"; shift 2 ;; + --mode) MODE="$2"; MODE_SET=1; shift 2 ;; + --delays) DELAYS="$2"; DELAYS_SET=1; shift 2 ;; + --max-runtime-s) MAX_RUNTIME_S="$2"; MAX_RUNTIME_S_SET=1; shift 2 ;; + --max-data-id) MAX_DATA_ID="$2"; MAX_DATA_ID_SET=1; shift 2 ;; + --num-trainers) NUM_TRAINERS="$2"; shift 2 ;; + --num-gpus) NUM_GPUS="$2"; shift 2 ;; + --c) SEL_C="$2"; shift 2 ;; + --c-async) SEL_C_ASYNC="$2"; shift 2 ;; + --k) SEL_K="$2"; shift 2 ;; + --agg-goal) AGG_GOAL="$2"; shift 2 ;; --min-initial-trainers) MIN_INIT_TRAINERS="$2"; shift 2 ;; - --avail-trace) AVAIL_TRACE="$2"; shift 2 ;; - --avail-traces) AVAIL_TRACES="$2"; shift 2 ;; - --stop-on-fail) STOP_ON_FAIL=1; shift ;; - --partition-method) PARTITION_METHOD="$2"; shift 2 ;; - --only) ONLY="$2"; shift 2 ;; - *) echo "usage: $0 [--max-runtime-s SECONDS] [--max-data-id N] [--num-trainers N] [--num-gpus N] [--c C] [--c-async C] [--k K] [--agg-goal N] [--min-initial-trainers N] [--avail-trace NAME | --avail-traces NAME1,NAME2,...] [--stop-on-fail] [--partition-method NAME] [--only name1,name2]" >&2; exit 2 ;; + --avail-trace) AVAIL_TRACE="$2"; shift 2 ;; + --avail-traces) AVAIL_TRACES="$2"; shift 2 ;; + --partition-method) PARTITION_METHOD="$2"; shift 2 ;; + --var-threshold) VAR_THRESHOLD="$2"; shift 2 ;; + --max-iter-per-data-id) MAX_ITER_PER_DATA_ID="$2"; shift 2 ;; + --var-stopping-policy) case "$2" in off|fixed_cap|plateau) ;; *) echo "ERROR: --var-stopping-policy must be off|fixed_cap|plateau (got '$2')" >&2; exit 2 ;; esac + VAR_STOPPING_POLICY="$2"; shift 2 ;; + --agg-rate-type) case "$2" in grad_aware|new|old) ;; *) echo "ERROR: --agg-rate-type must be grad_aware|new|old (got '$2')" >&2; exit 2 ;; esac + AGG_RATE_TYPE="$2"; shift 2 ;; + --target-acc) TARGET_ACC="$2"; shift 2 ;; + --converge-window) CONVERGE_WINDOW="$2"; shift 2 ;; + --stall-window-s) STALL_WINDOW_S="$2"; shift 2 ;; + --stall-window-h) # ergonomic hours alias -> seconds (e.g. --stall-window-h 6 => 21600) + case "$2" in ''|*[!0-9.]*|*.*.*) echo "ERROR: --stall-window-h needs a number of hours (got '$2')" >&2; exit 2 ;; esac + STALL_WINDOW_S="$(awk "BEGIN{printf \"%d\", ($2)*3600}")"; shift 2 ;; + --stall-min-delta) STALL_MIN_DELTA="$2"; shift 2 ;; + --stall-on) case "$2" in acc|loss|either) ;; *) echo "ERROR: --stall-on must be acc|loss|either (got '$2')" >&2; exit 2 ;; esac + STALL_ON="$2"; shift 2 ;; + --loss-min-rel-delta) LOSS_MIN_REL_DELTA="$2"; shift 2 ;; + --delay-factor) DELAY_FACTOR="$2"; shift 2 ;; + --run-set) RUN_SET="$2"; shift 2 ;; + --only) ONLY="$2"; shift 2 ;; + --after) AFTER="$2"; shift 2 ;; + --stop-on-fail) STOP_ON_FAIL=1; shift ;; + --dry-run) DRY_RUN=1; shift ;; + --yes) ASSUME_YES=1; shift ;; + --force) FORCE=1; shift ;; + --show-all) SHOW_ALL=1; shift ;; + --clean) CLEAN=1; shift ;; + *) echo "ERROR: unknown arg '$1'" >&2; usage ;; esac done +case "$MODE" in sim|real|both) ;; *) echo "ERROR: --mode must be sim|real|both (got '$MODE')" >&2; exit 2 ;; esac +case "$DELAYS" in on|off) ;; *) echo "ERROR: --delays must be on|off (got '$DELAYS')" >&2; exit 2 ;; esac -# --avail-traces takes precedence; otherwise fall back to the single -# --avail-trace (may be empty, meaning "leave each YAML's own trace"). -if [ -n "$AVAIL_TRACES" ]; then - IFS=',' read -ra TRACE_LIST <<< "$AVAIL_TRACES" -else - TRACE_LIST=("$AVAIL_TRACE") +# --run-set NAME: pull the SHARED condition from experiments.yaml so every node in a +# multi-node run launches the same condition from one source of truth (only --only +# differs per node). Explicit CLI flags WIN; the registry fills only unset knobs. +if [ -n "$RUN_SET" ]; then + REG="$(EXAMPLE_DIR="$EXAMPLE_DIR" RUN_SET="$RUN_SET" python3 - <<'PY' +import os, sys, yaml +p = os.path.join(os.environ["EXAMPLE_DIR"], "experiments.yaml") +try: + reg = yaml.safe_load(open(p, encoding="utf-8")) +except Exception as e: + sys.stderr.write(f"ERROR: cannot read {p}: {e}\n"); sys.exit(3) +rs = (reg.get("run_sets") or {}).get(os.environ["RUN_SET"]) +if not rs: + valid = list((reg.get("run_sets") or {}).keys()) + sys.stderr.write(f"ERROR: run_set '{os.environ['RUN_SET']}' not in {p}. Valid: {valid}\n"); sys.exit(3) +c = rs.get("condition", {}) or {} +d = reg.get("defaults", {}) or {} +C = c.get("C", {}) +def emit(k, v): + if v is not None: print(f"REG_{k}={v}") +emit("N", c.get("N")); emit("K", c.get("K")) +if isinstance(C, dict): + emit("C_SYNC", C.get("sync")); emit("C_ASYNC", C.get("async")) +elif C not in (None, {}): + emit("C_SYNC", C); emit("C_ASYNC", C) +emit("PART", c.get("partition_method")); emit("TRACE", c.get("avail_trace")) +_dl = c.get("delays") +if _dl is not None: + emit("DELAYS", "on" if _dl in (True, "on", "ON", "true", 1) else "off") +emit("DELAY_FACTOR", c.get("delay_factor")) +emit("TARGET_ACC", c.get("target_accuracy")); emit("CONVERGE_WINDOW", c.get("converge_window")) +emit("STALL_WINDOW_S", c.get("stall_window_s", d.get("stall_window_s"))) +emit("STALL_MIN_DELTA", c.get("stall_min_delta", d.get("stall_min_delta"))) +emit("STALL_ON", c.get("stall_on", d.get("stall_on"))) +emit("LOSS_MIN_REL_DELTA", c.get("loss_min_rel_delta", d.get("loss_min_rel_delta"))) +emit("MAX_RUNTIME_S", c.get("max_runtime_s", d.get("max_runtime_s"))) +emit("MAX_DATA_ID", c.get("max_data_id_progress", d.get("max_data_id_progress"))) +PY +)" + rc=$?; if [ "$rc" -ne 0 ]; then echo "$REG" >&2; exit "$rc"; fi + eval "$REG" # defines REG_* shell vars from the registry condition + # empty-default knobs: empty ⇒ operator didn't set ⇒ fill from registry + [ -z "$NUM_TRAINERS" ] && [ -n "${REG_N:-}" ] && NUM_TRAINERS="$REG_N" + [ -z "$SEL_K" ] && [ -n "${REG_K:-}" ] && SEL_K="$REG_K" + [ -z "$SEL_C" ] && [ -n "${REG_C_SYNC:-}" ] && SEL_C="$REG_C_SYNC" + [ -z "$SEL_C_ASYNC" ] && [ -n "${REG_C_ASYNC:-}" ] && SEL_C_ASYNC="$REG_C_ASYNC" + [ -z "$PARTITION_METHOD" ] && [ -n "${REG_PART:-}" ] && PARTITION_METHOD="$REG_PART" + [ -z "$AVAIL_TRACE" ] && [ -z "$AVAIL_TRACES" ] && [ -n "${REG_TRACE:-}" ] && AVAIL_TRACE="$REG_TRACE" + [ -z "$TARGET_ACC" ] && [ -n "${REG_TARGET_ACC:-}" ] && TARGET_ACC="$REG_TARGET_ACC" + [ -z "$CONVERGE_WINDOW" ] && [ -n "${REG_CONVERGE_WINDOW:-}" ] && CONVERGE_WINDOW="$REG_CONVERGE_WINDOW" + [ -z "$DELAY_FACTOR" ] && [ -n "${REG_DELAY_FACTOR:-}" ] && DELAY_FACTOR="$REG_DELAY_FACTOR" + [ -z "$STALL_WINDOW_S" ] && [ -n "${REG_STALL_WINDOW_S:-}" ] && STALL_WINDOW_S="$REG_STALL_WINDOW_S" + [ -z "$STALL_MIN_DELTA" ] && [ -n "${REG_STALL_MIN_DELTA:-}" ] && STALL_MIN_DELTA="$REG_STALL_MIN_DELTA" + [ -z "$STALL_ON" ] && [ -n "${REG_STALL_ON:-}" ] && STALL_ON="$REG_STALL_ON" + [ -z "$LOSS_MIN_REL_DELTA" ] && [ -n "${REG_LOSS_MIN_REL_DELTA:-}" ] && LOSS_MIN_REL_DELTA="$REG_LOSS_MIN_REL_DELTA" + # non-empty-default knobs: apply registry only when the operator didn't pass the flag + if [ "$DELAYS_SET" = "0" ] && [ -n "${REG_DELAYS:-}" ]; then DELAYS="$REG_DELAYS"; fi + if [ "$MAX_RUNTIME_S_SET" = "0" ] && [ -n "${REG_MAX_RUNTIME_S:-}" ]; then MAX_RUNTIME_S="$REG_MAX_RUNTIME_S"; fi + if [ "$MAX_DATA_ID_SET" = "0" ] && [ -n "${REG_MAX_DATA_ID:-}" ]; then MAX_DATA_ID="$REG_MAX_DATA_ID"; fi + echo "run-set '$RUN_SET' loaded from experiments.yaml (explicit CLI flags override registry)." fi -MULTI_TRACE=0 -[ "${#TRACE_LIST[@]}" -gt 1 ] && MULTI_TRACE=1 - -LOGDIR="$SCRIPT_DIR/smoke_logs/$(date '+%Y%m%d_%H%M%S')" -mkdir -p "$LOGDIR" -# Patch hyperparameters.max_runtime_s / max_data_id_progress, and optionally -# num_trainers / selector c+k+minInitialTrainers+agg_goal, in a copy of the -# YAML rather than the original -- keeps the checked-in smoke configs stable -# while letting this script's caller pick the scale per invocation. -patch_yaml() { - python - "$1" "$2" "$3" "$MAX_RUNTIME_S" "$MAX_DATA_ID" "$NUM_TRAINERS" "$NUM_GPUS" "$SEL_C" "$SEL_K" "$PARTITION_METHOD" "$SEL_C_ASYNC" "$AGG_GOAL" "$MIN_INIT_TRAINERS" "$AVAIL_TRACE" <<'PY' -import sys, yaml -(src, dst, run_key, max_runtime_s, max_data_id, num_trainers, num_gpus, sel_c, - sel_k, partition_method, sel_c_async, agg_goal, min_init_trainers, - avail_trace) = sys.argv[1:15] -# Only baseline in ALL_RUNS below that's async; --c-async targets it -# specifically so one invocation can decouple sync concurrency (==agg_goal) -# from async concurrency (overcommitted vs agg_goal) -- see async_oort.py. -IS_ASYNC_BASELINE = run_key == "fluxtune" -cfg = yaml.safe_load(open(src)) -for exp in cfg.get("experiments", []): - h = exp["aggregator"]["config_overrides"]["hyperparameters"] - h["max_runtime_s"] = int(max_runtime_s) - h["max_data_id_progress"] = int(max_data_id) - if partition_method: - h["partition_method"] = partition_method - exp["trainer"]["config_overrides"]["hyperparameters"]["partition_method"] = partition_method - if num_trainers: - exp["trainer"]["num_trainers"] = int(num_trainers) - # exp["name"] feeds the run directory name (run__); derive - # it from run_key + the actual trainer count rather than copying the - # source YAML's own checked-in name, which only reflects that file's - # default count. job.id must track exp["name"] (every checked-in - # YAML keeps them equal; it's the MQTT job/task id shared with - # trainers via runner.py). Include avail_trace when set so runs - # launched back-to-back under different traces (--avail-traces) - # don't produce identically-named run dirs/job ids. - new_name = ( - f"{run_key}_n{num_trainers}_{avail_trace}_smoke" - if avail_trace else f"{run_key}_n{num_trainers}_smoke" - ) - exp["name"] = new_name - exp["aggregator"]["config_overrides"]["job"]["id"] = new_name - if num_gpus: - exp["execution"]["num_gpus"] = int(num_gpus) - kwargs = exp["aggregator"]["config_overrides"]["selector"]["kwargs"] - if sel_c: - kwargs["c"] = int(sel_c) - if not min_init_trainers: - kwargs["minInitialTrainers"] = int(num_trainers) if num_trainers else int(sel_c) - if not agg_goal: - # legacy behavior: agg_goal matches c so no selected trainer goes - # uncounted/stranded. Superseded by --agg-goal below when given. - exp["aggregator"]["agg_goal"] = int(sel_c) - if sel_c_async and IS_ASYNC_BASELINE: - kwargs["c"] = int(sel_c_async) - if sel_k: - kwargs["k"] = int(sel_k) - if agg_goal: - exp["aggregator"]["agg_goal"] = int(agg_goal) - if min_init_trainers: - kwargs["minInitialTrainers"] = int(min_init_trainers) - if avail_trace: - # Cosmetic/consistency: trainer-side self-reported mode. - exp["trainer"].setdefault("availability", {})["mode"] = avail_trace - # Real signal for fluxtune (client_notify) and fwdllm/fwdllm_plus - # (dormant unless trackTrainerAvail below is ORACULAR). - t_hp = exp["trainer"].setdefault("config_overrides", {}).setdefault("hyperparameters", {}) - t_hp.setdefault("client_notify", {})["trace"] = avail_trace - # Real signal for fwdllm_plus (ORACULAR tracking reads this trace - # directly rather than waiting on trainer self-reports). - a_hp = exp["aggregator"]["config_overrides"]["hyperparameters"] - a_hp.setdefault("trackTrainerAvail", {})["trace"] = avail_trace -yaml.safe_dump(cfg, open(dst, "w"), sort_keys=False) -PY -} +# A convergence run (target-acc) is governed by accuracy, not the clock, so bump the +# wall ceiling to 48h when a target is set and --max-runtime-s is still the 600 default. +if [ -n "$TARGET_ACC" ] && [ "$MAX_RUNTIME_S_SET" = "0" ] && [ "$MAX_RUNTIME_S" = "600" ]; then + MAX_RUNTIME_S=172800 # 48h +fi -# Keys are plain baseline names -- independent of --num-trainers and of -# whatever scale is baked into each source YAML's own filename/checked-in -# default. The mapping to the actual YAML file lives only here. +# baseline -> (real yaml : sim yaml). Plain baseline names, independent of the +# "n10" baked into each source filename. ALL_RUNS=( - "fwdllm:$SCRIPT_DIR/fwdllm_n10_smoke.yaml" - "fwdllm_plus:$SCRIPT_DIR/fwdllm_plus_n10_smoke.yaml" - "fluxtune:$SCRIPT_DIR/fluxtune_n10_smoke.yaml" + "fwdllm:$SCRIPT_DIR/fwdllm_n10_smoke.yaml:$SCRIPT_DIR/fwdllm_n10_smoke_sim.yaml" + "fwdllm_plus:$SCRIPT_DIR/fwdllm_plus_n10_smoke.yaml:$SCRIPT_DIR/fwdllm_plus_n10_smoke_sim.yaml" + "fluxtune:$SCRIPT_DIR/fluxtune_n10_smoke.yaml:$SCRIPT_DIR/fluxtune_n10_smoke_sim.yaml" ) if [ -n "$ONLY" ]; then @@ -246,88 +266,634 @@ if [ -n "$ONLY" ]; then for want in "${ONLY_NAMES[@]}"; do found=0 for entry in "${ALL_RUNS[@]}"; do - if [ "${entry%%:*}" = "$want" ]; then - RUNS+=("$entry") - found=1 - break - fi + if [ "${entry%%:*}" = "$want" ]; then RUNS+=("$entry"); found=1; break; fi done if [ "$found" = "0" ]; then - echo "ERROR: --only name '$want' not recognized. Valid names: ${ALL_RUNS[*]%%:*}" >&2 - exit 2 + echo "ERROR: --only name '$want' not recognized. Valid: ${ALL_RUNS[*]%%:*}" >&2; exit 2 fi done else RUNS=("${ALL_RUNS[@]}") fi -declare -A RESULT -declare -A DURATION_S -ORDERED_KEYS=() # (name or name@trace) in the order actually run, for the summary +# trace list: --avail-traces wins; else single --avail-trace (may be empty). +if [ -n "$AVAIL_TRACES" ]; then TRACE_CSV="$AVAIL_TRACES"; else TRACE_CSV="$AVAIL_TRACE"; fi -CHILD_PID="" -cleanup() { - echo "" - echo "Interrupted. Killing child (PID=${CHILD_PID:-none})..." - [ -n "$CHILD_PID" ] && kill -- -"$CHILD_PID" 2>/dev/null - exit 130 +LOGDIR="$SCRIPT_DIR/smoke_logs/$(date '+%Y%m%d_%H%M%S')" +mkdir -p "$LOGDIR" +GPUS_VISIBLE="$( (command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi --query-gpu=index --format=csv,noheader 2>/dev/null | wc -l) || echo 0)" +MANIFEST="$LOGDIR/manifest.tsv" +SPEC_JSON="$LOGDIR/spec.json" + +# ---- PHASE A: generate all patched cfgs, build the tier/check spec, gate ---- +# One python step: the operator sees the whole matrix (baselines x traces x variants) +# and confirms once. Writes per-run cfgs + a launch manifest, renders the tiered +# table via expt_runner.render_and_gate, and exits 2 if any check is blocking. +RUN_TSV="$LOGDIR/_runs.tsv"; : > "$RUN_TSV" +for entry in "${RUNS[@]}"; do + bl="${entry%%:*}"; rest="${entry#*:}"; real_y="${rest%%:*}"; sim_y="${rest#*:}" + printf '%s\t%s\t%s\n' "$bl" "$real_y" "$sim_y" >> "$RUN_TSV" +done + +EXPT_RUNNER_DIR="$EXPT_RUNNER_DIR" \ +MODE="$MODE" DELAYS="$DELAYS" MAX_RUNTIME_S="$MAX_RUNTIME_S" MAX_DATA_ID="$MAX_DATA_ID" \ +NUM_TRAINERS="$NUM_TRAINERS" NUM_GPUS="$NUM_GPUS" SEL_C="$SEL_C" SEL_C_ASYNC="$SEL_C_ASYNC" \ +SEL_K="$SEL_K" AGG_GOAL="$AGG_GOAL" MIN_INIT_TRAINERS="$MIN_INIT_TRAINERS" \ +PARTITION_METHOD="$PARTITION_METHOD" TRACE_CSV="$TRACE_CSV" GPUS_VISIBLE="$GPUS_VISIBLE" \ +VAR_THRESHOLD="$VAR_THRESHOLD" MAX_ITER_PER_DATA_ID="$MAX_ITER_PER_DATA_ID" DELAY_FACTOR="$DELAY_FACTOR" \ +VAR_STOPPING_POLICY="$VAR_STOPPING_POLICY" AGG_RATE_TYPE="$AGG_RATE_TYPE" \ +TARGET_ACC="$TARGET_ACC" CONVERGE_WINDOW="$CONVERGE_WINDOW" \ +STALL_WINDOW_S="$STALL_WINDOW_S" STALL_MIN_DELTA="$STALL_MIN_DELTA" \ +STALL_ON="$STALL_ON" LOSS_MIN_REL_DELTA="$LOSS_MIN_REL_DELTA" \ +MODE_SET="$MODE_SET" DELAYS_SET="$DELAYS_SET" MAX_RUNTIME_S_SET="$MAX_RUNTIME_S_SET" MAX_DATA_ID_SET="$MAX_DATA_ID_SET" \ +LOGDIR="$LOGDIR" MANIFEST="$MANIFEST" RUN_TSV="$RUN_TSV" DRY_RUN="$DRY_RUN" SHOW_ALL="$SHOW_ALL" \ +EXAMPLE_DIR="$EXAMPLE_DIR" AC10_DIR="$AC10_DIR" \ +python - <<'PY' +import os, sys, copy, yaml, json, hashlib +sys.path.insert(0, os.environ["EXPT_RUNNER_DIR"]) +import expt_runner + +env = os.environ.get +MODE = env("MODE"); DELAYS = env("DELAYS") +MAX_RUNTIME_S = int(env("MAX_RUNTIME_S")); MAX_DATA_ID = int(env("MAX_DATA_ID")) +NUM_TRAINERS = env("NUM_TRAINERS") or "" +NUM_GPUS = env("NUM_GPUS") or "" +SEL_C = env("SEL_C") or ""; SEL_C_ASYNC = env("SEL_C_ASYNC") or ""; SEL_K = env("SEL_K") or "" +AGG_GOAL = env("AGG_GOAL") or ""; MIN_INIT = env("MIN_INIT_TRAINERS") or "" +PART = env("PARTITION_METHOD") or "" +VAR_THRESHOLD = env("VAR_THRESHOLD") or ""; MAX_ITER = env("MAX_ITER_PER_DATA_ID") or "" +VAR_STOPPING_POLICY = env("VAR_STOPPING_POLICY") or ""; AGG_RATE_TYPE = env("AGG_RATE_TYPE") or "" +DELAY_FACTOR = env("DELAY_FACTOR") or "" +STALL_ON = env("STALL_ON") or ""; LOSS_MIN_REL_DELTA = env("LOSS_MIN_REL_DELTA") or "" +TARGET_ACC = env("TARGET_ACC") or ""; CONVERGE_WINDOW = env("CONVERGE_WINDOW") or "" +STALL_WINDOW_S = env("STALL_WINDOW_S") or ""; STALL_MIN_DELTA = env("STALL_MIN_DELTA") or "" +# "was it passed on the command line?" (override -> green) for the defaulted flags +MODE_SET = env("MODE_SET") == "1"; DELAYS_SET = env("DELAYS_SET") == "1" +MAX_RUNTIME_S_SET = env("MAX_RUNTIME_S_SET") == "1"; MAX_DATA_ID_SET = env("MAX_DATA_ID_SET") == "1" +GPUS_VISIBLE = int(env("GPUS_VISIBLE") or "0") +LOGDIR = env("LOGDIR"); MANIFEST = env("MANIFEST") +DRY_RUN = env("DRY_RUN") == "1"; SHOW_ALL = env("SHOW_ALL") == "1" +delays_on = (DELAYS == "on") + +# Availability trace(s). Default to syn_0 (Phase-1, 100% availability) when the +# operator passes no --avail-trace, so patch() ALWAYS sets the mode EXPLICITLY on +# every baseline (trainer availability.mode + aggregator trackTrainerAvail + +# client_notify -- lines below) rather than silently inheriting each yaml's own +# `mode:`. This is what makes the printed "trace" row match what actually runs: +# the resolved value is patched into the launched cfg, not just displayed. +_trace_raw = [t for t in (env("TRACE_CSV") or "").replace(",", " ").split()] +trace_set = bool(_trace_raw) # operator passed --avail-trace(s)? +traces = _trace_raw or ["syn_0"] # Phase-1 default: 100% availability +multi_trace = len(traces) > 1 + +variants = {"real": 0, "sim": 1} if MODE == "both" else {MODE: (0 if MODE == "real" else 1)} + +# Baseline-distinguishing internals (selector algorithm / optimizer / sync|async) +# come from the shared catalog _metadata/baselines.yaml, merged at LAUNCH — NOT +# from the per-run YAML the operator edits. Surface them in the review table so a +# mis-picked baseline (e.g. a sync selector where async was intended) is caught +# BEFORE the run, not after. Best-effort: if the catalog can't be read, the +# columns show "?" rather than blocking. +_BL_INTERNALS = {} +try: + _bl_path = os.path.join(env("EXAMPLE_DIR"), "..", "_metadata", "baselines.yaml") + _bl = yaml.safe_load(open(_bl_path, encoding="utf-8")) + _bl = _bl.get("baselines", _bl) + for _name, _b in (_bl or {}).items(): + _agg = (_b or {}).get("aggregator", {}) or {} + _sel = _agg.get("selector", {}) or {} + _opt = _agg.get("optimizer", {}) or {} + _is_async = bool((_sel.get("kwargs", {}) or {}).get("is_async")) + _BL_INTERNALS[_name] = { + "selector": _sel.get("sort", "?"), + "optimizer": _opt.get("sort", "?"), + "async": "async" if _is_async else "sync", + } +except Exception: + _BL_INTERNALS = {} + +runs = [] # (baseline, real_yaml, sim_yaml) +with open(env("RUN_TSV")) as fh: + for line in fh: + line = line.rstrip("\n") + if line: + runs.append(line.split("\t")) + +manifest = [] # (name, cfg_path, variant, budget_s) +per_baseline = {} # baseline -> resolved knobs (for the tier ② rows) +checks = [] + + +def patch(exp, run_key, variant, trace): + h = exp["aggregator"]["config_overrides"]["hyperparameters"] + h["max_runtime_s"] = MAX_RUNTIME_S + h["max_data_id_progress"] = MAX_DATA_ID + # enable_training_delays: SAME on both sides of a pair (K-D8). + exp["trainer"]["enable_training_delays"] = delays_on + # training_delay_factor (simulate_fwdllm.md #12): divides the registry 4-18s + # delay (trainer_base default 10 => 0.4-1.8s). Set it on the TRAINER + # hyperparameters; runner.py fans the same value into the aggregator so both + # roles agree. Only patched when explicitly passed (else the base default). + if DELAY_FACTOR: + exp["trainer"].setdefault("hyperparameters", {}) + exp["trainer"]["hyperparameters"]["training_delay_factor"] = float(DELAY_FACTOR) + if PART: + h["partition_method"] = PART + exp["trainer"]["config_overrides"]["hyperparameters"]["partition_method"] = PART + # Variance-cadence knobs (review-every-run). Only patched when explicitly set, + # so an unset run keeps the code/trainer default (surfaced as "(D)" below). + if VAR_THRESHOLD: + h["var_threshold"] = float(VAR_THRESHOLD) + if MAX_ITER: + h["max_iterations_per_data_id"] = int(MAX_ITER) + # Opt-2/Opt-3 ablation toggles (charter 4-run 2x2). Written into the per-run + # config_overrides, which WIN over the baselines.yaml catalog at launch. A full + # agg_rate_conf dict is written per type so the result is deterministic + # regardless of merge depth (the felix `new` branch needs scale/a_exp/b_exp). + if VAR_STOPPING_POLICY: + h["var_stopping_policy"] = VAR_STOPPING_POLICY + if AGG_RATE_TYPE: + _opt = exp["aggregator"]["config_overrides"].setdefault("optimizer", {}) + _ok = _opt.setdefault("kwargs", {}) + if AGG_RATE_TYPE == "grad_aware": + _ok["agg_rate_conf"] = { + "type": "grad_aware", "base": "new", "align_gate": True, + "align_floor": 0.0, "inverse_var": False, "var_ref": 0.3, + "scale": 0.4, "a_exp": 0.25, "b_exp": 0.1, + } + else: # new (FeLiX) or old + _ok["agg_rate_conf"] = { + "type": AGG_RATE_TYPE, "scale": 0.4, "a_exp": 0.25, "b_exp": 0.1, + } + if NUM_TRAINERS: + exp["trainer"]["num_trainers"] = int(NUM_TRAINERS) + if NUM_GPUS: + exp["execution"]["num_gpus"] = int(NUM_GPUS) + kwargs = exp["aggregator"]["config_overrides"]["selector"]["kwargs"] + is_async = (run_key == "fluxtune") + if SEL_C: + kwargs["c"] = int(SEL_C) + if not MIN_INIT: + kwargs["minInitialTrainers"] = int(NUM_TRAINERS) if NUM_TRAINERS else int(SEL_C) + if not AGG_GOAL: + exp["aggregator"]["agg_goal"] = int(SEL_C) # legacy: agg_goal matches c + if SEL_C_ASYNC and is_async: + kwargs["c"] = int(SEL_C_ASYNC) + if SEL_K: + kwargs["k"] = int(SEL_K) + if AGG_GOAL: + exp["aggregator"]["agg_goal"] = int(AGG_GOAL) + if MIN_INIT: + kwargs["minInitialTrainers"] = int(MIN_INIT) + if trace: + exp["trainer"].setdefault("availability", {})["mode"] = trace + t_hp = exp["trainer"].setdefault("config_overrides", {}).setdefault("hyperparameters", {}) + t_hp.setdefault("client_notify", {})["trace"] = trace + h.setdefault("trackTrainerAvail", {})["trace"] = trace + # name / job id: carry a _real|_sim tag so scripts.parity.cli can glob the pair. + n = int(NUM_TRAINERS) if NUM_TRAINERS else exp["trainer"].get("num_trainers", 10) + parts = [run_key, f"n{n}", "smoke"] + if trace: + parts.append(trace) + parts.append(variant) + name = "_".join(parts) + exp["name"] = name + exp["aggregator"]["config_overrides"]["job"]["id"] = name + return name + + +for trace in traces: + for run_key, real_y, sim_y in runs: + for variant, _idx in variants.items(): + src = real_y if variant == "real" else sim_y + if not os.path.exists(src): + checks.append({"name": f"source yaml exists ({run_key} {variant})", + "level": "error", "detail": f"missing: {src}"}) + continue + cfg = yaml.safe_load(open(src, encoding="utf-8")) + exps = cfg.get("experiments", []) + for exp in exps: + name = patch(exp, run_key, variant, trace) + cfg["experiments"] = exps + out = os.path.join(LOGDIR, f"{name}.yaml") + yaml.safe_dump(cfg, open(out, "w", encoding="utf-8"), sort_keys=False) + manifest.append((name, out, variant, MAX_RUNTIME_S)) + + # record resolved knobs from the (first) patched experiment for display + e0 = exps[0] + h0 = e0["aggregator"]["config_overrides"]["hyperparameters"] + kw0 = e0["aggregator"]["config_overrides"]["selector"]["kwargs"] + per_baseline.setdefault(run_key, { + "c": kw0.get("c"), "k": kw0.get("k"), + "agg_goal": e0["aggregator"].get("agg_goal"), + "min_init": kw0.get("minInitialTrainers"), + "n_trainers": e0["trainer"].get("num_trainers"), + "n_gpus": e0.get("execution", {}).get("num_gpus"), + "partition": h0.get("partition_method"), + "delays": e0["trainer"].get("enable_training_delays"), + # RESOLVED availability mode read back from the PATCHED cfg (what + # actually launches), so the table can't show a stale default. + "avail": e0["trainer"].get("availability", {}).get("mode"), + "async": (run_key == "fluxtune"), + # baseline-distinguishing internals from the shared catalog + "selector": _BL_INTERNALS.get(run_key, {}).get("selector", "?"), + "optimizer": _BL_INTERNALS.get(run_key, {}).get("optimizer", "?"), + "sync_async": _BL_INTERNALS.get(run_key, {}).get("async", "?"), + }) + +with open(MANIFEST, "w") as fh: + for name, out, variant, budget in manifest: + fh.write(f"{name}\t{out}\t{variant}\t{budget}\n") + +# ---------------- build the tiered spec ---------------- +# Row colour convention (legend in the subtitle): +# 🟢 set = value came from a command-line flag -> OVERRIDES the yaml (even if +# the yaml happens to agree). Shown WITHOUT "(D)". +# 🟡 warn = review/attention (a review-every-run knob still on its default). +# (D) + dim = plain yaml/code default, not overridden this run. +def dflt(val, overridden): + # inline "(D)" marker for the bundled per-baseline knob strings (tier ②) + return f"{val}" if overridden else f"{val} (D)" + +def scalar_row(label, val, overridden, note=None, review=False): + if overridden: # operator passed the flag -> override + d = {"label": label, "value": f"{val}", "level": "set"} + elif review: # defaulted but must be eyeballed each run + d = {"label": label, "value": f"{val} (D)", "level": "warn"} + else: # plain default + d = {"label": label, "value": f"{val} (D)", "level": "ok"} + if note: + d["note"] = note + return d + +# --- condition fingerprint (TWO-NODE SAFETY) ---------------------------------- +# A short hash over the SHARED condition axes that MUST match for a valid cross- +# baseline comparison. agg_goal / selector / optimizer legitimately DIFFER per +# baseline, so they are EXCLUDED. Uses RESOLVED partition/trace (what actually +# runs), not just the flags. Print it on every node: if node A and node B show +# the SAME fingerprint, they launched the same condition — the single check that +# catches a mistyped flag on the second node before the runs diverge. +_res_parts = sorted({str(b.get("partition")) for b in per_baseline.values()}) +_res_traces = sorted({str(b.get("avail")) for b in per_baseline.values()}) +_cond = { + "N": NUM_TRAINERS or "yaml", "K": SEL_K or "yaml", + "C_sync": SEL_C or "yaml", "C_async": SEL_C_ASYNC or SEL_C or "yaml", + "partition": _res_parts, "trace": _res_traces, + "delays": "on" if delays_on else "off", + "delay_factor": DELAY_FACTOR or "base", + "target_acc": TARGET_ACC or "none", + "stall_window_s": STALL_WINDOW_S or "off", "stall_min_delta": STALL_MIN_DELTA or "off", + "stall_on": STALL_ON or "either", "loss_min_rel_delta": LOSS_MIN_REL_DELTA or "0.01", + "converge_window": (CONVERGE_WINDOW or "20") if TARGET_ACC else "none", + "max_runtime_s": MAX_RUNTIME_S, "max_data_id": MAX_DATA_ID, } -trap cleanup INT TERM +_cond_fp = hashlib.sha256(json.dumps(_cond, sort_keys=True).encode()).hexdigest()[:8] -cd "$REPO_ROOT" || exit 1 -echo "=== fwdllm sequential run: ${#RUNS[@]} runs (${RUNS[*]%%:*}) x ${#TRACE_LIST[@]} trace(s) (${TRACE_LIST[*]:-}), max_runtime_s=$MAX_RUNTIME_S max_data_id=$MAX_DATA_ID num_trainers=${NUM_TRAINERS:-} num_gpus=${NUM_GPUS:-} c=${SEL_C:-} k=${SEL_K:-} partition_method=${PARTITION_METHOD:-}, logs in $LOGDIR ===" +tiers = [] +# ① review every run +# The trace row reflects the RESOLVED per-baseline availability (read back from +# the patched cfgs), NOT a hardcoded default -- so "what is printed" == "what +# runs". If every baseline resolved to the same mode, show it; otherwise flag +# the divergence and defer to the per-baseline table (tier ②). +trace_overridden = trace_set +_resolved_avails = {b.get("avail") for b in per_baseline.values() if b.get("avail")} +if len(_resolved_avails) == 1: + trace_val = next(iter(_resolved_avails)) +elif _resolved_avails: + trace_val = "MIXED: " + ",".join(sorted(a or "?" for a in _resolved_avails)) + " (see ②)" +else: + trace_val = " ".join(traces) +# mode: single-sided always warns (parity needs both), regardless of override. +if MODE != "both": + mode_row = {"label": "mode", "value": MODE, "level": "warn", "note": "single-sided: parity needs both"} +else: + mode_row = scalar_row("mode", MODE, MODE_SET, note="real+sim pair (--mode)") +tier1 = {"name": "① REVIEW EVERY RUN", "rows": [ + {"label": "condition_fp", "value": _cond_fp, "level": "set", + "note": "TWO-NODE CHECK: same fingerprint on every node ⇒ same shared condition " + "(N/K/C/partition/trace/delays/target_acc/caps). Differs ⇒ a knob was mistyped."}, + mode_row, + {"label": "baselines", "value": " ".join(rk for rk, *_ in runs)}, + # The two similarly-named-but-DIFFERENT knobs, disambiguated + on their own rows: + scalar_row("max_runtime_s", MAX_RUNTIME_S, MAX_RUNTIME_S_SET, note="wall/vclock cap (--max-runtime-s)"), + scalar_row("max_data_id_progress", MAX_DATA_ID, MAX_DATA_ID_SET, + note="STOP condition: stop when data_id reaches this (--max-data-id)"), + scalar_row("trace", trace_val, trace_overridden, + note=("resolved availability actually patched into each launched cfg; " + + ("100% availability (Phase 1)" if _resolved_avails == {"syn_0"} + else "NON-syn_0 — unavailability (Phase 2+)"))), + scalar_row("enable_training_delays", str(delays_on).lower(), DELAYS_SET, + note=(f"modeled training delay {'ON (D>0)' if delays_on else 'OFF (D=0)'}; " + f"delay_factor={DELAY_FACTOR or 'base(10)'} (divides yaml base 4-18s delay); matched BOTH sides — K-D8")), + # var_threshold / max_iterations_per_data_id vary with data heterogeneity -> + # review-every-run (warn when defaulted). NOTE: max_iters_per_data_id is the + # FORCE-COMMIT cap and is NOT the same as max_data_id_progress (the stop) above. + scalar_row("var_threshold", VAR_THRESHOLD if VAR_THRESHOLD else "unset", + bool(VAR_THRESHOLD), review=True, + note="variance-pass gate; varies w/ data heterogeneity (--var-threshold). unset ⇒ trainer/code default"), + scalar_row("max_iters_per_data_id", MAX_ITER if MAX_ITER else "unset", + bool(MAX_ITER), review=True, + note="FORCE-COMMIT cap (--max-iter-per-data-id) — NOT the max_data_id_progress stop above. unset ⇒ code default"), + scalar_row("var_stopping_policy", VAR_STOPPING_POLICY if VAR_STOPPING_POLICY else "default", + bool(VAR_STOPPING_POLICY), review=True, + note="Opt-2 (--var-stopping-policy) off|fixed_cap|plateau. default ⇒ baselines.yaml (fluxtune=plateau)"), + scalar_row("agg_rate_type", AGG_RATE_TYPE if AGG_RATE_TYPE else "default", + bool(AGG_RATE_TYPE), review=True, + note="Opt-3 (--agg-rate-type) grad_aware|new. default ⇒ baselines.yaml (fluxtune=grad_aware); new = FeLiX"), + # Convergence stop (EXPERIMENTS.md WS2): terminate when the last W data bins + # are ALL >= target accuracy. When set, max_runtime_s/max_data_id become + # SAFETY CAPS (a non-converging run -> DID_NOT_CONVERGE). unset ⇒ time/data-id bound only. + scalar_row("target_acc", TARGET_ACC if TARGET_ACC else "unset", + bool(TARGET_ACC), review=True, + note=("convergence stop: last %s bins all >= this (--target-acc). " + "unset ⇒ NO accuracy stop, only max_runtime_s/max_data_id" + % (CONVERGE_WINDOW or "20"))), + # Stall guard: early-terminate a not-learning run before the wall ceiling. + scalar_row("stall_guard", + ((lambda _on, _h: { + "acc": f"<{STALL_MIN_DELTA or '0.01'} acc in {_h}h→STALLED", + "loss": f"<{float(LOSS_MIN_REL_DELTA or '0.01')*100:g}% loss in {_h}h→STALLED", + "either": f"<{STALL_MIN_DELTA or '0.01'} acc AND <{float(LOSS_MIN_REL_DELTA or '0.01')*100:g}% loss in {_h}h→STALLED", + }[_on])(STALL_ON or "either", int(float(STALL_WINDOW_S))//3600) + if STALL_WINDOW_S and float(STALL_WINDOW_S) > 0 else "off"), + bool(STALL_WINDOW_S), review=True, + note=("terminate EARLY if there is no PROGRESS within stall_window_s on the armed signal " + "(--stall-on acc|loss|either): acc gain < stall_min_delta (abs) and/or test-loss " + "drop < loss_min_rel_delta (rel vs running-best). off ⇒ run to convergence or the wall ceiling")), +]} +tiers.append(tier1) + +# ② per-baseline -- an ALIGNED TABLE (baselines = rows, knobs = columns). The +# renderer highlights any column whose value differs across the baselines (those +# are the ones to eyeball); columns identical across all 3 stay dim (expected). +tier2_cols = [ + ("sync_async", "mode"), ("selector", "selector"), ("optimizer", "optim"), + ("c", "c"), ("agg_goal", "agg_goal"), ("k", "k"), + ("min_init", "minInit"), ("n_trainers", "n_trainers"), + ("n_gpus", "n_gpus"), ("partition", "part"), ("avail", "avail"), +] +overridden2 = [] +if bool(SEL_C) or bool(SEL_C_ASYNC): overridden2.append("c") +if bool(AGG_GOAL) or bool(SEL_C): overridden2.append("agg_goal") +if bool(SEL_K): overridden2.append("k") +if bool(MIN_INIT): overridden2.append("min_init") +if bool(NUM_TRAINERS): overridden2.append("n_trainers") +if bool(NUM_GPUS): overridden2.append("n_gpus") +if bool(PART): overridden2.append("partition") +if trace_set: overridden2.append("avail") +rows2 = [] +for rk in (r[0] for r in runs): + b = per_baseline.get(rk, {}) + rows2.append({"name": rk, "cells": { + "sync_async": b.get("sync_async"), "selector": b.get("selector"), + "optimizer": b.get("optimizer"), + "c": b.get("c"), "agg_goal": b.get("agg_goal"), "k": b.get("k"), + "min_init": b.get("min_init"), "n_trainers": b.get("n_trainers"), + "n_gpus": b.get("n_gpus"), "partition": b.get("partition"), + "avail": b.get("avail"), + }}) +tiers.append({"name": "② PER-BASELINE (moderate)", + "table": {"columns": tier2_cols, "rows": rows2, + "overridden": overridden2}}) + +# ③ config-baked +tiers.append({"name": "③ RARELY CHANGED", "collapsed": True, "rows": [ + {"label": "env", "value": os.environ.get("CONDA_DEFAULT_ENV", "?")}, + {"label": "gpus_visible", "value": str(GPUS_VISIBLE)}, + {"label": "example_dir", "value": env("EXAMPLE_DIR")}, + {"label": "logdir", "value": LOGDIR}, +]}) +# ---------------- feasibility checks ---------------- +# D matched across each pair (by construction, but assert it visibly). +if MODE == "both": + checks.append({"name": "enable_training_delays matched across every real/sim pair", + "level": "ok", "detail": f"D={'>0' if delays_on else '0'} both sides"}) +# agg_goal <= c (more required than concurrently selected -> stall). +for rk in (r[0] for r in runs): + b = per_baseline.get(rk, {}) + c, g = b.get("c"), b.get("agg_goal") + if isinstance(c, int) and isinstance(g, int) and g > c: + checks.append({"name": f"agg_goal <= c ({rk})", "level": "error", + "detail": f"agg_goal={g} > c={c} — selected trainers would be stranded"}) + else: + checks.append({"name": f"agg_goal <= c ({rk})", "level": "ok", "detail": f"agg_goal={g} c={c}"}) +# agg_goal MATCHES across baselines (operator invariant 2026-07-06): the aggregation +# batch size should be identical for a fair head-to-head; a mismatch is almost always +# an unintended fan from --c/--c-async. Warn (visible), don't block (an experiment +# MAY intentionally vary it -- but then it's an eyeballed choice, not a silent one). +_goals = {rk: per_baseline.get(rk, {}).get("agg_goal") for rk in (r[0] for r in runs)} +_gset = {g for g in _goals.values() if g is not None} +if len(_gset) > 1: + checks.append({"name": "agg_goal matches across baselines", "level": "warn", + "detail": f"agg_goal differs: {_goals} — intended? (fair comparison expects one value)"}) +elif _gset: + checks.append({"name": "agg_goal matches across baselines", "level": "ok", + "detail": f"all baselines agg_goal={next(iter(_gset))}"}) +# Availability consistency + sync-barrier liveness: surface the RESOLVED per- +# baseline trace (what actually runs), and BLOCK a full-participation sync +# barrier under a non-syn_0 trace -- agg_goal == n_trainers can never assemble if +# any trainer is unavailable, so the real barrier waits to the wall cap (the +# fwdllm_plus / K-D20 stall; Stage C's wait bounds the log but still can't +# complete when full participation is required under scarcity). +for rk in (r[0] for r in runs): + b = per_baseline.get(rk, {}) + av, g, n, is_async = b.get("avail"), b.get("agg_goal"), b.get("n_trainers"), b.get("async") + if av and av != "syn_0": + if (not is_async) and isinstance(g, int) and isinstance(n, int) and g >= n: + checks.append({"name": f"availability liveness ({rk})", "level": "error", + "detail": f"trace={av} + sync agg_goal={g} >= n_trainers={n}: " + f"barrier can't assemble under unavailability → stall"}) + else: + checks.append({"name": f"availability ({rk})", "level": "warn", + "detail": f"trace={av} — non-syn_0 unavailability (Phase 2+); confirm intended"}) + else: + checks.append({"name": f"availability ({rk})", "level": "ok", "detail": f"trace={av}"}) +# (No k-vs-agg_goal check: in the random selector, send-side selection/concurrency +# is driven by `c` (required_trainers = min(len(ends), c - in_use)); `k` is the +# RECV-side batch size (num_ends_to_remove = min(..., self.k)), NOT a selection +# cap -- so k < agg_goal is fine, the barrier still collects agg_goal grads across +# RECV passes. c <= num_trainers and agg_goal <= c are the binding invariants.) +# num_gpus <= visible. +for rk in (r[0] for r in runs): + b = per_baseline.get(rk, {}) + g = b.get("n_gpus") + if isinstance(g, int) and GPUS_VISIBLE and g > GPUS_VISIBLE: + checks.append({"name": f"num_gpus <= gpus_visible ({rk})", "level": "error", + "detail": f"num_gpus={g} > visible={GPUS_VISIBLE}"}) +# num_trainers >= minInitialTrainers. +for rk in (r[0] for r in runs): + b = per_baseline.get(rk, {}) + n, mi = b.get("n_trainers"), b.get("min_init") + if isinstance(n, int) and isinstance(mi, int) and n < mi: + checks.append({"name": f"num_trainers >= minInitialTrainers ({rk})", "level": "error", + "detail": f"num_trainers={n} < minInitialTrainers={mi} — join barrier never clears"}) +# non-uniform partition: can't verify the H5 group from here. +if PART and PART != "uniform": + checks.append({"name": "partition group exists in agnews_partition.h5", "level": "warn", + "detail": f"verify group '{PART}' exists"}) +# parity pairing naming (only meaningful for a both-mode matrix). +if MODE == "both": + ok_pair = all(any(n.endswith("_real") for n, *_ in manifest) and + any(n.endswith("_sim") for n, *_ in manifest) for _ in [0]) + checks.append({"name": "run names carry _real/_sim tags for parity glob", + "level": "ok" if ok_pair else "error", + "detail": "scripts.parity.cli globs *baseline*real* / *baseline*sim*"}) + +goal_for_parity = AGG_GOAL or (SEL_C or "10") +next_cmd = (f"(cd {env('AC10_DIR')} && python -m scripts.parity.cli --batch " + f"--experiments-dir {env('EXAMPLE_DIR')}/experiments " + f"--baselines {' '.join(rk for rk, *_ in runs)} --agg-goal {goal_for_parity})") + +spec = { + "title": "FWDLLM RUN", + "subtitle": f"mode={MODE} {len(manifest)} run(s) · 🟢 set=flag override · (D)=yaml/code default", + "dry_run": DRY_RUN, + "tiers": tiers, + "checks": checks, + "next": next_cmd, +} +json.dump(spec, open(env("MANIFEST") + ".spec.json", "w"), indent=2) +rc = expt_runner.render_and_gate(spec, show_all=SHOW_ALL) +sys.exit(rc) +PY +GATE_RC=$? + +# ---- gate decision ---- +# render_and_gate returns 0 (ok) or 2 (blocking check). Anything else means the +# pre-flight step itself failed (e.g. a bad YAML or a spec-builder bug) -- abort +# rather than silently launch on an unvalidated config. +if [ "$GATE_RC" -ne 0 ] && [ "$GATE_RC" -ne 2 ]; then + echo "ERROR: pre-flight step failed (exit $GATE_RC) -- see traceback above. Nothing launched." >&2 + exit "$GATE_RC" +fi +if [ "$GATE_RC" -eq 2 ] && [ "$FORCE" != "1" ]; then + echo "Pre-flight BLOCKED (exit 2). Fix the config or pass --force to override. Nothing launched." >&2 + exit 2 +fi +if [ "$DRY_RUN" = "1" ]; then + echo "--dry-run: generated cfgs in $LOGDIR (manifest: $MANIFEST). Nothing launched." + exit 0 +fi +# Real (GPU) run confirmation unless --yes. +if [ "$ASSUME_YES" != "1" ]; then + read -r -p "Launch the runs above? [y/N] " _ans < /dev/tty || _ans="" + case "$_ans" in y|Y|yes|YES) ;; *) echo "Aborted (no --yes / declined). Nothing launched."; exit 0 ;; esac +fi + +# baseline names (for parity --batch and the summary), derived from RUNS. +RUNS_BASELINES="" +for _e in "${RUNS[@]}"; do RUNS_BASELINES="$RUNS_BASELINES ${_e%%:*}"; done +RUNS_BASELINES="${RUNS_BASELINES# }" + +# ---- post-launch hooks (--after ...), dispatched by expt_dispatch_after ---- +# Each is a shell function the shared harness calls by name; they own the +# fwdllm-specific command (parity CLI path / sanity extractor / plotter). +after_parity() { + # Real<->sim parity battery on the pairs just produced. The parity engine + # lives under async_cifar10/scripts (shared, fwdllm rungs registered in it). + ( cd "$AC10_DIR" && python -m scripts.parity.cli --batch \ + --experiments-dir "$EXAMPLE_DIR/experiments" \ + --baselines $RUNS_BASELINES --agg-goal "${AGG_GOAL:-${SEL_C:-10}}" \ + --json-out "$LOGDIR/parity_.json" ) +} +after_sanity() { + # Per-run sanity signals (selection / data_id / eval-per-data_id / partition). + local d + while IFS= read -r d; do + [ -d "$d" ] || continue + python "$SCRIPT_DIR/extract_sanity_checks.py" "$d" || true + done < <(find "$EXAMPLE_DIR/experiments" -maxdepth 1 -type d -name "run_*" -newer "$MANIFEST" 2>/dev/null) +} +after_plot() { + # Best-effort: analyze_run over the telemetry produced this session. + local ar="$REPO_ROOT/scripts/analysis/analyze_run.py" d + [ -f "$ar" ] || { echo " [after:plot] $ar not found — skipping" >&2; return 0; } + while IFS= read -r d; do + [ -d "$d/telemetry" ] || continue + python "$ar" "$d/telemetry" --out "$LOGDIR/plots_$(basename "$d")" || true + done < <(find "$EXAMPLE_DIR/experiments" -maxdepth 1 -type d -name "run_*" -newer "$MANIFEST" 2>/dev/null) +} + +# Convergence stop (EXPERIMENTS.md WS2): export so expt_launch arms the watcher. +# Only when --target-acc was passed; otherwise runs stay governed by their caps +# (default behavior unchanged). Window defaults to 20 bins. +if [ -n "$TARGET_ACC" ]; then + export EXPT_TARGET_ACC="$TARGET_ACC" + export EXPT_CONVERGE_WINDOW="${CONVERGE_WINDOW:-20}" + [ -n "$STALL_WINDOW_S" ] && export EXPT_STALL_WINDOW_S="$STALL_WINDOW_S" + [ -n "$STALL_MIN_DELTA" ] && export EXPT_STALL_MIN_DELTA="$STALL_MIN_DELTA" + [ -n "$STALL_ON" ] && export EXPT_STALL_ON="$STALL_ON" + [ -n "$LOSS_MIN_REL_DELTA" ] && export EXPT_LOSS_MIN_REL_DELTA="$LOSS_MIN_REL_DELTA" +fi + +# ---- PHASE B: launch each generated cfg sequentially ---- +declare -A RESULT DURATION_S +ORDERED_KEYS=() STOP_ALL=0 -for trace in "${TRACE_LIST[@]}"; do - AVAIL_TRACE="$trace" # read by patch_yaml() via the outer AVAIL_TRACE var - [ "$MULTI_TRACE" = "1" ] && echo "--- trace: ${trace:-} ---" - - for entry in "${RUNS[@]}"; do - name="${entry%%:*}" - src_cfg="${entry#*:}" - # Disambiguate by trace only when actually looping multiple traces, so a - # single-trace (or no-trace) invocation keeps today's exact file/key names. - if [ "$MULTI_TRACE" = "1" ]; then - key="${name}@${trace:-default}" - else - key="$name" - fi - cfg="$LOGDIR/${key}.yaml" - log="$LOGDIR/${key}.out" - patch_yaml "$src_cfg" "$cfg" "$name" - - start_ts=$(date +%s) - python -m flame.launch.run_experiment "$cfg" --example-dir "$EXAMPLE_DIR" \ - < /dev/null > "$log" 2>&1 & - CHILD_PID=$! - echo "[$(date '+%F %T')] START $key (PID=$CHILD_PID) -> $cfg (log: $log)" - echo " (to kill: kill -9 $CHILD_PID or Ctrl+C)" - wait "$CHILD_PID" - rc=$? - CHILD_PID="" - end_ts=$(date +%s) - DURATION_S[$key]=$((end_ts - start_ts)) - ORDERED_KEYS+=("$key") - if [ $rc -eq 0 ]; then - RESULT[$key]="PASS" - else - RESULT[$key]="FAIL(exit=$rc)" - fi - echo "[$(date '+%F %T')] DONE $key -> ${RESULT[$key]} (${DURATION_S[$key]}s)" +cd "$REPO_ROOT" || exit 1 - if [ $rc -ne 0 ] && [ "$STOP_ON_FAIL" = "1" ]; then - echo "--stop-on-fail set; aborting remaining runs (including remaining traces)." - STOP_ALL=1 - break - fi - done - [ "$STOP_ALL" = "1" ] && break -done +# Backstop trap for Ctrl+C landing OUTSIDE a run (between baselines, after-hooks); +# expt_launch installs its own thorough handler during each run. Sweep + exit. +_rs_interrupt() { + trap - INT TERM + echo "" >&2 + echo "[$(date '+%F %T')] INTERRUPT — aborting run-set, sweeping any workers ..." >&2 + pkill -TERM -f 'flame.launch.run_experiment' 2>/dev/null || true + pkill -9 -f 'trainer/forward_training' 2>/dev/null || true + pkill -9 -f 'trainer/pytorch/main.py' 2>/dev/null || true + pkill -9 -f 'aggregator/pytorch/main_' 2>/dev/null || true + pkill -9 -f converge_watch.py 2>/dev/null || true + exit 130 +} +trap _rs_interrupt INT TERM + +[ "$CLEAN" = "1" ] && export EXPT_AUTOCLEAN=1 # --clean -> preflight kills stragglers instead of aborting +while IFS=$'\t' read -r name cfg variant budget; do + [ -n "$name" ] || continue + trap _rs_interrupt INT TERM # re-arm: expt_launch clears its trap on return + # Clean-slate guard: refuse to launch on top of a prior run's stray workers. + if ! expt_assert_clean_slate "$name"; then + RESULT[$name]="DIRTY_ABORT"; ORDERED_KEYS+=("$name"); DURATION_S[$name]=0 + echo " [$name] aborting: node not clean (use --clean to auto-kill stragglers)." >&2 + STOP_ALL=1; break + fi + start_ts=$(date +%s) + expt_launch "$name" "$cfg" "$EXAMPLE_DIR" "$budget" 1 "$LOGDIR" + rc=$? + DURATION_S[$name]=$(( $(date +%s) - start_ts )) + ORDERED_KEYS+=("$name") + # Health verdict (COMPLETED / CRASH / NO_AGG_ROUNDS / WALL_CEILING) is the source + # of truth for the summary, NOT the launcher exit code -- which is 0 even when the + # aggregator subprocess crashed (run_experiment swallows the child's non-zero exit). + expt_assert_run "$EXAMPLE_DIR" "$EXPT_LAST_MARKER" "$name" + if [ "${EXPT_LAST_HEALTH:-}" = "CONVERGED" ] || [ "${EXPT_LAST_HEALTH:-}" = "STALLED" ]; then + # Watcher kills the run's process group -> rc is the SIGKILL code (expected), + # NOT a failure. Report the clean verdict (CONVERGED / STALLED) without exit noise. + RESULT[$name]="${EXPT_LAST_HEALTH}" + elif [ "$rc" -ne 0 ]; then + # Launcher itself failed: surface that, but keep the health word if the + # scan caught a more specific cause (e.g. CRASH) than a bare exit code. + case "${EXPT_LAST_HEALTH:-}" in + COMPLETED|NO_MARKER|"") RESULT[$name]="LAUNCH_EXIT=$rc" ;; + *) RESULT[$name]="${EXPT_LAST_HEALTH}(exit=$rc)" ;; + esac + else + RESULT[$name]="${EXPT_LAST_HEALTH:-COMPLETED}" + fi + if [ "$rc" -ne 0 ] && [ "$STOP_ON_FAIL" = "1" ]; then + echo "--stop-on-fail set; aborting remaining runs."; STOP_ALL=1; break + fi +done < "$MANIFEST" + +# ---- post-launch hooks ---- +[ -n "$AFTER" ] && [ "$STOP_ALL" != "1" ] && expt_dispatch_after "$AFTER" echo "" echo "=== Summary ===" for key in "${ORDERED_KEYS[@]}"; do - printf " %-35s %-15s %ss\n" "$key" "${RESULT[$key]:-SKIPPED}" "${DURATION_S[$key]:-0}" + printf " %-40s %-15s %ss\n" "$key" "${RESULT[$key]:-SKIPPED}" "${DURATION_S[$key]:-0}" done -echo "Logs: $LOGDIR" +echo "Logs: $LOGDIR" echo "Run dirs: $EXAMPLE_DIR/experiments/run_*" +echo "Parity: $(cd "$AC10_DIR" && echo "(cd $AC10_DIR && python -m scripts.parity.cli --batch --experiments-dir $EXAMPLE_DIR/experiments --baselines ${RUNS[*]%%:*} --agg-goal ${AGG_GOAL:-${SEL_C:-10}})")" diff --git a/lib/python/examples/fwdllm/fluxtune_contributions.md b/lib/python/examples/fwdllm/fluxtune_contributions.md new file mode 100644 index 000000000..94b2ea1cd --- /dev/null +++ b/lib/python/examples/fwdllm/fluxtune_contributions.md @@ -0,0 +1,268 @@ +# Fluxtune: Systems & ML Contributions for On-Device LLM Fine-Tuning + +**Thesis.** Fluxtune makes federated LLM fine-tuning practical on *memory- and hardware-constrained* +devices by training with **forward-mode (backprop-free) gradients** over **parameter-efficient adapters**, +aggregated **asynchronously**, with **informed (JVP-magnitude) perturbation selection**. Peak memory is +**independent of both model depth and perturbation count**; compute is **pure forward inference** (no +autograd graph, no backward GEMMs), so it maps onto the **inference-only accelerators** (mobile NPUs/DSPs) +that cannot run backpropagation at all. This doc states the contributions, quantifies them vs baselines and +alternatives, and separates *measured* from *argued*. + +> Numbers below are from `scripts/profile_jvp_opt.py` (reuses production `create_model` + `calculate_jvp`): +> DistilBERT-base + AdapterHub bottleneck adapters, batch 8, seq 192, fp16/autocast, NVIDIA A40. A *clean +> single-trainer* profile; a shared-GPU run multiplies wall-time by the contention factor, but **pass-counts, +> ratios, and memory transfer directly**. Mobile figures are *argued from structure*, not measured. + +--- + +## 1. The problem & why the baseline choice matters + +On-device transformer fine-tuning via **backpropagation** needs (i) a full **autograd graph** (stored +activations for the backward pass), (ii) **training-mode** kernels (backward GEMMs, transposes), (iii) +optimizer state. Mobile SoCs expose **inference-optimized** NPUs/DSPs (forward GEMMs, quantized, no autograd) +with tight memory — so backprop fine-tuning is often infeasible on the device that owns the data. + +**Forward-gradient** methods (FedFwd / FwdLLM family) replace the backward pass with **directional derivatives** +from *forward passes only*. Fluxtune builds on this. Its baselines here are the **sync** forward-grad variants +(`fwdllm`: cosine-similarity selection; `fwdllm++`: per-iteration reselection). This doc evaluates fluxtune's +distinguishing choices — **async aggregation + JVP-magnitude selection + adapter PEFT + server-side LR**. + +--- + +## 2. The computation, precisely + +Each perturbation's gradient estimate is a **central finite-difference JVP** (`fwdgrad_utils.calculate_jvp`): + +``` +jvp(v) = ( f(θ + h·v) − f(θ − h·v) ) / (2h), h = 0.01, under autocast + no_grad +``` + +- `f` = one full forward pass of the (frozen-backbone) model + cross-entropy loss. **1 perturbation = 2 forward + passes.** No backward pass, no stored graph. +- **Fluxtune** evaluates `P` candidate perturbations (`perturbation_count`, default 10) and **selects the one with + the largest |jvp|** (the steepest measured descent direction) → `2P` forward passes. The **sync `fwdllm` + baseline** selects by cosine similarity to a cached gradient → **0 forward passes** for selection, 1 final JVP. +- **PEFT:** the DistilBERT backbone (66.4M params) is **frozen**; only **bottleneck adapters in every layer + the + head** are trainable — **1.04M / 67.4M = 1.5%**. The perturbation `v` is non-zero only on trainable params. + +--- + +## 3. Systems contributions + +### 3.1 Memory is flat — in perturbation count *and* in depth +With **no autograd graph**, peak memory ≈ *model weights + one in-flight forward*. It does **not grow with +`P`** (perturbations evaluated one at a time, discarded) and does **not accumulate activations for a backward**. + +| regime | peak memory (measured) | +|---|---| +| forward-grad, any P (trainable-only, retained) | **3.19 GB** | +| forward-grad, any P (current) | 3.44 GB | +| backprop reference (1 fwd + 1 bwd) | 3.71 GB | + +At this favorable-to-backprop config (small batch, 98.5% frozen) the gap is ~14%; it **widens with batch size, +sequence length** (backprop's stored activations scale with both; forward-grad's do not) and **trainable +fraction**. The durable claim is structural: **forward-grad removes the autograd graph entirely**, so memory is +bounded by inference, not training. + +### 3.2 Compute is pure forward inference (hardware fit) +Every FLOP is a **forward-pass GEMM** — the exact operator set an inference NPU/DSP is built for. **No backward +GEMMs, no transposed weight matmuls, no autograd bookkeeping.** This is the **mobile-generality argument**: a +device that can *run* the model can *train* it under fluxtune, no autograd runtime. (Argued from operator +structure; not yet measured on-device.) + +### 3.3 Compute/latency cost — characterized honestly +Forward-grad trades memory for **time**: many forward passes instead of one forward+backward. The selection +surcharge is linear in `P`: + +| path | fwd passes | ms/batch (clean A40) | +|---|---|---| +| sync `fwdllm` (opt) | 2 | 16 | +| **fluxtune P=1** | 2 | 16 — *equals sync* | +| fluxtune P=5 | 10 | 80 | +| **fluxtune P=10** | 20 | 159 | +| backprop reference | 1f+1b | 17 | + +Cost is **`2P × per-pass`** → **10× sync compute at P=10**, collapsing to parity at `P=1`. JVP-selection is a +**tunable accuracy/compute knob**, not a fixed tax. Per-pass ≈ 8–10 ms here; a shared-GPU deployment multiplies +wall-time by the contention factor (the real 10-trainer run saw ~0.21 s/pass). + +### 3.4 Communication +Only the **1.5% adapter** parameters are exchanged per round (PEFT); async aggregation (`agg_goal=3 < K=10`) +commits **as stragglers arrive** with no synchronization barrier — targeting mobile intermittent connectivity +and device heterogeneity. + +--- + +## 4. Machine-learning contributions + +### 4.1 Informed perturbation selection (JVP vs cosine / random) +Random-direction forward-grad (MeZO-style) and the cosine-similarity baseline pick a perturbation *without* +measuring its loss effect. Fluxtune **measures** each candidate's directional derivative and keeps the steepest +— a better single-sample gradient estimate per round, improving sample/round efficiency at the cost of `2P` +forward passes (§3.3). `P` is an operator-controlled accuracy/compute trade. + +### 4.2 Finite-difference numerics & precision (a measured caution + an exactness result) +- The FD estimate subtracts two **O(1)** losses differing by **O(h)** ≈ 1e-3. Under **fp16/fp32** this is + **catastrophic cancellation**: the JVP retains ~1–2 significant figures, so the perturbation *ranking* is + mildly precision-limited — a genuine caveat for any deployment that lowers precision for the mobile NPU. +- **Batching is mathematically exact.** Vectorizing all `P` perturbations (`torch.func.vmap`) yields JVPs + **bit-identical to the sequential loop in fp64** and deterministic run-to-run — the fp32 divergence is *only* + the cancellation above, not a batching error. Bounds when the 2× batching speedup is safe to adopt. + +### 4.3 Tensor-operation profile +The FD perturbation touches **only trainable tensors** (`p ± h·v` with `v=0` on frozen params → `p−0=p` exactly); +the forward is otherwise identical inference. **No backward transpose-GEMM, no grad-accumulation kernel.** This +minimal op set enables the trainable-only optimization (§5) and the inference-hardware mapping (§3.2). + +--- + +## 5. Fidelity-preserving optimizations (retained) + +Validated **bit-identical** (`max|Δjvp| = 0`) by the profiler — they change *cost*, never the computed gradient, +so training fidelity and real↔sim simulator parity are untouched: + +1. **Trainable-only finite difference** — skip the `p − h·0 = p` arithmetic on the 98.5% frozen backbone: + **1.26× faster, −251 MB.** +2. **Remove redundant/diagnostic forward passes** — 3 passes that only fed a log line, plus (fluxtune) reusing the + selected perturbation's already-computed JVP: **fluxtune 25→20 passes, sync 5→2.** + +**Combined: sync −68%, fluxtune −37% GPU time, zero fidelity change.** Under the real 10-trainer / 8-GPU run +this brings fluxtune's **mean** per-batch compute (3.61s) under the 4.0s modeled mobile-delay budget, but the +**tail** (4.1–5.4s) still overruns on the two GPUs carrying 2 trainers each (10>8) plus the aggregator's eval +GPU — contention, not JVP cost. Clearing the tail needs the sim's GPU pipelining fix (keeps compute near the +~2.4s uncontended floor) and/or 1-trainer-per-GPU. + +**Deliberately *not* adopted** (change fidelity / inferior here): **vmap batching** (2× but re-baselines the +fp32 trajectory via §4.2 cancellation — exact only in fp64), **exact forward-mode AD** (0.5×, needs eager +attention, different math), **lowering `P`** (changes the algorithm — an accuracy knob, not a free optimization). + +--- + +## 6. Comparison summary + +| dimension | backprop FL | sync forward-grad (`fwdllm`/`++`) | **fluxtune** | +|---|---|---|---| +| training memory | activation graph (grows w/ batch·seq·depth) | flat (inference) | **flat (inference)** | +| hardware needed | autograd/training runtime | inference-only | **inference-only** | +| perturbation selection | n/a | cosine to cached grad (0 fwd) | **JVP magnitude (2P fwd, informed)** | +| aggregation | sync/async | **sync (barrier)** | **async (straggler-tolerant)** | +| compute/round | 1 fwd + 1 bwd | 1 JVP (2 fwd) | 2P fwd (tunable) | +| communication | full or PEFT | PEFT adapters | **PEFT adapters** | +| best when | server-class HW | homogeneous, fast clients | **memory/HW-constrained, heterogeneous, intermittent clients** | + +**Where fluxtune wins:** the on-device regime — constrained memory, inference-only accelerators, stragglers, +intermittent availability — where backprop is infeasible and a sync barrier stalls on the slowest phone. It buys +a better per-round gradient (JVP selection) and straggler tolerance (async) for a forward-pass compute cost that +§5 cuts ~40% without fidelity loss, and that `P` tunes directly. + +**Honest limits:** fluxtune is **compute-heavier** than the sync baselines (≈10× at P=10) and than backprop +per-round; its advantage is memory/hardware feasibility and robustness, not raw FLOPs. The FD JVP is +precision-sensitive in fp16 (§4.2). The memory gap over backprop is modest at small scale, growing with +batch/seq/trainable-fraction. + +--- + +## 7. Reproducibility + +All figures: `scripts/profile_jvp_opt.py` (env `test_fwdllm`), reusing the production model builder and JVP math +so gains transfer directly to the trainer. Per stage it reports forward-pass count, latency (mean±std, warmup + +`cuda.synchronize`), peak memory, speedup, and a BIT-IDENTICAL / WITHIN-TOL / DIVERGED verdict vs the +ground-truth sequential path, plus an fp64 cancellation diagnosis (`--fp64-check`). The real↔sim simulator that +validates fluxtune's *training dynamics* under a virtual clock is documented in `simulate_fwdllm.md` (compute +profile in its §L). + +--- + +## 8. Training-stability track — root cause & resolution (LIVE) + +> **DOC DISCIPLINE — STRICT. This section is *current truth*, not a log.** +> 1. **Edit in place, do NOT append.** When a finding/fix changes state, **rewrite its existing row** — no dated +> "update:" notes, no changelog, no sibling rows. Git holds history. (Opposite of `EXPERIMENTS.md`.) +> 2. **One row per finding, one per fix.** Superseded/refuted → overwrite. No duplicates. +> 3. **Every claim carries a status tag:** `VERIFIED` (evidence cited) · `SUSPECTED` (hypothesis) · `REFUTED` +> (checked false — keep the row so we don't re-chase) · `TODO` · `WIP` · `DONE` (landed **and** sanity-checked). +> 4. **No fix reaches `DONE` without its sanity check** named in the same row (telemetry / unit test / metric). +> 5. **Crisp only:** claim · evidence · status. No prose paragraphs in the ledgers. +> 6. **Flag-gate every change for A/B; the terminal lifecycle decision is the operator's.** Each fix lands +> **behind a named flag, default = old (byte-identical off)** — name it in "Code — how". States: **`A/B`** +> (both live) → **`PERMANENT`** (new default, old removed) · **`FLAGGED`** (new default, flag retained) · +> **`REVERTED`** (old reinstated). **Never pick the terminal state unilaterally — ASK the operator** with the +> A/B evidence; record the choice + deciding metric in the row. + +**The issue (one line).** On the N=100 α=1 runs the global model never converges — it **oscillates** with +recurring **single-class collapses** (acc 0.250, mcc 0.000 on balanced 4-class); peaks are transient and the +*same* positions collapse every epoch. Root cause (H0): an **undamped, high-variance forward-gradient +optimizer** (F6-F9) — each noisy JVP commit applied raw → random walk. Data/heterogeneity is **not** the driver +(F11 refuted class-bias; F13: the fixed schedule only *freezes* the noise → position-lock). Supersedes the +charter's I-1 "epoch-boundary bug" framing (F10 refuted). + +### 8.1 Verified-findings ledger (sanity checks / things found) + +| # | Finding | Evidence | Status | +|---|---|---|---| +| F1 | `data_id` selects the *actual training data*: at `data_id=k` a trainer trains on ONE fixed 8-sample batch `train_local_list[0][k]` | `FedSgdTrainer.py:500-505`, `:148-150`; `train_batch_size=8` (`configs/aggregator_base.json:37`); 1200 samples/client → 150 bins | VERIFIED | +| F2 | Per-client **train** loader is `shuffle=False` (SequentialSampler) → raw partition index order; the global **test** loader is `shuffle=True` | `base_data_manager.py:490` vs `:287` | VERIFIED | +| F3 | Bins materialized once, **never reshuffled across epochs**; round boundary resets `data_id=0` only → round 2 replays identical batch order | `FedSgdTrainer.py:148-150`; `fwdllm_aggregator.py:1940-1946` | VERIFIED | +| F4 | Down-swings are **single-class collapse**, not generic noise: acc = 0.250 **and** mcc = 0.000 exactly (32 of 150 R1 bins) | `agg_eval` telemetry, run `…025543…` | VERIFIED | +| F5 | Collapses are **position-locked and replay**: R1 and R2 collapse at the same `data_id` ranges (≈4-7, ≈63-79) | `agg_eval` telemetry | VERIFIED | +| F6 | Early variance-gate commits yield worse models: commit at it<15 → mean acc **0.289**; it≥20 → **0.595** | `agg_eval` telemetry | VERIFIED | +| F7 | Variance **floor (~0.45) > threshold (0.30)** → bins commit only on a *noise dip*, i.e. the noisiest estimates | charter M-3; `characterize_variance_curve.py` | VERIFIED | +| F8 | Global update is a **direct in-place SGD step** — `param.sub_(lr·Σ rate·gᵢ / N)` — with **no momentum, no EMA of global weights, no server optimizer state** across commits | `FedSgdAggregator.py:322-324`; §5 (aggregator trace) | VERIFIED | +| F9 | fedbuff "new" rate can exceed 1 (`beta_polynomial_upshift` adds +0.5) → some updates **amplified**, not damped | `fedbuff.py:100-101,139` | VERIFIED | +| F10 | Suspected round-boundary **staleness-reset bug is REFUTED**: `_model_version` is monotonic across the boundary, staleness stays ≥ 0. Only real discontinuity is the LR `ratio` warmup→decay step | `fwdllm_aggregator.py:1926-1929,1940-1946`; `FedSgdAggregator.py:234-242` | REFUTED | +| F11 | **REFUTED as the primary driver:** the per-`data_id` gradient the aggregator sees is **not** class-biased. Pooled over a realistic K=10 commit cohort at α=1 the class mix is near-balanced (dominant frac med **0.36**, entropy **0.95**, only **2.8%** of commits majority-one-class); pooled composition has **zero correlation** with which data_ids collapse (Spearman +0.002, p=0.98). A 0.36-dominant gradient cannot drive an exact single-class (25%/mcc=0) collapse. | H0 `diagnose_partition_binning.py` Layer C′ + overlay, `_diag_partition/` | REFUTED | +| F12 | Per-client class skew by α (100-client dist, exact): **α=0.1** dom-frac med **1.00** (½ of clients single-class) · **α=1** med **0.72**, 3 classes, entropy 0.50 · **α=100** med **0.30**, 4 classes, entropy 0.99 | H0 Layer A | VERIFIED | +| F13 | Collapses are **frozen-noise events, not class-structure events**: data (bins/cohort) is identical every epoch → the same high-variance JVP estimates recur at the same data_ids → position-locked collapse (F5). The *magnitude/variance* of the update, not its class direction, is the driver → points at F6-F9 (optimizer), not H1 (shuffle). | H0 (F11) + F6-F9 | VERIFIED | +| F14 | Partition is **quantity-balanced label-skew**: exactly **1200 samples/trainer at every α** (per-class totals 30k each) — α changes only the class *mix*, never the amount. Heterogeneity ladder (per-bin dom-frac mean / % bins ≥50%-one-class / #single-class trainers): α0.1 **0.94/96%/53** · α0.5 0.83/90%/0 · α1 **0.75/81%/1** · α5 0.57/46% · α10 0.52/32% · α100 **0.45/14%/0**. | H0 `--dist`, `trainer_class_heatmap.pdf` | VERIFIED | +| F15 | **Small-batch lumpiness is α-independent:** even at α=100 (near-IID clients) 14% of 8-sample bins are ≥50% one class — an artifact of the bin *size*, not heterogeneity. Within a trainer, bins are ≈ iid draws of that trainer's own class mix (the `executor.map` load already scrambles intra-trainer order — no class-sorted sequence). So **bin *composition/size* matters; intra-bin *order* is a no-op** (JVP/loss is a mean over the bin → permutation-invariant). | H0 `--dist`; F1 | VERIFIED | + +### 8.2 Resolution plan — by scope + +Two independent levers. Each item is tagged **cross-baseline hygiene** (applied identically to `fwdllm` / +`fwdllm++` / `fluxtune` — a fair-comparison correctness fix, **not** a fluxtune contribution) or a +**fluxtune-specific contribution** (claimed over baselines). Charter Opt-ladder cross-refs noted where they overlap. + +**Cross-baseline hygiene (H).** + +| ID | Fix | Why (finding) | Code — how | Sanity check to pass | Status | +|---|---|---|---|---|---| +| H0 | Diagnostic: per-client × per-`data_id` × per-K-cohort class distribution from `agnews_partition.h5` + frozen caches; overlay vs. observed collapses | Confirm/refute F1-F5 before any change | `expt_scripts/diagnose_partition_binning.py` (faithful cache order; α∈{0.1,1,100}) | ran; overlay Spearman +0.002 → **refuted** the data-bias story (F11) | **DONE** | +| H1 | ~~Shuffle each client's local train data before binning~~ **DE-PRIORITIZED** | F11: bins/cohort already near class-balanced at α=1 → little bias to remove. Only value is de-correlating F13 frozen noise across epochs (marginal) | `base_data_manager.py:490` `shuffle=True`+seed / per-round permute | would break F5 position-lock but not the collapse magnitude | PARKED (revisit if α=0.1) | +| H2 | **Larger data bins / higher agg_goal** (variance reduction) — ⚠ *not unambiguously good:* a large class-mixed bin averages conflicting per-sample gradients → the scalar JVP signal can sink below its noise floor and learning stalls. α-dependent optimum → sweep first (EXPERIMENTS.md **M1**). | F13,F6,F15: 8-sample bins × K=10 → high-variance, lumpy JVP is the driver, but too-large kills the signal | `train_batch_size` / bin regroup; `agg_goal` | M1 sweep finds the acc-maximizing bin size per α | DEFERRED | +| H3 | **Per-epoch, seeded, aggregator-orchestrated bin-order permutation** (visit bins in a fresh random order each round, once each) | F13,F15: visit order is irrelevant *in expectation* but the frozen 0→149 order replays the same noisy-estimate sequence every epoch → position-locked collapse. Randomizing de-correlates it across epochs. Intra-bin shuffle is NOT needed (F15). | aggregator drives `data_id` (`fwdllm_aggregator.py:1922`); broadcast a per-round seeded permutation — **must stay deterministic for real↔sim cohort_sequence parity** | F5 position-lock disappears across epochs; collapses stop recurring at fixed data_ids | TODO | + +**Fluxtune-specific contributions (S — stability).** Claimed over baselines; these are where we take credit. + +| ID | Contribution | Why (finding) | Code — how | Sanity check to pass | Status | +|---|---|---|---|---|---| +| S1 | **Server-side optimizer with momentum / EMA of the global model** (a damping / restoring force) | F8: undamped direct SGD → random walk on the loss surface | `FedSgdAggregator.py:322-324` add a server momentum buffer or global-weight EMA | loss envelope becomes monotone; peak is *sustained*, not transient | TODO | +| S2 | **Variance-gate recalibration** — commit on the *plateau*, not on a noise dip; align threshold to the achievable floor | F6,F7: gate commits the noisiest updates | `var_threshold` + plateau policy (extends charter **Opt-2**) | mean it-at-commit rises; early-commit collapse (F6) gone | TODO (Opt-2 partial) | +| S3 | **Aggregation-rate tempering** — cap rate ≤ 1; retune grad-aware to damp, not amplify | F9 + charter: R4 (full grad-aware) diverged *worst* | `fedbuff.py` beta upshift; grad_aware `align_floor`/`inverse_var` (retunes charter **Opt-3**) | R4 no longer the worst diverger; per-commit step magnitude bounded | TODO | + +**Order of attack — pick by IMPACT, not table order.** Re-rank each turn to whatever best fixes the problem: +- **NEXT → S1** (server optimizer: descent, not random walk) — the root cause (EXPERIMENTS.md M2). +- then **S2** (don't commit on noise dips), **S3** (rate cap). +- **DEFERRED:** H2 (bin size — needs the M1 sweep), H3 (bin-order permutation). **PARKED:** H1. + +### 8.3 Claim vs. correctness +- **Claimable (S1-S3):** fluxtune's async forward-grad aggregation is uniquely exposed to high-variance, + amplifiable, undamped updates (F8,F9) → a *server optimizer* (S1) and *signal-aware commit gate* (S2) that + stabilize forward-mode async FL are genuine contributions the sync baselines don't need. +- **Not claimable (H1-H3):** shuffle / bin-size / bin-order are correctness fixes any FL should have; applied to + all three baselines so E1 stays fair (P1). Reported as fixed, not as wins. + +### 8.4 Design Q&A +- **Bin vs. classical-FL round?** FedAvg updates from the *whole* local set (batch washed out pre-aggregation); + FwdLLM commits per **8-sample bin** ⇒ bin/cohort size *is* the per-update variance — a knob that matters here, + not in FedAvg (M1). +- **Shuffle *within* a bin?** No — JVP/loss is a mean over the bin ⇒ permutation-invariant (F15). No-op. +- **Why sequential bin order?** Order is irrelevant *in expectation*; the harm is the **frozen** 0→149 replaying + the same noisy sequence every epoch → position-locked collapse (F13). Fix = per-epoch **seeded, + aggregator-driven** permutation (H3), deterministic for sim parity. Reduces *repetition*, not *amplitude*. +- **Why does order matter if the model ignores sequence?** Only because the optimizer random-walks today (F8). At + a real minimum order won't matter ⇒ fix the optimizer (S1), don't lean on order. + +**S1 scoping (next task).** Flag `server_optimizer` (default off = raw SGD, byte-identical). Add momentum / +weight-EMA at `FedSgdAggregator.py:322-324` (`param.sub_(lr·Σg/N)`). Validate before A/B: (i) optimizer state +**deterministic** under frozen update order (sim parity); (ii) no **double-damp** with the fedbuff rate + variance +gate; (iii) A/B vs off at α=1 N=100 → sustained peak + W=20 window fills. Composable with C3/Opt-3, not a replacement. diff --git a/lib/python/examples/fwdllm/scripts/profile_jvp_opt.py b/lib/python/examples/fwdllm/scripts/profile_jvp_opt.py new file mode 100644 index 000000000..a81b79d2a --- /dev/null +++ b/lib/python/examples/fwdllm/scripts/profile_jvp_opt.py @@ -0,0 +1,435 @@ +#!/usr/bin/env python3 +"""Fluxtune/fwdllm forward-grad JVP optimization profiler (TEMP, real-code reuse). + +Isolated latency + correctness profiler for the perturbation/JVP compute that +dominates fluxtune GPU time (`_train_one_batch` in +`trainer/forward_training/tc_transformer_trainer_distribute.py`). Measures the +gain and validates the numerical fidelity of each optimization stage before +changing the real trainer. + +Deliberately REUSES the production code so a validated speedup transfers directly: +the real model (`expts.initializer.create_model` → DistilBERT-base + AdapterHub +bottleneck adapters, backbone frozen) and the real JVP math +(`fwdgrad_utils.calculate_jvp` / `functional_get_loss`). Targets only the current +fluxtune/fwdllm config (distilbert-base-uncased, agnews→4 labels, seq 192, +batch 8, adapter PEFT, perturbation_count 10) — not generic. + +STAGES (each vs Stage 0 as the numerical ground truth): + 0 baseline the real `calculate_jvp` in a Python loop over perturbations + (central finite difference, 2 forward passes each = the real code) + 1 vmap-FD the SAME finite-difference JVP, all perturbations batched via + torch.func.vmap (frozen params broadcast, only adapters vary) + 2 fwd-AD exact forward-mode-AD JVP (torch.func.jvp) — DIFFERENT math + (true directional derivative, not a finite difference); a small + systematic gap vs stage 0 is EXPECTED (the O(h^2) FD error). + +Per stage: forward-pass count, latency (mean±std over repeats, warmup + +cuda.sync), speedup vs baseline, max|Δjvp| / max rel-diff vs baseline, and a +verdict: BIT-IDENTICAL / WITHIN-TOL / DIVERGED. + +Run in the test_fwdllm env: + conda run -n test_fwdllm python scripts/profile_jvp_opt.py + conda run -n test_fwdllm python scripts/profile_jvp_opt.py --no-autocast +""" +from __future__ import annotations + +import argparse +import statistics +import sys +import time +from functools import partial +from pathlib import Path + +import torch + +# lib/python on the path so `examples.fwdllm...` imports resolve +_LIB_PYTHON = Path(__file__).resolve().parents[3] +if str(_LIB_PYTHON) not in sys.path: + sys.path.insert(0, str(_LIB_PYTHON)) + +import functorch as fc # noqa: E402 +from torch.func import jvp as func_jvp # noqa: E402 +from torch.func import vmap # noqa: E402 + +from examples.fwdllm.expts.initializer import create_model, set_seed # noqa: E402 +from examples.fwdllm.trainer.model.transformer.model_args import ( # noqa: E402 + ClassificationArgs, +) +from examples.fwdllm.trainer.forward_training.fwdgrad_utils import ( # noqa: E402 + calculate_jvp, + functional_get_loss, +) + + +# ─────────────────────────── real model (reused) ─────────────────────────── +def build_model(num_labels, seq, attn="sdpa"): + """Exactly the trainer's construction path (main.py): DistilBERT-base + + AdapterHub adapter, backbone frozen, only adapters + head trainable. + + `attn`: "sdpa" (the real default, fastest) or "eager" — forward-mode AD + (stage 2) is not implemented for scaled_dot_product_attention, so it needs + eager attention to run at all.""" + ma = ClassificationArgs() + ma.model_name = "distilbert-base-uncased" + ma.model_type = "distilbert" + ma.load(ma.model_name) # no-op for a hub name (no local model_args.json) + ma.num_labels = num_labels + ma.client_idx = 0 + ma.config["attn_implementation"] = attn + ma.update_from_dict( + { + "peft_method": "adapter", # BnConfig, output_adapter per layer + "do_lower_case": True, + "max_seq_length": seq, + "train_batch_size": 8, + "manual_seed": 42, + "var_control": True, + "perturbation_sampling": True, + "select_perturbation_using_jvp": True, # fluxtune JVP-selection path + "fp16": True, + "reprocess_input_data": False, + "overwrite_output_dir": True, + } + ) + ma.config["num_labels"] = num_labels + _, model, _ = create_model(ma, formulation="classification") + return model + + +# ─────────────────────────── JVP stage implementations ─────────────────────────── +def _loss_partial(fmodel, buffers, num_labels, x, labels): + # == _compute_forward_jvp's `f`: functional_get_loss(params, ...) + return partial( + functional_get_loss, + model=fmodel, + buffers=buffers, + num_classes=num_labels, + x=x, + t=labels, + ) + + +def stage0_baseline(f, params, V, h): + """The real code: calculate_jvp per perturbation (2 forward passes each).""" + jvps = [] + for v in V: # V: list[tuple(dir per param)] + _loss, jvp = calculate_jvp(f, params, v) + jvps.append(jvp) + return torch.stack([j.float().reshape(()) for j in jvps]) + + +def stage0b_trainable_only(fmodel, buffers, num_labels, x, labels, params, V, + train_idx, h, device): + """Same central-FD JVP but perturb ONLY the trainable params (frozen p-h*0=p + is skipped — the frozen tensors are reused as-is). Bit-identical to stage 0 + (p-0.0=p in IEEE), avoids copying the 98.5% frozen backbone twice/perturbation.""" + jvps = [] + with torch.no_grad(): + for v in V: + plus = list(params) + minus = list(params) + for i in train_idx: + plus[i] = params[i] + h * v[i] + minus[i] = params[i] - h * v[i] + with torch.autocast(device_type=device.type, enabled=(device.type == "cuda")): + lp = functional_get_loss(tuple(plus), model=fmodel, buffers=buffers, + num_classes=num_labels, x=x, t=labels) + lm = functional_get_loss(tuple(minus), model=fmodel, buffers=buffers, + num_classes=num_labels, x=x, t=labels) + jvps.append((lp - lm) / (2 * h)) + return torch.stack([j.float().reshape(()) for j in jvps]) + + +def stage1_vmap_fd(fmodel, buffers, num_labels, x, labels, params, train_idx, + V_train, h, autocast_on, device): + """Same central-FD JVP, all perturbations batched via vmap. Only the + trainable adapter dirs get a leading P dim (frozen params broadcast) so we + never materialise P copies of the 66M frozen backbone.""" + def jvp_one(vt_tuple): # vt_tuple: per-trainable dir (no P dim inside vmap) + plus = list(params) + minus = list(params) + for k, i in enumerate(train_idx): + plus[i] = params[i] + h * vt_tuple[k] + minus[i] = params[i] - h * vt_tuple[k] + with torch.autocast(device_type=device.type, enabled=autocast_on): + lp = functional_get_loss(tuple(plus), model=fmodel, buffers=buffers, + num_classes=num_labels, x=x, t=labels) + lm = functional_get_loss(tuple(minus), model=fmodel, buffers=buffers, + num_classes=num_labels, x=x, t=labels) + return ((lp - lm) / (2 * h)).float().reshape(()) + + in_dims = (tuple(0 for _ in train_idx),) + return vmap(jvp_one, in_dims=in_dims)(V_train) + + +def stage2_forward_mode(f, params, V, h): + """Exact forward-mode AD (torch.func.jvp): the true directional derivative, + not a finite difference. Different math -> a small O(h^2) gap vs stage 0 is + expected; answers 'is exact fwd-AD faster/more accurate than 2 FD passes?'.""" + jvps = [] + for v in V: + _loss, jv = func_jvp(f, (params,), (v,)) + jvps.append(jv) + return torch.stack([j.float().reshape(()) for j in jvps]) + + +# ─────────────────────────────── harness ─────────────────────────────── +def _time(fn, warmup, repeats, device): + for _ in range(warmup): + fn() + if device.type == "cuda": + torch.cuda.synchronize() + ts = [] + for _ in range(repeats): + t0 = time.perf_counter() + fn() + if device.type == "cuda": + torch.cuda.synchronize() + ts.append((time.perf_counter() - t0) * 1e3) + return statistics.mean(ts), (statistics.stdev(ts) if len(ts) > 1 else 0.0) + + +def _measure(fn, warmup, repeats, device): + """Latency (mean±std ms) AND peak CUDA memory (MB) for one call of fn.""" + for _ in range(warmup): + fn() + if device.type == "cuda": + torch.cuda.synchronize() + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + ts = [] + for _ in range(repeats): + t0 = time.perf_counter() + fn() + if device.type == "cuda": + torch.cuda.synchronize() + ts.append((time.perf_counter() - t0) * 1e3) + peak = (torch.cuda.max_memory_allocated() / 1024**2) if device.type == "cuda" else 0.0 + return statistics.mean(ts), (statistics.stdev(ts) if len(ts) > 1 else 0.0), peak + + +def backprop_ref(num_labels, seq, batch, h, seed, attn, autocast_on, device): + """Standard BACKPROP fine-tune step (1 forward + 1 backward on the trainable + adapters/head) — the memory/time contrast forward-grad is chosen to avoid. + Fresh model so functorch's make_functional does not interfere.""" + model = build_model(num_labels, seq, attn).to(device).eval() + g = torch.Generator(device="cpu").manual_seed(seed) + x = torch.randint(0, 30522, (batch, seq), generator=g).to(device) + labels = torch.randint(0, num_labels, (batch,), generator=g).to(device) + + def step(): + model.zero_grad(set_to_none=True) + with torch.autocast(device_type=device.type, enabled=autocast_on): + out = model(x) + logits = out[0] if not hasattr(out, "logits") else out.logits + loss = torch.nn.functional.cross_entropy(logits.float(), labels) + loss.backward() + return step + + +def _verdict(ref, got): + if ref.shape != got.shape: + return "SHAPE-MISMATCH", float("nan"), float("nan") + if torch.equal(ref, got): + return "BIT-IDENTICAL", 0.0, 0.0 + mad = (ref - got).abs().max().item() + rel = ((ref - got).abs() / (ref.abs() + 1e-12)).max().item() + tag = "WITHIN-TOL" if torch.allclose(ref, got, rtol=1e-3, atol=1e-4) else "DIVERGED" + return tag, mad, rel + + +def _fp64_check(attn, num_labels, seq, batch, P, h, seed, device): + """Prove the fp32 vmap-FD divergence is FINITE-DIFFERENCE CANCELLATION, not + a batching bug: in fp64 the sequential and vmap FD JVPs should agree to + ~machine eps. Also prints the JVP signal magnitude so 'DIVERGED' on tiny + (lp-lm) values is interpretable.""" + model = build_model(num_labels, seq, attn).to(device).double().eval() + mask = [p.requires_grad for p in model.parameters()] + fmodel, params, buffers = fc.make_functional_with_buffers(model) + train_idx = [i for i, m in enumerate(mask) if m] + g = torch.Generator(device="cpu").manual_seed(seed) + x = torch.randint(0, 30522, (batch, seq), generator=g).to(device) + labels = torch.randint(0, num_labels, (batch,), generator=g).to(device) + gp = torch.Generator(device=device).manual_seed(seed + 1) + V, per_train = [], [[] for _ in train_idx] + for _ in range(P): + v = [torch.randn(p.shape, generator=gp, device=device, dtype=p.dtype) + if mask[i] else torch.zeros_like(p) for i, p in enumerate(params)] + V.append(tuple(v)) + for k, i in enumerate(train_idx): + per_train[k].append(v[i]) + V_train = tuple(torch.stack(per_train[k]) for k in range(len(train_idx))) + f = _loss_partial(fmodel, buffers, num_labels, x, labels) + seq_jvp = stage0_baseline(f, params, V, h).double() + vmap_jvp = stage1_vmap_fd(fmodel, buffers, num_labels, x, labels, params, + train_idx, V_train, h, False, device).double() + tag, mad, rel = _verdict(seq_jvp, vmap_jvp) + print(f" fp64 seq-FD vs vmap-FD : {tag} max|Δ|={mad:.2e} rel={rel:.2e}") + print(f" jvp magnitude (fp64) : |jvp| in " + f"[{seq_jvp.abs().min():.2e}, {seq_jvp.abs().max():.2e}] " + f"(FD subtracts losses ~O(1) -> cancellation floors fp32 precision)") + + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--perturbations", type=int, default=10) # perturbation_count + ap.add_argument("--batch", type=int, default=8) + ap.add_argument("--seq", type=int, default=192) + ap.add_argument("--labels", type=int, default=4) # agnews + ap.add_argument("--h", type=float, default=0.01) + ap.add_argument("--repeats", type=int, default=20) + ap.add_argument("--warmup", type=int, default=5) + ap.add_argument("--no-autocast", action="store_true", + help="fp32 (default: autocast, matching the real fp16 code)") + ap.add_argument("--attn", choices=["sdpa", "eager"], default="sdpa", + help="eager needed for forward-mode AD (stage 2); sdpa is the real default") + ap.add_argument("--fp64-check", action="store_true", + help="run the fp64 seq-vs-vmap cancellation diagnosis") + ap.add_argument("--seed", type=int, default=42) + args = ap.parse_args(argv) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + autocast_on = (not args.no_autocast) and device.type == "cuda" + set_seed(args.seed) + + model = build_model(args.labels, args.seq, args.attn).to(device).eval() + trainable_mask = [p.requires_grad for p in model.parameters()] + fmodel, params, buffers = fc.make_functional_with_buffers(model) + train_idx = [i for i, m in enumerate(trainable_mask) if m] + + total = sum(p.numel() for p in params) + trainable = sum(params[i].numel() for i in train_idx) + + # fixed-seed inputs (the JVP-METHOD comparison is data-independent) — the + # real forward is model(params, buffers, x)[0]; x = token ids, no mask. + g = torch.Generator(device="cpu").manual_seed(args.seed) + x = torch.randint(0, 30522, (args.batch, args.seq), generator=g).to(device) + labels = torch.randint(0, args.labels, (args.batch,), generator=g).to(device) + + # perturbations: randn on trainable, zero on frozen (== _prepare_perturbation_tensors) + gp = torch.Generator(device=device).manual_seed(args.seed + 1) + V = [] + V_train = [] # trainable-only stacked view for vmap + per_train = [[] for _ in train_idx] + for _ in range(args.perturbations): + v = [] + for i, p in enumerate(params): + if trainable_mask[i]: + d = torch.randn(p.shape, generator=gp, device=device, dtype=p.dtype) + else: + d = torch.zeros_like(p) + v.append(d) + V.append(tuple(v)) + for k, i in enumerate(train_idx): + per_train[k].append(v[i]) + V_train = tuple(torch.stack(per_train[k]) for k in range(len(train_idx))) # each (P, *shape) + + f = _loss_partial(fmodel, buffers, args.labels, x, labels) + + print("=" * 94) + print("fluxtune/fwdllm JVP optimization profiler (TEMP — reuses create_model + calculate_jvp)") + print(f"device={device.type} autocast={autocast_on} seed={args.seed}") + print(f"model: distilbert-base + AdapterHub adapters | params={total/1e6:.1f}M " + f"trainable={trainable/1e6:.3f}M ({100*trainable/total:.1f}%) [{len(train_idx)} tensors]") + print(f"batch={args.batch} seq={args.seq} P(perturbations)={args.perturbations} " + f"h={args.h} repeats={args.repeats} (+{args.warmup} warmup)") + print("=" * 94) + print(f"{'stage':22s} {'fwd/batch':>9s} {'ms/batch':>13s} {'speedup':>8s} " + f"{'max|Δjvp|':>11s} {'max rel':>10s} verdict") + print("-" * 94) + + stages = [ + ("0 baseline (calc_jvp)", 2 * args.perturbations, + lambda: stage0_baseline(f, params, V, args.h)), + ("0b FD trainable-only", 2 * args.perturbations, + lambda: stage0b_trainable_only(fmodel, buffers, args.labels, x, labels, + params, V, train_idx, args.h, device)), + ("1 vmap-FD (all pert)", 2 * args.perturbations, + lambda: stage1_vmap_fd(fmodel, buffers, args.labels, x, labels, params, + train_idx, V_train, args.h, autocast_on, device)), + ("2 forward-mode AD", 1 * args.perturbations, + lambda: stage2_forward_mode(f, params, V, args.h)), + ] + + base = None + for name, n_pass, fn in stages: + try: + out = fn().detach().float() + ms, sd = _time(fn, args.warmup, args.repeats, device) + except Exception as e: + print(f"{name:22s} {n_pass:>9d} {'ERR':>13s} " + f"{type(e).__name__}: {str(e)[:44]}") + continue + if base is None: + base = (ms, out) + print(f"{name:22s} {n_pass:>9d} {ms:>9.2f}±{sd:>4.1f} {'1.00x':>8s} " + f"{'—':>11s} {'—':>10s} REFERENCE") + else: + verdict, mad, rel = _verdict(base[1], out) + print(f"{name:22s} {n_pass:>9d} {ms:>9.2f}±{sd:>4.1f} {base[0]/ms:>6.2f}x " + f"{mad:>11.2e} {rel:>10.2e} {verdict}") + + print("-" * 94) + + # ── COST COMPARISON: sync vs fluxtune (perturbation sweep) vs backprop ── + # Both baselines run the SAME _train_one_batch; the difference is the + # perturbation-SELECTION cost: fwdllm/sync uses cos-sim (0 forward passes, + # 1 final JVP=2 passes); fluxtune does JVP-selection (2*P passes) + 1 final. + # We retain ONLY fidelity-preserving (bit-identical) cuts: drop the 3 + # diagnostic-only passes (both) + reuse the winner JVP (fluxtune) => the + # per-batch forward-pass count below. + # Measure the two per-forward-pass unit costs + peak memory (memory is flat + # in P: no autograd graph, one forward held at a time). Totals are then + # passes x per-pass — transparent projection. + P = args.perturbations + fms, _, fmem = _measure(lambda: stage0_baseline(f, params, V[:P], args.h), + args.warmup, args.repeats, device) + oms, _, omem = _measure( + lambda: stage0b_trainable_only(fmodel, buffers, args.labels, x, labels, + params, V[:P], train_idx, args.h, device), + args.warmup, args.repeats, device) + full_pp, opt_pp = fms / (2 * P), oms / (2 * P) # ms per forward pass + print("COST MODEL — forward-grad (finite-difference JVP), per training batch") + print(f"per-pass: full-param={full_pp:.2f}ms trainable-only={opt_pp:.2f}ms " + f"(1.{round(100*(full_pp/opt_pp-1)):02d}x, BIT-IDENTICAL) | peak mem " + f"full={fmem:.0f}MB opt={omem:.0f}MB (flat in P)") + print(f"{'path':32s} {'fwd passes':>10s} {'ms/batch':>10s} note") + print("-" * 94) + # (passes, per-pass, note): current uses full-param; opt uses trainable-only + rows = [ + ("sync fwdllm (current)", 5, full_pp, "cos-sim select(0) + 1 JVP + 3 diag"), + ("sync fwdllm (opt)", 2, opt_pp, "1 JVP trainable-only; -3 diag [RETAIN]"), + ("fluxtune P=1 (opt)", 2, opt_pp, "== sync when P=1"), + ("fluxtune P=2 (opt)", 4, opt_pp, ""), + ("fluxtune P=5 (opt)", 10, opt_pp, ""), + (f"fluxtune P={P} (current)", 2 * P + 5, full_pp, "2P + 2 redundant + 3 diag"), + (f"fluxtune P={P} (opt)", 2 * P, opt_pp, "2P trainable-only; winner reused; -3 diag [RETAIN]"), + ] + for name, npass, pp, note in rows: + print(f"{name:32s} {npass:>10d} {npass*pp:>8.1f} {note}") + + bp = backprop_ref(args.labels, args.seq, args.batch, args.h, args.seed, + args.attn, autocast_on, device) + bms, bsd, bmem = _measure(bp, args.warmup, args.repeats, device) + print(f"{'backprop ref (1 fwd+1 bwd)':30s} {'1f+1b':>10s} {bms:>9.1f}±{bsd:>3.1f} {bmem:>9.0f}" + f" contrast: stores activations") + print("-" * 94) + print(f"forward-grad peak mem ≈ model+1 forward (no autograd graph); backprop " + f"stores activations ({bmem:.0f} MB).") + print("-" * 94) + + if args.fp64_check: + print("numerical diagnosis:") + _fp64_check(args.attn, args.labels, args.seq, args.batch, + args.perturbations, args.h, args.seed, device) + print("-" * 94) + print("stage 0 = numerical ground truth (the real calculate_jvp loop). Real fluxtune per") + print("batch also does +2 redundant + 3 diagnostic-only passes (removable, bit-identical)") + print("on top of P*2 selection passes. fwdllm uses cos-sim selection (no JVP-selection loop).") + print("forward-mode AD (stage 2) needs --attn eager (not impl for SDPA).") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/lib/python/examples/fwdllm/simulate_fwdllm.md b/lib/python/examples/fwdllm/simulate_fwdllm.md index adb27b247..f7ef67833 100644 --- a/lib/python/examples/fwdllm/simulate_fwdllm.md +++ b/lib/python/examples/fwdllm/simulate_fwdllm.md @@ -1,278 +1,570 @@ -# High-Fidelity Simulator for FwdLLM -- Design & Staged Build Plan - -**Design-only.** This document is the plan for building a high-fidelity **simulated-clock** runner -for the `fwdllm` example (FedFwd / forward-gradient FL) that reaches **real<->sim parity** across the -**fluxtune / fwdllm / fwdllm++** baselines, at **100% availability (syn_0)** and under -**unavailability (syn_20, syn_50, mobiperf)**. The implementation is a **future branch/PR**; nothing -here is built yet. - -**Prerequisites (read first):** -- [async_cifar10/PARITY.md](../async_cifar10/PARITY.md) -- the parity methodology (the ladder §1, roles/ - tiers/dependency-gating, workflow policy, run-length budget, landed sim mechanisms §3). **fwdllm's - own rung catalog is PARITY.md §F** (modified rungs, the variance-cadence layer). This doc references - §F for rung *definitions* and does not duplicate them; it states the fwdllm **build deltas**. -- [async_cifar10/UNAVAILABILITY_DESIGN.md](../async_cifar10/UNAVAILABILITY_DESIGN.md) -- the availability - substrate this reuses wholesale (ClientAvailability mixin, trace-read effect path, the two flag axes, - send-gate/deliver-late-stale, two-ledger discipline, starvation self-termination, the ground-truth - fidelity rungs A6/A7/A8/K11). fwdllm inherits it via the aggregator class chain (§A). - -**Starting reality (why this is real work, not a re-label):** FedFwd has **no simulated-clock support -today** -- every launcher yaml is `time_mode: real`, and the trainer emulates delay with a real -`time.sleep` (`trainer/forward_training/FedSgdTrainer.py:481` `_emulate_training_delay`, called at -`:523`; explicit note `trainer/main.py:57-78`). The aggregator machinery for a virtual clock is -*inherited but not driven* (§A). Adding the sim clock + wiring unavailability into the -variance-gated gradient loop is the build. - -**Decisions locked (kickoff):** -- A trusted **real** wall-clock run is the reference per baseline (as in async_cifar10). Parity != tuning - sim to real; real must pass `validate_real` admissibility first. -- The parity checker **extends the shared `async_cifar10/scripts/parity` engine** -- fwdllm rungs are - *added* (PARITY.md §F is their catalog), the engine is reused. -- The fluxtune / fwdllm / fwdllm++ -> config mapping is **resolved** from the landed launcher configs - (§C); it is no longer a blocking gap. +# High-Fidelity Simulator for FwdLLM — Real↔Sim Parity + +**Active build (branch `dg/fwdllm_sim_unavail`).** A simulated-clock runner for the `fwdllm` example +(FedFwd / forward-gradient FL) reaching **real↔sim parity** across the **fluxtune / fwdllm / fwdllm++** +baselines — at **100% availability (syn_0)** first (Phase 1), then **unavailability** (syn_20/50/mobiperf, +Phase 2), then **beyond syn_0** (Phase 3). fwdllm has no native sim clock; the build wires flame-core's virtual +clock + availability substrate into fwdllm's variance-gated gradient loop. It reuses async_cifar10's virtual +clock, sct reorder buffer, in-flight gate, availability substrate, and parity ladder where they transfer, and +deviates where the workload demands (fwdllm aggregates **gradients** not weights; **variance-gated dynamic-K** +commit cadence; **`data_id`** progress axis; one-message-per-call grad loop; rollback across agg-goal cycles). + +> ## PREAMBLE — how to maintain this doc (READ BEFORE EDITING) +> This is a **living status doc**, not a changelog. It must stay **rich but crisp**: enough to reconstruct *why* +> a decision was made, never a running history. Per-run history lives in git + the parity JSONs; the code is the +> source of truth for *what* the mechanism is. Every edit obeys: +> - **Current-truth only.** §A, the scoreboard, and the open-issues table describe the state **right now**. +> Rewrite them **in place** — never stack dated "UPDATE" blocks. When something closes, **delete it from the +> open list** and leave at most a one-line trace in §G (fixes) or §H (dead-ends). +> - **One line per landed item.** A fix that worked = one line in §G (**≤20 words problem + ≤20 words fix**). A +> belief that was refuted = one line in §H. A deviation = one line in §K (anchor + rationale). No paragraphs. +> - **Keep only what teaches.** Retain a big root-cause or a conceptual correction that would otherwise be +> re-litigated; **drop nitpicks** and mechanical scaffolding once landed (leave a pointer only if still +> referenced). If an entry no longer changes a future decision, delete it. +> - **No contradictions.** An issue is in exactly one place: OPEN (open table) **xor** CLOSED (§G/§H one-liner). +> Never both. Cross-link with §-anchors instead of restating. +> - **Decide forks explicitly.** When fwdllm diverges from async_cifar10, log the choice in §K and surface it in +> the §B.1 delta table — don't silently copy or silently invent. + +**Prerequisites:** [async_cifar10/PARITY.md](../async_cifar10/PARITY.md) — parity methodology (ladder, +roles/tiers/gating, run-length budget, landed sim mechanisms); **fwdllm's rung catalog is PARITY.md §F**. +[async_cifar10/UNAVAILABILITY_DESIGN.md](../async_cifar10/UNAVAILABILITY_DESIGN.md) — the availability substrate +fwdllm inherits via its aggregator class chain (ClientAvailability mixin, trace-read effect path, two-ledger +discipline, starvation self-termination, A6/A7/A8/K11 ground-truth rungs). --- -## §A Already on the class (inherited) vs. NOT yet wired - -The virtual clock is a **flame-core** capability. fwdllm's aggregator is -`examples/fwdllm/aggregator/FedSgdAggregator.py` -> `flame/mode/horizontal/syncfl/fwdllm_aggregator.py:: -TopAggregator(AsyncTopAgg)` -> `asyncfl.TopAggregator` -> `syncfl.TopAggregator(ClientAvailability, Role)`. -So both the vclock machinery **and** the availability substrate are already **on the instance** -- the -work is **wiring them into fwdllm's gradient loop + trainer**, not re-implementing them. - -### A.1 Aggregator spine (inherited, baseline-agnostic) -| Mechanism (from `asyncfl/top_aggregator.py`) | Status for fwdllm | -|---|---| -| Virtual clock `_vclock`, `_advance_sim_clock` (`vclock = max(vclock, sct)`) | inherited; **NOT yet driven** by the grad loop | -| §3.drain sct-ordered drain `_sim_recv_min` | inherited; **bypassed** (fwdllm's `_aggregate_grads_async` uses raw `recv_fifo`) | -| §3.resid one-in-flight residence `_sim_hold_busy_slots` | inherited; **bypassed** | -| §4.5/§4.9 sct-gated pool exclusion / carry-over (oort overlay) | available if a baseline selects via oort (fluxtune) | -| `ClientAvailability` substrate (trace-read gate, two ledgers, proactive evict, starvation advance) | inherited; **NOT yet wired** into the fwdllm grad loop | - -### A.2 Selector + duration signal (already wired) -- §S.dur duration helper `client_duration.py::real_client_task_train_duration` (intrinsic - `WALL_SEND_TS - WALL_RECV_TS`) is already called in `fwdllm_aggregator._process_single_trainer_message` - (`:745`) to set `PROP_CLIENT_TASK_TRAIN_DURATION`. No port needed. -- Oort selector (faithful §S.pacer two-branch pacer, §S.temporal UCB, D1-D8 fixes) is in - `flame/selector/oort.py`, usable verbatim by fluxtune (`async_oort`). -- Dynamic-KC controller/policy (`flame/selector/dynamic_kc_{controller,policy}.py`) is already integrated - (`_build_dynamic_kc_metrics`, `_dynamic_kc_controller.step()`); fluxtune leaves it **disabled** (fixed - K/C) per its config. See [docs/dynamic_kc_design.md](docs/dynamic_kc_design.md). -- Analysis: `scripts/analysis/analyze_run.py` is already fwdllm-aware via `telemetry_manifest.yaml` - (data_id/iteration_per_data_id progress hierarchy; per-category populate/partial status documented). - -### A.3 NOT reusable as-is -- the actual work -1. **Trainer sim path.** `FedSgdTrainer._emulate_training_delay` (`:481`) uses real `time.sleep`; it does - NOT stamp a modeled completion (`SIM_COMPLETION_TS` / `SIM_CLIENT_TASK_TRAIN_DURATION_S`) or - `WALL_SEND_TS`/`WALL_RECV_TS` the way async_cifar10's `trainer/pytorch/main.py` does. **Port required.** -2. **Aggregator grad loop is off the virtual clock.** `_aggregate_grads_async` (`fwdllm_aggregator.py:693`) - calls `channel.recv_fifo(...,1)` directly and commits on wall arrival -- never consults `_sim_recv_min`, - never advances `_vclock`, never holds in-flight slots. **The single biggest structural port.** -3. **Endogenous commit cadence** -- variance-gated dynamic-K, not fixed-K (PARITY.md §F.1): a new emergent - layer the async_cifar10 ladder does not model. -4. **Availability telemetry.** The fwdllm trainer emits **no `EVENT_AVAIL_CHANGE`** - (`telemetry_manifest.yaml`: availability = *partial*), so A6/A7/A8 ground-truth rungs can't run. - The `avail_change` / `agg_belief_change` / `send_gate_wait` builders exist - (`flame/telemetry/events.py:225/521`, `task_send` fields) and just need emitting from the fwdllm - trainer + aggregator. +## §A Current status + +> **⏸ SIM DEBUG PAUSED (2026-07-06) — real-experiment telemetry/impl next on this branch.** #15: the phantom +> commit-path stall is **FIXED + P3-validated** (`sim_compute_truthful_gate`; `STUCK_EVICT=0`, `phantom_skip` +> firing, 30s failsafes gone). `sim_rate` is still <1 for a **NEW, non-drain-gate** reason — (a) vclock omits +> `aggregate()` variance-compute wall (sim 32s / real 29s, symmetric, uncredited) + (b) GPU≈D no-headroom (#1d). +> So the "sim_rate=0.50 — COMMIT-PATH STALL / phantom drain gate" reading in the scoreboard + prose below is the +> **fixed** part; live residual + resume plan in [PARITY_LOGICAL_TASKS.md](PARITY_LOGICAL_TASKS.md) (PAUSED +> checkpoint at top + #15). Don't re-debug the sim until the real-experiment work lands. + +**All three baselines run end-to-end (#14/#1c/#13 fixed). SYNC baselines are good: with full registry delay +(`--delay-factor 1`) `sim_rate` is 2.9–3.0 (#12c RESOLVED for sync) and K-D31 makes bin-1 cohort order BIT-EXACT +(P2-7a validated). Two live fronts, both root-understood: (1) a SYNC float-nondeterminism wall at ~bin 7 (grads +not bit-reproducible → exact cadence parity unattainable past bin ~6 → parity target beyond bin 1 is +DISTRIBUTIONAL); (2) fluxtune `sim_rate = 0.50` — a COMMIT-PATH STALL: hold-to-commit is a CORRECTNESS check (a +trainer is freed only when its update commits), so commit RATE sets throughput; the sim's drain blocks real wall +on PHANTOM `_sim_inflight_expected` entries (30s failsafe) → correctly-held trainers idle far longer than their +GPU pass → 1.54× concurrency vs real's 3.37×. NOT a gate bug (felix's gate is INERT), NOT over-restrictive +hold-to-commit, NOT GPU under-provisioning (§H). Fix = fast/non-stalling commit path. Grounding + fix in +[PARITY_LOGICAL_TASKS.md](PARITY_LOGICAL_TASKS.md) (FELIX GROUNDING + #15).** + +Latest FULL pairs (`run_sequential.sh --mode both --delays on --delay-factor 1 --max-runtime-s 2700`, +`run_20260705_1924 → 2046`): + +| baseline | sim `sim_rate` | vclock / wall | verdict | +|---|---|---|---| +| **fwdllm** (sync) | **2.93** ✓ | 1139s / 389s | clean; sim ~3× faster than the virtual time it models | +| **fwdllm_plus** (sync) | **3.01** ✓ | 1137s / 378s | clean; sim ~3× faster | +| **fluxtune** (async) | **0.50** ⛔ | 1191s / 2425s | `sim_rate<1` = a commit-path stall (phantom drain gate, #15), not #12c; hold-to-commit is correct | + +**Why fluxtune `sim_rate = 0.50` — objective real↔sim telemetry (supersedes the #12c-delay-headroom reading).** +Real and sim do the SAME GPU work (~3.7–4.0k trainer-s) with the same ~480s 8-way pipeline floor. Real packs it +into 1210s wall at **3.37 trainers on the GPUs at once (98% busy)**; the sim takes **2425s** at **1.54× +concurrency (85% busy)**, converting real's correctly-skipped device-delay waits (4775s) into **19807s of +trainer idle in `recv`** (vs real 380s). Per-commit: real reaches agg-goal in **4.30s wall**; the sim models +**3.27s vclock** but spends **6.60s wall** → sim_rate 0.50. **ROOT (corrected 2026-07-06):** hold-to-commit is +a CORRECTNESS check — a trainer is freed only when committed (guards: no same-version re-dispatch, none while +computing, none while returned-but-uncommitted), so commit RATE sets throughput. In felix/cifar commit is +instant (~0.4s) so held trainers barely idle; fluxtune's commit path STALLS — the drain's `earlier_stuck` gate +blocks real wall on a PHANTOM `_sim_inflight_expected` entry (a trainer stamped expected-at-dispatch that isn't +computing, e.g. still waiting for weights the gate-blocked aggregator can't send) → 30s failsafe → correctly-held +trainers idle 30s instead of ~one GPU pass. Fix: fast/non-stalling commit path (gate waits only on a +genuinely-computing trainer); **hold-to-commit stays**. Grounding: +[PARITY_LOGICAL_TASKS.md](PARITY_LOGICAL_TASKS.md) FELIX GROUNDING F5/F6 + D1-D3. + +### Parity scoreboard — REFERENCE baseline (checker run on the pairs above; `expt_scripts/run_parity.py --yes`) +*The numbers held against until the open issues resolve. The stale pre-rewrite counts (49/4/21 etc.) are retired +— do NOT reference them. Below are the ENFORCED counts AFTER the logical-parity rungs landed (`cohort_sequence` +EXACT + V2 mean-guard, [PARITY_LOGICAL_TASKS.md](PARITY_LOGICAL_TASKS.md) P1) — these two now correctly FAIL the +order/grad divergence that used to be invisible.* + +| baseline | pass / fail / skip | JSON | +|---|---|---| +| **fwdllm/syn_0** | **41 / 13 / 22** | `experiments/_parity_reports/parity_fwdllm_syn_0_20260705_045724.json` | +| **fwdllm_plus/syn_0** | **36 / 17 / 22** | `parity_fwdllm_plus_syn_0_20260705_063350.json` | +| **fluxtune/syn_0** | **35 / 19 / 20** | `parity_fluxtune_syn_0_20260705_080931.json` | + +*(skip +1 each vs the prior ref = the new `timing_overrun` DIAG rung, which SKIPs on these banked logs — they +predate its P2-6 `training_overran` telemetry; it populates on the databin1 run. Pass/fail unchanged: the K-D30 +full-cohort gate flipped fwdllm's `selection`/`aggregation_sequence`/`utility` from a trivial GATED pass to GENUINE +enforcement without moving the counts — they pass truthfully; `utility` still fails on a pre-existing pooled KS.)* + +**Fails, categorized by blast radius (fix the SHARED roots first — principle #14).** +- **SHARED — all 3 (top priority):** `throughput` / `overhead_residual` / `per_round_advance` (the `sim_rate<1` + clock-rate family → #12c); `gpu_budget_real` / `gpu_budget_sim` (mis-applied invariant, symmetric real/sim + overrun — likely a rung-definition bug, not a divergence); `v1_iter_per_data_id` / `v5_variance_pass_ratio` / + `g2_grad_pool_size` / `utility` / `convergence` (cadence/emergent). **The bin-1 logical check (below) splits + these:** length-confound for fwdllm/fwdllm_plus (cadence identical on bin 1), but GENUINE for fluxtune (#1d — the + async cohort/cadence diverges even on bin 1). +- **fwdllm only:** `phase_gpu_compute` (marginal KS, means match). +- **fwdllm_plus only:** `eligibility` / `avail_timebase` / `selection_detail` (the #7 selection divergence); + `terminal_state` / `total_commits` (stop-mismatch from unmatched length). +- **fluxtune only:** `selection_detail` / `phase_gpu_compute` / `phase_weights_to_gpu`; `staleness`; + `terminal_state` / `total_commits` / `convergence_loss` (stop-mismatch + length). + +### STRATEGY — nail first-data-bin logical parity before any longer run +Prove parity by **logical determinism, not aggregate curve-matching**: for a matched scope the sim must take +**the same sequence of steps in the same order** as real — same trainers selected, same update-receipt order, +same aggregations/rollbacks — differing ONLY in wall-clock. **Scope to the first 1 data bin** (`--max-data-id 1` +/ `--max-bin 1`): bin 0 already contains many iterations, variance passes, and aggregations, so it exercises the +full cadence machinery while staying short, matched, and diffable. Prove bin 0/1 parity FIRST (from banked +telemetry — no re-run); only then chase the time dimension. This isolates length-confound from genuine logic bugs. + +### Logical-parity check (TIME-STRIPPED, all available bins) — CURRENT REFERENCE +Tool: `expt_scripts/logical_parity.py [--max-bin N]` — diffs the `agg_round` event stream (data_id, iteration, +receive-ordered contributors, variance decision) real vs sim with every timestamp removed. Real receive-order is +DETERMINISTIC in both real and sim by design (operator-confirmed) → exact match is the correct target. + +| baseline | receive-SET (to real's max bin) | cadence | logical parity HOLDS TO | verdict | +|---|---|---|---|---| +| **fwdllm** | 41/41 identical (K=10=all) | 22/41 | **data_id 6** — breaks at **7** | **BREAKS @ bin 7** — but receive-ORDER now 41/41 (K-D31) → the break is NOT order | +| **fwdllm_plus** | 14/14 identical | **14/14 identical** | **real's max (~13)** — no break | **PARITY ✓** over its (shorter) real run | +| **fluxtune** | **3/272 identical** | 29/272 | **breaks at bin 0** | thin-margin overrun (§ #1d) — GPU tail > min D on the contended-doubled GPUs | + +**ROOT — CORRECTED (2026-07-05 full-run investigation; supersedes the "order → var → RNG-desync" reading for the +SYNC full run).** With K-D31, fwdllm's receive-ORDER is **41/41 identical** on the full run — yet cadence still +breaks at data_id 7. So order is NOT the sync full-run cause. The residual is **grad non-reproducibility given +matched order**: `|Δvar|` ~1e-3 through bin 6 with *every* `var_good`/force decision matching, then at bin 7 +(all-10 cohort, identical order + RNG) grads diverge ~1e-3 (GPU fp16 non-reproducibility), which the split-half +variance — a ratio with a near-zero denominator at a bin's first iteration — **amplifies to a 0.26 var swing**, +flipping the `var<0.3` gate at (7,2). Both modes confirmed `jvp_perf_opt=False` (no config skew). This **answers +P0-2 empirically: grads are NOT bit-reproducible → exact cadence parity is unattainable past ~bin 6.** ⇒ **Parity +target (operator decision): cohort SET = HARD; `var_good`/cadence = HARD to bin 1, DISTRIBUTIONAL beyond; `var` +VALUE = SOFT (tolerance/KS); receive-ORDER within a set = SOFT for sync (fedavg order-invariant + K-D31 +canonicalizes). Keep `cohort_sequence` EXACT scoped to `--max-bin 1`; add a distributional cadence/var rung for +the full run.** For **fluxtune** the SET still genuinely diverges (#1d) — a timing/pipelining cause, not +nondeterminism. Full diagnosis: [PARITY_LOGICAL_TASKS.md](PARITY_LOGICAL_TASKS.md). + +> **STATUS (2026-07-05 — P2-7a VALIDATED + full-run bin-7 wall found).** Databin1 checks +> (`parity_fwdllm_syn_0_165256`, `_plus_165736`, `--max-bin 1`): sync `cohort_sequence` **ok=true, +> set/order/var/cadence = 1.0** — K-D31 closed the benign delay-tie; P2-7a DONE. Full runs exposed the bin-7 +> float-nondeterminism wall above (order matches 41/41, cadence still breaks) → the exact-cadence target is valid +> only ≤bin 1; beyond it must be distributional (pending a two-real-run P0-2 confirmation — strong single-run +> evidence already). **fluxtune (`--delay-factor 1` full run):** `jvp_perf_opt=True` cut GPU 7.57→3.61s MEAN +> (under the 4.0s min budget) but the TAIL (4.1–5.4s) still overruns on the two doubled GPUs (10 trainers / 8 +> GPUs) → `set_match=3/272`. A thin-margin contention effect, not the JVP algorithm. + +### Open issues (OPEN only — closed items live in §G/§H) +| # | issue | baseline(s) | next step | +|---|---|---|---| +| **#15** ⭐⭐ (⏸ PAUSED) | **fluxtune `sim_rate<1` — phantom COMMIT-PATH STALL FIXED (`sim_compute_truthful_gate`, P3 VALIDATED 2026-07-06: `STUCK_EVICT=0`, `phantom_skip`→65 ~1/commit, 30s failsafes gone). But `sim_rate` still <1** (steady ~0.49; per-commit `Δvclock/Δwall=0.465`). **Residual root (NEW, non-drain-gate):** (a) the vclock OMITS `aggregate()` variance-compute wall — sim 32s / real 29s, symmetric, uncredited (`fwdllm_aggregator.py:1783` folds eval only); (b) GPU≈D no-headroom (`max(gpu,D)≈gpu`, sim blocks in `drain_ready` on the real GPU pass while `buf_depth=6` waits) = #1d. Net: sim wall 145s > real 91s. Full evidence+plan: PARITY_LOGICAL_TASKS.md top (PAUSED checkpoint) + #15. | fluxtune | **PAUSED** (operator doing real-experiment telemetry/impl first). On resume: (1) fold the *non-overlapped* `aggregate()` time into the vclock (measure GPU-overlap first; flag-gated, operator sign-off); (2) GPU-vs-D headroom — 1 trainer/GPU or `--delay-factor` up; (3) longer flag-on run + flag-off A/B for parity. | +| **#N (bin-7 nondeterminism)** ⭐ | **SYNC exact-cadence parity has a float-nondeterminism wall at ~bin 7.** Order matches 41/41 (K-D31) yet cadence breaks: ~1e-3 GPU fp16 grad jitter, amplified by the split-half variance ratio, flips the `var<0.3` gate at (7,2). Not a sim bug. | fwdllm (fwdllm_plus latent) | Confirm with a 2-real-run diff (P0-2). Then land the parity-target relaxation: `cohort_sequence` EXACT scoped to `--max-bin 1`; distributional cadence/var rung (mean-band + KS + `var_good` fraction) for the full run. | +| **#1d** ⭐ | **fluxtune cohort SET diverges (thin-margin overrun).** `set_match=3/272`. `jvp_perf_opt` cut GPU to 3.61s MEAN (<4.0s budget) but the tail (4.1–5.4s) still overruns on the two GPUs that carry 2 trainers each (10/8) + the aggregator's eval GPU. Order flips → wrong 3-of-K commit. | fluxtune | Aggregator-GPU pin landed (K-D33); with #15's pipelining the compute drops toward the ~2.4s uncontended floor (<4.0s). If a residual tail remains: `perturbation_count`↓ (P2-5) or 1-trainer/GPU. | +| **#7** | fwdllm_plus real ~4× slower/round; at syn_0 real sees only ~4.9 eligible vs sim ~9.6. Not a sim bug. | fwdllm_plus | Profile per-iteration reselection + oracular-read cost from the banked per-phase log; explain the eligible-count gap at 100% avail. | +| **#11** | real-mode critical-path waste (`sleep(0.1)` MQTT-settle busy-waits; one-grad-per-poll drain tail) — real-only. | fwdllm, fwdllm_plus (real) | **Deferred to a validated pass** — ZERO parity impact (sim already skips them); removing them changes the working real reference + needs a real run (principle #8/#11c). | + +### Next roots — ranked (correctness before time; SHARED before per-baseline — principle #14) +1. **#15 fluxtune — make the commit path fast/non-stalling (time, per-baseline) — TOP, IN PROGRESS.** The one + thing keeping fluxtune `sim_rate<1`. Fix the phantom `_sim_inflight_expected` so the drain gate never blocks + real wall on a non-computing trainer → commits flow → held trainers freed. Hold-to-commit STAYS. Diagnose + D1/D2 first. Also lifts #1d. +2. **bin-7 nondeterminism → relax the parity target (SHARED, correctness-of-CHECK).** Land the distributional + cadence/var rung + scope `cohort_sequence` EXACT to bin 1, after P0-2 two-real-run confirmation. Un-reds the + sync full-run cadence fails that are float-noise, not bugs. +3. **#7 fwdllm_plus real speed / selection divergence (per-baseline).** Not a sim bug; explain from telemetry. +4. Then C1/C2 convergence (distributional target) at matched `data_id` per baseline → gate to Phase 2. + +### SKIP audit (20–22 skips; ~17 legit) +Legit at Phase-1 syn_0 + `random` selector: 7 availability ground-truth rungs + 4 delivery/withheld (Phase-2 +effect path, not built) + 3 DynamicKC (disabled by design) + 2 oort-only (`random` baselines) + `residence` +(async telemetry) + `timing_overrun` (banked logs predate its P2-6 telemetry; populates on the next run). The 12 +former rigor-gap skips (4 advance, 8 phase-timing) are un-skipped (K-D21). --- -## §B How FwdLLM differs structurally - -Full detail (the crux that makes the ladder applicable, and the variance-cadence rung layer it forces) is -**PARITY.md §F.1**. In one line: fwdllm aggregates **gradients** (JVPs) not weights, its commit cadence is -**endogenous** (variance-gated dynamic-K), and its progress axis is **`data_id`** (committed variance -passes), not raw update count. Gradient *values* are mode-invariant given identical input+perturbation -seed, so fwdllm parity reduces to clock+selection+ordering parity **plus** a new variance-cadence layer. - -Concrete anchors (`flame/mode/horizontal/syncfl/fwdllm_aggregator.py`): `aggregate()` var gate + -rollback/`cached_v` at each `_agg_goal` boundary; `total_data_bins=150` (`:271`); force-commit cap -`_max_iter_per_data_id` (`:276`, config key `max_iterations_per_data_id`); `_reselect_each_iteration` -(`:295`, per-iteration reselection for fwdllm++); sync path `_aggregate_grads_sync` (`:1354`). +## §B How fwdllm differs structurally + +fwdllm aggregates **gradients** (JVPs) not weights; commit cadence is **endogenous** (variance-gated dynamic-K); +progress axis is committed **`data_id`** (variance passes), not update count. Gradient values are mode-invariant +given identical input+perturbation seed, so parity reduces to **clock + selection + ordering parity plus a +variance-cadence layer**. Anchors (`flame/mode/horizontal/syncfl/fwdllm_aggregator.py`): `aggregate()` var gate + +rollback/`cached_v` at each `_agg_goal` boundary; `total_data_bins=150`; force-commit cap +`max_iterations_per_data_id`; `reselect_each_iteration` (fwdllm++ per-iteration reselection); sync path +`_aggregate_grads_sync`. Full detail: PARITY.md §F.1. + +### §B.1 Real↔sim design deltas vs async_cifar10 — **CURATED, KEEP CURRENT** +Separates an **intentional fwdllm choice** from an accidental discrepancy. The §K column points at the one-line +rationale. + +| # | Axis | async_cifar10 | fwdllm | Why | §K | +|---|---|---|---|---|---| +| 1 | Aggregated object | model **weights** | **gradients** (JVPs) | grad values mode-invariant → parity = clock+order+selection + variance-cadence | §F.1 | +| 2 | Progress axis | update/round count | committed **`data_id`** | cadence (updates-per-data_id) is an **output to match**, not an input | principle #2 | +| 3 | Commit cadence | fixed `agg_goal` | endogenous **variance-gated dynamic-K** | the emergent layer cifar doesn't model; V/DK/G rungs verify it | §F.1 | +| 4 | sct delay model | `send + max(gpu, D)` | `send + max(gpu, D)` (**remainder-wait, was additive**) | K-D29: real now sleeps `max(0,D−gpu)` (device wall = D, GPU hidden), so update order = per-trainer D order = deterministic & real↔sim identical. Reverses K-D2 | K-D2/**K-D29** | +| 5 | Per-eval sct | distinct, ~20× faster | **collapses to train sct** | eval lives on the aggregator; forward-grad "train" IS a forward pass (no 20× factor) | K-D3 | +| 6 | Slot residence | per-commit release | **hold slot to COMMIT** (felix port, K-D17b) — a trainer is freed only when its update commits (guards: no same-version dispatch, none while computing, none while returned-but-uncommitted) | Correct for BOTH sync & async — it is a correctness check, not a lever. fluxtune's `sim_rate<1` is a COMMIT-PATH STALL (K-D34/#15), NOT this rule. Commit RATE = throughput | K-D17b/**K-D34**, principle #4 | +| 7 | Surplus grad on rollback | carried | **carried** for async (fluxtune, c≫agg_goal); **drop** stays correct for sync (c≈agg_goal) | drop was benign only for sync; fluxtune dropped ~7/cycle → 2× passes | K-D12 | +| 8 | Async drain primitive | `_sim_recv_min` verbatim | purpose-built `_sim_recv_min_grad` / sync `_sync_sim_recv_first_k` | cifar's per-commit release + withheld paths key on WEIGHTS semantics | K-D4 | +| 9 | `time_mode` default | `"simulated"` | `"real"` (getattr fallback) | fwdllm's whole config corpus is `real`; a `simulated` default risks half-activating an unbuilt path | K-D1 | +| 10 | Cadence telemetry | n/a | **pre-mutation** cycle snapshot (`cycle_data_id`/`cycle_iteration`/`grad_pool_size`/`cached_v_size`) | post-mutation `data_id` advances before emit → off-by-one; snapshot makes V1 exact | K-D9 | +| 11 | Availability tracking (v1) | all `trace_read` | **mixed**: fwdllm unaware, fwdllm_plus `oracular`, fluxtune `client_notify`→approx `trace_read` | baselines carry different models; first-class `client_notify` deferred to Phase 2 | D1 | +| 12 | Launch tooling | single parity template (`debug_run.sh`) | per-baseline yamls + `_sim` files (`run_sequential.sh`) | different config models; both source the shared harness `examples/scripts/expt_runner.{sh,py}` | this work | --- -## §C Baseline matrix -- RESOLVED (from the landed launcher configs) - -Read from `expt_scripts/{fluxtune,fwdllm,fwdllm_plus}_n10_smoke.yaml`. Canonical copy + taxonomy mapping -is **PARITY.md §F.2**; repeated here for the build plan. - -| baseline | sync/async | selector | agg | tracking_mode / avail | reselection | K (agg_goal) | dynamic_kc | +## §C Baseline matrix (resolved from the launcher configs) +| baseline | sync/async | selector | agg | tracking / avail | reselection | K (agg_goal) | dynamic_kc | |---|---|---|---|---|---|---|---| -| **fluxtune** | async | `async_oort` | fedbuff (+server LR, JVP) | `client_notify` (3-tier, mobiperf_3st_50) | -- | 3 | disabled (fixed K/C) | -| **fwdllm** | sync | `random` | fedavg | `default` (unaware) | per-round | 10 (=c; all selected required) | -- | -| **fwdllm_plus** | sync | `random` | fedavg | `oracular` (`_metadata`, mobiperf_2st) | per-iteration (`reselect_each_iteration=True`) | 2 | -- | - -**Availability-taxonomy mapping** (UNAVAILABILITY_DESIGN.md two axes): -- **fwdllm** -- unaware at selection (`avail_select_filter=False`), reactive-90s in-flight - (`proactive_inflight_evict=False`). -- **fwdllm_plus** -- aware at selection via `oracular` trace read (`avail_select_filter=True`), - reactive-90s in-flight. -- **fluxtune** -- async, aware via `client_notify` (message-transport), reactive-90s in-flight. - `client_notify` is async_cifar10's **deferred Stage-H** tracking model -> **decision D1**. - -This decides: Stage 3 oort rungs run **only** for fluxtune; Stage 5 sync-barrier rungs run for -**fwdllm/fwdllm_plus**; V1/V4 distributions are shaped by each baseline's `max_iterations_per_data_id` / -`var_threshold`. Parity is staged **one baseline at a time** (PARITY.md workflow rule). +| **fluxtune** | async | `async_oort` | fedbuff (+server LR, JVP) | `client_notify`→`trace_read` v1 (aware-at-select, reactive-90s) | — | 3 | disabled | +| **fwdllm** | sync | `random` | fedavg | `default` unaware (reactive-90s) | per-round | 10 (=c) | — | +| **fwdllm_plus** | sync | `random` | fedavg | `oracular` (aware-at-select via trace read) | per-iteration | 10 (yaml is source of truth) | — | + +Phase order (locked): Phase-1 syn_0 → Phase-2 unavailability → Phase-3 beyond syn_0. One baseline at a time. +Stage 3 oort rungs run **only** for fluxtune; sync-barrier rungs run for **fwdllm/fwdllm_plus**. --- -## §D The parity ladder for fwdllm +## §D Parity ladder for fwdllm (rung defs live in PARITY.md §F) +- **Reused verbatim:** Stage 0 TC1/K10; Stage 1 P3/K1/K7; Stage 2 A1–A3 (+A6/A7/A8 once avail telemetry lands); + Stage 3 S3/4, A2c + the oort stack (fluxtune only); Stage 4 T2/K6/T_mqtt; Stage 8 C1; Stage 9 budget/stop. +- **Modified** for variance-gated dynamic-K: K3a/K3b/K2/U3/K8/U2 → re-keyed to the variance-pass boundary / + committed `data_id` (PARITY.md §F.3). +- **New** variance-cadence layer: V1–V5, DK1–DK3, G1–G2 (PARITY.md §F.4). Localize down; never fix an EMERGENT + rung directly. `var_threshold` / `max_iterations_per_data_id` are baseline knobs, not parity levers. +- **Availability** rungs (A1–A5, A6/A7/A8/K11, withheld/abandon/starvation) inherited; apply once Phase-2 wires + the effect path + telemetry. + +--- -**Rung definitions live in PARITY.md §F** (modified rungs §F.3; variance-cadence / dynamic-K/C / -forward-grad rungs §F.4). This section states only which rungs apply and how they gate. +## §E Roadmap — remaining phases -- **Reused verbatim** (PARITY.md §2): Stage 0 TC1/K10; Stage 1 P3/K1/K7; Stage 2 A1/A2 (+ A3 time-base - CONTROL, + A6/A7/A8 ground-truth once telemetry is ported); Stage 3 S3/4, A2c and the oort stack - (Sx/Sd/S2) **only for fluxtune**; Stage 4 T2, per-phase timing, K6, T_mqtt; Stage 8 C1; Stage 9 budget/stop. -- **Modified** for variance-gated dynamic-K: K3a/K3b/K2/U3/K8/U2 -> PARITY.md §F.3 (all re-keyed to the - **variance-pass boundary** / committed **data_id**). -- **New** -- the variance-cadence layer (the prize): V1-V5 (Stage 6'), DK1-DK3 (Stage 3'), G1-G2 (Stage 7') - -> PARITY.md §F.4. Localize down, never fix an EMERGENT rung directly; `var_threshold` / - `max_iterations_per_data_id` are config knobs, not parity levers. -- **Availability** rungs (A1-A5, A6/A7/A8/K11, withheld_delivery, abandon_timeout, starvation_advance, - eligible_pool_reduction) are inherited from the async_cifar10 substrate and apply once §Stage-Avail - (below) wires the effect path + telemetry. +**Phase 1 (syn_0) — CLOSE-OUT (near done):** #14/#1c/#13 fixed; the drain runs at speed. Remaining: (1) **#12c +`--delay-factor` run** to drive `sim_rate ≥ 1`, then re-run the checker + bank the three pairs; (2) **#7** +fwdllm_plus real slowness + selection divergence from telemetry; (3) C1/C2 convergence at matched `data_id` per +baseline → gate to Phase 2. ---- +**Phase 2 — unavailability (syn_20/50/mobiperf).** Wire the ClientAvailability effect path into the grad loop: +send-time gate (real) / `delivery_ts = max(sct, next_avail)` buffering (sim); two ledgers; reactive-90s in-flight; +starvation vclock-advance; per-baseline `avail_select_filter`. Emit `EVENT_AVAIL_CHANGE` + +`agg_belief_change`/`send_gate_wait` (builders in `flame/telemetry/events.py`) → unlocks A6/A7/A8/K11. **D3 watch:** +does a withheld/late grad roll into `cached_v` on rollback, or a stale-version late grad inflate `var` (thus +dynamic-K)? Over-instrument first. Exit: A1–A5 + A6/A7/A8 PASS; self-stops; withheld grads delivered not dropped. -## §E Staged implementation + validation - -Each stage: config-gated (flag-off => byte-identical), test-guarded, exit criteria + min run length -(PARITY.md budget table). Land one mechanism per run round. **Smoke (5 min) before any multi-hour run.** - -### Stage 0 -- Telemetry coverage + real admissibility (gate) -Confirm fwdllm emits every field the ladder reads: `vclock_now`, `sim_completion_ts`/`sct`, -`model_version`, `data_id`, `iteration_per_data_id`, `var`, `var_threshold`, `_agg_goal`(K), -`dynamic_c`(C), `agg_goal_count`, per-commit `train_duration`, `WALL_SEND_TS`/`WALL_RECV_TS`, grad/JVP -norm, **and `avail_change`/`agg_belief_change`/`send_gate_wait`** (§A.3 item 4). Implement **TC1** for -fwdllm; **partition eval vs train commits** on the variance-pass eval (async_cifar10 dead-end: eval commits -polluting `agg_rounds` broke monotone/staleness). Run `validate_real` admissibility. **Min:** 5-10 min -smoke. **Exit:** TC1 all-present both modes; real admissible; report JSON no SKIPs. - -### Stage 1 -- Trainer sim path (modeled completion, no wall sleep) -Add a `simulated` mode to `FedSgdTrainer`: replace `_emulate_training_delay` (real `time.sleep`) with -stamping `_sim_completion_ts = sim_send_ts + max(real_gpu_s, modeled_budget_D) + sim_completion_leg_s`; -send `SIM_COMPLETION_TS` + `SIM_CLIENT_TASK_TRAIN_DURATION_S` + `WALL_SEND_TS`/`WALL_RECV_TS`. Stamp a -**per-eval** `sct` (dead-end: reusing the last train `sct` past-dates every eval); eval delay ~= train -(forward pass), **not** async_cifar10's 20x speedup -- **decision D4**. **Min:** 10 min. **Exit:** P3 matches; -K6 advancing; T2 matched; no real `time.sleep` on the sim path. - -### Stage 2 -- Drive the virtual clock from the grad loop (THE structural port) -Replace the raw `channel.recv_fifo(...,1)` in `_aggregate_grads_async` (`:693`) with the inherited -**`_sim_recv_min`** path (sct-ordered reorder buffer + §3.drain) so a fast trainer's grad buffers as a -future and grads are consumed in completion order; wire **`_sim_hold_busy_slots`** (§3.resid) for -one-in-flight-per-trainer. **fwdllm processes one message per `_aggregate_grads_async` call**, so key the -slot release on the **agg-goal boundary** (where `_per_agg_trainer_list` clears), NOT per-message -- the -hold/release points differ from asyncfl's batch `_aggregate_weights`. Keep `simCommitOverheadSeconds=0` -(overhead on the vclock is a hard dead-end). **Risk:** a data_id can span **many** agg-goal cycles -(rollback path) -- slot-hold + sct-buffer must survive rollbacks without leaking or double-committing a -grad; over-instrument `inflight_residence` + per-cycle buffer occupancy. **Min:** 45 min. **Exit:** K1 -monotone; K3a/K3b (per variance-pass) ~= 0 residual; U5 inter-arrival; one-in-flight overlap ~= real; -no past-dating (commit_gap ~= 0). - -### Stage Avail -- Wire unavailability into the grad loop (syn_20/50/mobiperf) -Wire the inherited `ClientAvailability` effect path into fwdllm: send-time gate (real) / -`delivery_ts = max(sct, next_avail)` buffering (sim); the two ledgers; per-baseline in-flight timing -(reactive-90s for all three fwdllm baselines; no proactive-evict baseline here); starvation vclock-advance -under scarcity; and per-baseline `avail_select_filter` (fwdllm off, fwdllm_plus/fluxtune on). Emit the -availability telemetry (Stage 0) so A6/A7/A8/K11 run. **Interaction risk (decision D3):** a withheld/late -grad meets the variance gate -- does it roll into `cached_v` on a rollback, and does a late grad against an -old `model_version` inflate the variance signal? This is genuinely new vs async_cifar10 (which commits -weights, not a variance-gated pool). **Min:** 45-90 min. **Exit:** A1-A5 + A6/A7/A8 PASS; withheld-then- -delivered (not dropped); self-stops (`"stopping run"`, no `SIM_WALL_CEILING`); syn_0 byte-identical gate -ON vs OFF. - -### Stage 3 -- Variance-cadence parity (the fwdllm prize) -Implement V1-V5, G1-G2, DK1-DK3 (PARITY.md §F.4). Likely roots (confirm via the lowest broken rung): -contributing-set/order divergence (U5/S2 upstream) -> V1 -> K2; grad-pool accumulation order -(`cached_v` carry-over, `grad_pool.append` order) -> V2 with matched inputs = a true sim bug; force-commit -rate (V4) = chronic variance divergence, not a separate bug. **DynamicKC coupling:** validate **DK3** -(policy *input*) before DK1/DK2 (fluxtune leaves dynamic_kc off, so DK is inert there; relevant only if a -baseline enables it). **Min:** 90 min - 2 h (variance-feedback-compounding). **Exit:** V1/V2/V5 PASS; K2 -(committed-data_ids/vsec) PASS; DK tracks if enabled. - -### Stage 4 -- Selection fidelity (fluxtune only, `async_oort`) -Validate A2c/Sx/Sd/S2 + the §S.pacer/§S.temporal/§S.dur stack (all landed in `flame/selector/oort.py`). -For the two `random`-selector baselines this stage is inert (S-rungs WARN/skip). **Min:** 45 min - 3 h (a -selection-mix residual can surface late). **Exit:** A2c PASS, Sd binding real~=sim, no pacer ratchet. - -### Stage 5 -- Sync-path parity (fwdllm / fwdllm_plus, `_aggregate_grads_sync`) -Give the sync grad barrier the barrier-anchored visibility-lag treatment (§6.u6). fluxtune streams (async), -commits each grad at its own sct -> leave per-message correct. **Min:** 45 min. **Exit:** U6 barrier- -anchored lag real~=sim; no per-message past-dating. - -### Stage 6 -- Convergence sign-off -C1 accuracy, C2 loss, K8/U2 terminal-state @ matched **data_id**. **Min:** full (3-4 h+). **Exit:** curves -within tolerance at matched data_id; K8/U2 rel within bar. +**Phase 3 — beyond syn_0.** Full ladder under scarcity; bin V1/V2 by run-fraction to separate a constant mix bias +from a compounding variance-feedback loop. Exit: curves within tolerance at matched `data_id`; K8/U2 within bar; +V1/V2 binned residual flat. --- -## §F Key design decisions (open -- resolved at implementation) - -- **D1 -- tracking-mode strategy.** fluxtune's landed config uses `tracking_mode: client_notify`, which - UNAVAILABILITY_DESIGN.md treats as the **deferred Stage-H** message-transport model (async_cifar10 v1 is - all `trace_read`). Options: (a) implement `client_notify` first-class for fwdllm now (larger scope, but - it's baked into fluxtune's config and can't just be ignored), or (b) map fluxtune onto `trace_read` for a - v1 parity pass and treat `client_notify` as Stage H. **Lean (a)** -- fluxtune cannot run its intended - taxonomy without it. fwdllm_plus's `oracular` and fwdllm's unaware paths are already `trace_read`-shaped. -- **D2 -- availability telemetry port.** The fwdllm trainer emits no `EVENT_AVAIL_CHANGE`; the aggregator - emits no `agg_belief_change`/`send_gate_wait`. All three builders exist in `flame/telemetry/events.py` - -- port emission (trainer state machine + aggregator belief hooks) so A6/A7/A8/K11 light up. Prereq for - Stage Avail's exit. -- **D3 -- variance-cadence x withheld/late grads.** Does a withheld grad roll into `cached_v` on rollback? - Does a late grad against a stale `model_version` distort the `var` signal (and thus dynamic-K)? New to - fwdllm; over-instrument in Stage Avail before Stage 3. -- **D4 -- eval-delay factor.** FedFwd eval is a forward pass (~= train cost), unlike async_cifar10's - ~20x-faster eval. Confirm the factor per baseline before Stage 1's per-eval `sct` stamp. - -### Locked principles (from async_cifar10, carried over) +## §F Locked principles (from async_cifar10, carried over) 1. **Sim does real forward-grad compute, charges modeled time.** GPU runs the real JVP; the agg stamps - `sct = sim_send_ts + max(real_gpu_s, D) + leg`. Makes gradient values mode-invariant. - **Never** put overhead on the virtual clock (`vclock = max(vclock, sct)` only). -2. **Progress axis is `data_id`** (committed variance passes). Update count per data_id is the dynamic-K - random variable V1 validates -- an output to match, not an input to assume. -3. **Variance is an emergent gate; localize, never tune it.** V-rungs gate on U5/S2/V1. `var_threshold` is - a baseline-defining config knob. -4. **Slot residence must survive variance rollbacks.** Release keys on the agg-goal boundary; the sct - reorder buffer must not strand a grad across a rollback (the fwdllm-specific extension of §3.resid). -5. **DynamicKC: validate the input before the policy.** DK3 (CONTROL) before DK1/DK2 (MECHANISM). The - controller is shared/baseline-agnostic -- do not fork it per baseline. -6. **Real is the reference only after admissibility.** A real<->sim gap has two fix directions; if a - selection-mix residual appears, check whether the **real** input is the divergent side. -7. **syn_0 byte-identical gate OFF; unavailability config-gated.** Reuse the async_cifar10 flag axes and - the `simUnavailability` gate; do not delete eligibility plumbing. + `sct = sim_send_ts + max(real_gpu_s, D) + leg`. **Never** put overhead on the vclock (`vclock = max(vclock, sct)`). +2. **Progress axis is `data_id`.** Updates-per-data_id is the dynamic-K random variable V1 validates — an output + to match, not an input to assume. +3. **Variance is an emergent gate; localize, never tune it.** `var_threshold` / `max_iterations_per_data_id` are + baseline-defining config knobs. +4. **Slot residence must survive variance rollbacks.** Release keys on the agg-goal boundary; the sct reorder + buffer must not strand a grad across a rollback. +5. **DynamicKC: validate the input before the policy** (DK3 CONTROL before DK1/DK2). The controller is shared; + don't fork it per baseline. +6. **Real is the reference only after admissibility.** A real↔sim gap has two fix directions — check whether the + **real** input is the divergent side before tuning sim. +7. **syn_0 byte-identical gate OFF; unavailability config-gated.** Reuse the async_cifar10 flag axes; don't delete + eligibility plumbing. +8. **Fix the concept, not the symptom — never regress a working example.** Classify a mechanism as + **real-transport artifact** (incremental drain, `recv_fifo` timeouts, queue persistence — no sim analog; gate + `and not self.simulated`) vs **algorithmic property**. Scope-check before editing shared code: + `fwdllm_aggregator.py` = fwdllm blast radius; `top_aggregator.py` / shared `parity` engine / `_sim_recv_min` + can silently break async_cifar10. +9. **Match pytest scope to blast radius.** fwdllm-only edit → `pytest tests/mode -k fwdllm`; shared parity engine + → add `examples/async_cifar10/scripts/parity` + `tests/mode -k parity`; shared stack → full `pytest tests/`. +10. **Comment the WHY, crisply.** A conceptual choice, a real↔sim divergence + rationale, or a failure mode + why + the fix takes its shape — one or two tight sentences. Deeper rationale → §K. +11. **Telemetry-FIRST, then instrument, then (rarely) run.** (a) Validate/refute from telemetry ALREADY ON DISK + first — name the exact field/line. (b) Ship telemetry + plot + pytest IN THE SAME CHANGE as any new mechanism. + (c) A run is justified only to observe an EMERGENT quantity no stored telemetry can yield — then run the + SHORTEST length that exhibits it, smoke first, one mechanism per round. (d) fluxtune surfaces roots first. +12. **Consult PARITY.md vclock rules BEFORE any sim-clock change.** Clock is a monotone `max` + (`vclock = max(vclock, sct)`); NEVER put overhead on it; `sct = send + max(gpu,D) + leg` (leg NOT on + `trainer_speed`/utility); sync charges MAX-of-K sct, async the K-th-fastest; the sim SKIPS real waits and + reconstructs order from sct (`SimReorderBuffer`). +13. **The vclock is virtual wall-time; the sim MUST produce SPEEDUP (`sim_rate = vclock/wall ≥ 1`).** The vclock + advances exactly as the real wall-clock WOULD after the trainer's simulated waits — but faster, because the sim + **elapses through** the modeled time remaining after GPU compute instead of actually waiting it out. fwdllm + nuance: the forward-grad "train" is a REAL GPU pass (~2s) that MUST run in sim for grad mode-invariance — that + ONE GPU pass (parallel across trainers) is the only irreducible real wall; transport/inter-round/delay waits + are vclock jumps, never process sleeps. `sim_rate < 1` means the sim under-charges the vclock OR is stalling on + a real wait it should skip (#15). Emit it every run (`[VCLOCK_PROGRESS]`). +14. **Correctness before speed; SHARED roots before per-baseline.** Fix major logical-correctness divergences + (wrong selection order, wrong receive order, wrong cadence) before any throughput/wall tuning. Rank a root by + blast radius: a bug that fails rungs across ≥2 baselines outranks a single-baseline one — fix the shared cause + once, then narrow to per-baseline residuals. Never chase an aggregate-curve rung while a logic rung is red. +15. **Logical determinism is the parity definition.** For a matched scope the sim must take the SAME sequence of + steps in the SAME order as real — same trainers selected, same order of update receipt, same aggregations and + rollbacks — differing ONLY in wall-clock. Prove this on the **first data bin** (`--max-data-id 1`: multiple + iterations + variance passes + aggregations, short + diffable) before extending length. A matched, ordered + trace is the proof; a matched summary statistic is not. +16. **Do the right thing — no hacks.** Ask as many conceptual questions as it takes to understand the real↔sim + divergence, but solve it at the root. A hack that moves a number without a correct mechanism is a regression in + disguise — it will not close parity and it will mask the real bug. When unsure, stop and ask. + +**Open design decisions:** D2 (avail telemetry port — Phase 2); D3 (variance-cadence × withheld/late grads — +Phase 2); D4 (eval-delay factor — confirmed ~1× train cost, K-D3). D1/D6 resolved (§K). --- -## §G Open questions / risks / dead-ends - -**Open (toward Stage 0):** -- D1-D4 above. -- Which real datasets/traces to use for the fwdllm reference runs (agnews H5 is the smoke default). +## §G Fixes landed (what worked — ≤20-word problem + ≤20-word fix; do not redo) +- **#12c sync `sim_rate` (delay-factor).** No delay-headroom starved the vclock. Fix: `--delay-factor 1` (full + registry D) → sct = send + max(gpu,D) charges the skipped device wall → sync `sim_rate` 0.94→2.9–3.0. +- **K-D31 validated (P2-7a).** Databin1 checks: sync `cohort_sequence` ok=true, set/order/var/cadence=1.0 — the + benign delay-tie is canonicalized away; bin-1 cohort order is bit-exact. +- **Pinning visibility + aggregator GPU pin (K-D33).** Trainer emitted no actual CPU/GPU; aggregator eval defaulted + to GPU 0 (contended with trainers 1/9). Fix: trainer `[PIN]` self-report (`trainer/main.py`), `[LOAD_BALANCE]` + post-proc check (`spawner.spawn_all` — warns on imbalance / under-provisioning), aggregator `gpu_id` pin (least- + loaded/idle GPU). 120 launch tests green; async_cifar10 shares the launcher, unaffected. +- **#13 drain stall (K-D28/b/c).** fwdllm's `_sim_recv_min_grad` left a stuck straggler in the expected set → + re-fired the full 30s `RECV_TIMEOUT` every cycle → pipeline starved (`sim_rate` 0.06). Fix: felix-port + stuck-end eviction + recv-grace 2→5s + probe-ceiling/ready-gating + `drain_ready` direct ingest → `sim_rate` + 0.06→0.30, 30s stall gone, R1=0. (Step-4 staggering NEUTRAL, OFF — §H; residual is #12c-bound.) +- **#1c R1 two-ledger bridge (K-D27/b).** async_oort gated selection ONLY on `all_selected` (physical-event-pruned) + and never on the aggregator's virtual in-flight set → slow-sim returned-but-uncommitted trainer re-dispatched + (R1 62.9%). Fix: agg maintains `_sim_pending_commit` felix-style, binds it live to `sel._agg_pending_commit_ref`, + async_oort's `filtered_ends` excludes it; `outstanding = inflight ∪ buffer` (K-D27b, NOT `− _sim_committed`). + Sim-only → async_cifar10 byte-identical. `SIM_R1_DISPATCH` 238→0, confirmed on the latest full run. +- **#14 MQTT join-notify startup race.** `backend/mqtt.py` `join()` fired `notify(JOIN)` fire-and-forget; when the + MainThread beat async `on_connect`, the JOIN + subscriptions dropped with no retry → end never registered → + async_oort hung below `minInitialTrainers`. Fix: `join()` `_wait_for_connect()` (10s bound) before subscribe+notify; + no-op when connected → all examples unchanged. +- **#6 clock-rate anchor + Root B phase rungs (K-D25).** The clock-rate rungs compared sim-vclock against real's + FULL wall, which carries a ~7.7s/round localhost transport ARTIFACT — a checker-anchor bug, the sim vclock was + correct. Fix: agg emits `intrinsic_span_s` (barrier+eval, fedavg excluded); rungs + `wall_disparity` anchor real + on its intrinsic clock (async byte-identical). Validated throughput 0.47→0.048, wall_disparity 90→1.6. +- **Phase-4 stopping rule (K-D24).** `sim_wall_ceiling_s` was truncating a real-compute sim. Fix: decoupled to + `max_runtime_s × 20`; matched `--max-data-id` is the primary stop. Fixed root S1. +- **Residence: commit-then-carry + felix realign (K-D12/K-D14/K-D17b).** Boundary-drop stranded ~7 async + grads/cycle → 2× passes; slot-on-return undercounted in_flight 3×. Fix: carry the surplus, hold the compute slot + to COMMIT (so `len(selected_ends)` = virtual-time in-flight), source R1 from an immutable echoed interval. +- **Clock family re-keyed to `data_id` (#2).** Rungs keyed on `round` mismeasured fwdllm's variance-pass progress. + Fix: `_progress_axis`/`_per_progress_last_event`; async_cifar10 auto-detects `round` → byte-identical. +- **Foundational sim mechanisms (Batch 1–2, K-D2/K-D3/K-D4/K-D5/K-D9/K-D11).** Additive `sct = send + gpu + D`, + no-sleep sim path; `_sim_recv_min_grad` sct reorder buffer + in-flight gate + agg-goal rollback cleanup; sync + `_sync_sim_recv_first_k` first-k-smallest (real-only commit-1-per-pass clamp); V1–V5/DK1–DK3/G1–G2 rungs + the + pre-mutation cycle snapshot that makes V1 exact. +- **Scaffolding (landed, pointer only).** Pre-run instrumentation A–E (K-D21); availability params end-to-end, + print==run (K-D22); Phase-2 real-only speedup-leak skips (K-D23); `staleness_policy` wired from config (K-D13/15). -**Risks specific to fwdllm:** -- **Variance-cadence is feedback-compounding** -- a tiny per-cycle grad-pool order difference compounds into - a different iterations-per-data_id (like the async clock residuals that surfaced only at 3 h). Bin V1/V2 - by run-fraction to separate a *constant* mix bias from a *compounding* feedback loop. -- **`cached_v` carry-over** is stateful across cycles; a divergence looks like a variance bug but is - bookkeeping (V3 DIAG localizes it). -- **One-message-per-call grad loop** vs asyncfl's batch -- slot-hold/release wiring is genuinely different; - do not copy async_cifar10's release points blindly. -- **Withheld grad x variance gate** (D3) -- a stale late grad may distort `var` and the dynamic-K decision. +--- -**Pre-emptive dead-ends (from PARITY.md / UNAVAILABILITY_DESIGN.md):** overhead > 0 on the vclock; -prediction-only gates that never block; tuning a scalar redispatch/overhead knob instead of fixing the -mechanism; letting eval commits into the train/agg stream; expressing "busy" via the unavailable list. -**fwdllm-new:** do **not** tune `var_threshold`/`max_iterations_per_data_id` to force cadence parity -- -those are baseline-defining knobs; a cadence gap is an upstream set/order/clock divergence. +## §H Dead-ends & corrections — do NOT retry +- **fluxtune #15 — three superseded framings (all 2026-07-06, same session; final root = COMMIT-PATH STALL).** + (1) "GPU-PIPELINING loss; keep gate, decouple dispatch" — WRONG, felix's arrival gate is INERT (gate_holds=0). + (2) "re-dispatch on physical RETURN to keep GPUs busy" — WRONG, fedbuff never re-hands a returner the same + version, and real's 3.37 concurrency is a duty cycle `gpu/max(gpu,D)`, not under-use. (3) "hold-to-commit is a + SYNC barrier / over-restrictive; replace with model-advance re-dispatch" — WRONG, hold-to-commit is a + CORRECTNESS check (freed only on commit; guards no-same-version / not-while-computing / + not-while-returned-uncommitted). *Final root:* the commit path STALLS (drain gate blocks real wall on PHANTOM + `_sim_inflight_expected` entries) → correctly-held trainers idle. Fix = fast/non-stalling commit path; + hold-to-commit untouched. *Lesson:* commit RATE is the throughput lever, not the residence rule. See + PARITY_LOGICAL_TASKS.md FELIX GROUNDING F5/F6. +- **"The sync cadence break is a sim ORDER bug (order → split-half var → RNG desync) — exact cadence parity is + achievable once order matches."** CORRECT for bin ≤1, REFUTED for the full run (2026-07-05). K-D31 made + receive-ORDER 41/41 identical, yet fwdllm cadence STILL breaks at bin 7. Root is grad NON-reproducibility given + matched order: ~1e-3 GPU fp16 jitter, amplified by the split-half variance ratio, flips the `var<0.3` gate at a + sensitive bin. *Lesson:* exact `var` is only a valid target ≤bin 1; beyond it the target is DISTRIBUTIONAL. Don't + chase exact cadence past the nondeterminism wall (principle #16) — it reds on float noise, not a bug. +- **"The fluxtune runs are GPU under-provisioned (2–4 of 8 GPUs, 5 trainers/GPU) → contention is the root."** WRONG — + a misread of `gpu=4.9s` (GPU compute SECONDS in the delay log) as device IDs. The spawn table is authoritative: + **8 GPUs, balanced round-robin** (`spawner.py:305` `(tid-1)%num_gpus`), CPU `sched_setaffinity` 1 core/trainer, + threads capped. Only structural imbalance is 10 trainers > 8 GPUs (GPU 0,1 carry 2 each). Compute IS mode-invariant + at the floor (min 2.4s both modes). *Lesson:* verify a "device id" is a device id; confirm pinning from the spawn + table, not a grep of timing logs. +- **"#13 step 4 (freed-slot staggered re-dispatch) closes the residual 11-12s drain holds."** NEUTRAL/REFUTED + (K-D28d): `sim_rate` 0.289→0.274, holds persisted. The holds are the drain correctly waiting (strict sct order) + for the earliest-sct in-flight straggler while higher-sct grads buffer — INHERENT to real-GPU + strict-order + + trickle dispatch, not a dispatch-bunching artifact. Feeding an OLD freed-slot vclock as `send_ts` makes the gate + expect trainers even earlier → more holds. *Lesson:* the residual after steps 1-3 is **#12c** (no delay-headroom), + NOT the drain. Code kept flag-gated OFF; a reworked attempt would stamp the ACTUAL dispatch vclock. +- **"K-D19 fixed fluxtune R1 / the NONE reset is a red herring."** REFUTED (K-D26). R1 stayed 60.4%: K-D19 fixed + the wrong release path and dismissed the NONE hypothesis by checking the *selection filter*, but `all_selected` + MEMBERSHIP was deleted upstream (90s wall-timeout + async_oort reading `KEY_END_STATE=NONE` as "left"). + *Lessons:* (i) an R1 fix isn't done until the smoke banks R1≤2%; (ii) when a guard "should exclude" but doesn't, + check whether the member is *deleted* upstream, not just whether the filter reads it. +- **"Free the compute slot on physical RETURN" (K-D16 Option-A).** WRONG for virtual time — a returned-but- + uncommitted trainer is still in flight until the vclock reaches its sct; freeing undercounted `in_flight` 3×. + Reverted to hold-to-COMMIT (K-D17b). (The K-D16 re-run also DEADLOCKED, K-D17: drain gated on channel RECV not + the buffer; triplet stamped at DISPATCH froze the pool pre-commit.) +- **"Drop stranded grads at the agg-goal boundary" (K-D6).** REVERSED for async (K-D12) — the "|selected|≈agg_goal" + premise holds only for sync; fluxtune (c=10≫agg_goal=3) dropped ~7/cycle → 2× passes. Drop stays correct for sync. +- **Adding a scalar overhead to the vclock to close #6.** Forbidden (principle #1/#12) — the vclock is + `max(vclock, sct)`. Fold only genuine unmodeled compute (eval_s, straggler spread); artifacts stay off. #6 was a + checker-anchor bug (K-D25), not a missing fold. (Including FedAvg in `intrinsic_span_s` over-charged it → excluded.) +- **Tuning `var_threshold` / `max_iterations_per_data_id` to close a cadence gap.** Rejected — baseline-defining + knobs, not parity levers. A cadence gap is an upstream set/order/clock divergence. +- **"sim wall ≈ real, comparable" / "D=0 smoke ⇒ clock fails are artifacts".** WRONG — always check avg in-flight + concurrency (not just pass counts); the runs are D>0 and the fails traced to `round`-vs-`data_id` keying (#2). --- -## §H Status +## §K Deviation log — one line per decision (anchor + rationale; referenced from §B.1/§G) +- **K-D1** — `time_mode` default `"real"` (not cifar's `"simulated"`): fwdllm's whole corpus is `real`; a + `simulated` default risks half-activating an unbuilt path. +- **K-D2** — additive `sct = send + gpu + D` (not `max`): fwdllm's real mode sleeps D on top of GPU time. +- **K-D3** — per-eval sct collapses to the train sct (D4): eval lives on the aggregator; forward-grad "train" is a + forward pass (no 20× factor); the trainer eval message is a utility report, not a clocked commit. +- **K-D4** — purpose-built `_sim_recv_min_grad` (not `_sim_recv_min` verbatim): the latter's per-commit release + + withheld paths key on WEIGHTS semantics; reuse the primitives, fork the orchestration. +- **K-D5** — slot release + buffer clear on the AGG-GOAL boundary (not per-commit): a `data_id` spans many cycles + with rollbacks; per-commit release would strand a re-contributing trainer (principle #4). +- **K-D9** — cadence telemetry is a **pre-mutation** cycle snapshot: post-mutation `data_id` advances before emit + → off-by-one; snapshot makes V1 exact. Existing post-mutation fields unchanged (byte-identical real). +- **K-D11** — `ends_not_selected_yet` "commit-1-per-pass" clamp gated real-only: a real-transport draining + discipline; the sim barrier is single-pass (principle #8). +- **K-D12** — fluxtune async: commit-then-CARRY the surplus + hold residence (reverses K-D6); drop stays correct + for sync (c≈agg_goal). +- **K-D13/K-D15** — fluxtune `staleness_policy = fedbuff` staleness-weighted accept (was silently `none`); set + identically real+sim (a definition, not a lever). +- **K-D14** — R1/W1 sourced from an ECHOED per-contribution interval (not the agg's per-end dispatch stamp, which + is overwritten on re-dispatch — exactly when residence is broken). +- **K-D17b** — hold the compute slot to COMMIT (felix port): `len(selected_ends)` = virtual-time in-flight + (fixed in_flight 2.7→9.5). Reverts K-D16's slot-on-return. **CONFIRMED CORRECT for both sync AND async (K-D34):** a + trainer is freed only when its update commits — a correctness check, not a throughput lever. fluxtune's `sim_rate<1` + is a commit-path STALL (#15), NOT this rule; it stays untouched. +- **K-D21** — pre-run instrumentation A–E landed & pytest-green; un-skipped the 12 rigor-gap rungs (§G). +- **K-D22** — availability params respected end-to-end; Phase-1 syn_0 default; print==run (§G). +- **K-D24** — Phase-4 ceiling decouple (×20) + B1/B2 sct folds; fixed root S1 (§G). +- **K-D25** — #6 is a CHECKER-ANCHOR bug, not a sim under-charge: agg emits `intrinsic_span_s` (barrier+eval, + fedavg excluded), rungs anchor real on its intrinsic clock; Root B stamps `post_train_s` after the delay and + moves the B2 straggler into the sct only. Async byte-identical (§G). +- **K-D26** — #1c R1 root-caused to a physical-wall vs vclock desync in the shared `async_oort` re-pick guard + (`all_selected`), exposed by the slow sim. Fix (a): `_abandon_clock_now()` runs the 90s timeout on the vclock in + sim (None in real → byte-identical). Corrected K-D19. Sustained NONE-delete path → deeper root K-D27. +- **K-D27/b** — the two-ledger split: async_oort's selection eligibility never consulted the aggregator's virtual + in-flight truth. Fix: agg maintains `_sim_pending_commit` felix-style + binds `sel._agg_pending_commit_ref`; + `filtered_ends` excludes it; `outstanding = inflight ∪ buffer` (K-D27b — the `− _sim_committed` subtraction + dropped re-dispatched-after-commit trainers, R1 stuck 67.6%). Sim-only. `SIM_R1_DISPATCH` 238→0 (§G). +- **K-D28/b/c** — #13 drain-stall felix port (3 steps landed): stuck-end eviction + recv-grace floor 2→5s; + probe-ceiling + ready-gating; direct `drain_ready` ingest (flag `sim_sct_ordered_drain`). `sim_rate` 0.06→0.30, + 30s stall gone, R1=0, 0 recv_fifo timeouts (§G). +- **K-D28d** — #13 step-4 freed-slot staggered re-dispatch: IMPLEMENTED but NEUTRAL → flag DISABLED. The residual + holds are inherent strict-sct-order straggler waits, not a dispatch artifact; residual `sim_rate` is #12c-bound + (§H). +- **K-D29** — **REMAINDER-WAIT delay model (reverses K-D2's additive decision).** The order→split-half-var→RNG- + desync root (§A) requires the sim's commit order to equal real's arrival order. That needs a deterministic + per-trainer arrival order, which the flat-additive `gpu+D` model didn't give (order was GPU-jitter-dominated). + Fix (aligned with async_cifar10): real sleeps `max(0, D−gpu)` so the device wall = D (GPU hidden inside); sct = + `send + max(gpu, D)`; per-trainer registry delays supply the completion spread → order = D-order = deterministic + & real↔sim identical. Overrun (gpu>D) is flagged (`training_overran`, `[TIMING_OVERRUN]`) — it un-determinises + order, so it's the fluxtune watch (its JVP GPU 7.57s overruns D/2). Crc32 straggler offset disabled + (`sim_straggler_spread_s=0`); `perturbation_count` made a knob (default 10). See PARITY_LOGICAL_TASKS.md P2. + +- **K-D30** — **full-cohort determinism gate + `timing_overrun` signal (P1-5, closes the P1-4/P1-5 gap).** The + selection set/sequence rungs (`selection`/`aggregation_sequence`/`utility`) gated to a TRIVIAL pass for every + stochastic selector, so fwdllm's syn_0 selection was never actually checked. Fix: `_selection_is_deterministic` + un-gates when `num_chosen==num_candidates` in BOTH modes (K≥pool → set-deterministic) — data-driven, so it + enforces fwdllm (K=all), keeps fluxtune (agg_goal=3) + fwdllm_plus (#7 asymmetric eligible) gated, self-disables + under Phase-2 scarcity, and falls back to the old selector-name rule on count-less legacy telemetry (no + regression). `participation` EXCLUDED (round-keyed → mechanical KS on fwdllm's constant-`round`/`data_id` axis; + cohort_sequence is its per-cycle enforcement). New `timing_overrun` DIAG rung surfaces the K-D29 overrun tell + (P2-6 `training_overran` fraction + first-overrun bin → gpu>D flips order → cohort/var break is a timing-model + limit not a sim bug). P1-4 (async_cifar10 cohort adapter) assessed redundant, not built. Sim/checker-only → + async_cifar10 byte-identical; banked scoreboard stable. + +- **K-D31** — **canonical `(D, trainer_id)` cohort commit order (P2-7a; operator chose canonicalize-by-id).** The + databin1 run left ONE sync residual: two trainers sharing a registry delay (D=6.5) swap in receive order (real + breaks the sct tie by physical arrival, sim by sct-sort) — benign (both in the same split-half → `var`/grads + bit-identical) but the EXACT-order `cohort_sequence` rung flags it. Fix: trainer stamps pure `D` + (`MODELED_DELAY_S`, both modes, deterministic from the registry — unlike sct which folds in GPU jitter); agg + sorts each cohort by `(D, str(end))` in `_canonicalize_cohort_commit_order` before the telemetry snapshot + + `aggregate()`, reordering `_per_agg_trainer_list` + the trailing cohort slice of the grad/jvp lists in lockstep. + No-op when delays off / already canonical (real still receives in strict D-order, K-D29 — only ties move). + Sim/agg-only → async_cifar10 byte-identical; 428 mode + 9 canon + 115 parity green. + +- **K-D32** — **fluxtune JVP perf optimizations (`jvp_perf_opt`, §L; fluxtune-only, config-gated, all + BIT-IDENTICAL).** fluxtune's forward-grad JVP was ~10× the sync compute (2P=20 selection passes + 5 removable), + overrunning its modeled mobile delay → out-of-order commits (#1d). Fix (validated by + `scripts/profile_jvp_opt.py`): (a) trainable-only finite difference — `calculate_jvp(trainable_idx=…)` skips + `p−h·0=p` on the 98.5% frozen backbone (1.26× + −251 MB); (b) skip the 3 diagnostic-only forward passes; (c) + reuse the selected perturbation's cached JVP. Combined fluxtune −37% (sync would be −68% but is left OFF). + Gated on `jvp_perf_opt` (trainer_base default false = byte-identical; true in both fluxtune yamls, must match + real↔sim; revertible per-config). Startup `[JVP_PERF_OPT]` log confirms 10/10 trainers active. NOT adopted: vmap + (2× but fp32-diverges via FD cancellation), fwd-AD (slower). 13 pytests + 185 fwdllm mode + 115 parity green. + +- **K-D33** — **pinning self-report + load-balance check + aggregator GPU pin.** Trainers now emit a `[PIN]` line + (actual `CUDA_VISIBLE_DEVICES`/cuda device+name/`cpu_affinity`) at startup (`trainer/main.py` — NOT the dead + `fl_main.py`); `spawner.spawn_all` emits `[LOAD_BALANCE]` (WARN on GPU imbalance, `num_gpusnum_gpus`, else least-loaded) GPU so its eval stops contending GPU 0. Confirmed the + `client_idx%8` device arg is vestigial (`FedSgdTrainer:388` overwrites `self.device=torch.device("cuda")`=cuda:0 of + the CVD-masked view) → the spawner's pin is authoritative. Shared launcher; 120 launch tests green. + +- **K-D34** — **fluxtune #15 = a COMMIT-PATH STALL; `sim_compute_truthful_gate` fix LANDED (P1/P2), P3 pending + (2026-07-06).** Supersedes three same-session mis-framings (§H). Grounding (PARITY_LOGICAL_TASKS.md FELIX + GROUNDING F5/F6): hold-to-commit frees a trainer only on COMMIT (guards: no same-version dispatch, none while + computing, none while returned-but-uncommitted), so commit RATE is the throughput lever; felix's gate is INERT + and cifar's commit is instant (~0.4s) so held trainers barely idle. **D1/D2 CONFIRMED** (banked pair): + fluxtune's `_sim_recv_min_grad` `earlier_stuck` gate holds already-arrived grads (`buf_depth=7`, 99.7%) to wait + on trainers stamped-expected-at-DISPATCH but idle-in-recv (not computing) behind the single-threaded drain → + 82% of wall (~1974s) burned → concurrency sim 1.65 vs real 7.69 → sim_rate 0.50. **Fix (flag + `sim_compute_truthful_gate`, default off = byte-identical, fluxtune-yaml on):** stamp `_sim_dispatch_wall[end]` + at the real `channel.send`; the gate skips any expected entry whose last dispatch is older than + `sim_gate_compute_cap_s` (default 10s) — a stamped-but-idle phantom no longer blocks a ready commit, a genuine + in-window straggler is still held (commit order preserved). Hold-to-commit, sct-ordered drain, K-D12, K-D27 + UNTOUCHED; `fwdllm_aggregator`-only → async_cifar10 byte-identical. 25 pytests green. **P3 RAN 2026-07-06 + (`run_20260706_112114` sim / `_110555` real):** phantom fix **VALIDATED** — `STUCK_EVICT=0`, `phantom_skip`→65 + (~1/commit, load-bearing), 30s failsafes gone (max gap 12.8s). **But `sim_rate` still <1** (steady ~0.49; + per-commit `Δvclock/Δwall=0.465`) — the phantom stall was necessary but NOT the ceiling. **Residual root + (NEW):** (a) vclock omits `aggregate()` variance-compute wall (sim 32s / real 29s, symmetric, uncredited — + `:1783` folds eval only); (b) GPU≈D no-headroom (#1d) → sim blocks in `drain_ready` on the real GPU pass. Sim + wall (145s) > real (91s). **PAUSED** for operator real-experiment work; resume plan (fold non-overlapped + aggregate into vclock + GPU-vs-D headroom + longer flag-on run & flag-off A/B) at PARITY_LOGICAL_TASKS.md top. + jvp_perf_opt verified symmetric (trainers True, aggregator-eval False, both modes) — NOT a mismatch. + +*Retired/superseded anchors (kept only as pointers): K-D6 (→K-D12), K-D7/K-D8/K-D10/K-D16/K-D18/K-D19/K-D20/K-D23 +— landed scaffolding or corrections, folded into §G/§H; see git history for detail.* + +--- -*(empty -- first entry after Stage 0 smoke on the implementation branch. Record per-baseline run dirs, -lowest broken rung, root hypothesis, score X/N, JSON path; keep the run-length budget keyed. One section, -updated in place.)* +## §L Forward-grad JVP compute profile & retained fluxtune optimizations +*(2026-07-05, tool: `scripts/profile_jvp_opt.py` — reuses real `create_model` + `calculate_jvp`; distilbert-base ++ AdapterHub adapters, batch 8, seq 192, A40, fp16. Absolute ms are a CLEAN single-trainer profile; the real run +is ~10× from GPU contention across the 10 concurrent trainers, but pass-counts/ratios/memory transfer.)* + +**Mechanism.** Forward-grad trains via a **central finite-difference JVP** (`fwdgrad_utils.calculate_jvp`): each +perturbation = **2 forward passes** `f(θ±hv)`, h=0.01, autocast+no_grad → `jvp=(f(θ+hv)−f(θ−hv))/2h`. **fluxtune** +selects the best of `perturbation_count`(=10) perturbations by |jvp| (2P=**20 passes**); **fwdllm/sync** selects +by cos-sim (**0 forward passes**) + 1 final JVP. Only **~1.5% of params trainable** (bottleneck adapters in all 6 +layers + head, 1.04M/67.4M); backbone frozen. + +| path | fwd passes | ms/batch (clean) | +|---|---|---| +| sync fwdllm (current) | 5 | 50 | +| **sync fwdllm (opt)** | 2 | **16 (−68%)** | +| fluxtune P=1 (opt) | 2 | 16 (== sync) | +| fluxtune P=5 (opt) | 10 | 80 | +| fluxtune P=10 (current) | 25 | 251 | +| **fluxtune P=10 (opt)** | 20 | **159 (−37%)** | +| backprop ref (1 fwd+1 bwd) | — | 17 | + +- **Compute vs sync:** fluxtune = `2P × per-pass` → **10× sync at P=10**, linear in P, **equals sync at P=1**. + JVP-selection is the entire fluxtune surcharge; sync's cos-sim selection is free. +- **Memory:** forward-grad peak is **FLAT in P** (~3.2–3.4 GB = model + one held forward; **no autograd graph**). + Backprop stores activations (3.71 GB). fluxtune's extra JVP inferences cost **TIME, not memory** — same + footprint as sync (forward-grad's design tradeoff: many cheap forward passes, no backward, low memory). +- **Per-pass:** full-param FD 10.0 ms; trainable-only FD 7.95 ms. + +**RETAINED — bit-identical (fidelity-preserving; real↔sim parity untouched):** +1. **Trainable-only FD** — skip `p−h·0=p` on the 98.5% frozen params inside `calculate_jvp`: **1.26× + −251 MB**, + `max|Δjvp|=0`. +2. **Drop the 3 diagnostic-only forward passes** (`_train_one_batch:646-648`, loss before/after-update logging — + never feed grads/telemetry) **+ reuse the winner's cached JVP** (`:645`, fluxtune): fluxtune 25→20, sync 5→2. + +Combined: **sync −68%, fluxtune −37%**, all bit-identical → should clear the `delay_factor=1` overrun (4.2s → +~2.6s < min cohort D 4.0s) WITHOUT touching fidelity. **LANDED (K-D32), fluxtune-only & config-gated** +(`jvp_perf_opt`, default false = byte-identical; true in both fluxtune yamls; sync untouched). Startup log +`[JVP_PERF_OPT]`; 13 new pytests + 185 fwdllm mode + 115 parity green. + +**NOT retained — change fidelity:** +- **vmap-batching the perturbations** — **2.0×** (biggest single win), mathematically exact (fp64 seq==vmap + BIT-IDENTICAL, deterministic → *would* be real↔sim safe) BUT differs ~5% from sequential in fp16/fp32: the FD + subtracts two O(1) losses (catastrophic cancellation floors precision), so any reduction-order change + re-baselines the trajectory. Excluded per the fidelity bar; available if a re-baseline is accepted. +- **Forward-mode AD** (exact JVP) — slower (0.5×, needs eager attention; not impl for SDPA) + different math. +- **`perturbation_count`↓** — the direct lever, but changes the baseline algorithm. diff --git a/lib/python/examples/fwdllm/telemetry_manifest.yaml b/lib/python/examples/fwdllm/telemetry_manifest.yaml index 4a0bfabac..05e74c9ad 100644 --- a/lib/python/examples/fwdllm/telemetry_manifest.yaml +++ b/lib/python/examples/fwdllm/telemetry_manifest.yaml @@ -58,14 +58,17 @@ event_categories: # path) unlocks selected_utility_believed_vs_actual*/belief_gap*. # agg_observed_s (Part 6, reusing the same PROP_ROUND_DURATION already # read for agg_round's trainer_speed_s) unlocks runtime_agg_vs_trainer/ - # runtime_overhead_*. Deliberately NOT added (Part 6 design decision): - # training_budget_s/overran/remaining_time_s/wait_time_s/pre_train_s/ - # post_train_s -- fwdllm's delay is a flat additive sleep, not cifar10's - # budget-minus-actual sleep-to-fill model, so there is no faithful - # "overrun"/budget-slack concept to report; budget_slack_cdf/ - # trainer_late_fraction_hist/trainer_response_lateness_cdf/ - # trainer_runtime_expected_vs_actual stay not_populated rather than - # reporting a fabricated budget signal. + # runtime_overhead_*. Per-phase wall breakdown ADDED in Stage A1 for the + # #6 wall-decomposition (real vs sim s/round): mqtt_fetch_s/weights_to_ram_s/ + # weights_to_gpu_s/pre_train_s/gpu_compute_s/post_train_s/training_budget_s/ + # trainer_phase -- these un-skip the 8 phase rungs. NOTE the earlier Part-6 + # decision still holds for the DERIVED budget-slack plots: fwdllm's delay is + # a flat additive sleep (training_budget_s = _delay_s, the modeled D), not + # cifar10's budget-minus-actual sleep-to-fill model, so there is no faithful + # "overrun"/slack concept -- overran/remaining_time_s/wait_time_s stay + # unemitted and budget_slack_cdf/trainer_late_fraction_hist/ + # trainer_response_lateness_cdf/trainer_runtime_expected_vs_actual stay + # not_populated rather than reporting a fabricated budget signal. selection: populated # selection, emitted from flame/selector/random.py's select() (P5.5 fix # -- RandomSelector never called emit_selection before this session, so diff --git a/lib/python/examples/fwdllm/trainer/fl_main.py b/lib/python/examples/fwdllm/trainer/fl_main.py index f92d6d8a7..f12cfc84c 100644 --- a/lib/python/examples/fwdllm/trainer/fl_main.py +++ b/lib/python/examples/fwdllm/trainer/fl_main.py @@ -154,6 +154,13 @@ def post_complete_message(tc_args): "var_control": config.hyperparameters.var_control, "perturbation_sampling": config.hyperparameters.perturbation_sampling, "select_perturbation_using_jvp": config.hyperparameters.select_perturbation_using_jvp, + # P2-5: forward-pass count knob (default 10 = historical behavior). + "perturbation_count": getattr( + config.hyperparameters, "perturbation_count", 10), + # §L: fluxtune JVP perf-opt (bit-identical). Default False = + # byte-identical; enabled only in the fluxtune yamls. + "jvp_perf_opt": getattr( + config.hyperparameters, "jvp_perf_opt", False), } ) model_args.config["num_labels"] = num_labels diff --git a/lib/python/examples/fwdllm/trainer/forward_training/FedSgdTrainer.py b/lib/python/examples/fwdllm/trainer/forward_training/FedSgdTrainer.py index f64e9faf4..20104bf65 100755 --- a/lib/python/examples/fwdllm/trainer/forward_training/FedSgdTrainer.py +++ b/lib/python/examples/fwdllm/trainer/forward_training/FedSgdTrainer.py @@ -6,6 +6,7 @@ import json import hashlib import os +import zlib import numpy as np from datetime import datetime import ast @@ -17,6 +18,9 @@ from flame import telemetry from flame.telemetry.events import build_trainer_round +# Import under the same absolute path tc_transformer_trainer_distribute.py uses, +# so we read the SAME module-global forward-pass counters (WS3-b), not a copy. +from examples.fwdllm.trainer.forward_training import fwdgrad_utils logger = logging.getLogger(__name__) @@ -181,6 +185,30 @@ def __init__( ) self.speedup_factor = 1.0 + # --- Simulated-clock support (config-gated) --------------------------- + # time_mode is threaded into hyperparameters from the launcher's + # --time_mode CLI arg (see trainer/main.py). "simulated": _emulate_ + # training_delay() computes the modeled delay but does NOT sleep, and + # train_with_data_id() stamps a modeled completion timestamp + # (_sim_completion_ts) the aggregator orders updates by. "real" + # (default): unchanged wall-clock behavior (flag-off => byte-identical). + self.time_mode = getattr(self.config.hyperparameters, "time_mode", "real") + self.simulated = self.time_mode == "simulated" + _leg = getattr(self.config.hyperparameters, "sim_completion_leg_s", 0.0) + self.sim_completion_leg_s = float(_leg) if _leg is not None else 0.0 + # SIM_SEND_TS is stamped by the aggregator on each dispatch and read in + # the base trainer's _fetch_weights; the rest are stamped after training + # for _send_grads / telemetry. + self._sim_send_ts = None + self._sim_completion_ts = None + self._sim_round_duration_s = None + # Pure modeled delay D for this round: stamped in the grad message + # (MODELED_DELAY_S) so the aggregator orders cohort commits by + # (D, trainer_id) identically in real and sim. None until first round + # completes / when delays disabled. + self._modeled_delay_s = None + self._wall_recv_ts = None + self.trainer_start_ts = time.time() # TODO (ARM): Fix this to read traces better! # Storing synthetic avail traces @@ -434,6 +462,13 @@ def _check_availability(self): f"Trainer id {self.trainer_id} is not available to train. Waiting for it to be available" ) while self.avl_state != TrainerAvailState.AVL_TRAIN: + # Sim availability is enforced aggregator-side (send-gate / + # vclock-jump to next avail event), never by a trainer sleep: + # sim time can't advance while blocked on time.sleep, so this + # spin would freeze the virtual clock (#13). Sim proceeds and + # lets the agg-side gate withhold; real byte-identical. + if self.simulated: + break time.sleep(1) logger.info( f"Trainer id {self.trainer_id} is back to available to train." @@ -478,49 +513,138 @@ def _perform_training(self): self.jvp_for_snr_check = self.trainer.model_trainer.jvp_for_snr_check @timer_decorator - def _emulate_training_delay(self): - """Returns the seconds actually slept (0.0 if delay emulation is - disabled) -- unlike cifar10's trainer, this is a flat additive sleep - on top of GPU time, not a budget-minus-actual "sleep to fill" model, - so there is no meaningful overrun/remaining_time_s/training_budget_s - concept here (see ../../../MIGRATING_TO_LAUNCHER.md §9). The - caller adds this to real_gpu_time_s to report sim_round_duration_s.""" - if self.training_delay_enabled == "True": - # Eval is 3X faster than training on CPU - # Eval on NPUs is 10-50X is faster than training on CPUs. We could take 20X if we wanted to consider an all-NPU client cohort for Eval (NPUs don't support training) - eval_delay = self.training_delay_s / self.training_delay_factor - _sleep_s = eval_delay / self.speedup_factor - time.sleep(_sleep_s) + def _emulate_training_delay(self, gpu_time_s: float = 0.0): + """Remainder-wait delay model (aligned with async_cifar10). The modeled + mobile device takes `_delay_s = training_delay_s / factor / speedup`; our + GPU forward pass takes `gpu_time_s`, which should be << the device time. + + - REAL sleeps only the remainder max(0, _delay_s - gpu_time_s) → wall + this round ≈ _delay_s, with GPU compute hidden inside it. + - SIM skips the sleep (charged to the vclock); the round duration is + max(gpu, _delay_s) (see train_with_data_id), NOT gpu + _delay_s. The + per-trainer registry delays give the completion spread, so update + order = delay order = deterministic and identical real↔sim (enables + cohort_sequence parity). + + Overrun: if gpu_time_s > _delay_s the GPU is slower than the modeled + device (contention / delay_factor too big) — emulation unfaithful and + update order can flip, so we log [TIMING_OVERRUN] and flag it in + telemetry. Returns (modeled_delay_s, remaining_s, overran). + """ + # config schema types training_delay_enabled as bool (default False) + # but historical launcher yamls pass the string "True"; accept both so + # the modeled delay is not silently dropped to 0. + _enabled = self.training_delay_enabled in (True, "True", "true") + if not _enabled: + return 0.0, 0.0, False + _delay_s = (self.training_delay_s / self.training_delay_factor) / self.speedup_factor + _remaining_s = max(0.0, _delay_s - gpu_time_s) + _overran = gpu_time_s > _delay_s + if _overran: + logger.warning( + f"[TIMING_OVERRUN] trainer {self.trainer_id} data_id={self.data_id} " + f"iter={self.iteration_per_data_id}: gpu={gpu_time_s:.2f}s > " + f"budget={_delay_s:.2f}s (excess={gpu_time_s - _delay_s:.2f}s) — " + f"emulation unfaithful, update order may flip. Reduce trainers/GPU " + f"or raise training_delay_factor." + ) + if self.simulated: + logger.info( + f"time_mode=simulated: modeled delay for trainer {self.trainer_id} " + f"= {_delay_s:.3f}s (not slept; charged to vclock; gpu={gpu_time_s:.3f}s)." + ) + elif _remaining_s > 0: + time.sleep(_remaining_s) logger.info( - f"Delayed eval time for trainer " - f"{self.trainer_id} by {eval_delay}s. Sleeping for {_sleep_s}s." + f"Trainer {self.trainer_id} slept remainder {_remaining_s:.3f}s " + f"(budget {_delay_s:.3f}s - gpu {gpu_time_s:.3f}s)." ) - return _sleep_s - return 0.0 + return _delay_s, _remaining_s, _overran + + def _sim_straggler_offset_s(self) -> float: + """The modeled delay D is flat across trainers, so the sync barrier's + k-th-smallest sct under-spreads vs real's GPU-contention dispersion. In + SIM only, add a stable per-trainer offset in [0, simStragglerSpreadS) so + the cohort completion spread matches real. Deterministic in trainer_id + (crc32, not salted like hash()); spread 0 => byte-identical.""" + if not self.simulated: + return 0.0 + _hp = getattr(getattr(self, "config", None), "hyperparameters", None) + spread = float(getattr(_hp, "sim_straggler_spread_s", 0.0) or 0.0) + if spread <= 0.0: + return 0.0 + frac = (zlib.crc32(str(self.trainer_id).encode()) % 1000) / 1000.0 + return spread * frac @timer_decorator def train_with_data_id(self): - # Create FwdLLMStage for timing/metrics logging - self.fwd_llm_stage = FwdLLMStage( - self._round, self.data_id, self.iteration_per_data_id, self.trainer_id - ) - if self.abort_training == True: logger.info( f"Aborting training for trainer id: {self.trainer_id} because it has already sent updates for iteration_per_data_id: {self.iteration_per_data_id}" ) return + # Create FwdLLMStage for timing/metrics logging. Set AFTER the abort + # check so an aborted round stays a true no-op (no fwd_llm_stage => the + # @timer_decorator emits no step_timing record; see + # test_aborted_round_emits_nothing). + self.fwd_llm_stage = FwdLLMStage( + self._round, self.data_id, self.iteration_per_data_id, self.trainer_id + ) + + # Phase-timing entry: everything up to the compute loop is pre_train + # (avail check, FwdLLMStage setup, loader state). + _phase_entry = time.time() if not self._check_availability(): return _round_start_ts = time.time() + _pre_train_s = _round_start_ts - _phase_entry self._perform_training() _real_gpu_time_s = time.time() - _round_start_ts - # emulate delays in training (due to compute resource and/or - # dataset size and/or network latency) - _delay_s = self._emulate_training_delay() + # emulate the mobile-device delay via the remainder-wait model: real + # sleeps max(0, delay - gpu); sim skips it. Returns the modeled budget, + # the remainder actually waited, and whether the GPU overran the budget. + _delay_s, _remaining_s, _overran = self._emulate_training_delay(_real_gpu_time_s) + # Stash the pure modeled delay D so _send_grads can stamp it + # (MODELED_DELAY_S) for the aggregator's (D, trainer_id) commit ordering. + # 0.0 when delays off -> None (no ordering signal, arrival order). + self._modeled_delay_s = _delay_s if _delay_s else None + + # post_train phase starts AFTER the modeled delay: real SLEEPS _delay_s + # above, so stamping here excludes it -> post_train_s is pure + # post-processing, mode-comparable (~0 both modes; sim never slept). + _phase_post_start = time.time() + + # Sim-mode stamps: the modeled round duration is the remainder-wait + # max(real_gpu, D), matching real mode's sleep-the-remainder semantics + # (device wall = D, GPU hidden inside), NOT additive gpu + D. The sct + # (when this update commits on the virtual clock) = the aggregator's + # dispatch stamp (SIM_SEND_TS, read in _fetch_weights) + that duration + + # the optional pre-commit holding leg. _send_grads sends these so the + # aggregator can order updates by _sim_completion_ts. In real mode these + # stay None and the aggregator falls back to arrival order. + # WAN payload-transfer term (up+down): not measurable on localhost, so a + # documented knob left at 0 -- do NOT enable without a real WAN + # measurement. Byte-identical at 0. + _hp = getattr(getattr(self, "config", None), "hyperparameters", None) + _wan_s = ( + float(getattr(_hp, "sim_wan_transfer_s", 0.0) or 0.0) + if (self.simulated and _hp is not None) else 0.0 + ) + # Remainder-wait sct: modeled round wall is max(gpu, delay), NOT gpu + + # delay -- the device's compute is hidden inside its delay budget. The + # per-trainer registry delays supply the completion spread, so the crc32 + # straggler offset is redundant (kept flag-gated at 0 in sim yamls); + # _wan_s stays a documented knob at 0. + self._sim_round_duration_s = ( + max(_real_gpu_time_s, _delay_s) + self._sim_straggler_offset_s() + _wan_s + ) + if self.simulated: + _leg = self.sim_completion_leg_s + _base = self._sim_send_ts if self._sim_send_ts is not None else time.time() + self._sim_completion_ts = _base + self._sim_round_duration_s + _leg logger.info( f"completed training for trainer id: {self.trainer_id}, data_id = {self.data_id}" @@ -535,14 +659,29 @@ def train_with_data_id(self): _stat_utility = float(self._stat_utility) except (TypeError, ValueError): _stat_utility = None + # Post-compute overhead (telemetry build); last trainer-side phase + # term for the wall decomposition. Modeled delay is excluded via the + # _phase_post_start stamp position (after the delay). + _post_train_s = time.time() - _phase_post_start + # Forward-pass / perturbation accounting (WS3-b). Cumulative counters + # live in fwdgrad_utils (per-process = per-client); the delta since the + # last trainer_round is this iteration's cost. jvp_evals == scored + # perturbations. getattr defaults keep the first iteration correct + # without touching __init__. + _fp_total, _jvp_total = fwdgrad_utils.fwd_pass_counts() + _fp_iter = _fp_total - getattr(self, "_fwd_pass_last", 0) + _jvp_iter = _jvp_total - getattr(self, "_jvp_eval_last", 0) + self._fwd_pass_last = _fp_total + self._jvp_eval_last = _jvp_total ev, fields = build_trainer_round( round_num=int(self._round), real_gpu_time_s=_real_gpu_time_s, - # real GPU compute + the emulated delay (0 if disabled) -- - # NOT a budget-vs-actual quantity (fwdllm has no sleep-to- - # fill-budget model, unlike async_cifar10's trainer); this is - # simply the total wall time this round actually took. - sim_round_duration_s=_real_gpu_time_s + _delay_s, + # total modeled wall this round: max(real_gpu, D) under the + # remainder-wait model (device wall = D, GPU hidden) + the sim + # sct-model folds (straggler spread, WAN, both 0 by default). + # training_budget_s below carries D separately so the overrun + # (gpu > D) is recoverable from telemetry. + sim_round_duration_s=self._sim_round_duration_s, avail_state=self.avl_state.value, dataset_size=self.dataset_size, stat_utility=_stat_utility, @@ -550,6 +689,28 @@ def train_with_data_id(self): "data_id": self.data_id, "iteration_per_data_id": self.iteration_per_data_id, "model_version": self._model_version, + # Per-phase wall breakdown: pre/gpu/post are stamped here; + # mqtt_fetch_s + weights_to_{ram,gpu}_s ride in via + # _phase_times (populated in fwdllm_trainer._fetch_weights). + "pre_train_s": _pre_train_s, + "gpu_compute_s": _real_gpu_time_s, + "post_train_s": _post_train_s, + "training_budget_s": _delay_s, + # Remainder-wait model: what real actually slept + whether the + # GPU overran the modeled device budget (an overrun can flip + # update order => cohort_sequence break). + "remaining_time_s": _remaining_s, + "training_overran": _overran, + "trainer_phase": ( + f"{self._round}/{self.data_id}/{self.iteration_per_data_id}" + ), + # WS3-b: forward passes / scored perturbations this iteration + # (+ cumulative). Hardware-independent compute unit for Exp 3. + "forward_passes_iter": _fp_iter, + "forward_passes_total": _fp_total, + "perturbations_iter": _jvp_iter, + "perturbations_total": _jvp_total, + **getattr(self, "_phase_times", {}), }, ) telemetry.emit(ev, **fields) diff --git a/lib/python/examples/fwdllm/trainer/forward_training/fwdgrad_utils.py b/lib/python/examples/fwdllm/trainer/forward_training/fwdgrad_utils.py index b1511388f..34963c5fc 100644 --- a/lib/python/examples/fwdllm/trainer/forward_training/fwdgrad_utils.py +++ b/lib/python/examples/fwdllm/trainer/forward_training/fwdgrad_utils.py @@ -8,6 +8,24 @@ logger = logging.getLogger(__name__) +# --- forward-pass accounting (WS3-b) --------------------------------------- +# Each trainer runs in its own process, so these module-level counters are +# per-client cumulative. The three calculate_jvp* helpers below are the only +# places the functional model is evaluated (a "forward pass"): +# calculate_jvp -> 2 passes (loss + terbulence_loss) = 1 scored perturbation +# calculate_jvp_before_actual_update -> 1 pass +# calculate_jvp_after_actual_update -> 1 pass +# FedSgdTrainer reads fwd_pass_counts() into trainer_round telemetry, giving +# Exp 3 a hardware-independent compute denominator immune to GPU-contention. +# Pure counters: no behavior change, zero cost when unread. +_FWD_PASSES = 0 # total forward passes (func evaluations) this trainer +_JVP_EVALS = 0 # total calculate_jvp() calls (= perturbations scored) + + +def fwd_pass_counts() -> Tuple[int, int]: + """(forward_passes, jvp_evals) cumulative for THIS trainer process.""" + return _FWD_PASSES, _JVP_EVALS + def _get_loss(x: torch.Tensor, t: torch.Tensor, num_classes: int = 10) -> torch.Tensor: """Compute cross-entropy loss. @@ -70,16 +88,32 @@ def functional_get_loss( return _get_loss(y, t, num_classes) -def calculate_jvp(func, params, v): +def calculate_jvp(func, params, v, trainable_idx=None): """ - Calculations Jacobian-vector product using numerical differentiation + Calculations Jacobian-vector product using numerical differentiation. + + trainable_idx (fluxtune perf-opt, simulate_fwdllm.md §L): when given, only + those param indices are perturbed; the rest keep v=0 so `p - h*0 = p` + exactly -> bit-identical to perturbing every param, but skips copying the + frozen backbone twice per perturbation. None => legacy all-param path. """ + global _FWD_PASSES, _JVP_EVALS + _FWD_PASSES += 2 # loss + terbulence_loss forward passes below + _JVP_EVALS += 1 h = 0.01 with torch.no_grad(), autocast(): - loss = func(tuple([params[i] - h * v[i] for i in range(len(params))])) - terbulence_loss = func( - tuple([params[i] + h * v[i] for i in range(len(params))]) - ) + if trainable_idx is None: + minus = tuple([params[i] - h * v[i] for i in range(len(params))]) + plus = tuple([params[i] + h * v[i] for i in range(len(params))]) + else: + minus = list(params) + plus = list(params) + for i in trainable_idx: + minus[i] = params[i] - h * v[i] + plus[i] = params[i] + h * v[i] + minus, plus = tuple(minus), tuple(plus) + loss = func(minus) + terbulence_loss = func(plus) avg_loss = (terbulence_loss + loss) / 2 jvp = (terbulence_loss - loss) / (2 * h) return avg_loss, jvp @@ -89,6 +123,8 @@ def calculate_jvp_after_actual_update(func, params, v, jvp_scalar): """ Calculations Jacobian-vector product using numerical differentiation """ + global _FWD_PASSES + _FWD_PASSES += 1 h = 0.01 # learning rate factor with torch.no_grad(), autocast(): loss = func(tuple([params[i] - h * jvp_scalar * v[i] for i in range(len(params))])) @@ -98,6 +134,8 @@ def calculate_jvp_before_actual_update(func, params): """ Calculations Jacobian-vector product using numerical differentiation """ + global _FWD_PASSES + _FWD_PASSES += 1 with torch.no_grad(), autocast(): loss = func(tuple([params[i] for i in range(len(params))])) return loss diff --git a/lib/python/examples/fwdllm/trainer/forward_training/tc_transformer_trainer_distribute.py b/lib/python/examples/fwdllm/trainer/forward_training/tc_transformer_trainer_distribute.py index eba07e169..11f52f9b2 100755 --- a/lib/python/examples/fwdllm/trainer/forward_training/tc_transformer_trainer_distribute.py +++ b/lib/python/examples/fwdllm/trainer/forward_training/tc_transformer_trainer_distribute.py @@ -205,6 +205,32 @@ def __init__( if self.args.select_perturbation_using_jvp: self.select_perturbation_using_jvp = self.args.select_perturbation_using_jvp + # Number of candidate perturbations sampled per param. Drives the + # forward-pass count: the select_perturbation_using_jvp path does 2 JVP + # passes per perturbation, so N perturbations = ~2N passes (dominant + # fluxtune GPU cost). Default 10 (byte-identical to the historical + # hardcode); a knob so JVP cost can be tuned to keep GPU << the modeled + # mobile delay. Real and sim must use the same value (same config). + try: + self.perturbation_count = int(getattr(self.args, "perturbation_count", 10) or 10) + except (TypeError, ValueError): + self.perturbation_count = 10 + + # Fluxtune JVP perf optimizations (simulate_fwdllm.md §L) — all + # bit-identical to the current grads: + # (a) trainable-only finite difference (skip frozen p-h*0=p); + # (b) skip the 3 diagnostic-only forward passes (loss logging only); + # (c) reuse the selected perturbation's JVP computed in selection. + # Config-gated, default OFF (byte-identical); enabled only in the fluxtune + # yamls (`jvp_perf_opt: true`). Real and sim must match (same config). + self.jvp_perf_opt = bool(getattr(self.args, "jvp_perf_opt", False)) + self._sel_jvp_cache = {} + logging.info( + f"[JVP_PERF_OPT] jvp_perf_opt={self.jvp_perf_opt} " + f"(trainable-only FD + skip diagnostic passes + reuse winner JVP; " + f"all bit-identical — simulate_fwdllm.md §L)" + ) + # var control TODO: It is not layer id it is param id. Distilbert for eg # has only 6 layers. if self.args.model_type == "distilbert": @@ -303,9 +329,9 @@ def _select_optimal_perturbations(self, device, logging_state): if self.grad is not None and v.requires_grad: self.total_rng_iter += 1 shape = v.shape - candidate_v = _randn_wrapper((1 * 10, *shape), device="cpu", generator=self.torch_rng, logging_state=logging_state, param_name=k) + candidate_v = _randn_wrapper((self.perturbation_count, *shape), device="cpu", generator=self.torch_rng, logging_state=logging_state, param_name=k) logging.debug(f"Candidate v - random generation for layer - '{index}' layer shape {candidate_v.shape}") - # torch.randn((1 * 10, *shape), device="cpu", generator=self.torch_rng) + # torch.randn((self.perturbation_count, *shape), device="cpu", generator=self.torch_rng) target_grad = self.grad[index] target_grad = torch.flatten(target_grad) @@ -342,8 +368,8 @@ def _select_optimal_perturbations(device, logging_state): if self.grad is not None and v.requires_grad: self.total_rng_iter += 1 shape = v.shape - candidate_v = _randn_wrapper((1 * 10, *shape), device="cpu", generator=self.torch_rng, logging_state=logging_state, param_name=k) - # torch.randn((1 * 10, *shape), device="cpu", generator=self.torch_rng) + candidate_v = _randn_wrapper((self.perturbation_count, *shape), device="cpu", generator=self.torch_rng, logging_state=logging_state, param_name=k) + # torch.randn((self.perturbation_count, *shape), device="cpu", generator=self.torch_rng) target_grad = self.grad[index] target_grad = torch.flatten(target_grad) @@ -411,7 +437,7 @@ def _select_optimal_perturbations(device, logging_state ): if v.requires_grad: self.total_rng_iter += 1 shape = v.shape - candidate_v = _randn_wrapper((1 * 10, *shape), device="cpu", generator=self.torch_rng, logging_state=logging_state, param_name=k) + candidate_v = _randn_wrapper((self.perturbation_count, *shape), device="cpu", generator=self.torch_rng, logging_state=logging_state, param_name=k) candidate_v = torch.flatten(candidate_v, start_dim=1) logging.info(f"len of candidate_v {len(candidate_v)}") @@ -430,7 +456,7 @@ def _select_optimal_perturbations(device, logging_state ): del candidate_v, target_grad, cos_sim, sorted_indices, shape else: v_buffer[index] = [ - candidate_v[i].reshape(v.shape) for i in range(0, 10) + candidate_v[i].reshape(v.shape) for i in range(0, self.perturbation_count) ] del candidate_v, shape index += 1 @@ -440,11 +466,16 @@ def _select_optimal_perturbations(device, logging_state ): jvp_all_perturbations = [] - for i in range(0,10): + # perf-opt: cache each perturbation's (loss, jvp) so the winner's + # JVP is reused below instead of recomputed (2 fewer passes, §L). + self._sel_jvp_cache = {} + for i in range(0, self.perturbation_count): v_params = _prepare_perturbation_tensors(device, v_buffer, i) loss, jvp = _compute_forward_jvp(device, x, labels, v_params) # logging.info(f"Jvp of option: {jvp}") jvp_all_perturbations.append(jvp) + if self.jvp_perf_opt: + self._sel_jvp_cache[i] = (loss, jvp) logging.info(f"Number of pert and jvps: {len(jvp_all_perturbations)}") @@ -557,7 +588,11 @@ def _compute_forward_jvp(device, x, labels, v_params): t=labels, ) - loss, jvp = calculate_jvp(f, self.params, v_params) + # Perf-opt (fluxtune): perturb only trainable params — bit-identical + # since v=0 on frozen params (p-h*0=p). None => legacy all-param path. + _tidx = ([i for i, p in enumerate(self.params) if p.requires_grad] + if self.jvp_perf_opt else None) + loss, jvp = calculate_jvp(f, self.params, v_params, trainable_idx=_tidx) jvp = jvp.to(device) return loss, jvp @@ -630,11 +665,23 @@ def _accumulate_and_extract_grads(device, jvp, v_params): logging.debug(f"v_params hashes: {[(_calculate_hash(v), v.shape) for v in v_params if v.requires_grad]}") logging.info(f"params hashes: {[(_calculate_hash(p), p.shape) for p in self.params]}") - loss, jvp = _compute_forward_jvp(device, x, labels, v_params) - nonscaled_global_loss = _compute_loss_after_update(device, x, labels, v_params, jvp) - scaled_global_loss = _compute_loss_after_update(device, x, labels, v_params, jvp/15) - loss_before_update = _compute_loss_before_update(device, x, labels, v_params, jvp) - logging.info(f"At trainer: {self.trainer_id} - iteration: {logging_state.get('iteration')} - jvp_magnitude: {jvp} - loss before update: { loss_before_update } - loss after update (not downscaled): {nonscaled_global_loss} - loss after update (down scaled): {scaled_global_loss}") + # perf-opt: reuse the winner's JVP already computed during selection + # (same params + v_params for best_idx -> bit-identical), saving 2 + # passes. Falls back to compute when the cache is absent (cos-sim / sync + # path never ran the selection loop) or best_idx used the carried + # global-best v_params (the `best_idx == -1` branch). + if (self.jvp_perf_opt and best_idx != -1 + and best_idx in getattr(self, "_sel_jvp_cache", {})): + loss, jvp = self._sel_jvp_cache[best_idx] + else: + loss, jvp = _compute_forward_jvp(device, x, labels, v_params) + # 3 diagnostic-only passes: their losses ONLY feed the log below (never + # grads/telemetry), so skip them under perf-opt (bit-identical grads, §L). + if not self.jvp_perf_opt: + nonscaled_global_loss = _compute_loss_after_update(device, x, labels, v_params, jvp) + scaled_global_loss = _compute_loss_after_update(device, x, labels, v_params, jvp/15) + loss_before_update = _compute_loss_before_update(device, x, labels, v_params, jvp) + logging.info(f"At trainer: {self.trainer_id} - iteration: {logging_state.get('iteration')} - jvp_magnitude: {jvp} - loss before update: { loss_before_update } - loss after update (not downscaled): {nonscaled_global_loss} - loss after update (down scaled): {scaled_global_loss}") self.jvp_for_snr_check = abs(jvp) logging.info(f"JVP of the perturbation: {jvp}") diff --git a/lib/python/examples/fwdllm/trainer/main.py b/lib/python/examples/fwdllm/trainer/main.py index bef631a79..713afc993 100644 --- a/lib/python/examples/fwdllm/trainer/main.py +++ b/lib/python/examples/fwdllm/trainer/main.py @@ -52,13 +52,13 @@ def post_complete_message(tc_args): if __name__ == "__main__": config = load_config_from_argv() - # --time_mode is a launcher CLI-only arg (not in config JSON), unconditionally - # appended by TrainerSpawner.spawn_trainer() to every trainer's argv. FedFwd has - # no simulated-clock concept (FedSgdTrainer's _emulate_training_delay() always - # sleeps real wall-clock time) -- parsed here only so the launcher's argv - # injection doesn't crash with "unrecognized argument"; "simulated" is a - # documented no-op for now, deferred future work, not retrofitted in this - # migration. log_level is similarly launcher/manual-run CLI-only. + # --time_mode is a launcher CLI-only arg (not in config JSON), appended by + # TrainerSpawner.spawn_trainer() to every trainer's argv. Threaded into + # config.hyperparameters so FedSGDTrainer reads it uniformly with the + # aggregator. "simulated": skip the emulated-delay sleep and stamp a modeled + # completion timestamp the aggregator orders updates by; "real" (default): + # unchanged wall-clock behavior. log_level/battery_threshold are similarly + # launcher/manual-run CLI-only. _cli_parser = argparse.ArgumentParser(add_help=False) _cli_parser.add_argument("--time_mode", default="real") _cli_parser.add_argument("--log_level", default="INFO") @@ -72,13 +72,30 @@ def post_complete_message(tc_args): ) logging.debug(config) telemetry.configure(role="trainer", end_id=str(config.task_id)) - if _cli_args.time_mode != "real": - logging.warning( - f"--time_mode={_cli_args.time_mode!r} requested, but FedFwd has no " - "simulated-clock support yet; running in real wall-clock mode." - ) + config.hyperparameters.time_mode = _cli_args.time_mode set_seed(config.hyperparameters.manual_seed) + # Pinning self-report: confirm the CPU/GPU affinity the spawner intended + # (CUDA_VISIBLE_DEVICES + os.sched_setaffinity per trainer) took effect in + # this child. The trainer trains on cuda:0 of the CVD-masked single-GPU + # view, so one visible device here == correct pinning. Emitted as a grep-able + # [PIN] line so post-proc can verify balanced trainer->(gpu,core) placement. + try: + _cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") + _cores = sorted(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else [] + _gpu_name, _gpu_dev, _ndev = "", -1, 0 + if torch.cuda.is_available(): + _gpu_dev = torch.cuda.current_device() + _gpu_name = torch.cuda.get_device_name(_gpu_dev) + _ndev = torch.cuda.device_count() + logging.info( + f"[PIN] pid={os.getpid()} client_idx={config.hyperparameters.client_idx} " + f"CUDA_VISIBLE_DEVICES={_cvd} cuda_device={_gpu_dev} ({_gpu_name}) " + f"cuda_device_count={_ndev} cpu_affinity={_cores}" + ) + except Exception as _pin_exc: # never let self-report break a trainer + logging.warning(f"[PIN] self-report failed: {_pin_exc}") + # dataset attributes attributes = BaseDataManager.load_attributes(config.hyperparameters.data_file_path) num_labels = len(attributes["label_vocab"]) @@ -121,6 +138,13 @@ def post_complete_message(tc_args): "var_control": config.hyperparameters.var_control, "perturbation_sampling": config.hyperparameters.perturbation_sampling, "select_perturbation_using_jvp": config.hyperparameters.select_perturbation_using_jvp, + # forward-pass count knob (default 10 = historical behavior). + "perturbation_count": getattr( + config.hyperparameters, "perturbation_count", 10), + # fluxtune JVP perf-opt (§L, bit-identical). Default False = + # byte-identical; enabled only in the fluxtune yamls. + "jvp_perf_opt": getattr( + config.hyperparameters, "jvp_perf_opt", False), } ) model_args.config["num_labels"] = num_labels diff --git a/lib/python/examples/scripts/converge_watch.py b/lib/python/examples/scripts/converge_watch.py new file mode 100644 index 000000000..68202b2ca --- /dev/null +++ b/lib/python/examples/scripts/converge_watch.py @@ -0,0 +1,302 @@ +#!/usr/bin/env python3 +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Convergence-stop watcher (EXPERIMENTS.md). + +Polls the aggregator telemetry of the run that `expt_launch` just started and +terminates it the moment the convergence condition is met: + + the last W consecutive data bins are ALL >= target accuracy tau. + +Data bins complete in order (data_id = 0,1,2,...) and emit exactly one +``agg_eval`` per bin at completion, so "W continuous bins all >= tau" +is exactly a trailing run of W consecutive-data_id bins each with test-accuracy +>= tau. We track the trailing streak; a bin below tau resets it. + +On convergence we write ``converge.json`` (time-to-converge in wall + vclock + +data_id + round) and SIGTERM->SIGKILL the run's process group. If the run exits +on its own first (hit a safety cap / crashed), the pgid vanishes and we exit +WITHOUT writing converge.json -> the harness reports DID_NOT_CONVERGE. + +Entirely side-car: it only reads telemetry and signals the process group. When +``--target-acc`` is not passed the harness never launches this, so default runs +are byte-identical. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import signal +import sys +import time + + +def _newest_agg_jsonl(exp_dir: str, marker: str) -> str | None: + """Newest run_*/telemetry/aggregator_*.jsonl newer than the launch marker.""" + try: + marker_mtime = os.path.getmtime(marker) + except OSError: + marker_mtime = 0.0 + best, best_mtime = None, marker_mtime - 1.0 + for run_dir in glob.glob(os.path.join(exp_dir, "run_*")): + try: + if os.path.getmtime(run_dir) < marker_mtime - 5.0: + continue # pre-existing run dir, not ours + except OSError: + continue + for f in glob.glob(os.path.join(run_dir, "telemetry", "aggregator_*.jsonl")): + try: + m = os.path.getmtime(f) + except OSError: + continue + if m > best_mtime: + best, best_mtime = f, m + return best + + +def _scan(agg_path: str): + """Return (data_id->accuracy, data_id->loss, latest_vclock, max_round).""" + acc_by_bin: dict[int, float] = {} + loss_by_bin: dict[int, float] = {} + latest_vclock = None + max_round = 0 + with open(agg_path, "r", encoding="utf-8") as fh: + for line in fh: + line = line.strip() + if not line: + continue + try: + r = json.loads(line) + except json.JSONDecodeError: + continue # partial trailing line from a line-buffered writer + ev = r.get("event") + if r.get("round") is not None: + try: + max_round = max(max_round, int(r["round"])) + except (TypeError, ValueError): + pass + if ev == "agg_eval": + d = r.get("data_id") + a = r.get("test-accuracy") + lo = r.get("test-loss") + if d is not None and a is not None: + acc_by_bin[int(d)] = float(a) # last eval for the bin wins + if d is not None and lo is not None: + loss_by_bin[int(d)] = float(lo) + elif ev == "agg_round": + v = r.get("vclock_now") + if v is not None: + latest_vclock = float(v) + return acc_by_bin, loss_by_bin, latest_vclock, max_round + + +def _trailing_streak(acc_by_bin: dict[int, float], target: float): + """Length of the trailing run of consecutive data_ids all >= target, and the + first data_id in that run. Bins must be contiguous (no gap) to count.""" + if not acc_by_bin: + return 0, None + top = max(acc_by_bin) + streak, first = 0, None + d = top + while d in acc_by_bin and acc_by_bin[d] >= target: + streak += 1 + first = d + d -= 1 + return streak, first + + +def _pgid_alive(pgid: int) -> bool: + try: + os.killpg(pgid, 0) + return True + except (ProcessLookupError, PermissionError): + return False + except OSError: + return False + + +def _terminate(pgid: int, grace_s: float) -> None: + for sig in (signal.SIGTERM, signal.SIGKILL): + try: + os.killpg(pgid, sig) + except OSError: + return + if sig is signal.SIGTERM: + deadline = time.time() + grace_s + while time.time() < deadline and _pgid_alive(pgid): + time.sleep(1.0) + if not _pgid_alive(pgid): + return + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--exp-dir", required=True, help="/experiments") + ap.add_argument("--marker", required=True, help="expt_launch pre-launch marker file") + ap.add_argument("--target-acc", type=float, required=True) + ap.add_argument("--window", type=int, default=20) + ap.add_argument("--pgid", type=int, required=True, help="run process-group id to kill on convergence") + ap.add_argument("--converge-json", required=True) + ap.add_argument("--poll", type=float, default=15.0) + ap.add_argument("--grace", type=float, default=20.0) + # Stall guard: terminate early (before the wall ceiling) if the run is clearly + # not learning — no PROGRESS within any --stall-window-s wall window. 0 = off. + # Progress is measured on accuracy, loss, or EITHER (default), because a run + # can plateau in accuracy while test-loss keeps falling (still learning). + ap.add_argument("--stall-window-s", type=float, default=0.0, + help="terminate if no progress in this many wall s (0=off)") + ap.add_argument("--stall-min-delta", type=float, default=0.01, + help="ABSOLUTE accuracy gain that counts as progress (default 0.01 = 1%%)") + ap.add_argument("--stall-on", choices=("acc", "loss", "either"), default="either", + help="which signal resets the idle clock (default: either)") + ap.add_argument("--loss-min-rel-delta", type=float, default=0.01, + help="RELATIVE test-loss drop vs running-best that counts as progress " + "(default 0.01 = 1%%; loss is unbounded so it is relative, not absolute)") + ap.add_argument("--stall-json", default=None, help="written on a stall termination") + args = ap.parse_args() + + try: + start_wall = os.path.getmtime(args.marker) + except OSError: + start_wall = time.time() + + label = os.path.basename(args.converge_json) + _prog_desc = { + "acc": f"<{args.stall_min_delta:.3f} acc gain", + "loss": f"<{args.loss_min_rel_delta*100:.1f}% loss drop", + "either": f"<{args.stall_min_delta:.3f} acc gain AND <{args.loss_min_rel_delta*100:.1f}% loss drop", + }[args.stall_on] + _stall_desc = (f"; stall if {_prog_desc} in {args.stall_window_s/3600:.1f}h" + if args.stall_window_s > 0 else "") + print(f" [converge] watching for {args.window} consecutive bins >= {args.target_acc:.4f} acc" + f"{_stall_desc}", flush=True) + + # Stall tracking: milestones = best acc (max) / best loss (min) at the last + # recorded improvement; last_improve_ts = when it happened. Progress on the + # armed signal(s) resets the clock; the clock also runs from launch, so a run + # that produces NO eval (or no progress) within the window is caught too. + # acc progress: best_acc - milestone_acc >= stall_min_delta (ABSOLUTE — acc is in [0,1]) + # loss progress: (milestone_loss - best_loss)/milestone_loss >= loss_min_rel_delta + # (RELATIVE vs running-best — loss is unbounded/scale-dependent, so a + # fractional drop is scale-free and diminishing-returns aware) + # Both use the running best (max acc / min loss), so a single noisy eval can + # neither reset the clock nor fake progress. + milestone_acc = None + milestone_loss = None + last_improve_ts = time.time() + + while _pgid_alive(args.pgid): + agg_path = _newest_agg_jsonl(args.exp_dir, args.marker) + if agg_path is None: + time.sleep(args.poll) + continue + try: + acc_by_bin, loss_by_bin, vclock, max_round = _scan(agg_path) + except OSError: + time.sleep(args.poll) + continue + streak, first = _trailing_streak(acc_by_bin, args.target_acc) + best = max(acc_by_bin.values()) if acc_by_bin else None + best_loss = min(loss_by_bin.values()) if loss_by_bin else None + now = time.time() + # progress bookkeeping for the stall guard (per the armed signal). The + # milestone advances ONLY on a >=threshold move, so sub-threshold steps + # accumulate against a FIXED milestone until they cross it (a slow steady + # climb keeps resetting the clock; a true plateau does not). + acc_progress = ( + args.stall_on in ("acc", "either") and best is not None + and (milestone_acc is None or best - milestone_acc >= args.stall_min_delta)) + loss_progress = ( + args.stall_on in ("loss", "either") and best_loss is not None + and (milestone_loss is None + # relative drop vs running-best (test-loss is > 0 for real CE); if a + # milestone is somehow <= 0, fall back to any absolute improvement. + or (milestone_loss > 0 + and (milestone_loss - best_loss) / milestone_loss >= args.loss_min_rel_delta) + or (milestone_loss <= 0 and best_loss < milestone_loss))) + if acc_progress: + milestone_acc = best + if loss_progress: + milestone_loss = best_loss + if acc_progress or loss_progress: + last_improve_ts = now + no_improve_s = now - last_improve_ts + if acc_by_bin: + top = max(acc_by_bin) + _stall_note = (f" no_improve={no_improve_s/60:.0f}m/{args.stall_window_s/60:.0f}m" + if args.stall_window_s > 0 else "") + _loss_note = f" best_loss={best_loss:.4f}" if best_loss is not None else "" + print(f" [converge] bins={len(acc_by_bin)} top_bin={top} " + f"streak>={args.target_acc:.3f}={streak}/{args.window} " + f"best_acc={best:.4f}{_loss_note}{_stall_note}", flush=True) + if streak >= args.window: + wall_s = time.time() - start_wall + payload = { + "converged": True, + "target_accuracy": args.target_acc, + "window": args.window, + "trigger_data_id": max(acc_by_bin), + "first_bin_in_window": first, + "n_bins_completed": len(acc_by_bin), + "time_to_converge_wall_s": round(wall_s, 3), + "time_to_converge_vclock_s": vclock, + "rounds_at_converge": max_round, + "agg_telemetry": agg_path, + "window_accuracies": [acc_by_bin[b] for b in + range(first, max(acc_by_bin) + 1)], + } + try: + with open(args.converge_json, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + except OSError as e: + print(f" [converge] WARN could not write {args.converge_json}: {e}", flush=True) + print(f" [{label}] CONVERGED at data_id={max(acc_by_bin)} " + f"(wall={wall_s:.0f}s vclock={vclock}) -> terminating run", flush=True) + _terminate(args.pgid, args.grace) + return 0 + # Stall termination: no progress on the armed signal(s) within the window + # (and NOT converged). + if args.stall_window_s > 0 and no_improve_s >= args.stall_window_s: + _reason = {"acc": "no_accuracy_improvement", "loss": "no_loss_improvement", + "either": "no_accuracy_or_loss_improvement"}[args.stall_on] + payload = { + "stalled": True, + "reason": _reason, + "stall_on": args.stall_on, + "stall_window_s": args.stall_window_s, + "stall_min_delta": args.stall_min_delta, + "loss_min_rel_delta": args.loss_min_rel_delta, + "no_improve_s": round(no_improve_s, 1), + "best_accuracy": best, + "milestone_accuracy": milestone_acc, + "best_loss": best_loss, + "milestone_loss": milestone_loss, + "target_accuracy": args.target_acc, + "n_bins_completed": len(acc_by_bin), + "wall_s": round(now - start_wall, 1), + "rounds": max_round, + "agg_telemetry": agg_path, + } + if args.stall_json: + try: + with open(args.stall_json, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + except OSError as e: + print(f" [converge] WARN could not write {args.stall_json}: {e}", flush=True) + _bl = f" best_loss={best_loss:.4f}" if best_loss is not None else "" + print(f" [{label}] STALLED [{args.stall_on}]: best_acc={best}{_bl} — no progress " + f"({_prog_desc}) in {args.stall_window_s/3600:.1f}h -> terminating run (not learning)", + flush=True) + _terminate(args.pgid, args.grace) + return 0 + time.sleep(args.poll) + + print(f" [converge] run process group ended before convergence", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/lib/python/examples/scripts/expt_runner.py b/lib/python/examples/scripts/expt_runner.py new file mode 100644 index 000000000..f7dbbbdb5 --- /dev/null +++ b/lib/python/examples/scripts/expt_runner.py @@ -0,0 +1,285 @@ +"""Shared experiment-launch display + pre-flight gate for the flame examples. + +Generic and example-agnostic: it knows how to *render* a tiered hyperparameter +table and *gate* on a list of feasibility checks, but nothing about any specific +example's knobs. Each example's driver (fwdllm `run_sequential.sh`, +async_cifar10 `debug_run.sh`) builds a `spec` dict describing its own tiers and +checks and calls `render_and_gate(spec)`. + +Design goal (why this exists): a live GPU/MQTT run is expensive, so before we +fire one we (a) show the operator the hyperparameters organized by how often +they change -- ① review-every-run, ② per-baseline, ③ config-baked -- and +(b) run correctness checks that block infeasible configs. Colour/callout markers +draw the eye to anything overridden or dangerous. + +Spec schema +----------- + spec = { + "title": str, # e.g. "FWDLLM RUN" + "subtitle": str, # e.g. "mode=both baselines=fwdllm ..." + "dry_run": bool, # cosmetic tag in the header + "tiers": [ + {"name": "① REVIEW EVERY RUN", + "collapsed": False, # tier ③ sets True -> hidden unless EXPT_SHOW_ALL=1 + "bar": "44", # optional header-bar bg colour code (default cycles) + "rows": [ # label/value rows … + {"label": "stop", "value": "max_runtime_s=600 max_data_id=10"}, + {"label": "delays", "value": "D=0", "level": "warn", + "note": "both sides matched"}, + ]}, + {"name": "② PER-BASELINE", # … OR an aligned table (mutually exclusive with rows) + "table": { + "columns": ["c", "agg_goal", ("min_init", "minInit")], # key or (key, header) + "rows": [{"name": "fluxtune", "cells": {"c": 10, "agg_goal": 3}}], + "overridden": ["c"], # optional: columns set by a CLI flag (green •) + }}, + ], + "checks": [ + {"name": "D matched across real/sim pair", "level": "ok", + "detail": "..."}, + {"name": "agg_goal <= c (sync)", "level": "error", + "detail": "fwdllm_plus agg_goal=5 > c=2"}, + ], + "next": "python -m scripts.parity.cli --batch ...", # optional hand-off + } + +`level` is one of "ok" (default), "warn", "error". `render_and_gate` returns 2 if +any check (or any *row*) is level "error", else 0 -- the driver treats 2 as +"blocked unless --force". +""" + +from __future__ import annotations + +import json +import textwrap +import os +import sys + +# ---- colour / icon helpers ------------------------------------------------- + +def _use_colour(stream) -> bool: + if os.environ.get("NO_COLOR") is not None: + return False + if os.environ.get("EXPT_FORCE_COLOR") is not None: + return True + return hasattr(stream, "isatty") and stream.isatty() + + +class _Style: + def __init__(self, on: bool): + self.on = on + + def _w(self, code: str, s: str) -> str: + return f"\033[{code}m{s}\033[0m" if self.on else s + + def red(self, s): return self._w("1;31", s) + def yellow(self, s): return self._w("1;33", s) + def green(self, s): return self._w("32", s) + def cyan(self, s): return self._w("1;36", s) + def dim(self, s): return self._w("2", s) + def bold(self, s): return self._w("1", s) + + def bar(self, s: str, width: int, code: str = "44") -> str: + """A solid-background full-width header bar (bold bright-white on `code`), + so the ①②③ section markers read as distinct blocks, not plain text.""" + txt = (" " + s).ljust(width) + return self._w(f"1;97;{code}", txt) + + +# level -> (icon, colouriser name). "ok" rows stay quiet; set/warn/error shout. +# set = value explicitly overridden by a command-line flag (overrides yaml) -> green +# warn = review / attention -> yellow +# error = infeasible -> red +# NB: the emoji markers render 2 cells wide, so "ok" uses TWO spaces to keep the +# label column aligned with the set/warn/error rows. +_ICON = {"ok": " ", "set": "\U0001f7e2", "warn": "\U0001f7e1", "error": "\U0001f534"} # 🟢 🟡 🔴 +_CHECK_ICON = {"ok": "✓", "warn": "⚠", "error": "✗"} # ✓ ⚠ ✗ + + +def _colour_for(st: _Style, level: str): + return {"ok": st.dim, "set": st.green, "warn": st.yellow, "error": st.red}.get(level, st.dim) + + +# ---- renderer -------------------------------------------------------------- + +_RULE = "─" * 79 # ───── +_BAR_W = 79 # header-bar width (matches the rule) + +# Per-tier header-bar background colours so ①②③ are visually distinct blocks. +_TIER_BAR = ["44", "45", "100"] # blue, magenta, bright-black(grey) + + +def _cell(v) -> str: + """Table cell text: None -> '–', everything else str().""" + return "–" if v is None else str(v) + + +def _render_table(out, st, table: dict) -> None: + """Render an aligned per-entity table (tier ②'s per-baseline knobs). + + table = { + "columns": [key | (key, header), ...], # column order + "rows": [{"name": str, "cells": {key: value}}, ...], + "overridden": [key, ...], # optional: flag-overridden cols + } + + Columns whose value is NOT identical across every row are HIGHLIGHTED (bold + yellow header + cells) -- those are the knobs that differ between baselines + and must be eyeballed. Columns identical across all rows stay dim (expected). + Row (baseline) names are cyan. A flag-overridden column gets a green '•'. + """ + cols = [(c, c) if isinstance(c, str) else (c[0], c[1]) for c in table["columns"]] + rows = table["rows"] + overridden = set(table.get("overridden", [])) + + # Which columns differ across baselines? + differs = {} + for key, _h in cols: + vals = {_cell(r["cells"].get(key)) for r in rows} + differs[key] = len(vals) > 1 + + name_w = max([len("baseline")] + [len(_cell(r["name"])) for r in rows]) + col_w = {} + for key, hdr in cols: + col_w[key] = max(len(hdr), max((len(_cell(r["cells"].get(key))) for r in rows), + default=0)) + + # header row + hcells = [st.dim("baseline".ljust(name_w))] + for key, hdr in cols: + mark = st.green("•") if key in overridden else " " + htxt = hdr.ljust(col_w[key]) + hcells.append((st.yellow(htxt) if differs[key] else st.dim(htxt)) + mark) + out(" " + " ".join(hcells)) + + # data rows + for r in rows: + line = [st.cyan(_cell(r["name"]).ljust(name_w))] + for key, hdr in cols: + v = _cell(r["cells"].get(key)).ljust(col_w[key]) + line.append((st.yellow(v) if differs[key] else st.dim(v)) + " ") + out(" " + " ".join(line)) + # legend for the highlighting + diff_names = [h for k, h in cols if differs[k]] + legend = st.yellow("yellow") + st.dim(" = differs across baselines (review)") + if overridden: + legend += st.dim(" ") + st.green("•") + st.dim(" = flag override") + out(" " + st.dim("· ") + legend + + (st.dim(f" [{', '.join(diff_names)}]") if diff_names else "")) + + +def render_and_gate(spec: dict, show_all: bool | None = None, stream=None) -> int: + """Render the tiered table + checks; return 2 if any error, else 0. + + Rows/checks at level "error" mark the config infeasible; the driver should + refuse to launch (unless the operator passes --force). + """ + stream = stream or sys.stdout + st = _Style(_use_colour(stream)) + if show_all is None: + show_all = os.environ.get("EXPT_SHOW_ALL") not in (None, "", "0") + + def out(s=""): + print(s, file=stream) + + title = spec.get("title", "EXPERIMENT RUN") + subtitle = spec.get("subtitle", "") + tag = st.yellow("[DRY-RUN]") if spec.get("dry_run") else "" + out() + out(f" {st.bold(title)} {subtitle} {tag}".rstrip()) + out(" " + _RULE) + + n_err_rows = 0 + for ti, tier in enumerate(spec.get("tiers", [])): + name = tier.get("name", "") + rows = tier.get("rows", []) + table = tier.get("table") + collapsed = tier.get("collapsed", False) + bar_code = tier.get("bar", _TIER_BAR[ti % len(_TIER_BAR)]) + out(" " + st.bar(name, _BAR_W, bar_code)) + n_hidden = len(table["rows"]) if table else len(rows) + if collapsed and not show_all: + out(st.dim(f" [{n_hidden} row(s) hidden — set EXPT_SHOW_ALL=1]")) + continue + if table: + _render_table(out, st, table) + continue + # Three aligned columns: label | value | note. Labels pad to the widest + # label; the note lives in its OWN column and WRAPS to indented + # continuation lines (bounded), instead of trailing unbounded on one + # line. Value width is sized to the note-bearing rows so their notes all + # start at the same column (a long note-less value, e.g. the baselines + # list, may overflow harmlessly since nothing follows it). + label_w = min(max((len(r.get("label", "")) for r in rows), default=13), 24) + _noted = [r for r in rows if r.get("note")] + value_w = min(max((len(str(r.get("value", ""))) for r in _noted), default=0), 20) + # cell column at which the note starts: 2 lead + 2 icon + 1 + label + 1 + value + 1 + note_col = 2 + 2 + 1 + label_w + 1 + value_w + 1 + note_w = max(24, _BAR_W - note_col) + cont_pad = " " * (note_col + 2) # continuation lines align under the note text + for r in rows: + level = r.get("level", "ok") + if level == "error": + n_err_rows += 1 + icon = _ICON.get(level, " ") + col = _colour_for(st, level) + raw = r.get("label", "") + label = raw.ljust(label_w) + value = str(r.get("value", "")) + valpad = value.ljust(value_w) + # ok rows: label dim, value plain. warn/error: label bold, value coloured. + lbl_s = st.dim(label) if level == "ok" else st.bold(label) + val_s = valpad if level == "ok" else col(valpad) + note = r.get("note", "") + if not note: + out(f" {icon} {lbl_s} {val_s}".rstrip()) + continue + wrapped = textwrap.wrap(note, note_w) or [""] + out(f" {icon} {lbl_s} {val_s} {st.dim('· ' + wrapped[0])}") + for cont in wrapped[1:]: + out(cont_pad + st.dim(cont)) + + out(" " + _RULE) + + # ---- checks / gate ---- + checks = spec.get("checks", []) + n_ok = sum(1 for c in checks if c.get("level", "ok") == "ok") + n_warn = sum(1 for c in checks if c.get("level") == "warn") + n_err = sum(1 for c in checks if c.get("level") == "error") + summary = (f" PRE-FLIGHT: {len(checks)} check(s) … " + f"{st.green(str(n_ok) + ' ✓')} " + f"{st.yellow(str(n_warn) + ' ⚠')} " + f"{st.red(str(n_err) + ' ✗')}") + out(summary) + for c in checks: + level = c.get("level", "ok") + if level == "ok" and not show_all: + continue # keep the passing checks quiet unless asked + icon = _CHECK_ICON.get(level, "?") + col = _colour_for(st, level) + detail = c.get("detail", "") + detail_s = f" {st.dim('— ' + detail)}" if detail else "" + out(f" {col(icon)} {c.get('name', '')}{detail_s}") + + blocked = (n_err + n_err_rows) > 0 + if blocked: + out(" " + st.red("BLOCKED: infeasible config — fix, or re-run with --force to override.")) + if spec.get("next"): + out(" " + st.bold("NEXT: ") + spec["next"]) + out() + return 2 if blocked else 0 + + +def main(argv=None) -> int: + """CLI form: `expt_runner.py ` -> render + gate, exit 0/2.""" + argv = argv if argv is not None else sys.argv[1:] + if not argv: + print("usage: expt_runner.py ", file=sys.stderr) + return 2 + with open(argv[0], encoding="utf-8") as fh: + spec = json.load(fh) + return render_and_gate(spec) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/lib/python/examples/scripts/expt_runner.sh b/lib/python/examples/scripts/expt_runner.sh new file mode 100644 index 000000000..bb7c2ae06 --- /dev/null +++ b/lib/python/examples/scripts/expt_runner.sh @@ -0,0 +1,375 @@ +#!/bin/bash +# Shared experiment-launch harness for the flame examples: robust conda activation, +# PYTHONPATH pinning, the launch+progress-ticker loop, and post-run log-health +# assertions. Sourced by each driver (async_cifar10 debug_run.sh, fwdllm +# run_sequential.sh); config discovery / YAML patching / knob sets stay per-driver. +# +# Usage (from a driver): +# source "/lib/python/examples/scripts/expt_runner.sh" +# expt_activate_conda [default_env] # activate + report python +# expt_pin_pythonpath "$REPO_ROOT" # this checkout's flame ahead of any editable install +# expt_launch label cfg example_dir budget_s n_exps logdir [experiments_dir] +# expt_assert_log "$logdir/label.out" label +# +# Pre-flight display + gate lives on the Python side (expt_runner.py): a driver +# builds a `spec` of tiers+checks and calls render_and_gate(). $EXPT_RUNNER_DIR / +# $EXPT_RUNNER_PY are exported so a driver's python heredoc can import expt_runner. + +# Absolute dir of THIS harness (works when sourced), so drivers/python can find +# the Python renderer next to it regardless of the caller's cwd. +EXPT_RUNNER_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +EXPT_RUNNER_PY="$EXPT_RUNNER_DIR/expt_runner.py" +export EXPT_RUNNER_DIR EXPT_RUNNER_PY + +# --- robust conda activation ------------------------------------------------ +# Env choice: FLAME_CONDA_ENV overrides; else the shell's already-active env +# (CONDA_DEFAULT_ENV); else the caller-supplied default (arg $1). No hardcoded +# fallback -- cifar passes dg_flame, fwdllm passes nothing (requires an active env). +expt_activate_conda() { + local default_env="${1:-}" + local envname="${FLAME_CONDA_ENV:-${CONDA_DEFAULT_ENV:-$default_env}}" + if [ -z "$envname" ]; then + echo "ERROR: no conda env active and FLAME_CONDA_ENV not set." >&2 + echo " Activate an env first (conda activate ), set FLAME_CONDA_ENV=," >&2 + echo " or have the driver pass a default to expt_activate_conda." >&2 + return 1 + fi + local cb="" + if command -v conda >/dev/null 2>&1; then + cb="$(conda info --base 2>/dev/null)" + elif [ -n "${CONDA_EXE:-}" ]; then + cb="$(dirname "$(dirname "$CONDA_EXE")")" + fi + if [ -z "$cb" ] || [ ! -f "$cb/etc/profile.d/conda.sh" ]; then + local c + for c in "$HOME/miniconda3" "/coc/scratch/${USER%??}/miniconda3" \ + "/coc/scratch/$USER/miniconda3" "$HOME/anaconda3" /opt/conda; do + [ -f "$c/etc/profile.d/conda.sh" ] && cb="$c" && break + done + fi + if [ -z "$cb" ] || [ ! -f "$cb/etc/profile.d/conda.sh" ]; then + echo "ERROR: conda not found. Activate '$envname' yourself or set CONDA_EXE." >&2 + return 1 + fi + # shellcheck disable=SC1091 + source "$cb/etc/profile.d/conda.sh" + conda activate "$envname" || { echo "ERROR: 'conda activate $envname' failed" >&2; return 1; } + echo "conda: base=$cb env=$envname python=$(which python)" +} + +# Force this checkout's flame package ahead of anything already on sys.path +# (e.g. a stale `pip install -e` editable pointing at a different clone). +expt_pin_pythonpath() { + local repo_root="$1" + export PYTHONPATH="$repo_root/lib/python${PYTHONPATH:+:$PYTHONPATH}" +} + +# FL worker process patterns — single source of truth for interrupt teardown and +# the clean-slate preflight (run orchestrator, trainers, aggregator, watcher). +EXPT_WORKER_PATS=( + 'flame.launch.run_experiment' + 'trainer/forward_training' + 'trainer/pytorch/main.py' + 'aggregator/pytorch/main_' + 'converge_watch.py' +) + +# expt_assert_clean_slate [label] -- refuse to launch on top of stray FL workers +# from a prior/crashed/Ctrl+C'd run (own procs only). EXPT_AUTOCLEAN=1 kills them +# and re-checks; default ABORTS and prints the kill command. Returns 1 if dirty. +# EXPT_GPU_FREE_MB (default 500) warns on residual GPU memory; EXPT_GPU_STRICT=1 aborts. +expt_assert_clean_slate() { + local label="${1:-preflight}" uid; uid="$(id -u)" + _ecs_scan() { + local p pids out="" + for p in "${EXPT_WORKER_PATS[@]}"; do + pids="$(pgrep -u "$uid" -f "$p" 2>/dev/null | tr '\n' ' ')" + [ -n "$pids" ] && out+=" ${p} -> ${pids}\n" + done + printf '%b' "$out" + } + local dirty; dirty="$(_ecs_scan)" + if [ -n "$dirty" ]; then + echo " [$label] NOT CLEAN — stray FL workers from a previous run:" >&2 + printf '%b' "$dirty" >&2 + if [ "${EXPT_AUTOCLEAN:-0}" = "1" ]; then + echo " [$label] EXPT_AUTOCLEAN=1 → killing and re-checking ..." >&2 + # TERM the orchestrator first (lets it tear down its own group), then KILL all. + pkill -TERM -u "$uid" -f 'flame.launch.run_experiment' 2>/dev/null || true + sleep "${EXPT_CLEAN_GRACE_S:-3}" + local p + for p in "${EXPT_WORKER_PATS[@]}"; do + pkill -9 -u "$uid" -f "$p" 2>/dev/null || true + done + sleep 2 + dirty="$(_ecs_scan)" + if [ -n "$dirty" ]; then + echo " [$label] STILL NOT CLEAN after autoclean — aborting:" >&2 + printf '%b' "$dirty" >&2; return 1 + fi + echo " [$label] cleaned." >&2 + else + echo " [$label] refusing to launch. Clean it, or re-run with EXPT_AUTOCLEAN=1 (or --clean):" >&2 + echo " pkill -TERM -f 'flame.launch.run_experiment'; sleep 3; pkill -9 -f 'trainer/forward_training'; pkill -9 -f 'trainer/pytorch/main.py'; pkill -9 -f 'aggregator/pytorch/main_'; pkill -9 -f converge_watch.py" >&2 + return 1 + fi + fi + # Residual GPU memory here is a peer's job (your workers are gone) -> warn only. + if command -v nvidia-smi >/dev/null 2>&1; then + local maxused; maxused="$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | sort -n | tail -1)" + local thresh="${EXPT_GPU_FREE_MB:-500}" + if [ -n "$maxused" ] && [ "$maxused" -gt "$thresh" ] 2>/dev/null; then + echo " [$label] WARNING: a GPU shows ${maxused}MB used (> ${thresh}MB) — a peer job may be resident." >&2 + if [ "${EXPT_GPU_STRICT:-0}" = "1" ]; then + echo " [$label] EXPT_GPU_STRICT=1 → aborting." >&2; return 1 + fi + fi + fi + echo " [$label] clean slate verified." >&2 + return 0 +} + +# expt_launch label cfg example_dir budget_s n_exps logdir [experiments_dir] +# Runs one flame experiment-config file in the foreground with a 30s progress +# ticker (elapsed/remaining/percent + how many run_* dirs have appeared). Sets +# EXPT_LAST_RC and returns the launcher's exit code. +expt_launch() { + local label="$1" cfg="$2" example_dir="$3" budget_s="${4:-0}" n_exps="${5:-1}" logdir="$6" + local exp_dir="${7:-$example_dir/experiments}" + mkdir -p "$logdir" + # Marker touched BEFORE launch so expt_assert_run can find exactly the run + # dir(s)/logs this launch produced (find -newer "$marker"). Exported for the + # caller's health check. + EXPT_LAST_MARKER="$logdir/.marker_${label}"; : > "$EXPT_LAST_MARKER"; export EXPT_LAST_MARKER + local start_ts; start_ts=$(date +%s) + local initial_runs; initial_runs=$(find "$exp_dir" -maxdepth 1 -name "run_*" -type d 2>/dev/null | wc -l) + + echo "[$(date '+%F %T')] START $label ($n_exps exp(s), ~${budget_s}s budget)" | tee -a "$logdir/expt_runner.log" + + # Launch the run in its OWN process group (set -m -> bg job's pgid == its pid; + # run_experiment's Popen children inherit it, no setsid — same pattern as + # expt_timed_run) so the convergence watcher can signal the whole tree. + set -m + python -m flame.launch.run_experiment "$cfg" --example-dir "$example_dir" \ + < /dev/null >> "$logdir/${label}.out" 2>&1 & + local run_pid=$! + set +m + + # Progress ticker. SELF-TERMINATES the instant the run process is gone, so it + # never outlives the run (watcher kill, crash, Ctrl+C). Torn down below BY PID + # (+ its in-flight `sleep` child), NEVER by process group: job control doesn't + # reliably place a backgrounded subshell in a fresh group here, so a group + # signal would miss it and a following `wait` would hang the harness. + ( + while kill -0 "$run_pid" 2>/dev/null; do + sleep 30 + kill -0 "$run_pid" 2>/dev/null || break # run gone -> stop; never outlive it + local now elapsed pct=0 + now=$(date +%s); elapsed=$(( now - start_ts )) + if [ "$budget_s" -gt 0 ]; then + pct=$(( elapsed * 100 / budget_s )); [ "$pct" -gt 100 ] && pct=100 + fi + local curr started + curr=$(find "$exp_dir" -maxdepth 1 -name "run_*" -type d 2>/dev/null | wc -l) + started=$(( curr - initial_runs )); [ "$started" -lt 0 ] && started=0 + printf " [%s] %s | %ds elapsed / ~%ds (%d%%) | exp started: %d/%d\n" \ + "$(date '+%T')" "$label" "$elapsed" "$budget_s" "$pct" "$started" "$n_exps" + done + ) & + local ticker_pid=$! + + # PID-based ticker teardown (NOT group-based — see above). Idempotent: the + # ticker may already have self-exited when run_pid died. Kills the in-flight + # `sleep` child first, then the subshell, then reaps so no `wait` can hang. + _expt_stop_ticker() { + [ -n "$ticker_pid" ] || return 0 + pkill -P "$ticker_pid" 2>/dev/null || true # its in-flight `sleep` + kill "$ticker_pid" 2>/dev/null || true + wait "$ticker_pid" 2>/dev/null || true + } + + # Optional convergence-stop watcher (EXPERIMENTS.md WS2) -- ONLY when the driver + # set EXPT_TARGET_ACC. Side-car: on "EXPT_CONVERGE_WINDOW consecutive data bins + # all >= EXPT_TARGET_ACC" it writes converge.json + kills the run's process + # group. Absent -> the run is governed only by its max_runtime_s/max_data_id caps. + local watcher_pid="" cj="" sj="" + EXPT_LAST_CONVERGE_JSON=""; export EXPT_LAST_CONVERGE_JSON + if [ -n "${EXPT_TARGET_ACC:-}" ]; then + cj="$logdir/converge_${label}.json" + sj="$logdir/stall_${label}.json" + # EXPT_STALL_WINDOW_S=0 (default) disables the stall guard -> pure wall/window run. + python "$EXPT_RUNNER_DIR/converge_watch.py" \ + --exp-dir "$exp_dir" --marker "$EXPT_LAST_MARKER" \ + --target-acc "$EXPT_TARGET_ACC" --window "${EXPT_CONVERGE_WINDOW:-20}" \ + --pgid "$run_pid" --converge-json "$cj" --stall-json "$sj" \ + --stall-window-s "${EXPT_STALL_WINDOW_S:-0}" --stall-min-delta "${EXPT_STALL_MIN_DELTA:-0.01}" \ + --stall-on "${EXPT_STALL_ON:-either}" --loss-min-rel-delta "${EXPT_LOSS_MIN_REL_DELTA:-0.01}" \ + --poll "${EXPT_CONVERGE_POLL_S:-15}" 2>&1 | tee -a "$logdir/expt_runner.log" & + watcher_pid=$! + EXPT_LAST_CONVERGE_JSON="$cj"; export EXPT_LAST_CONVERGE_JSON + fi + + # Ctrl+C/SIGTERM teardown: the run is in its OWN process group (set -m) so a + # terminal SIGINT never reaches it -- without this trap it (and the watcher and + # ticker) would orphan and keep the GPU pinned. TERM run group + watcher + ticker, + # escalate to KILL after a grace, sweep stragglers, exit 130. + _expt_interrupt() { + trap - INT TERM + echo "" >&2 + echo "[$(date '+%F %T')] INTERRUPT — tearing down '$label' (run pgid=$run_pid) ..." >&2 + kill -TERM -"$run_pid" 2>/dev/null || true # run's process group ... + kill -TERM "$run_pid" 2>/dev/null || true # ... and its leader (group may not have formed) + _expt_stop_ticker # PID-based; never a group signal + # $watcher_pid is the tee of the `converge_watch.py | tee` pipeline, so also + # kill the poller by name. + [ -n "$watcher_pid" ] && kill -TERM "$watcher_pid" 2>/dev/null + pkill -TERM -f converge_watch.py 2>/dev/null || true + sleep "${EXPT_INT_GRACE_S:-5}" + kill -KILL -"$run_pid" 2>/dev/null || true + kill -KILL "$run_pid" 2>/dev/null || true + pkill -9 -f 'flame.launch.run_experiment' 2>/dev/null || true + pkill -9 -f 'trainer/forward_training' 2>/dev/null || true + pkill -9 -f 'trainer/pytorch/main.py' 2>/dev/null || true + pkill -9 -f 'aggregator/pytorch/main_' 2>/dev/null || true + pkill -9 -f converge_watch.py 2>/dev/null || true + echo "[$(date '+%F %T')] INTERRUPT — teardown complete for '$label'. GPU/RAM freed." >&2 + exit 130 + } + trap _expt_interrupt INT TERM + + wait "$run_pid" + local rc=$? + trap - INT TERM + + _expt_stop_ticker # PID-based (+ its sleep child); may already have self-exited + if [ -n "$watcher_pid" ]; then kill "$watcher_pid" 2>/dev/null; wait "$watcher_pid" 2>/dev/null; fi + + # Watcher verdict: converge.json = CONVERGED, stall.json = STALLED. On a + # watcher-driven kill (either), sweep any orphaned workers so GPU memory frees. + EXPT_LAST_CONVERGED=0; EXPT_LAST_STALLED=0; export EXPT_LAST_CONVERGED EXPT_LAST_STALLED + [ -n "$cj" ] && [ -f "$cj" ] && EXPT_LAST_CONVERGED=1 + [ -n "$sj" ] && [ -f "$sj" ] && EXPT_LAST_STALLED=1 + if [ "$EXPT_LAST_CONVERGED" = "1" ] || [ "$EXPT_LAST_STALLED" = "1" ]; then + pkill -9 -f "trainer/forward_training" 2>/dev/null || true + pkill -9 -f "aggregator/pytorch/main_" 2>/dev/null || true + fi + + local elapsed=$(( $(date +%s) - start_ts )) + echo "[$(date '+%F %T')] DONE $label exit=$rc converged=${EXPT_LAST_CONVERGED} stalled=${EXPT_LAST_STALLED} (took ${elapsed}s / ~${budget_s}s budget)" | tee -a "$logdir/expt_runner.log" + EXPT_LAST_RC=$rc + return $rc +} + +# expt_assert_run example_dir marker [label] -- post-run health check. FL signals +# do NOT land in the runner .out: agg_round events live in +# experiments/run_*/telemetry/aggregator_*.jsonl; "stopping run" / SIM_WALL_CEILING +# / SIM_STARVATION / tracebacks live in experiments/run_*/*_aggregator.log. Only +# files newer than `marker` are scanned, so signals attribute to THIS run. Prints a +# one-line verdict and exports EXPT_LAST_HEALTH (COMPLETED / CRASH / NO_AGG_ROUNDS / +# WALL_CEILING / NO_MARKER); returns 0 only for agg_rounds && no wall-ceiling && no +# crash. A healthy run is "COMPLETED", not "PASS" (PASS/FAIL is for actual checks). +expt_assert_run() { + local example_dir="$1" marker="$2" label="${3:-run}" + local expdir="$example_dir/experiments" + if [ ! -e "$marker" ]; then + EXPT_LAST_HEALTH="NO_MARKER"; export EXPT_LAST_HEALTH + echo " [$label] NO_MARKER (cannot scope health check)"; return 1 + fi + local agg=0 stop=0 wall=0 starv=0 crash=0 nlogs=0 f n + while IFS= read -r f; do + n=$(grep -c "agg_round" "$f" 2>/dev/null); agg=$(( agg + ${n:-0} )) + done < <(find "$expdir" -name "aggregator_*.jsonl" -newer "$marker" 2>/dev/null) + while IFS= read -r f; do + nlogs=$(( nlogs + 1 )) + n=$(grep -ic "stopping run" "$f" 2>/dev/null); stop=$(( stop + ${n:-0} )) + n=$(grep -c "SIM_WALL_CEILING" "$f" 2>/dev/null); wall=$(( wall + ${n:-0} )) + n=$(grep -c "\[SIM_STARVATION\]" "$f" 2>/dev/null); starv=$(( starv + ${n:-0} )) + n=$(grep -cE "Traceback \(most recent call last\)|CUDA error|Segmentation fault" "$f" 2>/dev/null); crash=$(( crash + ${n:-0} )) + done < <(find "$expdir" -name "*_aggregator.log" -newer "$marker" 2>/dev/null) + # A healthy, fully-run experiment is COMPLETED (a process outcome), not PASS. + local status="COMPLETED" + [ "$agg" -eq 0 ] && status="NO_AGG_ROUNDS" + [ "$wall" -gt 0 ] && status="WALL_CEILING" + [ "$crash" -gt 0 ] && status="CRASH" + # Convergence-mode verdict (EXPERIMENTS.md WS2): only with a target accuracy. A + # watcher-driven stop is CONVERGED; running out the safety caps is DID_NOT_CONVERGE. + # A crash/wall-ceiling still wins (a failure, not a verdict). + if [ -n "${EXPT_TARGET_ACC:-}" ] && [ "$crash" -eq 0 ] && [ "$wall" -eq 0 ] && [ "$agg" -gt 0 ]; then + if [ "${EXPT_LAST_CONVERGED:-0}" = "1" ]; then status="CONVERGED" + elif [ "${EXPT_LAST_STALLED:-0}" = "1" ]; then status="STALLED" + else status="DID_NOT_CONVERGE"; fi + fi + EXPT_LAST_HEALTH="$status"; export EXPT_LAST_HEALTH + printf " [%s] %-14s agg_round=%s stopping_run=%s wall_ceiling=%s starvation=%s crash=%s (%s agg log(s))\n" \ + "$label" "$status" "$agg" "$stop" "$wall" "$starv" "$crash" "$nlogs" + [ "$agg" -gt 0 ] && [ "$wall" -eq 0 ] && [ "$crash" -eq 0 ] +} + +# expt_timed_run label runtime_s buffer_s shell_log -- cmd... -- run a command in +# its OWN process group under a wall-clock timeout, killing the whole tree +# (SIGTERM grace -> SIGKILL) on overrun, then sweeping stragglers so no orphan +# trainers/aggregator survive and GPU memory drains. The campaign primitive shared +# by smoke_suite.sh. Emits a 5s elapsed/kill-in ticker. Returns cmd's rc, or 124. +expt_timed_run() { + local label="$1" runtime_s="$2" buffer_s="$3" shell_log="$4"; shift 4 + [ "${1:-}" = "--" ] && shift + local settle_s="${EXPT_KILL_SETTLE_S:-20}" grace_s="${EXPT_KILL_GRACE_S:-20}" + local wall_timeout=$(( runtime_s + buffer_s )) + local ts_start; ts_start=$(date +%s) + local deadline=$(( ts_start + wall_timeout )) timed_out=0 + # set -m: assign the background job its own PGID so kill - hits the tree + # (flame's spawner uses plain Popen with no setsid, so descendants inherit it). + set -m + "$@" > "$shell_log" 2>&1 & + local pid=$! + set +m + local ela kin + while kill -0 "$pid" 2>/dev/null; do + sleep 5 + # Re-test liveness: if the child finished DURING the sleep, exit normally + # rather than force-killing a dead pid and misreporting TIMEOUT (this bites + # whenever a run completes near the deadline). + kill -0 "$pid" 2>/dev/null || break + ela=$(( $(date +%s) - ts_start )); kin=$(( deadline - $(date +%s) )); [ "$kin" -lt 0 ] && kin=0 + printf '\r [%s] %ds/%ds kill in %ds ' "$label" "$ela" "$runtime_s" "$kin" >&2 + if [ "$(date +%s)" -ge "$deadline" ]; then + printf '\n' >&2 + echo " [$label] TIMEOUT after ${wall_timeout}s — killing process group $pid" >&2 + kill -TERM -"$pid" 2>/dev/null || true + sleep "$grace_s" + kill -KILL -"$pid" 2>/dev/null || true + wait "$pid" 2>/dev/null + timed_out=1; break + fi + done + printf '\n' >&2 + [ "$timed_out" = "0" ] && wait "$pid" 2>/dev/null + local rc=$? + if [ "$timed_out" = "1" ]; then + pkill -9 -f "trainer/pytorch/main.py" 2>/dev/null || true + pkill -9 -f "trainer/forward_training" 2>/dev/null || true + pkill -9 -f "aggregator/pytorch/main_" 2>/dev/null || true + sleep "$settle_s" + EXPT_LAST_RC=124; return 124 + fi + EXPT_LAST_RC=$rc; return "$rc" +} + +# expt_dispatch_after "hook1,hook2" -- run the requested post-launch hooks. Each +# hook `foo` is delegated to a shell function `after_foo` the DRIVER defines (e.g. +# after_parity, after_plot, after_sanity), so the mechanism is shared but each +# example supplies its own command. Unknown hooks warn and are skipped (not fatal). +expt_dispatch_after() { + local csv="${1:-}"; [ -n "$csv" ] || return 0 + local hook + IFS=',' read -ra _hooks <<< "$csv" + for hook in "${_hooks[@]}"; do + hook="${hook// /}"; [ -n "$hook" ] || continue + if declare -F "after_${hook}" >/dev/null 2>&1; then + echo "=== after: ${hook} ===" + "after_${hook}" || echo " [after:${hook}] returned non-zero (continuing)" + else + echo " [after:${hook}] not supported for this example — skipping" >&2 + fi + done +} diff --git a/lib/python/flame/backend/mqtt.py b/lib/python/flame/backend/mqtt.py index 8142d6d82..696a586e4 100644 --- a/lib/python/flame/backend/mqtt.py +++ b/lib/python/flame/backend/mqtt.py @@ -194,8 +194,31 @@ def _topics_for_notify(self, channel: Channel) -> list[str]: return notify_topics + def _wait_for_connect(self, timeout: float = DEFAULT_RUN_ASYNC_WAIT_TIME) -> bool: + """Block until the MQTT client reports connected (bounded by timeout). + + connect() is async: if join() wins the race against on_connect, + subscribe() targets an unconnected client and notify(JOIN) is dropped by + the connection guard, so this end never registers with peers and hangs. + Waiting defers subscription+JOIN until the connection is up; returns + immediately when on_connect already fired. + """ + deadline = time.time() + timeout + while not self._is_connected and time.time() < deadline: + time.sleep(0.01) + if not self._is_connected: + logger.warning( + f"join proceeding without a confirmed MQTT connection " + f"after {timeout}s; presence notify may be dropped" + ) + return self._is_connected + def join(self, channel: Channel) -> None: """Join a channel by subscribing to topics.""" + # subscribe/announce only after the connection is up, else both can be + # silently dropped (startup race) + self._wait_for_connect() + for topic in self._topics_for_notify(channel): self.subscribe(topic) diff --git a/lib/python/flame/config.py b/lib/python/flame/config.py index dc8f28fb1..2578ffd16 100644 --- a/lib/python/flame/config.py +++ b/lib/python/flame/config.py @@ -273,6 +273,22 @@ class Hyperparameters(FlameSchema, extra=Extra.allow): sim_unavailability: t.Optional[bool] = Field( alias="simUnavailability", default=False ) + # sct-model folds (fwdllm Stage B / K-D20 #6). Default OFF ⇒ byte-identical. + # simModelEvalTime (B1): charge the measured server-eval wall to the vclock + # after each committed data_id -- the largest GENUINE unmodeled term (~3.34 + # s/round). simStragglerSpreadS (B2): per-trainer straggler dispersion (s) + # added to the modeled delay so the sync barrier's k-th sct reflects real + # trainer_speed_s spread (~2.3 s/round). simWanTransferS (B3): WAN payload- + # transfer term -- NOT measurable on localhost, documented knob, leave 0. + sim_model_eval_time: t.Optional[bool] = Field( + alias="simModelEvalTime", default=False + ) + sim_straggler_spread_s: t.Optional[float] = Field( + alias="simStragglerSpreadS", default=0.0 + ) + sim_wan_transfer_s: t.Optional[float] = Field( + alias="simWanTransferS", default=0.0 + ) # Per-baseline: aware baselines free stalled slots proactively at the next # selection boundary (Stage D); unaware wait for the 90s vclock abandon. # Kept for backward compat — new code reads proactive_inflight_evict first. diff --git a/lib/python/flame/launch/aggregator_spawner.py b/lib/python/flame/launch/aggregator_spawner.py index e109062c8..04961f431 100644 --- a/lib/python/flame/launch/aggregator_spawner.py +++ b/lib/python/flame/launch/aggregator_spawner.py @@ -28,6 +28,7 @@ def spawn( log_to_wandb: bool = False, wandb_run_name: Optional[str] = None, cpu_cores: Optional[set] = None, + gpu_id: Optional[int] = None, ) -> subprocess.Popen: """Spawn aggregator process. @@ -88,6 +89,12 @@ def spawn( # math libs use exactly that many threads (it benefits from a few cores # for chunk reassembly / aggregation, unlike a 1-core-pinned trainer). env = os.environ.copy() + # GPU pin: give the aggregator its own device so its eval forward pass + # does not time-slice a trainer's GPU. Without it the aggregator defaults + # to GPU 0, inflating that trainer's compute over its delay budget. + if gpu_id is not None: + env["CUDA_VISIBLE_DEVICES"] = str(gpu_id) + print(f" ✓ Aggregator pinned to GPU {gpu_id}") preexec_fn = None if cpu_cores: _cores = {int(c) for c in cpu_cores} @@ -119,12 +126,43 @@ def is_running(self) -> bool: return False return self.process.poll() is None + # Detect a startup crash from its traceback so wait_until_ready fails fast, + # instead of the old "alive => ready" heuristic waving a crash through and + # then wasting the EOT grace on trainers that never get one. + _CRASH_MARKER = "Traceback (most recent call last)" + + def _read_log_tail(self, max_bytes: int = 65536) -> str: + if not self.log_file: + return "" + p = Path(self.log_file) + if not p.exists(): + return "" + try: + with open(p, "r", errors="replace") as f: + data = f.read() + return data[-max_bytes:] + except OSError: + return "" + + def _print_crash_tail(self, n_lines: int = 25) -> None: + """Surface the aggregator log tail so the operator sees the crash cause + without opening the log.""" + tail = self._read_log_tail() + if not tail: + return + lines = tail.rstrip().splitlines()[-n_lines:] + print(" ── aggregator log tail ──") + for ln in lines: + print(f" {ln}") + print(" ─────────────────────────") + def wait_until_ready(self, timeout: int = 30) -> bool: """ - Wait for aggregator to be ready. + Wait for aggregator to be ready, failing fast on a startup crash. - Simple implementation: just wait fixed time and check process is alive. - Could be enhanced with log monitoring for "ready" message. + Ready: process alive for >=5s with no traceback in the log. Returns False + immediately if the process exits during the window or a traceback appears + -- the run has failed, so the caller must not spawn trainers. Args: timeout: Maximum time to wait in seconds @@ -139,14 +177,23 @@ def wait_until_ready(self, timeout: int = 30) -> bool: while time.time() - start_time < timeout: if not self.is_running(): - print(f" ✗ Aggregator process died") + print(f" ✗ Aggregator process died during startup") + self._print_crash_tail() + return False + if self._CRASH_MARKER in self._read_log_tail(): + print(f" ✗ Aggregator logged a traceback during startup") + self._print_crash_tail() return False time.sleep(check_interval) elapsed = time.time() - start_time - # Simple heuristic: if process is alive for 5 seconds, assume ready + # Alive for 5s AND no startup traceback => assume ready. if elapsed >= 5: + if self._CRASH_MARKER in self._read_log_tail(): + print(f" ✗ Aggregator logged a traceback during startup") + self._print_crash_tail() + return False print(f" ✓ Aggregator ready (process alive for {elapsed:.1f}s)") return True diff --git a/lib/python/flame/launch/execution_config_generator.py b/lib/python/flame/launch/execution_config_generator.py index 61afc1c2e..eaa54c64b 100644 --- a/lib/python/flame/launch/execution_config_generator.py +++ b/lib/python/flame/launch/execution_config_generator.py @@ -130,6 +130,15 @@ def create_execution_config( }, "battery_threshold": exp_config.trainer.battery_threshold, "time_mode": exp_config.trainer.time_mode, + # Bank effective training-delay config for post-hoc audit (#12). + # training_delay_factor is the trainer override if set, else the + # trainer_base default applies. + "enable_training_delays": exp_config.trainer.enable_training_delays, + "training_delay_factor": ( + (exp_config.trainer.hyperparameters or {}).get( + "training_delay_factor" + ) + ), }, "execution": { "num_gpus": exp_config.execution.num_gpus, diff --git a/lib/python/flame/launch/runner.py b/lib/python/flame/launch/runner.py index 86eba2de3..ab3758098 100644 --- a/lib/python/flame/launch/runner.py +++ b/lib/python/flame/launch/runner.py @@ -202,12 +202,28 @@ def run_experiment(self, exp: ExperimentConfig) -> None: }, ) + # Dedicated aggregator GPU: prefer a physical GPU the trainer pool + # does NOT use (visible > num_gpus → the first idle one); else the + # least-loaded trainer GPU (highest index under (tid-1)%num_gpus). + _num_gpus = exp.execution.num_gpus + try: + import torch as _torch + _visible = _torch.cuda.device_count() + except Exception: + _visible = 0 + if _visible > _num_gpus: + _agg_gpu = _num_gpus # a fully idle physical GPU + elif _num_gpus > 0: + _agg_gpu = _num_gpus - 1 # least-loaded trainer GPU + else: + _agg_gpu = None self.aggregator_spawner.spawn( paths["aggregator_main"], config_json=json.dumps(agg_cfg), log_to_wandb=exp.aggregator.log_to_wandb, wandb_run_name=exp.aggregator.wandb_run_name, cpu_cores=reserved_cores, + gpu_id=_agg_gpu, ) if not self.aggregator_spawner.wait_until_ready( exp.execution.aggregator_warmup_time @@ -323,8 +339,15 @@ def run_experiment(self, exp: ExperimentConfig) -> None: self.aggregator_spawner.terminate() agg_rc = getattr(self.aggregator_spawner.process, "returncode", None) rc_msg = f"exit={agg_rc}" if agg_rc == 0 else f"exit={agg_rc} ⚠" - print(f" aggregator done ({rc_msg}), waiting for trainers to exit...") - self.trainer_spawner.wait_all(timeout_per_trainer=30.0) + if agg_rc not in (0, None): + # On a crash the trainers never get an EOT, so the per-trainer + # grace below is wasted -- terminate them now instead. + print(f" aggregator FAILED ({rc_msg}); terminating trainers now " + f"(skipping EOT grace).") + self.trainer_spawner.terminate_all() + else: + print(f" aggregator done ({rc_msg}), waiting for trainers to exit...") + self.trainer_spawner.wait_all(timeout_per_trainer=30.0) print("\nexperiment completed.") # Auto post-run analysis: parse the telemetry JSONL and emit plots. @@ -535,7 +558,39 @@ def _build_aggregator_config( }, )) + # Single source of truth: exp.trainer's delay flags govern the whole run. + # Fan them into the aggregator hyperparameters as the final (highest- + # precedence) layer so both roles agree; else the aggregator keeps its + # pydantic default (False), a real<->sim desync risk for examples whose + # aggregator reads it (async_cifar10). See #12 / #13. + _delay_fan: dict = { + "trainingDelayEnabled": bool(exp.trainer.enable_training_delays), + } + # If the experiment set training_delay_factor on the trainer, fan the same + # value to the aggregator so a launcher knob reaches both roles. + _tr_hp = exp.trainer.hyperparameters or {} + if "training_delay_factor" in _tr_hp: + _delay_fan["trainingDelayFactor"] = _tr_hp["training_delay_factor"] + layers.append(( + "experiment.trainer.training_delay (fanned to aggregator)", + {"hyperparameters": _delay_fan}, + )) + merged, provenance = merge_with_provenance(layers) + + # Tripwire: the delay fan is the final layer, so the merged config must + # reflect the requested flag. If a refactor reorders layers or a higher- + # precedence override shadows it, fail loudly instead of running stale (#12). + _eff = merged.get("hyperparameters", {}).get("trainingDelayEnabled") + _req = bool(exp.trainer.enable_training_delays) + if _eff is not None and bool(_eff) != _req: + raise ValueError( + "training-delay config did not flow to the aggregator: requested " + f"enable_training_delays={_req} but merged aggregator " + f"trainingDelayEnabled={_eff!r}. Check the config-override layer " + "order in _build_aggregator_config (simulate_fwdllm.md #12)." + ) + return merged, provenance def _create_experiment_directory(self, exp: ExperimentConfig) -> Path: diff --git a/lib/python/flame/launch/spawner.py b/lib/python/flame/launch/spawner.py index de1c06f35..e2f86846a 100644 --- a/lib/python/flame/launch/spawner.py +++ b/lib/python/flame/launch/spawner.py @@ -428,6 +428,56 @@ def spawn_all( print(f" {'Trainer':>8} {'GPU':>4} {'CPU core':>9} {'PID':>7}") for p in self.processes: print(f" {p['trainer_id']:>8} {p['gpu_id']:>4} {str(p.get('cpu_core', 'N/A')):>9} {p['process'].pid:>7}") + self._assert_load_balanced() + + def _assert_load_balanced(self) -> None: + """Post-spawn load-balance check relative to available hardware. + + Verifies round-robin placement spread work evenly and left no hardware + idle (under-provisioned num_gpus is a silent contention source). Emits one + grep-able [LOAD_BALANCE] verdict; WARNs, never raises. + """ + procs = self.processes + if not procs: + return + # trainers per assigned GPU + per CPU core + from collections import Counter + gpu_counts = Counter(p["gpu_id"] for p in procs) + core_counts = Counter(p.get("cpu_core") for p in procs if p.get("cpu_core") is not None) + # visible physical GPUs on the node (independent of num_gpus we chose) + try: + import torch + visible_gpus = torch.cuda.device_count() + except Exception: + visible_gpus = 0 + + issues = [] + # 1. Even GPU spread: round-robin guarantees max-min <= 1; flag otherwise. + if gpu_counts: + spread = max(gpu_counts.values()) - min(gpu_counts.values()) + if spread > 1: + issues.append(f"GPU imbalance: per-GPU trainer counts {dict(sorted(gpu_counts.items()))} (spread={spread}>1)") + # 2. Under-provisioning: physical GPUs left completely idle. + if visible_gpus and self.num_gpus < visible_gpus: + issues.append(f"under-provisioned: num_gpus={self.num_gpus} < visible={visible_gpus} " + f"→ {visible_gpus - self.num_gpus} GPU(s) idle; raise execution.num_gpus") + if visible_gpus and self.num_gpus > visible_gpus: + issues.append(f"over-subscribed: num_gpus={self.num_gpus} > visible={visible_gpus}") + # 3. Even CPU-core spread among pinned trainers. + if core_counts: + cspread = max(core_counts.values()) - min(core_counts.values()) + if cspread > 1: + issues.append(f"CPU imbalance: core reuse spread={cspread}>1") + + per_gpu = ", ".join(f"gpu{g}={n}" for g, n in sorted(gpu_counts.items())) + if issues: + print(f" ⚠ [LOAD_BALANCE] WARN ({len(procs)} trainers, num_gpus={self.num_gpus}, " + f"visible={visible_gpus}): {per_gpu}") + for it in issues: + print(f" - {it}") + else: + print(f" ✓ [LOAD_BALANCE] balanced: {len(procs)} trainers over {self.num_gpus} GPU(s) " + f"(visible={visible_gpus}): {per_gpu}; CPU cores 1/trainer") def wait_all(self, timeout_per_trainer: float = 30.0): """Wait for all trainer processes to complete. diff --git a/lib/python/flame/mode/horizontal/syncfl/fwdllm_aggregator.py b/lib/python/flame/mode/horizontal/syncfl/fwdllm_aggregator.py index 8423004d2..cb53e701d 100644 --- a/lib/python/flame/mode/horizontal/syncfl/fwdllm_aggregator.py +++ b/lib/python/flame/mode/horizontal/syncfl/fwdllm_aggregator.py @@ -32,7 +32,7 @@ from flame.common.constants import DeviceType from flame.common.util import weights_to_device, weights_to_model_device from flame.config import OptimizerType, TrainerAvailState -from flame.end import PROP_END_AVL_STATE +from flame.end import KEY_END_STATE, PROP_END_AVL_STATE, VAL_END_STATE_NONE from flame.mode.composer import CloneComposer import pickle from flame.mode.horizontal.syncfl.top_aggregator import ( @@ -47,6 +47,8 @@ from flame.mode.horizontal.asyncfl.top_aggregator import ( RECV_TIMEOUT_WAIT_S, TopAggregator as AsyncTopAgg, + _SIM_GATE_MAX_PASSES, + _SIM_ORDER_SLACK_S, ) from flame.mode.message import MessageType from flame.mode.horizontal.client_duration import real_client_task_train_duration @@ -61,6 +63,7 @@ PROP_STAT_UTILITY, PROP_UPDATE_COUNT, ) +from flame.selector.properties import PROP_SIM_SEND_TS import functorch as fc import torch @@ -70,7 +73,7 @@ import math from flame import telemetry -from flame.telemetry.events import build_agg_eval, build_agg_round, build_utility_belief +from flame.telemetry.events import build_agg_eval, build_agg_round, build_utility_belief, build_comm logger = logging.getLogger(__name__) @@ -236,6 +239,20 @@ class TopAggregator(AsyncTopAgg): """Top level Aggregator implements an ML aggregation role.""" + # The sim runs real forward-grad GPU + server eval, so its physical wall + # legitimately exceeds the vclock budget (#6/#13); a 1x wall ceiling would + # truncate it before vclock reached the budget. Decoupled to a generous + # multiple so it is a runaway OUTER safety, not the primary stop (primary + # stops: max_data_id_progress and vclock >= max_runtime_s). Override with an + # explicit `sim_wall_ceiling_s` for a tighter bound. + SIM_WALL_CEILING_FACTOR = 20.0 + + # Per-pass recv_fifo window the drain waits for an in-flight grad to arrive + # (fwdllm override of SyncTopAgg's 2.0). fwdllm's forward-grad pass is a real + # GPU pass (~4s), so 2s can close before the grad reassembles. + # TUNABLE: keep as low as possible without missing genuinely-in-flight grads. + SIM_RECV_GRACE_FLOOR_S = 5.0 + def internal_init(self) -> None: """Initialize internal state for role.""" super().internal_init() @@ -249,6 +266,11 @@ def internal_init(self) -> None: self._updates_in_queue = 0 self._updates_received = {} self._per_agg_trainer_list = [] + # end -> canonical commit-order key (modeled_delay D, str(end)) for the + # current cycle's cohort. Populated per contribution in aggregate_weights; + # consumed by _canonicalize_cohort_commit_order to break equal-D ties by + # trainer_id identically in real and sim. + self._commit_key_by_end = {} self._model_version_unique_trainers = set() self._model_version_trainer_stats = { "train_duration": [], @@ -280,12 +302,94 @@ def internal_init(self) -> None: logger.info( f"[MaxIterBypass] max_iterations_per_data_id={self._max_iter_per_data_id}" ) + + # Opt-2 (charter §5c): variance-plateau stopping policy. At α=1 the + # variance floor sits above the gate, so a data-bin crosses only on a + # noise dip and grinds many iterations while the denoised estimate has + # plateaued. This commits early once the variance-decay curve flattens. + # fluxtune-only; absent/'off' => legacy max-iter cap governs. + # * fixed_cap: rely on max_iterations_per_data_id as the ceiling. + # * plateau: ALSO commit when the relative var drop over the last N + # cycles < rel_delta while var is still above threshold. + # Commit reason (natural/cap/plateau) is recorded for telemetry. + self._var_stopping_policy = getattr( + self.config.hyperparameters, "var_stopping_policy", None + ) + self._var_plateau_patience = int( + getattr(self.config.hyperparameters, "var_plateau_patience", 3) + ) + self._var_plateau_rel_delta = float( + getattr(self.config.hyperparameters, "var_plateau_rel_delta", 0.10) + ) + self._force_commit_reason = None # {natural, cap, plateau}; for agg_round + self._grad_aware_gated_total = 0 # Opt-3: anti-aligned updates down-weighted + if self._var_stopping_policy not in (None, "off", "fixed_cap", "plateau"): + raise ValueError( + f"unknown var_stopping_policy={self._var_stopping_policy!r}; " + "expected one of off/fixed_cap/plateau" + ) + if self._var_stopping_policy in ("fixed_cap", "plateau"): + if ( + self._var_stopping_policy == "fixed_cap" + and self._max_iter_per_data_id is None + ): + logger.warning( + "[VarStopPolicy] policy=fixed_cap but " + "max_iterations_per_data_id is unset -> no cap will fire." + ) + logger.info( + f"[VarStopPolicy] policy={self._var_stopping_policy} " + f"cap(max_iter)={self._max_iter_per_data_id} " + f"plateau_N={self._var_plateau_patience} " + f"plateau_rel_delta={self._var_plateau_rel_delta}" + ) + + # Opt-1 (charter §5c): suppress byte-identical intra-databin weight + # re-sends. Within a databin the WEIGHTS+GRAD_POOL payload is identical + # across iterations, yet a trainer pulled in to refill concurrency reads + # as "stale" (_trainer_last_model_version is written only on grad-RETURN) + # and gets the identical payload again. When on, downgrade such a + # re-dispatch to the tiny VAR=bad "keep training" message (the trainer + # caches its weights). Async path only; flag OFF => byte-identical. + self._suppress_redundant_weights = bool( + getattr(self.config.hyperparameters, "suppress_redundant_weights", False) + ) + # Ends already sent the CURRENT model_version's full payload this data-bin + # (cleared on every model_version advance). Separate from + # _trainer_last_model_version so staleness accounting stays return-driven. + self._weights_sent_this_cycle: set = set() + self._redundant_weights_suppressed_total = 0 + if self._suppress_redundant_weights: + logger.info( + "[SuppressRedundantWeights] ON — intra-databin identical weight " + "re-sends will be downgraded to VAR=bad." + ) self.grad_pool = [] self.cached_shared_grad_pool_trainable = None self.var = None self.ends_not_selected_yet = False self.iteration_per_data_id = 0 + # end_id -> {dispatch_ts, commit_ts} of its last accepted contribution, + # emitted per-cycle as contributor_intervals for the R1/W1 rungs (§L.3). + self._sim_contrib_intervals = {} + + # #15 compute-truthful commit gate (flag-gated; default off). The + # earlier_stuck gate in _sim_recv_min_grad blocks on _sim_inflight_expected + # entries stamped at DISPATCH. A trainer whose payload was sent long ago + # but has not returned is idle-in-recv, not computing; blocking on it burns + # the grace/failsafe and starves re-dispatch. When on, the gate only blocks + # on a trainer still within its modeled compute window, so a stamped-but- + # idle phantom no longer defers a ready commit. Hold-to-commit untouched. + self._sim_compute_truthful_gate = bool(getattr( + self.config.hyperparameters, "sim_compute_truthful_gate", False)) + # Wall seconds a dispatched grad may still plausibly be computing before it + # is treated as idle/phantom. Only consulted when the gate flag is on. + _cap = getattr(self.config.hyperparameters, "sim_gate_compute_cap_s", 10.0) + self._sim_gate_compute_cap_s = float(_cap) if _cap is not None else 10.0 + # end -> wall time its weights/VAR=bad payload was last sent (sim only). + self._sim_dispatch_wall = {} + # Selection granularity for the sync path (fwdllm/fwdllm_plus): # True (default, preserves pre-existing behavior) = re-select # trainers on every SEND-state call, i.e. every iteration of every @@ -305,11 +409,25 @@ def internal_init(self) -> None: # but not formally departed (see ROUND_CACHE_STUCK_TIMEOUT_S). self._round_cache_activity_ts: dict = {} + # Wire staleness_policy from config into an instance attr; without this + # the message handler's getattr fell back to "none" for every run. + self.staleness_policy = getattr( + self.config.hyperparameters, "staleness_policy", None + ) or "none" + logger.info(f"staleness_policy = {self.staleness_policy}") + self._optimizer_sort_value = self.config.optimizer.sort OPTIMIZERS_SUPPORTING_GRAD_AGGREGATION = (OptimizerType.FEDBUFF,) self._weighted_aggregation_enabled = ( self._optimizer_sort_value in OPTIMIZERS_SUPPORTING_GRAD_AGGREGATION ) + if self.staleness_policy == "fedbuff" and not self._weighted_aggregation_enabled: + # fedbuff accepts stale grads expecting the optimizer to down-weight + # them by (V'-V); without a weighting optimizer the rate stays 1.0. + logger.warning( + f"staleness_policy=fedbuff but optimizer.sort={self._optimizer_sort_value} " + "does not down-weight staleness (rate stays 1.0)." + ) if not self._weighted_aggregation_enabled: logger.info( f"Setting rate=1.0 for all updates because optimizer.sort is " @@ -522,7 +640,7 @@ def _process_trainer_heartbeat(self, msg, end) -> None: logger.warning(f"Got invalid {msg} while processing heartbeat") def read_trainer_unavailability( - self, trace=None, metadata_dir: Optional[Union[str, Path]] = None + self, trace=None, base_dir: Optional[Union[str, Path]] = None ) -> dict: """Build task_id -> SortedDict(timestamp -> state) for `trace`. @@ -530,10 +648,17 @@ def read_trainer_unavailability( traces), mirroring async_cifar10/aggregator/pytorch/main_oort_sync_agg.py's pattern -- not from the legacy per-trainer json_scripts/trainer_*.json files. + + NOTE: the parameter MUST be named `base_dir` to match the caller in + ClientAvailability._init_availability, which passes base_dir=...; this + method shadows the mixin wrapper, so a mismatched name raises TypeError at + aggregator init. Canonical impl now in + flame.availability.trace.read_trainer_unavailability; this override is a + deletion candidate once mobiperf/syn parity with load_trace is confirmed. """ logger.info(f"Reading trainer unavailability for trace: {trace}") - metadata_dir = Path(metadata_dir) if metadata_dir is not None else _METADATA_DIR + metadata_dir = Path(base_dir) if base_dir is not None else _METADATA_DIR registry_path = metadata_dir / "trainer_registry.yaml" with open(registry_path) as f: @@ -642,6 +767,52 @@ def aggregate_grads_from_trainers( ) rate = 1.0 + elif self.optimizer.agg_rate_conf["type"] == "grad_aware": + # Opt-3 (charter §5c): gradient-aware aggregation. The "new" rate + # rescales magnitude by staleness×utility but can't refuse a wrong + # direction, so anti-aligned JVP estimates still get averaged. + # grad_aware weights by direction (align gate) and optionally + # reliability (inverse-variance), bounded <= base so the effective + # LR never inflates. OFF unless type set here. + conf = self.optimizer.agg_rate_conf + if conf.get("base", "new") == "new": + try: + base = self.optimizer.weight_factor( + scale=conf.get("scale", 0.4), staleness=staleness_val, + a_exp=conf.get("a_exp", 0.25), loss=stat_utility, + b_exp=conf.get("b_exp", 0.1), alpha_type="polynomial", + beta_type="polynomial_upshift", + ) + except Exception: + base = 1.0 + else: + base = 1.0 + cos = self._cosine_flat( + trainer_grad, self.grad, list(self.model.named_parameters()) + ) + var_i = None + if conf.get("inverse_var", False) and grad_for_var_check is not None: + try: + var_i = float(torch.stack(list(grad_for_var_check)).var()) + except Exception: + var_i = None + _align_floor = conf.get("align_floor", 0.0) + rate = self._grad_aware_rate( + base, cos, var_i, + var_ref=getattr(self, "var_threshold", conf.get("var_ref", 0.3)), + align_gate=conf.get("align_gate", True), + inverse_var=conf.get("inverse_var", False), + align_floor=_align_floor, var_eps=conf.get("var_eps", 1e-8), + ) + if cos is not None and cos < _align_floor: + self._grad_aware_gated_total = ( + getattr(self, "_grad_aware_gated_total", 0) + 1 + ) + logger.info( + f"[GradAware] anti-aligned update gated: cos={cos:.3f} " + f"< floor={_align_floor} base={base:.3f} -> rate={rate:.4f}" + ) + if rate != 1.0: logger.info( f"Weighted received gradients by rate: {rate} with staleness: {staleness_val}, stat utility: {stat_utility}" @@ -690,6 +861,291 @@ def aggregate_grad_pool(self, grad_list): self.print_trainable_params_stats(location="[end,aggregate_grad_pool()]") return grad + def _sim_recv_min_grad(self, channel, recv_ends): + """Grad-loop analog of asyncfl._sim_recv_min (§J.2). + + Ingest every ready grad into the sct-ordered reorder buffer (keyed by + SIM_COMPLETION_TS), hold the commit while an in-flight trainer is expected + to complete before the buffered minimum, then pop and commit the single + smallest-sct message, advancing the virtual clock to it. Returns one + (msg, metadata) matching next(channel.recv_fifo(...,1)), or (None, ("", + now)) when nothing is committable. + + Not _sim_recv_min verbatim because fwdllm commits gradients one-per-call + and releases slots on the agg-goal boundary (not per message), and must + survive variance-FAIL rollbacks without stranding or double-committing a + grad (§I.5). Slot/selector state is cleared at the boundary, not here. + """ + live = [e for e in (recv_ends or []) if channel.has(e)] + deadline = time.time() + RECV_TIMEOUT_WAIT_S + for _pass in range(_SIM_GATE_MAX_PASSES): + grace = self._sim_recv_grace_s() + # Base probe: the live recv_ends (always drained), minus anything + # already buffered or committed this cycle. + _base = [ + e for e in live + if not self._sim_buffer.has(e) and e not in self._sim_committed + ] + _seen = set(_base) + if getattr(self, "_sim_sct_ordered_drain", False): + # #13: direct sct-ordered ingest. Drain each live in-flight end's + # rx queue directly (no recv_fifo streamer) over the full in-flight + # set. Two wins: (1) the buffer is a complete snapshot of every + # arrived grad -- the streamer's background task/shared queue can + # strand a delivered grad and let the clock lap it (past-dated + # commit); (2) drain_ready sweeps ready rxqs non-blocking, so it + # doesn't burn the full grace per not-ready end. + to_probe = _base + [ + e for e in self._sim_inflight_expected + if e not in _seen and channel.has(e) + and not self._sim_buffer.has(e) and e not in self._sim_committed + ] + if to_probe: + for m, md in channel.drain_ready(to_probe, timeout=grace): + _e = md[0] + _s = m.get(MessageType.SIM_COMPLETION_TS) + _s = float(_s) if _s is not None else self._vclock.now + self._sim_buffer.add(_e, _s, (m, md)) + else: + # #13: probe-ceiling + ready-gating. Probe an in-flight-expected + # end that is not already a recv_end only if it is physically ready + # or its modeled completion `exp` is at/before the buffered minimum + # (+slack); else the drain blocked the full grace window every pass + # on far-future stragglers. Safe against the HOLD gate: any end that + # could trigger `earlier_stuck` also satisfies exp <= bmin + slack, + # so the gate's stuck end is always in this probe set. + _bmin = self._sim_buffer.peek_min_ts() + _probe_ceiling = ( + _bmin + _SIM_ORDER_SLACK_S if _bmin is not None else float("inf") + ) + to_probe = _base + [ + e for e, exp in self._sim_inflight_expected.items() + if e not in _seen and channel.has(e) + and not self._sim_buffer.has(e) and e not in self._sim_committed + and (self._sim_end_has_ready_msg(channel, e) or exp <= _probe_ceiling) + ] + if to_probe: + for m, md in channel.recv_fifo( + to_probe, first_k=len(to_probe), timeout=grace + ): + if m is None: # no more ready (grace expired or set drained) + break + _e = md[0] + _s = m.get(MessageType.SIM_COMPLETION_TS) + _s = float(_s) if _s is not None else self._vclock.now + self._sim_buffer.add(_e, _s, (m, md)) + # Gate: earliest expected completion among un-buffered in-flight ends. + bmin = self._sim_buffer.peek_min_ts() + _stuck_end, min_stuck = None, None + # #15 compute-truthful gate: only a trainer still within its modeled + # compute window can be a live earlier-sct straggler. One whose last + # dispatch is older than the compute cap (or never dispatched) is + # idle-in-recv/phantom and won't produce a grad until the drain yields, + # so blocking on it deadlocks to the failsafe. Skip it so the buffered + # commit proceeds and the loop can re-dispatch. + _truthful = getattr(self, "_sim_compute_truthful_gate", False) + _now_wall = time.time() + _cap = getattr(self, "_sim_gate_compute_cap_s", 10.0) + for e, exp in self._sim_inflight_expected.items(): + if self._sim_buffer.has(e) or e in self._sim_committed: + continue + if _truthful: + _dw = self._sim_dispatch_wall.get(e) + if _dw is None or (_now_wall - _dw) > _cap: + self._sim_gate_phantom_skip = getattr( + self, "_sim_gate_phantom_skip", 0) + 1 + continue # not genuinely computing -> can't block a commit + if min_stuck is None or exp < min_stuck: + min_stuck, _stuck_end = exp, e + earlier_stuck = ( + bmin is not None and min_stuck is not None + and min_stuck + _SIM_ORDER_SLACK_S < bmin + ) + if bmin is None and not to_probe: + break # nothing to commit and nothing in flight + if not earlier_stuck: + break # the buffered minimum is the true next completion + if time.time() >= deadline: + # #13 failsafe: the earliest-expected in-flight trainer never + # arrived within the window. Without eviction it stays in + # _sim_inflight_expected forever, so `earlier_stuck` re-fires the + # full deadline every drain cycle -> pipeline starvation. Treat it + # as lost: drop it from the expected set and commit the buffered + # min now. Sim-only. + self._sim_gate_failsafe = getattr(self, "_sim_gate_failsafe", 0) + 1 + self._sim_inflight_expected.pop(_stuck_end, None) + logger.info( + f"[SIM_GRAD_STUCK_EVICT] round={getattr(self, '_round', -1)} " + f"end={str(_stuck_end)[-4:]} exp={min_stuck} bmin={bmin} " + f"failsafe={self._sim_gate_failsafe}" + ) + break # failsafe: a stuck trainer never arrived; commit buffered + popped = self._sim_buffer.pop_min() + if popped is None: + return None, ("", datetime.now()) + _end, sct, (m, md) = popped + # Clock-jump clamp: a far-future straggler commit must not lap a fresh + # in-flight cohort still expected to complete earlier. + _now = self._vclock.now + _min_future = None + for e, exp in self._sim_inflight_expected.items(): + if e == _end or e in self._sim_committed: + continue + if exp > _now and (_min_future is None or exp < _min_future): + _min_future = exp + _advance_to = ( + sct if _min_future is None + else max(_now, min(sct, _min_future + _SIM_ORDER_SLACK_S)) + ) + self._advance_sim_clock(_advance_to) + self._sim_committed.add(_end) + self._sim_inflight_expected.pop(_end, None) + # #13 freed-slot refill stamp: this commit frees a compute slot; record + # the just-advanced vclock so the trainer refilling the slot rides THIS + # vclock (not the round-start frontier). Spreads each cohort's expected + # completions across the timeline (matching real's staggered returns) + # instead of collapsing them at one frozen `_round_now`, so the drain gate + # stops holding for a batch of same-expected stragglers. Sim-only + gated. + if getattr(self, "_sim_staggered_redispatch", False): + self._sim_free_slot_ts.append(self._vclock.now) + # Grad committed -> trainer no longer in flight in virtual time -> drop it + # from pending so it is re-pickable (_sim_hold_busy_slots reconciles too). + self._sim_pending_commit.discard(_end) + # Learn this end's MODELED budget (contention-free lower bound) so the + # gate fires on genuine stragglers for not-yet-observed trainers. + _b = m.get(MessageType.TRAINING_BUDGET_S) if isinstance(m, dict) else None + if _b is not None: + self._sim_trainer_budget[_end] = float(_b) + self._sim_budget_min = min(self._sim_budget_min, float(_b)) + # Reassert selected_ends == the virtual-time in-flight set after this + # commit: recv_fifo just marked freshly-buffered ends RECVD (stripping + # their slots), but they are still in flight until THEY commit; else the + # in_flight telemetry undercounts. + if getattr(self, "_sim_inflight_residence", False): + self._sim_hold_busy_slots(channel) + # in_flight (virtual, dispatched-not-committed) vs physically-computing + # selected_ends; kept for concurrency parity debugging. + _sel = getattr(getattr(channel, "_selector", None), "selected_ends", None) + if isinstance(_sel, dict): + _sel_n = sum(len(v) for v in _sel.values()) + elif isinstance(_sel, (set, list)): + _sel_n = len(_sel) + else: + _sel_n = -1 + logger.info( + f"[SIM_GRAD_RECV] round={getattr(self, '_round', -1)} " + f"end={str(_end)[-4:]} sct={sct:.1f} T_v={self._vclock.now:.1f} " + f"buf_depth={len(self._sim_buffer)} " + f"inflight_exp={len(self._sim_inflight_expected)} sel_ends={_sel_n} " + f"phantom_skip={getattr(self, '_sim_gate_phantom_skip', 0)}" + ) + return m, md + + def _release_sim_slots_at_agg_goal(self, channel, is_async): + """Sim slot release at the agg-goal boundary. Two policies (§L): + + - async + sim_inflight_residence (fluxtune, c >> agg_goal): commit-then- + carry. Hold still-busy trainers before clearing and carry the surplus + buffer to the next fedbuff step. `_sim_hold_busy_slots` holds every + dispatched-but-not-committed trainer (computing ∪ carried) in both its + compute slot and re-pick guard until its grad commits. + - else (sync barriers c ≈ agg_goal, or residence off): legacy drop -- no + surplus, so clearing is correct and flag-off is byte-identical. + + `_sim_committed` always clears so a variance-FAIL re-contributor on the + rolled-back data_id isn't skipped. + """ + if not self.simulated: + return + if is_async and getattr(self, "_sim_inflight_residence", False): + self._sim_hold_busy_slots(channel) # reads buffer/in-flight -> hold before clear + self._sim_committed.clear() + return # keep _sim_buffer / _sim_inflight_expected -> carry surplus + self._sim_committed.clear() + self._sim_buffer.clear() + self._sim_inflight_expected.clear() + # legacy-drop path abandons all in-flight -> nothing is pending. (async + # path clears via _sim_hold_busy_slots' reconcile; sync barriers use the + # random selector where the ref is inert, but keep the set bounded.) + self._sim_pending_commit.clear() + if is_async: + self._sim_hold_busy_slots(channel) + + def _sim_hold_busy_slots(self, channel) -> None: + """Assert `selected_ends` == the virtual-time in-flight set: every + dispatched-but-not-committed trainer (`_sim_inflight_expected` ∪ buffered + surplus), holding both its compute slot (`selected_ends`, drives + `extra = c − len(selected_ends)`) and its re-pick guard (`all_selected`) + until its grad commits. + + WHY: a returned-but-uncommitted trainer is still in flight in virtual time + (its grad commits only when the vclock reaches its sct), so its slot is + genuinely occupied. Freeing on physical RETURN collapsed `selected_ends` + to the physically-computing few while many were truly in flight, so the + in_flight telemetry undercounted and `extra` read false free capacity. + Holding to COMMIT also keeps one-in-flight-per-trainer across the carry + boundary + variance-FAIL rollbacks. Idempotent -- safe per-commit and at + the boundary. fwdllm-class override only (adds the triplet prune). + """ + sel = getattr(channel, "_selector", None) + if sel is None: + return + requester = getattr(sel, "requester", None) + all_selected = getattr(sel, "all_selected", None) + selected_ends = getattr(sel, "selected_ends", None) + + buffered = set(self._sim_buffer.pending_ends()) # returned, grad carried + # Outstanding = still in flight in virtual time: membership in + # `_sim_inflight_expected` or `_sim_buffer` (both popped on commit) IS the + # "not yet committed" truth. Do NOT subtract `_sim_committed`: it is a + # stale cross-cycle marker cleared only at the agg-goal boundary, so a + # trainer that committed then got re-picked + re-dispatched (re-added to + # _sim_inflight_expected) would be wrongly dropped -> re-pickable while its + # new dispatch is in flight -> R1 violation. One that committed THIS cycle + # is already absent from both sets, so the subtraction was redundant. + outstanding = set(self._sim_inflight_expected) | buffered + # `_sim_pending_commit` is the authoritative virtual in-flight set; + # reconcile it to `outstanding` in place (clear+update, never rebind -- the + # selector holds a live reference) so a committed trainer drops out + # (re-pickable) while a dispatched-but-uncommitted one stays. Bind the ref + # so async_oort's eligibility filter (`_pending`) excludes it regardless of + # all_selected churn. `|=` would never shrink -> committed trainer starved. + self._sim_pending_commit.clear() + self._sim_pending_commit.update(outstanding) + sel._agg_pending_commit_ref = self._sim_pending_commit + # Prune the (model_version, data_id, iteration) triplet guard to the + # still-outstanding set (a trainer whose grad committed is re-pickable). + self._trainer_state_dict = { + e: v for e, v in getattr(self, "_trainer_state_dict", {}).items() + if e in outstanding + } + # recv_fifo marked every delivered end RECVD, which the channel would + # strip from selected_ends; reset the still-buffered (returned-but- + # uncommitted) ends to NONE so they KEEP their slot until they commit. + for eid in buffered: + if channel is not None and channel.has(eid): + channel._ends[eid].set_property(KEY_END_STATE, VAL_END_STATE_NONE) + + if isinstance(all_selected, dict): + # Release slot + guard for every trainer no longer outstanding + # (its grad committed this cycle) -> re-pickable next cycle. + for eid in [e for e in list(all_selected.keys()) if e not in outstanding]: + del all_selected[eid] + if channel is not None and channel.has(eid): + channel._ends[eid].set_property(KEY_END_STATE, VAL_END_STATE_NONE) + if isinstance(selected_ends, dict) and requester in selected_ends: + selected_ends[requester].discard(eid) + # Hold every still-outstanding trainer in the re-pick guard. + for eid in outstanding: + if eid not in all_selected: + all_selected[eid] = time.time() + + # Compute slot: held by EVERY outstanding trainer (computing OR returned- + # but-uncommitted) — both are in flight in virtual time until commit. + if isinstance(selected_ends, dict) and requester in selected_ends: + for eid in outstanding: + selected_ends[requester].add(eid) + def _aggregate_grads_async(self, tag: str) -> None: """ Aggregate local model GRADIENTS asynchronously for FwdLLM. @@ -702,7 +1158,20 @@ def _aggregate_grads_async(self, tag: str) -> None: logger.info("No channel found") return - if channel.ends(VAL_CH_STATE_RECV) is None: + recv_ends = channel.ends(VAL_CH_STATE_RECV) + # Sim: the sct reorder buffer -- NOT the transient channel RECV state -- + # is the source of truth for what is committable. _sim_recv_min_grad + # greedily drains ALL ready channel messages into _sim_buffer on its + # first call (which empties RECV), so gating loop re-entry on RECV strands + # the already-received grads in the buffer: 10 grads received, 1 popped, + # then "no ends yet" forever -> agg_goal never met from a full cohort + # (fluxtune deadlock, §K-D17). Keep draining while the buffer holds grads + # or in-flight trainers are still expected. Real path unchanged + # (principle #8: RECV gating is a real-transport artifact, no sim analog). + sim_has_pending = self.simulated and ( + len(self._sim_buffer) > 0 or bool(self._sim_inflight_expected) + ) + if recv_ends is None and not sim_has_pending: logger.info("no ends yet") return # time.sleep(0.1) # Slight delay to allow messages to arrive @@ -716,10 +1185,19 @@ def _aggregate_grads_async(self, tag: str) -> None: # to run -- the run hangs past its configured budget until manually # killed. Matches the same pattern asyncfl/top_aggregator.py's # _aggregate_weights already uses for this exact reason. - msg, metadata = next( - channel.recv_fifo(channel.ends(VAL_CH_STATE_RECV), 1, - timeout=RECV_TIMEOUT_WAIT_S) - ) + # + # Sim branch (Batch 1): instead of committing on wall arrival, order the + # commit by the modeled completion sct via _sim_recv_min_grad (sct-ordered + # reorder buffer + in-flight gate + virtual-clock advance). It returns ONE + # (msg, metadata) matching next(recv_fifo(...,1)) -- one grad per call, as + # this loop expects. Real branch byte-identical to before. + if self.simulated: + msg, metadata = self._sim_recv_min_grad(channel, recv_ends or []) + else: + msg, metadata = next( + channel.recv_fifo(channel.ends(VAL_CH_STATE_RECV), 1, + timeout=RECV_TIMEOUT_WAIT_S) + ) end, timestamp = metadata if not msg: logger.debug(f"No data from {end}; skipping it") @@ -742,6 +1220,24 @@ def _aggregate_grads_async(self, tag: str) -> None: self._process_aggregation_goal_met(tag, channel, is_async=True) @timer_decorator + def _release_end_on_return(self, channel, end) -> None: + """Release a returned trainer's compute slot + re-pick guard on grad + RETURN -- except on the async sim residence path, where the grad is + carried and commits later in virtual time. There the guard/slot lifetime + is owned by `_sim_hold_busy_slots` (held to COMMIT); releasing + `all_selected` on physical RETURN would re-eligible a trainer whose + carried grad hasn't committed -> re-dispatch-while-in-flight -> R1 + violation. Real / non-residence: return ~= commit, release immediately. + """ + if self.is_async: + if getattr(self, "simulated", False) and getattr( + self, "_sim_inflight_residence", False + ): + return # guard/slot held to COMMIT by _sim_hold_busy_slots + channel.cleanup_provided_ends(end) + else: + channel.cleanup_recvd_end(end) + def _process_single_trainer_message(self, channel, msg, end, timestamp): # An end may only contribute once per (round, data_id, # iteration_per_data_id) collection cycle -- _per_agg_trainer_list is @@ -767,11 +1263,12 @@ def _process_single_trainer_message(self, channel, msg, end, timestamp): logger.info( f"Received grad with staleness={self._model_version-version}." ) - # Mode semantics: Hyperparameters.staleness_policy (flame/config.py). - # Implementation note: model_version alone already identifies - # (round, data_id) when inc_model_version_per_data_id is set (it - # only advances on a data_id transition), so "round_data_id" needs - # no extra fields; "exact" additionally checks iteration_per_data_id. + # staleness_policy. REJECT: round_data_id, exact. ACCEPT: none (no + # gate); fedbuff (consume + down-weight by V'-V) -- the async default, + # since its carried surplus grads (§L) are stale by construction. + # model_version alone identifies (round, data_id) when + # inc_model_version_per_data_id is set, so round_data_id needs no extra + # field; exact also checks iteration_per_data_id. policy = getattr(self, "staleness_policy", "none") stale, stale_reason = False, None if policy == "round_data_id": @@ -790,7 +1287,7 @@ def _process_single_trainer_message(self, channel, msg, end, timestamp): f"iteration_per_data_id={msg_iter} != " f"agg iteration_per_data_id={self.iteration_per_data_id}" ) - elif policy != "none": + elif policy not in ("none", "fedbuff"): logger.warning( f"Unrecognized staleness_policy={policy!r}; treating as 'none' " f"(no staleness gate)." @@ -813,6 +1310,9 @@ def _process_single_trainer_message(self, channel, msg, end, timestamp): f"with model version {msg[MessageType.MODEL_VERSION]}" ) self._agg_goal_cnt += 1 + # Wall of the most-recent accepted grad -> barrier_wait_s / drain_ + # tail_s in the per-round wall decomposition. + self._last_grad_wall_ts = time.time() channel.set_end_property( end, PROP_LAST_SELECTED_ROUND, msg[MessageType.MODEL_VERSION] @@ -841,6 +1341,22 @@ def _process_single_trainer_message(self, channel, msg, end, timestamp): logger.info( f"Set PROP_CLIENT_TASK_TRAIN_DURATION for {end}: {round_duration.total_seconds():.3f}s" ) + + # Record this contribution's [dispatch, commit] interval for R1/W1 + # (§L.3): sim uses the vclock interval the trainer echoes back (exact + # per contribution); real uses the wall interval (dispatch .. receipt). + if getattr(self, "_sim_contrib_intervals", None) is not None: + if self.simulated: + _disp = msg.get(MessageType.SIM_SEND_TS) + _comm = msg.get(MessageType.SIM_COMPLETION_TS) + else: + _disp = (round_start_time_tup[1].timestamp() + if round_start_time_tup is not None else None) + _comm = timestamp.timestamp() if hasattr(timestamp, "timestamp") else None + self._sim_contrib_intervals[end] = { + "dispatch_ts": float(_disp) if _disp is not None else None, + "commit_ts": float(_comm) if _comm is not None else None, + } else: logger.error( f"Invalid message received from {end} in aggregate_weights: {msg}" @@ -860,6 +1376,32 @@ def _process_single_trainer_message(self, channel, msg, end, timestamp): self._updates_in_queue += 1 self._per_agg_trainer_list.append(end) + # Canonical commit-order key: the trainer's pure modeled delay D + # (deterministic from the registry) + str(end) as tie-break. Lets + # _canonicalize_cohort_commit_order reproduce real's D-ordered arrival and + # break equal-D ties by trainer_id identically in both modes. None when + # delays off / D unstamped -> that cycle falls back to arrival order. + _md = msg.get(MessageType.MODELED_DELAY_S) + self._commit_key_by_end = getattr(self, "_commit_key_by_end", {}) + self._commit_key_by_end[end] = ( + (float(_md), str(end)) if _md is not None else None + ) + + # Re-pick guard (async_oort triplet filter): record the (model_version, + # data_id, iteration) this trainer contributed at so it is not re-selected + # until the agg version advances (the commit boundary prunes committed ends + # in _sim_hold_busy_slots). Stamped on RETURN, not dispatch: stamping the + # whole cohort at dispatch froze the eligible pool before any commit could + # advance the version -> re-dispatch deadlock. The compute slot already + # guards an in-flight trainer; the triplet only covers the return->commit + # carry window. Residence-only (sync baselines keep an empty map). + if getattr(self, "simulated", False) and getattr( + self, "_sim_inflight_residence", False + ): + _agg_ver = getattr(self, "_curr_agg_version", None) + if _agg_ver is not None: + self._trainer_state_dict[end] = _agg_ver + if end not in self._updates_received.keys(): self._updates_received[end] = 1 else: @@ -955,15 +1497,11 @@ def _process_single_trainer_message(self, channel, msg, end, timestamp): logger.info( f"Received grads from {end}. It was trained on model version {version}, with {count} samples" ) - # Not remove_from_selected_ends(): it never clears the selector's - # all_selected set, permanently blocking this trainer from future - # reselection. Async selectors only implement the batch - # _cleanup_recvd_ends/_cleanup_provided_ends path; cleanup_recvd_end() - # is sync-only (random selector). - if self.is_async: - channel.cleanup_provided_ends(end) - else: - channel.cleanup_recvd_end(end) + # Release the slot/guard on return -- but the async sim residence path + # defers that release to COMMIT; see _release_end_on_return. + # (cleanup_recvd_end() is sync-only, random selector; async selectors only + # implement the batch _cleanup_provided_ends path.) + self._release_end_on_return(channel, end) return True def _log_and_reset_model_version_stats(self): @@ -1077,12 +1615,62 @@ def _build_dynamic_kc_metrics(self, channel) -> dict: "target_iter_per_data_id": target_iter, } + def _canonicalize_cohort_commit_order(self): + """Reorder this cycle's cohort commits to a canonical (D, trainer_id) + order, identically in real and sim. + + Real receives updates in modeled-delay (D) order; sim commits in sct order + (= D order). The sole residual divergence is the tie-break when two + trainers share a D: real breaks it by physical arrival, sim by sct-sort. + Both tied members land in the same split-half so `var` is unchanged, but + the exact-order `cohort_sequence` rung flags the swap. Sort by (D, str(end)) + so equal-D ties break by trainer_id in both modes; var/grads stay + bit-identical (the aggregated grad is an order-independent sum). + + Scope = this cycle only. `grad_for_var_check_list` accumulates across a + data_id's iterations while `_per_agg_trainer_list` is per-cycle, so reorder + just the trailing len(cohort) slice (the sync barrier appends this cohort + contiguously in `_per_agg_trainer_list` order). + + No-op unless every contributor stamped a modeled delay AND a tie actually + changes the order. + """ + ends = self._per_agg_trainer_list + n = len(ends) + if n < 2: + return + keys = [self._commit_key_by_end.get(e) for e in ends] + if any(k is None for k in keys): + return # delays off / a contributor without a stamp → arrival order + perm = sorted(range(n), key=lambda i: keys[i]) + if perm == list(range(n)): + return # already canonical (the common, non-tie path) + self._per_agg_trainer_list = [ends[i] for i in perm] + # Reorder the trailing cohort slice of the accumulating var/jvp lists in + # lockstep. Length guard: only when the slice aligns 1:1 with this cohort + # (a mismatch means a non-grad message slipped in -> leave lists untouched). + for lst in (self.grad_for_var_check_list, self.jvp_for_snr_check_list): + if len(lst) >= n: + tail = lst[-n:] + lst[-n:] = [tail[i] for i in perm] + logger.info( + f"[COMMIT_CANON] equal-D tie → reordered {n}-cohort to (D,id) order " + f"(perm={perm}); var/grads unchanged, receive order now real↔sim identical." + ) + @timer_decorator def _process_aggregation_goal_met(self, tag, channel, is_async=False): logger.info( f"Aggregation goal {self._agg_goal} reached. Performing FwdLLM aggregation." ) + # Canonicalize this cohort's commit order to (D, trainer_id) before the + # telemetry snapshot and aggregate() so the recorded receive order and the + # split-half var use the same deterministic order in real and sim. Gated + # on commit-key state -> skipped when delays are off (arrival order). + if getattr(self, "_commit_key_by_end", None): + self._canonicalize_cohort_commit_order() + # Snapshot for this cycle's agg_round telemetry (emitted further down, # after self._per_agg_trainer_list is cleared and self._model_version # may have advanced -- see build_agg_round call below). @@ -1150,7 +1738,18 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): f"if variance check fails in aggregate()." ) + # Per-round wall decomposition (#6): the FedAvg merge, the dispatch->last- + # grad barrier wait, and the last-grad->commit drain tail (a real-transport + # artifact the sim's all-k barrier doesn't model). eval_s measured below. + # All wall (real); the sim advances the vclock, so these read ~0 there. + _agg_start_wall = time.time() self.aggregate(self._round) + _aggregate_fedavg_s = time.time() - _agg_start_wall + _disp = getattr(self, "_round_dispatch_wall_ts", None) + _lastg = getattr(self, "_last_grad_wall_ts", None) + _barrier_wait_s = (_lastg - _disp) if (_disp and _lastg) else None + _drain_tail_s = (_agg_start_wall - _lastg) if _lastg else None + _eval_s = None # set below only when the variance gate passes (eval runs) _var_thr = getattr(self, "var_threshold", None) _ratio = ( @@ -1167,6 +1766,32 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): f"force_commit_planned={_force_commit_planned}" ) + # Snapshot the cycle identity BEFORE the pass/fail branch mutates + # data_id/iteration_per_data_id. The emitted `data_id`/`iteration` fields + # are post-mutation (a commit advances data_id and zeroes iteration) -- + # fine for the progress axis but ambiguous for the variance-cadence rungs. + # `cycle_data_id`/`cycle_iteration` unambiguously identify the data_id this + # cycle worked on and its 0-based attempt index. + _cycle_data_id = self.data_id + _cycle_iteration = self.iteration_per_data_id + # Pool sizes at the variance gate (before a commit clears grad_pool): + # grad_pool = realized contributions this data_id (G2); cached_v = carried + # aggregated pool across variance-FAIL rollbacks (V3). getattr-guarded so + # test doubles without the pools still emit. + _grad_pool = getattr(self, "grad_pool", None) + _grad_pool_size = len(_grad_pool) if _grad_pool is not None else None + _cached_v = getattr(self, "cached_shared_grad_pool_trainable", None) + _cached_v_size = len(_cached_v) if _cached_v is not None else 0 + # Per-contributor [dispatch, commit] intervals for R1/W1: one entry per + # end that committed into this cycle. History preserved since each cycle + # emits its own list. + _contrib_map = getattr(self, "_sim_contrib_intervals", None) or {} + _contributor_intervals = [ + {"end": str(_e), **_contrib_map.get(_e, {"dispatch_ts": None, + "commit_ts": None})} + for _e in _cycle_contributors + ] + if self.var_good_enough: _pass_kind = ( "FORCE-COMMITTED (max_iter bypass)" @@ -1177,7 +1802,17 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): f"Variance check {_pass_kind}. Evaluating model and advancing data_id." ) self.iteration_per_data_id += 1 + _eval_start_wall = time.time() result, _, _ = self.eval_model() + _eval_s = time.time() - _eval_start_wall # genuine server-eval term + # #6: the server eval is genuine algorithmic time between committed + # data_ids that the sct model omits, so the vclock under-counts by + # ~eval_s/commit. When enabled, charge the measured eval wall to the + # vclock (genuine compute, not transport overhead). Config-gated OFF. + if self.simulated and getattr( + self.config.hyperparameters, "sim_model_eval_time", False + ): + self._vclock.advance(self._vclock.now + _eval_s) logger.info( f"Round {self._round}, Data ID {self.data_id} Eval Loss: {result['eval_loss']}" ) @@ -1211,6 +1846,12 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): else: self._model_version = self._round + # Opt-1: the data-bin (and model_version) just advanced -> every + # trainer is genuinely stale, so the "already sent this cycle" set must + # start empty. getattr-guarded for test doubles that bypass __init__. + if getattr(self, "_weights_sent_this_cycle", None) is not None: + self._weights_sent_this_cycle.clear() + self._log_and_reset_model_version_stats() if self.data_id == self.total_data_bins: @@ -1242,6 +1883,52 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): self._is_model_updated = False if telemetry.is_enabled(): + # Speedup instrumentation (#13): the sim must run virtual time faster + # than wall (sim_rate >= 1). fwdllm emitted vclock_now but no paired + # wall stamp, so a slowdown was invisible. Emit wall_elapsed_s in both + # modes (real needs it for wall_speedup = real_wall/sim_wall) and + # sim_rate = vclock/wall in sim only. agg_start_time_ts is re-anchored + # past the join wait, so this excludes the initial join stall. + _wall_elapsed_s = time.time() - getattr( + self, "agg_start_time_ts", time.time() + ) + _sim_rate = None + if self.simulated and getattr(self, "_vclock", None) is not None: + _sim_rate = ( + float(self._vclock.now) / _wall_elapsed_s + if _wall_elapsed_s > 0 + else None + ) + # Live speedup log (throttled ~30s). The inherited base + # [VCLOCK_PROGRESS] lives in increment_round, which fwdllm's + # composer loop bypasses, so emit it here. + _last = getattr(self, "_last_vclock_log_wall_ts", 0.0) + if time.time() - _last >= 30.0 and _sim_rate is not None: + logger.info( + f"[VCLOCK_PROGRESS] vclock={float(self._vclock.now):.1f}s " + f"wall={_wall_elapsed_s:.1f}s sim_rate={_sim_rate:.3f} " + f"(virtual-s/wall-s; <1 == SLOWDOWN) round={self._round} " + f"data_id={self.data_id}" + ) + self._last_vclock_log_wall_ts = time.time() + + # Intrinsic algorithmic span of this cycle (#6 anchor): the genuine + # per-cycle work the sim charges to the vclock -- the barrier (slowest + # committed trainer's intrinsic compute+delay) plus the server eval + # (commit cycles only). Mirrors the sim vclock composition exactly: + # the FedAvg merge is excluded because the sim does not charge it to + # the vclock. Emitted in both modes so the clock-rate rungs anchor real + # on intrinsic time instead of raw wall Δts (real's wall bundles a + # ~constant inter-round transport artifact the sim omits). The barrier + # uses the trainer intrinsic duration (PROP_CLIENT_TASK_TRAIN_DURATION, + # _cycle_speed_s), not the agg-side barrier_wait_s (which reads ~0 in + # real because trainers pipeline). max() = the sync barrier. + _barrier_span_s = max(_cycle_speed_s) if _cycle_speed_s else None + _intrinsic_span_s = ( + _barrier_span_s + (_eval_s or 0.0) + if _barrier_span_s is not None + else None + ) try: ev, fields = build_agg_round( round_num=self._round, @@ -1254,13 +1941,54 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): contributing_trainers=_cycle_contributors, agg_observed_s=_cycle_agg_observed_s, extra={ + # Virtual clock at commit (sim only). The parity engine + # gates its whole clock/throughput/convergence family on + # this; fwdllm never emitted it. + "vclock_now": (self._vclock.now if self.simulated + and getattr(self, "_vclock", None) else None), "data_id": self.data_id, "iteration_per_data_id": self.iteration_per_data_id, "var": self.var, "var_threshold": getattr(self, "var_threshold", None), "var_good_enough": self.var_good_enough, "force_commit_planned": _force_commit_planned, + # Opt-2: active stopping policy and why this cycle committed + # (natural gate / max-iter cap / plateau). None on non-commit + # cycles. Lets the reducer split commits by cause. + "stopping_policy": getattr(self, "_var_stopping_policy", None), + "commit_reason": getattr(self, "_force_commit_reason", None), + # Opt-3: active aggregation rate type + running count of + # anti-aligned updates the grad-aware gate down-weighted. + "agg_rate_type": (self.optimizer.agg_rate_conf.get("type") + if getattr(self, "optimizer", None) and + getattr(self.optimizer, "agg_rate_conf", None) + else None), + "grad_aware_gated_total": getattr( + self, "_grad_aware_gated_total", 0), "is_async": is_async, + # Variance-cadence rung inputs: cycle-relative identity for + # V1, pool sizes for V3/G2. + "cycle_data_id": _cycle_data_id, + "cycle_iteration": _cycle_iteration, + "grad_pool_size": _grad_pool_size, + "cached_v_size": _cached_v_size, + # R1/W1 residence rungs: per-contributor [dispatch_ts, + # commit_ts] intervals for this cycle. + "contributor_intervals": _contributor_intervals, + # Per-round wall decomposition (#6): barrier wait + drain + # tail (artifact) + fedavg + eval. + "barrier_wait_s": _barrier_wait_s, + "drain_tail_s": _drain_tail_s, + "aggregate_fedavg_s": _aggregate_fedavg_s, + "eval_s": _eval_s, + # #6 anchor: real's genuine per-cycle algorithmic time, + # the like-for-like counterpart to the sim's Δvclock; lets + # the clock-rate rungs exclude real's transport artifact. + "intrinsic_span_s": _intrinsic_span_s, + # Speedup metric (#13): wall in both modes; sim_rate = + # vclock/wall (sim only). sim_rate < 1 means slowdown. + "wall_elapsed_s": _wall_elapsed_s, + "sim_rate": _sim_rate, }, ) telemetry.emit(ev, **fields) @@ -1269,6 +1997,7 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): self._updates_in_queue -= self._agg_goal self._per_agg_trainer_list = [] + self._commit_key_by_end = {} # cohort-scoped logger.info( f"====== aggregation finished for round {self._round}, " @@ -1307,6 +2036,13 @@ def _process_aggregation_goal_met(self, tag, channel, is_async=False): ) channel.cleanup_recvd_ends() + # Sim rollback-safety: clear this cycle's committed marks + stranded + # reorder-buffer entries + in-flight gate, and (async) release held slots, + # so the next agg-goal cycle of a rolled-back data_id starts clean (§I.5). + # Reached by both the variance-PASS and -FAIL branches. Inert in real. + if self.simulated: + self._release_sim_slots_at_agg_goal(channel, is_async) + # Centralized cleanup # self._force_cuda_memory_cleanup() @@ -1322,7 +2058,10 @@ def sync_collect_and_accumulate_grads(self, tag, channel): num_min_req = self._agg_goal # change hardcoding, set it to aggGoal logger.info(f"Total ends: {len(recv_ends)}, required : {num_min_req}") num_min_req = min(num_min_req, len(recv_ends)) - if self.ends_not_selected_yet: + # Real-only: "commit 1 per pass" relies on uncommitted msgs persisting in + # the queue. The sim barrier drains + drops past first_k, so clamping to 1 + # strands the cohort -> deadlock (§F #8). + if self.ends_not_selected_yet and not self.simulated: logger.info(f"We are waiting to clear up queue") num_min_req = min(num_min_req, 1) @@ -1333,20 +2072,56 @@ def sync_collect_and_accumulate_grads(self, tag, channel): # re-check _check_early_stop_conditions() (max_runtime_s/ # max_data_id_progress), and the run hangs past its configured # budget until manually killed. Same fix as _aggregate_grads_async. - for msg, metadata in channel.recv_fifo(channel.ends(), num_min_req, - timeout=RECV_TIMEOUT_WAIT_S): - end, timestamp = metadata - if not msg: - logger.info(f"No data from {end}; skipping it") - continue + # + # Sim barrier: instead of committing by physical arrival, + # _sync_sim_recv_first_k commits the num_min_req trainers with the smallest + # modeled sim_completion_ts (the k that would finish first in real), + # advancing the vclock to the k-th smallest -- immune to arrival jitter. It + # returns an ascending-sct list and stamps PROP_CLIENT_TASK_TRAIN_DURATION + # per commit; each is fed through the same per-message path. Real unchanged. + if self.simulated: + committed = self._sync_sim_recv_first_k( + channel, channel.ends(), num_min_req + ) + _barrier_durs = [] + for msg, metadata in committed: + end, timestamp = metadata + if not msg: + continue + # dispatch-relative completion for the U6 barrier anchor = the + # modeled round duration. + _barrier_durs.append( + msg.get(MessageType.SIM_CLIENT_TASK_TRAIN_DURATION_S) + ) + self._process_single_trainer_message(channel, msg, end, timestamp) + if self._agg_goal_cnt >= self._agg_goal: + logger.info( + f"Reached agg_goal of {self._agg_goal} (sim barrier); " + f"proceeding to aggregate." + ) + break + # U6 visibility lag on a strict sync barrier: lag_i = max_completion - + # completion_i over the committed cohort. Stashed for the U6 rung. + self._sync_barrier_lags_s = self._barrier_anchored_lags(_barrier_durs) + logger.info( + f"[SYNC_SIM_BARRIER] round={self._round} committed={len(committed)} " + f"T_v={self._vclock.now:.1f} barrier_lags_s={self._sync_barrier_lags_s}" + ) + else: + for msg, metadata in channel.recv_fifo(channel.ends(), num_min_req, + timeout=RECV_TIMEOUT_WAIT_S): + end, timestamp = metadata + if not msg: + logger.info(f"No data from {end}; skipping it") + continue - self._process_single_trainer_message(channel, msg, end, timestamp) + self._process_single_trainer_message(channel, msg, end, timestamp) - if self._agg_goal_cnt >= self._agg_goal: - logger.info( - f"Reached agg_goal of {self._agg_goal} since agg_goal_count is {self._agg_goal_cnt}. Breaking from for loop, proceeding to aggregate." - ) - break + if self._agg_goal_cnt >= self._agg_goal: + logger.info( + f"Reached agg_goal of {self._agg_goal} since agg_goal_count is {self._agg_goal_cnt}. Breaking from for loop, proceeding to aggregate." + ) + break # Second loop @@ -1586,9 +2361,14 @@ def check_trainer_availability(self, end: str) -> bool: @timer_decorator def _prepare_distribution_payload(self, task_to_perform: str, force_weights: bool = False): """Build a WEIGHTS payload (always) or a VAR=bad payload (when var fails and force_weights=False).""" - if not self.var_good_enough and not force_weights: + if not force_weights: + # force_weights=False always means "build the tiny VAR=bad keep- + # training message". Previously gated on `not var_good_enough`, which + # returned WEIGHTS when var_good_enough=True, so Opt-1 could not + # downgrade a commit-branch re-send. The message carries the current + # model_version, so a trainer that cached it just keeps training. logger.info( - f"[PreparePayload/VAR=bad] var={self.var} > " + f"[PreparePayload/VAR=bad] var={self.var} vs " f"thr={self.var_threshold}; payload asks trainer to keep " f"training on current model_version={self._model_version}." ) @@ -1773,6 +2553,137 @@ def _select_ends_respecting_reselect_gate(self, channel, task_to_perform: str): return merged return new_ends + def _await_dispatchable_under_scarcity(self, task_to_perform: str) -> None: + """Real-mode sync-barrier liveness under availability scarcity. + + When a trace keeps the eligible pool below `agg_goal`, the plain loop + hot-re-dispatches every pass, burning the wall budget with no progress and + ballooning the log. Parity-faithful fix: keep the cohort == `agg_goal` and + wait for availability to recover (sleep-to-next-avail) instead of spinning + -- the sim assembles the same full cohort by jumping its vclock past the + unavailable window, so cohort size stays identical. Self-terminates via + `_check_early_stop_conditions`. No-op when availability tracking is off or + on the sim path (the vclock, not a wall sleep, models the wait). + """ + if self.simulated or getattr(self, "trainer_event_dict", None) is None: + return + # Only a genuine post-join scarcity, not startup join-lag: wait only once + # at least `agg_goal` trainers have joined. + if len(getattr(self, "all_trainers", ())) < self._agg_goal: + return + poll_s = float(getattr(self.config.hyperparameters, "scarcity_poll_s", 2.0)) + warned = False + while not self._work_done: + unavail = set(self.get_curr_unavail_trainers()) + contributed = set(self._per_agg_trainer_list or []) + # A trainer can still advance this cycle's barrier iff it is available + # AND has not already contributed to it. + dispatchable = [ + e for e in self.all_trainers + if e not in unavail and e not in contributed + ] + if dispatchable or self._agg_goal_cnt >= self._agg_goal: + return # progress possible this pass (or barrier already met) + if not warned: # one line per stall, not one per spin + logger.warning( + f"[SYNC_SCARCITY_WAIT] round={self._round} data_id={self.data_id} " + f"committed={self._agg_goal_cnt}/{self._agg_goal} " + f"unavail={len(unavail)}/{len(self.all_trainers)}; waiting for " + f"availability (poll={poll_s}s) instead of spin-dispatching." + ) + warned = True + time.sleep(poll_s) + self._check_early_stop_conditions() # self-terminate at max_runtime_s + + def _should_send_full_weights(self, end, is_stale: bool) -> bool: + """Decide WEIGHTS vs the tiny VAR=bad 'keep training' message for one end. + + Shared by both the sync and async distribute loops so the two paths stay + identical. Legacy behavior (flag OFF): send full WEIGHTS on a commit + (`var_good_enough`) or when the end is stale. + + Opt-1 (charter §5c, flag ON): within a data-bin the WEIGHTS+GRAD_POOL + payload is byte-identical across iterations; an end already sent it this + data-bin has it cached, so a re-send is pure redundancy -> downgrade to + VAR=bad. The legacy guard missed this because the redundancy is repeated + distribute calls on the var_good_enough=True branch (which sends WEIGHTS + unconditionally, so the stale check never ran). Fix: a send-time set + (cleared on model_version advance) checked before both branches. + """ + # Checked FIRST so it gates both the commit branch and the stale branch: + # the dominant redundancy is repeated distribute calls within one data-bin, + # each re-shipping the byte-identical model to trainers that already hold + # it. Once an end got this model_version's payload this cycle -> VAR=bad. + if self._suppress_redundant_weights and end in self._weights_sent_this_cycle: + return False + if self.var_good_enough: + # Commit / first distribution of this model_version: send it once. + return True + # Legacy stale path (kept so flag OFF is byte-identical): a trainer the + # return-map still shows on an older version gets the weights. + return is_stale + + def _should_force_commit_on_plateau(self) -> bool: + """Opt-2 (charter §5c) variance-plateau rule -- pure decision (reads only + instance attrs, unit-testable). Called from FedSGDAggregator.aggregate() + once this cycle's var has been appended to `var_prev_iter_list`. + + Returns True iff the 'plateau' policy is active AND the per-bin variance + curve has flattened (relative drop over the last N cycles < rel_delta) + while var is still above the commit threshold -- the "more denoising buys + nothing" signal. Policy off/absent, fewer than N+1 samples, or a + non-positive baseline => False. + """ + if getattr(self, "_var_stopping_policy", None) != "plateau": + return False + hist = self.var_prev_iter_list + n = getattr(self, "_var_plateau_patience", 3) + eps = getattr(self, "_var_plateau_rel_delta", 0.10) + if len(hist) <= n or hist[-1 - n] <= 0: + return False + rel_drop = (hist[-1 - n] - hist[-1]) / hist[-1 - n] + return rel_drop < eps and hist[-1] > self.var_threshold + + @staticmethod + def _grad_aware_rate(base_rate, cos, var_i, var_ref, *, align_gate=True, + inverse_var=False, align_floor=0.0, var_eps=1e-8): + """Opt-3 gradient-aware aggregation weight -- pure scalar math. Replaces + the scalar staleness×utility rate with a direction/reliability-aware + weight, bounded so it only ever down-weights (result <= base_rate), so it + never inflates the effective server LR. + * align_gate (primary): down-weight an update whose direction opposes the + running aggregate. factor = 1 for cos >= align_floor, linearly -> 0 at + cos = -1. Fixes averaging anti-aligned JVP estimates (mid-run instability). + * inverse_var (optional): down-weight an update noisier than the reference + var_ref: factor = min(1, var_ref / (var_i + eps)). + cos / var_i == None => that factor is 1 (e.g. first update of a cycle). + """ + rate = float(base_rate) + if align_gate and cos is not None and cos < align_floor: + denom = align_floor + 1.0 + rate *= (max(0.0, (cos + 1.0) / denom) if denom > 0 else 0.0) + if inverse_var and var_i is not None: + rate *= min(1.0, var_ref / (var_i + var_eps)) + return rate + + @staticmethod + def _cosine_flat(grad_named, running_grad, named_params): + """Cosine between a trainer's gradient (dict name→tensor) and the running + aggregate `self.grad` (list indexed by `named_params` order), flattened over + all trainable params. Returns None if either side has ~zero norm (e.g. the + running aggregate before any update has landed this cycle).""" + dot = nt = ng = 0.0 + for i, (name, _p) in enumerate(named_params): + if name in grad_named: + t = grad_named[name] + g = running_grad[i].to(t.device) + dot += float((t * g).sum()) + nt += float((t * t).sum()) + ng += float((g * g).sum()) + if nt <= 0.0 or ng <= 0.0: + return None + return dot / (math.sqrt(nt) * math.sqrt(ng)) + @timer_decorator def _distribute_weights_sync( self, tag: str, task_to_perform: str = "train" @@ -1792,6 +2703,9 @@ def _distribute_weights_sync( return channel.await_join() + # Sleep-to-next-avail under real-mode scarcity instead of hot-spinning + # (Stage C); no-op on the sim path / when availability tracking is off. + self._await_dispatchable_under_scarcity(task_to_perform) global_model_params = self.get_global_model_params() format_hash = lambda d: {k: _calculate_hash(v)[:8] for k, v in d.items()} logging.info( @@ -1799,9 +2713,14 @@ def _distribute_weights_sync( ) self.weights = global_model_params - logger.debug(f"Starting busy wait at time {time.time()}") - time.sleep(0.1) - logger.debug(f"Ended busy wait at time {time.time()}") + # Real-transport pad to let just-distributed messages settle before the + # selection read (real-mode MQTT artifact, #8). No sim analog: the sim + # orders by sct, not physical arrival, so this is pure wall overhead there + # -- skip it. Real unchanged. + if not self.simulated: + logger.debug(f"Starting busy wait at time {time.time()}") + time.sleep(0.1) + logger.debug(f"Ended busy wait at time {time.time()}") if self.trainer_event_dict is not None: curr_unavail_trainer_list = self.get_curr_unavail_trainers() @@ -1840,24 +2759,43 @@ def _distribute_weights_sync( payload_without_weights = None payload_with_weights = self._prepare_distribution_payload(task_to_perform, force_weights=True) - if not self.var_good_enough: + # Opt-1: also need the VAR=bad variant when suppression is on (a commit- + # branch re-send to an already-served end downgrades to VAR=bad). + if not self.var_good_enough or self._suppress_redundant_weights: payload_without_weights = self._prepare_distribution_payload(task_to_perform, force_weights=False) - + self._update_state_after_payload_prepared() + # Sim-clock dispatch stamp: in the sync barrier every selected end + # dispatches at the same virtual instant, so one _round_now stamp is + # injected into each payload variant (the trainer bases its modeled sct on + # SIM_SEND_TS) and recorded as a per-end property. Inert in real mode + # (SIM_SEND_TS absent -> arrival order). + _round_now = self._vclock.now if self.simulated else None + if self.simulated: + for _p in (payload_with_weights, payload_without_weights): + if _p is not None: + _p[MessageType.SIM_SEND_TS] = _round_now + for end in ends: trainer_version = self._trainer_last_model_version.get(end, -1) is_stale = (trainer_version != self._model_version) - if self.var_good_enough: + # Opt-1: shared decision (identical in the async path). Flag OFF -> + # legacy (var_good_enough or is_stale). Flag ON -> suppress a + # byte-identical intra-databin re-send to VAR=bad. + send_weights = self._should_send_full_weights(end, is_stale) + if send_weights: payload = payload_with_weights - else: - if is_stale: - payload = payload_with_weights + if self._suppress_redundant_weights: + self._weights_sent_this_cycle.add(end) + if not self.var_good_enough and is_stale: logger.info(f"Trainer {end} hasn't received weights for model_version {self._model_version} (has {trainer_version}). Sending WEIGHTS payload instead of VAR=bad.") - else: - payload = payload_without_weights + else: + payload = payload_without_weights + if self._suppress_redundant_weights and is_stale: + self._redundant_weights_suppressed_total += 1 logger.debug( f"Setting channel property {PROP_ROUND_START_TIME} for " @@ -1866,8 +2804,13 @@ def _distribute_weights_sync( channel.set_end_property( end, PROP_ROUND_START_TIME, (self._round, datetime.now()) ) + if self.simulated: + channel.set_end_property(end, PROP_SIM_SEND_TS, _round_now) - if self.var_good_enough: + # Label/size by the actual payload (send_weights), not var_good_enough: + # keying off the latter mislabeled a var_good_enough=False is_stale + # WEIGHTS send as var_bad, undercounting sync weight bytes. + if send_weights: logger.info( f"sending weights to {end} with model_version: {self._model_version}, data_id: {self.data_id} for task: {task_to_perform}" ) @@ -1880,6 +2823,8 @@ def _distribute_weights_sync( for key, value in payload.items() } total_size_mb = sum(sizes_mb.values()) + _send_bytes = int(round(total_size_mb * 1024 * 1024)) + _payload_kind = "weights" logger.info( f"[DEBUG] Payload size breakdown for {end}: " @@ -1892,14 +2837,34 @@ def _distribute_weights_sync( ) msg_bytes = pickle.dumps(payload) + _send_bytes = len(msg_bytes) + _payload_kind = "var_bad" logger.info( f"[DEBUG] Payload size for {end}: {len(msg_bytes) / (1024 * 1024):.2f} MB" ) + # Network telemetry: one dispatch message onto the wire. Size reuses + # the value already computed for the debug log above. + try: + ev, f = build_comm( + direction="agg_to_trainer", size_bytes=_send_bytes, peer_id=str(end), + round_num=int(self._round), data_id=self.data_id, + iteration=self.iteration_per_data_id, payload_kind=_payload_kind, + n_tensors=len(payload), + ) + telemetry.emit(ev, **f) + except Exception as e: + logger.debug(f"comm telemetry emit failed (agg send): {e}") + channel.send(end, payload) logger.info(f"Sent weights to {end}") # self.invoke_gc() + # Cohort-dispatch wall -> barrier_wait_s anchor. Sync barrier: all `ends` + # dispatch in this one pass, marking the start of the dispatch->last-grad + # window. + self._round_dispatch_wall_ts = time.time() + @timer_decorator def _distribute_weights_async( self, tag: str, task_to_perform: str = "train" @@ -1914,9 +2879,14 @@ def _distribute_weights_async( channel.await_join() global_model_params = self.get_global_model_params() self.weights = global_model_params - logger.debug(f"Starting busy wait at time {time.time()}") - time.sleep(0.1) - logger.debug(f"Ended busy wait at time {time.time()}") + # Real-transport pad to let just-distributed messages settle before the + # selection read (real-mode MQTT artifact, #8). No sim analog: the sim + # orders by sct, not physical arrival, so this is pure wall overhead there + # -- skip it. Real unchanged. + if not self.simulated: + logger.debug(f"Starting busy wait at time {time.time()}") + time.sleep(0.1) + logger.debug(f"Ended busy wait at time {time.time()}") if self.trainer_event_dict is not None: curr_unavail_trainer_list = self.get_curr_unavail_trainers() channel.set_curr_unavailable_trainers( @@ -1970,20 +2940,52 @@ def _distribute_weights_async( payload_var_bad = None payload_weights = self._prepare_distribution_payload(task_to_perform, force_weights=True) - if not self.var_good_enough: + # Opt-1: also need VAR=bad when suppression is on (commit-branch re-sends + # to already-served ends downgrade to VAR=bad). + if not self.var_good_enough or self._suppress_redundant_weights: payload_var_bad = self._prepare_distribution_payload(task_to_perform, force_weights=False) self._update_state_after_payload_prepared() + # Network telemetry: serialized size of each payload variant, computed once + # per distribute (the async path has no per-end debug-size log), reused per + # end below. + try: + _bytes_weights = len(pickle.dumps(payload_weights)) if payload_weights is not None else 0 + _bytes_var_bad = len(pickle.dumps(payload_var_bad)) if payload_var_bad is not None else 0 + except Exception as e: + _bytes_weights = _bytes_var_bad = 0 + logger.debug(f"comm telemetry size calc failed (async): {e}") + + # Sim-clock dispatch stamp: the async path arms the in-flight gate + # (_sim_inflight_expected[end] = send vclock + a lower-bound budget) so the + # reorder-buffer drain can't lap a trainer still expected to complete. + # Inert in real mode. + # #13 (staggered): when on, each end's SEND vclock is popped from the + # freed-slot FIFO instead of the shared round frontier, spreading a cohort's + # expected completions across the timeline. Off/real => one shared + # _round_now injected into the two payload variants. + _round_now = self._vclock.now if self.simulated else None + _staggered = self.simulated and getattr(self, "_sim_staggered_redispatch", False) + if self.simulated and not _staggered: + for _p in (payload_weights, payload_var_bad): + if _p is not None: + _p[MessageType.SIM_SEND_TS] = _round_now + _n_weights_sent = 0 _n_var_bad_sent = 0 for end in ends: trainer_version = self._trainer_last_model_version.get(end, -1) is_stale = (trainer_version != self._model_version) - - if self.var_good_enough or is_stale: + # Opt-1: shared decision with the sync path (_should_send_full_weights). + # Flag OFF ⇒ (var_good_enough or is_stale) ⇒ byte-identical to legacy. + send_weights = self._should_send_full_weights(end, is_stale) + if send_weights: payload = payload_weights + _pk = "weights" # kind set here (survives staggered rebuild below) _n_weights_sent += 1 + if self._suppress_redundant_weights: + self._weights_sent_this_cycle.add(end) if not self.var_good_enough and is_stale: logger.debug( f"[Distribute] Trainer {end} stale " @@ -1992,7 +2994,12 @@ def _distribute_weights_async( ) else: payload = payload_var_bad + _pk = "var_bad" _n_var_bad_sent += 1 + # Opt-1: this end read stale but already got this version's payload + # this cycle -> a redundant full re-send we just avoided. + if self._suppress_redundant_weights and is_stale: + self._redundant_weights_suppressed_total += 1 logger.debug( f"[Distribute] Trainer {end} has current v{trainer_version}; " f"sending VAR=bad (keep training)." @@ -2005,11 +3012,63 @@ def _distribute_weights_async( channel.set_end_property( end, PROP_ROUND_START_TIME, (self._round, datetime.now()) ) + if self.simulated: + # R1 tripwire: dispatching a trainer already outstanding (in flight + # or grad carried-but-uncommitted) IS the residence violation -- + # surface it at dispatch instead of reconstructing it offline. + if end in self._sim_inflight_expected or self._sim_buffer.has(end): + logger.warning( + f"[SIM_R1_DISPATCH] end={end} re-dispatched while still " + f"outstanding (vclock={_round_now}) -- R1 residence violation" + ) + # #13: SEND vclock = the freed-slot stamp (staggered) else the + # shared round frontier. _pop_free_slot_ts is FIFO + clamped <= now; + # an empty queue falls back to _round_now. + _sst = self._pop_free_slot_ts(_round_now) if _staggered else _round_now + channel.set_end_property(end, PROP_SIM_SEND_TS, _sst) + # Expected completion = SEND vclock + a lower-bound budget + # (this end's own last-observed TRAINING_BUDGET_S, else the + # running min) so the gate never laps a not-yet-committed trainer. + _budget = self._sim_trainer_budget.get(end, self._sim_budget_min) + self._sim_inflight_expected[end] = _sst + _budget + # Staggered: this end's payload must carry its OWN SIM_SEND_TS, so + # rebuild a shallow copy (weights shared by ref; small vs GPU cost). + if _staggered and isinstance(payload, dict): + payload = dict(payload) + payload[MessageType.SIM_SEND_TS] = _sst + # Once dispatched a trainer is in flight in virtual time -> add to + # the pending-commit set so the selector's eligibility filter + # excludes it until its grad COMMITS (discarded in _sim_recv_min_grad). + self._sim_pending_commit.add(end) + # NOTE: the async_oort re-pick triplet is stamped on grad RETURN, + # not here at dispatch -- stamping the whole cohort at the current + # _curr_agg_version froze the eligible pool before any commit could + # advance the version (re-dispatch deadlock). An in-flight trainer + # is already guarded by its compute slot. + # One dispatch message onto the wire (async path). + try: + _sz = _bytes_weights if _pk == "weights" else _bytes_var_bad + ev, f = build_comm( + direction="agg_to_trainer", size_bytes=_sz, peer_id=str(end), + round_num=int(self._round), data_id=self.data_id, + iteration=self.iteration_per_data_id, payload_kind=_pk, + n_tensors=len(payload) if isinstance(payload, dict) else None, + ) + telemetry.emit(ev, **f) + except Exception as e: + logger.debug(f"comm telemetry emit failed (agg async send): {e}") channel.send(end, payload) + # #15 compute-truthful gate: stamp the wall time this end was dispatched + # so the drain can tell a live straggler from an idle-in-recv phantom. + # Sim-only; inert unless the gate flag is on. + if self.simulated: + self._sim_dispatch_wall[end] = time.time() logger.info( f"[Distribute] Done. Sent {_n_weights_sent} WEIGHTS + " f"{_n_var_bad_sent} VAR=bad payloads to {len(ends)} trainers " - f"(model_version={self._model_version}, data_id={self.data_id})." + f"(model_version={self._model_version}, data_id={self.data_id}, " + f"iter={self.iteration_per_data_id}, " + f"redundant_weights_suppressed_total={self._redundant_weights_suppressed_total})." ) ends_in_recv_state = channel.ends(VAL_CH_STATE_RECV) @@ -2051,13 +3110,45 @@ def _check_early_stop_conditions(self) -> None: max_runtime_s = getattr(self.config.hyperparameters, "max_runtime_s", None) if max_runtime_s is not None: - elapsed = time.time() - self.agg_start_time_ts + # Mode-dependent clock (mirrors base syncfl `max_experiment_runtime_s`): + # real -> wall seconds; sim -> virtual seconds (`vclock.now`). One budget + # then means the same amount of real work AND modeled work -- what the + # matched-budget convergence parity needs. Until #6 (vclock models real + # wall) lands the sim vclock under-counts, so the wall failsafe below + # fires first -- itself the #6 signal. + if self.simulated and hasattr(self, "_vclock"): + elapsed = float(self._vclock.now) + clock_label = "sim/vclock" + else: + elapsed = time.time() - self.agg_start_time_ts + clock_label = "real/wall" if elapsed >= float(max_runtime_s): logger.info( f"max_runtime_s={max_runtime_s}s reached " - f"(elapsed={elapsed:.0f}s); stopping run." + f"({clock_label}_elapsed={elapsed:.0f}s); stopping run." ) self._work_done = True + return + # Sim wall failsafe: the sim does real GPU + server eval, so its wall + # legitimately exceeds the vclock budget -- a 1x ceiling truncated every + # sim before its paired vclock stop. Decoupled to max_runtime_s × + # SIM_WALL_CEILING_FACTOR (a runaway OUTER safety), overridable by + # `sim_wall_ceiling_s`. Primary stops stay max_data_id_progress / + # vclock >= max_runtime_s. + if self.simulated: + _wall = time.time() - self.agg_start_time_ts + _ceil = getattr(self.config.hyperparameters, "sim_wall_ceiling_s", None) + _ceil = (float(_ceil) if _ceil is not None + else float(max_runtime_s) * self.SIM_WALL_CEILING_FACTOR) + if _wall > _ceil: + logger.warning( + f"[SIM_WALL_CEILING] wall={_wall:.0f}s > ceiling={_ceil:.0f}s " + f"(={self.SIM_WALL_CEILING_FACTOR:g}× budget {max_runtime_s}s; " + f"vclock={self._vclock.now:.0f}s) -- runaway outer safety " + f"tripped (a healthy sim stops on max_data_id/vclock first); " + f"stopping run." + ) + self._work_done = True def _async_inner_loop_done(self) -> bool: """Exit condition for the async/hybrid compose path's inner diff --git a/lib/python/flame/mode/horizontal/syncfl/fwdllm_trainer.py b/lib/python/flame/mode/horizontal/syncfl/fwdllm_trainer.py index af5e2b537..f984da9f7 100644 --- a/lib/python/flame/mode/horizontal/syncfl/fwdllm_trainer.py +++ b/lib/python/flame/mode/horizontal/syncfl/fwdllm_trainer.py @@ -19,6 +19,7 @@ import logging import math import time +from contextlib import contextmanager import torch from flame.channel import VAL_CH_STATE_HTBT_SEND, VAL_CH_STATE_RECV, VAL_CH_STATE_SEND @@ -39,6 +40,8 @@ from flame.mode.composer import Composer from flame.mode.message import MessageType from flame.mode.role import Role +from flame import telemetry +from flame.telemetry.events import build_task_recv, build_comm from flame.mode.tasklet import Loop, Tasklet from flame.optimizers import optimizer_provider from flame.privacies import privacy_provider @@ -159,6 +162,22 @@ def internal_init(self) -> None: self.abort_training = False self._stat_utility = 0 + # Per-round phase-timing accumulator (#6 wall decomposition). Reset each + # fetch; drained into the trainer_round telemetry `extra`. Mirrors the + # base syncfl trainer's _phase/_phase_times. + self._phase_times: dict = {} + + @contextmanager + def _phase(self, name: str): + """Time a named phase and accumulate into self._phase_times.""" + t0 = time.time() + try: + yield + finally: + self._phase_times[name] = self._phase_times.get(name, 0.0) + ( + time.time() - t0 + ) + def get(self, tag: str) -> None: """Get data from remote role(s).""" if tag == TAG_FETCH: @@ -172,6 +191,8 @@ def _fetch_weights(self, tag: str) -> None: ) self.fetch_success = False + # Reset per-round phase accumulator at the round boundary. + self._phase_times = {} channel = self.cm.get_by_tag(tag) if not channel: logger.info( @@ -192,7 +213,10 @@ def _fetch_weights(self, tag: str) -> None: # one aggregator is sufficient end = channel.one_end(VAL_CH_STATE_RECV) + _recv_start = time.time() msg, _ = recv_wrapper(self, channel, end) + # agg->trainer delivery + payload transfer (leg i); the first phase term. + self._phase_times["mqtt_fetch_s"] = time.time() - _recv_start if not msg: logger.info(f"NO msg received for trainer_id {self.trainer_id}") @@ -207,9 +231,35 @@ def _fetch_weights(self, tag: str) -> None: logger.info(f"New message received for trainer_id {self.trainer_id}") + # Sim-clock stamps: SIM_SEND_TS is the aggregator's virtual-clock "now" at + # dispatch, the base for the trainer's modeled completion sct; + # _wall_recv_ts is the real receipt wall time, echoed back so the + # aggregator can derive the intrinsic (server-overhead-free) task duration + # (WALL_SEND - WALL_RECV). Both inert in real mode (SIM_SEND_TS absent). + self._sim_send_ts = msg.get(MessageType.SIM_SEND_TS) + self._wall_recv_ts = time.time() + if MessageType.ROUND in msg: self._round = msg[MessageType.ROUND] + # Emit task_recv carrying sim_send_ts (#8): fwdllm overrode _fetch_weights + # and dropped the base trainer's emission; restore it here. None in real + # mode (SIM_SEND_TS absent). + if telemetry.is_enabled(): + try: + _avl = getattr(getattr(self, "avl_state", None), "value", None) + ev, fields = build_task_recv( + round_num=int(self._round), + trainer_id=str(getattr(self, "trainer_id", "")), + time_mode=getattr(self, "time_mode", "real"), + sim_send_ts=(float(self._sim_send_ts) + if self._sim_send_ts is not None else None), + avl_state=_avl, + ) + telemetry.emit(ev, **fields) + except Exception as e: # telemetry must never break training + logger.debug(f"task_recv telemetry emit failed: {e}") + logger.info( f"Checking DataID: {self.data_id}| MessageType.DATA_ID in msg: {msg.get(MessageType.DATA_ID)}| IterationPerDataID: {self.iteration_per_data_id}| MessageType.ITERATION_PER_DATA_ID in msg: {msg.get(MessageType.ITERATION_PER_DATA_ID)}" ) @@ -290,13 +340,15 @@ def _fetch_weights(self, tag: str) -> None: # Update the model logger.info(f"Weights received: # {msg[MessageType.WEIGHTS]}") self.weights = # weights_to_model_device(msg[MessageType.WEIGHTS], self.model) - trainable_weights = weights_to_model_device( - msg[MessageType.WEIGHTS], self.model - ) - full_state_dict = self.model.state_dict() - full_state_dict.update(trainable_weights) - self.weights = full_state_dict - self._update_model() + with self._phase("weights_to_ram_s"): + trainable_weights = weights_to_model_device( + msg[MessageType.WEIGHTS], self.model + ) + full_state_dict = self.model.state_dict() + full_state_dict.update(trainable_weights) + self.weights = full_state_dict + with self._phase("weights_to_gpu_s"): + self._update_model() if MessageType.MODEL_VERSION in msg: self._model_version = msg[MessageType.MODEL_VERSION] @@ -499,8 +551,25 @@ def _send_grads(self, tag: str) -> None: format_hash = lambda d: {k: _calculate_hash(v)[:8] for k, v in d.items()} logger.info(f"Sending grads from Trainer: {self.trainer_id} - model version: {self._model_version} - grad: {format_hash(grad_dict)} - grad_for_var_check: {_calculate_hash(self.grad_for_var_check)}") else: + total_bytes = 0 logger.info("No gradients exist; sending an empty dictionary.") + # Network telemetry: the update this trainer uploads. total_bytes is + # the gradient payload the debug log already reports (fluxtune's comm + # story is that this is small). + if telemetry.is_enabled(): + try: + ev, f = build_comm( + direction="trainer_to_agg", size_bytes=total_bytes, + peer_id=str(end), round_num=int(self._round), + data_id=self.data_id, iteration=self.iteration_per_data_id, + payload_kind="gradients", n_tensors=len(grad_dict), + trainer_id=self.trainer_id, + ) + telemetry.emit(ev, **f) + except Exception as e: + logger.debug(f"comm telemetry emit failed (trainer send): {e}") + logger.debug(f"self.jvp_for_snr_check on trainer before sending message: {self.jvp_for_snr_check}") msg = { @@ -518,11 +587,36 @@ def _send_grads(self, tag: str) -> None: MessageType.STAT_UTILITY: self._stat_utility, # - rn FedSgdTrainer has no utility MessageType.TOTAL_DATA_BINS: self.total_data_bins, + # Sim-clock stamps: the modeled completion sct the aggregator's + # reorder buffer keys on, the additive modeled round duration + # (completion budget for the async in-flight gate), and the real + # wall send/recv pair for the intrinsic task duration. All None in + # real mode -> aggregator uses arrival order. + MessageType.SIM_COMPLETION_TS: self._sim_completion_ts, + # Pure modeled delay D: deterministic from the registry (unlike + # SIM_COMPLETION_TS, which folds in GPU jitter), so the aggregator + # orders this cohort's commits by (D, trainer_id) identically in + # real and sim. Stamped in BOTH modes; None when delays are off. + MessageType.MODELED_DELAY_S: getattr(self, "_modeled_delay_s", None), + # Echo the dispatch stamp so the aggregator can reconstruct this + # contribution's [dispatch, completion] interval for R1. + MessageType.SIM_SEND_TS: self._sim_send_ts, + MessageType.SIM_CLIENT_TASK_TRAIN_DURATION_S: self._sim_round_duration_s, + MessageType.TRAINING_BUDGET_S: self._sim_round_duration_s, + MessageType.WALL_SEND_TS: time.time(), + MessageType.WALL_RECV_TS: self._wall_recv_ts, } else: + # fwdllm eval lives on the aggregator; this eval message is only a + # utility report, not a separately-clocked commit. train_with_data_id + # ran just before in the same loop iteration, so its fresh + # _sim_completion_ts is a valid (not past-dated) sct to echo. msg = { MessageType.MODEL_VERSION: self._model_version, MessageType.STAT_UTILITY: self._stat_utility, + MessageType.SIM_COMPLETION_TS: self._sim_completion_ts, + MessageType.WALL_SEND_TS: time.time(), + MessageType.WALL_RECV_TS: self._wall_recv_ts, } channel.send(end, msg) @@ -755,7 +849,12 @@ def reset_stat_utility(self) -> None: @timer_decorator def pause_execution(self): - time.sleep(1) + # Per-round MQTT throttle chained at the tail of the trainer loop. A + # real-transport artifact with no sim analog (#8): the sim's inter-round + # barrier is the blocking recv in _fetch_weights + the sct reorder buffer, + # so charging 1 wall-s/round to the sim is pure slowdown. Gate off in sim. + if not getattr(self, "simulated", False): + time.sleep(1) return def compose(self) -> None: diff --git a/lib/python/flame/mode/message.py b/lib/python/flame/mode/message.py index 89d1ac8bf..d3a1aaf2f 100644 --- a/lib/python/flame/mode/message.py +++ b/lib/python/flame/mode/message.py @@ -96,3 +96,11 @@ class MessageType(Enum): # skew, same class of bug as B2.0.3). Sim mode never sends this -- its # shared origin is the single _vclock, already exposed via SIM_SEND_TS. AGG_START_TS = 42 + + # Modeled per-trainer mobile delay D = training_delay_s / factor / speedup, + # stamped by the trainer in BOTH modes. Unlike SIM_COMPLETION_TS / + # TRAINING_BUDGET_S (which fold in measured GPU time and differ run-to-run), + # D is deterministic from the registry, so the aggregator orders a cohort's + # commits by (D, trainer_id) identically in real and sim. None when delays + # are disabled (aggregator then falls back to arrival order). + MODELED_DELAY_S = 43 diff --git a/lib/python/flame/monitor/runtime.py b/lib/python/flame/monitor/runtime.py index 714079989..dc0452edb 100644 --- a/lib/python/flame/monitor/runtime.py +++ b/lib/python/flame/monitor/runtime.py @@ -40,6 +40,23 @@ def wrapper(*args, **kwargs): f"[decorator] Runtime of {func.__name__}: {duration:.6f}s " f"(Round={stage.round_id}, DataId={stage.data_id}, Iter={stage.iteration}, TrainerId={stage.trainer_id})" ) + # Structured companion to the log line: attribute step wall time to + # (func, data_id, iteration) for GPU-cost decomposition. No-op when + # telemetry is off. The `stage` guard keeps nested helpers whose + # args[0] is not the trainer self from emitting mis-attributed + # records. Best-effort: a telemetry hiccup must never break training. + try: + from flame import telemetry + if telemetry.is_enabled(): + from flame.telemetry.events import build_step_timing + ev, fields = build_step_timing( + func=func.__name__, duration_s=duration, + round_num=stage.round_id, data_id=stage.data_id, + iteration=stage.iteration, trainer_id=stage.trainer_id, + ) + telemetry.emit(ev, **fields) + except Exception: # pragma: no cover - telemetry must never fault training + logger.debug("step_timing telemetry emit failed", exc_info=True) else: logger.info( f"[decorator] Runtime of {func.__name__}: {duration:.6f}s (no stage info)" diff --git a/lib/python/flame/selector/async_oort.py b/lib/python/flame/selector/async_oort.py index c14bbe030..7b7c6f6aa 100644 --- a/lib/python/flame/selector/async_oort.py +++ b/lib/python/flame/selector/async_oort.py @@ -63,6 +63,11 @@ def __init__(self, **kwargs): """Initailize instance.""" super().__init__(**kwargs) + # #1c: the abandon-timeout clock — set per-select() from + # channel_props["vclock_now"] (sim) or left None (real -> wall). See + # _abandon_clock_now. + self._sim_now_s = None + ml_framework_in_use = get_ml_framework_in_use() if ml_framework_in_use != MLFramework.PYTORCH: raise NotImplementedError( @@ -323,6 +328,14 @@ def select( if self.requester not in self.selected_ends: self.selected_ends[self.requester] = set() + # #1c: the in-flight abandon-timeout (SEND_TIMEOUT_WAIT_S) must run on the + # same clock the trainer commits on -- the virtual clock in sim + # (vclock_now via channel_props), physical wall in real. In a slow sim a + # wall-keyed timeout evicts a still-outstanding trainer -> re-dispatch -> + # R1 residence violation. Stash it so the dispatch STAMP and the CHECK + # (_handle_send_state) use it consistently; None in real -> time.time(). + self._sim_now_s = channel_props.get("vclock_now") + # TODO: (DG) Is explicit round tracking required here? round = # channel_props["round"] if "round" in channel_props else 0 # logger.debug(f"let's select {num_of_ends} ends for new round @@ -1392,6 +1405,15 @@ def _select_candidates_prioritiseUnavail( return candidates, exploit_end_ids + def _abandon_clock_now(self) -> float: + """Clock for the in-flight abandon-timeout (SEND_TIMEOUT_WAIT_S): the + virtual clock in sim (vclock_now stashed per-select), physical wall in + real. Keeping the STAMP (all_selected[end]) and the CHECK on the same + clock makes the timeout mean virtual seconds in sim, so a slow sim no + longer evicts a still-outstanding trainer from the re-pick guard (#1c).""" + sim_now = getattr(self, "_sim_now_s", None) + return sim_now if sim_now is not None else time.time() + def _handle_send_state( self, ends: dict[str, End], @@ -1425,7 +1447,7 @@ def _handle_send_state( # deadlock this caused). curr_all_selected_ends = list(self.all_selected.keys()) for end in curr_all_selected_ends: - current_time_s = time.time() + current_time_s = self._abandon_clock_now() # vclock in sim, wall in real (#1c) if end in self.all_selected.keys(): # Check again to avoid possible case of race condition # when all_selected has been updated from another @@ -1572,19 +1594,30 @@ def _handle_send_state( count_avl_eval = 0 count_ineligible = 0 + # SIM R1 guard: async_oort releases `all_selected` on physical events + # (recv-fifo re-select loop, RECVD/NONE cleanup). In a slow sim a grad + # stays returned-but-uncommitted for a long virtual window while the + # aggregator still models the trainer as in flight; a physical prune then + # frees a still-outstanding trainer -> select() re-dispatches it -> R1 + # violation. Also exclude the aggregator's virtual in-flight set (bound + # via `_agg_pending_commit_ref`) so a trainer is un-re-pickable until its + # grad COMMITS. Empty (default) in real -> unchanged. + _pending = getattr(self, "_agg_pending_commit_ref", None) or set() + # Check the eligible set first. Out of the ends, how many are # not in all_selected? Only those are eligible since the rest # have weights already sent to them for either train/eval # task. count_eligible_set_to_check = [ - end for end in ends if end not in self.all_selected + end for end in ends + if end not in self.all_selected and end not in _pending ] logger.debug( f"Before creating filtered_ends. count_eligible_set_to_check: {len(count_eligible_set_to_check)} from total {len(ends)} ends." ) for end_id in ends: - if end_id not in self.all_selected.keys(): + if end_id not in self.all_selected.keys() and end_id not in _pending: logger.debug( f"Creating filtered ends. Checking end id {end_id}, avl_state = {ends[end_id].get_property(PROP_AVL_STATE)}" ) @@ -2010,7 +2043,7 @@ def _handle_recv_state( for selected_end in selected_ends: # Add to all_selected. {key: end, val: TS epoch (s)} - self.all_selected[selected_end] = time.time() + self.all_selected[selected_end] = self._abandon_clock_now() # #1c logging.debug( f"self.all_selected {self.all_selected} after combining with " f"selected_ends {selected_ends}" @@ -2084,7 +2117,7 @@ def process_chosen_candidate_dict( for candidate_end in candidates: # Add to all_selected. {key: end, val: TS epoch (s)} - self.all_selected[candidate_end] = time.time() + self.all_selected[candidate_end] = self._abandon_clock_now() # #1c logging.debug( f"self.all_selected {self.all_selected} after combining" f" with candidates {candidates}" diff --git a/lib/python/flame/telemetry/events.py b/lib/python/flame/telemetry/events.py index 237f09d8c..d666c49e5 100644 --- a/lib/python/flame/telemetry/events.py +++ b/lib/python/flame/telemetry/events.py @@ -29,6 +29,8 @@ EVENT_WITHHELD_DELIVERY = "withheld_delivery" # late stale delivery of a send-gated update EVENT_ABANDON_TIMEOUT = "abandon_timeout" # 90s vclock slot-free of a stalled trainer EVENT_AGG_BELIEF_CHANGE = "agg_belief_change" # aggregator's belief about a trainer's avail state +EVENT_STEP_TIMING = "step_timing" # per-function wall duration of a timed compute step +EVENT_COMM = "comm" # one message put on the wire (byte-size accounting) KNOWN_EVENTS = frozenset( { @@ -47,6 +49,8 @@ EVENT_WITHHELD_DELIVERY, EVENT_ABANDON_TIMEOUT, EVENT_AGG_BELIEF_CHANGE, + EVENT_STEP_TIMING, + EVENT_COMM, } ) @@ -188,6 +192,73 @@ def build_trainer_round( return EVENT_TRAINER_ROUND, fields +def build_step_timing( + *, + func: str, + duration_s: float, + round_num: Optional[int] = None, + data_id: Optional[int] = None, + iteration: Optional[int] = None, + trainer_id: Optional[str] = None, +) -> tuple[str, dict[str, Any]]: + """Per-function wall duration of one timed compute step (`timer_decorator`). + + Fine-grained companion to `trainer_round`'s coarse phase split: attributes + wall time to individual forward-grad steps (functional-model setup, + perturbation selection, per-batch JVP, delay emulation) for GPU-cost + decomposition. Keyed by (data_id, iteration) to track cost across the cadence. + """ + fields: dict[str, Any] = {"func": func, "duration_s": duration_s} + for k, v in ( + ("round", round_num), + ("data_id", data_id), + ("iteration_per_data_id", iteration), + ("trainer_id", trainer_id), + ): + if v is not None: + fields[k] = v + return EVENT_STEP_TIMING, fields + + +def build_comm( + *, + direction: str, + size_bytes: int, + peer_id: Optional[str] = None, + round_num: Optional[int] = None, + data_id: Optional[int] = None, + iteration: Optional[int] = None, + payload_kind: Optional[str] = None, + n_tensors: Optional[int] = None, + trainer_id: Optional[str] = None, +) -> tuple[str, dict[str, Any]]: + """One message placed on the wire, for network-cost accounting (Experiment 4). + + Emitted by BOTH roles so total bytes / message counts / per-message size + distributions are comparable across baselines (fluxtune sends perturbation + seeds/scalars, not full gradients, so real wire size differs from the static + model_param_count reconstruction). + + direction: "agg_to_trainer" (dispatch) | "trainer_to_agg" (update upload). + peer_id: the other end (may be None trainer-side). payload_kind: "weights" / + "var_bad" / "gradients", to split dispatch vs update and full-weight vs + var-signal. size_bytes: serialized message size. + """ + fields: dict[str, Any] = {"direction": direction, "size_bytes": int(size_bytes)} + for k, v in ( + ("peer_id", peer_id), + ("round", round_num), + ("data_id", data_id), + ("iteration_per_data_id", iteration), + ("payload_kind", payload_kind), + ("n_tensors", n_tensors), + ("trainer_id", trainer_id), + ): + if v is not None: + fields[k] = v + return EVENT_COMM, fields + + def build_util_disparity( *, round_num: int, diff --git a/lib/python/tests/launch/test_config_generator.py b/lib/python/tests/launch/test_config_generator.py index 8686fc595..21c4d50c1 100644 --- a/lib/python/tests/launch/test_config_generator.py +++ b/lib/python/tests/launch/test_config_generator.py @@ -412,3 +412,39 @@ def test_generate_trainer_config_differs_across_trainers(self, gen): cfg_054["hyperparameters"]["avl_events_syn_20"] != cfg_005["hyperparameters"]["avl_events_syn_20"] ) + + +class TestExecutionConfigBanksDelays: + """#12: the effective training-delay config is banked in the execution config + so a run can be audited post-hoc (previously absent -> a `--delays on` run + "looked" off because nothing recorded it).""" + + def _exp(self, enable, factor=None): + from flame.launch.experiment_config import ExperimentConfig, TrainerConfig + tr_hp = {"training_delay_factor": factor} if factor is not None else None + return ExperimentConfig( + name="bank_test", + trainer=TrainerConfig( + enable_training_delays=enable, hyperparameters=tr_hp + ), + ) + + def test_banks_enable_and_factor(self, tmp_path): + from flame.launch.execution_config_generator import create_execution_config + cfg = create_execution_config( + self._exp(enable=True, factor=3), + aggregator_config_path=tmp_path / "agg.json", + ) + tr = cfg["experiment"]["trainer"] + assert tr["enable_training_delays"] is True + assert tr["training_delay_factor"] == 3 + + def test_banks_disabled(self, tmp_path): + from flame.launch.execution_config_generator import create_execution_config + cfg = create_execution_config( + self._exp(enable=False), + aggregator_config_path=tmp_path / "agg.json", + ) + assert cfg["experiment"]["trainer"]["enable_training_delays"] is False + # factor unset -> banked as None (trainer_base default applies) + assert cfg["experiment"]["trainer"]["training_delay_factor"] is None diff --git a/lib/python/tests/launch/test_runner_paths.py b/lib/python/tests/launch/test_runner_paths.py index 012ace1c5..f07a71054 100644 --- a/lib/python/tests/launch/test_runner_paths.py +++ b/lib/python/tests/launch/test_runner_paths.py @@ -11,6 +11,7 @@ ExampleConfig, ExperimentConfig, MetadataPaths, + TrainerConfig, ) from flame.launch.runner import ExperimentRunner @@ -24,6 +25,50 @@ def fake_example_dir(tmp_path): return ex +class TestTrainingDelayFan: + """#12: the trainer's enable_training_delays is a single source of truth + fanned into the AGGREGATOR config too, so both roles agree (previously the + aggregator defaulted to False -- a real<->sim desync risk).""" + + def _agg_hp(self, fake_example_dir, enable, tr_hp=None, agg_overrides=None): + runner = ExperimentRunner(fake_example_dir) + exp = ExperimentConfig( + name="x", + trainer=TrainerConfig( + enable_training_delays=enable, hyperparameters=tr_hp + ), + aggregator=AggregatorConfig(config_overrides=agg_overrides), + ) + merged, _ = runner._build_aggregator_config(exp, None, None) + return merged.get("hyperparameters", {}) + + def test_enabled_fans_to_aggregator(self, fake_example_dir): + hp = self._agg_hp(fake_example_dir, enable=True) + assert hp["trainingDelayEnabled"] is True + + def test_disabled_fans_to_aggregator(self, fake_example_dir): + # The key case: without the fan the aggregator would default False + # regardless of the request; with it, it reflects the actual value. + hp = self._agg_hp(fake_example_dir, enable=False) + assert hp["trainingDelayEnabled"] is False + + def test_factor_fans_when_set(self, fake_example_dir): + hp = self._agg_hp( + fake_example_dir, enable=True, + tr_hp={"training_delay_factor": 3}, + ) + assert hp["trainingDelayFactor"] == 3 + + def test_fan_wins_over_conflicting_config_override(self, fake_example_dir): + # The fan is the final, highest-precedence layer: a stale aggregator + # config_override cannot shadow the run-wide request. + hp = self._agg_hp( + fake_example_dir, enable=True, + agg_overrides={"hyperparameters": {"trainingDelayEnabled": False}}, + ) + assert hp["trainingDelayEnabled"] is True + + class TestPathResolution: def test_defaults(self, fake_example_dir): runner = ExperimentRunner(fake_example_dir) diff --git a/lib/python/tests/mode/test_async_sim_ordering.py b/lib/python/tests/mode/test_async_sim_ordering.py index 8d0d33654..edbc99d62 100644 --- a/lib/python/tests/mode/test_async_sim_ordering.py +++ b/lib/python/tests/mode/test_async_sim_ordering.py @@ -427,6 +427,26 @@ def test_fresh_end_not_reclaimed_still_short_circuits(self): assert result == {} assert not pacer_called + + def test_sim_vclock_keeps_virtually_recent_end_despite_old_wall(self): + """#1c: in sim the abandon-timeout runs on the VCLOCK, not physical wall. + A trainer stamped at a small vclock (recent in VIRTUAL time) is NOT + reclaimed even though a slow sim has burned >90 physical wall-seconds -- + which was evicting still-outstanding trainers from all_selected and + causing R1 re-dispatch. Contrast test_stale_end_* (wall) which reclaims.""" + sel = self._stub_selector(stale=True) # wall-old stamp from the helper... + # ...but drive the timeout on the VIRTUAL clock instead: stamp = vclock 5 + # at dispatch, sim-now = vclock 10 -> delta 5 < SEND_TIMEOUT_WAIT_S (90). + sel.all_selected = {self.STALE_END: 5.0} + sel._sim_now_s = 10.0 # would come from channel_props["vclock_now"] + pacer_called = [] + sel.pacer = lambda: pacer_called.append(True) + + result = self._call(sel, concurrency=1) + + # Guard held: still in all_selected, concurrency saturated -> short-circuit. + assert self.STALE_END in sel.all_selected + assert result == {} and not pacer_called assert self.STALE_END in sel.all_selected assert self.STALE_END in sel.selected_ends["agg"] diff --git a/lib/python/tests/mode/test_fwdllm_agg_telemetry.py b/lib/python/tests/mode/test_fwdllm_agg_telemetry.py index c5300be37..786ed7a5b 100644 --- a/lib/python/tests/mode/test_fwdllm_agg_telemetry.py +++ b/lib/python/tests/mode/test_fwdllm_agg_telemetry.py @@ -13,6 +13,7 @@ it isn't. """ +import time from datetime import timedelta import torch @@ -71,6 +72,8 @@ class _FakeAggregator: def __init__(self, contributors, var_good_enough, staleness_map=None, total_data_bins=150): + # Real path skips the sim boundary hook, so telemetry is byte-identical. + self.simulated = False self._per_agg_trainer_list = list(contributors) self._agg_goal_cnt = len(contributors) self._agg_goal = len(contributors) or 1 @@ -199,6 +202,82 @@ def test_emitted_every_cycle_regardless_of_variance_outcome(self, tmp_path): finally: telemetry.shutdown() + def test_speedup_fields_emitted(self, tmp_path): + """#13: agg_round carries wall_elapsed_s in both modes and, in sim, + sim_rate = vclock/wall so the slowdown is observable in telemetry.""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=False) + # Put the aggregator on the sim path with a virtual clock ahead of + # a known wall span. + agg.simulated = True + agg._vclock = type("V", (), {"now": 120.0})() + agg.agg_start_time_ts = time.time() - 60.0 # ~60 wall-s elapsed + # The sim boundary hook is exercised elsewhere; no-op it here so this + # test isolates the speedup-telemetry emission. + agg._release_sim_slots_at_agg_goal = lambda *a, **k: None + channel = _FakeChannel( + durations={"t1": timedelta(seconds=1)}, utilities={"t1": 0.1} + ) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert r["wall_elapsed_s"] > 0 + # vclock 120 over ~60 wall-s -> sim_rate ~2 (a speedup); must be present + assert r["sim_rate"] is not None and r["sim_rate"] > 1.0 + finally: + telemetry.shutdown() + + def test_intrinsic_span_is_barrier_plus_eval_excludes_fedavg(self, tmp_path): + """#6 anchor: intrinsic_span_s = the barrier (MAX committed-cohort + duration) + eval_s (commit only), EXCLUDING the FedAvg merge -- it must + mirror the sim vclock's composition (barrier sct + eval fold) so the + clock-rate rungs anchor REAL like-for-like. A variance-FAIL cycle runs no + eval, so intrinsic collapses to exactly the barrier -- which also proves + max() (not min/sum) and that fedavg is not folded in.""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1", "t2"], + var_good_enough=False) + channel = _FakeChannel(durations={"t1": timedelta(seconds=5), + "t2": timedelta(seconds=3)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + # barrier = max(5, 3) = 5; no eval on the fail path; fedavg (stubbed + # ~0) is excluded regardless -> intrinsic == the barrier. + assert r["intrinsic_span_s"] is not None + assert abs(r["intrinsic_span_s"] - 5.0) < 0.5, r + finally: + telemetry.shutdown() + + def test_wall_elapsed_emitted_in_real_mode_sim_rate_none(self, tmp_path): + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=False) + # default fake is real mode (simulated=False) + agg.agg_start_time_ts = time.time() - 5.0 + channel = _FakeChannel( + durations={"t1": timedelta(seconds=1)}, utilities={"t1": 0.1} + ) + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert r["wall_elapsed_s"] > 0 + assert r["sim_rate"] is None # real mode has no virtual clock rate + finally: + telemetry.shutdown() + def test_staleness_computed_against_pre_cycle_model_version(self, tmp_path): """staleness = the model version this cycle aggregated against, minus each contributor's last-known trained-on version -- captured BEFORE @@ -223,6 +302,50 @@ def test_staleness_computed_against_pre_cycle_model_version(self, tmp_path): finally: telemetry.shutdown() + def test_cadence_fields_snapshot_pre_mutation(self, tmp_path): + """Variance-cadence inputs: cycle_data_id/cycle_iteration identify the + data_id this cycle WORKED on (pre-advance), and the pool sizes are captured + at the variance gate. On a variance FAIL data_id does not advance, so + cycle_data_id == the emitted (post) data_id == 3.""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=False) + channel = _FakeChannel(durations={"t1": timedelta(seconds=2)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert r["cycle_data_id"] == 3 and r["cycle_iteration"] == 0 + # grad_pool got this cycle's grad appended before the snapshot; the + # fake sets no cached_v, so cached_v_size defaults to 0. + assert r["grad_pool_size"] == 1 and r["cached_v_size"] == 0 + finally: + telemetry.shutdown() + + def test_cycle_data_id_is_pre_advance_on_commit(self, tmp_path): + """On a variance PASS the emitted (post) data_id advances to 4, but + cycle_data_id stays 3 -- the data_id this cycle committed. This is the + off-by-one V1 relies on: bin cadence cycles by cycle_data_id, not the + post-mutation data_id (which would attribute a commit to the next bin).""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=True) + channel = _FakeChannel(durations={"t1": timedelta(seconds=2)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert r["cycle_data_id"] == 3 # worked-on data_id + assert r["data_id"] == 4 # post-commit advance + finally: + telemetry.shutdown() + def test_contributor_list_captured_before_reset(self, tmp_path): """_per_agg_trainer_list is cleared at the end of this method -- agg_round's contributing_trainers must reflect this cycle's @@ -275,6 +398,146 @@ def test_agg_observed_s_keyed_by_end_id(self, tmp_path): telemetry.shutdown() +class TestContributorIntervalsEmission: + """R1/W1 residence rungs read a per-contributor [dispatch, commit] interval + list off each agg_round event. It must land once per contributor, carrying + the ts captured in _process_single_trainer_message.""" + + def test_intervals_emitted_per_contributor(self, tmp_path): + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1", "t2"], var_good_enough=False) + # Simulate the per-contribution capture done in the message handler. + agg._sim_contrib_intervals = { + "t1": {"dispatch_ts": 1.0, "commit_ts": 6.0}, + "t2": {"dispatch_ts": 2.0, "commit_ts": 9.0}, + } + channel = _FakeChannel( + durations={"t1": timedelta(seconds=5), "t2": timedelta(seconds=7)}, + ) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + ci = {d["end"]: d for d in r["contributor_intervals"]} + assert set(ci) == {"t1", "t2"} + assert ci["t1"]["dispatch_ts"] == 1.0 and ci["t1"]["commit_ts"] == 6.0 + assert ci["t2"]["dispatch_ts"] == 2.0 and ci["t2"]["commit_ts"] == 9.0 + finally: + telemetry.shutdown() + + def test_intervals_present_even_without_captured_ts(self, tmp_path): + """When no interval was captured (e.g. a test double / real run with the + dict unpopulated) the field is still emitted with null ts, so the rung + SKIPs cleanly rather than the field being absent.""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=False) + channel = _FakeChannel(durations={"t1": timedelta(seconds=5)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert r["contributor_intervals"] == [ + {"end": "t1", "dispatch_ts": None, "commit_ts": None}] + finally: + telemetry.shutdown() + + +class TestPerRoundWallDecomposition: + """agg_round carries the per-round wall breakdown feeding #6 -- + aggregate_fedavg_s + eval_s always; barrier_wait_s/drain_tail_s when the + dispatch/last-grad wall stamps were captured (else null, rung SKIPs).""" + + def test_fedavg_and_eval_present_on_pass(self, tmp_path): + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=True) + channel = _FakeChannel(durations={"t1": timedelta(seconds=2)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + # fedavg wall is measured around aggregate() and is non-negative + assert r["aggregate_fedavg_s"] is not None + assert r["aggregate_fedavg_s"] >= 0.0 + # eval ran (variance passed) -> eval_s measured, non-negative + assert r["eval_s"] is not None and r["eval_s"] >= 0.0 + finally: + telemetry.shutdown() + + def test_eval_s_null_on_variance_fail(self, tmp_path): + """eval_model only runs on the pass path, so eval_s is null on a FAIL.""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=False) + channel = _FakeChannel(durations={"t1": timedelta(seconds=2)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert r["eval_s"] is None + assert r["aggregate_fedavg_s"] is not None # aggregate always runs + finally: + telemetry.shutdown() + + def test_barrier_and_drain_from_wall_stamps(self, tmp_path): + """When the dispatch + last-grad wall stamps exist, barrier_wait_s = + last_grad - dispatch and drain_tail_s = commit - last_grad.""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=False) + # Stamps the barrier collection would have set (dispatch then last + # grad, both in the past relative to the commit inside the method). + import time as _t + now = _t.time() + agg._round_dispatch_wall_ts = now - 5.0 + agg._last_grad_wall_ts = now - 2.0 + channel = _FakeChannel(durations={"t1": timedelta(seconds=2)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert abs(r["barrier_wait_s"] - 3.0) < 0.5 # (now-2) - (now-5) + assert r["drain_tail_s"] >= 0.0 # commit is after last grad + finally: + telemetry.shutdown() + + def test_barrier_drain_null_without_stamps(self, tmp_path): + """No wall stamps captured (test double / async path) -> null, not a + crash; the rung SKIPs rather than the field being absent.""" + telemetry.configure(role="aggregator", run_dir=str(tmp_path)) + try: + agg = _FakeAggregator(contributors=["t1"], var_good_enough=False) + channel = _FakeChannel(durations={"t1": timedelta(seconds=2)}) + + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + + import json + events = [json.loads(l) for l in + (tmp_path / "aggregator.jsonl").read_text().splitlines()] + r = [e for e in events if e["event"] == "agg_round"][0] + assert r["barrier_wait_s"] is None + assert r["drain_tail_s"] is None + finally: + telemetry.shutdown() + + class _UtilityFakeChannel: """Generic fake for _process_single_trainer_message's channel calls -- stores per-end properties in a dict, doesn't care about specific PROP_* @@ -317,6 +580,7 @@ class _UtilityFakeAggregator: trainers is a no-op) since it's irrelevant to the telemetry under test.""" process = TopAggregator._process_single_trainer_message + _release_end_on_return = TopAggregator._release_end_on_return def __init__(self, model_version=5, data_id=3, iteration_per_data_id=0, is_async=False): @@ -441,6 +705,42 @@ def test_noop_when_telemetry_disabled(self, tmp_path): assert not (tmp_path / "aggregator.jsonl").exists() +class TestStalenessPolicy: + """staleness_policy gate in _process_single_trainer_message. + + REJECT policies (round_data_id/exact) drop a stale grad; ACCEPT policies + (none/fedbuff) consume it -- fedbuff is the async baseline default so its + carried surplus grads (commit-then-carry) are accepted + down-weighted by + (V'-V) in aggregate_grads_from_trainers, never silently dropped.""" + + def _run(self, policy, msg_version, agg_version=5): + agg = _UtilityFakeAggregator(model_version=agg_version, is_async=True) + agg.staleness_policy = policy + channel = _UtilityFakeChannel() + agg.process(channel, _msg(model_version=msg_version, stat_utility=0.5), + "t1", timestamp=0) + return agg + + def test_fedbuff_accepts_stale_grad(self): + # msg trained on v3, agg now at v5 -> stale by 2, but fedbuff ACCEPTS it. + agg = self._run("fedbuff", msg_version=3) + assert agg._agg_goal_cnt == 1 # grad consumed + assert agg._per_agg_trainer_list == ["t1"] + + def test_none_accepts_stale_grad(self): + agg = self._run("none", msg_version=3) + assert agg._agg_goal_cnt == 1 + + def test_round_data_id_rejects_stale_grad(self): + agg = self._run("round_data_id", msg_version=3) + assert agg._agg_goal_cnt == 0 # dropped + assert agg._per_agg_trainer_list == [] + + def test_fedbuff_accepts_fresh_grad(self): + agg = self._run("fedbuff", msg_version=5) # not stale + assert agg._agg_goal_cnt == 1 + + class TestRoundCacheActivityResetOnContribution: """A real accepted contribution must restart the round-cache stuck-timeout clock (see TestStuckCachePruning in diff --git a/lib/python/tests/mode/test_fwdllm_commit_canon.py b/lib/python/tests/mode/test_fwdllm_commit_canon.py new file mode 100644 index 000000000..0ef198049 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_commit_canon.py @@ -0,0 +1,137 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Canonical (D, trainer_id) commit ordering. + +When two trainers share a registry delay (equal D), real breaks the sct tie by +physical arrival and sim by sct-sort, so the receive ORDER can swap even though +`var` and the cohort SET match -- and the EXACT-order `cohort_sequence` rung +flags it. `_canonicalize_cohort_commit_order` reorders THIS cycle's cohort by +(D, str(end)) so equal-D ties break by trainer_id IDENTICALLY in real and sim. +These tests pin: (1) two different input orders (real-physical vs sim-sct) +canonicalize to the SAME sequence; (2) the grad/jvp trailing slice reorders in +lockstep; (3) accumulated earlier-iteration entries are untouched; (4) no-op +when keys are missing (delays off) or the cohort is already canonical. +""" + +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator + + +class _CanonAgg: + """Minimal stand-in exposing only what _canonicalize_cohort_commit_order + touches. `end` values are the last-3-hex trainer tokens for readability.""" + + def __init__(self, ends, delays, grad_list=None, jvp_list=None): + self._per_agg_trainer_list = list(ends) + # end -> (D, str(end)) key, mirroring aggregate_weights' capture; a None + # delay models a contributor that did not stamp D (delays off). + self._commit_key_by_end = { + e: ((float(d), str(e)) if d is not None else None) + for e, d in zip(ends, delays) + } + self.grad_for_var_check_list = list( + grad_list if grad_list is not None else ends + ) + self.jvp_for_snr_check_list = list( + jvp_list if jvp_list is not None else ends + ) + + canon = TopAggregator._canonicalize_cohort_commit_order + + +# D-values keyed by the last-3-hex token; trainers 372 & 378 both drew +# training_delay_s=13.0 -> D=6.5, the tie. +_D = { + "370": 2.0, "375": 2.5, "373": 3.5, "376": 5.0, "379": 5.5, + "378": 6.5, "372": 6.5, "371": 8.0, "377": 8.5, "374": 9.0, +} +_CANON_ORDER = ["370", "375", "373", "376", "379", "372", "378", "371", "377", "374"] + + +def _mk(ends): + return _CanonAgg(ends, [_D[e] for e in ends]) + + +class TestCanonicalizesToOneOrder: + def test_real_physical_and_sim_sct_orders_converge(self): + # real broke the 6.5 tie 378-before-372; sim broke it 372-before-378. + real_order = ["370", "375", "373", "376", "379", "378", "372", "371", "377", "374"] + sim_order = ["370", "375", "373", "376", "379", "372", "378", "371", "377", "374"] + ra, sa = _mk(real_order), _mk(sim_order) + ra.canon() + sa.canon() + assert ra._per_agg_trainer_list == _CANON_ORDER + assert sa._per_agg_trainer_list == _CANON_ORDER + assert ra._per_agg_trainer_list == sa._per_agg_trainer_list + + def test_tie_breaks_by_trainer_id_not_arrival(self): + # only the two tied (D=6.5) members may move, and by id order (372<378) + a = _mk(["379", "378", "372", "371"]) + a.canon() + assert a._per_agg_trainer_list == ["379", "372", "378", "371"] + + def test_grad_and_jvp_reorder_in_lockstep(self): + a = _mk(["379", "378", "372", "371"]) + a.canon() + assert a.grad_for_var_check_list == ["379", "372", "378", "371"] + assert a.jvp_for_snr_check_list == ["379", "372", "378", "371"] + + +class TestSliceScope: + def test_only_trailing_cohort_slice_reorders(self): + # grad list ACCUMULATES: an earlier iteration's 2 entries precede this + # cycle's 4-cohort. Only the trailing 4 may move. + prev = ["p0", "p1"] + cohort = ["379", "378", "372", "371"] + a = _CanonAgg(cohort, [_D[e] for e in cohort], grad_list=prev + cohort, + jvp_list=prev + cohort) + a.canon() + assert a.grad_for_var_check_list == prev + ["379", "372", "378", "371"] + assert a._per_agg_trainer_list == ["379", "372", "378", "371"] + + def test_misaligned_grad_list_left_untouched(self): + # a shorter-than-cohort grad list (a non-grad message slipped in) must + # not be sliced/corrupted -- reorder the contributor list only. + cohort = ["379", "378", "372", "371"] + a = _CanonAgg(cohort, [_D[e] for e in cohort], grad_list=["x"], + jvp_list=["x"]) + a.canon() + assert a._per_agg_trainer_list == ["379", "372", "378", "371"] + assert a.grad_for_var_check_list == ["x"] # untouched + + +class TestNoOps: + def test_already_canonical_is_unchanged(self): + a = _mk(_CANON_ORDER) + before = list(a._per_agg_trainer_list) + a.canon() + assert a._per_agg_trainer_list == before + + def test_missing_delay_key_falls_back_to_arrival_order(self): + # a contributor without a stamped D (delays off) -> arrival order kept + ends = ["379", "378", "372", "371"] + a = _CanonAgg(ends, [5.5, None, 6.5, 8.0]) + a.canon() + assert a._per_agg_trainer_list == ends # no reorder + + def test_single_contributor_is_noop(self): + a = _mk(["370"]) + a.canon() + assert a._per_agg_trainer_list == ["370"] + + +class TestVarInvariance: + def test_tie_swap_stays_within_split_half(self): + # the tie (372/378) sits at indices 5-6 of the 10-cohort; the split-half + # boundary is n//2 = 5, so both are in the SECOND half in BOTH orders -> + # the reorder cannot move a grad across the boundary -> calculate_var + # (mean of first-half vs second-half) is invariant. + n = len(_CANON_ORDER) + real_order = ["370", "375", "373", "376", "379", "378", "372", "371", "377", "374"] + first_ids = lambda seq: set(seq[: n // 2]) + second_ids = lambda seq: set(seq[n // 2:]) + # half-MEMBERSHIP is identical before and after canonicalization + a = _mk(real_order) + pre_first, pre_second = first_ids(real_order), second_ids(real_order) + a.canon() + assert first_ids(a._per_agg_trainer_list) == pre_first + assert second_ids(a._per_agg_trainer_list) == pre_second diff --git a/lib/python/tests/mode/test_fwdllm_early_stop_conditions.py b/lib/python/tests/mode/test_fwdllm_early_stop_conditions.py index 460dcd6b3..bbea82a6e 100644 --- a/lib/python/tests/mode/test_fwdllm_early_stop_conditions.py +++ b/lib/python/tests/mode/test_fwdllm_early_stop_conditions.py @@ -30,10 +30,15 @@ def __init__( max_runtime_s=None, agg_start_time_ts=None, is_async=False, + simulated=False, + vclock_now=0.0, + sim_wall_ceiling_s=None, ): self._work_done = False self.data_id = data_id self.is_async = is_async + self.simulated = simulated + self._vclock = SimpleNamespace(now=vclock_now) self.agg_start_time_ts = ( agg_start_time_ts if agg_start_time_ts is not None else time.time() ) @@ -41,6 +46,7 @@ def __init__( hyperparameters=SimpleNamespace( max_data_id_progress=max_data_id_progress, max_runtime_s=max_runtime_s, + sim_wall_ceiling_s=sim_wall_ceiling_s, ) ) @@ -50,6 +56,8 @@ def _distribute_weights_sync(self, tag, task_to_perform="train"): def _distribute_weights_async(self, tag, task_to_perform="train"): pass + # Phase 4a: the ceiling default reads self.SIM_WALL_CEILING_FACTOR. + SIM_WALL_CEILING_FACTOR = TopAggregator.SIM_WALL_CEILING_FACTOR check = TopAggregator._check_early_stop_conditions _check_early_stop_conditions = TopAggregator._check_early_stop_conditions distribute = TopAggregator._distribute_weights @@ -105,6 +113,69 @@ def test_whichever_fires_first_runtime_before_data_id(self): agg.check() assert agg._work_done is True + def test_sim_max_runtime_uses_vclock_not_wall(self): + """SIM mode (root #6 methodology): max_runtime_s is a VIRTUAL-clock + budget, not wall. vclock below budget => keep running even though wall + has long passed it (the sim runs faster than real).""" + agg = _FakeAggregator( + max_runtime_s=3600.0, simulated=True, vclock_now=100.0, + agg_start_time_ts=time.time(), # ~0 wall elapsed + ) + agg.check() + assert agg._work_done is False # vclock 100 < budget 3600 + + def test_sim_max_runtime_stops_when_vclock_reaches_budget(self): + agg = _FakeAggregator( + max_runtime_s=3600.0, simulated=True, vclock_now=3600.0, + agg_start_time_ts=time.time(), + ) + agg.check() + assert agg._work_done is True + + def test_sim_wall_below_decoupled_ceiling_keeps_running(self): + """Phase 4a (root S1): the default ceiling is now a GENEROUS multiple of + the budget (SIM_WALL_CEILING_FACTOR), NOT 1×. A sim whose WALL exceeds + the budget but is well under budget×factor must KEEP running -- the sim + legitimately uses more wall than vclock (real GPU + eval). This is the + S1 truncation the old 1× ceiling caused.""" + agg = _FakeAggregator( + max_runtime_s=3600.0, simulated=True, vclock_now=80.0, # under budget + agg_start_time_ts=time.time() - 7200.0, # wall 2× budget + ) + agg.check() + assert agg._work_done is False # 7200s < 20× × 3600s ceiling + + def test_sim_wall_ceiling_stops_true_runaway(self): + """The outer safety still fires on a genuine runaway: wall beyond + budget × SIM_WALL_CEILING_FACTOR stops the sim.""" + over = 3600.0 * TopAggregator.SIM_WALL_CEILING_FACTOR + 10.0 + agg = _FakeAggregator( + max_runtime_s=3600.0, simulated=True, vclock_now=80.0, + agg_start_time_ts=time.time() - over, + ) + agg.check() + assert agg._work_done is True + + def test_explicit_sim_wall_ceiling_overrides_default(self): + """An explicit `sim_wall_ceiling_s` is honored verbatim (tighter outer + bound) -- not multiplied by the factor.""" + agg = _FakeAggregator( + max_runtime_s=3600.0, simulated=True, vclock_now=80.0, + sim_wall_ceiling_s=100.0, + agg_start_time_ts=time.time() - 101.0, # wall > explicit ceiling + ) + agg.check() + assert agg._work_done is True + + def test_real_max_runtime_still_uses_wall(self): + """Real mode unchanged: wall-clock elapsed drives the cap.""" + agg = _FakeAggregator( + max_runtime_s=1.0, simulated=False, + agg_start_time_ts=time.time() - 2.0, + ) + agg.check() + assert agg._work_done is True + def test_already_work_done_is_a_noop(self): """Once stopped, the check must not re-derive or overwrite state.""" agg = _FakeAggregator(data_id=5, max_data_id_progress=10) diff --git a/lib/python/tests/mode/test_fwdllm_grad_aware_agg.py b/lib/python/tests/mode/test_fwdllm_grad_aware_agg.py new file mode 100644 index 000000000..48437dbbc --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_grad_aware_agg.py @@ -0,0 +1,124 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Opt-3: gradient-aware aggregation. + +The default fedbuff `new` rate rescales an update's MAGNITUDE by staleness×utility +but cannot refuse a wrong DIRECTION, so anti-aligned JVP estimates still get +averaged. The `grad_aware` rate weights by direction (alignment gate) and +optionally reliability (inverse-variance), BOUNDED so the result never exceeds the +base rate -> the effective server LR never inflates, so no LR re-tune is needed +for stability. + +`_grad_aware_rate` (pure scalar) and `_cosine_flat` (tensor cosine vs the running +aggregate) are the testable primitives. These pin: +(1) both mechanisms off -> passthrough (== base rate); +(2) alignment gate: no-op when aligned, linear down-weight when anti-aligned, 0 at + cos=-1, None cos -> no-op, and a non-zero align_floor gates weakly-aligned too; +(3) inverse-variance: no-op at/under the reference, shrinks above it, None -> no-op; +(4) the result is always <= base (never inflates LR); +(5) cosine sign/degenerate behavior. +""" + +import torch + +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator + +rate = TopAggregator._grad_aware_rate +cosf = TopAggregator._cosine_flat + + +# --- (1) both off -> passthrough ------------------------------------------- + +def test_both_off_is_passthrough(): + assert rate(0.7, cos=-1.0, var_i=99.0, var_ref=0.3, + align_gate=False, inverse_var=False) == 0.7 + + +# --- (2) alignment gate (S2) ----------------------------------------------- + +def test_align_noop_when_aligned(): + # cos >= align_floor -> factor 1. + assert rate(0.8, cos=0.5, var_i=None, var_ref=0.3, align_gate=True) == 0.8 + assert rate(0.8, cos=0.0, var_i=None, var_ref=0.3, align_gate=True) == 0.8 + + +def test_align_downweights_anti_aligned(): + # floor 0: factor = (cos+1)/1 for cos<0. cos=-0.5 -> 0.5, cos=-1 -> 0. + assert rate(1.0, cos=-0.5, var_i=None, var_ref=0.3, align_gate=True) == 0.5 + assert rate(1.0, cos=-1.0, var_i=None, var_ref=0.3, align_gate=True) == 0.0 + assert rate(0.6, cos=-0.5, var_i=None, var_ref=0.3, align_gate=True) == 0.3 + + +def test_align_none_cos_is_noop(): + assert rate(0.9, cos=None, var_i=None, var_ref=0.3, align_gate=True) == 0.9 + + +def test_align_floor_gates_weakly_aligned(): + # floor 0.2: cos=0.1 is below the floor -> gated. factor=(0.1+1)/(0.2+1)=1.1/1.2. + r = rate(1.0, cos=0.1, var_i=None, var_ref=0.3, align_gate=True, align_floor=0.2) + assert abs(r - (1.1 / 1.2)) < 1e-9 + # cos above the floor -> untouched. + assert rate(1.0, cos=0.3, var_i=None, var_ref=0.3, align_gate=True, + align_floor=0.2) == 1.0 + + +# --- (3) inverse-variance (S1) --------------------------------------------- + +def test_inverse_var_noop_at_or_under_ref(): + # var_i <= var_ref -> min(1, ref/var) == 1. + assert rate(1.0, cos=None, var_i=0.2, var_ref=0.3, + align_gate=False, inverse_var=True) == 1.0 + + +def test_inverse_var_shrinks_above_ref(): + # var_i=0.6, ref=0.3 -> ~0.5 (eps negligible). + r = rate(1.0, cos=None, var_i=0.6, var_ref=0.3, + align_gate=False, inverse_var=True, var_eps=0.0) + assert abs(r - 0.5) < 1e-9 + + +def test_inverse_var_none_is_noop(): + assert rate(0.7, cos=None, var_i=None, var_ref=0.3, + align_gate=False, inverse_var=True) == 0.7 + + +def test_align_and_inverse_var_compose(): + # base 1.0 * align(cos=-0.5 -> .5) * invvar(var .6/ref .3 -> .5) = .25 + r = rate(1.0, cos=-0.5, var_i=0.6, var_ref=0.3, + align_gate=True, inverse_var=True, var_eps=0.0) + assert abs(r - 0.25) < 1e-9 + + +# --- (4) never inflates the LR --------------------------------------------- + +def test_never_exceeds_base(): + for cos in (-1.0, -0.3, 0.0, 0.9, None): + for var_i in (0.0, 0.3, 5.0, None): + r = rate(0.5, cos=cos, var_i=var_i, var_ref=0.3, + align_gate=True, inverse_var=True) + assert r <= 0.5 + 1e-12 + + +# --- (5) cosine primitive --------------------------------------------------- + +def _np(names): # minimal (name, param) list; param unused by _cosine_flat + return [(n, None) for n in names] + + +def test_cosine_aligned_and_opposite(): + g = {"w": torch.tensor([1.0, 0.0, 0.0])} + run_same = [torch.tensor([2.0, 0.0, 0.0])] + run_opp = [torch.tensor([-2.0, 0.0, 0.0])] + assert abs(cosf(g, run_same, _np(["w"])) - 1.0) < 1e-6 + assert abs(cosf(g, run_opp, _np(["w"])) + 1.0) < 1e-6 + + +def test_cosine_none_on_zero_running(): + g = {"w": torch.tensor([1.0, 1.0])} + assert cosf(g, [torch.zeros(2)], _np(["w"])) is None + + +def test_cosine_multi_tensor_orthogonal(): + g = {"a": torch.tensor([1.0, 0.0]), "b": torch.tensor([0.0, 0.0])} + run = [torch.tensor([0.0, 1.0]), torch.tensor([0.0, 0.0])] + assert abs(cosf(g, run, _np(["a", "b"]))) < 1e-6 diff --git a/lib/python/tests/mode/test_fwdllm_jvp_perf_opt.py b/lib/python/tests/mode/test_fwdllm_jvp_perf_opt.py new file mode 100644 index 000000000..f3cf47504 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_jvp_perf_opt.py @@ -0,0 +1,66 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Fluxtune JVP perf-opt (simulate_fwdllm.md §L) — the retained optimizations +must be BIT-IDENTICAL to the un-optimized grads, so enabling `jvp_perf_opt` on +fluxtune never changes training fidelity or real↔sim parity. + +The core claim is `calculate_jvp(..., trainable_idx=T)` == `calculate_jvp(...)` +when the perturbation `v` is zero on the non-T (frozen) params: `p - h·0 = p` +exactly, so both paths feed the model identical param values. These tests pin +that (the trainer-level gating just chooses which path to call).""" + +import pytest + +torch = pytest.importorskip("torch") + +from examples.fwdllm.trainer.forward_training.fwdgrad_utils import calculate_jvp + + +def _linear_func(params): + # a deterministic scalar of the params (order-independent) so the JVP is + # well-defined and precision-stable for an exact-equality check + return sum(p.sum() for p in params) + + +def _params_and_v(n=6, frozen=(1, 3, 4), shape=(4, 4), seed=0): + g = torch.Generator().manual_seed(seed) + params = [torch.randn(shape, generator=g) for _ in range(n)] + v = [torch.randn(shape, generator=g) for _ in range(n)] + for i in frozen: # frozen params get a ZERO perturbation + v[i] = torch.zeros(shape) + trainable_idx = [i for i in range(n) if i not in frozen] + return params, v, trainable_idx + + +class TestTrainableOnlyFD: + def test_bit_identical_to_all_param_fd(self): + params, v, tidx = _params_and_v() + l0, j0 = calculate_jvp(_linear_func, params, v) # legacy + l1, j1 = calculate_jvp(_linear_func, params, v, trainable_idx=tidx) # perf-opt + assert torch.equal(j0, j1), (j0, j1) + assert torch.equal(l0, l1) + + def test_none_is_the_legacy_path(self): + params, v, _ = _params_and_v() + a = calculate_jvp(_linear_func, params, v) + b = calculate_jvp(_linear_func, params, v, trainable_idx=None) + assert torch.equal(a[1], b[1]) and torch.equal(a[0], b[0]) + + def test_frozen_params_untouched_by_perturbation(self): + # a func that reads a FROZEN param heavily — its contribution must be + # identical in both signs (v=0 there), so it cancels out of the JVP + params, v, tidx = _params_and_v(frozen=(0,)) + def f(p): + return (p[0] ** 3).sum() + p[2].sum() # p[0] frozen, nonlinear + _, j_all = calculate_jvp(f, params, v) + _, j_opt = calculate_jvp(f, params, v, trainable_idx=tidx) + assert torch.equal(j_all, j_opt) + + def test_covers_extra_zero_indices_safely(self): + # trainable_idx may include an index whose v happens to be 0 (a trainable + # param that drew a zero dir) — still p-h·0=p, so no divergence + params, v, tidx = _params_and_v(frozen=(2,)) + v[0] = torch.zeros_like(v[0]) # a "trainable" idx with zero dir + l0, j0 = calculate_jvp(_linear_func, params, v) + l1, j1 = calculate_jvp(_linear_func, params, v, trainable_idx=tidx) + assert torch.equal(j0, j1) and torch.equal(l0, l1) diff --git a/lib/python/tests/mode/test_fwdllm_oracular_avail.py b/lib/python/tests/mode/test_fwdllm_oracular_avail.py index aeb0875a1..f7e16e53f 100644 --- a/lib/python/tests/mode/test_fwdllm_oracular_avail.py +++ b/lib/python/tests/mode/test_fwdllm_oracular_avail.py @@ -45,7 +45,7 @@ def test_reads_from_metadata_bundle(tmp_path): metadata_dir = _write_metadata(tmp_path) result = TopAggregator.read_trainer_unavailability( - None, trace="mobiperf_2st", metadata_dir=metadata_dir + None, trace="mobiperf_2st", base_dir=metadata_dir ) assert set(result.keys()) == {"task-aaa", "task-bbb"} @@ -62,6 +62,6 @@ def fail_glob(*args, **kwargs): monkeypatch.setattr(glob_module, "glob", fail_glob) result = TopAggregator.read_trainer_unavailability( - None, trace="mobiperf_2st", metadata_dir=metadata_dir + None, trace="mobiperf_2st", base_dir=metadata_dir ) assert len(result) == 2 diff --git a/lib/python/tests/mode/test_fwdllm_recv_timeout.py b/lib/python/tests/mode/test_fwdllm_recv_timeout.py index bcd487fb8..701ae934a 100644 --- a/lib/python/tests/mode/test_fwdllm_recv_timeout.py +++ b/lib/python/tests/mode/test_fwdllm_recv_timeout.py @@ -52,6 +52,8 @@ class _FakeAggregator: _aggregate_grads_async = TopAggregator._aggregate_grads_async def __init__(self, channel): + # Batch 1 flag-off: real path uses next(recv_fifo(...,1,timeout=...)). + self.simulated = False self.cm = _FakeChannelManager(channel) def test_recv_fifo_called_with_recv_timeout_wait_s(self): @@ -74,6 +76,8 @@ class _FakeAggregator: ) def __init__(self, channel, agg_goal=3): + # Batch 1 flag-off: real barrier path (unchanged recv_fifo drain). + self.simulated = False self._agg_goal = agg_goal self.ends_not_selected_yet = False self._round = 1 diff --git a/lib/python/tests/mode/test_fwdllm_rounds_termination.py b/lib/python/tests/mode/test_fwdllm_rounds_termination.py index 88cd9b975..1577a72d4 100644 --- a/lib/python/tests/mode/test_fwdllm_rounds_termination.py +++ b/lib/python/tests/mode/test_fwdllm_rounds_termination.py @@ -36,6 +36,10 @@ class _FakeAggregator: machinery stubbed out.""" def __init__(self, round_, rounds, total_data_bins=150, var_good_enough=True): + # Batch 1 flag-off regression: real path (time_mode: real). With + # simulated=False the new sim boundary hook (_release_sim_slots_at_agg_goal) + # is skipped, so _process_aggregation_goal_met behaves byte-identically. + self.simulated = False self._agg_goal = 2 self._agg_goal_cnt = 2 self._per_agg_trainer_list = [] diff --git a/lib/python/tests/mode/test_fwdllm_sct_model.py b/lib/python/tests/mode/test_fwdllm_sct_model.py new file mode 100644 index 000000000..9bef9dda4 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_sct_model.py @@ -0,0 +1,104 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Stage B -- sct-model folds (K-D20 #6), all config-gated OFF => byte-identical. + +B1 eval_s onto the vclock (aggregator): after a committed data_id's eval, + advance the vclock by the measured eval wall when simModelEvalTime is set. +B2 per-trainer straggler spread (trainer): a stable offset in [0, spread) added + to the modeled delay in SIM only. +B3 WAN transfer knob: documented, default 0 (verified inert here). + +Only the MECHANISM + the flag-off byte-identical invariant is unit-tested; +whether the folds drive wall_disparity->~0 is an EMERGENT run quantity. +""" + +import os +import sys +from types import SimpleNamespace + +sys.path.insert( + 0, + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "..", "examples", "fwdllm", + "trainer", "forward_training", + ), +) +from FedSgdTrainer import FedSGDTrainer # noqa: E402 + + +class _Host: + """Binds the straggler-offset helper onto a minimal stand-in.""" + + _sim_straggler_offset_s = FedSGDTrainer._sim_straggler_offset_s + + def __init__(self, trainer_id="3", simulated=True, spread=0.0): + self.trainer_id = trainer_id + self.simulated = simulated + self.config = SimpleNamespace( + hyperparameters=SimpleNamespace(sim_straggler_spread_s=spread) + ) + + +class TestStragglerSpreadB2: + def test_zero_spread_is_no_offset(self): + assert _Host(spread=0.0)._sim_straggler_offset_s() == 0.0 + + def test_real_mode_never_offsets(self): + # real mode gets dispersion from GPU contention, not this knob + assert _Host(simulated=False, spread=5.0)._sim_straggler_offset_s() == 0.0 + + def test_offset_within_spread_and_stable(self): + h = _Host(trainer_id="7", spread=2.3) + o1 = h._sim_straggler_offset_s() + o2 = h._sim_straggler_offset_s() + assert 0.0 <= o1 < 2.3 + assert o1 == o2 # deterministic per trainer (reproducible) + + def test_offset_varies_across_trainers(self): + offs = { + _Host(trainer_id=str(i), spread=2.3)._sim_straggler_offset_s() + for i in range(1, 11) + } + # crc32-derived fractions spread the cohort (not all identical) + assert len(offs) > 1 + + +class TestEvalOnVclockB1: + """B1 mechanism: the aggregator advances the vclock by the measured eval + wall only when simModelEvalTime is set (byte-identical off).""" + + def _run(self, flag): + from flame import telemetry + from tests.mode.test_fwdllm_agg_telemetry import ( + _FakeAggregator, _FakeChannel, + ) + from datetime import timedelta + + agg = _FakeAggregator(contributors=["t1"], var_good_enough=True) + agg.simulated = True + # sim path calls the boundary slot-release hook; not under test here. + agg._release_sim_slots_at_agg_goal = lambda *a, **k: None + agg._vclock = SimpleNamespace( + now=100.0, + advance=lambda ts: setattr(agg._vclock, "now", max(agg._vclock.now, ts)), + ) + agg.config.hyperparameters.sim_model_eval_time = flag + # eval_model in the fake returns instantly, so measured eval_s ~ 0; force + # a nonzero measured eval by making eval_model sleep a hair. + import time as _t + orig_eval = agg.eval_model + + def _slow_eval(): + _t.sleep(0.02) + return orig_eval() + + agg.eval_model = _slow_eval + channel = _FakeChannel(durations={"t1": timedelta(seconds=2)}) + agg._process_aggregation_goal_met(tag="aggregate", channel=channel) + return agg._vclock.now + + def test_flag_off_does_not_advance_vclock(self): + assert self._run(flag=False) == 100.0 # byte-identical: vclock untouched + + def test_flag_on_charges_eval_wall(self): + assert self._run(flag=True) > 100.0 # vclock advanced by measured eval_s diff --git a/lib/python/tests/mode/test_fwdllm_sim_drain_and_repick.py b/lib/python/tests/mode/test_fwdllm_sim_drain_and_repick.py new file mode 100644 index 000000000..ecfda27bc --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_sim_drain_and_repick.py @@ -0,0 +1,203 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Regression for the fluxtune sim deadlock (two bugs froze the async sim after +one cohort). + + Bug A -- drain gated on channel RECV state, not on the sct reorder buffer. + _aggregate_grads_async early-returned whenever the channel reported no end in + RECV. But _sim_recv_min_grad greedily drains ALL ready channel messages into + _sim_buffer on its first call (emptying RECV), so the already-received grads + were stranded in the buffer and never popped: agg_goal never met. The sim + commit path's readiness must key on _sim_buffer / _sim_inflight_expected, not + on the real transport's RECV bookkeeping. + + Bug B -- the async_oort re-pick triplet was stamped at DISPATCH. + Stamping the whole cohort at the current _curr_agg_version made every + dispatched trainer match the aggregator's version; since the version advances + only at a commit boundary (never reached, see Bug A), async_oort's filter + excluded the entire pool -> no re-dispatch, ever. The triplet is now stamped + on grad RETURN, so an in-flight-but-not-returned trainer stays eligible and + the pool is never frozen before the first commit. +""" + +import torch + +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator +from flame.mode.message import MessageType +from flame.sim.virtual_clock import SimReorderBuffer + + +# --------------------------------------------------------------------------- # +# Bug A -- the sim drain loop must not stall on empty RECV while the buffer is +# non-empty. +# --------------------------------------------------------------------------- # +class _GateChannel: + def __init__(self, recv_ends): + self._recv_ends = recv_ends + + def ends(self, state=None): + return self._recv_ends + + +class _CM: + def __init__(self, channel): + self._channel = channel + + def get_by_tag(self, tag): + return self._channel + + +class _DrainGateAgg: + """Binds the real _aggregate_grads_async and spies _sim_recv_min_grad so we + observe ONLY the entry-guard decision (drain vs. early-return).""" + + _aggregate_grads_async = TopAggregator._aggregate_grads_async + + def __init__(self, recv_ends, buffer_scts=(), inflight=(), simulated=True): + self.simulated = simulated + self._sim_buffer = SimReorderBuffer() + for i, s in enumerate(buffer_scts): + self._sim_buffer.add(f"b{i}", float(s), None) + self._sim_inflight_expected = {e: 1.0 for e in inflight} + self.cm = _CM(_GateChannel(recv_ends)) + self.drain_calls = [] + + def _sim_recv_min_grad(self, channel, recv_ends): + # Spy: record the call and return "nothing committable" so + # _aggregate_grads_async returns at `if not msg` without needing the + # heavy _process_single_trainer_message path. + self.drain_calls.append(list(recv_ends)) + return None, ("", 0) + + +class TestDrainGateNotBlockedByEmptyRecv: + def test_buffered_grads_drain_with_no_recv_ends(self): + """The core deadlock: RECV empty but the reorder buffer still holds grads + -> the loop MUST proceed to drain (call _sim_recv_min_grad), not stall.""" + agg = _DrainGateAgg(recv_ends=None, buffer_scts=(10.0, 20.0, 30.0)) + agg._aggregate_grads_async("param-channel") + assert agg.drain_calls == [[]], ( + "buffered grads must drain even when the channel reports no RECV ends" + ) + + def test_inflight_expected_keeps_loop_alive_with_no_recv_ends(self): + """A trainer still computing (in the in-flight gate) must also keep the + loop draining so the buffered minimum is not stranded behind it.""" + agg = _DrainGateAgg(recv_ends=None, inflight=("A", "B")) + agg._aggregate_grads_async("param-channel") + assert agg.drain_calls == [[]] + + def test_early_return_when_truly_idle(self): + """RECV empty AND buffer empty AND nothing in flight -> genuinely nothing + to do, so the loop still early-returns (no busy-spin on the drain).""" + agg = _DrainGateAgg(recv_ends=None) + agg._aggregate_grads_async("param-channel") + assert agg.drain_calls == [] + + def test_recv_ends_present_always_drains(self): + agg = _DrainGateAgg(recv_ends=["A"]) + agg._aggregate_grads_async("param-channel") + assert agg.drain_calls == [["A"]] + + def test_real_mode_ignores_sim_buffer(self): + """Real path must stay byte-identical: a populated _sim_buffer must NOT + make the real loop proceed when RECV is empty.""" + agg = _DrainGateAgg( + recv_ends=None, buffer_scts=(10.0,), simulated=False + ) + agg._aggregate_grads_async("param-channel") + assert agg.drain_calls == [] + + +# --------------------------------------------------------------------------- # +# Bug B -- the re-pick triplet is stamped on grad RETURN, not at dispatch. +# --------------------------------------------------------------------------- # +class _RepickChannel: + def __init__(self): + self._props = {} + self.provided_cleaned_up = [] + + class _Sel: + ordered_updates_recv_ends = [] + + self._selector = _Sel() + + def set_end_property(self, end, key, value): + self._props[(end, key)] = value + + def get_end_property(self, end, key): + return self._props.get((end, key)) + + def cleanup_provided_ends(self, end): + self.provided_cleaned_up.append(end) + + def cleanup_recvd_end(self, end): + pass + + +class _RepickAgg: + """Drives the real _process_single_trainer_message far enough to observe the + triplet-stamp side effect. aggregate_grads_from_trainers is stubbed (the + stamp lands before it) and telemetry stays disabled by default.""" + + process = TopAggregator._process_single_trainer_message + _release_end_on_return = TopAggregator._release_end_on_return + + def __init__(self, residence=True, simulated=True, curr_ver=(5, 2, 1)): + self.simulated = simulated + self._sim_inflight_residence = residence + self._curr_agg_version = curr_ver + self._trainer_state_dict = {} + self._per_agg_trainer_list = [] + self._agg_goal_cnt = 0 + self._round = 1 + self.data_id = 2 + self.iteration_per_data_id = 1 + self.is_async = True + self._model_version = 5 + self._sim_contrib_intervals = None + self._round_cache_activity_ts = {} + self._updates_in_queue = 0 + self._updates_received = {} + self._trainer_last_model_version = {} + self.grad_pool = [] + + def aggregate_grads_from_trainers(self, *a, **k): + pass # stub: the stamp executes before this call + + +def _valid_grad_msg(): + return { + MessageType.MODEL_VERSION: 5, + MessageType.GRADIENTS: torch.zeros(1), + MessageType.GRADIENTS_FOR_VAR_CHECK: torch.zeros(1), + MessageType.STAT_UTILITY: 1.0, + MessageType.DATASET_SIZE: 8, + } + + +class TestRepickTripletStampedOnReturn: + def test_returning_trainer_is_stamped_at_current_agg_version(self): + agg = _RepickAgg(residence=True, curr_ver=(5, 2, 1)) + ch = _RepickChannel() + + assert agg.process(ch, _valid_grad_msg(), "t1", timestamp=0) is True + # A RETURNED trainer records the triplet it just contributed at, so + # async_oort skips re-picking it until the agg version advances. + assert agg._trainer_state_dict == {"t1": (5, 2, 1)} + + def test_residence_off_leaves_triplet_map_empty(self): + """Sync baselines (residence off): empty map -> async_oort filter inert + -> byte-identical to Batch 1.""" + agg = _RepickAgg(residence=False) + ch = _RepickChannel() + + assert agg.process(ch, _valid_grad_msg(), "t1", timestamp=0) is True + assert agg._trainer_state_dict == {} + + def test_not_simulated_leaves_triplet_map_empty(self): + agg = _RepickAgg(residence=True, simulated=False) + ch = _RepickChannel() + + assert agg.process(ch, _valid_grad_msg(), "t1", timestamp=0) is True + assert agg._trainer_state_dict == {} diff --git a/lib/python/tests/mode/test_fwdllm_sim_grad_loop.py b/lib/python/tests/mode/test_fwdllm_sim_grad_loop.py new file mode 100644 index 000000000..63ccd7196 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_sim_grad_loop.py @@ -0,0 +1,597 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""The async grad-loop sim drive. + +In simulated mode fwdllm_aggregator._aggregate_grads_async commits a grad by the +modeled sim_completion_ts via _sim_recv_min_grad -- an sct-ordered reorder buffer ++ a one-in-flight gate + a virtual-clock advance -- so ordering is immune to +physical arrival jitter. These tests drive _sim_recv_min_grad directly with +synthetic out-of-order grad messages (a fake channel), and cover the +fwdllm-specific rollback-safety of _release_sim_slots_at_agg_goal (a data_id +spans many agg-goal cycles; a variance-FAIL rolls back to the same data_id and +must NOT strand or double-commit a grad). +""" + +import time +from collections import deque +from datetime import datetime + +from flame.mode.horizontal.asyncfl.top_aggregator import TopAggregator as _AsyncBase +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator +from flame.mode.horizontal.syncfl.top_aggregator import TopAggregator as _SyncBase +from flame.mode.message import MessageType +from flame.sim.virtual_clock import SimReorderBuffer, VirtualClock + + +class _FakeGradChannel: + """Serves at most one grad message per end (as fwdllm's one-message-per-call + loop expects). A message with release_at=k is only yielded once recv_fifo + has been probed more than k times -- this models a straggler whose modeled + completion is earlier but whose physical arrival is later, which is exactly + what the in-flight gate must wait for.""" + + def __init__(self, ends): + self._ends_set = set(ends) + self._msgs = {} # end -> message dict + self._release_at = {} # end -> probe index gating physical arrival + self._delivered = set() + self._probe = 0 + + def add_msg(self, end, sct, budget=None, release_at=0): + m = {MessageType.SIM_COMPLETION_TS: sct} + if budget is not None: + m[MessageType.TRAINING_BUDGET_S] = budget + self._ends_set.add(end) + self._msgs[end] = m + self._release_at[end] = release_at + + def ends(self, state=None): + return list(self._ends_set) + + def has(self, end): + return end in self._ends_set + + def recv_fifo(self, end_ids, first_k=None, timeout=None): + self._probe += 1 + for e in end_ids: + if e in self._delivered or e not in self._msgs: + continue + if self._probe > self._release_at.get(e, 0): + self._delivered.add(e) + yield self._msgs[e], (e, datetime.now()) + + def drain_ready(self, end_ids, timeout=None): + """Non-blocking snapshot: return every arrived-but-undelivered message + for the given ends (mirrors channel.drain_ready's direct rxq sweep).""" + self._probe += 1 + out = [] + for e in end_ids: + if e in self._delivered or e not in self._msgs: + continue + if self._probe > self._release_at.get(e, 0): + self._delivered.add(e) + out.append((self._msgs[e], (e, datetime.now()))) + return out + + +class _FakeGradAgg: + """Binds the real sim grad-loop methods onto a minimal stand-in.""" + + _sim_recv_min_grad = TopAggregator._sim_recv_min_grad + _release_sim_slots_at_agg_goal = TopAggregator._release_sim_slots_at_agg_goal + _advance_sim_clock = _SyncBase._advance_sim_clock + _sim_recv_grace_s = _SyncBase._sim_recv_grace_s + # #13 step 2 ready-gating helper (inherited by the real fwdllm agg from asyncfl). + _sim_end_has_ready_msg = staticmethod(TopAggregator._sim_end_has_ready_msg) + # #13 step 4 freed-slot FIFO consumer (inherited from asyncfl). + _pop_free_slot_ts = _AsyncBase._pop_free_slot_ts + # fwdllm overrides the hold with the two-lifetime split. + _sim_hold_busy_slots = TopAggregator._sim_hold_busy_slots + # The return-path guard/slot release (defers to COMMIT in sim residence). + _release_end_on_return = TopAggregator._release_end_on_return + # grace-window class knobs the base method reads off self + SIM_RECV_GRACE_FLOOR_S = 0.0 + SIM_RECV_GRACE_FACTOR = 0.0 + + def __init__(self): + self.simulated = True + self._round = 1 + self._vclock = VirtualClock() + self._sim_buffer = SimReorderBuffer() + self._sim_committed = set() + self._sim_inflight_expected = {} + self._sim_trainer_budget = {} + self._sim_budget_min = 12.0 + self._sim_fill_ema = 0.0 + self._sim_pending_commit = set() + self._sim_inflight_residence = False + self._sim_staggered_redispatch = False # #13 step 4 (default off) + self._sim_free_slot_ts = deque(maxlen=128) + self._trainer_state_dict = {} + + def _drain(self, channel, recv_ends, n): + """Commit n grads, returning the ordered list of committed scts.""" + scts = [] + for _ in range(n): + msg, md = self._sim_recv_min_grad(channel, recv_ends) + if not msg: + break + scts.append(msg[MessageType.SIM_COMPLETION_TS]) + return scts + + +class TestSctOrderedCommit: + def test_commits_in_sct_order_not_arrival_order(self): + agg = _FakeGradAgg() + ch = _FakeGradChannel([]) + # Physical arrival order A,B,C but modeled completion order B no grace burned waiting on it. + assert all("FAR" not in call for call in ch.probe_calls) + + def test_near_expected_straggler_is_probed(self): + agg = _FakeGradAgg() + # NEAR's exp (11) is within bmin(10)+slack -> the gate may still wait for + # it, so it MUST be probed. + agg._sim_inflight_expected = {"NEAR": 11.0} + ch = _RecordingChannel([]) + ch.add_msg("NEAR", sct=11.0, release_at=0) + agg._sim_buffer.add("A", 10.0, ({MessageType.SIM_COMPLETION_TS: 10.0}, ("A", datetime.now()))) + + agg._sim_recv_min_grad(ch, []) + assert any("NEAR" in call for call in ch.probe_calls) + + def test_physically_ready_straggler_is_probed_regardless_of_exp(self): + agg = _FakeGradAgg() + # READY is expected far in the future BUT its grad has physically arrived + # (ready rxq) -> drain it now so it buffers as a future rather than being + # committed past-dated later. + agg._sim_inflight_expected = {"READY": 1000.0} + ch = _RecordingChannel([], ready={"READY"}) + ch.add_msg("READY", sct=1000.0, release_at=0) + agg._sim_buffer.add("A", 10.0, ({MessageType.SIM_COMPLETION_TS: 10.0}, ("A", datetime.now()))) + + agg._sim_recv_min_grad(ch, []) + assert any("READY" in call for call in ch.probe_calls) + + +class TestSctOrderedDrain: + """#13 step 3: with sim_sct_ordered_drain ON, the drain ingests via + channel.drain_ready (direct rxq sweep, no per-end grace timeout) instead of + the blocking recv_fifo streamer, while keeping sct-ordered commit.""" + + def test_commits_via_drain_ready_not_recv_fifo(self): + agg = _FakeGradAgg() + agg._sim_sct_ordered_drain = True + ch = _RecordingChannel([]) + ch.add_msg("A", sct=30.0) + ch.add_msg("B", sct=10.0) + + scts = agg._drain(ch, ["A", "B"], 2) + assert scts == [10.0, 30.0] # still committed in sct order + assert ch.drain_calls # drain_ready was used... + assert not ch.probe_calls # ...and recv_fifo was NOT + + def test_recv_fifo_path_when_flag_off(self): + agg = _FakeGradAgg() # flag defaults off + ch = _RecordingChannel([]) + ch.add_msg("A", sct=10.0) + + agg._drain(ch, ["A"], 1) + assert ch.probe_calls # recv_fifo used + assert not ch.drain_calls # drain_ready NOT used + + def test_drain_ready_still_honors_inflight_gate(self): + """A (sct=100) arrives; B (sct=50) is expected earlier and arrives on the + 2nd sweep -- the gate must still hold A and commit B first.""" + agg = _FakeGradAgg() + agg._sim_sct_ordered_drain = True + agg._sim_inflight_expected = {"A": 98.0, "B": 48.0} + ch = _RecordingChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) + ch.add_msg("B", sct=50.0, release_at=1) + + first, _ = agg._sim_recv_min_grad(ch, ["A", "B"]) + assert first[MessageType.SIM_COMPLETION_TS] == 50.0 # B, not A + second, _ = agg._sim_recv_min_grad(ch, ["A", "B"]) + assert second[MessageType.SIM_COMPLETION_TS] == 100.0 + + +class TestFreedSlotRefill: + """#13 step 4: a grad commit stamps the freed-slot vclock into the FIFO; a + later dispatch pops it (oldest-first, clamped) as the re-dispatched trainer's + SEND vclock -- spreading expected completions instead of collapsing the cohort + at one round frontier.""" + + def test_commit_stamps_freed_slot_vclock_when_staggered(self): + agg = _FakeGradAgg() + agg._sim_staggered_redispatch = True + agg._sim_inflight_expected = {"X": 10.0} + ch = _FakeGradChannel([]) + ch.add_msg("X", sct=10.0) + + agg._sim_recv_min_grad(ch, ["X"]) + # vclock advanced to sct=10 on commit -> that freed-slot vclock is stamped. + assert list(agg._sim_free_slot_ts) == [10.0] + + def test_commit_does_not_stamp_when_flag_off(self): + agg = _FakeGradAgg() # staggered off (default) + agg._sim_inflight_expected = {"X": 10.0} + ch = _FakeGradChannel([]) + ch.add_msg("X", sct=10.0) + + agg._sim_recv_min_grad(ch, ["X"]) + assert len(agg._sim_free_slot_ts) == 0 + + def test_pop_free_slot_ts_is_fifo_and_clamped(self): + agg = _FakeGradAgg() + agg._sim_free_slot_ts.extend([5.0, 8.0]) + assert agg._pop_free_slot_ts(100.0) == 5.0 # oldest first + assert agg._pop_free_slot_ts(100.0) == 8.0 + assert agg._pop_free_slot_ts(100.0) == 100.0 # empty -> live frontier + agg._sim_free_slot_ts.append(50.0) + assert agg._pop_free_slot_ts(30.0) == 30.0 # min(stamp, round_now) + + def test_staggered_expected_completions_spread_not_bunched(self): + """End-to-end at the unit level: two commits free slots at vclock 10 and + 25; a subsequent 2-trainer refill pops those as SEND vclocks, so expected + completions (sst + budget) SPREAD (14, 29) instead of bunching at one + round frontier (which would give the same expected for both).""" + agg = _FakeGradAgg() + agg._sim_staggered_redispatch = True + agg._sim_budget_min = 4.0 + for e, s in [("A", 10.0), ("B", 25.0)]: + agg._sim_inflight_expected = {e: s} + ch = _FakeGradChannel([]) + ch.add_msg(e, sct=s) + agg._sim_recv_min_grad(ch, [e]) + assert list(agg._sim_free_slot_ts) == [10.0, 25.0] + + sst1 = agg._pop_free_slot_ts(100.0) + sst2 = agg._pop_free_slot_ts(100.0) + exp1 = sst1 + agg._sim_budget_min + exp2 = sst2 + agg._sim_budget_min + assert (exp1, exp2) == (14.0, 29.0) # spread, not (round_now+budget)×2 + + +class TestStuckEndEviction: + """#13 step 1: a trainer EXPECTED to complete earlier than the buffered + minimum but never physically arrives must be evicted from + _sim_inflight_expected on the recv failsafe deadline -- else `earlier_stuck` + re-fires the full 30s deadline every drain cycle and the composer freezes.""" + + def _immediate_deadline(self, monkeypatch): + # Fire the recv failsafe on the first pass (no real 30s wait). + import flame.mode.horizontal.syncfl.fwdllm_aggregator as fa + monkeypatch.setattr(fa, "RECV_TIMEOUT_WAIT_S", 0.0) + + def test_stuck_end_evicted_on_deadline_and_buffered_min_commits(self, monkeypatch): + self._immediate_deadline(monkeypatch) + agg = _FakeGradAgg() + # STUCK is expected early (10) but has no message; A arrived at sct=100. + agg._sim_inflight_expected = {"STUCK": 10.0, "A": 98.0} + ch = _FakeGradChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) + + msg, _md = agg._sim_recv_min_grad(ch, ["A", "STUCK"]) + + # The buffered min commits despite the earlier-expected straggler... + assert msg[MessageType.SIM_COMPLETION_TS] == 100.0 + # ...and STUCK is dropped so it can't block future cycles. + assert "STUCK" not in agg._sim_inflight_expected + assert agg._sim_gate_failsafe == 1 + + def test_evicted_end_does_not_block_the_next_cycle(self, monkeypatch): + self._immediate_deadline(monkeypatch) + agg = _FakeGradAgg() + agg._sim_inflight_expected = {"STUCK": 10.0} + ch = _FakeGradChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) + ch.add_msg("B", sct=200.0, release_at=0) + + # Cycle 1: STUCK holds the gate, hits the failsafe, is evicted; A commits. + m1, _ = agg._sim_recv_min_grad(ch, ["A", "B", "STUCK"]) + assert m1[MessageType.SIM_COMPLETION_TS] == 100.0 + assert "STUCK" not in agg._sim_inflight_expected + + # Cycle 2: with STUCK gone, B commits WITHOUT re-arming the failsafe. + m2, _ = agg._sim_recv_min_grad(ch, ["A", "B", "STUCK"]) + assert m2[MessageType.SIM_COMPLETION_TS] == 200.0 + assert agg._sim_gate_failsafe == 1 # not re-incremented + + def test_no_spurious_eviction_when_buffered_min_is_the_true_next(self, monkeypatch): + # STUCK expected LATER than the buffered min -> not `earlier_stuck` -> + # commit proceeds without touching the failsafe or the expected set. + self._immediate_deadline(monkeypatch) + agg = _FakeGradAgg() + agg._sim_inflight_expected = {"LATE": 500.0} + ch = _FakeGradChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) + + msg, _ = agg._sim_recv_min_grad(ch, ["A", "LATE"]) + assert msg[MessageType.SIM_COMPLETION_TS] == 100.0 + assert agg._sim_inflight_expected == {"LATE": 500.0} # untouched + assert getattr(agg, "_sim_gate_failsafe", 0) == 0 + + +class TestRollbackSafety: + def test_no_double_commit_within_a_cycle(self): + agg = _FakeGradAgg() + agg._sim_inflight_expected = {"X": 10.0} + ch = _FakeGradChannel([]) + ch.add_msg("X", sct=10.0) + + msg1, _ = agg._sim_recv_min_grad(ch, ["X"]) + assert msg1[MessageType.SIM_COMPLETION_TS] == 10.0 + assert agg._sim_committed == {"X"} + + # Second call in the SAME cycle: X is already committed and no new + # message exists -> nothing committable, no re-commit. + msg2, _ = agg._sim_recv_min_grad(ch, ["X"]) + assert msg2 is None + + def test_boundary_clear_lets_same_end_recontribute_after_rollback(self): + agg = _FakeGradAgg() + # --- cycle 1: X commits on data_id d, iteration 0 --- + agg._sim_inflight_expected = {"X": 10.0} + ch1 = _FakeGradChannel([]) + ch1.add_msg("X", sct=10.0) + m1, _ = agg._sim_recv_min_grad(ch1, ["X"]) + assert m1[MessageType.SIM_COMPLETION_TS] == 10.0 + assert agg._sim_committed == {"X"} + + # --- variance FAIL -> rollback: agg-goal boundary cleanup --- + agg._release_sim_slots_at_agg_goal(ch1, is_async=False) + assert agg._sim_committed == set() # not stranded as committed + assert len(agg._sim_buffer) == 0 # no stranded grad + assert agg._sim_inflight_expected == {} + + # --- cycle 2: SAME data_id, iteration 1, X re-contributes --- + agg._sim_inflight_expected = {"X": 25.0} # re-armed at dispatch + ch2 = _FakeGradChannel([]) + ch2.add_msg("X", sct=25.0) + m2, _ = agg._sim_recv_min_grad(ch2, ["X"]) + # X is committed again (NOT skipped as a stale committed mark). + assert m2[MessageType.SIM_COMPLETION_TS] == 25.0 + assert agg._sim_committed == {"X"} + assert agg._vclock.now == 25.0 + + +class _FakeEnd: + def set_property(self, *_args, **_kwargs): + pass + + +class _FakeSelector: + def __init__(self, selected): + self.requester = "agg" + self.all_selected = {e: 0.0 for e in selected} + self.selected_ends = {"agg": set(selected)} + + +class _FakeSelChannel(_FakeGradChannel): + def __init__(self, ends): + super().__init__(ends) + self._selector = _FakeSelector(ends) + self._ends = {e: _FakeEnd() for e in ends} + + def add_msg(self, end, sct, budget=None, release_at=0): + # Keep the real channel invariant has(e) <=> e in _ends (the slot-hold + # path does channel._ends[e] guarded only by channel.has(e)). + super().add_msg(end, sct, budget, release_at) + self._ends.setdefault(end, _FakeEnd()) + + # Faithful mirror of channel.cleanup_provided_ends -> selector + # _cleanup_provided_ends: drop the end from the re-pick guard + slot. + def cleanup_provided_ends(self, end): + self._selector.all_selected.pop(end, None) + self._selector.selected_ends[self._selector.requester].discard(end) + + def cleanup_recvd_end(self, end): # sync path (random selector) + self.cleanup_provided_ends(end) + + +class TestAsyncBoundaryReleasesSlots: + def test_async_boundary_frees_every_committed_slot(self): + agg = _FakeGradAgg() + agg._sim_pending_commit = set() + agg._sim_inflight_residence = False + ch = _FakeSelChannel(["X", "Y"]) + agg._sim_committed = {"X", "Y"} + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + + # buffer empty -> nothing held -> both slots released for re-selection. + assert ch._selector.all_selected == {} + assert ch._selector.selected_ends["agg"] == set() + assert agg._sim_committed == set() + + +class TestFlagOffNoOp: + def test_release_is_noop_when_not_simulated(self): + agg = _FakeGradAgg() + agg.simulated = False + agg._sim_committed = {"X"} + agg._sim_inflight_expected = {"X": 1.0} + agg._sim_buffer.add("X", 1.0, None) + + # Real mode: the boundary hook must not touch any sim state. + agg._release_sim_slots_at_agg_goal(channel=None, is_async=False) + + assert agg._sim_committed == {"X"} + assert agg._sim_inflight_expected == {"X": 1.0} + assert len(agg._sim_buffer) == 1 + + +class TestComputeTruthfulGate: + """#15: the compute-truthful commit gate (`sim_compute_truthful_gate`, + flag-gated, default off) must not block a ready commit on a trainer that is + NOT actually computing -- one stamped 'expected' at dispatch but idle-in-recv + behind the single-threaded drain. That phantom wait burns the grace floor / + 30s failsafe and, via hold-to-commit, starves re-dispatch. The guard stays + SELECTIVE: a genuine in-window straggler is still waited for, so sct-commit + order is preserved.""" + + def test_idle_phantom_is_skipped_and_ready_grad_commits(self): + agg = _FakeGradAgg() + agg._sim_compute_truthful_gate = True + agg._sim_gate_compute_cap_s = 10.0 + agg._sim_dispatch_wall = {} # PHANTOM: never dispatched -> no wall stamp + agg._sim_inflight_expected = {"PHANTOM": 10.0, "A": 98.0} + ch = _FakeGradChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) # A arrived; PHANTOM never will + + # No monkeypatched deadline: flag OFF would spin to the 30s failsafe; + # flag ON skips the phantom and commits A at once. + msg, _md = agg._sim_recv_min_grad(ch, ["A", "PHANTOM"]) + + assert msg[MessageType.SIM_COMPLETION_TS] == 100.0 + assert agg._sim_gate_phantom_skip >= 1 + assert getattr(agg, "_sim_gate_failsafe", 0) == 0 # never hit the failsafe + + def test_stale_dispatch_is_treated_as_phantom(self): + agg = _FakeGradAgg() + agg._sim_compute_truthful_gate = True + agg._sim_gate_compute_cap_s = 10.0 + # dispatched 30s ago -> older than the 10s cap -> not genuinely computing. + agg._sim_dispatch_wall = {"STALE": time.time() - 30.0} + agg._sim_inflight_expected = {"STALE": 10.0, "A": 98.0} + ch = _FakeGradChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) + + msg, _md = agg._sim_recv_min_grad(ch, ["A", "STALE"]) + assert msg[MessageType.SIM_COMPLETION_TS] == 100.0 + assert agg._sim_gate_phantom_skip >= 1 + + def test_in_window_straggler_is_still_held(self): + """Selective guard: a trainer dispatched WITHIN the compute window is a + genuine straggler and must still be waited for -- commit order preserved.""" + agg = _FakeGradAgg() + agg._sim_compute_truthful_gate = True + agg._sim_gate_compute_cap_s = 10.0 + agg._sim_dispatch_wall = {"B": time.time()} # just dispatched -> computing + agg._sim_inflight_expected = {"A": 98.0, "B": 48.0} + ch = _FakeGradChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) # arrives first + ch.add_msg("B", sct=50.0, release_at=1) # arrives on the 2nd probe + + first_msg, _md = agg._sim_recv_min_grad(ch, ["A", "B"]) + assert first_msg[MessageType.SIM_COMPLETION_TS] == 50.0 # B still held-for + assert agg._vclock.now == 50.0 + + def test_flag_off_is_byte_identical(self, monkeypatch): + """Flag OFF (default): a phantom still blocks to the failsafe (unchanged) + and phantom_skip stays 0.""" + import flame.mode.horizontal.syncfl.fwdllm_aggregator as fa + monkeypatch.setattr(fa, "RECV_TIMEOUT_WAIT_S", 0.0) + agg = _FakeGradAgg() # flag defaults off (never set) + agg._sim_inflight_expected = {"PHANTOM": 10.0, "A": 98.0} + ch = _FakeGradChannel([]) + ch.add_msg("A", sct=100.0, release_at=0) + + msg, _md = agg._sim_recv_min_grad(ch, ["A", "PHANTOM"]) + assert msg[MessageType.SIM_COMPLETION_TS] == 100.0 + assert getattr(agg, "_sim_gate_phantom_skip", 0) == 0 + assert agg._sim_gate_failsafe == 1 diff --git a/lib/python/tests/mode/test_fwdllm_sim_grad_residence.py b/lib/python/tests/mode/test_fwdllm_sim_grad_residence.py new file mode 100644 index 000000000..d24af4ce9 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_sim_grad_residence.py @@ -0,0 +1,338 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Async grad-path residence + commit-then-carry. + +For the async path with `sim_inflight_residence` on, `_release_sim_slots_at_agg_goal` +HOLDs the still-busy trainers (surplus buffered u not-yet-arrived in-flight) in +their slots BEFORE clearing anything, releases only the committed subset, and +CARRIEs the surplus buffer to the next cycle (never dropped) -- otherwise the +boundary re-dispatches busy trainers (2x forward passes) and drops arrived-but- +uncommitted grads. These tests drive the boundary directly and assert (a) surplus +carried, (b) busy trainers not re-selected, (c) R1 one-in-flight residence holds, +and (d) flag-off => byte-identical to the legacy drop behavior (sync baselines + +async-without-residence unchanged). +""" + +from flame.mode.horizontal.asyncfl.top_aggregator import ( + TopAggregator as _AsyncBase, +) +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator +from flame.mode.horizontal.syncfl.top_aggregator import TopAggregator as _SyncBase +from flame.mode.message import MessageType +from flame.sim.virtual_clock import SimReorderBuffer, VirtualClock + +from tests.mode.test_fwdllm_sim_grad_loop import ( + _FakeGradAgg, + _FakeSelChannel, +) + + +def _residence_agg(residence: bool) -> _FakeGradAgg: + agg = _FakeGradAgg() + agg._sim_pending_commit = set() + agg._sim_inflight_residence = residence + return agg + + +class TestCommitThenCarryResidenceOn: + """async + sim_inflight_residence=True: hold-before-clear + carry surplus.""" + + def test_surplus_carried_and_busy_held(self): + agg = _residence_agg(residence=True) + ch = _FakeSelChannel(["A", "B", "C", "D", "E"]) + # Committed this cycle: A, B (already popped from the buffer). Surplus + # arrived: C, D (still buffered). Still in flight (not arrived): E. + agg._sim_committed = {"A", "B"} + agg._sim_buffer.add("C", 30.0, None) + agg._sim_buffer.add("D", 40.0, None) + agg._sim_inflight_expected = {"E": 50.0} + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + + # Surplus NOT dropped -- carried to the next cycle. + assert set(agg._sim_buffer.pending_ends()) == {"C", "D"} + # Not-yet-arrived trainer still in flight. + assert agg._sim_inflight_expected == {"E": 50.0} + # Per-cycle committed marks cleared so a re-contributor isn't skipped. + assert agg._sim_committed == set() + # Every OUTSTANDING trainer (carried surplus C,D u still-computing E) holds + # BOTH its re-pick guard (all_selected) AND its compute slot (selected_ends) + # until it commits; the two committed ones (A, B) are released. + assert set(ch._selector.all_selected) == {"C", "D", "E"} + assert ch._selector.selected_ends["agg"] == {"C", "D", "E"} + + def test_end_to_end_surplus_commits_next_cycle_no_refetch(self): + agg = _residence_agg(residence=True) + ends = ["A", "B", "C", "D", "E"] + ch = _FakeSelChannel([]) + for e, sct in zip(ends, (10.0, 20.0, 30.0, 40.0, 50.0)): + ch.add_msg(e, sct) + agg._sim_inflight_expected = {e: sct + 0.1 for e, sct + in zip(ends, (10, 20, 30, 40, 50))} + + # agg_goal = 2 commits this cycle (A, B by sct order); C, D, E arrive + # into the buffer as surplus. + scts = agg._drain(ch, ends, 2) + assert scts == [10.0, 20.0] + assert set(agg._sim_buffer.pending_ends()) == {"C", "D", "E"} + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + # Carried, not dropped. + assert set(agg._sim_buffer.pending_ends()) == {"C", "D", "E"} + + # Next cycle: the carried surplus commits from the buffer WITHOUT a new + # dispatch/forward pass (no re-fetch), in sct order. + scts2 = agg._drain(ch, ends, 3) + assert scts2 == [30.0, 40.0, 50.0] + + def test_no_same_trainer_reselected_while_in_flight(self): + """R1 residence: a held (still-busy) trainer must not reappear as a fresh + selection slot -- that is exactly the re-dispatch-while-in-flight bug.""" + agg = _residence_agg(residence=True) + ch = _FakeSelChannel(["A", "B", "C"]) + agg._sim_committed = {"A"} + agg._sim_buffer.add("B", 20.0, None) # surplus, arrived + agg._sim_inflight_expected = {"C": 30.0} # still computing + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + + held = set(ch._selector.all_selected) + # The busy trainers stay held (not freed for a duplicate dispatch); + # only the committed A is released. + assert "B" in held and "C" in held and "A" not in held + assert "B" in agg._sim_pending_commit and "C" in agg._sim_pending_commit + # Both the carried B and still-computing C keep their compute slot (in + # flight in virtual time until commit); neither can be re-picked. + assert ch._selector.selected_ends["agg"] == {"B", "C"} + + +class TestReturnPathGuardHeldToCommit: + """The guard release on grad RETURN (`_release_end_on_return`, called from + `_process_single_trainer_message`). The async accept path must not call + `channel.cleanup_provided_ends(end)` on physical return -- that tears the + trainer out of `all_selected` while its carried grad has not committed in + virtual time -> re-selectable -> re-dispatch-while-in-flight. + """ + + def test_guard_held_on_return_in_sim_residence(self): + """async + sim + residence: return must NOT release the re-pick guard + (held to COMMIT by _sim_hold_busy_slots). This is the regression guard.""" + agg = _residence_agg(residence=True) + agg.is_async = True + ch = _FakeSelChannel(["A", "B", "C"]) # all dispatched + in flight + agg._release_end_on_return(ch, "A") # A's grad returns (carried) + # A stays in the guard -> cannot be re-picked while still outstanding. + assert "A" in ch._selector.all_selected + assert "A" in ch._selector.selected_ends["agg"] + + def test_guard_released_on_return_when_residence_off(self): + """async WITHOUT residence: legacy behavior -- release immediately + (return ~= commit), so the fix is byte-identical off the residence path.""" + agg = _residence_agg(residence=False) + agg.is_async = True + ch = _FakeSelChannel(["A", "B", "C"]) + agg._release_end_on_return(ch, "A") + assert "A" not in ch._selector.all_selected + + def test_sync_return_uses_recvd_cleanup(self): + """sync (random selector, is_async=False): unchanged cleanup_recvd_end + path -- releases on return (barrier re-selects the whole cohort).""" + agg = _residence_agg(residence=True) # residence flag is inert when sync + agg.is_async = False + ch = _FakeSelChannel(["A", "B", "C"]) + agg._release_end_on_return(ch, "A") + assert "A" not in ch._selector.all_selected + + +class TestVirtualInflightSlotHold: + """A returned-but-uncommitted trainer is still in flight in VIRTUAL time (its + grad commits when the vclock reaches its sct), so it KEEPS its compute slot + (selected_ends, drives `extra`) until COMMIT, not on physical return. Both + ledgers track the same virtual-time in-flight set until commit.""" + + def test_returned_trainer_keeps_slot_until_commit(self): + agg = _residence_agg(residence=True) + ch = _FakeSelChannel(["A", "B", "C", "D"]) + # A,B still computing; C,D returned (buffered surplus, not yet consumed). + agg._sim_inflight_expected = {"A": 10.0, "B": 20.0, "C": 30.0, "D": 40.0} + agg._sim_buffer.add("C", 30.0, None) + agg._sim_buffer.add("D", 40.0, None) + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + + # Slot ledger (extra = c - len(selected_ends)): ALL four outstanding + # trainers occupy a slot -- carried C,D are still in flight in virtual + # time, so `in_flight` telemetry (= len(selected_ends)) counts them. + assert ch._selector.selected_ends["agg"] == {"A", "B", "C", "D"} + # Guard ledger: all four un-re-pickable until they commit. + assert set(ch._selector.all_selected) == {"A", "B", "C", "D"} + + def test_selected_ends_tracks_inflight_across_a_commit(self): + """The invariant that fixes the concurrency mis-measurement: after each + commit, selected_ends == the still-outstanding (dispatched-not-committed) + set. Committing one trainer releases exactly its slot; the rest stay.""" + agg = _residence_agg(residence=True) + ch = _FakeSelChannel([]) + for e, sct in zip(["A", "B", "C"], (10.0, 20.0, 30.0)): + ch.add_msg(e, sct) + agg._sim_inflight_expected = {"A": 10.0, "B": 20.0, "C": 30.0} + + # Commit the smallest-sct grad (A). _sim_recv_min_grad reasserts the hold. + scts = agg._drain(ch, ["A", "B", "C"], 1) + assert scts == [10.0] + # A committed -> slot released; B, C still in flight -> keep their slots. + assert ch._selector.selected_ends["agg"] == {"B", "C"} + assert "A" not in ch._selector.all_selected + + def test_triplet_guard_pruned_to_busy_on_commit(self): + agg = _residence_agg(residence=True) + ch = _FakeSelChannel(["A", "B", "C"]) + # A committed (popped) this cycle; B carried; C computing. + agg._trainer_state_dict = { + "A": (1, 0, 0), "B": (1, 0, 0), "C": (1, 0, 0), + } + agg._sim_committed = {"A"} + agg._sim_buffer.add("B", 20.0, None) + agg._sim_inflight_expected = {"C": 30.0} + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + + # The committed A is dropped from the triplet guard (re-pickable); the + # still-outstanding B, C remain so async_oort's filter keeps skipping them. + assert set(agg._trainer_state_dict) == {"B", "C"} + + +class TestFlagOffByteIdentical: + """Residence OFF (default) => the legacy drop, unchanged. Sync baselines + never take the carry path, so their boundary is untouched.""" + + def test_async_residence_off_drops_and_releases_all(self): + agg = _residence_agg(residence=False) + ch = _FakeSelChannel(["X", "Y"]) + agg._sim_committed = {"X", "Y"} + agg._sim_buffer.add("Z", 5.0, None) + agg._sim_inflight_expected = {"W": 9.0} + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + + # Legacy drop: buffer + in-flight cleared, every slot released. + assert len(agg._sim_buffer) == 0 + assert agg._sim_inflight_expected == {} + assert agg._sim_committed == set() + assert ch._selector.all_selected == {} + + def test_sync_barrier_always_drops_even_with_residence_flag(self): + # The sync path (fwdllm/fwdllm_plus, c ~= agg_goal) has no surplus, so it + # keeps the drop regardless of the residence flag. + agg = _residence_agg(residence=True) + agg._sim_committed = {"X"} + agg._sim_buffer.add("X", 1.0, None) + agg._sim_inflight_expected = {"X": 1.0} + + agg._release_sim_slots_at_agg_goal(channel=None, is_async=False) + + assert agg._sim_committed == set() + assert len(agg._sim_buffer) == 0 + assert agg._sim_inflight_expected == {} + + def test_not_simulated_is_noop(self): + agg = _residence_agg(residence=True) + agg.simulated = False + agg._sim_committed = {"X"} + agg._sim_buffer.add("X", 1.0, None) + agg._sim_inflight_expected = {"X": 1.0} + + agg._release_sim_slots_at_agg_goal(channel=None, is_async=True) + + assert agg._sim_committed == {"X"} + assert len(agg._sim_buffer) == 1 + assert agg._sim_inflight_expected == {"X": 1.0} + + +class TestPendingCommitBridge: + """The fwdllm aggregator maintains its VIRTUAL in-flight set + (`_sim_pending_commit`) and BINDS it to the selector's + `_agg_pending_commit_ref` so async_oort's eligibility filter excludes a + returned-but-uncommitted trainer regardless of `all_selected` churn. The + reconcile must SHRINK (a committed trainer becomes re-pickable); a + `|= outstanding` accumulate would starve every committed trainer forever. + """ + + def test_hold_reconciles_pending_to_outstanding_and_binds_ref(self): + agg = _residence_agg(residence=True) + ch = _FakeSelChannel(["A", "B", "C"]) + agg._sim_committed = {"A"} # committed this cycle + agg._sim_buffer.add("B", 20.0, None) # returned, carried (uncommitted) + agg._sim_inflight_expected = {"C": 30.0} # still computing + + agg._sim_hold_busy_slots(ch) + + # pending == still-outstanding (carried B ∪ computing C); committed A dropped. + assert agg._sim_pending_commit == {"B", "C"} + # bound to the SAME live object the selector reads (never rebind). + assert ch._selector._agg_pending_commit_ref is agg._sim_pending_commit + + def test_committed_trainer_drops_out_not_starved(self): + agg = _residence_agg(residence=True) + ch = _FakeSelChannel(["A", "B"]) + agg._sim_inflight_expected = {"A": 10.0, "B": 20.0} + agg._sim_hold_busy_slots(ch) + assert agg._sim_pending_commit == {"A", "B"} + + # A commits -> leaves in-flight, marked committed. + agg._sim_inflight_expected = {"B": 20.0} + agg._sim_committed = {"A"} + agg._sim_hold_busy_slots(ch) + + # A is re-pickable again (dropped from pending); `|=` would have kept it. + assert agg._sim_pending_commit == {"B"} + + def test_commit_discards_from_pending(self): + agg = _residence_agg(residence=True) + ch = _FakeSelChannel([]) + for e, sct in zip(["A", "B"], (10.0, 20.0)): + ch.add_msg(e, sct) + agg._sim_inflight_expected = {"A": 10.0, "B": 20.0} + agg._sim_pending_commit = {"A", "B"} + + agg._drain(ch, ["A", "B"], 1) # commit the smallest-sct grad (A) + + assert "A" not in agg._sim_pending_commit # discarded on COMMIT + assert "B" in agg._sim_pending_commit # still in flight + + def test_legacy_drop_path_clears_pending(self): + # residence OFF -> the boundary drops all in-flight; pending must clear + # too (else the sync-barrier set grows unbounded). + agg = _residence_agg(residence=False) + ch = _FakeSelChannel([]) + agg._sim_committed = {"X"} + agg._sim_buffer.add("Y", 5.0, None) + agg._sim_inflight_expected = {"Z": 9.0} + agg._sim_pending_commit = {"X", "Y", "Z"} + + agg._release_sim_slots_at_agg_goal(ch, is_async=True) + + assert agg._sim_pending_commit == set() + + def test_recommitted_trainer_stays_pending_despite_stale_committed(self): + """`_sim_committed` is a STALE cross-cycle marker (cleared only at the + boundary). A trainer that committed then was re-picked + re-dispatched is + back in `_sim_inflight_expected`; the reconcile must NOT drop it from + pending just because it lingers in `_sim_committed` -- else it is + re-pickable while its NEW dispatch is still in flight (R1 overlap). + `outstanding` therefore keys on inflight/buffer membership only, never + `- _sim_committed`.""" + agg = _residence_agg(residence=True) + ch = _FakeSelChannel(["A", "B"]) + # A committed earlier this cycle (still in the stale marker) AND has been + # re-dispatched -> back in flight. B is a first-time in-flight trainer. + agg._sim_committed = {"A"} + agg._sim_inflight_expected = {"A": 40.0, "B": 20.0} + + agg._sim_hold_busy_slots(ch) + + # A is genuinely in flight again -> stays pending + held (un-re-pickable). + assert "A" in agg._sim_pending_commit + assert "B" in agg._sim_pending_commit + assert "A" in ch._selector.all_selected + assert "A" in ch._selector.selected_ends["agg"] diff --git a/lib/python/tests/mode/test_fwdllm_sim_speedup_waits.py b/lib/python/tests/mode/test_fwdllm_sim_speedup_waits.py new file mode 100644 index 000000000..52636997f --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_sim_speedup_waits.py @@ -0,0 +1,119 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Phase 2 (speedup leak, root #13): the simulated clock must NOT pay +real-transport wall waits that have no fidelity value -- the sim's job is to +advance the virtual clock FASTER than physical wall (sim_rate >= 1), so every +skippable per-round sleep on the sim critical path is a slowdown. + +Telemetry localization (banked run_20260704_134801 sim): the trainer's +inter-round `mqtt_fetch_s` is barrier-wait realized by the blocking recv in +_fetch_weights (which MUST stay -- it delivers the real weights the forward-grad +pass needs for grad mode-invariance; async_cifar10 keeps real MQTT recv in sim +too). The only additive, fidelity-free per-round wall the sim can skip is: + + (1) the trainer's `pause_execution` throttle -- a `time.sleep(1)` chained at + the tail of EVERY trainer loop iteration ("don't overwhelm mqtt", a + real-transport artifact, principle #8); and + (2) the `_check_availability` avail-spin -- a `while UN_AVL: time.sleep(1)` + that, in sim, would freeze the virtual clock (sim time can't advance while + a trainer blocks on time.sleep); sim availability is enforced agg-side. + +Both are gated so REAL mode is byte-identical and SIM mode never wall-sleeps. +""" + +import os +import sys + +sys.path.insert( + 0, + os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "..", "examples", "fwdllm", + "trainer", "forward_training", + ), +) + +import FedSgdTrainer as _fst_module # noqa: E402 +from FedSgdTrainer import FedSGDTrainer # noqa: E402 +from flame.config import TrainerAvailState # noqa: E402 +from flame.mode.horizontal.syncfl import fwdllm_trainer as _tr_module # noqa: E402 + + +class _PauseTrainer: + """Minimal stand-in binding the base Trainer.pause_execution under test. + timer_decorator reads its own runtime.time (untouched here) and only a + best-effort self.fwd_llm_stage, so no other state is needed.""" + + pause_execution = _tr_module.Trainer.pause_execution + + def __init__(self, simulated): + self.simulated = simulated + + +class TestPauseExecutionGatedInSim: + def test_real_mode_pauses_one_second(self, monkeypatch): + slept = [] + monkeypatch.setattr(_tr_module.time, "sleep", lambda s: slept.append(s)) + _PauseTrainer(simulated=False).pause_execution() + assert slept == [1] # real byte-identical: the MQTT throttle stays + + def test_sim_mode_does_not_pause(self, monkeypatch): + slept = [] + monkeypatch.setattr(_tr_module.time, "sleep", lambda s: slept.append(s)) + _PauseTrainer(simulated=True).pause_execution() + assert slept == [] # root #13: no per-round wall charged to the sim + + def test_sim_default_missing_attr_is_treated_real(self, monkeypatch): + """A stand-in with no `simulated` attr (getattr default False) must keep + the real throttle -- the gate must never silently skip in real mode.""" + slept = [] + monkeypatch.setattr(_tr_module.time, "sleep", lambda s: slept.append(s)) + t = _PauseTrainer(simulated=False) + del t.simulated + t.pause_execution() + assert slept == [1] + + +class _AvailTrainer: + """Binds the real _check_availability onto a minimal stand-in.""" + + _check_availability = FedSGDTrainer._check_availability + + def __init__(self, simulated, avl_state, wait_until_next_avl=True): + self.simulated = simulated + self.avl_state = avl_state + self.wait_until_next_avl = wait_until_next_avl + self.trainer_id = "t1" + + +class TestCheckAvailabilityGatedInSim: + def test_sim_never_spins_and_proceeds(self, monkeypatch): + """UN_AVL trainer in sim: must NOT time.sleep (would freeze the vclock) + and must return True so it proceeds -- the agg-side send-gate handles + the withhold.""" + def _boom(_s): + raise AssertionError("time.sleep called on the simulated avail path") + monkeypatch.setattr(_fst_module.time, "sleep", _boom) + t = _AvailTrainer(simulated=True, avl_state=TrainerAvailState.UN_AVL) + assert t._check_availability() is True + + def test_real_mode_still_spins_until_available(self, monkeypatch): + """Real mode is byte-identical: it wall-sleeps while UN_AVL and returns + True once the trainer flips back to AVL_TRAIN.""" + t = _AvailTrainer(simulated=False, avl_state=TrainerAvailState.UN_AVL) + slept = [] + + def _sleep_then_avail(s): + slept.append(s) + t.avl_state = TrainerAvailState.AVL_TRAIN # becomes available after 1 tick + + monkeypatch.setattr(_fst_module.time, "sleep", _sleep_then_avail) + assert t._check_availability() is True + assert slept == [1] + + def test_real_unavailable_no_wait_exits(self, monkeypatch): + """wait_until_next_avl=False is unchanged in both modes: skip training.""" + monkeypatch.setattr(_fst_module.time, "sleep", + lambda s: (_ for _ in ()).throw(AssertionError("no sleep"))) + t = _AvailTrainer(simulated=False, avl_state=TrainerAvailState.UN_AVL, + wait_until_next_avl=False) + assert t._check_availability() is False diff --git a/lib/python/tests/mode/test_fwdllm_sim_sync_barrier.py b/lib/python/tests/mode/test_fwdllm_sim_sync_barrier.py new file mode 100644 index 000000000..39546f60b --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_sim_sync_barrier.py @@ -0,0 +1,150 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Batch 1 (simulate_fwdllm.md §I.8/§J.3): the SYNC-barrier sim drive. + +fwdllm/fwdllm_plus are sync -- their primary commit path is the barrier +(sync_collect_and_accumulate_grads). In simulated mode it commits the k +trainers with the SMALLEST modeled sim_completion_ts (the k that would +physically finish first in real), advancing the virtual clock to the k-th +smallest -- immune to physical arrival jitter -- via _sync_sim_recv_first_k. +The per-update U6 visibility lag on a strict barrier (lag_i = max_completion - +completion_i) is computed by _barrier_anchored_lags. These tests drive both +directly with the availability gate OFF (byte-identical to no-avail). +""" + +from datetime import datetime, timedelta + +from flame.mode.horizontal.syncfl.top_aggregator import TopAggregator as _SyncBase +from flame.mode.message import MessageType +from flame.sim.virtual_clock import SimReorderBuffer, VirtualClock +from flame.selector.properties import ( + PROP_CLIENT_TASK_TRAIN_DURATION, + PROP_SIM_SEND_TS, +) + + +class _FakeBarrierChannel: + """Drains every selected end in one recv_fifo pass (as the barrier does), + yielding each end's single grad message out of sct order to prove the + commit set is chosen by sct, not arrival.""" + + def __init__(self): + self._msgs = {} # end -> message dict + self._delivered = set() + self._end_props = {} # (end, prop) -> value + + def add_msg(self, end, sct, dur=None): + m = {MessageType.SIM_COMPLETION_TS: sct} + if dur is not None: + m[MessageType.SIM_CLIENT_TASK_TRAIN_DURATION_S] = dur + self._msgs[end] = m + + def ends(self, state=None): + return list(self._msgs) + + def has(self, end): + return end in self._msgs + + def recv_fifo(self, end_ids, first_k=None, timeout=None): + for e in end_ids: + if e in self._delivered or e not in self._msgs: + continue + self._delivered.add(e) + yield self._msgs[e], (e, datetime.now()) + + def get_end_property(self, end, prop): + return self._end_props.get((end, prop)) + + def set_end_property(self, end, prop, value): + self._end_props[(end, prop)] = value + + +class _FakeBarrierAgg: + """Binds the real sync-barrier sim methods onto a minimal stand-in with the + availability gate OFF (pending_withheld / trainer_event_dict unset).""" + + _sync_sim_recv_first_k = _SyncBase._sync_sim_recv_first_k + _advance_sim_clock = _SyncBase._advance_sim_clock + _sim_recv_grace_s = _SyncBase._sim_recv_grace_s + _note_sim_fill = _SyncBase._note_sim_fill + _sim_reinject_ready_withheld = _SyncBase._sim_reinject_ready_withheld + _sim_withhold_if_unavail = _SyncBase._sim_withhold_if_unavail + _barrier_anchored_lags = staticmethod(_SyncBase._barrier_anchored_lags) + SIM_RECV_GRACE_FLOOR_S = 0.0 + SIM_RECV_GRACE_FACTOR = 0.0 + + def __init__(self): + self._round = 1 + self._vclock = VirtualClock() + self._sim_buffer = SimReorderBuffer() + self._sim_fill_ema = 0.0 + + +class TestSyncFirstKSmallestSct: + def test_commits_the_k_smallest_sct_regardless_of_arrival(self): + agg = _FakeBarrierAgg() + ch = _FakeBarrierChannel() + # arrival order A,B,C,D but modeled completion 40,10,30,20. + ch.add_msg("A", sct=40.0, dur=40.0) + ch.add_msg("B", sct=10.0, dur=10.0) + ch.add_msg("C", sct=30.0, dur=30.0) + ch.add_msg("D", sct=20.0, dur=20.0) + + committed = agg._sync_sim_recv_first_k(ch, ch.ends(), first_k=2) + + scts = [m[MessageType.SIM_COMPLETION_TS] for m, _md in committed] + assert scts == [10.0, 20.0] # the two SMALLEST, ascending + + def test_vclock_advances_to_kth_smallest_no_past_dating(self): + agg = _FakeBarrierAgg() + ch = _FakeBarrierChannel() + for e, s in [("A", 40.0), ("B", 10.0), ("C", 30.0), ("D", 20.0)]: + ch.add_msg(e, sct=s, dur=s) + + agg._sync_sim_recv_first_k(ch, ch.ends(), first_k=2) + + # advanced to the k-th (2nd) smallest = 20.0, never past-dated below it. + assert agg._vclock.now == 20.0 + + def test_stamps_client_task_train_duration_for_committed(self): + agg = _FakeBarrierAgg() + ch = _FakeBarrierChannel() + ch.add_msg("B", sct=10.0, dur=10.0) + ch.add_msg("A", sct=40.0, dur=40.0) + + agg._sync_sim_recv_first_k(ch, ch.ends(), first_k=1) + + # B committed -> its intrinsic train duration is exposed to the selector. + assert ch.get_end_property("B", PROP_CLIENT_TASK_TRAIN_DURATION) == ( + timedelta(seconds=10.0) + ) + # A dropped (past the first_k quota) -> not stamped. + assert ch.get_end_property("A", PROP_CLIENT_TASK_TRAIN_DURATION) is None + + def test_falls_back_to_sct_minus_send_when_no_duration_field(self): + agg = _FakeBarrierAgg() + ch = _FakeBarrierChannel() + ch.add_msg("B", sct=10.0) # no SIM_CLIENT_TASK_TRAIN_DURATION_S + ch.set_end_property("B", PROP_SIM_SEND_TS, 3.0) + + agg._sync_sim_recv_first_k(ch, ch.ends(), first_k=1) + + assert ch.get_end_property("B", PROP_CLIENT_TASK_TRAIN_DURATION) == ( + timedelta(seconds=7.0) # sct(10) - sim_send(3) + ) + + +class TestBarrierAnchoredLags: + def test_lag_is_max_completion_minus_each(self): + lags = _SyncBase._barrier_anchored_lags([2.0, 5.0, 3.0]) + assert lags == [3.0, 0.0, 2.0] # barrier = 5.0 + + def test_none_safe_for_missing_completions(self): + lags = _SyncBase._barrier_anchored_lags([2.0, None, 5.0]) + assert lags == [3.0, None, 0.0] + + def test_all_none_yields_all_none(self): + assert _SyncBase._barrier_anchored_lags([None, None]) == [None, None] + + def test_empty_is_empty(self): + assert _SyncBase._barrier_anchored_lags([]) == [] diff --git a/lib/python/tests/mode/test_fwdllm_suppress_redundant_weights.py b/lib/python/tests/mode/test_fwdllm_suppress_redundant_weights.py new file mode 100644 index 000000000..6255002b6 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_suppress_redundant_weights.py @@ -0,0 +1,139 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Opt-1: intra-databin weight-resend suppression. + +Within a data-bin the model_version is constant and the full WEIGHTS+GRAD_POOL +payload is byte-identical across iterations, yet the aggregator re-sent it to the +same trainers every iteration because the WEIGHTS-vs-VAR=bad guard keyed off the +return-driven `_trainer_last_model_version` map, which never marks an +actively-training trainer current -> `is_stale` stays True -> full re-send. + +`_should_send_full_weights` is the SHARED decision used by both the sync and the +async distribute loops (their parity is the regression guard). These tests pin: +(1) flag OFF is byte-identical to the legacy `var_good_enough or is_stale` rule; +(2) flag ON downgrades a within-cycle repeat to VAR=bad and never a first send; +(3) a commit (`var_good_enough`) always ships weights; +(4) the full distribute-loop contract: exactly ONE weights send per trainer per + data-bin, reverting to weights after the cycle set is cleared on a commit. +""" + +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator + + +class _DistribAgg: + """Minimal stand-in exposing only what `_should_send_full_weights` reads.""" + + def __init__(self, suppress, var_good, sent=None): + self._suppress_redundant_weights = suppress + self.var_good_enough = var_good + self._weights_sent_this_cycle = set(sent or []) + + decide = TopAggregator._should_send_full_weights + + +# --- (1) flag OFF == legacy (var_good_enough or is_stale) ------------------ + +def test_flag_off_is_legacy(): + off_varbad = _DistribAgg(suppress=False, var_good=False) + assert off_varbad.decide("a", is_stale=True) is True # stale -> weights + assert off_varbad.decide("a", is_stale=False) is False # current -> var_bad + off_commit = _DistribAgg(suppress=False, var_good=True) + assert off_commit.decide("a", is_stale=False) is True # commit -> weights + assert off_commit.decide("a", is_stale=True) is True + + +def test_flag_off_ignores_sent_set(): + # Even if an end is in the set, flag OFF must not suppress (byte-identical). + off = _DistribAgg(suppress=False, var_good=False, sent={"a"}) + assert off.decide("a", is_stale=True) is True + + +# --- (2)/(3) flag ON semantics --------------------------------------------- + +def test_commit_branch_first_send_weights_repeat_suppressed(): + # The dominant redundancy: repeated distribute calls in the SAME data-bin all + # take the var_good_enough=True branch. First send per end -> weights; a repeat + # within the cycle -> VAR=bad (the set is checked BEFORE the var_good branch). + on = _DistribAgg(suppress=True, var_good=True, sent=set()) + assert on.decide("a", is_stale=False) is True # first send -> weights + on._weights_sent_this_cycle.add("a") # caller records it + assert on.decide("a", is_stale=False) is False # repeat same cycle -> VAR=bad + assert on.decide("a", is_stale=True) is False # still suppressed if stale too + assert on.decide("b", is_stale=False) is True # a NEW end still gets weights + + +def test_flag_off_commit_always_weights(): + # With the flag OFF the set is never consulted -> legacy commit-branch behavior. + off = _DistribAgg(suppress=False, var_good=True, sent={"a"}) + assert off.decide("a", is_stale=False) is True + assert off.decide("a", is_stale=True) is True + + +def test_first_send_is_weights_repeat_is_varbad(): + on = _DistribAgg(suppress=True, var_good=False, sent=set()) + # first time this cycle, stale -> weights (nothing in the set yet) + assert on.decide("a", is_stale=True) is True + on._weights_sent_this_cycle.add("a") # caller records the send + # re-dispatched same cycle, still stale by the return-map -> VAR=bad now + assert on.decide("a", is_stale=True) is False + # a different, not-yet-sent end still gets weights + assert on.decide("b", is_stale=True) is True + + +def test_current_by_return_map_is_varbad(): + # not stale and not a commit -> var_bad regardless of flag + on = _DistribAgg(suppress=True, var_good=False, sent=set()) + assert on.decide("a", is_stale=False) is False + + +# --- (4) full distribute-loop contract: exactly one weights send / bin ----- + +def _run_cycle(agg, dispatches): + """Mimic the distribute loop's set-management for a sequence of dispatches + within ONE data-bin. Returns the list of 'weights'/'var_bad' decisions.""" + kinds = [] + for end, is_stale in dispatches: + if agg.decide(end, is_stale): + kinds.append("weights") + if agg._suppress_redundant_weights: + agg._weights_sent_this_cycle.add(end) + else: + kinds.append("var_bad") + return kinds + + +def test_exactly_one_weights_per_trainer_per_databin(): + # fwdllm pattern: same K=3 trainers, 4 iterations, always is_stale by the map. + agg = _DistribAgg(suppress=True, var_good=False, sent=set()) + dispatches = [(e, True) for _ in range(4) for e in ("a", "b", "c")] + kinds = _run_cycle(agg, dispatches) + # first 3 (one per trainer) are weights; the remaining 9 are var_bad. + assert kinds[:3] == ["weights", "weights", "weights"] + assert set(kinds[3:]) == {"var_bad"} + assert kinds.count("weights") == 3 # exactly once per unique trainer + # legacy (flag off) would have sent 12 weights -> 9 redundant. + off = _DistribAgg(suppress=False, var_good=False, sent=set()) + assert _run_cycle(off, dispatches).count("weights") == 12 + + +def test_exactly_one_weights_per_trainer_var_good_branch(): + # The fwdllm pattern: sync loop distribute>>aggregate re-runs ~10x per data-bin, + # each distribute taking the var_good_enough=True branch for the SAME K=3 + # trainers. Legacy shipped weights every time (10x redundant); the fix ships + # each trainer the model exactly once per data-bin. + agg = _DistribAgg(suppress=True, var_good=True, sent=set()) + dispatches = [(e, False) for _ in range(10) for e in ("a", "b", "c")] + kinds = _run_cycle(agg, dispatches) + assert kinds[:3] == ["weights", "weights", "weights"] + assert set(kinds[3:]) == {"var_bad"} + assert kinds.count("weights") == 3 + off = _DistribAgg(suppress=False, var_good=True, sent=set()) + assert _run_cycle(off, dispatches).count("weights") == 30 # legacy: all redundant + + +def test_weights_resume_after_cycle_cleared_on_commit(): + agg = _DistribAgg(suppress=True, var_good=False, sent=set()) + _run_cycle(agg, [("a", True)]) + assert agg.decide("a", is_stale=True) is False # suppressed mid-bin + agg._weights_sent_this_cycle.clear() # model_version advanced + assert agg.decide("a", is_stale=True) is True # new bin -> weights again diff --git a/lib/python/tests/mode/test_fwdllm_sync_scarcity_wait.py b/lib/python/tests/mode/test_fwdllm_sync_scarcity_wait.py new file mode 100644 index 000000000..e501a308b --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_sync_scarcity_wait.py @@ -0,0 +1,182 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Real-mode sync-barrier liveness under availability scarcity. + +`_await_dispatchable_under_scarcity` keeps the cohort == `agg_goal` and +sleep-to-next-avail (parity-faithful, matches the sim vclock-jump) instead of +hot-spin-dispatching when a trace keeps the eligible pool below `agg_goal`. + +Invariants under test: +- no-op on the sim path and when availability tracking is off (byte-identical), +- no-op during startup join-lag (fewer than `agg_goal` trainers joined), +- returns immediately when an available un-contributed trainer exists, +- waits (bounded) while the pool is scarce, then proceeds once it recovers, +- self-terminates via `_check_early_stop_conditions` at the wall budget. +""" + +from types import SimpleNamespace + +import pytest + +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator + + +class _FakeAggregator: + """Minimal stand-in exposing only what the scarcity gate touches.""" + + def __init__( + self, + *, + simulated=False, + trainer_event_dict=None, + all_trainers=None, + agg_goal=10, + agg_goal_cnt=0, + per_agg_trainer_list=None, + unavail_sequence=None, + ): + self.simulated = simulated + self.trainer_event_dict = trainer_event_dict + self.all_trainers = set(all_trainers or []) + self._agg_goal = agg_goal + self._agg_goal_cnt = agg_goal_cnt + self._per_agg_trainer_list = list(per_agg_trainer_list or []) + self._round = 0 + self.data_id = 0 + self._work_done = False + self.config = SimpleNamespace(hyperparameters=SimpleNamespace(scarcity_poll_s=0.0)) + # Each poll pops the next unavailable-set from this sequence (the last + # one repeats), letting a test model availability recovering over time. + self._unavail_sequence = list(unavail_sequence or [[]]) + self.slept = 0 + self.stop_checks = 0 + + def get_curr_unavail_trainers(self): + if len(self._unavail_sequence) > 1: + return self._unavail_sequence.pop(0) + return self._unavail_sequence[0] + + def _check_early_stop_conditions(self): + self.stop_checks += 1 + + _await_dispatchable_under_scarcity = ( + TopAggregator._await_dispatchable_under_scarcity + ) + + +@pytest.fixture(autouse=True) +def _no_real_sleep(monkeypatch): + import flame.mode.horizontal.syncfl.fwdllm_aggregator as mod + + def _count_sleep(_s): + _count_sleep.calls += 1 + + _count_sleep.calls = 0 + monkeypatch.setattr(mod.time, "sleep", _count_sleep) + return _count_sleep + + +class TestScarcityWaitNoOp: + def test_sim_path_never_waits(self, _no_real_sleep): + agg = _FakeAggregator( + simulated=True, + trainer_event_dict={"1": []}, + all_trainers=[str(i) for i in range(10)], + unavail_sequence=[[str(i) for i in range(10)]], # all unavail + ) + agg._await_dispatchable_under_scarcity("train") + assert _no_real_sleep.calls == 0 + + def test_tracking_off_never_waits(self, _no_real_sleep): + agg = _FakeAggregator( + trainer_event_dict=None, # availability tracking disabled + all_trainers=[str(i) for i in range(10)], + unavail_sequence=[[str(i) for i in range(10)]], + ) + agg._await_dispatchable_under_scarcity("train") + assert _no_real_sleep.calls == 0 + + def test_join_lag_does_not_wait(self, _no_real_sleep): + # Only 3 of the 10 have joined -> startup, not scarcity: do not block. + agg = _FakeAggregator( + trainer_event_dict={"1": []}, + all_trainers=["0", "1", "2"], + agg_goal=10, + unavail_sequence=[["0", "1", "2"]], + ) + agg._await_dispatchable_under_scarcity("train") + assert _no_real_sleep.calls == 0 + + def test_available_trainer_returns_without_wait(self, _no_real_sleep): + agg = _FakeAggregator( + trainer_event_dict={"1": []}, + all_trainers=[str(i) for i in range(10)], + unavail_sequence=[["0", "1", "2"]], # 7 available, none contributed + ) + agg._await_dispatchable_under_scarcity("train") + assert _no_real_sleep.calls == 0 + + def test_barrier_already_met_returns(self, _no_real_sleep): + agg = _FakeAggregator( + trainer_event_dict={"1": []}, + all_trainers=[str(i) for i in range(10)], + agg_goal=3, + agg_goal_cnt=3, # cohort complete + unavail_sequence=[[str(i) for i in range(10)]], + ) + agg._await_dispatchable_under_scarcity("train") + assert _no_real_sleep.calls == 0 + + +class TestScarcityWaitEngages: + def test_waits_then_proceeds_when_availability_recovers(self, _no_real_sleep): + allt = [str(i) for i in range(10)] + # Scarce for two polls (all unavail), then trainer "5" comes back. + agg = _FakeAggregator( + trainer_event_dict={"1": []}, + all_trainers=allt, + unavail_sequence=[list(allt), list(allt), [x for x in allt if x != "5"]], + ) + agg._await_dispatchable_under_scarcity("train") + assert _no_real_sleep.calls == 2 # slept exactly across the two scarce polls + assert agg.stop_checks == 2 # re-checked early-stop each scarce poll + + def test_already_contributed_pool_still_scarce(self, _no_real_sleep): + # The available trainers have ALL already contributed this cycle, and the + # one still-uncontributed trainer ("5") is unavailable -> no dispatchable + # end -> must wait, not spin-redispatch the already-used ones. On the next + # poll "5" recovers -> dispatchable -> proceed. + allt = [str(i) for i in range(10)] + used = [x for x in allt if x != "5"] # the 9 available already contributed + agg = _FakeAggregator( + trainer_event_dict={"1": []}, + all_trainers=allt, + agg_goal=10, + agg_goal_cnt=9, + per_agg_trainer_list=used, + unavail_sequence=[["5"], []], # "5" unavail, then recovers + ) + agg._await_dispatchable_under_scarcity("train") + assert _no_real_sleep.calls == 1 # waited one poll, then "5" recovered + + def test_self_terminates_at_budget(self, _no_real_sleep): + allt = [str(i) for i in range(10)] + + agg = _FakeAggregator( + trainer_event_dict={"1": []}, + all_trainers=allt, + unavail_sequence=[list(allt)], # permanently scarce + ) + + # Model the wall budget firing after 3 polls. + orig = agg._check_early_stop_conditions + + def _stop_after_3(): + orig() + if agg.stop_checks >= 3: + agg._work_done = True + + agg._check_early_stop_conditions = _stop_after_3 + agg._await_dispatchable_under_scarcity("train") + assert agg._work_done is True + assert _no_real_sleep.calls == 3 # bounded, then stopped diff --git a/lib/python/tests/mode/test_fwdllm_trainer_phase_timing.py b/lib/python/tests/mode/test_fwdllm_trainer_phase_timing.py new file mode 100644 index 000000000..ba3e44255 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_trainer_phase_timing.py @@ -0,0 +1,183 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Trainer per-phase wall timing. + +The fwdllm trainer emits a per-phase breakdown feeding the 8 phase rungs +(mqtt_fetch_s / weights_to_ram_s / weights_to_gpu_s / pre_train_s / +gpu_compute_s / post_train_s / training_budget_s / trainer_phase). This covers: +- the `_phase` context manager accumulates into `_phase_times` (fwdllm_trainer), +- `FedSgdTrainer.train_with_data_id` drains those + the pre/gpu/post/budget/ + phase terms into the emitted `trainer_round` event. +""" + +import json +import os +import sys + +import pytest + +from flame import telemetry +from flame.mode.horizontal.syncfl.fwdllm_trainer import Trainer as FwdLLMTrainer + +_EXAMPLE = os.path.join( + os.path.dirname(__file__), + "..", "..", "examples", "fwdllm", "trainer", "forward_training", +) +sys.path.insert(0, os.path.abspath(_EXAMPLE)) +import FedSgdTrainer as _fedsgd_mod # noqa: E402 + +FedSgdTrainer = _fedsgd_mod.FedSGDTrainer + +_PHASE_KEYS = [ + "mqtt_fetch_s", + "weights_to_ram_s", + "weights_to_gpu_s", + "pre_train_s", + "gpu_compute_s", + "post_train_s", + "training_budget_s", + "trainer_phase", +] + + +class _PhaseHost: + """Minimal host for the (abstract) fwdllm Trainer's _phase method.""" + + _phase = FwdLLMTrainer._phase + + def __init__(self): + self._phase_times = {} + + +class TestPhaseContextManager: + def test_accumulates(self): + t = _PhaseHost() + with t._phase("gpu_compute_s"): + pass + with t._phase("gpu_compute_s"): + pass + assert "gpu_compute_s" in t._phase_times + # two entries under one name accumulate, not overwrite + assert t._phase_times["gpu_compute_s"] >= 0.0 + assert isinstance(t._phase_times["gpu_compute_s"], float) + + +class _AvlState: + value = "AVL_TRAIN" + + +class _FakeFedSgd: + """Binds the real train_with_data_id onto a stand-in that stubs only the + heavy compute (perform_training / delay / availability); the phase-timing + + telemetry-emit logic under test runs for real.""" + + train_with_data_id = FedSgdTrainer.train_with_data_id + # train_with_data_id folds the B2 straggler into the sct (#6); no config -> + # spread 0 -> offset 0 -> phase/duration values preserved. + _sim_straggler_offset_s = FedSgdTrainer._sim_straggler_offset_s + + def __init__(self, phase_times=None, delay_s=1.5): + self._round = 7 + self.data_id = 3 + self.iteration_per_data_id = 2 + self._model_version = 5 + self.trainer_id = "t1" + self.simulated = False + self.abort_training = False + self.avl_state = _AvlState() + self.dataset_size = 128 + self._stat_utility = 0.9 + self._sim_send_ts = None + self.sim_completion_leg_s = 0.0 + self._sim_completion_ts = None + self._sim_round_duration_s = None + # Populated by _fetch_weights in a live run; pre-seeded here. + self._phase_times = dict(phase_times or {}) + self._delay_s = delay_s + + def _check_availability(self): + return True + + def _perform_training(self): + pass + + def _emulate_training_delay(self, gpu_time_s=0.0): + # remainder-wait signature: (modeled_delay, remaining, overran). + return self._delay_s, max(0.0, self._delay_s - gpu_time_s), False + + +class TestTrainWithDataIdEmitsPhases: + def test_phase_fields_present(self, tmp_path): + telemetry.configure(role="trainer", run_dir=str(tmp_path)) + try: + t = _FakeFedSgd( + phase_times={ + "mqtt_fetch_s": 0.4, + "weights_to_ram_s": 0.2, + "weights_to_gpu_s": 0.1, + } + ) + t.train_with_data_id() + + lines = (tmp_path / "trainer.jsonl").read_text().splitlines() + events = [json.loads(l) for l in lines] + rounds = [e for e in events if e["event"] == "trainer_round"] + assert len(rounds) == 1 + ev = rounds[0] + for k in _PHASE_KEYS: + assert k in ev, f"missing phase field {k}" + # the three _fetch_weights phases flowed through verbatim + assert ev["mqtt_fetch_s"] == 0.4 + assert ev["weights_to_ram_s"] == 0.2 + assert ev["weights_to_gpu_s"] == 0.1 + # training_budget_s is the modeled additive delay + assert ev["training_budget_s"] == 1.5 + # post_train excludes the delay (stamped after it) -> pure post-proc. + assert ev["post_train_s"] >= 0.0 and ev["post_train_s"] < 0.5 + # trainer_phase encodes round/data_id/iteration identity + assert ev["trainer_phase"] == "7/3/2" + # pre/post are real non-negative wall slivers + assert ev["pre_train_s"] >= 0.0 + assert ev["post_train_s"] >= 0.0 + finally: + telemetry.shutdown() + + def test_straggler_in_sct_not_in_training_budget(self, tmp_path): + """#6: the B2 straggler spread is folded into the sct + (sim_round_duration_s) but NOT into the emitted training_budget_s -- so + training_budget stays a mode-invariant INPUT (T2 passes) while the sync + barrier still gets its per-trainer dispersion.""" + from types import SimpleNamespace + telemetry.configure(role="trainer", run_dir=str(tmp_path)) + try: + t = _FakeFedSgd(delay_s=1.5) + t.simulated = True + t._sim_send_ts = 100.0 + t.config = SimpleNamespace( + hyperparameters=SimpleNamespace(sim_straggler_spread_s=0.9)) + offset = t._sim_straggler_offset_s() + assert offset > 0.0, "t1 should have a non-zero stable offset" + t.train_with_data_id() + + ev = [json.loads(l) for l in + (tmp_path / "trainer.jsonl").read_text().splitlines() + if json.loads(l)["event"] == "trainer_round"][0] + # training_budget = the BASE delay; straggler EXCLUDED. + assert ev["training_budget_s"] == 1.5 + # sct carries the straggler: duration - budget == gpu(~0) + offset. + assert (ev["sim_round_duration_s"] - ev["training_budget_s"]) == \ + pytest.approx(offset, abs=0.05) + finally: + telemetry.shutdown() + + def test_aborted_round_emits_nothing(self, tmp_path): + telemetry.configure(role="trainer", run_dir=str(tmp_path)) + try: + t = _FakeFedSgd() + t.abort_training = True + t.train_with_data_id() + assert not (tmp_path / "trainer.jsonl").exists() or not ( + tmp_path / "trainer.jsonl" + ).read_text().strip() + finally: + telemetry.shutdown() diff --git a/lib/python/tests/mode/test_fwdllm_trainer_sim_duration.py b/lib/python/tests/mode/test_fwdllm_trainer_sim_duration.py index 9e5863e4f..ffa497d78 100644 --- a/lib/python/tests/mode/test_fwdllm_trainer_sim_duration.py +++ b/lib/python/tests/mode/test_fwdllm_trainer_sim_duration.py @@ -1,16 +1,19 @@ # Copyright 2026 Cisco Systems, Inc. and its affiliates # SPDX-License-Identifier: Apache-2.0 -"""fwdllm's trainer (FedSgdTrainer.py) reported only real_gpu_time_s, not -sim_round_duration_s -- unlike async_cifar10's trainer, which reports total -round wall time (gpu + modeled delay). fwdllm has no budget-vs-actual -contention model (its delay is a flat additive sleep, not a sleep-to-fill- -budget pattern), so only sim_round_duration_s is added here -- NOT -training_budget_s/overran/remaining_time_s, which would need a budget -concept fwdllm doesn't have (see ../../examples/MIGRATING_TO_LAUNCHER.md §9). - -This covers _emulate_training_delay()'s return-value change: it now returns -the seconds actually slept (0.0 if delay emulation is disabled), which the -caller adds to real_gpu_time_s to report sim_round_duration_s. +"""fwdllm's trainer sim-duration / delay model (K-D29 REMAINDER-WAIT, replacing +the earlier flat-additive model, K-D2). + +The modeled mobile device takes ``_delay_s = training_delay_s/factor/speedup``. +On our GPU the forward pass takes ``gpu_time_s`` (SHOULD be << device time). So: + - REAL mode sleeps only the remainder ``max(0, _delay_s - gpu)`` -> real wall + ≈ _delay_s, GPU hidden inside it. + - SIM mode skips the sleep; the sct round duration is ``max(gpu, _delay_s)`` + (NOT gpu + _delay_s). Per-trainer registry delays supply the completion + SPREAD -> update order = delay order = deterministic + identical real↔sim. + - OVERRUN: gpu > _delay_s => emulation unfaithful; flagged (remaining==0). + +``_emulate_training_delay(gpu_time_s)`` returns +``(modeled_delay_s, remaining_s, overran)``. """ import os @@ -24,6 +27,7 @@ ), ) +import FedSgdTrainer as _fst_module # noqa: E402 from FedSgdTrainer import FedSGDTrainer # noqa: E402 @@ -34,34 +38,159 @@ class _FakeTrainer: _emulate_training_delay = FedSGDTrainer._emulate_training_delay def __init__(self, training_delay_enabled, training_delay_s=0.0, - training_delay_factor=1.0, speedup_factor=1.0): + training_delay_factor=1.0, speedup_factor=1.0, simulated=False): self.training_delay_enabled = training_delay_enabled self.training_delay_s = training_delay_s self.training_delay_factor = training_delay_factor self.speedup_factor = speedup_factor + self.simulated = simulated self.trainer_id = "t1" + self.data_id = 3 + self.iteration_per_data_id = 0 -class TestEmulateTrainingDelayReturnsSleptSeconds: - def test_returns_zero_when_disabled(self): +class TestEmulateTrainingDelayRemainderWait: + def test_returns_zero_tuple_when_disabled(self): t = _FakeTrainer(training_delay_enabled="False", training_delay_s=10.0) - assert t._emulate_training_delay() == 0.0 - - def test_returns_computed_delay_when_enabled(self): - t = _FakeTrainer( - training_delay_enabled="True", training_delay_s=3.0, - training_delay_factor=1.0, speedup_factor=1.0, - ) - assert t._emulate_training_delay() == 3.0 - - def test_speedup_factor_scales_the_returned_delay(self): - """The returned value must match what was actually slept (eval_delay - / speedup_factor), not the unscaled eval_delay -- otherwise - sim_round_duration_s would overstate the real wall time under a - speedup.""" - t = _FakeTrainer( - training_delay_enabled="True", training_delay_s=10.0, - training_delay_factor=2.0, speedup_factor=5.0, - ) - # eval_delay = 10.0 / 2.0 = 5.0; slept = 5.0 / 5.0 = 1.0 - assert t._emulate_training_delay() == 1.0 + assert t._emulate_training_delay(0.2) == (0.0, 0.0, False) + + def test_modeled_delay_and_remainder_when_gpu_below_budget(self): + # delay = 4.0/2.0/1.0 = 2.0; gpu = 0.5 -> remaining = 1.5, no overrun. + t = _FakeTrainer(training_delay_enabled="True", training_delay_s=4.0, + training_delay_factor=2.0, speedup_factor=1.0) + modeled, remaining, overran = t._emulate_training_delay(0.5) + assert modeled == 2.0 and remaining == 1.5 and overran is False + + def test_speedup_factor_scales_the_modeled_delay(self): + # eval_delay = 10/2 = 5; modeled = 5/5 = 1.0; gpu 0.25 -> remaining 0.75. + t = _FakeTrainer(training_delay_enabled="True", training_delay_s=10.0, + training_delay_factor=2.0, speedup_factor=5.0) + modeled, remaining, overran = t._emulate_training_delay(0.25) + assert modeled == 1.0 and remaining == 0.75 and overran is False + + def test_overrun_when_gpu_exceeds_budget(self): + # gpu 3.0 > budget 2.0 -> overran, remaining clamped to 0. + t = _FakeTrainer(training_delay_enabled="True", training_delay_s=4.0, + training_delay_factor=2.0, speedup_factor=1.0) + modeled, remaining, overran = t._emulate_training_delay(3.0) + assert modeled == 2.0 and remaining == 0.0 and overran is True + + +class TestSleepOnlyTheRemainderInRealMode: + def test_real_mode_sleeps_the_remainder(self, monkeypatch): + slept = [] + monkeypatch.setattr(_fst_module.time, "sleep", lambda s: slept.append(s)) + t = _FakeTrainer(training_delay_enabled="True", training_delay_s=4.0, + training_delay_factor=2.0, speedup_factor=1.0, + simulated=False) + modeled, remaining, _ = t._emulate_training_delay(0.5) + assert modeled == 2.0 and remaining == 1.5 + assert slept == [1.5] # ONLY the remainder, not the full delay + + def test_real_mode_overrun_sleeps_nothing(self, monkeypatch): + slept = [] + monkeypatch.setattr(_fst_module.time, "sleep", lambda s: slept.append(s)) + t = _FakeTrainer(training_delay_enabled="True", training_delay_s=2.0, + training_delay_factor=1.0, speedup_factor=1.0, + simulated=False) + t._emulate_training_delay(5.0) # gpu > budget + assert slept == [] # nothing to sleep; overran + + def test_sim_mode_does_not_sleep(self, monkeypatch): + slept = [] + monkeypatch.setattr(_fst_module.time, "sleep", lambda s: slept.append(s)) + t = _FakeTrainer(training_delay_enabled="True", training_delay_s=4.0, + training_delay_factor=2.0, speedup_factor=1.0, + simulated=True) + modeled, remaining, _ = t._emulate_training_delay(0.5) + assert modeled == 2.0 and remaining == 1.5 # same modeled math as real + assert slept == [] # but NOTHING slept in sim + + def test_disabled_sim_no_sleep(self, monkeypatch): + slept = [] + monkeypatch.setattr(_fst_module.time, "sleep", lambda s: slept.append(s)) + t = _FakeTrainer(training_delay_enabled="False", training_delay_s=9.0, + simulated=True) + assert t._emulate_training_delay(0.1) == (0.0, 0.0, False) + assert slept == [] + + +class _FakeTime: + """Scripted time source so the sct arithmetic is deterministic. Rebound only + onto FedSgdTrainer's `time` name (not the shared module).""" + + def __init__(self, ticks): + self._ticks = list(ticks) + self._i = 0 + + def time(self): + v = self._ticks[self._i] + self._i = min(self._i + 1, len(self._ticks) - 1) + return v + + def sleep(self, _s): # must never be called on the sim path + raise AssertionError("time.sleep called on the simulated path") + + +class _StampTrainer: + """Binds the real train_with_data_id onto a minimal stand-in, stubbing the + heavy compute so only the sim-stamp arithmetic is exercised. No config -> + straggler spread 0 -> offset 0.""" + + train_with_data_id = FedSGDTrainer.train_with_data_id + _sim_straggler_offset_s = FedSGDTrainer._sim_straggler_offset_s + + def __init__(self, sim_send_ts, delay_d, leg_s=0.0): + self.simulated = True + self.abort_training = False + self.trainer_id = "t1" + self._round = 7 + self.data_id = 3 + self.iteration_per_data_id = 0 + self._sim_send_ts = sim_send_ts + self.sim_completion_leg_s = leg_s + self._sim_completion_ts = None + self._sim_round_duration_s = None + self._delay_d = delay_d + + def _check_availability(self): + return True + + def _perform_training(self): + pass # no GPU work; wall time is scripted via _FakeTime + + def _emulate_training_delay(self, gpu_time_s=0.0): + # modeled D, remainder (irrelevant in sim), no overrun + return self._delay_d, max(0.0, self._delay_d - gpu_time_s), False + + +class TestSimCompletionStampIsMaxGpuDelay: + """K-D29: the sct the aggregator orders by is ``max(gpu, D)`` (the mobile + device wall, GPU hidden inside), NOT the old additive gpu + D, and + _sim_completion_ts = _sim_send_ts + max(gpu, D) + leg.""" + + def test_max_round_duration_and_completion_ts(self, monkeypatch): + monkeypatch.setattr(_fst_module.telemetry, "is_enabled", lambda: False) + # real_gpu = 0.5s; D = 2.0 -> max(0.5, 2.0) = 2.0 (additive would be 2.5). + monkeypatch.setattr(_fst_module, "time", _FakeTime([100.0, 100.0, 100.5])) + t = _StampTrainer(sim_send_ts=10.0, delay_d=2.0, leg_s=0.0) + t.train_with_data_id() + assert t._sim_round_duration_s == 2.0 + assert t._sim_completion_ts == 12.0 # 10.0 + 2.0 + 0.0 + + def test_gpu_dominates_when_over_budget(self, monkeypatch): + monkeypatch.setattr(_fst_module.telemetry, "is_enabled", lambda: False) + # real_gpu = 3.0s; D = 2.0 -> max = 3.0 (the overrun case). + monkeypatch.setattr(_fst_module, "time", _FakeTime([100.0, 100.0, 103.0])) + t = _StampTrainer(sim_send_ts=10.0, delay_d=2.0, leg_s=0.0) + t.train_with_data_id() + assert t._sim_round_duration_s == 3.0 + assert t._sim_completion_ts == 13.0 + + def test_completion_ts_includes_leg(self, monkeypatch): + monkeypatch.setattr(_fst_module.telemetry, "is_enabled", lambda: False) + monkeypatch.setattr(_fst_module, "time", _FakeTime([100.0, 100.0, 100.5])) + t = _StampTrainer(sim_send_ts=10.0, delay_d=2.0, leg_s=1.5) + t.train_with_data_id() + assert t._sim_round_duration_s == 2.0 + assert t._sim_completion_ts == 13.5 # 10.0 + 2.0 + 1.5 diff --git a/lib/python/tests/mode/test_fwdllm_trainer_task_recv.py b/lib/python/tests/mode/test_fwdllm_trainer_task_recv.py new file mode 100644 index 000000000..076ca8491 --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_trainer_task_recv.py @@ -0,0 +1,104 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Stage D1 (§H #8): the fwdllm trainer overrode _fetch_weights and dropped the +base trainer's task_recv emission, so field_coverage's sim_send_ts INV rung + K6 +had no field to read (null/null both modes). Restored here -- task_recv carries +sim_send_ts (the aggregator's dispatch vclock), null in real mode. +""" + +import json + +from flame import telemetry +from flame.mode.horizontal.syncfl.fwdllm_trainer import Trainer +from flame.mode.message import MessageType + + +class _FakeSelector: + def __init__(self): + self.ordered_updates_recv_ends = [] + + +class _FakeChannel: + def __init__(self, msg): + self._msg = msg + self._selector = _FakeSelector() + + def await_join(self): + pass + + def one_end(self, state): + return "end_1" + + def recv(self, end_id): + return self._msg, None + + def cleanup_recvd_ends(self): + pass + + +class _FakeChannelManager: + def __init__(self, channel): + self._channel = channel + + def get_by_tag(self, tag): + return self._channel + + +class _FakeTrainer: + _fetch_weights = Trainer._fetch_weights + + def __init__(self, channel, time_mode="real"): + self.cm = _FakeChannelManager(channel) + self.trainer_id = "trainer_1" + self.fetch_success = False + self._work_done = False + self.data_id = None + self.iteration_per_data_id = None + self._round = 1 + self._model_version = 0 + self.time_mode = time_mode + + +def _task_recv_events(tmp_path): + lines = (tmp_path / "trainer.jsonl").read_text().splitlines() + return [json.loads(l) for l in lines if json.loads(l)["event"] == "task_recv"] + + +class TestTaskRecvEmission: + def test_sim_send_ts_emitted_in_sim(self, tmp_path): + telemetry.configure(role="trainer", run_dir=str(tmp_path)) + try: + channel = _FakeChannel( + {MessageType.ROUND: 3, MessageType.SIM_SEND_TS: 42.5} + ) + t = _FakeTrainer(channel, time_mode="simulated") + t._fetch_weights("fetch") + + evs = _task_recv_events(tmp_path) + assert len(evs) == 1 + assert evs[0]["sim_send_ts"] == 42.5 + assert evs[0]["time_mode"] == "simulated" + assert evs[0]["trainer_id"] == "trainer_1" + finally: + telemetry.shutdown() + + def test_sim_send_ts_null_in_real(self, tmp_path): + telemetry.configure(role="trainer", run_dir=str(tmp_path)) + try: + # real mode: aggregator stamps no SIM_SEND_TS -> field is null, which + # is what makes real~=sim task_recv directly comparable. + channel = _FakeChannel({MessageType.ROUND: 3}) + t = _FakeTrainer(channel, time_mode="real") + t._fetch_weights("fetch") + + evs = _task_recv_events(tmp_path) + assert len(evs) == 1 + assert evs[0]["sim_send_ts"] is None + finally: + telemetry.shutdown() + + def test_noop_when_telemetry_disabled(self, tmp_path): + assert not telemetry.is_enabled() + channel = _FakeChannel({MessageType.ROUND: 3, MessageType.SIM_SEND_TS: 1.0}) + _FakeTrainer(channel)._fetch_weights("fetch") + assert not (tmp_path / "trainer.jsonl").exists() diff --git a/lib/python/tests/mode/test_fwdllm_var_stopping_policy.py b/lib/python/tests/mode/test_fwdllm_var_stopping_policy.py new file mode 100644 index 000000000..6dda8973c --- /dev/null +++ b/lib/python/tests/mode/test_fwdllm_var_stopping_policy.py @@ -0,0 +1,105 @@ +# Copyright 2026 Cisco Systems, Inc. and its affiliates +# SPDX-License-Identifier: Apache-2.0 +"""Opt-2: variance-plateau stopping policy. + +At the α=1 operating point the achievable variance floor (~0.45) sits ABOVE the +commit gate (0.30), so a data-bin crosses the gate only on a noise dip and grinds +many iterations while the denoised estimate has long plateaued. The 'plateau' +policy commits a bin early once its per-bin variance curve flattens (relative +drop over the last N cycles < rel_delta) while var is still above threshold -- +shipping the denoised estimate instead of a lucky noise sample. + +`_should_force_commit_on_plateau` is the PURE decision (reads only instance +attrs, newest-var-last history in `var_prev_iter_list`), driven from +FedSGDAggregator.aggregate(). These tests pin: +(1) policy off / None / fixed_cap never forces a plateau commit (byte-identical off); +(2) it fires only when the curve has flattened AND var is still above threshold; +(3) it stays quiet while the curve is still dropping fast, or before N+1 samples; +(4) N (patience) and rel_delta (tolerance) sensitivity flips the decision; +(5) degenerate guards (non-positive baseline) return False. +""" + +from flame.mode.horizontal.syncfl.fwdllm_aggregator import TopAggregator + + +class _PlateauAgg: + """Minimal stand-in exposing only what `_should_force_commit_on_plateau` reads.""" + + def __init__(self, policy, hist, N=3, eps=0.10, thr=0.30): + self._var_stopping_policy = policy + self.var_prev_iter_list = list(hist) # newest var last + self._var_plateau_patience = N + self._var_plateau_rel_delta = eps + self.var_threshold = thr + + decide = TopAggregator._should_force_commit_on_plateau + + +# --- (1) policy off/None/fixed_cap => never a plateau commit (byte-identical) --- + +def test_policy_off_never_fires(): + flat = [0.50, 0.49, 0.485, 0.484] # a clearly flattened curve + for policy in (None, "off", "fixed_cap"): + agg = _PlateauAgg(policy=policy, hist=flat, N=3, eps=0.10) + assert agg.decide() is False + + +# --- (2) plateau fires on a flat curve with var still above threshold ---------- + +def test_plateau_fires_on_flat_curve_above_threshold(): + # last-N (N=3) window: 0.50 -> 0.484, rel drop = 0.032 < eps=0.10 ; var>thr. + agg = _PlateauAgg(policy="plateau", hist=[0.9, 0.50, 0.49, 0.485, 0.484], + N=3, eps=0.10, thr=0.30) + assert agg.decide() is True + + +def test_plateau_quiet_while_still_dropping_fast(): + # last-N window: 1.0 -> 0.55, rel drop = 0.45 >= eps=0.10 -> still improving. + agg = _PlateauAgg(policy="plateau", hist=[2.0, 1.0, 0.8, 0.55], + N=3, eps=0.10, thr=0.30) + assert agg.decide() is False + + +def test_plateau_quiet_when_var_already_under_threshold(): + # Curve flat AND var<=thr -> the natural gate commits; the plateau rule must NOT + # claim it (avoids double-attribution / a spurious 'plateau' reason). + agg = _PlateauAgg(policy="plateau", hist=[0.40, 0.30, 0.29, 0.285], + N=3, eps=0.10, thr=0.30) + assert agg.decide() is False + + +# --- (3) insufficient history -------------------------------------------------- + +def test_no_fire_before_n_plus_one_samples(): + # len == N -> not enough to look back N steps (need > N). + agg = _PlateauAgg(policy="plateau", hist=[0.49, 0.485, 0.484], N=3, eps=0.10) + assert agg.decide() is False + # one more sample -> now len == N+1, and it's flat -> fires. + agg2 = _PlateauAgg(policy="plateau", hist=[0.50, 0.49, 0.485, 0.484], + N=3, eps=0.10, thr=0.30) + assert agg2.decide() is True + + +# --- (4) N (patience) and eps (tolerance) sensitivity -------------------------- + +def test_eps_sensitivity_flips_decision(): + # window 0.50 -> 0.46 over N=3, rel drop = 0.08. + hist = [0.7, 0.50, 0.48, 0.47, 0.46] + assert _PlateauAgg("plateau", hist, N=3, eps=0.05).decide() is False # 0.08 !< 0.05 + assert _PlateauAgg("plateau", hist, N=3, eps=0.10).decide() is True # 0.08 < 0.10 + + +def test_patience_window_sensitivity(): + # A curve that fell early then flattened. Short window (N=2) sees only the flat + # tail -> fires; long window (N=4) still spans the early steep drop -> quiet. + hist = [1.0, 0.5, 0.47, 0.465, 0.462] + assert _PlateauAgg("plateau", hist, N=2, eps=0.10).decide() is True + assert _PlateauAgg("plateau", hist, N=4, eps=0.10).decide() is False + + +# --- (5) degenerate guards ----------------------------------------------------- + +def test_nonpositive_baseline_guard(): + # A zero N-steps-back value would divide by zero; guard returns False. + agg = _PlateauAgg(policy="plateau", hist=[0.0, 0.1, 0.1, 0.1], N=3, eps=0.10) + assert agg.decide() is False diff --git a/lib/python/tests/mode/test_parity_checks.py b/lib/python/tests/mode/test_parity_checks.py index 0bfe4ef08..a79cbb100 100644 --- a/lib/python/tests/mode/test_parity_checks.py +++ b/lib/python/tests/mode/test_parity_checks.py @@ -213,6 +213,107 @@ def test_no_vclock_fails(self): assert not r["ok"] and "K10" in r.get("note", "") +class TestIntrinsicSpanAnchor: + """#6: clock-rate rungs anchor REAL on `intrinsic_span_s` (barrier+eval) + instead of raw wall ts (which bundles a transport artifact the sim omits). + With vs without the field must flip the verdict; async_cifar10 (no field) + stays byte-identical via the raw-ts fallback.""" + + def _pair(self, with_intrinsic): + # 10 data_ids: 20 s/data_id genuine work; REAL wall adds +15 s/data_id + # transport (35/unit), but intrinsic_span_s reports the clean 20. + real_rounds, sim_rounds = [], [] + for d in range(1, 11): + re = {"event": "agg_round", "round": 1, "ts": float(d * 35), + "data_id": d, "cycle_data_id": d, + "contributing_trainers": ["a"], "staleness": [0], + "agg_goal_count": 1} + if with_intrinsic: + re["intrinsic_span_s"] = 20.0 + real_rounds.append(re) + sim_rounds.append({"event": "agg_round", "round": 1, "ts": float(d), + "vclock_now": float(d * 20), "data_id": d, + "cycle_data_id": d, "contributing_trainers": ["a"], + "staleness": [0], "agg_goal_count": 1, + "intrinsic_span_s": 20.0}) + return _agg(agg_rounds=real_rounds), _agg(agg_rounds=sim_rounds) + + def test_intrinsic_anchor_passes(self): + real, sim = self._pair(with_intrinsic=True) + assert pc.throughput_parity(real, sim, tol_rel=0.05)["ok"] + assert pc.per_round_advance_parity(real, sim)["ok"] + assert pc.total_commits_parity(real, sim)["ok"] + assert pc.terminal_state_parity(real, sim)["ok"] + wd = pc.wall_disparity(real, sim) + assert wd["anchor"] == "intrinsic_span" + assert wd["mean_abs_disparity_s"] < 1.0, wd + + def test_raw_wall_fallback_fails_and_is_byte_identical(self): + # Field absent -> real anchors on raw ts (35/unit) vs sim vclock + # (20/unit): the #6 gap. Also the async_cifar10 path. + real, sim = self._pair(with_intrinsic=False) + assert not pc.throughput_parity(real, sim, tol_rel=0.05)["ok"] + assert not pc.per_round_advance_parity(real, sim)["ok"] + wd = pc.wall_disparity(real, sim) + assert wd["anchor"] == "wall_ts" + assert wd["max_abs_disparity_s"] > 100.0, wd # 15 s/unit artifact, cumulative + + +class TestSimSpeedup: + """sim_speedup [DIAG] (#13): sim must run virtual time at least as fast as + wall (sim_rate >= 1). K7's sane-range [0.01,100] check passes a slowdown; + this rung catches it.""" + + def test_healthy_speedup_passes(self): + # sim: vclock 0..100 virtual-s in 0..10 wall-s (10x speedup); + # real: same work took 0..100 wall-s. + sim = _agg(agg_rounds=[ + _round(r, ["a"], [0], vclock=float(r * 10), ts=float(r)) + for r in range(1, 11)]) + real = _agg(agg_rounds=[ + _round(r, ["a"], [0], ts=float(r * 10)) for r in range(1, 11)]) + r = pc.sim_speedup(real, sim) + assert r["ok"] and r["is_speedup"], r + assert r["sim_rate"] >= 1.0 + assert r["wall_speedup"] > 1.0 # sim finished faster than real + + def test_slowdown_fails(self): + # Slowdown shape: vclock reaches ~213 while wall burns ~566 + # (sim_rate ~0.38 < 1) -> sim is broken (#13). + sim = _agg(agg_rounds=[ + _round(r, ["a"], [0], vclock=float(r * 213.0 / 24), + ts=float(r * 566.0 / 24)) + for r in range(1, 25)]) + real = _agg(agg_rounds=[ + _round(r, ["a"], [0], ts=float(r * 558.0 / 14)) + for r in range(1, 15)]) + r = pc.sim_speedup(real, sim) + assert not r["ok"] and not r["is_speedup"], r + assert r["sim_rate"] < 1.0 + assert "SLOWDOWN" in r["note"] + + def test_prefers_wall_elapsed_s_over_ts_span(self): + # When agg_round carries wall_elapsed_s (the re-anchored measure), it is + # used instead of the ts epoch span. + sim_rounds = [ + _round(r, ["a"], [0], vclock=float(r * 10), ts=float(r * 1000 + r)) + for r in range(1, 11)] + for e in sim_rounds: + e["wall_elapsed_s"] = float(e["round"]) # 1..10, unlike the ts span + sim = _agg(agg_rounds=sim_rounds) + real = _agg(agg_rounds=[ + _round(r, ["a"], [0], ts=float(r * 10)) for r in range(1, 11)]) + r = pc.sim_speedup(real, sim) + assert r["sim_wall_s"] == 10.0 # from wall_elapsed_s, not the huge ts span + assert r["sim_rate"] >= 1.0 + + def test_no_vclock_skips(self): + real = _agg(agg_rounds=[_round(1, ["a"], [0], ts=10.0)]) + sim = _agg(agg_rounds=[_round(1, ["a"], [0], ts=5.0)]) # no vclock + r = pc.sim_speedup(real, sim) + assert r.get("status") == "SKIP" + + class TestPerRoundAdvanceParity: def test_matched_passes(self): # Both advance 10s per round @@ -312,6 +413,241 @@ def test_diverged_fails(self): assert not r["ok"], r +def _fwd_round(data_id, contributing, vclock=None, ts=0.0): + """fwdllm-style agg_round: `round` static, progress on the committed + `data_id` axis (cycle_data_id).""" + e = {"event": "agg_round", "round": 1, "ts": ts, + "cycle_data_id": data_id, "var_good_enough": True, + "contributing_trainers": contributing, "staleness": [0], + "agg_goal_count": 1} + if vclock is not None: + e["vclock_now"] = vclock + return e + + +class TestProgressAxisRekey: + """The clock family must measure progress on the axis the run advances. + fwdllm keeps `round` at 1 and advances committed `data_id`, so a rung keyed + on `round` divides by a stuck counter. The re-key auto-detects the axis; + async_cifar10 (round-advancing) stays byte-identical.""" + + def test_throughput_counts_data_ids_not_static_round(self): + # 10 committed data_ids, round pinned at 1, matched 10 units / 100s. + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float((d + 1) * 10)) + for d in range(10)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float((d + 1) * 10), + ts=float(d + 1)) + for d in range(10)]) + r = pc.throughput_parity(real, sim, tol_rel=0.10) + # Re-keyed to data_id: 10 units, NOT collapsed to 1 stuck round. + assert r["sim_rounds"] == 10 and r["real_rounds"] == 10, r + assert r["ok"], r + + def test_throughput_divergence_caught_on_data_id_axis(self): + # sim crawls (10 data_ids in 200s vclock) vs real (10 in 100s wall) -> 2x. + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float((d + 1) * 10)) + for d in range(10)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float((d + 1) * 20), + ts=float(d + 1)) + for d in range(10)]) + r = pc.throughput_parity(real, sim, tol_rel=0.10) + assert not r["ok"], r + + def test_terminal_state_data_ids_at_V_nonzero(self): + real = _agg(agg_rounds=[_fwd_round(d, ["a", "b"], ts=float((d + 1) * 10)) + for d in range(10)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a", "b"], vclock=float(d * 10), + ts=float(d + 1)) + for d in range(10)]) + r = pc.terminal_state_parity(real, sim) + # Previously real_rounds_at_V==0 (round static); now counts data_ids. + assert r["real_rounds_at_V"] > 0 and r["sim_rounds_at_V"] > 0, r + assert r["ok"], r + + def test_total_commits_counts_distinct_data_ids(self): + # A variance-FAIL retry emits 2 cycles on the SAME data_id; the commit + # count must be distinct data_ids (2), not raw cycles (3). + real = _agg(agg_rounds=[ + _fwd_round(0, ["a"], ts=0.0), + _fwd_round(0, ["a"], ts=5.0), # retry, same data_id + _fwd_round(1, ["a"], ts=10.0), + ]) + sim = _agg(agg_rounds=[ + _fwd_round(0, ["a"], vclock=0.0, ts=1.0), + _fwd_round(0, ["a"], vclock=5.0, ts=2.0), + _fwd_round(1, ["a"], vclock=10.0, ts=3.0), + ]) + r = pc.total_commits_parity(real, sim, tol_rel=0.05) + assert r["n_sim_commits"] == 2 and r["n_real_commits"] == 2, r + assert r["ok"], r + + def test_normal_fl_still_keyed_on_round(self): + # >1 distinct round -> axis stays "round"; identical to pre-re-key. + real = _agg(agg_rounds=[_round(r, ["a"], [0], ts=float(r * 10)) + for r in range(1, 11)]) + sim = _agg(agg_rounds=[_round(r, ["a"], [0], vclock=float(r * 10), + ts=float(r)) for r in range(1, 11)]) + r = pc.throughput_parity(real, sim, tol_rel=0.10) + assert r["sim_rounds"] == 10, r + assert r["ok"], r + + # --- The 4 advance rungs re-key too via _per_round_advances: keyed on + # `round` they SKIP fwdllm ("<2 rounds"); on data_id they measure advances. --- + + def test_advance_rung_measures_on_data_id_axis(self): + # Matched 10 s/data_id both sides -> advances measured, PASSes (not SKIP). + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float(d * 10)) + for d in range(10)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 10), + ts=float(d)) for d in range(10)]) + r = pc.per_round_advance_parity(real, sim, ks_tol=0.2, mean_tol_rel=0.15) + assert "run too short" not in r.get("note", ""), r + assert r["ok"], r + + def test_advance_rung_catches_data_id_rate_gap(self): + # sim 26 s/data_id vclock vs real 15 s/data_id wall -> divergence caught. + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float(d * 15)) + for d in range(10)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 26), + ts=float(d)) for d in range(10)]) + r = pc.per_round_advance_parity(real, sim, ks_tol=0.2, mean_tol_rel=0.15) + assert not r["ok"], r + + def test_advance_mean_on_round_axis_unchanged(self): + # async_cifar10 (round-advancing): the sim mean advance is measured on + # `round` exactly as before the re-key (+7 vclock per round). + real = _agg(agg_rounds=[_round(r, ["a"], [0], ts=float(r * 7)) + for r in range(1, 8)]) + sim = _agg(agg_rounds=[_round(r, ["a"], [0], vclock=float(r * 7), + ts=float(r)) for r in range(1, 8)]) + r = pc.per_round_advance_parity(real, sim, ks_tol=0.2, mean_tol_rel=0.15) + assert r["sim_mean_advance_s"] == 7.0, r + assert r["real_mean_advance_s"] == 7.0, r + + def test_advance_mean_on_data_id_when_round_static(self): + # fwdllm: round pinned at 1 -> mean advance is measured per data_id + # (+4 vclock per committed data_id), not collapsed to a single stuck bin. + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float(d * 4)) + for d in range(5)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 4), + ts=float(d)) for d in range(5)]) + r = pc.per_round_advance_parity(real, sim, ks_tol=0.2, mean_tol_rel=0.15) + assert r["sim_mean_advance_s"] == 4.0, r + + +class TestWallDisparity: + """#6: a DIAG rung reporting |real_wall - sim_vclock| per matched progress + unit. Never gates; surfaces the residual to drive to ~0.""" + + def test_zero_disparity_when_clocks_match(self): + # real advances 10 wall-s/data_id, sim 10 vclock-s/data_id -> residual 0. + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float(d * 10)) + for d in range(10)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 10), + ts=float(d)) for d in range(10)]) + r = pc.wall_disparity(real, sim) + assert r["axis"] == "data_id" + assert r["mean_abs_disparity_s"] == 0.0, r + assert r["max_abs_disparity_s"] == 0.0, r + assert r["n_matched_units"] == 10 + + def test_surfaces_the_rate_gap(self): + # real 35 wall-s/data_id vs sim 8 vclock-s/data_id -> growing residual, + # but the rung still "ok" (DIAG never fails). + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float(d * 35)) + for d in range(10)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 8), + ts=float(d)) for d in range(10)]) + r = pc.wall_disparity(real, sim) + assert r["ok"] is True # DIAG: informational only + # by the last of 10 data_ids: |9*35 - 9*8| = 243 + assert r["max_abs_disparity_s"] > 100.0, r + assert r["mean_abs_disparity_s"] > 0.0 + + def test_aligns_when_run_does_not_start_at_unit_zero(self): + # matched units start at data_id 3; cumulative-from-first-matched keeps + # residual 0 despite the nonzero vclock/ts origin. + real = _agg(agg_rounds=[_fwd_round(d, ["a"], ts=float(d * 10)) + for d in range(3, 8)]) + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 10), + ts=float(d)) for d in range(3, 8)]) + r = pc.wall_disparity(real, sim) + assert r["mean_abs_disparity_s"] == 0.0, r + + def test_skips_when_too_few_matched_units(self): + real = _agg(agg_rounds=[_fwd_round(0, ["a"], ts=0.0)]) + sim = _agg(agg_rounds=[_fwd_round(0, ["a"], vclock=0.0, ts=0.0)]) + r = pc.wall_disparity(real, sim) + assert r.get("status") == "SKIP", r + assert r["ok"] is True + + def test_round_axis_for_normal_fl(self): + real = _agg(agg_rounds=[_round(r, ["a"], [0], ts=float(r * 5)) + for r in range(1, 6)]) + sim = _agg(agg_rounds=[_round(r, ["a"], [0], vclock=float(r * 5), + ts=float(r)) for r in range(1, 6)]) + r = pc.wall_disparity(real, sim) + assert r["axis"] == "round" + assert r["mean_abs_disparity_s"] == 0.0, r + + +class TestFailsafeRealComputeSim: + """A real-compute sim (fwdllm runs the real forward-grad pass in sim mode) + has sim wall >> vclock by construction, so K5 must compare wall against the + RUN wall budget, not the vclock (else it false-fails).""" + + def test_real_compute_sim_skips_without_run_budget(self): + # data_id axis auto-detects real_compute_sim; wall ≫ vclock but no run + # budget passed -> SKIP, not a false INV failure. + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 8), + ts=float(d * 35)) for d in range(6)]) + r = pc.failsafe_ok(sim, budget_s=None) + assert r.get("status") == "SKIP", r + assert r["ok"] is True + + def test_real_compute_sim_uses_run_budget_when_given(self): + # 6 data_ids, wall ends at 5*35=175s; run budget 3600s -> within budget. + sim = _agg(agg_rounds=[_fwd_round(d, ["a"], vclock=float(d * 8), + ts=float(d * 35)) for d in range(6)]) + r = pc.failsafe_ok(sim, budget_s=3600.0) + assert r.get("status") != "SKIP", r + assert r["real_compute_sim"] is True + assert r["ok"] is True # 175s wall << 3600s run budget + + def test_cheap_compute_sim_keeps_vclock_fallback(self): + # round-advancing (async_cifar10): wall≈vclock, no run budget -> vclock + # fallback unchanged (byte-identical). + sim = _agg(agg_rounds=[_round(r, ["a"], [0], vclock=float(r * 10), + ts=float(r * 10)) for r in range(1, 7)]) + r = pc.failsafe_ok(sim, budget_s=None) + assert r["real_compute_sim"] is False + assert r.get("status") != "SKIP" + assert r["ok"] is True # wall == vclock -> 0 overshoot + + +class TestFieldCoverageAlias: + """fwdllm's trainer emits gpu/budget under different names; the coverage + matrix must accept either spelling instead of false-FAILing.""" + + def test_fwdllm_field_aliases_count_as_covered(self): + agg = _agg(agg_rounds=[{"event": "agg_round", "round": 1, "ts": 0.0, + "vclock_now": 1.0, "trainer_speed_s": [1.0], + "staleness": [0], "stat_utility": [1.0], + "contributing_trainers": ["t1"]}]) + sel = {"num_eligible": 5, "avail_composition": {"a": 1}, + "num_chosen": 3} + agg["selection_train"] = [sel] + # trainer emits the fwdllm spellings only. + tr = {"t1": {"trainer_round": [ + {"real_gpu_time_s": 1.2, "sim_round_duration_s": 2.3}]}} + r = pc.field_coverage(agg, agg, tr, tr) + gpu = r["matrix"]["trainer_round.gpu_compute_s"] + bud = r["matrix"]["trainer_round.training_budget_s"] + assert gpu["real"] == 1.0 and gpu["sim"] == 1.0, r + assert bud["real"] == 1.0 and bud["sim"] == 1.0, r + assert "trainer_round.gpu_compute_s(real)" not in r["violations"], r + + class TestTrainerSpeedParity: def test_identical_passes(self): rounds_with_speed = [ @@ -467,3 +803,559 @@ def test_builder_paired_commit_class(self): round_num=1, time_mode="sim", in_flight_before=13, in_flight_after=13, residence_rounds=[1]) assert "residence_staleness" not in f2 + + +# ── FwdLLM variance-cadence layer (V/DK/G rungs) ────────── + +def _cadence(cycle_data_id, cycle_iteration, var, var_threshold, + var_good_enough, force_commit_planned=False, grad_pool_size=None, + cached_v_size=None, agg_goal=None, round_=1, ts=0.0, **extra): + """One fwdllm agg-goal-boundary (variance gate) agg_round event.""" + e = {"event": "agg_round", "round": round_, "ts": ts, + "cycle_data_id": cycle_data_id, "cycle_iteration": cycle_iteration, + "var": var, "var_threshold": var_threshold, + "var_good_enough": var_good_enough, + "force_commit_planned": force_commit_planned} + for k, v in (("grad_pool_size", grad_pool_size), + ("cached_v_size", cached_v_size), ("agg_goal", agg_goal)): + if v is not None: + e[k] = v + e.update(extra) + return e + + +def _cadence_run(iters_per_data, var_threshold=1.0, pass_var=0.5, fail_var=2.0, + agg_goal=3): + """Series where data_id d takes iters_per_data[d] cycles: (n-1) variance + FAILs (var>thr) then one PASS (var<=thr). grad_pool grows each retry and is + at its max on the committing cycle; cached_v grows across the FAIL rollbacks. + """ + events = [] + for d, n in enumerate(iters_per_data): + for it in range(n): + committed = (it == n - 1) + events.append(_cadence( + cycle_data_id=d, cycle_iteration=it, + var=(pass_var if committed else fail_var), + var_threshold=var_threshold, var_good_enough=committed, + grad_pool_size=(it + 1) * agg_goal, cached_v_size=it, + agg_goal=agg_goal)) + return events + + +class TestV1IterPerDataId: + def test_matched_passes(self): + real = _agg(agg_rounds=_cadence_run([1, 2, 1, 3, 1])) + sim = _agg(agg_rounds=_cadence_run([1, 2, 1, 3, 1])) + r = pc.iters_per_data_id_parity(real, sim) + assert r["ok"] and r["real_mean_iters"] == r["sim_mean_iters"], r + + def test_diverged_fails(self): + # sim compounds: every data_id needs 3 cycles vs real's 1 (a + # contributing-set/order divergence surfacing as more retries). + real = _agg(agg_rounds=_cadence_run([1, 1, 1, 1, 1])) + sim = _agg(agg_rounds=_cadence_run([3, 3, 3, 3, 3])) + r = pc.iters_per_data_id_parity(real, sim) + assert not r["ok"] and r["sim_mean_iters"] > r["real_mean_iters"], r + + def test_iters_binning_exact_for_force_commit(self): + # A data_id that force-commits after 2 fails still counts 3 cycles: the + # commit event carries cycle_data_id of the SAME data_id (pre-advance). + cycles = [ + _cadence(0, 0, 2.0, 1.0, False), + _cadence(0, 1, 2.0, 1.0, False), + _cadence(0, 2, 2.0, 1.0, True, force_commit_planned=True), # forced + ] + assert pc._iters_per_data_id(cycles) == {0: 3} + + def test_non_fwdllm_skips(self): + a = _agg(agg_rounds=[_round(1, ["a"], [0], vclock=1.0)]) # no cadence fields + r = pc.iters_per_data_id_parity(a, a) + assert r["ok"] and r.get("status") == "SKIP", r + + +class TestV2VarTrajectory: + def test_matched_passes(self): + real = _agg(agg_rounds=_cadence_run([2, 2, 2])) + sim = _agg(agg_rounds=_cadence_run([2, 2, 2])) + assert pc.var_trajectory_parity(real, sim)["ok"] + + def test_diverged_fails(self): + # Same iteration counts but the variance *signal* differs (grad-pool + # accumulation-order bug): sim's var values sit far from real's. + real = _agg(agg_rounds=[_cadence(d, 0, 0.4, 1.0, True) for d in range(8)]) + sim = _agg(agg_rounds=[_cadence(d, 0, 5.0, 1.0, True) for d in range(8)]) + assert not pc.var_trajectory_parity(real, sim)["ok"] + + +class TestV4ForceCommitRate: + def test_matched_passes(self): + real = _agg(agg_rounds=_cadence_run([1, 1, 1, 1])) + sim = _agg(agg_rounds=_cadence_run([1, 1, 1, 1])) + r = pc.force_commit_rate_parity(real, sim) + assert r["ok"] and r["sim_force_commit_rate"] == 0.0, r + + def test_diverged_fails(self): + # sim hits the iteration cap far more often (chronic variance divergence). + real = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True) for d in range(10)]) + sim = _agg(agg_rounds=[ + _cadence(d, 0, 2.0, 1.0, True, force_commit_planned=True) + for d in range(10)]) + r = pc.force_commit_rate_parity(sim, real) # order-agnostic + assert not r["ok"], r + + +class TestV5VariancePassRatio: + def test_matched_passes(self): + real = _agg(agg_rounds=_cadence_run([1, 2, 1, 2])) + sim = _agg(agg_rounds=_cadence_run([1, 2, 1, 2])) + assert pc.variance_pass_ratio_parity(real, sim)["ok"] + + def test_force_commit_excluded_from_pass(self): + # A force-commit (var>thr, committed only because the cap fired) is NOT a + # genuine variance pass -> real (genuine) and sim (forced) diverge on V5 + # even though both "committed" every cycle. + real = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True) for d in range(10)]) + sim = _agg(agg_rounds=[ + _cadence(d, 0, 2.0, 1.0, True, force_commit_planned=True) + for d in range(10)]) + r = pc.variance_pass_ratio_parity(real, sim) + assert not r["ok"] + assert r["real_pass_ratio"] == 1.0 and r["sim_pass_ratio"] == 0.0, r + + +class TestV3CachedVPool: + def test_matched_passes(self): + real = _agg(agg_rounds=_cadence_run([2, 3, 2])) + sim = _agg(agg_rounds=_cadence_run([2, 3, 2])) + assert pc.cached_v_pool_parity(real, sim)["ok"] + + def test_absent_skips(self): + real = _agg(agg_rounds=[_cadence(0, 0, 0.5, 1.0, True)]) + sim = _agg(agg_rounds=[_cadence(0, 0, 0.5, 1.0, True)]) + r = pc.cached_v_pool_parity(real, sim) + assert r["ok"] and r.get("status") == "SKIP", r + + +class TestDK1AggGoalTrajectory: + def test_constant_k_skips(self): + real = _agg(agg_rounds=_cadence_run([1, 1, 1], agg_goal=3)) + sim = _agg(agg_rounds=_cadence_run([1, 1, 1], agg_goal=3)) + r = pc.agg_goal_trajectory_parity(real, sim) + assert r["ok"] and r.get("status") == "SKIP" and "disabled" in r["note"], r + + def test_varying_k_matched_passes(self): + real = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True, agg_goal=k) + for d, k in enumerate([3, 3, 4, 5, 4])]) + sim = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True, agg_goal=k) + for d, k in enumerate([3, 3, 4, 5, 4])]) + assert pc.agg_goal_trajectory_parity(real, sim)["ok"] + + +class TestDK3EligibleEndsMetric: + def test_absent_skips_with_note(self): + real = _agg(agg_rounds=_cadence_run([1, 1])) + sim = _agg(agg_rounds=_cadence_run([1, 1])) + r = pc.eligible_ends_metric_parity(real, sim) + assert r["ok"] and r.get("status") == "SKIP" and "deferred" in r["note"], r + + def test_present_diverged_fails(self): + real = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True, n_eligible_train=10) + for d in range(6)]) + sim = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True, n_eligible_train=2) + for d in range(6)]) + assert not pc.eligible_ends_metric_parity(real, sim)["ok"] + + +class TestG1GradNorm: + def test_absent_skips_with_note(self): + real = _agg(agg_rounds=_cadence_run([1, 1])) + sim = _agg(agg_rounds=_cadence_run([1, 1])) + r = pc.grad_norm_parity(real, sim) + assert r["ok"] and r.get("status") == "SKIP" and "deferred" in r["note"], r + + +class TestG2GradPoolSize: + def test_matched_passes(self): + real = _agg(agg_rounds=_cadence_run([1, 2, 1, 2])) + sim = _agg(agg_rounds=_cadence_run([1, 2, 1, 2])) + assert pc.grad_pool_size_parity(real, sim)["ok"] + + def test_diverged_fails(self): + # sim accumulates far larger pools before committing (K x V1 rollup). + real = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True, grad_pool_size=3) + for d in range(8)]) + sim = _agg(agg_rounds=[_cadence(d, 0, 0.5, 1.0, True, grad_pool_size=30) + for d in range(8)]) + assert not pc.grad_pool_size_parity(real, sim)["ok"] + + +class TestRunAllParityFwdllm: + _CADENCE_KEYS = [ + "v1_iter_per_data_id", "v2_var_trajectory", "v3_cached_v_pool", + "v4_force_commit_rate", "v5_variance_pass_ratio", + "dk1_agg_goal_trajectory", "dk2_dynamic_c", "dk3_eligible_ends_metric", + "g1_grad_norm", "g2_grad_pool_size", + ] + + def test_keys_present_and_identical_passes(self): + a = _agg(selection=[_sel(1, ["a", "b"])], + agg_rounds=_cadence_run([1, 2, 1, 3, 1])) + tr = {"aa": {"task_recv": [], "trainer_round": []}} + res = pc.run_all_parity(a, a, tr, tr, agg_goal=3) + for key in self._CADENCE_KEYS: + assert key in res, f"missing cadence key: {key}" + v = res[key] + # identical real==sim ⇒ every cadence rung PASSes or SKIPs cleanly + assert v["ok"], f"{key}: {v}" + + def test_non_fwdllm_run_skips_all_cadence(self): + a = _agg(agg_rounds=[_round(r, ["a"], [0], vclock=float(r)) + for r in range(1, 4)]) + tr: dict = {} + res = pc.run_all_parity(a, a, tr, tr) + for key in self._CADENCE_KEYS: + assert res[key].get("status") == "SKIP", f"{key} should SKIP: {res[key]}" + + def test_cadence_divergence_fails_verdict(self): + real = _agg(agg_rounds=_cadence_run([1, 1, 1, 1, 1, 1])) + sim = _agg(agg_rounds=_cadence_run([3, 3, 3, 3, 3, 3])) + tr: dict = {} + res = pc.run_all_parity(real, sim, tr, tr, agg_goal=3) + passed, roots, downstream, _w = pc.overall_verdict(res) + assert not passed + assert "v1_iter_per_data_id" in (set(roots) | set(downstream)) + + +class TestAggRoundCadenceEmission: + """The agg_round builder carries the cadence fields through `extra` (the + aggregator snapshots them pre-mutation).""" + + def test_cadence_fields_land_in_event(self): + from flame.telemetry.events import build_agg_round, EVENT_AGG_ROUND + ev, f = build_agg_round( + round_num=4, agg_goal=3, agg_goal_count=3, + extra={"cycle_data_id": 7, "cycle_iteration": 2, "var": 0.8, + "var_threshold": 1.0, "var_good_enough": True, + "force_commit_planned": False, "grad_pool_size": 9, + "cached_v_size": 2}) + assert ev == EVENT_AGG_ROUND + assert f["cycle_data_id"] == 7 and f["cycle_iteration"] == 2 + assert f["grad_pool_size"] == 9 and f["cached_v_size"] == 2 + + +# ── FwdLLM async residence rungs R1 / W1 ── + +def _cyc(cycle_data_id, intervals, contributing=None): + """A committed cadence cycle carrying per-contributor [dispatch, commit] + intervals. `intervals` = list of (end, dispatch_ts, commit_ts).""" + ci = [{"end": e, "dispatch_ts": d, "commit_ts": c} for (e, d, c) in intervals] + return {"event": "agg_round", "round": 1, "ts": 0.0, + "cycle_data_id": cycle_data_id, "var_good_enough": True, + "agg_goal_count": len(ci), + "contributing_trainers": contributing or [e for (e, _, _) in intervals], + "contributor_intervals": ci} + + +def _trainers_with_rounds(counts): + """{short_id: {"trainer_round": [ ...n events ]}} for W1 forward-pass counts.""" + return {sid: {"trainer_round": [{"event": "trainer_round"} for _ in range(n)]} + for sid, n in counts.items()} + + +class TestR1InflightOverlap: + """R1 [INV]: per-trainer dispatch->commit intervals must not overlap + (one-in-flight residence).""" + + def test_non_overlapping_intervals_pass(self): + # Each trainer's two contributions are strictly sequential (commit before + # the next dispatch) in BOTH modes. + agg = _agg(agg_rounds=[ + _cyc(0, [("A", 0.0, 5.0), ("B", 0.0, 5.0)]), + _cyc(1, [("A", 6.0, 11.0), ("B", 6.0, 11.0)]), + ]) + r = pc.inflight_overlap_parity(agg, agg) + assert r["ok"] + assert r["real_overlap_frac"] == 0.0 and r["sim_overlap_frac"] == 0.0 + + def test_sim_overlap_fails_with_clean_real(self): + # Real: A's 2nd dispatch (6.0) is after its 1st commit (5.0) -> clean. + real = _agg(agg_rounds=[ + _cyc(0, [("A", 0.0, 5.0)]), + _cyc(1, [("A", 6.0, 11.0)]), + ]) + # Sim: A re-dispatched at 2.0 while its 1st contribution (commit 5.0) was + # still in flight -> overlap = the residence violation. + sim = _agg(agg_rounds=[ + _cyc(0, [("A", 0.0, 5.0)]), + _cyc(1, [("A", 2.0, 7.0)]), + ]) + r = pc.inflight_overlap_parity(real, sim) + assert not r["ok"] + assert r["real_overlap_frac"] == 0.0 + assert r["sim_overlap_frac"] > 0.0 + + def test_skips_without_contributor_intervals(self): + # Sync baselines / non-fwdllm runs don't emit contributor_intervals. + agg = _agg(agg_rounds=[_round(1, ["a"], [0])]) + r = pc.inflight_overlap_parity(agg, agg) + assert r.get("status") == "SKIP" and r["ok"] + + +class TestW1ComputeConservation: + """W1 [DIAG]: forward passes per committed grad; a large sim excess over + real = wasted recompute (the residence violation), localizes to R1.""" + + def test_matched_ratio_passes(self): + real_agg = _agg(agg_rounds=[_cyc(0, [("A", 0.0, 5.0), ("B", 0.0, 5.0)])]) + sim_agg = _agg(agg_rounds=[_cyc(0, [("A", 0.0, 5.0), ("B", 0.0, 5.0)])]) + # 2 committed each; ~1.5 forward passes per commit in both modes. + real_tr = _trainers_with_rounds({"A": 2, "B": 1}) + sim_tr = _trainers_with_rounds({"A": 2, "B": 1}) + r = pc.compute_conservation_parity(real_agg, sim_agg, real_tr, sim_tr) + assert r["ok"] + assert r["real_fwd_per_commit"] == r["sim_fwd_per_commit"] + + def test_sim_recompute_excess_fails(self): + real_agg = _agg(agg_rounds=[_cyc(0, [("A", 0.0, 5.0), ("B", 0.0, 5.0)])]) + sim_agg = _agg(agg_rounds=[_cyc(0, [("A", 0.0, 5.0), ("B", 0.0, 5.0)])]) + # Same 2 commits both modes, but sim ran ~2x the forward passes (the + # residence-bug direction: sim OVER-computes -> violation). + real_tr = _trainers_with_rounds({"A": 1, "B": 1}) # 2 fwd / 2 commit = 1.0 + sim_tr = _trainers_with_rounds({"A": 3, "B": 3}) # 6 fwd / 2 commit = 3.0 + r = pc.compute_conservation_parity(real_agg, sim_agg, real_tr, sim_tr) + assert not r["ok"] + assert r["sim_excess_rel"] > 0 + + def test_sim_under_computing_is_ok(self): + # W1 is ASYMMETRIC: sim doing FEWER forward passes than real (the async + # start-tail, not wasted recompute) must NOT fail. + real_agg = _agg(agg_rounds=[_cyc(0, [("A", 0.0, 5.0), ("B", 0.0, 5.0)])]) + sim_agg = _agg(agg_rounds=[_cyc(0, [("A", 0.0, 5.0), ("B", 0.0, 5.0)])]) + real_tr = _trainers_with_rounds({"A": 3, "B": 3}) # 3.0 fwd/commit + sim_tr = _trainers_with_rounds({"A": 1, "B": 1}) # 1.0 fwd/commit + r = pc.compute_conservation_parity(real_agg, sim_agg, real_tr, sim_tr) + assert r["ok"] and r["sim_excess_rel"] < 0 + + def test_skips_without_data(self): + empty = _agg(agg_rounds=[]) + r = pc.compute_conservation_parity(empty, empty, {}, {}) + assert r.get("status") == "SKIP" and r["ok"] + + +class TestR1W1Registered: + """Both rungs run in run_all_parity and are wired into the causal registry + (R1 upstream of V1 -- the dep chain proving cadence is downstream).""" + + def test_present_in_run_all_and_meta(self): + assert "r1_inflight_overlap" in pc.CHECK_META + assert "w1_compute_conservation" in pc.CHECK_META + assert "r1_inflight_overlap" in pc.CHECK_META["v1_iter_per_data_id"]["deps"] + assert pc.CHECK_META["w1_compute_conservation"]["deps"] == ("r1_inflight_overlap",) + + +def _lcyc(data_id, iteration, cohort, var, var_good=False, force=False, goal=3): + """One fwdllm variance-cadence cycle event (cohort in receive/commit order).""" + return {"event": "agg_round", "round": 1, + "cycle_data_id": data_id, "iteration_per_data_id": iteration, + "contributing_trainers": list(cohort), "var": var, + "var_good_enough": var_good, "force_commit_planned": force, + "agg_goal_count": goal, "staleness": [0] * len(cohort)} + + +class TestCohortSequence: + """L1 cohort_sequence_parity: the ordered per-aggregation logical sequence + (set + receive-ORDER + cadence + var value) must be IDENTICAL. EXACT, ungated + (real receive-order is deterministic in both modes by design).""" + + def test_identical_sequence_passes(self): + cyc = [_lcyc(0, 1, ["a", "b", "c"], 0.9), + _lcyc(0, 2, ["a", "b", "c"], 0.28, var_good=True), + _lcyc(1, 1, ["a", "b", "c"], 0.5)] + r = pc.cohort_sequence_parity(_agg(agg_rounds=list(cyc)), + _agg(agg_rounds=list(cyc))) + assert r["ok"], r + assert r["order_match_frac"] == 1.0 and r["var_match_frac"] == 1.0 + + def test_reordered_cohort_same_set_FAILS(self): + # Same SET each cycle, different receive ORDER -> feeds split-half var. + real = _agg(agg_rounds=[_lcyc(0, 1, ["a", "b", "c"], 0.9)]) + sim = _agg(agg_rounds=[_lcyc(0, 1, ["c", "a", "b"], 0.9)]) + r = pc.cohort_sequence_parity(real, sim) + assert not r["ok"], r + assert r["set_match_frac"] == 1.0 and r["order_match_frac"] == 0.0 + assert r["first_divergence"]["order_ok"] is False + + def test_different_cohort_FAILS(self): + real = _agg(agg_rounds=[_lcyc(0, 1, ["a", "b", "c"], 0.9)]) + sim = _agg(agg_rounds=[_lcyc(0, 1, ["a", "b", "d"], 0.9)]) + r = pc.cohort_sequence_parity(real, sim) + assert not r["ok"] and r["set_match_frac"] == 0.0 + + def test_var_divergence_FAILS_even_with_matched_order(self): + # Identical cohort+order, var off by >0.1% -> the RNG-desync tell. + real = _agg(agg_rounds=[_lcyc(0, 1, ["a", "b", "c"], 0.371605)]) + sim = _agg(agg_rounds=[_lcyc(0, 1, ["a", "b", "c"], 0.371067)]) + r = pc.cohort_sequence_parity(real, sim) + assert not r["ok"] and r["var_match_frac"] == 0.0 + assert r["first_divergence"]["var_ok"] is False + + def test_cadence_shift_FAILS(self): + real = _agg(agg_rounds=[_lcyc(0, 1, ["a"], 0.5), _lcyc(1, 0, ["a"], 0.2, var_good=True)]) + sim = _agg(agg_rounds=[_lcyc(0, 1, ["a"], 0.5), _lcyc(0, 2, ["a"], 0.4)]) # extra iter + r = pc.cohort_sequence_parity(real, sim) + assert not r["ok"] and r["cadence_match_frac"] < 1.0 + + def test_tier_exact_and_enforced(self): + # EXACT tier -> enforced FAIL even under --lenient (not a DIST/DIAG warn). + r = pc.cohort_sequence_parity( + _agg(agg_rounds=[_lcyc(0, 1, ["a"], 0.5)]), + _agg(agg_rounds=[_lcyc(0, 1, ["b"], 0.5)])) + assert r["tier"] == "EXACT" and r["ok"] is False + passed, roots, _down, _warn = pc.overall_verdict( + {"cohort_sequence": r}, lenient=True) + assert not passed and "cohort_sequence" in roots + + def test_skips_on_non_fwdllm(self): + # No cadence fields (async_cifar10 shape) -> clean SKIP, byte-identical. + plain = {"event": "agg_round", "round": 0, "ts": 0.0, + "contributing_trainers": ["a"], "staleness": [0]} + rd = _agg(agg_rounds=[plain]) + r = pc.cohort_sequence_parity(rd, rd) + assert r.get("status") == "SKIP" and r["ok"] + + def test_max_bin_windows_to_first_bin(self): + # Cohorts match on bin 0, diverge on bin 1 -> --max-bin 0 passes. + real = _agg(agg_rounds=[_lcyc(0, 1, ["a", "b"], 0.5), _lcyc(1, 1, ["a", "b"], 0.5)]) + sim = _agg(agg_rounds=[_lcyc(0, 1, ["a", "b"], 0.5), _lcyc(1, 1, ["b", "a"], 0.5)]) + assert pc.cohort_sequence_parity(real, sim, max_bin=0)["ok"] + assert not pc.cohort_sequence_parity(real, sim)["ok"] + + def test_present_in_run_all_and_meta(self): + assert "cohort_sequence" in pc.CHECK_META + assert pc.CHECK_META["cohort_sequence"]["deps"] + + +class TestVarTrajectoryMeanGuard: + """V2 must fail a systematic mean offset that KS alone misses (a uniform ~1% + shift barely moves the CDF -> KS~0 but grads have desynced).""" + + def test_systematic_offset_fails_despite_low_ks(self): + base = [0.9, 0.5, 0.42, 0.31, 0.6, 0.48] + real = _agg(agg_rounds=[_lcyc(0, i, ["a"], v) for i, v in enumerate(base)]) + sim = _agg(agg_rounds=[_lcyc(0, i, ["a"], v * 1.05) for i, v in enumerate(base)]) + r = pc.var_trajectory_parity(real, sim) + assert r["mean_rel_diff"] > r["mean_tol_rel"], r + assert not r["ok"], r + + def test_matched_var_passes(self): + base = [0.9, 0.5, 0.42, 0.31] + agg = _agg(agg_rounds=[_lcyc(0, i, ["a"], v) for i, v in enumerate(base)]) + assert pc.var_trajectory_parity(agg, agg)["ok"] + + +def _selc(round_, chosen, num_candidates, selector="random", ts=0.0): + """Selection event WITH the num_chosen/num_candidates count fields the + full-cohort determinism gate reads (real runs always emit these).""" + return {"event": "selection", "task": "train", "round": round_, "ts": ts, + "selector": selector, "chosen": list(chosen), + "num_chosen": len(chosen), "num_candidates": num_candidates} + + +class TestSelectionDeterminismGate: + """Set/sequence selection rungs ENFORCE when selection is provably + deterministic (full-cohort K>=pool, or a declared deterministic selector) + and gate to a trivial pass otherwise -- data-driven from num_chosen/ + num_candidates so it self-splits fwdllm (enforce) vs fluxtune / fwdllm_plus + (gate) and self-disables under scarcity.""" + + def test_full_cohort_helper(self): + full = _agg(selection=[_selc(1, ["a", "b", "c"], 3), + _selc(2, ["a", "b", "c"], 3)]) + subset = _agg(selection=[_selc(1, ["a", "b", "c"], 10)]) # fluxtune-like + assert pc._full_cohort_selection(full) + assert not pc._full_cohort_selection(subset) + + def test_full_cohort_ungates_and_enforces(self): + # fwdllm syn_0: everyone selected -> deterministic -> ENFORCED. A genuine + # per-round contributing-set divergence must now FAIL (was invisible). + real = _agg(selection=[_selc(1, ["a", "b", "c"], 3)], + agg_rounds=[_round(1, ["a", "b", "c"], [0, 0, 0])]) + sim = _agg(selection=[_selc(1, ["a", "b", "c"], 3)], + agg_rounds=[_round(1, ["a", "b", "x"], [0, 0, 0])]) # x != c + assert pc._selection_is_deterministic(real, sim) + agg_seq = pc.aggregation_sequence_parity(real, sim) + assert agg_seq["gated"] is False and agg_seq["ok"] is False + + def test_full_cohort_matched_is_genuine_pass(self): + a = _agg(selection=[_selc(1, ["a", "b", "c"], 3)], + agg_rounds=[_round(1, ["a", "b", "c"], [0, 0, 0])]) + agg_seq = pc.aggregation_sequence_parity(a, a) + assert agg_seq["gated"] is False and agg_seq["ok"] is True + + def test_subset_selection_stays_gated(self): + # fluxtune agg_goal=3 < pool=10: stochastic subset -> gated trivial pass + # even when the sim picks a DIFFERENT subset (cohort_sequence owns that). + real = _agg(selection=[_selc(1, ["a", "b", "c"], 10)], + agg_rounds=[_round(1, ["a", "b", "c"], [0, 0, 0])]) + sim = _agg(selection=[_selc(1, ["d", "e", "f"], 10)], + agg_rounds=[_round(1, ["d", "e", "f"], [0, 0, 0])]) + assert not pc._selection_is_deterministic(real, sim) + agg_seq = pc.aggregation_sequence_parity(real, sim) + assert agg_seq["gated"] is True and agg_seq["ok"] is True + + def test_asymmetric_eligibility_stays_gated(self): + # fwdllm_plus #7: real sees fewer eligible than the pool -> not full + # cohort in real -> gated (don't hard-fail a real-side eligibility gap). + real = _agg(selection=[_selc(1, ["a", "b"], 5)]) # 2 of 5 chosen + sim = _agg(selection=[_selc(1, ["a", "b", "c", "d", "e"], 5)]) # full + assert not pc._selection_is_deterministic(real, sim) + + def test_legacy_no_counts_falls_back_to_selector_rule(self): + # No count telemetry + unknown selector -> old behavior: ENFORCED (so the + # pre-existing disjoint-fails test semantics are preserved, no regression). + real = _agg(selection=[_sel(1, ["a", "b"])]) + sim = _agg(selection=[_sel(1, ["c", "d"])]) + assert pc._selection_is_deterministic(real, sim) # enforce on unknown + assert not pc.selection_parity(real, sim)["ok"] + + +def _tr_round(overran, data_id=0, it=0, gpu=1.0, budget=2.0): + return {"real_gpu_time_s": gpu, "training_budget_s": budget, + "training_overran": overran, "data_id": data_id, + "iteration_per_data_id": it} + + +class TestTimingOverrun: + """Ovr: order-determinism tell. A high overrun fraction means gpu > modeled + D -> arrival order can flip -> cohort_sequence/v2 breaks are a TIMING-MODEL + limit, not a sim ordering bug. DIAG (never fails).""" + + def test_no_overrun_verdict(self): + tr = {"a": {"trainer_round": [_tr_round(False), _tr_round(False)]}} + r = pc.timing_overrun(tr, tr) + assert r["ok"] and r["tier"] == "DIAG" + assert r["real_overrun_frac"] == 0.0 + assert "attainable" in r["verdict"] + + def test_overrun_flags_and_reports_first_bin(self): + # fluxtune-like: some rounds overrun; earliest at (data_id, iter). + real = {"a": {"trainer_round": [ + _tr_round(False, 0, 0), _tr_round(True, 0, 1), _tr_round(True, 2, 0)]}} + sim = {"a": {"trainer_round": [_tr_round(False, 0, 0)]}} + r = pc.timing_overrun(real, sim) + assert r["ok"] # DIAG never fails the scoreboard + assert r["real_overrun_frac"] == pytest.approx(2 / 3, abs=1e-3) + assert r["real_first_overrun"] == [0, 1] + assert "OVERRUN" in r["verdict"] + + def test_skips_without_telemetry(self): + tr = {"a": {"trainer_round": [{"real_gpu_time_s": 1.0}]}} # no overran field + r = pc.timing_overrun(tr, tr) + assert r.get("status") == "SKIP" and r["ok"] + + def test_present_in_run_all_and_meta(self): + assert "timing_overrun" in pc.CHECK_META + tr = {"a": {"trainer_round": [_tr_round(False)]}} + res = pc.run_all_parity(_agg(), _agg(), tr, tr) + assert "timing_overrun" in res diff --git a/lib/python/tests/selector/test_oort_selector.py b/lib/python/tests/selector/test_oort_selector.py index 5fe76243a..0956e4b67 100644 --- a/lib/python/tests/selector/test_oort_selector.py +++ b/lib/python/tests/selector/test_oort_selector.py @@ -404,3 +404,68 @@ def test_fallback_to_eligible_when_no_connected_pool( ends=make_ends(["t1"]), concurrency=1, channel_props={}, ) assert async_oort.selected_ends["agg"] == {"t1"} + + +class TestPendingCommitExcludedFromSelection: + """async_oort must exclude the aggregator's VIRTUAL in-flight set + (`_agg_pending_commit_ref`, bound live to the fwdllm aggregator's + `_sim_pending_commit`) from selection eligibility -- so a returned-but- + uncommitted trainer is NEVER re-dispatched even when `all_selected` has been + pruned by a physical event (the recv-fifo re-select loop / RECVD-NONE cleanup + a slow sim triggers). This hardens over the all_selected-only guard (an unset + ref -> byte-identical). + """ + + def test_pending_end_excluded_even_when_all_selected_empty( + self, async_oort, make_ends + ): + async_oort.requester = "agg" + async_oort.selected_ends = {"agg": set()} + async_oort.all_selected = {} # guard pruned by a physical event + async_oort._agg_pending_commit_ref = {"t3"} # agg still holds t3 in flight + ends = make_ends(count=5, prefix="t") + + result = async_oort._handle_send_state( + ends=ends, concurrency=5, channel_props={"round": 1}, + trainer_unavail_list=[], task_to_perform="train", + agg_version_state=(1, 0, 0), trainer_version_states={}, + connected_ends=ends, + ) + + # t3 is still outstanding in virtual time -> must not be re-picked. + assert "t3" not in result + # the other four are eligible -> selection still works (no over-restriction). + assert len(result) >= 1 + assert set(result).issubset({"t0", "t1", "t2", "t4"}) + + def test_all_pending_yields_no_selection(self, async_oort, make_ends): + async_oort.requester = "agg" + async_oort.selected_ends = {"agg": set()} + async_oort.all_selected = {} + ends = make_ends(count=4, prefix="t") + async_oort._agg_pending_commit_ref = set(ends) # every trainer in flight + + result = async_oort._handle_send_state( + ends=ends, concurrency=4, channel_props={"round": 1}, + trainer_unavail_list=[], task_to_perform="train", + agg_version_state=(1, 0, 0), trainer_version_states={}, + connected_ends=ends, + ) + assert result == {} # nobody eligible -> no re-dispatch-while-in-flight + + def test_empty_ref_is_byte_identical_noop(self, async_oort, make_ends): + """Real / async_cifar10 never populate the ref -> selection is unchanged.""" + async_oort.requester = "agg" + async_oort.selected_ends = {"agg": set()} + async_oort.all_selected = {} + # no _agg_pending_commit_ref attribute at all -> getattr default {} + ends = make_ends(count=5, prefix="t") + + result = async_oort._handle_send_state( + ends=ends, concurrency=5, channel_props={"round": 1}, + trainer_unavail_list=[], task_to_perform="train", + agg_version_state=(1, 0, 0), trainer_version_states={}, + connected_ends=ends, + ) + assert len(result) >= 1 + assert set(result).issubset(set(ends)) diff --git a/lib/python/tests/telemetry/test_event_schema.py b/lib/python/tests/telemetry/test_event_schema.py index 86748daa3..029dd687d 100644 --- a/lib/python/tests/telemetry/test_event_schema.py +++ b/lib/python/tests/telemetry/test_event_schema.py @@ -73,6 +73,18 @@ def test_agg_round_lists(self): assert f["staleness"] == [0, 1, 2] assert f["agg_goal"] == 3 + def test_step_timing_required_and_optional(self): + from flame.telemetry.events import EVENT_STEP_TIMING, build_step_timing + ev, f = build_step_timing( + func="_train_one_batch", duration_s=1.25, + round_num=1, data_id=0, iteration=2, trainer_id="t3") + assert ev == EVENT_STEP_TIMING + assert f["func"] == "_train_one_batch" and f["duration_s"] == 1.25 + assert f["data_id"] == 0 and f["iteration_per_data_id"] == 2 + # optional context omitted when absent + _ev2, f2 = build_step_timing(func="x", duration_s=0.1) + assert "data_id" not in f2 and "trainer_id" not in f2 + def test_util_disparity_ratio_and_fraction(self): ev, f = build_util_disparity( round_num=1, @@ -117,6 +129,64 @@ def test_full_pool_higher_utility_gives_ratio_below_one(self): assert f["utility_ratio"] < 1.0 +# --- timer_decorator step_timing emission ---------------------------------- + + +class TestTimerDecoratorTelemetry: + """timer_decorator emits a step_timing record for a genuine trainer step + (self has fwd_llm_stage) and stays silent for a nested helper whose first + arg is not the trainer (e.g. a torch device) — so no mis-attributed record. + """ + + def _timed(self): + from flame.monitor.runtime import timer_decorator + + @timer_decorator + def _train_one_batch(obj, x): + return x + 1 + + return _train_one_batch + + def test_emits_with_stage(self, tmp_path): + from flame.monitor.runtime import FwdLLMStage + + telemetry.configure(role="trainer", end_id="t9", run_dir=str(tmp_path)) + + class _Trainer: + fwd_llm_stage = FwdLLMStage(1, 0, 2, trainer_id="t9") + + assert self._timed()(_Trainer(), 41) == 42 + telemetry.shutdown() + recs = [r for r in _read(tmp_path / "trainer_t9.jsonl") + if r["event"] == "step_timing"] + assert len(recs) == 1 + r = recs[0] + assert r["func"] == "_train_one_batch" and r["data_id"] == 0 + assert r["iteration_per_data_id"] == 2 and "duration_s" in r + + def test_silent_without_stage(self, tmp_path): + telemetry.configure(role="trainer", end_id="t0", run_dir=str(tmp_path)) + + class _NoStage: # e.g. a nested helper's args[0] (a device) — no stage + pass + + self._timed()(_NoStage(), 1) + telemetry.shutdown() + recs = [r for r in _read(tmp_path / "trainer_t0.jsonl") + if r["event"] == "step_timing"] + assert recs == [] + + def test_noop_when_telemetry_disabled(self): + from flame.monitor.runtime import FwdLLMStage + + telemetry.shutdown() # unconfigured -> emit is a no-op, must not raise + + class _Trainer: + fwd_llm_stage = FwdLLMStage(1, 0, 0) + + assert self._timed()(_Trainer(), 7) == 8 + + # --- selector emission produces a consistent schema ------------------------