Skip to content

fix(measure): every measurement leaves the model in the mode it came in - #10

Merged
nathanhubens merged 1 commit into
masterfrom
fix/measure-restores-training-mode
Sep 4, 2026
Merged

nathanhubens merged 1 commit into
masterfrom
fix/measure-restores-training-mode

Conversation

@nathanhubens

Copy link
Copy Markdown
Contributor

Summary

Same family as #9: every measurement put the model in eval mode and never restored training, so a model in training mode came back in eval mode after compute_speed (and the sweeps), compute_memory, compute_energy, LayerProfiler.profile or RooflineAnalyzer.profile. Only benchmark() restored the flag.

One context manager, _measuring(model, dev) in fasterbench.core, records model.training, wraps _on_device, yields the model in eval mode, and restores the mode in a finally. The six sites go from with _on_device(model.eval(), dev) as model: to with _measuring(model, dev) as model:, one line each; _on_device itself is unchanged. benchmark() keeps its own restore because compute_compute never sets eval mode and benchmark()'s model.eval() is what covers it.

Red → green

On origin/master (a git archive snapshot), a .train() model measured on CPU:

FAIL  compute_speed: training=False (expected True)
FAIL  compute_memory: training=False (expected True)
FAIL  compute_energy: training=False (expected True)
FAIL  LayerProfiler.profile: training=False (expected True)
FAIL  RooflineAnalyzer.profile: training=False (expected True)

On this branch all five read training=True, and each site also checks that a model already in eval mode stays in eval mode. Unlike the device tests of #9, these run on CPU, so CI covers them (the energy one stays #| notest, next to the existing codecarbon cell, so CI does not gain a second tracker run); a mutation of the assertion fails the default suite. The two-line CUDA check in core's #| notest cell also asserts device and mode together.

Tests

Default suite: Success. The seven touched notebooks with flags='notest' on a GPU: Success. each. CI gate (nbdev-clean then nbdev-export with CI's nbdev 3.3.15): git status --porcelain -uno empty on two passes. __version__ 0.1.0, _modidx.py gains _measuring, __init__.py untouched.

https://claude.ai/code/session_0177tJZDQTTjQdycXV1EZtdJ

Measuring put the model in eval mode and never put it back, so a model
being trained came back in eval mode (dropout/BN silently off for the
rest of training). Same family as the device bug: `_on_device` restored
the device, nothing restored `model.training`.

`_measuring(model, dev)` composes `_on_device` with the mode round-trip
(record `model.training`, `eval()`, restore in `finally`), and replaces
the six `_on_device(model.eval(), dev)` sites:

- speed.py `_forward_latencies` (compute_speed, sweep_*)
- memory.py `_gpu_metrics`
- memory.py `_cpu_metrics`
- energy.py `compute_energy`
- profiling.py `_profile_layers` (LayerProfiler.profile)
- roofline.py `RooflineAnalyzer.profile`

`benchmark()` keeps its own eval/restore: it is not redundant, because
compute_compute() never sets eval and relies on it. It now says so and
uses the same `model.train(was_training)` idiom.

Tests: per touched notebook, a `.train()` model must come back with
`training is True` (red on master for all five entry points) and an
`.eval()` one must stay False (catches an over-eager restore). Those
cells are CPU-only, except energy's, which sits with the existing
codecarbon cell under `#|notest` to keep host power metering out of CI.
Core also covers restore-on-exception and the CUDA device+mode round-trip.

Claude-Session: https://claude.ai/code/session_0177tJZDQTTjQdycXV1EZtdJ
@nathanhubens
nathanhubens merged commit dfd2b93 into master Sep 4, 2026
2 checks passed
@nathanhubens
nathanhubens deleted the fix/measure-restores-training-mode branch September 4, 2026 12:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant