DeepSpeed
f4ee787d - Count each module object once when aggregating flops profiler totals (#8320)

Commit
4 days ago
Count each module object once when aggregating flops profiler totals (#8320) `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 Signed-off-by: Ehsan Barkhordar <realbarkhordar@gmail.com>
Author
Parents
Loading