GPU-Aware Training Configuration

GPU-Aware Training Configuration

Experiment Overview

Context

Training was extremely slow on A100 Colab GPU despite using "quick_test" mode. Investigation revealed that get_auto_config(training_mode="quick_test") was returning a generic config with n_envs=256 and torch.compile=False, completely ignoring GPU capabilities.

Root Cause

The original get_auto_config() function had training modes that completely bypassed GPU detection:

# WRONG - ignores GPU capabilities
def get_auto_config(total_timesteps, training_mode="auto"):
    if training_mode == "quick_test":
        return NativePPOConfig(
            n_envs=256,           # Too low for A100!
            compile_policy=False,  # Missing 3-6x speedup!
            # ... generic settings
        )

Verified Solution

Training modes must layer on top of GPU-specific settings, not replace them:

def get_auto_config(total_timesteps=1_000_000, training_mode="auto"):
    # Step 1: ALWAYS detect GPU first
    gpu_tier = _detect_gpu_tier()  # "h100", "a100", "high", "medium", "low"

    # Step 2: Get GPU-appropriate base config
    if gpu_tier == "h100":
        config = _get_h100_base_config()
    elif gpu_tier == "a100":
        config = _get_a100_base_config()
    # ... etc

    # Step 3: Apply training mode ADJUSTMENTS (not replacements)
    if training_mode == "quick_test":
        config.total_timesteps = 10_000_000
        config.validation_interval = 25
        # BUT KEEP GPU-specific n_envs, compile_policy, etc!

GPU Configuration Matrix

Training Mode Adjustments

Training modes should ONLY adjust these parameters:

GPU Detection Code

def _detect_gpu_tier() -> str:
    """Detect GPU tier for optimal configuration."""
    if not torch.cuda.is_available():
        return "cpu"

    gpu_name = torch.cuda.get_device_name(0).lower()
    vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9

    # Check for H100 (compute capability 9.0+)
    compute_cap = torch.cuda.get_device_capability(0)
    if compute_cap[0] >= 9:
        return "h100"

    # Check for A100
    if "a100" in gpu_name:
        return "a100"

    # Tier by VRAM
    if vram_gb >= 40:
        return "high"
    elif vram_gb >= 20:
        return "medium"
    else:
        return "low"

Failed Attempts (Critical)

Diagnostic Checklist

If training is slow, check these in order:

1. FPS < 10,000 on A100? → Check n_envs (should be 1024+)
2. torch.compile: False? → Enable it (3-6x speedup after warmup)
3. GPU util < 20%? → Increase n_envs
4. Memory errors? → Decrease n_envs or minibatch_size
5. H100 with FP8=False? → Enable FP8 for additional speedup

Key Insights

• GPU detection must happen FIRST, before applying training modes
• Research shows 1000+ parallel environments needed for GPU saturation
• torch.compile provides 3-6x speedup but takes 10+ min to warmup
• FP8 is only available on Hopper architecture (H100, compute capability 9.0+)
• Training modes should adjust timesteps/​epochs, NOT hardware-specific params

Quick Fix Command

If you see slow training on A100, the config should show:

n_envs: 1024+
torch.compile: True
compile_mode: reduce-overhead

If any of these are wrong, the get_auto_config() function isn't detecting the GPU properly.

References

• LeanRL: GPU-Native RL
• Isaac Gym: High-Performance Simulation
• torch.compile documentation