# Train and benchmark AI workloads with DeepSpeed on Google Cloud C4A Axion VMs

## In this learning path

- [Introduction](https://learn.arm.com/learning-paths/servers-and-cloud-computing/deepspeed-on-axion/)
- [Understand DeepSpeed and Google Axion C4A for AI training](https://learn.arm.com/learning-paths/servers-and-cloud-computing/deepspeed-on-axion/background/)
- [Create a Google Axion C4A virtual machine for DeepSpeed](https://learn.arm.com/learning-paths/servers-and-cloud-computing/deepspeed-on-axion/instance/)
- [Set up PyTorch and DeepSpeed on a Google Axion C4A virtual machine](https://learn.arm.com/learning-paths/servers-and-cloud-computing/deepspeed-on-axion/install-deepspeed-arm/)
- [Train and benchmark AI workloads on an Arm-based Google Axion virtual machine](https://learn.arm.com/learning-paths/servers-and-cloud-computing/deepspeed-on-axion/train-benchmark-deepspeed-arm/)
- [Next Steps](https://learn.arm.com/learning-paths/servers-and-cloud-computing/deepspeed-on-axion/_next-steps/)

## About this Learning Path

| Skill level:        | Introductory           |
|---------------------|------------------------|
| Reading time:       | 30 min                 |
| Last updated:       | 31 Jul 2026            |

| Author:          | Pareena Verma, Arm [GitHub](https://github.com/pareenaverma) [LinkedIn](https://linkedin.com/in/pareena-verma-7853607) |
|------------------|------------------------------------------------------------------------------------------------|
| Arm IP:          | [Neoverse](https://support.arm.com/?tab=compute-ip&Product%20Type=Infrastructure%20Processors)  |
| Tags:            | [ML](https://learn.arm.com/tag/ml) [Google Cloud](https://learn.arm.com/tag/google-cloud) [Linux](https://learn.arm.com/tag/linux) [DeepSpeed](https://learn.arm.com/tag/deepspeed) [PyTorch](https://learn.arm.com/tag/pytorch) [Python](https://learn.arm.com/tag/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:

- A [Google Cloud Platform (GCP)](https://cloud.google.com/free) account with billing enabled
- Basic familiarity with Python and machine learning concepts

### Summary

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

#### 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.