Package: ml-goodput-measurement
ML Goodput MeasurementOverviewML Goodput Measurement is a library intended to be used with Cloud TPU to log the
necessary information and query a job's Goodput. It can be pip installed to
import its modules, and retrieve information about a training job's overall
productive Goodput. The package exposes API interfaces to log useful
information from the user application and query Goodput for the job run, gain
insight into the productivity of ML workloads and utilization of compute
resources.ComponentsThe ML Goodput Measurement library consists of two main components:
theGoodputRecorderand theGoodputCalculator. TheGoodputRecorderexposes APIs to the client to export key timestamps while a training job makes
progress, namely APIs that allow logging of productive step time and total job
run time. The library will serialize and store this data in Google Cloud
Logging. TheGoodputCalculatorexposes APIs to compute Goodput based on the
recorded data. Cloud Logging handles its internal operations asynchronously.
The recommended way to compute Goodput is to run an analysis program separate
from the training application, either on a CPU instance or on the users'
development machine.InstallationTo install the ML Goodput Measurement package, run the following command on TPU VM:pipinstallml-goodput-measurementUsageThe usage of this package requires the setup of a Google Cloud project with
billing enabled to properly use Google Cloud Logging. If you don't have a Google
Cloud project, or if you don't have billing enabled for your Google Cloud
project, then do the following:In the Google Cloud console, on the project selector page,select or create a Google Cloud project.Make sure that billing is enabled for your Google Cloud project. Instructions can be foundhereTo run your training on Cloud TPU, set up the Cloud TPU environment by following
instructionshere.To learn more about Google Cloud Logging, visit thispage.ImportTo use this package, import thegoodputmodule:fromml_goodput_measurementimportgoodputDefine the name of the Google Cloud Logging logger bucketCreate a run-specific logger bucket where Cloud Logging entries can be written and read from.For example:goodput_logger_name=f'goodput_{config.run_name}'Create aGoodputRecorderobjectNext, create a recorder object with the following parameters:job_name: The full run name of the job.logger_name: The name of the Cloud Logging logger object (created in the previous step).logging_enabled: Whether or not this process has Cloud Logging enabled.NOTE:For a multi-worker setup, please ensure that only one worker
writes the logs to avoid the duplication. In JAX, for example, the check
could beif jax.process_index() == 0NOTE:logging_enableddefaults toFalseand Goodput computations cannot be completed if no logs are ever written.For example:goodput_recorder=goodput.GoodputRecorder(job_name=config.run_name,logger_name=goodput_logger_name,logging_enabled=(jax.process_index()==0))Record Data withGoodputRecorderRecord Job Start and End TimeUse the recorder object to record the job's overall start and end time.For example:defmain(argv:Sequence[str])->None:# Initialize configs…goodput_recorder.record_job_start_time(datetime.datetime.now())# TPU Initialization and device scanning…# Set up other things for the main training loop…# Main training looptrain_loop(config)goodput_recorder.record_job_end_time(datetime.datetime.now())Record Step TimeUse the recorder object to record a step's start time usingrecord_step_start_time(step_count):For example:deftrain_loop(config,state=None):# Set up mesh, model, state, checkpoint manager…# Initialize functional train arguments and model parameters…# Define the compilationforstepinnp.arange(start_step,config.steps):goodput_recorder.record_step_start_time(step)# Training step…returnstateRetrieve Goodput withGoodputCalculatorIn order to retrieve the Goodput of a job run, all you need to do is instantiate
aGoodputCalculatorobject with the job's run name and the Cloud Logging
logger name used to record data for that job run. Then call theget_job_goodputAPI to get the computed Goodput for the job run.It is recommended to make theget_job_goodputcalls for a job run from an
instance that runs elsewhere from your training machine.Create aGoodputCalculatorobjectCreate the calculator object:goodput_logger_name=f'goodput_{config.run_name}'# You can choose your own logger name.goodput_calculator=goodput.GoodputCalculator(job_name=config.run_name,logger_name=goodput_logger_name)Retrieve GoodputFinally, call theget_job_goodputAPI to retrieve Goodput for the entire job run.total_goodput=goodput_calculator.get_job_goodput()print(f"Total job goodput:{total_goodput:.2f}%")