pytscope 0.2.1__py3-none-any.whl

pytscope/__init__.py

@@ -0,0 +1,57 @@
+"""pytscope — an intelligence layer for ML training.
+
+One telemetry backbone (aligned per-step records) feeds pluggable analyzers
+(timing today; memory, convergence, reproducibility next) and a diagnosis engine
+that turns raw numbers into ranked, actionable findings.
+"""
+
+from __future__ import annotations
+
+from .analyzers.convergence import ConvergenceSummary, analyze_convergence
+from .analyzers.distributed import (
+    DistributedSummary,
+    analyze_distributed,
+    load_multirank,
+)
+from .analyzers.efficiency import EfficiencyBudget, analyze_efficiency
+from .analyzers.memory import MemorySummary, analyze_memory
+from .analyzers.pipeline import PipelineSummary, analyze_pipeline
+from .analyzers.timing import TimingSummary, analyze_timing
+from .analyzers.trace import TraceSummary, analyze_trace, analyze_trace_file
+from .auto import AutoProfiler
+from .core.events import StepRecord
+from .core.store import RunStore
+from .diagnosis.engine import DiagnosisContext, Finding, run_diagnosis
+from .hardware import measure_flops, peak_flops_for
+from .profiler import Profiler
+
+__version__ = "0.2.1"
+
+__all__ = [
+    "Profiler",
+    "AutoProfiler",
+    "RunStore",
+    "StepRecord",
+    "analyze_timing",
+    "TimingSummary",
+    "analyze_memory",
+    "MemorySummary",
+    "analyze_convergence",
+    "ConvergenceSummary",
+    "analyze_distributed",
+    "DistributedSummary",
+    "load_multirank",
+    "analyze_pipeline",
+    "PipelineSummary",
+    "analyze_trace",
+    "analyze_trace_file",
+    "TraceSummary",
+    "analyze_efficiency",
+    "EfficiencyBudget",
+    "measure_flops",
+    "peak_flops_for",
+    "run_diagnosis",
+    "DiagnosisContext",
+    "Finding",
+    "__version__",
+]

pytscope/analyzers/__init__.py

@@ -0,0 +1,15 @@
+from .convergence import ConvergenceSummary, analyze_convergence
+from .memory import MemorySummary, analyze_memory
+from .repro import RunDiff, diff_runs
+from .timing import TimingSummary, analyze_timing
+
+__all__ = [
+    "analyze_timing",
+    "TimingSummary",
+    "analyze_memory",
+    "MemorySummary",
+    "analyze_convergence",
+    "ConvergenceSummary",
+    "diff_runs",
+    "RunDiff",
+]

pytscope/analyzers/convergence.py

