Repository navigation
Conversation
`evaluate_on_dataset` and `induction_loss` defaulted `device` to the literal string "cuda", and `evaluate` never forwarded a device at all and exposed no parameter for one. On a CPU-only or MPS machine `evaluate(model)` therefore moved batches to a device the model was not on, and the caller had no way to correct it. `cfg.device` is always populated by TransformerLensConfig.__post_init__, which resolves to cuda when it is available and cpu otherwise. `mmlu_eval` in this same module already uses it. The three helpers now default to None and fall back to `model.cfg.device`, so an explicit argument still wins and callers who were passing one are unaffected. `evaluate` gains the `device` parameter it was missing and forwards it. Adds tests/unit/test_evals_device.py: the helpers move tensors to the model's device, an explicit device still overrides it, and none of the three defaults to a fixed device any more. Uses a synthetic stub, so it loads no model and stays at the unit tier.
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for fixing this bug @abhishekcpp007, just a couple small comments about testing effectiveness & a quick comment cleanup. Please let me know if you have any questions
| """An explicit argument still wins over cfg.device.""" | ||
| model = StubModel(device="cpu") | ||
|
|
||
| evaluate_on_dataset(model, make_loader(), truncate=1, device="cpu") |
There was a problem hiding this comment.
The stub's cfg.device and the explicit argument are both "cpu", so this passes even if the argument is ignored. Could you give the stub a different cfg.device, and do the same for test_induction_loss_builds_tokens_on_the_models_device below on line 68?
| assert all(device.type == "cpu" for device in model.seen_devices) | ||
|
|
||
|
|
||
| def test_eval_helpers_no_longer_hardcode_cuda() -> None: |
There was a problem hiding this comment.
This asserts the declared defaults rather than what the helpers do, and nothing in the file calls evaluate, so the device=device forwarding inside it has no test. Could you replace this with a test that monkeypatches evals.DATASET_NAMES and evals.DATASET_LOADERS to a stub loader and checks that evaluate(model, device="cpu") reaches the stub on CPU?
| @torch.inference_mode() | ||
| def evaluate_on_dataset(model, data_loader, truncate=100, device="cuda"): | ||
| def evaluate_on_dataset(model, data_loader, truncate=100, device=None): | ||
| # Default to where the model actually is. cfg.device is always populated, so |
There was a problem hiding this comment.
Can we this comment to something like # None defaults to the model's own device, a hardcoded "cuda" would crash CPU/MPS callers?
Give the stub a "meta" cfg.device in the explicit-device and induction tests so they can tell the argument apart from cfg.device, add an explicit device test for induction_loss, and replace the signature check with a test that runs evaluate against a monkeypatched stub loader. Also shorten the comment in evaluate_on_dataset.
|
Thanks for the review @jlarson4! Pushed 0d2b1eb. The explicit-device tests now give the stub a "meta" cfg.device, so they fail if the argument gets ignored, and I added the same kind of test for induction_loss. I swapped the signature check for a test that monkeypatches DATASET_NAMES and DATASET_LOADERS with a stub loader and checks that evaluate(model, device="cpu") reaches it on CPU. Comment's updated too. I checked by dropping the forwarding locally and the 3 new tests fail. |
|
Great work on this, thanks for the updates. Approved, I'll merge once this it passes CI |
Description
transformer_lens/evals.pysends work to the wrong device on any machine that is not CUDA.evaluate_on_dataset(..., device="cuda")moves every batch to a literal"cuda".induction_loss(..., device="cuda")allocates its random sequence there too.evaluate()never forwarded a device at all and exposed no parameter for one, so a callerhad no way to correct it.
On a CPU-only or MPS machine
evaluate(model)therefore moves batches to a device the model isnot on.
mmlu_eval, in this same module, already does the right thing (.to(model.cfg.device)).The right value is already computed for us:
TransformerLensConfig.__post_init__setsself.device = "cuda" if torch.cuda.is_available() else "cpu", socfg.deviceis alwayspopulated.
Fix
The three helpers now default
devicetoNoneand fall back tomodel.cfg.device.device=argument still wins, so anyone already passing one is unaffected."cuda"now gets CUDA anyway whenever it is available,since that is what
cfg.deviceresolves to.evaluategains thedeviceparameter it was missing and forwards it.Type of change
Tests
tests/unit/test_evals_device.py, unit tier:evaluate_on_datasetmoves batches to the model's devicedevicestill overrides itinduction_lossallocates its sequence on the model's devicedeviceto a fixed device any moreIt uses a synthetic stub that records the device of each tensor it is handed, so it loads no
model and mocks no model load, per
tests/AGENTS.md.Verification I was able to run, and what I was not
Being straight about this rather than claiming a green run, per
AGENTS.md§10.isortandblackover both changed files with the repo'spyproject.tomlconfig.Both report no changes, so
make check-formatis clean on them.[tool.mypy] excludecoverstests, so the new test file is outsidemypy's surface. The change in
evals.pyadds no annotations.make unit-testanduv run mypy .. My environment could not install theproject's dependencies, so I could not execute the suite, and I would rather say so than
report a pass I did not see. CI is the real run here. If anything in the new test goes red I
will fix it straight away.