Fix Microbenchmark Profiling Memory Issues #597
Conversation
morgandu
force-pushed
the
mor--inference
branch
5 times, most recently
from
8253dbb
to
3f28ee3
Compare
|
|
||
|
|
||
| def main(config): | ||
| engine = maxengine.MaxEngine(config) | ||
| params = engine.load_params() | ||
| prefill_lengths = [64, 128, 256, 512, 1024] | ||
| benchmark_loop_iters = 10 | ||
| prefill_lengths = [int(l) for l in config.inference_microbenchmark_prefill_lengths.split(",")] |
There was a problem hiding this comment.
Does this work if you pass in a command line param like --inference-microbenchmark-prefill-lengths="512,1024" or something similar?
There was a problem hiding this comment.
Yeah, example commands:
to run a single prefill length:
inference_microbenchmark_prefill_lengths=1024
to run a single stage:
inference_microbenchmark_stages=generate
MaxText/inference_utils.py
Outdated
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
inference_utils.py is an existing file, but I can move these functions to max_utils.py
morgandu
changed the title
I am writing a batch inference, which need some common utility functions |
| benchmark_results["AutoRegressive"], decode_state = ar_benchmark( | ||
| config, engine, params, decode_state, iters=benchmark_loop_iters, cache_size=cache_size, model_size=model_size) |
There was a problem hiding this comment.
For running just generate benchmark, you still need to populate kv cache to produce proper perf numbers?
There was a problem hiding this comment.
You will still need to initialize a decode_state for generate step calculation
| @@ -261,3 +261,8 @@ vertex_tensorboard_project: "" | |||
| # Region to create Vertex AI Tensorboard in for GCE, blank if running via XPK | |||
| # Vertex AI supported regions: https://cloud.google.com/vertex-ai/docs/general/locations#available-regions | |||
| vertex_tensorboard_region: "" | |||
|
|
|||
There was a problem hiding this comment.
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
force-pushed
the
mor--inference
branch
2 times, most recently
from
210d7b2
to
38831be
Compare
morgandu
force-pushed
the
mor--inference
branch
from
f522aba
to
50ae199
Compare
morgandu
force-pushed
the
mor--inference
branch
from
3a3a1c5
to
c8ed884
Compare
- allow run specified stages - allow run specific prefill length(s) - delete prefill result - printout prefill result added funcs in max_utils
morgandu
force-pushed
the
mor--inference
branch
from
c8ed884
to
b46783c
Compare
Below are changes in this PR