Properly training machine learning models can be a difficult and time-consuming task. Unlike typical software development, training results depend on both the training code and the input data. This can make debugging training jobs a complicated process, even when running on your local machine. Running code on remote infrastructure can make this task even more difficult.
Debugging code that runs in a managed cloud environment can be a tedious and error-prone process because the standard tools used to debug programs locally are not available in a managed environment. Also, if logs and metrics aren’t visible, training jobs can get bogged down and stop progressing. Interactive access to jobs can make the entire debugging process much easier.
This article introduces a new tool available to users of Vertex AI custom training jobs: Interactive Shell. This feature gives you direct, shell-like access to the VM running your code, allowing you to run arbitrary commands to profile and debug issues that cannot be resolved with logs or monitoring metrics. You can also run commands using the same credentials as your training code to investigate permission issues or other issues that you can’t reproduce locally. Access to the interactive shell is authenticated using the same IAM permission set used for regular custom training jobs, providing a secure interface to the Vertex AI training environment.
Example: TensorFlow distributed training
Let’s look at one example where using Vertex AI’s interactive shell can help you debug your training program. In this case, you intentionally submit a job to Vertex AI Training that deadlocks and stops progressing. Use py-spy in an interactive shell to understand the root cause of the problem.
Vertex AI is a managed ML platform that provides a convenient way to scale up your training jobs to take advantage of additional computing resources. To run a TensorFlow trainer across multiple nodes or accelerators, you can use TensorFlow’s distributed strategy API, the TensorFlow module for performing distributed computations. To use multiple workers, each with one or more GPUs, use: tf.distribute.MultiWorkerMirroredStrategyuses an all-reduce algorithm to synchronize gradient updates across multiple devices.
Set the code
Use the example from the Vertex AI Multi-Worker Training codelab. In this codelab, you will train an image classification model on the Tensorflow Cassava dataset using a ResNet50 model pretrained on Imagenet. Run a training job on multiple nodes using: tf.distribute.MultiWorkerMirroredStrategy.
Codelab creates a custom container for your training code and pushes it to Google Container Registry (GCR) in your GCP project.
