[maxtext] improve profiling in decode #1863
Conversation
![copybara-service[bot]](https://avatars.githubusercontent.com/in/44061?s=60&v=4){kind=small}
copybara-service
bot
requested review from
gobbleturk,
khatwanimohit,
bvandermoon,
vipannalla,
RissyRan,
richjames0,
gagika,
shralex,
yangyuwei,
SurbhiJainUSC,
hengtaoguo,
A9isha and
aireenmei
as code owners
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
MaxText/decode.py
Outdated
| @@ -33,6 +33,7 @@ | |||
|
|
|||
| # Number of text sequences to process in a single batch. | |||
| _NUM_STREAMS = 1 | |||
| _MAX_STEPS_TO_PROFILE = 10 | |||
There was a problem hiding this comment.
Why not just use the general MaxText profiling config?
| @@ -92,3 +92,6 @@ def should_activate_periodic_profile(self, step): | |||
|
|
|||
| def should_deactivate_periodic_profile(self, step): | |||
| return self.profile_period > 0 and (step - self.finished_initial_profile_step) % self.profile_period == 0 | |||
|
|
|||
| def post_process(self): | |||
| pass | |||
There was a problem hiding this comment.
The plan is to fill this in later, right?
There was a problem hiding this comment.
Why is profiling even added to decode.py? Looks like the PR got merged without any review: #1814.
@SurbhiJainUSC As FYI - the internal cl number is in the branch name at the top of the PR: |
copybara-service
bot
force-pushed
the
test_773929904
branch
2 times, most recently
from
cb6d96c to
a7213ce
Compare
ohh, I got it. Thank you for pointing that out. |
copybara-service
bot
force-pushed
the
test_773929904
branch
from
a7213ce to
6539cb9
Compare
Limit the profiling of decode to `config.profiler_steps` steps. User may request generate many tokens, but since each generation is identical, it's not necessary to profile all of them. Profiling only the first few should be sufficient. Also added a method in Profiler to do post-pocessing on the collected trace. PiperOrigin-RevId: 775799587
copybara-service
bot
force-pushed
the
test_773929904
branch
from
6539cb9 to
3a78db6
Compare
[maxtext] improve profiling in decode
Limit the profiling of decode to
config.profiler_stepssteps. User may request generate many tokens, but since each generation is identical, it's not necessary to profile all of them. Profiling only the first few should be sufficient.Also added a method in Profiler to do post-pocessing on the collected trace.