@@ -0,0 +1,77 @@
+"""Convergence analyzer (vertical #3) — reads per-step ``scalars``.
+
+Consumes loss / grad_norm / lr that the user logs via ``prof.log(...)`` (or the
+Lightning callback's loss). Pure functions over the timeline.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+from ..core.events import StepRecord
+from .stats import local_spikes, median
+
+
+@dataclass
+class ConvergenceSummary:
+    has_loss: bool = False
+    has_grad_norm: bool = False
+    n_steps: int = 0
+    final_loss: float | None = None
+    best_loss: float | None = None
+    loss_trend: str = "unknown"  # improving | plateau | worsening | diverged | unknown
+    diverged_at: int | None = None  # step index of first non-finite loss
+    loss_spikes: list[int] = field(default_factory=list)  # step indices
+    grad_norm_spikes: list[int] = field(default_factory=list)
+
+
+def _trend(losses: list[float], tol: float = 0.02) -> str:
+    n = len(losses)
+    if n < 4:
+        return "unknown"
+    w = max(1, n // 10)
+    first = median(losses[:w])
+    last = median(losses[-w:])
+    denom = abs(first) if first != 0 else 1.0
+    rel = (last - first) / denom
+    if rel < -tol:
+        return "improving"
+    if rel > tol:
+        return "worsening"
+    return "plateau"
+
+
+def analyze_convergence(steps: list[StepRecord]) -> ConvergenceSummary:
+    loss_pairs = [(s.step, s.scalars["loss"]) for s in steps if "loss" in s.scalars]
+    grad_series = [s.scalars.get("grad_norm") for s in steps]
+    has_grad = any(g is not None for g in grad_series)
+
+    if not loss_pairs:
+        return ConvergenceSummary(
+            has_loss=False, has_grad_norm=has_grad, n_steps=len(steps)
+        )
+
+    steps_idx = [st for st, _ in loss_pairs]
+    losses = [lv for _, lv in loss_pairs]
+
+    diverged_at = next((st for st, lv in loss_pairs if not math.isfinite(lv)), None)
+    finite = [lv for lv in losses if math.isfinite(lv)]
+
+    # Spike indices map back to the original step numbers.
+    loss_spike_pos = local_spikes(losses)
+    loss_spikes = [steps_idx[i] for i in sorted(loss_spike_pos)]
+    grad_spike_pos = local_spikes(grad_series) if has_grad else set()
+    grad_spikes = [steps[i].step for i in sorted(grad_spike_pos)]
+
+    return ConvergenceSummary(
+        has_loss=True,
+        has_grad_norm=has_grad,
+        n_steps=len(steps),
+        final_loss=finite[-1] if finite else None,
+        best_loss=min(finite) if finite else None,
+        loss_trend="diverged" if diverged_at is not None else _trend(finite),
+        diverged_at=diverged_at,
+        loss_spikes=loss_spikes,
+        grad_norm_spikes=grad_spikes,
+    )

pytscope/analyzers/distributed.py

@@ -0,0 +1,230 @@
+"""Distributed (data-parallel) analyzer — the headline vertical.
+
+In synchronous data-parallel training every rank must reach the gradient
+all-reduce before *any* rank can proceed. The step is therefore gated by the
+**slowest** rank's compute (the critical path); every faster rank sits idle at
+the barrier. That idle time is pure waste, and it is invisible to any
+single-rank profiler — you only see it by putting all ranks on one timeline.
+
+This module aligns the per-rank step timelines and computes:
+
+- **Critical-path wall loss** — wall time lost because the step waits for the
+  slowest rank instead of the average rank.
+- **Straggler attribution with a persistence test** — *which* rank is slow and,
+  crucially, whether it is a *consistent* straggler (a bad GPU/node) or just
+  noise. We test the count of steps a rank is the critical path against the
+  null hypothesis that the slowest rank is uniformly random (Binomial(S, 1/N)),
+  via a one-sided normal approximation z-score.
+- **Load imbalance** — robust coefficient of variation of per-rank compute.
+- **Communication fraction** — share of step time spent in collectives.
+- **Sync skew** — how far ahead the fastest ranks arrive at the barrier.
+
+Everything is pure-stdlib and numerically careful (``math.fsum``, robust
+medians from :mod:`pytscope.analyzers.stats`).
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from ..core.events import COMM, COMPUTE_PHASES, StepRecord
+from ..core.store import RunStore
+from .stats import median
+
+
+@dataclass
+class RankStat:
+    rank: int
+    mean_compute: float  # mean per-step local compute (non-comm), seconds
+    median_compute: float
+    slowest_count: int  # steps where this rank was the critical path
+    slowest_fraction: float
+    straggler_z: float  # z-score of slowest_count vs Binomial(S, 1/N)
+    rel_slowdown: float  # median(this_rank_compute / step_median_compute) - 1
+
+
+@dataclass
+class DistributedSummary:
+    world_size: int
+    n_steps: int  # number of aligned steps analyzed
+    mean_step_wall: float  # mean critical-path step time (max compute + comm)
+    mean_comm_fraction: float  # comm / step wall
+    imbalance_cv: float  # median over steps of CV(per-rank compute)
+    sync_skew: float  # median over steps of (max - median) compute, seconds
+    wall_frac_lost_to_imbalance: float  # (max-mean)/max summed over steps
+    ranks: list[RankStat] = field(default_factory=list)
+    straggler: RankStat | None = None  # worst *persistent* straggler, if any
+
+    @property
+    def has_straggler(self) -> bool:
+        return self.straggler is not None
+
+
+# --- loading -------------------------------------------------------------
+
+
+def load_multirank(run_dir) -> dict[int, RunStore]:
+    """Load a distributed run written by ``Profiler(distributed=True)``.
+
+    Layout is ``run_dir/rank{k}/{steps.jsonl,run.json}``. Returns a mapping of
+    rank -> RunStore. Returns ``{}`` if the directory has no per-rank subdirs.
+    """
+    run_dir = Path(run_dir)
+    ranks: dict[int, RunStore] = {}
+    if not run_dir.is_dir():
+        return ranks
+    for child in sorted(run_dir.iterdir()):
+        if child.is_dir() and child.name.startswith("rank"):
+            suffix = child.name[len("rank") :]
+            if suffix.isdigit():
+                ranks[int(suffix)] = RunStore.load(child)
+    return ranks
+
+
+def is_multirank(run_dir) -> bool:
+    run_dir = Path(run_dir)
+    if not run_dir.is_dir():
+        return False
+    return any(
+        c.is_dir() and c.name.startswith("rank") and c.name[len("rank") :].isdigit()
+        for c in run_dir.iterdir()
+    )
+
+
+# --- helpers -------------------------------------------------------------
+
+
+def _compute_seconds(rec: StepRecord) -> float:
+    """Local compute time for a step: everything that is not communication."""
+    return math.fsum(v for p, v in rec.phases.items() if p in COMPUTE_PHASES)
+
+
+def _comm_seconds(rec: StepRecord) -> float:
+    return float(rec.phases.get(COMM, 0.0))
+
+
+def _mean(xs: list[float]) -> float:
+    return math.fsum(xs) / len(xs) if xs else 0.0
+
+
+def _cv(xs: list[float]) -> float:
+    """Coefficient of variation (std / mean), population std. 0 if mean<=0."""
+    n = len(xs)
+    if n < 2:
+        return 0.0
+    m = _mean(xs)
+    if m <= 0:
+        return 0.0
+    var = math.fsum((x - m) ** 2 for x in xs) / n
+    return math.sqrt(var) / m
+
+
+# --- the analyzer --------------------------------------------------------
+
+
+def analyze_distributed(
+    ranks: dict[int, RunStore],
+    straggler_z: float = 3.0,
+    min_rel_slowdown: float = 0.05,
+) -> DistributedSummary | None:
+    """Analyze aligned per-rank timelines.
+
+    ``straggler_z``: z-score threshold for calling a rank a *persistent*
+    straggler (3.0 ≈ p<0.002, one-sided). ``min_rel_slowdown``: also require the
+    flagged rank to be at least this much slower than the per-step median, so we
+    don't flag a statistically-consistent-but-negligible straggler.
+    """
+    if len(ranks) < 2:
+        return None
+
+    rank_ids = sorted(ranks)
+    n = len(rank_ids)
+    # Per-rank step index -> record, so we can align by the shared step number.
+    by_rank: dict[int, dict[int, StepRecord]] = {
+        r: {rec.step: rec for rec in ranks[r].steps} for r in rank_ids
+    }
+    common_steps = sorted(set.intersection(*[set(by_rank[r]) for r in rank_ids]))
+    if not common_steps:
+        return None
+
+    # Accumulators.
+    step_walls: list[float] = []
+    comm_fracs: list[float] = []
+    imbalance_cvs: list[float] = []
+    sync_skews: list[float] = []
+    slowest_count = {r: 0 for r in rank_ids}
+    rel_slowdowns: dict[int, list[float]] = {r: [] for r in rank_ids}
+    compute_samples: dict[int, list[float]] = {r: [] for r in rank_ids}
+    sum_max = 0.0
+    sum_lost = 0.0  # sum of (max - mean) compute
+
+    for s in common_steps:
+        comps = {r: _compute_seconds(by_rank[r][s]) for r in rank_ids}
+        comms = {r: _comm_seconds(by_rank[r][s]) for r in rank_ids}
+        cvals = [comps[r] for r in rank_ids]
+        cmax = max(cvals)
+        cmean = _mean(cvals)
+        cmed = median(cvals)
+        # Critical path: slowest compute, plus the max comm observed that step.
+        wall = cmax + max(comms.values())
+        step_walls.append(wall)
+        comm_fracs.append((max(comms.values()) / wall) if wall > 0 else 0.0)
+        imbalance_cvs.append(_cv(cvals))
+        sync_skews.append(cmax - cmed)
+        sum_max += cmax
+        sum_lost += cmax - cmean
+
+        # Critical rank (argmax compute); ties -> lowest rank id (deterministic).
+        crit = min(rank_ids, key=lambda r: (-comps[r], r))
+        slowest_count[crit] += 1
+        for r in rank_ids:
+            compute_samples[r].append(comps[r])
+            if cmed > 0:
+                rel_slowdowns[r].append(comps[r] / cmed - 1.0)
+
+    s_total = len(common_steps)
+    # Binomial(S, 1/N) null for "is rank r the slowest".
+    p0 = 1.0 / n
+    sd = math.sqrt(s_total * p0 * (1.0 - p0)) or 1.0
+    expected = s_total * p0
+
+    rank_stats: list[RankStat] = []
+    for r in rank_ids:
+        z = (slowest_count[r] - expected) / sd
+        rank_stats.append(
+            RankStat(
+                rank=r,
+                mean_compute=_mean(compute_samples[r]),
+                median_compute=median(compute_samples[r]),
+                slowest_count=slowest_count[r],
+                slowest_fraction=slowest_count[r] / s_total,
+                straggler_z=z,
+                rel_slowdown=median(rel_slowdowns[r]) if rel_slowdowns[r] else 0.0,
+            )
+        )
+
+    # A *single* persistent straggler exists only if exactly one rank clears both
+    # the statistical-persistence bar (z) and the practical-magnitude bar
+    # (rel_slowdown). If several ranks do (e.g. two slow nodes alternating as the
+    # critical path), that's load imbalance, not one straggler — leave it to the
+    # imbalance rule rather than fingering one rank misleadingly.
+    significant = [
+        rs
+        for rs in rank_stats
+        if rs.straggler_z >= straggler_z and rs.rel_slowdown >= min_rel_slowdown
+    ]
+    straggler = significant[0] if len(significant) == 1 else None
+
+    return DistributedSummary(
+        world_size=n,
+        n_steps=s_total,
+        mean_step_wall=_mean(step_walls),
+        mean_comm_fraction=_mean(comm_fracs),
+        imbalance_cv=median(imbalance_cvs),
+        sync_skew=median(sync_skews),
+        wall_frac_lost_to_imbalance=(sum_lost / sum_max) if sum_max > 0 else 0.0,
+        ranks=rank_stats,
+        straggler=straggler,
+    )

