Advanced Model Training Strategies

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import os

class DistributedTrainer:
    def __init__(self, model, train_dataset, val_dataset, config):
        self.config = config
        self.setup_distributed()
        
        # Move model to GPU and wrap with DDP
        self.model = model.to(self.device)
        self.model = DDP(self.model, device_ids=[self.local_rank])
        
        # Setup distributed data loaders
        self.train_sampler = DistributedSampler(train_dataset)
        self.train_loader = torch.utils.data.DataLoader(
            train_dataset,
            batch_size=config.batch_size,
            sampler=self.train_sampler,
            num_workers=config.num_workers,
            pin_memory=True
        )
        
        # Optimizer with scaled learning rate
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=config.learning_rate * self.world_size,  # Scale LR
            weight_decay=config.weight_decay
        )
        
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=config.num_epochs
        )
        
    def setup_distributed(self):
        """Initialize distributed training environment"""
        self.local_rank = int(os.environ.get('LOCAL_RANK', 0))
        self.world_size = int(os.environ.get('WORLD_SIZE', 1))
        self.rank = int(os.environ.get('RANK', 0))
        
        # Initialize process group
        dist.init_process_group(
            backend='nccl',  # Use NCCL for GPU communication
            init_method='env://',
            world_size=self.world_size,
            rank=self.rank
        )
        
        # Set device
        torch.cuda.set_device(self.local_rank)
        self.device = torch.device(f'cuda:{self.local_rank}')
        
    def train_epoch(self, epoch):
        """Train for one epoch with gradient synchronization"""
        self.model.train()
        self.train_sampler.set_epoch(epoch)  # Shuffle data differently each epoch
        
        total_loss = 0
        num_batches = 0
        
        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.device), target.to(self.device)
            
            self.optimizer.zero_grad()
            
            # Forward pass
            output = self.model(data)
            loss = nn.CrossEntropyLoss()(output, target)
            
            # Backward pass - gradients automatically synchronized
            loss.backward()
            
            # Gradient clipping
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            
            self.optimizer.step()
            
            total_loss += loss.item()
            num_batches += 1
            
            if batch_idx % 100 == 0 and self.rank == 0:
                print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}')
        
        # Average loss across all processes
        avg_loss = total_loss / num_batches
        dist.all_reduce(torch.tensor(avg_loss).to(self.device), op=dist.ReduceOp.AVG)
        
        return avg_loss

# Launch distributed training
def launch_distributed_training():
    # torchrun --nproc_per_node=4 --nnodes=2 train_script.py
    trainer = DistributedTrainer(model, train_dataset, val_dataset, config)
    
    for epoch in range(config.num_epochs):
        train_loss = trainer.train_epoch(epoch)
        
        if trainer.rank == 0:  # Only log on main process
            print(f"Epoch {epoch}: Average Loss = {train_loss:.4f}")

import optuna
from sklearn.model_selection import cross_val_score

def bayesian_optimization_example():
    def objective(trial):
        # Define hyperparameter search space
        lr = trial.suggest_float('learning_rate', 1e-5, 1e-1, log=True)
        batch_size = trial.suggest_categorical('batch_size', [16, 32, 64, 128])
        dropout = trial.suggest_float('dropout', 0.1, 0.5)
        hidden_size = trial.suggest_int('hidden_size', 64, 512)
        
        # Build and train model with these hyperparameters
        model = create_model(hidden_size=hidden_size, dropout=dropout)
        
        # Train and evaluate
        train_model(model, lr=lr, batch_size=batch_size)
        val_accuracy = evaluate_model(model, val_loader)
        
        return val_accuracy
    
    # Create study and optimize
    study = optuna.create_study(direction='maximize')
    study.optimize(objective, n_trials=100)
    
    print(f"Best hyperparameters: {study.best_params}")
    print(f"Best accuracy: {study.best_value}")
    
    return study.best_params