Skip to content

Count each module object once when aggregating flops profiler totals - #8320

Open
ebarkhordar wants to merge 1 commit into
deepspeedai:masterfrom
ebarkhordar:fix/7256-flops-profiler-shared-module-dedup
Open

Count each module object once when aggregating flops profiler totals#8320
ebarkhordar wants to merge 1 commit into
deepspeedai:masterfrom
ebarkhordar:fix/7256-flops-profiler-shared-module-dedup

Conversation

@ebarkhordar

Copy link
Copy Markdown
Contributor

get_module_flops and get_module_macs aggregate by walking children() 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 one Linear(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 shared Linear there 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_flops ruled out modules() because it "returns duplicate modules only once". That is accurate about modules(), and given a per-object accumulator it is the behaviour the aggregation wants, so this replaces the walk with a sum over modules() 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

  • On the issue's own script the total goes from 180 KFLOPs / 90 KMACs to 60 KFLOPs / 30 KMACs, which is what torch.profiler reports for the same model on both the shared and unshared variants.
  • Added test_flops_profiler_counts_shared_module_once, which fails on master at assert 180000 == 60000 and passes here for both shared=True and shared=False. It is marked sequential, so it runs in the second cpu-torch-latest pytest leg, where the file is 4 passed.
  • pre-commit run --files passes on both changed files under Python 3.10, the formatting job's environment.
  • Not checked: the RNN flop hooks write the same accumulator but no RNN module was exercised. get_module_duration and the per-depth table in print_model_profile still aggregate per tree position rather than per object; whether they should change too is a separate question I have not measured.

Fixes #7256

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>

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 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())

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@delock
delock requested review from delock and removed request for delock August 26, 2026 01:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Inaccurate FLOPs for shared weights

1 participant