debug:pytorch

What this skill does

# PyTorch Debugging Guide

This guide provides systematic approaches to debugging PyTorch models, from common tensor errors to complex training issues.

## Common Error Patterns

### 1. CUDA Out of Memory (OOM)

**Error Message:**
```
RuntimeError: CUDA out of memory. Tried to allocate X.XX GiB
```

**Causes:**
- Batch size too large for GPU memory
- Accumulating gradients without clearing
- Storing tensors on GPU unnecessarily
- Memory leaks from not detaching tensors

**Solutions:**
```python
# Check current memory usage
print(torch.cuda.memory_summary(device=None, abbreviated=False))
print(f"Allocated: {torch.cuda.memory_allocated() / 1e9:.2f} GB")
print(f"Cached: {torch.cuda.memory_reserved() / 1e9:.2f} GB")

# Clear cache
torch.cuda.empty_cache()

# Reduce batch size
batch_size = batch_size // 2

# Use gradient checkpointing for large models
from torch.utils.checkpoint import checkpoint
output = checkpoint(self.heavy_layer, input)

# Use mixed precision training
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
with autocast():
    output = model(input)
    loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# Detach tensors when storing for logging
logged_loss = loss.detach().cpu().item()

# Use gradient accumulation instead of large batches
accumulation_steps = 4
for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()
```

### 2. Tensor Size/Shape Mismatch

**Error Message:**
```
RuntimeError: size mismatch, m1: [32 x 512], m2: [256 x 10]
RuntimeError: The size of tensor a (64) must match the size of tensor b (32)
```

**Causes:**
- Incorrect layer dimensions
- Wrong tensor reshaping
- Mismatched batch sizes
- Incorrect input preprocessing

**Solutions:**
```python
# Debug by printing shapes at each layer
class DebugModel(nn.Module):
    def forward(self, x):
        print(f"Input shape: {x.shape}")
        x = self.layer1(x)
        print(f"After layer1: {x.shape}")
        x = self.layer2(x)
        print(f"After layer2: {x.shape}")
        return x

# Add shape assertions as contracts
def forward(self, x):
    assert x.dim() == 4, f"Expected 4D input, got {x.dim()}D"
    assert x.shape[1] == 3, f"Expected 3 channels, got {x.shape[1]}"
    # ... rest of forward pass

# Use einops for clearer reshaping
from einops import rearrange
x = rearrange(x, 'b c h w -> b (c h w)')

# Calculate dimensions programmatically
def _get_conv_output_size(self, shape):
    with torch.no_grad():
        dummy = torch.zeros(1, *shape)
        output = self.conv_layers(dummy)
        return output.numel()
```

### 3. NaN in Gradients/Loss

**Error Message:**
```
Loss is nan
RuntimeError: Function 'XXXBackward' returned nan values
```

**Causes:**
- Learning rate too high
- Numerical instability in operations
- Division by zero
- Log of zero or negative numbers
- Exploding gradients

**Solutions:**
```python
# Enable anomaly detection to find the source
torch.autograd.set_detect_anomaly(True)

# Check for NaN in tensors
def check_nan(tensor, name="tensor"):
    if torch.isnan(tensor).any():
        print(f"NaN detected in {name}")
        print(f"Shape: {tensor.shape}")
        print(f"NaN count: {torch.isnan(tensor).sum()}")
        raise ValueError(f"NaN in {name}")

# Add epsilon for numerical stability
eps = 1e-8
log_probs = torch.log(probs + eps)
normalized = x / (x.norm() + eps)

# Gradient clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
torch.nn.utils.clip_grad_value_(model.parameters(), clip_value=1.0)

# Check gradients after backward
def check_gradients(model):
    for name, param in model.named_parameters():
        if param.grad is not None:
            if torch.isnan(param.grad).any():
                print(f"NaN gradient in {name}")
            if torch.isinf(param.grad).any():
                print(f"Inf gradient in {name}")
            grad_norm = param.grad.norm()
            print(f"{name}: grad_norm = {grad_norm:.4f}")

# Use stable loss functions
# BAD: nn.CrossEntropyLoss on softmax output
# GOOD: nn.CrossEntropyLoss on logits (raw scores)
loss = nn.CrossEntropyLoss()(logits, targets)  # Not softmax(logits)
```

### 4. Device Mismatch (CPU/GPU)

**Error Message:**
```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same
```

**Causes:**
- Model on GPU, data on CPU (or vice versa)
- Loading model saved on different device
- Creating new tensors without specifying device
- Mixing tensors from different GPUs

**Solutions:**
```python
# Always explicitly move to device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
inputs = inputs.to(device)
targets = targets.to(device)

# Load model with map_location
model.load_state_dict(torch.load('model.pt', map_location=device))

# Create tensors on the correct device
new_tensor = torch.zeros(10, device=device)
new_tensor = torch.zeros_like(existing_tensor)  # Same device as existing

# Check device of all model parameters
def check_model_device(model):
    devices = {p.device for p in model.parameters()}
    print(f"Model parameters on devices: {devices}")

# Debug device issues
print(f"Model device: {next(model.parameters()).device}")
print(f"Input device: {inputs.device}")
print(f"Target device: {targets.device}")
```

### 5. Autograd Graph Issues

**Error Message:**
```
RuntimeError: Trying to backward through the graph a second time
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
```

**Causes:**
- Calling backward() twice without retain_graph
- In-place operations on tensors requiring gradients
- Detaching tensors incorrectly
- Not enabling gradients on input tensors

**Solutions:**
```python
# For multiple backward passes
loss.backward(retain_graph=True)  # Use sparingly - memory intensive

# Avoid in-place operations on tensors with gradients
# BAD:
x += 1
x[0] = 0
x.add_(1)

# GOOD:
x = x + 1
x = x.clone()
x[0] = 0

# Ensure requires_grad is set
input_tensor = torch.randn(10, requires_grad=True)

# Clone before in-place modification
x_modified = x.clone()
x_modified[0] = 0

# Check if tensor has gradient function
print(f"requires_grad: {tensor.requires_grad}")
print(f"grad_fn: {tensor.grad_fn}")
print(f"is_leaf: {tensor.is_leaf}")

# Properly detach for logging/storage
logged_value = tensor.detach().cpu().numpy()
```

### 6. DataLoader Problems

**Error Message:**
```
RuntimeError: DataLoader worker (pid X) is killed
BrokenPipeError: [Errno 32] Broken pipe
RuntimeError: Cannot re-initialize CUDA in forked subprocess
```

**Causes:**
- Too many workers
- Memory issues in workers
- CUDA operations before DataLoader fork
- Shared memory issues

**Solutions:**
```python
# Reduce number of workers
dataloader = DataLoader(dataset, batch_size=32, num_workers=2)

# Use spawn instead of fork for CUDA
import multiprocessing
multiprocessing.set_start_method('spawn', force=True)

# Pin memory for faster GPU transfer (but uses more memory)
dataloader = DataLoader(dataset, pin_memory=True)

# Increase shared memory for Docker
# In docker-compose.yml: shm_size: '2gb'

# Debug DataLoader issues
dataloader = DataLoader(dataset, num_workers=0)  # Single process for debugging

# Use persistent workers to avoid respawning
dataloader = DataLoader(dataset, num_workers=4, persistent_workers=True)

# Custom collate function with error handling
def safe_collate(batch):
    try:
        return torch.utils.data.dataloader.default_collate(batch)
    except Exception as e:
        print(f"