AI Workloads on amplitUDE

amplitUDE provides state-of-the-art GPU resources for artificial intelligence and machine learning workloads. This guide helps you get started with deep learning, large language models, and other AI applications on amplitUDE’s GPU infrastructure.

---

Overview

amplitUDE is optimized for AI/ML workloads with:

• NVIDIA H200 GPUs - 141 GB HBM3 memory, latest generation

• NVIDIA H100 GPUs - 80 GB HBM2e memory

• High-speed Lustre storage - Optimized for large dataset access

• Python/Conda environments - Pre-configured scientific computing stack

• Apptainer containers - Reproducible workflows with NVIDIA NGC support

---

Quick Start

1. Set Up Python Environment

See: Python Environments on amplitUDE

# Load Anaconda module
module load anaconda3/2023

# Create PyTorch environment
mamba create -n pytorch-gpu python=3.11 \
    pytorch pytorch-cuda=12.1 -c pytorch -c nvidia

# Activate
conda activate pytorch-gpu

2. Request GPU Resources

# Interactive session
salloc --partition=GPU-H200 --gres=gpu:1 --time=02:00:00 --mem=64G
srun --pty bash

# Batch job
sbatch train_model.sh

3. Verify GPU Access

import torch
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU: {torch.cuda.get_device_name(0)}")
# Output: CUDA available: True
#         GPU: NVIDIA H200

---

Getting Started Guides

Working with amplitUDE Filesystems and AI

Storage Best Practices for AI Workloads

amplitUDE provides two main storage areas for AI work:

Filesystem	Path	Quota	Use For	Performance
Home	$HOME	0.5 TB	Code, environments, configs	Medium
Scratch	$SCRATCH//lustre/scratch/$USER	10 TB	Datasets, checkpoints, outputs	High (parallel I/O)

Filesystem Strategy for AI

# Project structure recommendation
$HOME/
├── projects/my-ml-project/     # Code and scripts
│   ├── train.py
│   ├── models/                  # Model definitions
│   └── configs/                 # Configuration files
└── envs/                        # Conda environments

$SCRATCH/
├── datasets/                    # Large training datasets
│   ├── imagenet/
│   └── squad/
├── checkpoints/                 # Model checkpoints
│   └── my-model/
└── outputs/                     # Training outputs/logs

Optimize Dataset Loading

Problem: AI datasets often have millions of small files, which is inefficient on Lustre.

Solutions:

1. Use archive formats (recommended for datasets >10k files)

   # Create tar archive
   tar -cf imagenet.tar imagenet/
   
   # Or use WebDataset format
   pip install webdataset
   

2. Pre-load to node-local storage (for small datasets)

   #!/bin/bash
   #SBATCH --gres=gpu:1
   
   # Copy dataset to node's /tmp
   cp -r $SCRATCH/datasets/small_dataset /tmp/
   
   # Train from /tmp (much faster random access)
   python train.py --data-path /tmp/small_dataset
   

3. Use HDF5 or LMDB for random access

   import h5py
   import lmdb
   
   # Store dataset in single file
   with h5py.File('dataset.h5', 'w') as f:
       f.create_dataset('images', data=images)
       f.create_dataset('labels', data=labels)
   

Set Cache Directories

# Add to ~/.bashrc or job script
export HF_HOME=$SCRATCH/.cache/huggingface
export TORCH_HOME=$SCRATCH/.cache/torch
export TRANSFORMERS_CACHE=$SCRATCH/.cache/transformers
export HF_DATASETS_CACHE=$SCRATCH/.cache/datasets

---

Installing Python Software for AI

See: Python Environments on amplitUDE for complete guide.

Quick Setup:

module load anaconda3/2023

# Install mamba (faster than conda)
conda install mamba -c conda-forge

# Create environment for deep learning
mamba create -n deep-learning python=3.11 \
    pytorch torchvision torchaudio pytorch-cuda=12.1 \
    numpy scipy pandas matplotlib \
    -c pytorch -c nvidia -c conda-forge

conda activate deep-learning

