Fix repeated FlopsProfiler metric accumulation - #8246
Conversation
Sorry, something went wrong.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 997dabee8f
ℹ️ About Codex in GitHubCodex 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".
Sorry, something went wrong.
| F.interpolate = old_functions[F.interpolate.__str__] | ||
| F.softmax = old_functions[F.softmax.__str__] | ||
| F.embedding = old_functions[F.embedding.__str__] | ||
| if required_torch_version(min_version=2.0): |
There was a problem hiding this comment.
Add the required Signed-off-by trailer
This is a non-merge commit, but its message has no Signed-off-by: trailer, so it violates the workspace commit requirement and can be rejected by CI/review tooling; please recreate or amend the commit with --signoff.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
Sorry, something went wrong.
ebarkhordar
left a comment
There was a problem hiding this comment.
Reproduced both sides on CPU (clean python:3.11-slim, torch 2.13.0+cpu), base 1d580d65 against head 997dabee, three rounds on the SDPA model from your test:
base: [(65536, 32768), (131072, 65536), (196608, 98304)] head: [(65536, 32768), (65536, 32768), (65536, 32768)]
The bmm line is more than a metrics bug, and I think the description should say so. torch.bmm = old_functions[torch.matmul.__str__] runs after torch.matmul has already been restored on the line above, so it rebinds torch.bmm to torch.matmul for the rest of the process. matmul broadcasts where bmm requires 3-D: after profiling a plain nn.Linear once, torch.bmm(4x4, 4x4) returns a 4x4 instead of raising RuntimeError: batch1 must be a 3D tensor. Any profiled model reaches it, not only attention ones. Head restores it correctly.
To check nothing of the same shape is left, I parsed _patch_* and _reload_*: base is 51 patched against 49 reloaded (exactly the two you add) plus the mis-keyed bmm; head is 51/51 with matching version guards. So that is the complete set.
One suggestion. The module is skipped entirely when the accelerator does not report fp16 (line 17), so the new test does not collect on a CPU-only runner (1 skipped, collected 0 items here) even though it needs neither fp16 nor a device. Hoisting it above that guard would let the cpu-torch-latest -m sequential leg run it.
Sorry, something went wrong.
|
Thanks for the thorough verification. Agreed that the incorrect torch.bmm restoration has broader process-wide correctness implications, not just inaccurate profiler metrics. I’ll make that explicit in the PR description. I also addressed the CPU collection issue. Since a module-level pytest.skip(..., allow_module_level=True) aborts collection of the entire file even if the test is defined above it, I moved the FP16 capability check into the existing test that actually uses FP16. The repeated-profile regression now collects and passes in the CPU-only sequential test path. |
Sorry, something went wrong.
Signed-off-by: Vedant Chauhan <staranonymous1011@gmail.com>
tohtana
left a comment
There was a problem hiding this comment.
@baremetaldevx86 Thank you for the fix! This looks good to me.
Sorry, something went wrong.
Summary
Problem
FlopsProfiler temporarily replaces PyTorch operations with FLOP-counting wrappers. Its cleanup path did not restore F.scaled_dot_product_attention or
Tensor.__matmul__, causing wrappers to accumulate across profiling sessions. As a result, identical model executions reported progressively increasing FLOPs and
MACs.
The torch.bmm cleanup was also incorrect:
Because torch.matmul had already been restored, this rebound torch.bmm to torch.matmul for the rest of the process. This changed normal PyTorch behavior after
profiling: torch.bmm began accepting inputs supported by broadcasting matmul but invalid for bmm.
Fix
Restore every patched operation from its matching saved original function. The patch and cleanup paths now cover the same set of operations with matching PyTorch
version guards.
Testing
A CPU regression test profiles the same scaled-dot-product-attention operation three times and verifies:
Observed totals:
Before:
[(65536, 32768), (131072, 65536), (196608, 98304)]
After:
[(65536, 32768), (65536, 32768), (65536, 32768)]
All pre-commit checks pass.
Fixes #7413