train-doctor 0.1.0__tar.gz

train_doctor-0.1.0/.gitignore

@@ -0,0 +1,11 @@
+.venv/
+__pycache__/
+*.pyc
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+dist/
+build/
+*.egg-info/
+.train-doctor/
+.DS_Store

train_doctor-0.1.0/CHANGELOG.md

@@ -0,0 +1,13 @@
+# Changelog
+
+## 0.1.0 (unreleased)
+
+First version.
+
+- `profile`: runs a training command for a bounded window (warmup plus at least N steps and T seconds) with an in-process recorder injected through `sitecustomize`. Records step times, DataLoader wait, host/device sync calls, `torch.save` and log-write time, samples per step, loss, DataLoader and precision config, process CPU and memory, and GPU telemetry (nvidia-smi, or ioreg on Apple silicon). Optional torch.profiler steps after the window and py-spy sampling.
+- `diagnose`: deterministic rules for DataLoader-bound runs, expensive per-sample preprocessing, pin_memory on CUDA, host syncs in the loop, small batches, missing mixed precision, missing torch.compile, checkpoints and logging in the hot loop, gradient accumulation (including DDP without no_sync), cuDNN autotuning, and CPU training with an accelerator available.
+- `compare`: discarded warmup runs, K repeats per arm in randomized pair order, same seed, median and spread, bootstrap interval on the ratio of medians, a verdict that refuses to call changes inside the noise, and a loss-trajectory check.
+- `report`: Markdown and JSON bundles with machine, versions, commands, per-repeat numbers, findings, verdict and limits.
+- `setup`: registers the MCP server with Claude Code, Codex and Cursor and installs the skill, with a dry run, backups and idempotent re-runs.
+- `--label` names run directories; unstartable commands give a clear error instead of a traceback.
+- MCP server over stdio with `profile`, `diagnose`, `compare`, `report` and `version` tools.

train_doctor-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Abel Yagubyan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

train_doctor-0.1.0/PKG-INFO

