Train and benchmark AI workloads with DeepSpeed on Google Cloud C4A Axion VMs
Introduction
Understand DeepSpeed and Google Axion C4A for AI training
Create a Google Axion C4A virtual machine for DeepSpeed
Set up PyTorch and DeepSpeed on a Google Axion C4A virtual machine
Train and benchmark AI workloads on an Arm-based Google Axion virtual machine
Next Steps

  Log an issue

  Fork and edit

  Discuss on Discord

Set up the virtual machine

Create a Google Axion C4A Arm-based virtual machine (VM) on Google Cloud Platform. You’ll use the c4a-standard-4 machine type with 4 vCPUs and 16 GB of memory. This VM will host PyTorch and DeepSpeed training and benchmarking workloads.

Note

For help with Google Cloud Platform setup, see the Learning Path Getting started with Google Cloud Platform .

To create a C4A virtual machine in the Google Cloud console:

  1. Navigate to the Google Cloud console .
  2. Go to Compute Engine > VM Instances and select Create Instance.
  3. Under Machine configuration, populate fields such as Instance name, Region, and Zone.
  4. Set Series to C4A, then select c4a-standard-4 for Machine type.

Image Alt Text:Screenshot of the Google Cloud console showing the Machine configuration section. The Series dropdown is set to C4A and the machine type c4a-standard-4 is selectedConfiguring machine type to C4A in Google Cloud Console

  1. Under OS and storage, select Change and then choose an Arm64-based operating system image. For this Learning Path, select SUSE Linux Enterprise Server.
  2. For the license type, choose Pay as you go.
  3. Increase Size (GB) from 10 to 100 to allocate sufficient disk space, and then select Choose.
  4. Select Create to launch the virtual machine.

After the instance starts, select SSH next to the VM in the instance list to open a browser-based terminal session.

Image Alt Text:Google Cloud console VM instances page displaying running instance with green checkmark and SSH button in the Connect columnConnecting to a running C4A VM using SSH

A new browser window opens with a terminal connected to your VM.

Image Alt Text:Browser-based SSH terminal connected to the Google Axion C4A VM. The shell prompt confirms that the instance is running and ready for the next step, where you’ll install DeepSpeed and its dependencies.Terminal session connected to the VM

What you’ve accomplished and what’s next

You’ve now provisioned a Google Axion C4A Arm VM and connected to it using SSH.

Next, you’ll install PyTorch and DeepSpeed and configure the Python environment for AI training.

Back
Next