# Install Hugging Face ecosystem
pip install transformers datasets accelerate evaluate
pip install tensorboard wandb  # Monitoring tools

---

PyTorch on amplitUDE

Hardware-Specific Optimizations

amplitUDE’s H200 and H100 GPUs support advanced features. Enable them for maximum performance:

import torch

# Enable TF32 for faster training (H100/H200)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# Enable cuDNN benchmarking (finds fastest algorithms)
torch.backends.cudnn.benchmark = True

# Check GPU compute capability
print(torch.cuda.get_device_capability())  # (9, 0) for H100/H200

Common Issues and Solutions

Issue 1: CUDA Out of Memory

# Solutions:
# 1. Reduce batch size
batch_size = 16  # Instead of 32

# 2. Use gradient accumulation
for i, batch in enumerate(dataloader):
    loss = model(batch)
    loss = loss / accumulation_steps
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

# 3. Use mixed precision training
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    loss = model(batch)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# 4. Enable gradient checkpointing
model.gradient_checkpointing_enable()

Issue 2: Slow Data Loading

# Use multiple workers (set to number of CPU cores allocated)
from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=32,
    num_workers=8,  # Use your --cpus-per-task value
    pin_memory=True,  # Faster GPU transfer
    persistent_workers=True  # Keep workers alive between epochs
)

Issue 3: Multi-GPU Training Not Working

# In Slurm script, request multiple GPUs:
#SBATCH --gres=gpu:2

# Use PyTorch DistributedDataParallel
python -m torch.distributed.launch \
    --nproc_per_node=2 \
    --nnodes=1 \
    train.py

Example: Complete PyTorch Training Script

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Enable optimizations for H200/H100
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True

# Check GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
print(f"GPU: {torch.cuda.get_device_name(0)}")

# Data
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

train_dataset = datasets.MNIST('/lustre/scratch/$USER/datasets', 
                               train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, 
                         num_workers=4, pin_memory=True)

# Model
model = nn.Sequential(
    nn.Flatten(),
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
).to(device)

# Training
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}')

# Save model
torch.save(model.state_dict(), '/lustre/scratch/$USER/models/mnist_model.pth')

---

Distributed Training on amplitUDE

Single-Node Multi-GPU Training

#!/bin/bash
#SBATCH --job-name=multi-gpu-train
#SBATCH --partition=GPU-Big
#SBATCH --gres=gpu:4           # Request 4 GPUs
#SBATCH --time=08:00:00
#SBATCH --mem=256G
#SBATCH --cpus-per-task=32

module load anaconda3/2023
conda activate pytorch-gpu

# PyTorch DDP
torchrun --nproc_per_node=4 train_ddp.py

Training Script (train_ddp.py):

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def main():
    # Initialize distributed training
    dist.init_process_group(backend='nccl')
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    
    # Create model and move to GPU
    model = MyModel().to(local_rank)
    model = DDP(model, device_ids=[local_rank])
    
    # Use DistributedSampler for data loading
    train_sampler = DistributedSampler(train_dataset)
    train_loader = DataLoader(train_dataset, sampler=train_sampler, 
                             batch_size=32, num_workers=4)
    
    # Training loop
    for epoch in range(num_epochs):
        train_sampler.set_epoch(epoch)  # Shuffle data differently each epoch
        
        for batch in train_loader:
            # Training step
            pass
    
    dist.destroy_process_group()

if __name__ == '__main__':
    main()

Multi-Node Training (Advanced)

#!/bin/bash
#SBATCH --nodes=2              # 2 nodes
#SBATCH --gres=gpu:4           # 4 GPUs per node = 8 total GPUs
#SBATCH --ntasks-per-node=4    # 4 tasks (1 per GPU)

module load anaconda3/2023
conda activate pytorch-gpu

# Get master node address
MASTER_ADDR=$(scontrol show hostname $SLURM_NODELIST | head -n1)
MASTER_PORT=29500

srun torchrun \
    --nnodes=$SLURM_NNODES \
    --nproc_per_node=4 \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    train_ddp.py

---

Handling Datasets with Many Files

Problem: ImageNet has 1.2M images. Loading from Lustre is slow with many small files.