@@ -0,0 +1,163 @@
+Metadata-Version: 2.5
+Name: train-doctor
+Version: 0.1.0
+Summary: Give your agent a way to find out why a training run is slow, fix it, and prove the speedup.
+Project-URL: Homepage, https://github.com/Abelo9996/train-doctor
+Author: Abel Yagubyan
+License-Expression: MIT
+License-File: LICENSE
+Keywords: agents,benchmark,mcp,profiling,pytorch,training
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Testing
+Requires-Python: >=3.11
+Requires-Dist: mcp>=1.2
+Requires-Dist: psutil>=5.9
+Description-Content-Type: text/markdown
+
+# train-doctor
+
+Give your agent a way to find out why your training run is slow, fix it, and prove the speedup.
+
+```bash
+uvx train-doctor profile -- python train.py        # where does step time go?
+uvx train-doctor compare --baseline "python train.py" --candidate "python train.py --num-workers 4"
+uvx train-doctor setup                             # give Claude Code, Codex and Cursor the MCP tools
+```
+
+> Not on PyPI yet. Until the first release, run it straight from GitHub by replacing `uvx train-doctor` with
+> `uvx --from git+https://github.com/Abelo9996/train-doctor train-doctor`. `setup` registers `uvx train-doctor mcp`, so it works once the
+> package is on PyPI.
+
+## Example output
+
+From a real run in this repo: [examples/cnn_images](examples/cnn_images), a small CNN written the way many
+first scripts are (per-sample augmentation in `__getitem__`, `num_workers=0`, `.item()` and a print on every
+step). Apple M4 (10 cores, 16 GB), torch 2.14.1 on MPS, with other jobs running on the machine (1-minute
+load average between 4.6 and 10.4 across all runs).
+
+Profile output (trimmed where marked `...`):
+
+```text
+$ train-doctor profile -- python examples/cnn_images/train.py
+source: hook, device: mps, steps measured: 87
+step time: median 20.88 ms (p10 14.32, p90 32.84)
+throughput: 2766.8 samples/s
+split: data_wait 39.1%, host_sync 26.4%, checkpoint_save 0.0%, log_write 0.1%, other 34.4%
+findings:
+  1. [dataloader_bound] Training waits on the DataLoader (num_workers=0) (score 0.391, risk low)
+  2. [cpu_bound_preprocessing] Per-sample preprocessing in SyntheticImages is expensive (score 0.319, risk low)
+     evidence: 39% of step time is data wait and 91% of that is inside SyntheticImages.__getitem__ (0.13 ms per sample).
+  3. [host_sync_in_loop] Host waits on the device inside every step (score 0.132, risk none)
+     evidence: 2.0 blocking reads of mps tensors per step (...), 6.1 ms per step blocked in them (26% of step time).
+  4. [no_mixed_precision] ...  5. [logging_in_hot_loop] ...  6. [no_torch_compile] ...
+```
+
+Then one change at a time, each checked with `train-doctor compare` (5 measured repeats per side after one
+discarded warmup run, randomized order, 95% bootstrap interval on the ratio of median samples/s):
+
+| Change | Median samples/s, before to after | Ratio (95% interval) | Loss check | Decision |
+|---|---|---|---|---|
+| `num_workers=5` | 2,249 to 2,750 | 1.22x (0.66 to 1.68) | within tolerance | inconclusive, dropped |
+| crop/flip/normalize per batch on the device | 2,232 to 2,855 | 1.28x (1.10 to 1.35) | within tolerance | keep |
+| read metrics every 50 steps instead of every step | 2,420 to 3,029 | 1.25x (1.11 to 1.35) | identical | keep |
+| autocast float16 | 2,648 to 2,842 | 1.07x (0.81 to 1.28) | within tolerance | inconclusive, dropped |
+| original vs the two kept changes | 2,091 to 3,381 | 1.62x (1.37 to 3.48) | within tolerance | keep |
+
+Two of the four candidates looked faster by their medians and were still not kept, because their intervals
+included 1.0. In the second example ([examples/tabular_mlp](examples/tabular_mlp)), checkpointing less often
+was inconclusive with 5 repeats (1.54x, interval 0.62 to 2.67) and clearly faster with 9 (2.01x, 1.54 to 2.52);
+original vs final there is 3.05x (2.15 to 3.40) with an identical loss trajectory. Every report, per-repeat
+number and raw trace is committed under `examples/*/runs/`, and
+[examples/run_validation.sh](examples/run_validation.sh) is the exact command sequence.
+
+## How it works
+
+A "2x faster" claim is only useful if you can trust it. train-doctor excludes warmup, brackets the
+measured window with a device synchronize, runs baseline and candidate several times in randomized
+interleaved order with the same seed, reports medians, spread and a bootstrap interval, and calls a
+change "faster" only when the whole interval clears a threshold. It also checks that the loss
+trajectory didn't move, so a speedup that broke training is rejected.
+
+train-doctor runs your real training command as a child process and stops it after a bounded window
+(warmup steps, then at least N steps and T seconds; defaults 5, 50 and 2 s). No code changes are needed
+for PyTorch:
+
+- **Step recorder.** A `sitecustomize.py` on the child's PYTHONPATH loads a small standard-library
+  recorder that patches torch right after it is imported. It marks a step at every `optimizer.step()`
+  (a global post-hook), times `DataLoader.__next__`, times blocking reads of GPU/MPS tensors
+  (`.item()`, `.cpu()`, `.tolist()`, `float()`, `bool()`, `synchronize()`), times `torch.save` and
+  stdout/stderr writes, counts `backward()` calls, and keeps the loss tensor passed to `backward()`
+  (converted to a number only at the end, so recording it adds no sync). It also notes DataLoader
+  settings, autocast, GradScaler, `torch.compile`, DDP, parameter count and dtype. Any other
+  `sitecustomize` you have still runs. Loops without an optimizer (or other frameworks) can call
+  `train_doctor.step(samples=..., loss=...)`, or just print a line containing `loss` each step.
+- **torch.profiler** traces a few extra steps after the measured window (so it doesn't distort the
+  timings) and saves a Chrome trace you can open in Perfetto.
+- **py-spy** (optional, `--py-spy`) samples Python stacks for any framework when it is installed.
+- **System telemetry**: psutil for CPU and memory of the whole process tree (including DataLoader
+  workers), `nvidia-smi` for NVIDIA GPUs, and `ioreg` for Apple GPU utilization (no sudo needed).
+- **Diagnosis** is a fixed set of rules over that evidence: DataLoader-bound, expensive per-sample
+  preprocessing, pin_memory off on CUDA, host syncs in the loop, small batch on an idle device, no
+  mixed precision on capable hardware, no torch.compile, checkpoints or logging in the hot loop,
+  gradient accumulation (including DDP without `no_sync()`), cuDNN autotuning off, and training on CPU
+  while an accelerator is available. Each finding carries the numbers it fired on, a suggested change,
+  the expected effect (an upper bound where one can be computed) and the risk to results.
+- **Compare** runs one discarded warmup run per command, then K repeats per command (default 5) in
+  randomized pair order. Throughput is samples/s over the synchronized window. The verdict uses a 95%
+  percentile bootstrap interval on the ratio of medians: "faster" only if the whole interval is above
+  1 + min-effect (default 2%), "slower" if it's below 1 - min-effect, otherwise "no clear difference".
+  The loss check aligns both loss curves on samples seen, smooths them, and requires the mean and final
+  relative difference to stay within a tolerance (default 5%). The decision is `keep` only when both pass.
+- **Reports**: every run directory gets `report.md` and `report.json` with the machine, versions,
+  commands, file hashes, git commit, per-repeat numbers, findings, verdict and limits, plus the raw
+  files (per-step events, resource samples, stdout and stderr, profiler output).
+
+The tool is deterministic and never calls an LLM. Your coding agent reads the findings, makes one
+change, and uses `compare` to decide whether to keep it.
+
+## Setup for agents
+
+```bash
+uvx train-doctor setup          # shows what it would change
+uvx train-doctor setup --yes    # applies it
+```
+
+`setup` detects Claude Code, Codex and Cursor. It runs `claude mcp add --scope user train-doctor -- uvx train-doctor mcp`,
+adds `[mcp_servers.train-doctor]` to `~/.codex/config.toml`, adds the server to `~/.cursor/mcp.json`, and copies the
+skill to `~/.claude/skills/train-doctor/` and `~/.codex/skills/train-doctor/`. It backs up every file it edits and
+does nothing when an entry is already there. `--project DIR` also writes a project `.mcp.json`.
+
+MCP tools: `profile`, `diagnose`, `compare`, `report`, `version`. The skill
+([skills/train-doctor/SKILL.md](skills/train-doctor/SKILL.md)) tells the agent the workflow: profile, read
+the findings, change one thing, compare, keep only clear wins that don't move the loss, report.
+
+CLI: `train-doctor profile|diagnose|compare|report|setup|mcp`, each with `--help` and `--json`.
+
+## What it can't do
+
+- It doesn't edit your code. The agent (or you) makes the change; train-doctor measures it.
+- Timings come from the host side. On CUDA and MPS, time inside a sync call includes waiting for device
+  work that has to happen anyway, so the "host sync" share is an upper bound on what removing syncs saves.
+- Apple GPU utilization from `ioreg` is system-wide, not per process. There is no per-kernel timing on MPS.
+- `__getitem__` time is measured only with `num_workers=0`; with workers it runs in other processes.
+- py-spy needs root on macOS, so `--py-spy` usually fails there. It was only tested with a fake py-spy.
+- CUDA support (nvidia-smi parsing, CUDA sync and memory paths, CUDA-specific rules) is implemented and
+  tested with fakes. It has not been run on an NVIDIA GPU yet. The real runs in this repo are on Apple MPS.
+- Multi-GPU: with DDP it reads the process that recorded the most steps (usually rank 0); it doesn't merge ranks.
+- The loss check covers the first steps only. It catches changes that break or shift training early; it
+  can't tell you about final accuracy. Batch size and learning-rate changes are flagged as high risk for that reason.
+- Short windows on a busy machine are noisy. Reports record the load average and system CPU per run; when
+  the spread is high, compare says so and is more likely to answer "no clear difference".
+
+## Privacy and safety
+
+Everything runs locally. train-doctor makes no network requests and sends nothing anywhere. It runs the
+command you give it with your permissions, and writes to `.train-doctor/runs/` (or `--out`). Reports
+replace your home directory with `~`. `setup` changes only the files it lists, and only after `--yes`.
+
+## License
+
+MIT, see [LICENSE](LICENSE).

