1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

import torch
import torch.nn as nn
import time
 
def benchmark_batch_parallelism():
    """Demonstrate that RNNs parallelize over batch dimension."""
    hidden_dim = 256
    input_dim = 128
    seq_len = 100
    
    rnn = nn.LSTM(input_dim, hidden_dim, batch_first=True).cuda()
    
    results = []
    for batch_size in [1, 8, 32, 128]:
        x = torch.randn(batch_size, seq_len, input_dim).cuda()
        
        # Warmup
        for _ in range(5):
            _ = rnn(x)
        
        torch.cuda.synchronize()
        start = time.perf_counter()
        for _ in range(50):
            _ = rnn(x)
        torch.cuda.synchronize()
        elapsed = (time.perf_counter() - start) / 50 * 1000
        
        per_sequence = elapsed / batch_size
        results.append((batch_size, elapsed, per_sequence))
        print(f"Batch {batch_size:3d}: {elapsed:.2f}ms total, {per_sequence:.3f}ms/seq")
    
    # Larger batches amortize sequential overhead
    speedup = results[0][2] / results[-1][2]
    print(f"\nPer-sequence speedup (batch 1 vs 128): {speedup:.1f}x")
 
if __name__ == "__main__":
    if torch.cuda.is_available():
        benchmark_batch_parallelism()
    else:
        print("CUDA required for benchmark")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

import torch
import torch.nn as nn
 
def compare_implementations():
    """Compare manual loop vs cuDNN-fused implementation."""
    hidden_dim = 512
    input_dim = 256
    seq_len = 200
    batch_size = 64
    
    # cuDNN-optimized LSTM
    lstm_fused = nn.LSTM(input_dim, hidden_dim, batch_first=True).cuda()
    
    # Manual cell-by-cell (not cuDNN optimized)
    lstm_cell = nn.LSTMCell(input_dim, hidden_dim).cuda()
    
    x = torch.randn(batch_size, seq_len, input_dim).cuda()
    
    # Benchmark fused
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    
    start.record()
    for _ in range(20):
        out_fused, _ = lstm_fused(x)
    end.record()
    torch.cuda.synchronize()
    fused_time = start.elapsed_time(end) / 20
    
    # Benchmark manual loop
    start.record()
    for _ in range(20):
        h = torch.zeros(batch_size, hidden_dim).cuda()
        c = torch.zeros(batch_size, hidden_dim).cuda()
        outputs = []
        for t in range(seq_len):
            h, c = lstm_cell(x[:, t, :], (h, c))
            outputs.append(h)
        out_manual = torch.stack(outputs, dim=1)
    end.record()
    torch.cuda.synchronize()
    manual_time = start.elapsed_time(end) / 20
    
    print(f"cuDNN fused: {fused_time:.2f}ms")
    print(f"Manual loop: {manual_time:.2f}ms")
    print(f"Speedup: {manual_time/fused_time:.1f}x")
 
# Mixed precision training
def mixed_precision_rnn():
    """Demonstrate mixed precision for RNN training."""
    model = nn.LSTM(256, 512, num_layers=2, batch_first=True).cuda()
    optimizer = torch.optim.Adam(model.parameters())
    scaler = torch.cuda.amp.GradScaler()
    
    x = torch.randn(32, 100, 256).cuda()
    target = torch.randn(32, 100, 512).cuda()
    
    # Training step with mixed precision
    with torch.cuda.amp.autocast():
        output, _ = model(x)
        loss = nn.functional.mse_loss(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    
    print(f"Output dtype: {output.dtype}")  # float16 in autocast
    print(f"Loss: {loss.item():.4f}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

import torch
import torch.nn as nn
 
class StableRNNTrainer:
    """Trainer with numerical stability safeguards."""
    
    def __init__(self, model, lr=1e-3, clip_grad=1.0):
        self.model = model
        self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        self.clip_grad = clip_grad
        
    def train_step(self, x, y):
        self.model.train()
        self.optimizer.zero_grad()
        
        # Forward with NaN detection
        output, _ = self.model(x)
        
        if torch.isnan(output).any():
            print("WARNING: NaN in forward pass!")
            return None
        
        loss = nn.functional.mse_loss(output, y)
        
        if torch.isnan(loss):
            print("WARNING: NaN loss!")
            return None
        
        # Backward
        loss.backward()
        
        # Check gradients before clipping
        total_norm = 0
        for p in self.model.parameters():
            if p.grad is not None:
                total_norm += p.grad.data.norm(2).item() ** 2
        total_norm = total_norm ** 0.5
        
        if torch.isnan(torch.tensor(total_norm)):
            print("WARNING: NaN in gradients!")
            return None
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_grad)
        
        self.optimizer.step()
        
        return {'loss': loss.item(), 'grad_norm': total_norm}
 
 
def demonstrate_gradient_clipping():
    """Show importance of gradient clipping for RNNs."""
    torch.manual_seed(42)
    
    # Intentionally bad initialization to cause gradient explosion
    model = nn.RNN(64, 128, batch_first=True)
    with torch.no_grad():
        model.weight_hh_l0.mul_(3.0)  # Large weights
    
    x = torch.randn(16, 50, 64)
    y = torch.randn(16, 50, 128)
    
    # Without clipping
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    out, _ = model(x)
    loss = nn.functional.mse_loss(out, y)
    loss.backward()
    
    grad_norm_before = sum(p.grad.norm().item()**2 for p in model.parameters())**0.5
    print(f"Gradient norm before clipping: {grad_norm_before:.2e}")
    
    # Apply clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    
    grad_norm_after = sum(p.grad.norm().item()**2 for p in model.parameters())**0.5
    print(f"Gradient norm after clipping: {grad_norm_after:.2f}")
 
if __name__ == "__main__":
    demonstrate_gradient_clipping()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

import torch
import torch.nn as nn
 
class StreamingRNNInference:
    """RNN inference handler for streaming applications."""
    
    def __init__(self, model_path: str):
        self.model = torch.jit.load(model_path)
        self.model.eval()
        self.hidden = None
    
    def process_chunk(self, input_chunk: torch.Tensor) -> torch.Tensor:
        """Process a chunk of streaming input, maintaining state."""
        with torch.no_grad():
            output, self.hidden = self.model(input_chunk, self.hidden)
        return output
    
    def reset_state(self):
        """Reset state for new sequence."""
        self.hidden = None
 
 
def export_for_deployment(model: nn.Module, example_input: torch.Tensor):
    """Export model for production deployment."""
    model.eval()
    
    # TorchScript tracing
    traced = torch.jit.trace(model, example_input)
    traced.save("model_traced.pt")
    
    # Optimize for inference
    optimized = torch.jit.optimize_for_inference(traced)
    optimized.save("model_optimized.pt")
    
    # ONNX export
    torch.onnx.export(
        model, example_input,
        "model.onnx",
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch', 1: 'seq_len'}}
    )
    
    print("Exported: model_traced.pt, model_optimized.pt, model.onnx")