[DO NOT MERGE] Goodput prototype v1 #337
Conversation
| @@ -0,0 +1,127 @@ | |||
| """ | |||
| Goodput Prototype (TODO: Move to packaged library) | |||
There was a problem hiding this comment.
Also why is this writing to GCS? Is there some metric sservice we can use?
There was a problem hiding this comment.
Writing to GCS is a lot more complicated that it seems so we'd need to think about that carefully.
There was a problem hiding this comment.
+1, I think we can readily use Cloud Logging to stream logs: https://cloud.google.com/logging/docs/write-query-log-entries-python#linux
| with mesh, nn_partitioning.axis_rules(config.logical_axis_rules): | ||
| state, metrics, nextrng = p_train_step( | ||
| state, example_batch, nextrng | ||
| ) | ||
| step_stop_time = datetime.datetime.now() |
There was a problem hiding this comment.
This code doesn't do what you think it does -- Jax calls are non blocking. I think it would better to provide a goodput.goodput_log_step_time call that took step time (and inside of it you can get a datetime).
| @@ -0,0 +1,127 @@ | |||
| """ | |||
| Goodput Prototype (TODO: Move to packaged library) | |||
There was a problem hiding this comment.
Writing to GCS is a lot more complicated that it seems so we'd need to think about that carefully.
| storage_client = storage.Client() | ||
| bucket = storage_client.bucket(bucket_name) | ||
| blob = bucket.blob(blob_name) | ||
| blob.download_to_filename(destination_file_name) |
There was a problem hiding this comment.
BTW this type of logic is unwelcome on our control thread! We need to be writing asynchronously to something that gracefully handles appends.
(GCS really doesn't like appends. And this code seems to be synchronous on our control thread.)
| with diagnostic.diagnose(diagnostic_config): | ||
| train_loop(config) | ||
|
|
||
| job_end_time = datetime.datetime.now() | ||
| goodput.goodput_log_job_runtime(config, job_end_time - job_start_time) |
There was a problem hiding this comment.
Also you really don't want to pass a config outside of MaxText into your generic interface! Pass what you need!
| job_end_time = datetime.datetime.now() | ||
| goodput.goodput_log_job_runtime(config, job_end_time - job_start_time) | ||
| total_goodput = goodput.get_job_goodput(config) | ||
| max_logging.log(f"Total job goodput: {total_goodput:.2f}%") |
There was a problem hiding this comment.
Are we sure we want to print this? Because this run doesn't know the overall goodput -- the goodput depends on fusing multiple runs together.
Quick prototype to compute Goodput based on total step time and job time on MaxText.