train_doctor-0.1.0/README.md

@@ -0,0 +1,144 @@
+# train-doctor
+
+Give your agent a way to find out why your training run is slow, fix it, and prove the speedup.
+
+```bash
+uvx train-doctor profile -- python train.py        # where does step time go?
+uvx train-doctor compare --baseline "python train.py" --candidate "python train.py --num-workers 4"
+uvx train-doctor setup                             # give Claude Code, Codex and Cursor the MCP tools
+```
+
+> Not on PyPI yet. Until the first release, run it straight from GitHub by replacing `uvx train-doctor` with
+> `uvx --from git+https://github.com/Abelo9996/train-doctor train-doctor`. `setup` registers `uvx train-doctor mcp`, so it works once the
+> package is on PyPI.
+
+## Example output
+
+From a real run in this repo: [examples/cnn_images](examples/cnn_images), a small CNN written the way many
+first scripts are (per-sample augmentation in `__getitem__`, `num_workers=0`, `.item()` and a print on every
+step). Apple M4 (10 cores, 16 GB), torch 2.14.1 on MPS, with other jobs running on the machine (1-minute
+load average between 4.6 and 10.4 across all runs).
+
+Profile output (trimmed where marked `...`):
+
+```text
+$ train-doctor profile -- python examples/cnn_images/train.py
+source: hook, device: mps, steps measured: 87
+step time: median 20.88 ms (p10 14.32, p90 32.84)
+throughput: 2766.8 samples/s
+split: data_wait 39.1%, host_sync 26.4%, checkpoint_save 0.0%, log_write 0.1%, other 34.4%
+findings:
+  1. [dataloader_bound] Training waits on the DataLoader (num_workers=0) (score 0.391, risk low)
+  2. [cpu_bound_preprocessing] Per-sample preprocessing in SyntheticImages is expensive (score 0.319, risk low)
+     evidence: 39% of step time is data wait and 91% of that is inside SyntheticImages.__getitem__ (0.13 ms per sample).
+  3. [host_sync_in_loop] Host waits on the device inside every step (score 0.132, risk none)
+     evidence: 2.0 blocking reads of mps tensors per step (...), 6.1 ms per step blocked in them (26% of step time).
+  4. [no_mixed_precision] ...  5. [logging_in_hot_loop] ...  6. [no_torch_compile] ...
+```
+
+Then one change at a time, each checked with `train-doctor compare` (5 measured repeats per side after one
+discarded warmup run, randomized order, 95% bootstrap interval on the ratio of median samples/s):
+
+| Change | Median samples/s, before to after | Ratio (95% interval) | Loss check | Decision |
+|---|---|---|---|---|
+| `num_workers=5` | 2,249 to 2,750 | 1.22x (0.66 to 1.68) | within tolerance | inconclusive, dropped |
+| crop/flip/normalize per batch on the device | 2,232 to 2,855 | 1.28x (1.10 to 1.35) | within tolerance | keep |
+| read metrics every 50 steps instead of every step | 2,420 to 3,029 | 1.25x (1.11 to 1.35) | identical | keep |
+| autocast float16 | 2,648 to 2,842 | 1.07x (0.81 to 1.28) | within tolerance | inconclusive, dropped |
+| original vs the two kept changes | 2,091 to 3,381 | 1.62x (1.37 to 3.48) | within tolerance | keep |
+
+Two of the four candidates looked faster by their medians and were still not kept, because their intervals
+included 1.0. In the second example ([examples/tabular_mlp](examples/tabular_mlp)), checkpointing less often
+was inconclusive with 5 repeats (1.54x, interval 0.62 to 2.67) and clearly faster with 9 (2.01x, 1.54 to 2.52);
+original vs final there is 3.05x (2.15 to 3.40) with an identical loss trajectory. Every report, per-repeat
+number and raw trace is committed under `examples/*/runs/`, and
+[examples/run_validation.sh](examples/run_validation.sh) is the exact command sequence.
+
+## How it works
+
+A "2x faster" claim is only useful if you can trust it. train-doctor excludes warmup, brackets the
+measured window with a device synchronize, runs baseline and candidate several times in randomized
+interleaved order with the same seed, reports medians, spread and a bootstrap interval, and calls a
+change "faster" only when the whole interval clears a threshold. It also checks that the loss
+trajectory didn't move, so a speedup that broke training is rejected.
+
+train-doctor runs your real training command as a child process and stops it after a bounded window
+(warmup steps, then at least N steps and T seconds; defaults 5, 50 and 2 s). No code changes are needed
+for PyTorch:
+
+- **Step recorder.** A `sitecustomize.py` on the child's PYTHONPATH loads a small standard-library
+  recorder that patches torch right after it is imported. It marks a step at every `optimizer.step()`
+  (a global post-hook), times `DataLoader.__next__`, times blocking reads of GPU/MPS tensors
+  (`.item()`, `.cpu()`, `.tolist()`, `float()`, `bool()`, `synchronize()`), times `torch.save` and
+  stdout/stderr writes, counts `backward()` calls, and keeps the loss tensor passed to `backward()`
+  (converted to a number only at the end, so recording it adds no sync). It also notes DataLoader
+  settings, autocast, GradScaler, `torch.compile`, DDP, parameter count and dtype. Any other
+  `sitecustomize` you have still runs. Loops without an optimizer (or other frameworks) can call
+  `train_doctor.step(samples=..., loss=...)`, or just print a line containing `loss` each step.
+- **torch.profiler** traces a few extra steps after the measured window (so it doesn't distort the
+  timings) and saves a Chrome trace you can open in Perfetto.
+- **py-spy** (optional, `--py-spy`) samples Python stacks for any framework when it is installed.
+- **System telemetry**: psutil for CPU and memory of the whole process tree (including DataLoader
+  workers), `nvidia-smi` for NVIDIA GPUs, and `ioreg` for Apple GPU utilization (no sudo needed).
+- **Diagnosis** is a fixed set of rules over that evidence: DataLoader-bound, expensive per-sample
+  preprocessing, pin_memory off on CUDA, host syncs in the loop, small batch on an idle device, no
+  mixed precision on capable hardware, no torch.compile, checkpoints or logging in the hot loop,
+  gradient accumulation (including DDP without `no_sync()`), cuDNN autotuning off, and training on CPU
+  while an accelerator is available. Each finding carries the numbers it fired on, a suggested change,
+  the expected effect (an upper bound where one can be computed) and the risk to results.
+- **Compare** runs one discarded warmup run per command, then K repeats per command (default 5) in
+  randomized pair order. Throughput is samples/s over the synchronized window. The verdict uses a 95%
+  percentile bootstrap interval on the ratio of medians: "faster" only if the whole interval is above
+  1 + min-effect (default 2%), "slower" if it's below 1 - min-effect, otherwise "no clear difference".
+  The loss check aligns both loss curves on samples seen, smooths them, and requires the mean and final
+  relative difference to stay within a tolerance (default 5%). The decision is `keep` only when both pass.
+- **Reports**: every run directory gets `report.md` and `report.json` with the machine, versions,
+  commands, file hashes, git commit, per-repeat numbers, findings, verdict and limits, plus the raw
+  files (per-step events, resource samples, stdout and stderr, profiler output).
+
+The tool is deterministic and never calls an LLM. Your coding agent reads the findings, makes one
+change, and uses `compare` to decide whether to keep it.
+
+## Setup for agents
+
+```bash
+uvx train-doctor setup          # shows what it would change
+uvx train-doctor setup --yes    # applies it
+```
+
+`setup` detects Claude Code, Codex and Cursor. It runs `claude mcp add --scope user train-doctor -- uvx train-doctor mcp`,
+adds `[mcp_servers.train-doctor]` to `~/.codex/config.toml`, adds the server to `~/.cursor/mcp.json`, and copies the
+skill to `~/.claude/skills/train-doctor/` and `~/.codex/skills/train-doctor/`. It backs up every file it edits and
+does nothing when an entry is already there. `--project DIR` also writes a project `.mcp.json`.
+
+MCP tools: `profile`, `diagnose`, `compare`, `report`, `version`. The skill
+([skills/train-doctor/SKILL.md](skills/train-doctor/SKILL.md)) tells the agent the workflow: profile, read
+the findings, change one thing, compare, keep only clear wins that don't move the loss, report.
+
+CLI: `train-doctor profile|diagnose|compare|report|setup|mcp`, each with `--help` and `--json`.
+
+## What it can't do
+
+- It doesn't edit your code. The agent (or you) makes the change; train-doctor measures it.
+- Timings come from the host side. On CUDA and MPS, time inside a sync call includes waiting for device
+  work that has to happen anyway, so the "host sync" share is an upper bound on what removing syncs saves.
+- Apple GPU utilization from `ioreg` is system-wide, not per process. There is no per-kernel timing on MPS.
+- `__getitem__` time is measured only with `num_workers=0`; with workers it runs in other processes.
+- py-spy needs root on macOS, so `--py-spy` usually fails there. It was only tested with a fake py-spy.
+- CUDA support (nvidia-smi parsing, CUDA sync and memory paths, CUDA-specific rules) is implemented and
+  tested with fakes. It has not been run on an NVIDIA GPU yet. The real runs in this repo are on Apple MPS.
+- Multi-GPU: with DDP it reads the process that recorded the most steps (usually rank 0); it doesn't merge ranks.
+- The loss check covers the first steps only. It catches changes that break or shift training early; it
+  can't tell you about final accuracy. Batch size and learning-rate changes are flagged as high risk for that reason.
+- Short windows on a busy machine are noisy. Reports record the load average and system CPU per run; when
+  the spread is high, compare says so and is more likely to answer "no clear difference".
+
+## Privacy and safety
+
+Everything runs locally. train-doctor makes no network requests and sends nothing anywhere. It runs the
+command you give it with your permissions, and writes to `.train-doctor/runs/` (or `--out`). Reports
+replace your home directory with `~`. `setup` changes only the files it lists, and only after `--yes`.
+
+## License
+
+MIT, see [LICENSE](LICENSE).

