PyTorch

26min
Running PyTorch on Vast.ai: A Complete Guide
Introduction
This guide walks you through setting up and running PyTorch workloads on Vast.ai, a marketplace for renting GPU compute power. Whether you're training large models or running inference, this guide will help you get started efficiently.
Prerequisites
A Vast.ai account
Basic familiarity with PyTorch
﻿Install TLS Certificate for Jupyter﻿
﻿(Optional) SSH client installed on your local machine and SSH public key added in Account tab at cloud.vast.ai 
﻿(Optional) Install and use vast-cli 
﻿(Optional) Docker knowledge for custom environments﻿
Setting Up Your Environment
1. Selecting PyTorch Template
Navigate to the Templates tab to view available templates. Before choosing a specific instance, you'll need to select the appropriate PyTorch template for your needs:
Choose recommended PyTorch (cuDNN Runtime) template if:
You're running standard training and inference workloads
You're using pre-built PyTorch functions and layers
You don't need to compile custom CUDA kernels
You want a smaller container size and faster instance startup
You're running production inference workloads
Choose recommended PyTorch (cuDNN Devel) template if:
You need to build custom CUDA extensions
You're developing new GPU operations
You're using libraries that require CUDA compilation (like some versions of Flash Attention)
You need to modify or compile PyTorch from source
You're doing PyTorch development or research requiring low-level GPU access
2. Choosing an Instance
Click the play button to select the template and see GPUs you can rent. For PyTorch workloads, consider:
GPU Memory: Minimum 8GB for most models
CUDA Version: PyTorch 2.0+ works best with CUDA 11.7 or newer
Disk Space: Minimum 50GB for datasets and checkpoints
Internet Speed: Look for instances with >100 Mbps for dataset downloads
Rent the GPU of your choice.
3. Connecting to Your Instance
Click blue button on instance card in Instances tab when it says "Open" to access Jupyter.
Setting Up Your PyTorch Environment
1. Basic Environment Check
Verify your setup by executing these commands in Python's Interactive Shell in a Jupyter terminal:
Python
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU device: {torch.cuda.get_device_name(0)}")
import torch
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU device: {torch.cuda.get_device_name(0)}")
﻿
2. Data Management
For efficient data handling:
a) Fast local storage:
Bash
mkdir /workspace/data
cd /workspace/data
mkdir /workspace/data
cd /workspace/data
﻿
b) Dataset downloads:
Bash
# Using wget
wget your_dataset_url

# Using git lfs for larger files: https://git-lfs.com/
sudo apt-get install git-lfs
git lfs install
git clone your_dataset_repo
# Using wget
wget your_dataset_url

# Using git lfs for larger files: https://git-lfs.com/
sudo apt-get install git-lfs
git lfs install
git clone your_dataset_repo
﻿
Training Best Practices
Checkpoint Management
Always save checkpoints to prevent data loss:
Python
checkpoint_dir = '/workspace/checkpoints'
os.makedirs(checkpoint_dir, exist_ok=True)

checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, f'{checkpoint_dir}/checkpoint_{epoch}.pt')
checkpoint_dir = '/workspace/checkpoints'
os.makedirs(checkpoint_dir, exist_ok=True)

checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}
torch.save(checkpoint, f'{checkpoint_dir}/checkpoint_{epoch}.pt')
﻿
Resource Monitoring
Monitor GPU usage:
Bash
watch -n 1 nvidia-smi
watch -n 1 nvidia-smi
﻿
Or in Python:
Python
def print_gpu_utilization():
    print(torch.cuda.memory_allocated() / 1024**2, "MB Allocated")
    print(torch.cuda.memory_reserved() / 1024**2, "MB Reserved")
def print_gpu_utilization():
    print(torch.cuda.memory_allocated() / 1024**2, "MB Allocated")
    print(torch.cuda.memory_reserved() / 1024**2, "MB Reserved")
﻿
Cost Optimization
Instance Selection
Use vast cli search offers command to search for machines that fit your budget 
Monitor your spending in Vast.ai's Billing tab
Resource Utilization
Use appropriate batch sizes to maximize GPU utilization
Enable gradient checkpointing for large models
Implement early stopping to avoid unnecessary compute time
Troubleshooting
Common Issues and Solutions
Out of Memory (OOM) Errors
Reduce batch size
Enable gradient checkpointing
Use mixed precision training
Python
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)
scaler.scale(loss).backward()
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    outputs = model(inputs)
    loss = criterion(outputs, labels)
scaler.scale(loss).backward()
﻿
Slow Training
Check GPU utilization
Verify data loading pipeline
Consider using torch.compile() for PyTorch 2.0+
Python
model = torch.compile(model)
model = torch.compile(model)
﻿
Connection Issues
Use tmux or screen for persistent sessions
Set up automatic reconnection in your SSH config
Best Practices
Environment Management
Document your setup and requirements
Keep track of software versions
Data Management
Use data versioning tools
Implement proper data validation
Set up efficient data loading pipelines
Training Management
Implement logging (e.g., WandB, TensorBoard)
Set up experiment tracking
Use configuration files for hyperparameters
Advanced Topics
Multi-GPU Training
For distributed training:
Python
model = torch.nn.DataParallel(model)
model = torch.nn.DataParallel(model)
﻿
Mixed Precision Training
Enable AMP for faster training:
Python
from torch.cuda.amp import autocast

with autocast():
    outputs = model(inputs)
from torch.cuda.amp import autocast

with autocast():
    outputs = model(inputs)
﻿
Custom Docker Images
Create a custom Docker image from your own Dockerfile and create your own template as needed:
Dockerfile
FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime

# Install additional dependencies
RUN pip install wandb tensorboard

# Add your custom requirements
COPY requirements.txt .
RUN pip install -r requirements.txt
FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime

# Install additional dependencies
RUN pip install wandb tensorboard

# Add your custom requirements
COPY requirements.txt .
RUN pip install -r requirements.txt
﻿
Conclusion
Running PyTorch on Vast.ai provides a cost-effective way to rent cheap GPUs and accelerate deep learning workloads. By following this guide and best practices, you can efficiently set up and manage your PyTorch workloads while optimizing costs and performance.
Additional Resources
﻿PyTorch Documentation 
﻿Vast.ai Documentation 
﻿PyTorch Performance Tuning Guide 
﻿
Updated 21 Jan 2025
Did this page help you?
Example 2 - How to Create Template for GROBID
Linux Virtual Machines