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

About this Learning Path

Skill level:  Introductory
Reading time:  30 min
Last updated:  31 Jul 2026
Skill level:
  Introductory
Reading time:
  30 min
Last updated:
  31 Jul 2026
Author:Pareena Verma, Arm
Arm IP:Neoverse
Tags: ML Google Cloud Linux DeepSpeed PyTorch Python
Author:
  • Pareena Verma, Arm
Arm IP:
Neoverse
Tags:
ML Google Cloud Linux DeepSpeed PyTorch Python

Who is this for?

This is an introductory topic for DevOps engineers, ML engineers, and software developers who want to run AI training and benchmarking workloads using PyTorch and DeepSpeed on SUSE Linux Enterprise Server (SLES) Arm64, validate CPU-based neural network execution, and benchmark AI performance on Arm processors.

What will you learn?

Upon completion of this Learning Path, you will be able to:

  • Install and configure PyTorch and DeepSpeed on Arm-based Google Cloud C4A Axion virtual machines (VMs).
  • Create and execute neural network training workloads using PyTorch.
  • Benchmark CPU-based AI workloads on Arm64 processors.
  • Validate scalable AI execution and workload performance on Google Axion Arm VMs.

Prerequisites

Before starting, you will need the following:

Summary

AI-assisted

This summary was drafted with an approved AI-assisted workflow and reviewed by Arm contributors before publication. Human technical review remains part of the process so the final page reflects engineering rigor, accuracy, and Arm editorial standards.

Close
?
You’ll provision a Google Cloud C4A Arm VM running SUSE Linux, set up Python 3.11, and install PyTorch and DeepSpeed for CPU training and benchmarking. You’ll verify the aarch64 architecture and Neoverse-V2 cores with uname and lscpu, create a dedicated virtual environment, and run a baseline model followed by a larger benchmark. From the training logs and benchmark output, you’ll confirm correct execution.

Frequently asked questions

AI-assisted

These FAQs were drafted with an approved AI-assisted workflow and reviewed by Arm contributors before publication. Human technical review remains part of the process so the final page reflects engineering rigor, accuracy, and Arm editorial standards.

Close
?
How do I know the VM is Arm64 and running on Axion cores?
Run uname -m and confirm that the output is aarch64. Then run lscpu and check that the model name reports Neoverse-V2.
Which Google Cloud VM configuration should I use for the steps?
Use the c4a-standard-4 machine type with 4 vCPUs and 16 GB of memory. This configuration hosts the PyTorch and DeepSpeed workloads.
Which Python version do I need and what is the virtual environment called?
Install Python 3.11 and create a virtual environment for the project. The path uses an environment named deepspeed-env.
I opened a new SSH session. What should I do before running workloads?
Re-activate the deepspeed-env virtual environment and navigate to the ~/deepspeed-demo directory. This ensures the correct dependencies and paths are active.
What result should I expect from the baseline and benchmark runs?
The baseline model should run without errors and print training logs, confirming that PyTorch and DeepSpeed are correctly installed. The larger benchmark runs longer and lets you observe CPU scaling behavior by reviewing its logs and runtime characteristics.
Next