AI Runtime Notebook quickstart

Important

This feature is in Public Preview.

This quickstart shows the fastest way to run code on AI Runtime, the serverless GPU compute for deep learning. You attach a notebook to a GPU and run a short PyTorch training step, with no cluster setup.

Tip

  • Attach a notebook to Serverless GPU to start running on a GPU in minutes.
  • Start with 1xA10 for development, then scale up for larger models or distributed training.
  • No cluster setup is required, and the Standard environment already includes torch.

Before you begin

Your workspace must have AI Runtime enabled and be in a supported region. See Requirements.

Step 1: Attach a notebook to a GPU

  1. Create or open a notebook.
  2. From the compute drop-down at the top of the notebook, select Serverless GPU.
  3. In the Accelerator field, select 1xA10.
  4. Click Apply, and then Confirm.

Tip

Start with 1xA10 for development and small workloads. Move to 1xH100 or 8xH100 when you need more GPU memory or multi-GPU distributed training. See Hardware options.

Step 2: Run your first GPU code

Run the following in a notebook cell to confirm the GPU is attached and run one training step:

import torch

# Confirm a GPU is attached.
print(torch.cuda.get_device_name(0))

# Run a minimal training step on the GPU.
model = torch.nn.Linear(10, 1).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
x, y = torch.randn(32, 10).cuda(), torch.randn(32, 1).cuda()

loss = torch.nn.functional.mse_loss(model(x), y)
loss.backward()
optimizer.step()
print("loss:", loss.item())

You are now running on serverless GPU compute.

Next steps

  • Last updated on