Mixed Precision Training

1Mixed Precision Training Flow:
2
3┌─────────────────────────────────────────────────────────┐
4│                    FP32 Master Weights                   │
5│                         (stored)                         │
6└────────────────────────┬────────────────────────────────┘
7                         │ Cast to FP16
8                         ▼
9┌─────────────────────────────────────────────────────────┐
10│              FP16 Working Weights                        │
11│                (for forward/backward)                    │
12└────────────────────────┬────────────────────────────────┘
13                         │
14         ┌───────────────┴───────────────┐
15         ▼                               ▼
16    ┌─────────┐                    ┌─────────┐
17    │ Forward │                    │Backward │
18    │  Pass   │───→ FP32 Loss ───→│  Pass   │
19    │ (FP16)  │                    │ (FP16)  │
20    └─────────┘                    └────┬────┘
21                                        │ FP16 gradients
22                                        │ (scaled)
23                                        ▼
24                              ┌──────────────────┐
25                              │  Unscale to FP32 │
26                              │  Update master   │
27                              │    weights       │
28                              └──────────────────┘

1import torch
2from torch.cuda.amp import autocast, GradScaler
3
4def train_mixed_precision(
5    model: nn.Module,
6    train_loader: DataLoader,
7    optimizer: torch.optim.Optimizer,
8    criterion: nn.Module,
9    epochs: int,
10    device: torch.device = torch.device('cuda')
11) -> dict:
12    """
13    Training with automatic mixed precision.
14
15    Uses FP16 for forward/backward, FP32 for weight updates.
16    """
17    # Initialize gradient scaler
18    scaler = GradScaler()
19
20    history = {'train_loss': []}
21
22    for epoch in range(epochs):
23        model.train()
24        epoch_loss = 0.0
25
26        for batch in train_loader:
27            x, y = batch
28            x, y = x.to(device), y.to(device)
29
30            optimizer.zero_grad()
31
32            # Forward pass with autocast (FP16)
33            with autocast():
34                pred = model(x)
35                loss = criterion(pred, y)
36
37            # Backward pass with scaled gradients
38            scaler.scale(loss).backward()
39
40            # Unscale gradients and update weights
41            scaler.step(optimizer)
42
43            # Update scale factor for next iteration
44            scaler.update()
45
46            epoch_loss += loss.item()
47
48        history['train_loss'].append(epoch_loss / len(train_loader))
49
50    return history

1def train_mixed_precision_with_clipping(
2    model: nn.Module,
3    train_loader: DataLoader,
4    optimizer: torch.optim.Optimizer,
5    criterion: nn.Module,
6    max_grad_norm: float = 1.0,
7    device: torch.device = torch.device('cuda')
8):
9    """
10    Mixed precision training with gradient clipping.
11
12    The key is to unscale gradients BEFORE clipping.
13    """
14    scaler = GradScaler()
15
16    for batch in train_loader:
17        x, y = batch
18        x, y = x.to(device), y.to(device)
19
20        optimizer.zero_grad()
21
22        # Forward with autocast
23        with autocast():
24            pred = model(x)
25            loss = criterion(pred, y)
26
27        # Backward with scaled gradients
28        scaler.scale(loss).backward()
29
30        # Unscale gradients BEFORE clipping
31        scaler.unscale_(optimizer)
32
33        # Now clip in FP32 space
34        grad_norm = torch.nn.utils.clip_grad_norm_(
35            model.parameters(),
36            max_norm=max_grad_norm
37        )
38
39        # Step (skips if gradients contain inf/NaN)
40        scaler.step(optimizer)
41
42        # Update scale
43        scaler.update()

1def train_epoch_complete(
2    model: nn.Module,
3    train_loader: DataLoader,
4    optimizer: torch.optim.Optimizer,
5    criterion: nn.Module,
6    scaler: GradScaler,
7    ema: ExponentialMovingAverage,
8    max_grad_norm: float = 1.0,
9    accumulation_steps: int = 2,
10    device: torch.device = torch.device('cuda')
11) -> dict:
12    """
13    Complete training epoch with all enhancements:
14    - Mixed precision (autocast + GradScaler)
15    - Gradient clipping
16    - Gradient accumulation
17    - EMA weight tracking
18    """
19    model.train()
20    total_loss = 0.0
21    num_steps = 0
22
23    for batch_idx, (x, y) in enumerate(train_loader):
24        x, y = x.to(device), y.to(device)
25
26        # Zero gradients at start of accumulation cycle
27        if batch_idx % accumulation_steps == 0:
28            optimizer.zero_grad()
29
30        # Forward pass with mixed precision
31        with autocast():
32            pred = model(x)
33            loss = criterion(pred, y)
34            # Scale loss for accumulation
35            loss = loss / accumulation_steps
36
37        # Backward with scaled gradients
38        scaler.scale(loss).backward()
39
40        # Update weights after accumulation cycle
41        if (batch_idx + 1) % accumulation_steps == 0:
42            # Unscale gradients before clipping
43            scaler.unscale_(optimizer)
44
45            # Clip gradients in FP32 space
46            grad_norm = torch.nn.utils.clip_grad_norm_(
47                model.parameters(),
48                max_norm=max_grad_norm
49            )
50
51            # Optimizer step (skipped if inf/NaN gradients)
52            scaler.step(optimizer)
53            scaler.update()
54
55            # Update EMA weights
56            ema.update(model)
57
58            num_steps += 1
59
60        total_loss += loss.item() * accumulation_steps
61
62    return {
63        'loss': total_loss / len(train_loader),
64        'steps': num_steps
65    }