Solution 1: WebDataset Format (Recommended)

import webdataset as wds

# Create WebDataset (once, during preprocessing)
with wds.ShardWriter("imagenet-%06d.tar", maxcount=10000) as sink:
    for image, label in dataset:
        sink.write({
            "__key__": f"sample{i:07d}",
            "jpg": image,
            "cls": label
        })

# Load during training (fast!)
dataset = wds.WebDataset("imagenet-{000000..000119}.tar") \
    .decode("pil") \
    .to_tuple("jpg", "cls") \
    .batched(32)

Solution 2: LMDB (For Random Access)

import lmdb
import pickle

# Create LMDB (once)
env = lmdb.open('imagenet.lmdb', map_size=1099511627776)  # 1TB
with env.begin(write=True) as txn:
    for i, (image, label) in enumerate(dataset):
        txn.put(f'{i}'.encode(), pickle.dumps((image, label)))

# Load during training
env = lmdb.open('imagenet.lmdb', readonly=True)
with env.begin() as txn:
    data = pickle.loads(txn.get(f'{idx}'.encode()))

Solution 3: Copy to /tmp (For Small Datasets)

#!/bin/bash
#SBATCH --gres=gpu:1

# Copy to node-local storage
cp -r $SCRATCH/datasets/cifar10 /tmp/

# Train from /tmp
python train.py --data /tmp/cifar10

# Clean up
rm -rf /tmp/cifar10

---

Monitoring and Visualization

TensorBoard

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter(log_dir='/lustre/scratch/$USER/runs/experiment1')

# Log metrics
for epoch in range(num_epochs):
    writer.add_scalar('Loss/train', train_loss, epoch)
    writer.add_scalar('Accuracy/train', train_acc, epoch)
    writer.add_scalar('Learning_rate', lr, epoch)

writer.close()

View TensorBoard:

# On compute node
module load anaconda3/2023
conda activate pytorch-gpu

tensorboard --logdir=/lustre/scratch/$USER/runs --port=6006 --bind_all

# SSH tunnel from your laptop
ssh -L 6006:ndgh2001:6006 username@login2.amplitude.hpc.uni-due.de

# Open browser to http://localhost:6006

Weights & Biases

import wandb

wandb.login()  # Set API key once

# Initialize
wandb.init(project="my-project", name="experiment1")

# Log during training
for epoch in range(num_epochs):
    wandb.log({
        "train_loss": train_loss,
        "train_acc": train_acc,
        "learning_rate": lr
    })

---

Using Apptainer for AI Workloads

See: Apptainer Containers on amplitUDE for complete guide.

Quick Start with NVIDIA NGC Containers

module load apptainer

# Pull PyTorch container from NVIDIA
apptainer pull docker://nvcr.io/nvidia/pytorch:24.01-py3

# Run training
apptainer exec --nv \
    --bind $SCRATCH:/workspace \
    pytorch_24.01-py3.sif \
    python /workspace/train.py

Popular NGC Containers for AI:

Container	Use Case	Pull Command
PyTorch	Deep learning	docker://nvcr.io/nvidia/pytorch:24.01-py3
TensorFlow	Deep learning	docker://nvcr.io/nvidia/tensorflow:24.01-tf2-py3
RAPIDS	Data science (GPU)	docker://nvcr.io/nvidia/rapidsai/base:24.02-cuda12.0-py3.11
TensorRT	Inference optimization	docker://nvcr.io/nvidia/tensorrt:24.01-py3
Triton	Model serving	docker://nvcr.io/nvidia/tritonserver:24.01-py3

---

Advanced Topics

Large Language Model Training and Fine-Tuning

See: Fine-Tuning LLMs on amplitUDE for complete guide.

Quick Example: LoRA Fine-Tuning

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model

# Load base model
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# Configure LoRA
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05
)

# Apply LoRA
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Output: trainable params: 8.4M || all params: 8000M || trainable%: 0.11

# Train (only LoRA weights are updated)
trainer.train()

---

Model Inference and Serving

vLLM for Fast LLM Inference

