Skip to content

perf(callbacks): avoid O(G^2) duplicate scan and redundant string replacements in LearningRateMonitor - #22037

Open
deep1blue wants to merge 1 commit into
Lightning-AI:masterfrom
deep1blue:perf/lr-monitor-duplicate-pg-names
Open

deep1blue wants to merge 1 commit into
Lightning-AI:masterfrom
deep1blue:perf/lr-monitor-duplicate-pg-names

Conversation

@deep1blue

Copy link
Copy Markdown

What does this PR do?

Optimizes LearningRateMonitor (src/lightning/pytorch/callbacks/lr_monitor.py) on the per-step / per-epoch logging path (_extract_stats):

  1. Single-pass $O(G)$ duplicate parameter-group name detection in _duplicate_param_group_names:
    • Previously, _duplicate_param_group_names built both an intermediate list (names) and a set (unique), and when duplicates were present, executed {n for n in names if names.count(n) > 1} which scans names via list.count for each element ($O(G^2)$ where $G = \text{len(param_groups)}$).
    • Replaced with a single-pass seen / duplicates set loop that runs in $O(G)$ worst-case time and avoids allocating the intermediate names list on every logged training step.
  2. Avoid redundant string replacements on every logged step:
    • In _get_optimizer_stats, pass name=f"{name}-momentum" directly instead of calling name.replace(name, f"{name}-momentum") on every parameter group (matching f"{name}-weight_decay" on the following line).
    • In _remap_keys, check if token in new_name: before calling new_name.replace(token, "") so parameter groups without /pg1 do not perform an unnecessary string replacement on every logged step.

Benchmark (_duplicate_param_group_names)

  • $G = 200$ parameter groups: 17.17 ms $\to$ 1.13 ms (15.2x speedup)
  • $G = 1,000$ parameter groups: 394.92 ms $\to$ 5.43 ms (72.8x speedup)
  • $G = 5,000$ parameter groups: 9,417.50 ms $\to$ 29.46 ms (319.7x speedup)

Credits & Acknowledgments

Before submitting
  • Was this discussed/agreed via a GitHub issue? (not for typos and docs)
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure your PR does only one thing, instead of bundling different changes together?
  • Did you make sure to update the documentation with your changes? (if necessary)
  • Did you write any new necessary tests? (not for typos and docs)
  • Did you verify new and existing tests pass locally with your changes? (pytest tests/tests_pytorch/callbacks/test_lr_monitor.py — 24 passed, 1 skipped)
  • Did you list all the breaking changes introduced by this pull request?
  • Did you update the CHANGELOG? (not for typos, docs, test updates, or minor internal changes/refactors)

…lacements in LearningRateMonitor

Signed-off-by: Parth Vadera <parthvadera@gmail.com>

This branch has not been deployed

No deployments
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.

1 participant