GPU infrastructure planning in Prompt Engineering / GenAI - ML Experiment: Train & Evaluate

Experiment - GPU infrastructure planning

Problem:You want to train a deep learning model that requires a lot of GPU power. Currently, you have a single GPU setup, but training takes too long and sometimes runs out of memory.

Current Metrics:Training time per epoch: 45 minutes; GPU memory usage: 95%; Validation accuracy: 82%

Issue:The model training is slow and sometimes crashes due to GPU memory limits. This limits experimentation and model improvements.

Your Task

Plan and implement a GPU infrastructure setup that reduces training time per epoch to under 20 minutes and avoids out-of-memory errors, while maintaining or improving validation accuracy above 82%.

You can only use up to 2 GPUs.

You must keep the same model architecture and dataset.

You cannot reduce batch size below 32.

Hint 1

Hint 2

Hint 3

Solution

Prompt Engineering / GenAI

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.cuda.amp import GradScaler, autocast

# Define model (example: simple CNN)
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 32, 3, 1)
        self.fc = nn.Linear(5408, 10)
    def forward(self, x):
        x = self.conv(x)
        x = torch.relu(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

# Setup device and model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleCNN()
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
model.to(device)

# Data loaders
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST('.', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)

# Optimizer and loss
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Mixed precision scaler
scaler = GradScaler()

# Training loop
model.train()
for epoch in range(1):
    total_loss = 0
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        with autocast():
            output = model(data)
            loss = criterion(output, target)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        total_loss += loss.item()
    print(f'Epoch {epoch+1}, Loss: {total_loss/len(train_loader):.4f}')

Used nn.DataParallel to utilize 2 GPUs for parallel training.

Increased batch size to 64 to better use GPU memory across devices.

Implemented mixed precision training with torch.cuda.amp to reduce memory usage and speed up computation.

Corrected the input feature size of the fully connected layer from 21632 to 5408 to match the output of the convolutional layer.

Results Interpretation

Before: Training time 45 min, GPU memory 95%, Accuracy 82%

After: Training time 18 min, GPU memory 80%, Accuracy 83%

Using multiple GPUs with data parallelism and mixed precision training can reduce training time and memory usage, helping avoid crashes and improving efficiency without losing accuracy.

Bonus Experiment

Try using gradient accumulation to simulate larger batch sizes on a single GPU and compare training time and accuracy.

💡 Hint

Accumulate gradients over multiple smaller batches before updating model weights to fit larger effective batch size in limited memory.

Practice

1. Why is it important to plan GPU infrastructure before starting a GenAI project?

easy

A. To reduce the size of the AI model automatically

B. To ensure the GPU has enough memory and speed for the AI model

C. Because GPUs are always cheaper than CPUs

D. To avoid using any GPUs and rely only on CPUs

Solution

1. Step 1: Understand GPU role in AI projects

   GPUs speed up AI model training and need enough memory to handle data.

2. Step 2: Importance of matching GPU specs to model needs

   Choosing a GPU with insufficient memory or speed will slow down or fail the project.

3. Final Answer:

   To ensure the GPU has enough memory and speed for the AI model -> Option B

4. Quick Check:

   GPU specs must match AI needs = D [OK]

Hint: Match GPU memory and speed to your AI model size [OK]

Common Mistakes:

• Thinking CPUs can replace GPUs for heavy AI tasks
• Assuming all GPUs have the same performance
• Ignoring GPU memory limits

2. Which of the following is the correct way to check GPU memory using Python's PyTorch library?

easy

A. torch.cuda.memory_size()

B. torch.gpu.memory.total()

C. torch.cuda.get_device_properties(0).total_memory

D. torch.device.memory()

Solution

1. Step 1: Recall PyTorch GPU memory query syntax

   The correct method is torch.cuda.get_device_properties(device_id).total_memory.

2. Step 2: Check each option for correctness

   Only torch.cuda.get_device_properties(0).total_memory uses the correct PyTorch function and attribute.

3. Final Answer:

   torch.cuda.get_device_properties(0).total_memory -> Option C

4. Quick Check:

   Correct PyTorch GPU memory call = C [OK]

Hint: Use torch.cuda.get_device_properties(0).total_memory to check GPU memory [OK]

Common Mistakes:

• Using non-existent PyTorch functions
• Confusing device and memory functions
• Missing the device index argument

3. Given this Python code snippet using PyTorch, what will be printed?

import torch
if torch.cuda.is_available():
    mem = torch.cuda.get_device_properties(0).total_memory
    print(mem > 8_000_000_000)
else:
    print(False)

medium

A. True if GPU memory is more than 8GB, else False

B. Always True

C. Always False

D. Raises an error if no GPU

Solution

1. Step 1: Understand the code logic

   The code checks if a GPU is available, then compares its memory to 8GB (8 billion bytes).

2. Step 2: Determine output based on GPU memory

   If GPU memory is greater than 8GB, it prints True; otherwise, False. If no GPU, prints False.

3. Final Answer:

   True if GPU memory is more than 8GB, else False -> Option A

4. Quick Check:

   GPU memory check > 8GB = A [OK]

Hint: Check GPU memory size condition to predict output [OK]

Common Mistakes:

• Assuming always True regardless of GPU
• Expecting error if no GPU instead of False
• Confusing bytes with gigabytes

4. Identify the error in this GPU memory check code and select the fix:

import torch
if torch.cuda.is_available():
    mem = torch.cuda.get_device_properties().total_memory
    print(mem)
else:
    print('No GPU')

medium

A. Add device index 0 in get_device_properties: get_device_properties(0)

B. Replace torch.cuda.is_available() with torch.has_cuda()

C. Use torch.cuda.memory_allocated() instead of get_device_properties()

D. No error, code is correct

Solution

1. Step 1: Check get_device_properties usage

   The function requires a device index argument, e.g., 0 for the first GPU.

2. Step 2: Identify the fix

   Adding (0) fixes the error. Other options are incorrect or unnecessary.

3. Final Answer:

   Add device index 0 in get_device_properties: get_device_properties(0) -> Option A

4. Quick Check:

   Missing device index causes error = B [OK]

Hint: Always provide device index to get_device_properties() [OK]

Common Mistakes:

• Omitting device index argument
• Using non-existent torch.has_cuda()
• Confusing memory functions

5. You plan to train a large GenAI model requiring 24GB GPU memory. Your local GPUs have 16GB each. Which is the best GPU infrastructure planning choice?

hard

A. Ignore memory limits and expect training to succeed

B. Reduce the model size to fit 16GB GPU and train locally

C. Train on CPU only to avoid GPU memory limits

D. Use multiple GPUs with model parallelism or switch to cloud GPUs with 24GB+ memory

Solution

1. Step 1: Analyze GPU memory requirement vs available hardware

   The model needs 24GB, but local GPUs have only 16GB, so one GPU is insufficient.

2. Step 2: Consider solutions for insufficient GPU memory

   Using multiple GPUs with model parallelism or cloud GPUs with enough memory solves the problem effectively.

3. Final Answer:

   Use multiple GPUs with model parallelism or switch to cloud GPUs with 24GB+ memory -> Option D

4. Quick Check:

   Match GPU memory to model needs with parallelism or cloud = A [OK]

Hint: Use multi-GPU or cloud GPUs for models needing more memory [OK]

Common Mistakes:

• Trying to train large models on insufficient GPU memory
• Ignoring cloud GPU options
• Assuming CPU can replace GPU for large models