benchmark fix - #94
Conversation
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: de233443db
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| torch.cuda.synchronize() | ||
| start = perf_counter() | ||
| for _ in range(steps): | ||
| set_grads() |
There was a problem hiding this comment.
When users benchmark many small tensors (or increase the tensor count), this line times the Python set_grads() loop once per measured step, so the reported median scales with per-parameter attribute assignment before optimizer.step() even starts. Since HeavyBall exposes consume_grad=False to keep .grad attached, the no-op issue can be fixed without adding this assignment loop to the timed region; otherwise this benchmark no longer reports just optimizer-step latency.
Useful? React with 👍 / 👎.
This reverts commit de23344.
|
Should be good : Results with current fix : root@C.41943479:~/HeavyBall$ python benchmarks/bench_optimizer_step.py |
|
looks good to me, thank you! does this meaningfully change the results compared to what's reported in the readme? -- i believe it was profiled before the bug was introduced |
|
this fix gives about the same result as the previous working version (e387c8a commit) Results with current fix : There's a 1-2 us noise difference, both tests were made on the same hardware (RTX 3090) |
|
awesome, thank you for pushing this through! that's reassuring. |
Fixed the #93 issue
Current results :
python benchmarks/bench_optimizer_step.py
Version
heavyball: main branch (latest)
PyTorch: 2.12.1+cu130
Python: 3.12
RTX 3090 24GB