Skip to content

Fix Microbenchmark Profiling Memory Issues#597

Merged
copybara-service[bot] merged 1 commit into
mainfrom
mor--inference
Apr 19, 2024
Merged

Fix Microbenchmark Profiling Memory Issues#597
copybara-service[bot] merged 1 commit into
mainfrom
mor--inference

Conversation

@morgandu

@morgandu morgandu commented Apr 16, 2024

Copy link
Copy Markdown
Collaborator

Below are changes in this PR

  • allow microbenchmark with prefill lengths
  • allow microbenchmark with loop iters
  • allow benchmark prefill / generate stage only
  • delete prefill_results to that limited the batch size
  • added / move pytree util funcs

@morgandu morgandu requested a review from patemotter April 16, 2024 20:05
@morgandu morgandu force-pushed the mor--inference branch 5 times, most recently from 8253dbb to 3f28ee3 Compare April 18, 2024 21:15

@patemotter patemotter left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Overall LGTM, but I don't think we need the inference_utils.py file.

Comment thread MaxText/inference_microbenchmark.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Does this work if you pass in a command line param like --inference-microbenchmark-prefill-lengths="512,1024" or something similar?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Yeah, example commands:
to run a single prefill length:

  inference_microbenchmark_prefill_lengths=1024

to run a single stage:

  inference_microbenchmark_stages=generate

Comment thread MaxText/inference_utils.py Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I don't think we need to make a file solely for inference_utils at this time. Especially since these functions are not unique to inference. I would add them to max_utils.py since they cover fairly generic usecases.

@morgandu morgandu Apr 18, 2024

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

inference_utils.py is an existing file, but I can move these functions to max_utils.py

@morgandu morgandu changed the title Enable Inference mode and Fix Microbenchmark Profiling / Memory Issues Fix Microbenchmark Profiling Memory Issues Apr 18, 2024
@morgandu

Copy link
Copy Markdown
Collaborator Author

Overall LGTM, but I don't think we need the inference_utils.py file.

I am writing a batch inference, which need some common utility functions

Comment thread MaxText/inference_microbenchmark.py Outdated
Comment on lines 188 to 205

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

For running just generate benchmark, you still need to populate kv cache to produce proper perf numbers?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

You will still need to initialize a decode_state for generate step calculation

Comment thread MaxText/configs/base.yml Outdated

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Not in this PR, but I see a need for separate inference specific config files in future -- both base.yml and model specific config.

@morgandu morgandu force-pushed the mor--inference branch 2 times, most recently from 210d7b2 to 38831be Compare April 19, 2024 00:01

@rwitten rwitten left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

looks good

@rwitten rwitten left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Please remember to squash your commits.

  - allow run specified stages
  - allow run specific prefill length(s)
  - delete prefill result
  - printout prefill result

added funcs in max_utils
@copybara-service copybara-service Bot merged commit 0e1c078 into main Apr 19, 2024
@copybara-service copybara-service Bot deleted the mor--inference branch April 19, 2024 17:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants