#427 base_eval.py: hellaswag gets progressively slower and leaks memory on small models (Mac Studio)

I’m running base_eval.py on a very small model trained with base_train.py. Evaluation starts normally, but on the hellaswag task each example becomes slower and slower, memory usage grows continuously, and the script will take massive amount time to finish. Other tasks finish in minutes, but hellaswag runs for hours.

Environment - Mac Studio, Python 3.10.19

I trained "much, much smaller" model using a command from comment in base_train.py

python -m scripts.base_train --depth=4 --max_seq_len=512 --device_batch_size=1 --eval_tokens=512 --core_metric_every=-1 --total_batch_size=512 --num_iterations=20

I gave script few hours to run but it didn't finish. When interrupted with Ctrl+C, it always stops here:

Traceback (most recent call last):
  File "/Users/user/.pyenv/versions/3.10.19/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/Users/user/.pyenv/versions/3.10.19/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/Users/user/Projects/nanochat/scripts/base_eval.py", line 219, in <module>
    main()
  File "/Users/user/Projects/nanochat/scripts/base_eval.py", line 182, in main
    out = evaluate_model(model, tokenizer, device, max_per_task=args.max_per_task)
  File "/Users/user/Projects/nanochat/scripts/base_eval.py", line 107, in evaluate_model
    accuracy = evaluate_task(model, tokenizer, data, device, task_meta)
  File "/Users/user/Projects/nanochat/nanochat/core_eval.py", line 254, in evaluate_task
    is_correct = evaluate_example(idx, model, tokenizer, data, device, task_meta)
  File "/Users/user/Projects/nanochat/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
  File "/Users/user/Projects/nanochat/nanochat/core_eval.py", line 234, in evaluate_example
    mean_losses = [losses[i, si-1:ei-1].mean().item()
  File "/Users/user/Projects/nanochat/nanochat/core_eval.py", line 234, in <listcomp>
    mean_losses = [losses[i, si-1:ei-1].mean().item()
KeyboardInterrupt
Evaluating: hellaswag (10-shot, type: multiple_choice)... ⏎                
(nanochat) user@users-Mac-Studio ~/P/nanochat (master) [SIGINT]>

What happens

First ~100 examples are fast
After that, some examples are getting more and more time to run
These slow spikes become more frequent
RAM usage keeps growing until the machine runs out of memory

I added simple timing using time.time() inside evaluate_example() and logged per-example runtimes:

1 0.1499180793762207
2 0.0699930191040039
3 0.07904505729675293
4 0.06212663650512695
5 0.16261887550354004
6 0.0742349624633789
7 0.06836700439453125
8 0.07444477081298828
9 0.15993094444274902
10 0.0733938217163086
11 0.1757950782775879
12 0.06647109985351562
13 0.0805962085723877
14 0.06577205657958984
15 0.07175707817077637
16 0.07457876205444336
17 0.13344931602478027
18 0.167036771774292
19 0.07505202293395996
20 0.055587053298950195
...
180 0.10489702224731445
181 1.0350978374481201
182 0.9883811473846436
183 0.08685803413391113
184 0.10237693786621094
185 1.0098748207092285
186 7.256951093673706
187 0.10192418098449707
188 3.0714337825775146
189 4.08650016784668
190 0.0888969898223877
191 1.0460128784179688
192 0.07270097732543945
193 0.06732010841369629
194 0.08224320411682129
195 0.5282318592071533
196 0.8869402408599854
197 0.09626007080078125
198 1.0446851253509521
199 0.07490086555480957
200 2.118406057357788

I’m not 100% sure this is a bug rather than an expected limitation, but the behaviour looks suspicious enough that I thought it was worth flagging and reporting.