pytscope/analyzers/efficiency.py

@@ -0,0 +1,149 @@
+"""The Training Efficiency Budget — one accounting identity for the whole run.
+
+Every other analyzer produces a *finding*. This one produces a **budget**: it
+decomposes the wall time of training into named line items that, by construction,
+sum back to the measured wall time. Anchored at the top by ``useful_compute`` —
+the time the math would take at the hardware's peak throughput — the rest of the
+budget is, line by line, *recoverable* time:
+
+    wall = useful_compute            (irreducible: the FLOPs, at peak)
+         + compute_overhead          (your kernels don't hit peak)
+         + data_stall                (waiting on the dataloader)
+         + communication             (collective time on the timeline)
+         + other                     (everything else attributed)
+
+``MFU`` (Model FLOPs Utilization) falls straight out: ``useful_compute / wall``.
+Because the phase timeline already partitions each step, the decomposition is
+**exact** — the line items sum to the attributed wall with no fudge factor, which
+makes the model falsifiable (a wrong term shows up as a non-zero residual).
+
+This turns a profiler into an advisor: each recoverable line is a number of
+seconds you could win back, so fixes rank themselves by payoff.
+
+FLOPs and peak are optional. With them you get a true MFU anchor; without them
+the budget still decomposes wall time, with ``useful_compute`` falling back to
+measured compute (MFU unknown).
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+
+from ..core.events import (
+    COMM,
+    COMPUTE,
+    DATA,
+    OTHER,
+    StepRecord,
+)
+
+# Phases that count as local compute (where FLOPs are spent).
+_COMPUTE_PHASES = ("forward", "backward", "optimizer", COMPUTE)
+
+
+@dataclass
+class BudgetLine:
+    name: str
+    seconds: float
+    fraction: float  # of attributed wall
+    recoverable: bool  # is this time you could win back?
+
+
+@dataclass
+class EfficiencyBudget:
+    wall: float  # attributed wall (sum of all phases over the run)
+    n_steps: int
+    compute_measured: float  # forward+backward+optimizer+compute, seconds
+    ideal_compute: float | None  # FLOPs-anchored minimum, seconds (if known)
+    mfu: float | None  # useful_compute / wall, in [0, 1] (if FLOPs+peak known)
+    flops_per_step: float | None
+    peak_flops: float | None
+    lines: list[BudgetLine] = field(default_factory=list)
+
+    @property
+    def efficiency(self) -> float:
+        """Useful fraction of wall (== MFU when FLOPs+peak are known)."""
+        useful = next(
+            (ln.seconds for ln in self.lines if ln.name == "useful_compute"), 0.0
+        )
+        return useful / self.wall if self.wall > 0 else 0.0
+
+    @property
+    def recoverable_lines(self) -> list[BudgetLine]:
+        """Recoverable line items, largest first — the ranked fix list."""
+        return sorted(
+            (ln for ln in self.lines if ln.recoverable and ln.seconds > 0),
+            key=lambda ln: ln.seconds,
+            reverse=True,
+        )
+
+    @property
+    def top_recoverable(self) -> BudgetLine | None:
+        rec = self.recoverable_lines
+        return rec[0] if rec else None
+
+
+def _phase_total(steps: Sequence[StepRecord], phase: str) -> float:
+    return math.fsum(s.phases.get(phase, 0.0) for s in steps)
+
+
+def analyze_efficiency(
+    steps: Sequence[StepRecord],
+    *,
+    flops_per_step: float | None = None,
+    peak_flops: float | None = None,
+) -> EfficiencyBudget | None:
+    """Build the efficiency budget from the per-step phase timeline.
+
+    ``flops_per_step``: training FLOPs per step (forward+backward; see
+    ``hardware.measure_flops``). ``peak_flops``: device peak FLOP/s. Both
+    optional — supply them for a true MFU anchor.
+    """
+    steps = list(steps)
+    if not steps:
+        return None
+
+    data = _phase_total(steps, DATA)
+    comm = _phase_total(steps, COMM)
+    other = _phase_total(steps, OTHER)
+    compute_measured = math.fsum(_phase_total(steps, p) for p in _COMPUTE_PHASES)
+    wall = compute_measured + data + comm + other
+    if wall <= 0:
+        return None
+
+    n = len(steps)
+    ideal_compute: float | None = None
+    mfu: float | None = None
+    if flops_per_step and peak_flops and peak_flops > 0:
+        ideal_compute = flops_per_step * n / peak_flops
+
+    # Useful compute is the irreducible work. If we know the FLOPs-anchored ideal,
+    # it caps useful at the measured compute (you can't be "more than 100%
+    # efficient"; if the estimate exceeds measured, treat compute as all-useful
+    # and surface no negative overhead).
+    if ideal_compute is not None:
+        useful = min(ideal_compute, compute_measured)
+        mfu = useful / wall
+    else:
+        useful = compute_measured
+    overhead = compute_measured - useful
+
+    lines = [
+        BudgetLine("useful_compute", useful, useful / wall, recoverable=False),
+        BudgetLine("compute_overhead", overhead, overhead / wall, recoverable=True),
+        BudgetLine("data_stall", data, data / wall, recoverable=True),
+        BudgetLine("communication", comm, comm / wall, recoverable=True),
+        BudgetLine("other", other, other / wall, recoverable=True),
+    ]
+    return EfficiencyBudget(
+        wall=wall,
+        n_steps=n,
+        compute_measured=compute_measured,
+        ideal_compute=ideal_compute,
+        mfu=mfu,
+        flops_per_step=flops_per_step,
+        peak_flops=peak_flops,
+        lines=lines,
+    )

pytscope/analyzers/memory.py

@@ -0,0 +1,58 @@
+"""Memory analyzer (vertical #2) — reads the per-step ``memory`` block.
+
+Operates on whatever the memory collector stored (bytes): ``alloc``,
+``reserved``, ``peak_alloc``, ``peak_reserved``. Pure functions over the
+timeline, like the timing analyzer.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+from ..core.events import StepRecord
+from .stats import robust_slope
+
+
+@dataclass
+class MemorySummary:
+    has_memory: bool = False
+    n_steps: int = 0
+    peak_alloc_bytes: float = 0.0
+    peak_reserved_bytes: float = 0.0
+    mean_alloc_bytes: float = 0.0
+    # Fraction of reserved memory not actually allocated — fragmentation/slack.
+    fragmentation: float = 0.0
+    # Robust per-step growth of allocated memory — a leak signal.
+    growth_bytes_per_step: float = 0.0
+    alloc_series: list[float] = field(default_factory=list)
+
+
+def analyze_memory(steps: list[StepRecord]) -> MemorySummary:
+    allocs = [s.memory["alloc"] for s in steps if "alloc" in s.memory]
+    if not allocs:
+        return MemorySummary(has_memory=False, n_steps=len(steps))
+
+    reserved = [s.memory["reserved"] for s in steps if "reserved" in s.memory]
+    peak_alloc = [
+        s.memory.get("peak_alloc", s.memory.get("alloc", 0.0)) for s in steps if s.memory
+    ]
+    peak_reserved = [
+        s.memory.get("peak_reserved", s.memory.get("reserved", 0.0))
+        for s in steps
+        if s.memory
+    ]
+
+    frag_samples = [(r - a) / r for a, r in zip(allocs, reserved) if r > 0]
+    fragmentation = math.fsum(frag_samples) / len(frag_samples) if frag_samples else 0.0
+
+    return MemorySummary(
+        has_memory=True,
+        n_steps=len(steps),
+        peak_alloc_bytes=max(peak_alloc) if peak_alloc else max(allocs),
+        peak_reserved_bytes=max(peak_reserved) if peak_reserved else 0.0,
+        mean_alloc_bytes=math.fsum(allocs) / len(allocs),
+        fragmentation=fragmentation,
+        growth_bytes_per_step=robust_slope(allocs),
+        alloc_series=allocs,
+    )