# Install vLLM
pip install vllm

# Start inference server
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-3.1-8B \
    --gpu-memory-utilization 0.9 \
    --max-model-len 8192

Batch Inference Example:

from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-3.1-8B")
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

prompts = [
    "Explain quantum computing",
    "Write a Python function for quicksort"
]

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    print(f"Prompt: {output.prompt}")
    print(f"Output: {output.outputs[0].text}")

---

Example Job Scripts

Basic Training Job

#!/bin/bash
#SBATCH --job-name=train-model
#SBATCH --partition=GPU-Big
#SBATCH --gres=gpu:1
#SBATCH --time=08:00:00
#SBATCH --mem=64G
#SBATCH --cpus-per-task=8
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err

mkdir -p logs

# Environment
module load anaconda3/2023
conda activate pytorch-gpu

# Verify GPU
nvidia-smi
python -c "import torch; print(f'GPU: {torch.cuda.get_device_name(0)}')"

# Train
python train.py \
    --data-path $SCRATCH/datasets/imagenet \
    --epochs 100 \
    --batch-size 256 \
    --output-dir $SCRATCH/checkpoints/resnet50

Hyperparameter Tuning with Optuna

#!/bin/bash
#SBATCH --job-name=hpo
#SBATCH --partition=GPU-Big
#SBATCH --gres=gpu:1
#SBATCH --array=0-9    # 10 parallel trials
#SBATCH --time=04:00:00
#SBATCH --mem=32G

module load anaconda3/2023
conda activate pytorch-gpu

# Each array task runs one trial
python hpo_search.py --trial-id $SLURM_ARRAY_TASK_ID

---

Best Practices

Resource Allocation

1. Start small, scale up

   • Test on 1 GPU first

   • Then scale to multi-GPU

   • Finally multi-node if needed

2. Match resources to task

   # Small model (<1B params): 1 GPU
   #SBATCH --gres=gpu:1
   
   # Medium model (7-13B): 1 H200 GPU
   #SBATCH --partition=GPU-H200
   
   # Large model (70B+): Multiple GPUs
   #SBATCH --gres=gpu:4
   

3. Set appropriate time limits

   # Too short: Job killed before completion
   # Too long: Longer queue wait
   
   # Rule of thumb: request 20% more than estimated
   #SBATCH --time=10:00:00  # If you estimate 8 hours
   

Code Organization

project/
├── data/                    # Dataset loading code
│   ├── __init__.py
│   └── dataset.py
├── models/                  # Model definitions
│   ├── __init__.py
│   └── resnet.py
├── utils/                   # Helper functions
│   ├── __init__.py
│   └── metrics.py
├── configs/                 # Configuration files
│   └── default.yaml
├── scripts/                 # Slurm job scripts
│   ├── train.sh
│   └── evaluate.sh
├── train.py                 # Main training script
├── evaluate.py              # Evaluation script
└── requirements.txt         # Dependencies

Checkpointing

# Save checkpoints regularly
if epoch % save_every == 0:
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
    }
    torch.save(checkpoint, f'$SCRATCH/checkpoints/model_epoch_{epoch}.pth')

# Resume from checkpoint
if os.path.exists(checkpoint_path):
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    start_epoch = checkpoint['epoch'] + 1

---

Troubleshooting

Common Issues

GPU not visible

# Check on compute node (not login node!)
salloc --partition=GPU-H200 --gres=gpu:1
srun --pty bash
nvidia-smi

Out of memory

• Reduce batch size

• Enable gradient checkpointing

• Use mixed precision (FP16/BF16)

• Use gradient accumulation

Slow training

• Check data loading (use profiler)

• Increase num_workers in DataLoader

• Use pin_memory=True

• Pre-process data offline

Multi-GPU not working

• Use torchrun instead of python

• Check NCCL_DEBUG=INFO for errors

• Verify all GPUs visible with nvidia-smi

---

Further Reading

• Python Environments - Environment setup

• LLM Fine-Tuning - Complete LLM guide

• Apptainer Containers - Container workflows

• amplitUDE Hardware - System specifications

---

Last updated: May 2026