Google Colab Mastery

Before you learn anything else, master your environment.

Every minute you spend fighting Colab is a minute not spent understanding transformers. Every GPU error you debug is cognitive load stolen from actual learning.

This chapter eliminates that friction.

The 5-Minute Setup

Step 1: Create a new Colab notebook

Go to colab.research.google.com and create a new notebook.

Step 2: Enable GPU

Click Runtime → Change runtime type → T4 GPU → Save

Do this EVERY TIME you open a notebook. Without GPU, nothing works.

Step 3: Run the setup cell

Copy this into your first cell and run it:

# ARENA Environment Setup
import os
import sys

# Check GPU
import torch
if not torch.cuda.is_available():
    raise SystemExit("❌ No GPU! Go to Runtime → Change runtime type → T4 GPU")

print(f"✅ GPU: {torch.cuda.get_device_name(0)}")
print(f"✅ Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

# Mount Google Drive
from google.colab import drive
drive.mount('/content/drive', force_remount=False)

# Create working directory
ARENA_DIR = '/content/drive/MyDrive/ARENA_3.0'
os.makedirs(ARENA_DIR, exist_ok=True)
print(f"✅ Working directory: {ARENA_DIR}")

print("\n🚀 Ready!")

If you see three green checkmarks, you're ready.

The 5 Skills That Matter

Skill 1: Installing Packages 

!pip install einops transformers transformer_lens -q

The -q flag suppresses output. Always use it.

Run package installation in its own cell. Wait for it to complete before running the next cell.

Skill 2: Saving to Google Drive 

import shutil
from datetime import datetime

def save_to_drive(local_path, folder="checkpoints"):
    """Save file to Google Drive with timestamp."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    drive_path = f"/content/drive/MyDrive/ARENA_3.0/{folder}"
    os.makedirs(drive_path, exist_ok=True)

    name, ext = os.path.splitext(os.path.basename(local_path))
    dest = f"{drive_path}/{name}_{timestamp}{ext}"
    shutil.copy(local_path, dest)
    print(f"✅ Saved: {dest}")

# Usage:
save_to_drive("model.pt", "milestone_1")

Colab sessions disconnect. Google Drive persists. Save early, save often.

Skill 3: Memory Management 

def clear_memory():
    """Free GPU memory."""
    import gc
    gc.collect()
    torch.cuda.empty_cache()

    allocated = torch.cuda.memory_allocated() / 1e9
    print(f"✅ Memory cleared. Using: {allocated:.2f} GB")

# Use when you see "CUDA out of memory"
clear_memory()

Skill 4: Preventing Disconnection 

from IPython.display import Javascript

# Run once at session start
display(Javascript('''
    setInterval(() => {
        document.querySelector("colab-connect-button")?.click();
        console.log("Keeping alive...");
    }, 60000);
'''))

Colab disconnects after ~90 minutes of inactivity. This keeps it alive.

Skill 5: Common Error Fixes

Error Solution
"No GPU" Runtime → Change runtime type → T4 GPU
"CUDA out of memory" clear_memory() or reduce batch size
"Session crashed (RAM)" Runtime → Restart runtime
"No module named X" !pip install X -q
"Drive already mounted" This is fine. Continue.

Keyboard Shortcuts

Shortcut Action
Ctrl+Enter Run cell
Shift+Enter Run cell, move to next
Ctrl+M B Insert cell below
Ctrl+S Save notebook

The Colab Mastery Test

Before proceeding, complete these tasks:

Target time: Under 10 minutes

If you can do all five in under 10 minutes, Colab is no longer a distraction. It's invisible. All your cognitive resources go to learning.

That's the goal.