Count each module object once when aggregating flops profiler totals - #8320
Count each module object once when aggregating flops profiler totals#8320ebarkhordar wants to merge 1 commit into
Conversation
A module aliased at several positions in the tree was reached once per position by the children() walk, so its flops were added that many times. Each module object owns one __flops__ accumulator that the forward hook adds into on every call, so visiting it once already counts all of its calls. Fixes deepspeedai#7256 Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9111391a12
ℹ️ 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".
| return sum | ||
| # Each module object owns a single __flops__ accumulator that every forward call adds | ||
| # into, so visiting an aliased module once already counts all of its calls. | ||
| return sum(m.__flops__ for m in module.modules()) |
There was a problem hiding this comment.
Add the required Signed-off-by trailer
This is a non-merge commit, but its message ends at Fixes #7256 without a Signed-off-by trailer, so it does not satisfy the repository's commit requirements and may be rejected by contribution checks. Recreate the commit with --signoff using the configured Git identity.
AGENTS.md reference: AGENTS.md:L6-L8
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The trailer is already there. This PR has one commit, 9111391, and its message ends with Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>, which is why the DCO check on this PR is green. No change needed on this point.
get_module_flopsandget_module_macsaggregate by walkingchildren()and adding each child's subtree into a running total. A submodule object that is aliased at several positions in the tree is reached once per position, so its flops are added that many times. The issue's model shares oneLinear(100, 100)across three parents and reports 180 KFLOPs where the work done is 60 KFLOPs.The per-module counters are already correct. Each module object owns one
__flops__accumulator and the forward hook adds into it on every call, so the sharedLinearthere ends up holding 60000, the flops of all three calls. Visiting that object once is what counts every call; the extra visits re-add the same number.The comment above
get_module_flopsruled outmodules()because it "returns duplicate modules only once". That is accurate aboutmodules(), and given a per-object accumulator it is the behaviour the aggregation wants, so this replaces the walk with a sum overmodules()instead of leaving the comment standing next to code that no longer follows it.Note that
named_children()already skips a repeat among immediate siblings, so this only ever showed up when the shared module sat under distinct parents, as it does in the report.Verification
torch.profilerreports for the same model on both the shared and unshared variants.test_flops_profiler_counts_shared_module_once, which fails on master atassert 180000 == 60000and passes here for bothshared=Trueandshared=False. It is markedsequential, so it runs in the secondcpu-torch-latestpytest leg, where the file is 4 passed.pre-commit run --filespasses on both changed files under Python 3.10, the formatting job's environment.get_module_durationand the per-depth table inprint_model_profilestill aggregate per tree position rather than per object; whether they should change too is a separate question I have not measured.Fixes #7256