Skip to content

perf: remove Mooncake warmup during preparation#772

Merged
gdalle merged 2 commits intomainfrom
gd/remove_mooncake_warmup
Apr 10, 2025
Merged

perf: remove Mooncake warmup during preparation#772
gdalle merged 2 commits intomainfrom
gd/remove_mooncake_warmup

Conversation

@gdalle
Copy link
Copy Markdown
Member

@gdalle gdalle commented Apr 10, 2025

Mooncake preparation used to run the pullback or gradient once in order to compile the rule. This led to a x2 slowdown when used without preparation, because

DI.gradient(f, backend, x)

falls back on

prep = DI.prepare_gradient(f, backend, x)  # used to include one gradient call
DI.gradient(f, prep, x)

This PR removes the duplicate, so the first run of DI.gradient will be slower but only due to compilation.

Ping @penelopeysm who had asked about this

@gdalle gdalle requested a review from willtebbutt April 10, 2025 07:13
@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 10, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.88%. Comparing base (ba8a2fa) to head (eadb50e).
Report is 1 commits behind head on main.

Additional details and impacted files
@@            Coverage Diff             @@
##             main     #772      +/-   ##
==========================================
- Coverage   97.88%   97.88%   -0.01%     
==========================================
  Files         128      128              
  Lines        7463     7460       -3     
==========================================
- Hits         7305     7302       -3     
  Misses        158      158              
Flag Coverage Δ
DI 98.94% <100.00%> (-0.01%) ⬇️
DIT 95.20% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Copy Markdown
Member

@willtebbutt willtebbutt left a comment

Choose a reason for hiding this comment

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

LGTM -- thanks for looking into this.

@gdalle gdalle merged commit a94c9e9 into main Apr 10, 2025
50 checks passed
@gdalle gdalle deleted the gd/remove_mooncake_warmup branch April 10, 2025 08:44
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.

2 participants