train_doctor-0.1.0/examples/README.md

@@ -0,0 +1,57 @@
+# Examples
+
+Two small PyTorch training scripts with realistic slow spots, and the real train-doctor runs on them.
+Every directory under `*/runs/` is untouched tool output: `report.md` for reading, `report.json` and
+`compare.json` for the numbers, and the raw per-step events, resource samples, logs and profiler output.
+
+Machine: Apple M4 MacBook (10 cores, 16 GB), macOS 26 (Darwin 25.2), Python 3.12.13, torch 2.14.1, device MPS.
+Other jobs were running on the machine during these runs (1-minute load average 4.6 to 10.4 at the start
+of each run); the per-repeat tables in the compare reports list the load and system CPU of every run.
+
+Each command-line flag on the scripts stands in for one code edit, so a single change can be compared
+on its own. [run_validation.sh](run_validation.sh) is the exact sequence of commands.
+
+## cnn_images
+
+A 3-layer CNN on 20,000 synthetic 32x32 images: per-sample crop, flip and normalize in `__getitem__`,
+`num_workers=0`, and loss and accuracy read with `.item()` and printed every step.
+
+| Run | What happened |
+|---|---|
+| `profile-baseline` | 39.1% of step time waiting on the DataLoader (91% of that in `__getitem__`), 26.4% blocked in 2 `.item()` calls per step. Top finding: `dataloader_bound`. |
+| `compare-workers-5` | `--num-workers 5`: 1.22x by medians but the interval is 0.66 to 1.68, spread up to 33%. Inconclusive, dropped. |
+| `compare-batch-augment` | Next finding, `cpu_bound_preprocessing`: augmentation moved to one vectorized call per batch on the device. 1.28x (1.10 to 1.35), loss within tolerance (1.7% mean difference). Kept. |
+| `profile-batch-augment` | Data wait down to 6.0%; host sync now 37.1% of step time. Top finding: `host_sync_in_loop`. |
+| `compare-log-every-50` | Metrics kept on the device and read every 50 steps. 1.25x (1.11 to 1.35), loss identical. Kept. |
+| `profile-log-every-50` | Host sync 0.2%, compute and Python 95.8%. Top finding: `no_mixed_precision`. |
+| `compare-amp` | `--amp` (autocast float16 on MPS): 1.07x (0.81 to 1.28). Inconclusive, dropped. |
+| `compare-original-vs-final` | Original vs `--batch-augment --log-every 50`: 1.62x (1.37 to 3.48), loss within tolerance. |
+
+## tabular_mlp
+
+A 4-layer MLP on 60,000 synthetic CSV-like rows, featurized in plain Python inside `__getitem__`
+(parse, normalize, one-hot, hashed text buckets), with a `torch.save` checkpoint every 5 steps.
+
+| Run | What happened |
+|---|---|
+| `profile-baseline` | 52.8% of step time in `torch.save`, 28.2% waiting on the DataLoader. Top finding: `checkpoint_in_hot_loop`. |
+| `compare-ckpt-every-200` | Checkpoint every 200 steps: 1.54x by medians, interval 0.62 to 2.67 (spread up to 35%). Inconclusive. |
+| `compare-ckpt-every-200-9-repeats` | Same change with 9 repeats: 2.01x (1.54 to 2.52), loss identical. Kept. |
+| `profile-ckpt-every-200` | Data wait 34.3% with `num_workers=0`. Top finding: `dataloader_bound`. |
+| `compare-workers-5` | `--num-workers 5`: 1.93x (1.54 to 2.86), loss identical. Kept. The candidate needs 6.1 s more to reach its first step (spawned workers import torch); the report states that separately. |
+| `profile-workers-5` | Data wait 7.3%; the window caught one checkpoint (8.8%). Next findings: mixed precision, torch.compile. Stopped here. |
+| `compare-original-vs-final` | Original vs `--ckpt-every 200 --num-workers 5`: 3.05x (2.15 to 3.40), loss identical. |
+
+Throughput medians differ between compares of the same command (for example the CNN baseline at 2,091
+to 2,249 samples/s) because each compare ran at a different time on a shared machine. That is why every
+decision comes from baseline and candidate runs interleaved in the same compare, never from numbers
+taken at different times.
+
+## Reproduce
+
+```bash
+uv venv --python 3.12 && uv pip install -e . torch numpy
+PATH="$PWD/.venv/bin:$PATH" sh examples/run_validation.sh
+```
+
+Your numbers will differ on a different machine or under different load, and so can the decisions.

