Count the operator spelling of add and mul in the flops profiler - #8509
ebarkhordar wants to merge 1 commit into
Conversation
The profiler patches torch.add, torch.Tensor.add, torch.mul and torch.Tensor.mul, but not the operator forms, so an elementwise op written as x + y reported zero flops while the same op written as torch.add(x, y) reported one flop per element. Every residual connection in a profiled model was invisible. torch.Tensor.__matmul__ is already patched for this reason, so this adds the same treatment for __add__, __radd__, __mul__ and __rmul__, and the matching entries in _reload_tensor_methods so the names are restored when profiling ends. Refs deepspeedai#3087 Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: df5f2284a7
ℹ️ 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.mul = wrapFunc(torch.mul, _mul_flops_compute) | ||
| torch.Tensor.mul = wrapFunc(torch.Tensor.mul, _mul_flops_compute) | ||
| torch.Tensor.__mul__ = wrapFunc(torch.Tensor.__mul__, _mul_flops_compute) |
There was a problem hiding this comment.
This is a non-merge commit with one parent, but its message contains no Signed-off-by trailer. The root AGENTS.md requires every non-merge commit to be signed off, so this commit does not meet the repository's commit/CI requirements until the author trailer is added.
AGENTS.md reference: AGENTS.md:L6-L10
Useful? React with 👍 / 👎.
This branch has not been deployed
1 participant
get_model_profilereports 0 flops for an elementwise op written asx + y, and one flop per element for the same op written astorch.add(x, y). The profiler patchestorch.add,torch.Tensor.add,torch.mulandtorch.Tensor.mul, but not the operator forms, so every residual connection in a profiled model is counted as free. On the reporter's script atc102bbd1,x + xover 20 elements gives 0 wheretorch.add(x, x)gives 20.torch.Tensor.__matmul__is already patched for this reason (#4724), so the list is incomplete rather than deliberately limited. This adds__add__,__radd__,__mul__and__rmul__beside it, with matching entries in_reload_tensor_methodsso the names are restored when profiling ends.The reflected forms are here so that
2.0 * xandx * 2.0agree. In-place__iadd__and__imul__are left out on purpose:Tensor.add_andTensor.mul_are not counted either, and adding one without the other recreates the same asymmetry a level down.x - yandx / yalso report 0, buttorch.subandtorch.divare missing from the patch list in every spelling, so that is a separate change and this one isRefsrather thanFixes.How I checked it, on CPU with torch 2.8.0:
torch.add(x, x)andx.add(x)still report 20, so reaching the op through the operator does not count it twice.x + self.fc(x)moves from 200 to 210, which is the 10-element add.pytest -m sequential unit/profiling/flops_profiler/passes 23, and everypre-commithook passes on both changed files, with flake8 run under Python 3.10 because its pinned version needs one.There is no GPU on the machine I used, so I have not exercised the CUDA or precompile paths, and I tested only against torch 2.8.0.
Refs #3087