train_doctor-0.1.0/examples/cnn_images/train.py

@@ -0,0 +1,165 @@
+"""Small CNN on a synthetic 10-class image dataset.
+
+Written the way a lot of first training scripts are: per-sample augmentation
+in __getitem__, num_workers=0, and the loss and accuracy read back with
+.item() and printed on every step.
+
+Each command-line flag stands in for one code edit, so train-doctor can
+compare a single change at a time:
+
+  --num-workers N   DataLoader workers (default 0)
+  --log-every N     read metrics and print every N steps (default 1);
+                    between logs, metrics stay on the device
+  --amp             autocast to float16 on mps/cuda
+  --batch-augment   do the crop, flip and normalize once per batch on the
+                    device instead of per sample in __getitem__
+  --batch-size N    default 64
+"""
+
+import argparse
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+
+MEAN = torch.tensor([0.5, 0.5, 0.5]).view(3, 1, 1)
+STD = torch.tensor([0.25, 0.25, 0.25]).view(3, 1, 1)
+
+
+class SyntheticImages(Dataset):
+    """10 classes told apart by a faint colored square in noise. Stored as uint8 like decoded JPEGs."""
+
+    def __init__(self, n: int = 20000, seed: int = 0, augment: bool = True):
+        self.augment = augment
+        g = torch.Generator().manual_seed(seed)
+        self.labels = torch.randint(0, 10, (n,), generator=g)
+        base = torch.randint(0, 256, (n, 3, 32, 32), generator=g, dtype=torch.uint8)
+        # a faint class-colored square blended into the noise, plus 10% label noise
+        colors = torch.linspace(0, 255, 10)
+        for c in range(10):
+            idx = (self.labels == c).nonzero().squeeze(1)
+            patch = base[idx, c % 3, 8:24, 8:24].float()
+            base[idx, c % 3, 8:24, 8:24] = (0.75 * patch + 0.25 * colors[c]).to(torch.uint8)
+        flip = torch.rand(n, generator=g) < 0.1
+        self.labels[flip] = torch.randint(0, 10, (int(flip.sum()),), generator=g)
+        self.images = base
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, i):
+        if not self.augment:
+            return self.images[i], self.labels[i]
+        img = self.images[i].float() / 255.0
+        # random crop with 4 px padding
+        img = F.pad(img, (4, 4, 4, 4), mode="reflect")
+        x, y = torch.randint(0, 9, (2,)).tolist()
+        img = img[:, y : y + 32, x : x + 32]
+        if torch.rand(1).item() < 0.5:
+            img = torch.flip(img, dims=[2])
+        img = (img - MEAN) / STD
+        return img, self.labels[i]
+
+
+def augment_batch(x: torch.Tensor) -> torch.Tensor:
+    """Same augmentation as SyntheticImages.__getitem__, vectorized over a uint8 batch on the device."""
+    b = x.shape[0]
+    x = F.pad(x.float() / 255.0, (4, 4, 4, 4), mode="reflect")
+    oy = torch.randint(0, 9, (b,), device=x.device)
+    ox = torch.randint(0, 9, (b,), device=x.device)
+    ar = torch.arange(32, device=x.device)
+    rows = (oy[:, None] + ar)[:, None, :, None].expand(b, 3, 32, 40)
+    x = torch.gather(x, 2, rows)
+    cols = (ox[:, None] + ar)[:, None, None, :].expand(b, 3, 32, 32)
+    x = torch.gather(x, 3, cols)
+    flip = torch.rand(b, device=x.device) < 0.5
+    x = torch.where(flip[:, None, None, None], x.flip(3), x)
+    return (x - MEAN.to(x.device)) / STD.to(x.device)
+
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.c1 = nn.Conv2d(3, 32, 3, padding=1)
+        self.c2 = nn.Conv2d(32, 64, 3, padding=1)
+        self.c3 = nn.Conv2d(64, 128, 3, padding=1)
+        self.bn1 = nn.BatchNorm2d(32)
+        self.bn2 = nn.BatchNorm2d(64)
+        self.bn3 = nn.BatchNorm2d(128)
+        self.fc = nn.Linear(128 * 4 * 4, 10)
+
+    def forward(self, x):
+        x = F.max_pool2d(F.relu(self.bn1(self.c1(x))), 2)
+        x = F.max_pool2d(F.relu(self.bn2(self.c2(x))), 2)
+        x = F.max_pool2d(F.relu(self.bn3(self.c3(x))), 2)
+        return self.fc(x.flatten(1))
+
+
+def pick_device(name: str) -> torch.device:
+    if name != "auto":
+        return torch.device(name)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--device", default="auto")
+    ap.add_argument("--batch-size", type=int, default=64)
+    ap.add_argument("--num-workers", type=int, default=0)
+    ap.add_argument("--log-every", type=int, default=1)
+    ap.add_argument("--amp", action="store_true")
+    ap.add_argument("--batch-augment", action="store_true")
+    ap.add_argument("--max-steps", type=int, default=300)
+    ap.add_argument("--lr", type=float, default=0.002)
+    args = ap.parse_args()
+
+    torch.manual_seed(0)
+    device = pick_device(args.device)
+    ds = SyntheticImages(augment=not args.batch_augment)
+    loader = DataLoader(
+        ds,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        persistent_workers=args.num_workers > 0,
+        drop_last=True,
+    )
+    model = Net().to(device)
+    opt = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9)
+
+    step = 0
+    run_loss = torch.zeros((), device=device)
+    run_correct = torch.zeros((), device=device)
+    t0 = time.time()
+    while step < args.max_steps:
+        for xb, yb in loader:
+            xb, yb = xb.to(device), yb.to(device)
+            if args.batch_augment:
+                xb = augment_batch(xb)
+            with torch.autocast(device.type, dtype=torch.float16, enabled=args.amp):
+                out = model(xb)
+                loss = F.cross_entropy(out, yb)
+            opt.zero_grad(set_to_none=True)
+            loss.backward()
+            opt.step()
+            step += 1
+            run_loss += loss.detach()
+            run_correct += (out.argmax(1) == yb).sum()
+            if step % args.log_every == 0:
+                n = args.log_every * args.batch_size
+                print(f"step {step} loss {run_loss.item() / args.log_every:.4f} acc {run_correct.item() / n:.3f}")
+                run_loss.zero_()
+                run_correct.zero_()
+            if step >= args.max_steps:
+                break
+    print(f"done: {step} steps in {time.time() - t0:.1f} s")
+
+
+if __name__ == "__main__":
+    main()

train_doctor-0.1.0/examples/run_validation.sh

@@ -0,0 +1,30 @@
+#!/bin/sh
+# The exact sequence that produced examples/*/runs, in order.
+# Run from the repo root with a venv that has torch and train-doctor installed:
+#   uv venv --python 3.12 && uv pip install -e . torch numpy
+#   PATH="$PWD/.venv/bin:$PATH" sh examples/run_validation.sh
+# Each flag on the training scripts stands in for one code edit. Which edit
+# came next was decided from the previous profile's top finding and the
+# previous compare's decision (see examples/README.md).
+set -e
+
+CNN="python examples/cnn_images/train.py"
+OUT=examples/cnn_images/runs
+train-doctor profile --out $OUT --label baseline -- $CNN
+train-doctor compare --out $OUT --label workers-5 --baseline "$CNN" --candidate "$CNN --num-workers 5"
+train-doctor compare --seconds 4 --out $OUT --label batch-augment --baseline "$CNN" --candidate "$CNN --batch-augment"
+train-doctor profile --out $OUT --label batch-augment -- $CNN --batch-augment
+train-doctor compare --seconds 4 --out $OUT --label log-every-50 --baseline "$CNN --batch-augment" --candidate "$CNN --batch-augment --log-every 50"
+train-doctor profile --out $OUT --label log-every-50 -- $CNN --batch-augment --log-every 50
+train-doctor compare --seconds 4 --out $OUT --label amp --baseline "$CNN --batch-augment --log-every 50" --candidate "$CNN --batch-augment --log-every 50 --amp"
+train-doctor compare --seconds 4 --out $OUT --label original-vs-final --baseline "$CNN" --candidate "$CNN --batch-augment --log-every 50"
+
+MLP="python examples/tabular_mlp/train.py"
+OUT=examples/tabular_mlp/runs
+train-doctor profile --out $OUT --label baseline -- $MLP
+train-doctor compare --seconds 4 --out $OUT --label ckpt-every-200 --baseline "$MLP" --candidate "$MLP --ckpt-every 200"
+train-doctor compare --seconds 4 --repeats 9 --out $OUT --label ckpt-every-200-9-repeats --baseline "$MLP" --candidate "$MLP --ckpt-every 200"
+train-doctor profile --out $OUT --label ckpt-every-200 -- $MLP --ckpt-every 200
+train-doctor compare --seconds 4 --out $OUT --label workers-5 --baseline "$MLP --ckpt-every 200" --candidate "$MLP --ckpt-every 200 --num-workers 5"
+train-doctor profile --out $OUT --label workers-5 -- $MLP --ckpt-every 200 --num-workers 5
+train-doctor compare --seconds 4 --out $OUT --label original-vs-final --baseline "$MLP" --candidate "$MLP --ckpt-every 200 --num-workers 5"