4.4 Cost Optimization & Sustainability in Large-Scale Training

import torchimport torch.nn as nnimport torch.optim as optimimport timefrom torch.cuda.amp import autocast, GradScaler # Define a simple modelclass SimpleModel(nn.Module):    def __init__(self, dim=2048):        super().__init__()        self.layers = nn.Sequential(            nn.Linear(dim, dim*2),            nn.ReLU(),            nn.Linear(dim*2, dim*2),            nn.ReLU(),            nn.Linear(dim*2, dim)        )        def forward(self, x):        return self.layers(x) # Set random seed for reproducibilitytorch.manual_seed(42) # Create model and move to GPUmodel = SimpleModel().cuda()print(f"Model has {sum(p.numel() for p in model.parameters())} parameters") # Choose optimizeroptimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-2) # Create gradient scaler for mixed precision trainingscaler = GradScaler() # Training parametersbatch_size = 32input_dim = 2048epochs = 5 # Track metricstimes = []losses = [] # Training loopfor epoch in range(epochs):    epoch_start = time.time()    epoch_losses = []        # Inner training loop (simplified)    for i in range(10):        # Generate random data (in real scenarios, use DataLoader)        x = torch.randn(batch_size, input_dim).cuda()        y = torch.randn(batch_size, input_dim).cuda()                # Reset gradients        optimizer.zero_grad()                # Forward pass with autocast for mixed precision        with autocast():            out = model(x)            loss = ((out - y) ** 2).mean()  # MSE loss                # Backward pass with scaling        scaler.scale(loss).backward()                # Optimizer step with unscaling        scaler.step(optimizer)                # Update scaler for next iteration        scaler.update()                # Record loss        epoch_losses.append(loss.item())        # Calculate epoch statistics    epoch_time = time.time() - epoch_start    times.append(epoch_time)    avg_loss = sum(epoch_losses) / len(epoch_losses)    losses.append(avg_loss)        print(f"Epoch {epoch+1}/{epochs}: Loss={avg_loss:.6f}, Time={epoch_time:.3f}s") # Report final statisticsprint(f"Average epoch time: {sum(times)/len(times):.3f}s")print(f"Final loss: {losses[-1]:.6f}")print(f"Loss reduction: {(losses[0] - losses[-1])/losses[0]*100:.2f}%") 

import torchimport torch.nn as nnfrom torch.utils.checkpoint import checkpointimport timeimport matplotlib.pyplot as pltimport numpy as np # Define a more complex model that represents a transformer-like blockclass TransformerBlock(nn.Module):    def __init__(self, dim, expansion_factor=4):        super().__init__()        # Self-attention component (simplified)        self.attention = nn.Sequential(            nn.Linear(dim, dim),            nn.ReLU(),            nn.Linear(dim, dim)        )                # Feed-forward network        self.ffn = nn.Sequential(            nn.Linear(dim, dim * expansion_factor),            nn.ReLU(),            nn.Linear(dim * expansion_factor, dim)        )                self.layer_norm1 = nn.LayerNorm(dim)        self.layer_norm2 = nn.LayerNorm(dim)            def forward(self, x):        # Residual connection with layer norm        residual = x        x = self.layer_norm1(x)        x = self.attention(x)        x = x + residual                # Second residual connection        residual = x        x = self.layer_norm2(x)        x = self.ffn(x)        x = x + residual                return x # Create a deep model with multiple transformer blocksclass DeepTransformer(nn.Module):    def __init__(self, dim, depth):        super().__init__()        self.blocks = nn.ModuleList([TransformerBlock(dim) for _ in range(depth)])            def forward(self, x, use_checkpointing=False):        for block in self.blocks:            if use_checkpointing:                x = checkpoint(block, x)            else:                x = block(x)        return x # Benchmark function to compare memory and time with and without checkpointingdef benchmark_checkpointing(batch_size=16, dim=1024, depth=12, seq_len=512):    # Create input tensor    x = torch.randn(batch_size, seq_len, dim).cuda()        # Create model and move to GPU    model = DeepTransformer(dim, depth).cuda()        results = {}        # Test without checkpointing    torch.cuda.empty_cache()    torch.cuda.reset_peak_memory_stats()    start_time = time.time()        # Forward pass    with torch.cuda.amp.autocast():        try:            model(x, use_checkpointing=False)                        # Record results            results['standard_time'] = time.time() - start_time            results['standard_memory'] = torch.cuda.max_memory_allocated() / (1024 ** 3)  # Convert to GB            results['standard_success'] = True        except RuntimeError as e:            if "out of memory" in str(e).lower():                results['standard_success'] = False                results['standard_memory'] = None                results['standard_time'] = None                print("Standard forward pass ran out of memory")            else:                raise e        # Test with checkpointing    torch.cuda.empty_cache()    torch.cuda.reset_peak_memory_stats()    start_time = time.time()        # Forward pass with checkpointing    with torch.cuda.amp.autocast():        try:            model(x, use_checkpointing=True)                        # Record results            results['checkpointed_time'] = time.time() - start_time            results['checkpointed_memory'] = torch.cuda.max_memory_allocated() / (1024 ** 3)  # Convert to GB            results['checkpointed_success'] = True        except RuntimeError as e:            if "out of memory" in str(e).lower():                results['checkpointed_success'] = False                results['checkpointed_memory'] = None                results['checkpointed_time'] = None                print("Checkpointed forward pass ran out of memory")            else:                raise e        return results # Run the benchmarkresults = benchmark_checkpointing() # Print resultsprint("\n--- BENCHMARK RESULTS ---")if results.get('standard_success'):    print(f"Standard forward pass:")    print(f"  Time: {results['standard_time']:.4f} seconds")    print(f"  Memory: {results['standard_memory']:.2f} GB")else:    print("Standard forward pass: OUT OF MEMORY") if results.get('checkpointed_success'):    print(f"\nCheckpointed forward pass:")    print(f"  Time: {results['checkpointed_time']:.4f} seconds")    print(f"  Memory: {results['checkpointed_memory']:.2f} GB")else:    print("\nCheckpointed forward pass: OUT OF MEMORY") # If both methods succeeded, show comparisonif results.get('standard_success') and results.get('checkpointed_success'):    memory_reduction = (results['standard_memory'] - results['checkpointed_memory']) / results['standard_memory'] * 100    time_increase = (results['checkpointed_time'] - results['standard_time']) / results['standard_time'] * 100        print("\nComparison:")    print(f"  Memory reduction with checkpointing: {memory_reduction:.1f}%")    print(f"  Time increase with checkpointing: {time_increase:.1f}%")        # Create a visualization    if plt:        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))                # Memory plot        bars1 = ax1.bar(['Standard', 'Checkpointed'],                        [results['standard_memory'], results['checkpointed_memory']],                       color=['blue', 'green'])        ax1.set_ylabel('Memory Usage (GB)')        ax1.set_title('Peak Memory Usage')        ax1.bar_label(bars1, fmt='%.2f GB')                # Time plot        bars2 = ax2.bar(['Standard', 'Checkpointed'],                        [results['standard_time'], results['checkpointed_time']],                       color=['blue', 'green'])        ax2.set_ylabel('Time (seconds)')        ax2.set_title('Forward Pass Time')        ax2.bar_label(bars2, fmt='%.4f s')                plt.tight_layout()        plt.savefig('checkpointing_benchmark.png')        print("\nBenchmark visualization saved as 'checkpointing_benchmark.png'") # Example of checkpointing with backward passdef demonstrate_backward_pass():    # Set up a simple example    dim = 1024    batch_size = 16    model = TransformerBlock(dim).cuda()    x = torch.randn(batch_size, dim, requires_grad=True).cuda()    target = torch.randn(batch_size, dim).cuda()    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)        # Without checkpointing    optimizer.zero_grad()    out1 = model(x)    loss1 = ((out1 - target) ** 2).mean()    loss1.backward()    grad1 = {name: param.grad.clone() for name, param in model.named_parameters()}        # Reset gradients    optimizer.zero_grad()        # With checkpointing    out2 = checkpoint(model, x)    loss2 = ((out2 - target) ** 2).mean()    loss2.backward()    grad2 = {name: param.grad.clone() for name, param in model.named_parameters()}        # Verify gradients are the same    all_close = True    for name in grad1:        if not torch.allclose(grad1[name], grad2[name], atol=1e-5):            all_close = False            break        print("\n--- GRADIENT VERIFICATION ---")    print(f"Gradients match between standard and checkpointed versions: {all_close}")    print(f"Output values match: {torch.allclose(out1, out2, atol=1e-5)}") # Run gradient verificationdemonstrate_backward_pass() # Demonstrate a concrete exampledef run_concrete_example():    # Create a simple block and input    block = TransformerBlock(1024).cuda()    x = torch.randn(16, 1024).cuda()        # Run without checkpointing    y1 = block(x)        # Run with checkpointing    y2 = checkpoint(block, x)        # Check shapes and values    print("\n--- CONCRETE EXAMPLE ---")    print(f"Output shape: {y1.shape}")    print(f"Outputs are identical: {torch.allclose(y1, y2)}") run_concrete_example() 

import osimport timeimport signalimport argparseimport torchimport torch.nn as nnimport torch.distributed as distfrom torch.nn.parallel import DistributedDataParallel as DDPfrom transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizerfrom transformers import get_linear_schedule_with_warmupfrom datasets import load_datasetfrom torch.utils.data import DataLoader, DistributedSamplerimport boto3from botocore.exceptions import ClientError class SpotTrainingManager:    def __init__(self, model, optimizer, scheduler, args):        self.model = model        self.optimizer = optimizer        self.scheduler = scheduler        self.args = args        self.epoch = 0        self.global_step = 0        self.best_val_loss = float('inf')        self.checkpoint_dir = args.checkpoint_dir        self.s3_bucket = args.s3_bucket                # Create local checkpoint directory if it doesn't exist        os.makedirs(self.checkpoint_dir, exist_ok=True)                # Set up termination signal handler        signal.signal(signal.SIGTERM, self._termination_handler)            def _termination_handler(self, signum, frame):        """Handle spot instance termination notice"""        print("⚠️ Termination signal received! Saving checkpoint before shutdown...")        self.save_checkpoint(is_emergency=True)        print("Emergency checkpoint saved. Shutting down...")        exit(0)        def save_checkpoint(self, is_best=False, is_emergency=False):        """Save model checkpoint locally and to S3"""        if dist.get_rank() != 0:            return  # Only save checkpoint from the main process                    checkpoint = {            'epoch': self.epoch,            'global_step': self.global_step,            'model_state_dict': self.model.module.state_dict() if hasattr(self.model, 'module') else self.model.state_dict(),            'optimizer_state_dict': self.optimizer.state_dict(),            'scheduler_state_dict': self.scheduler.state_dict() if self.scheduler else None,            'best_val_loss': self.best_val_loss        }                # Determine checkpoint path        if is_emergency:            checkpoint_path = os.path.join(self.checkpoint_dir, 'emergency_checkpoint.pt')        elif is_best:            checkpoint_path = os.path.join(self.checkpoint_dir, 'best_checkpoint.pt')        else:            checkpoint_path = os.path.join(self.checkpoint_dir, f'checkpoint_epoch_{self.epoch}.pt')                    # Save locally        torch.save(checkpoint, checkpoint_path)        print(f"Checkpoint saved locally to {checkpoint_path}")                # Upload to S3        if self.s3_bucket:            try:                s3_client = boto3.client('s3')                s3_path = os.path.basename(checkpoint_path)                s3_client.upload_file(checkpoint_path, self.s3_bucket, f"checkpoints/{s3_path}")                print(f"Checkpoint uploaded to s3://{self.s3_bucket}/checkpoints/{s3_path}")            except ClientError as e:                print(f"S3 upload failed: {e}")        def load_latest_checkpoint(self):        """Load the most recent checkpoint from S3 or local storage"""        # First try to download from S3        if self.s3_bucket:            try:                s3_client = boto3.client('s3')                objects = s3_client.list_objects_v2(Bucket=self.s3_bucket, Prefix="checkpoints/")                if 'Contents' in objects:                    checkpoints = [obj for obj in objects['Contents'] if obj['Key'].endswith('.pt')]                    if checkpoints:                        # Sort by last modified time                        latest = sorted(checkpoints, key=lambda x: x['LastModified'], reverse=True)[0]                        local_path = os.path.join(self.checkpoint_dir, os.path.basename(latest['Key']))                        s3_client.download_file(self.s3_bucket, latest['Key'], local_path)                        print(f"Downloaded checkpoint from S3: {latest['Key']}")                        return self._load_checkpoint_file(local_path)            except ClientError as e:                print(f"S3 download failed: {e}")                # If S3 fails or no S3 bucket, try local checkpoints        checkpoint_files = [f for f in os.listdir(self.checkpoint_dir) if f.endswith('.pt')]        if checkpoint_files:            # Check for emergency checkpoint first            if 'emergency_checkpoint.pt' in checkpoint_files:                checkpoint_path = os.path.join(self.checkpoint_dir, 'emergency_checkpoint.pt')                print("Found emergency checkpoint, loading...")                return self._load_checkpoint_file(checkpoint_path)                        # Then check for best checkpoint            if 'best_checkpoint.pt' in checkpoint_files:                checkpoint_path = os.path.join(self.checkpoint_dir, 'best_checkpoint.pt')                print("Found best checkpoint, loading...")                return self._load_checkpoint_file(checkpoint_path)                        # Otherwise, load latest epoch checkpoint            epoch_checkpoints = [f for f in checkpoint_files if f.startswith('checkpoint_epoch_')]            if epoch_checkpoints:                # Extract epoch numbers and find the latest                epochs = [int(f.split('_')[-1].split('.')[0]) for f in epoch_checkpoints]                latest_epoch = max(epochs)                checkpoint_path = os.path.join(self.checkpoint_dir, f'checkpoint_epoch_{latest_epoch}.pt')                print(f"Loading checkpoint from epoch {latest_epoch}")                return self._load_checkpoint_file(checkpoint_path)                print("No checkpoints found. Starting from scratch.")        return False        def _load_checkpoint_file(self, checkpoint_path):        """Load a specific checkpoint file"""        try:            checkpoint = torch.load(checkpoint_path, map_location='cpu')                        # Load model state            if hasattr(self.model, 'module'):                self.model.module.load_state_dict(checkpoint['model_state_dict'])            else:                self.model.load_state_dict(checkpoint['model_state_dict'])                            # Load optimizer and scheduler states            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])            if self.scheduler and checkpoint['scheduler_state_dict']:                self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])                            # Restore training state            self.epoch = checkpoint['epoch']            self.global_step = checkpoint['global_step']            self.best_val_loss = checkpoint['best_val_loss']                        print(f"Resumed from epoch {self.epoch}, global step {self.global_step}")            return True        except Exception as e:            print(f"Failed to load checkpoint: {e}")            return False def setup_distributed_training(rank, world_size):    """Initialize distributed training environment"""    os.environ['MASTER_ADDR'] = 'localhost'    os.environ['MASTER_PORT'] = '12355'    dist.init_process_group("nccl", rank=rank, world_size=world_size)    torch.cuda.set_device(rank) def load_and_prepare_data(args, tokenizer):    """Load and prepare dataset for training"""    # Load dataset    dataset = load_dataset('wikitext', 'wikitext-103-v1')        # Tokenize function    def tokenize_function(examples):        return tokenizer(examples['text'], truncation=True, max_length=args.max_seq_length)        # Apply tokenization    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=['text'])        # Create DataLoaders    train_sampler = DistributedSampler(tokenized_dataset['train']) if dist.is_initialized() else None    val_sampler = DistributedSampler(tokenized_dataset['validation']) if dist.is_initialized() else None        train_loader = DataLoader(        tokenized_dataset['train'],         batch_size=args.batch_size,        sampler=train_sampler,        shuffle=train_sampler is None    )        val_loader = DataLoader(        tokenized_dataset['validation'],        batch_size=args.batch_size,        sampler=val_sampler,        shuffle=False    )        return train_loader, val_loader, train_sampler def train_model(rank, world_size, args):    """Main training function for each process"""    if world_size > 1:        setup_distributed_training(rank, world_size)        # Load model, tokenizer    config = GPT2Config.from_pretrained(args.model_name)    model = GPT2LMHeadModel.from_pretrained(args.model_name, config=config)    tokenizer = GPT2Tokenizer.from_pretrained(args.model_name)        # Move model to GPU    model = model.to(rank)        # Set up distributed model if needed    if world_size > 1:        model = DDP(model, device_ids=[rank])        # Prepare optimizer and scheduler    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)    train_loader, val_loader, train_sampler = load_and_prepare_data(args, tokenizer)        total_steps = len(train_loader) * args.num_epochs    scheduler = get_linear_schedule_with_warmup(        optimizer,         num_warmup_steps=args.warmup_steps,        num_training_steps=total_steps    )        # Initialize the spot training manager    trainer = SpotTrainingManager(model, optimizer, scheduler, args)        # Try to load checkpoint    resumed = trainer.load_latest_checkpoint()        # Main training loop    model.train()    for epoch in range(trainer.epoch, args.num_epochs):        trainer.epoch = epoch        if train_sampler:            train_sampler.set_epoch(epoch)                    # Track time for each epoch        epoch_start_time = time.time()                # Training loop        for step, batch in enumerate(train_loader):            # Move batch to device            batch = {k: v.to(rank) for k, v in batch.items()}                        # Forward pass            outputs = model(**batch, labels=batch['input_ids'])            loss = outputs.loss                        # Backward pass            loss.backward()                        # Gradient clipping            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)                        # Update parameters            optimizer.step()            scheduler.step()            optimizer.zero_grad()                        trainer.global_step += 1                        # Periodic logging            if rank == 0 and step % args.logging_steps == 0:                print(f"Epoch: {epoch}, Step: {step}, Loss: {loss.item():.4f}")                        # Periodic checkpoint            if (rank == 0 and                 trainer.global_step % args.save_steps == 0 and                 trainer.global_step > 0):                trainer.save_checkpoint()                        # Periodically check for spot instance termination            if step % args.termination_check_steps == 0:                if check_for_termination_notice():                    # This will trigger the signal handler                    print("Termination notice detected, preparing for shutdown...")                    trainer.save_checkpoint(is_emergency=True)                    exit(0)                # End of epoch        epoch_time = time.time() - epoch_start_time        if rank == 0:            print(f"Epoch {epoch} completed in {epoch_time:.2f} seconds")                # Validation at end of epoch        if rank == 0:            val_loss = validate(model, val_loader, rank)            print(f"Validation loss: {val_loss:.4f}")                        # Save if best model            if val_loss < trainer.best_val_loss:                trainer.best_val_loss = val_loss                trainer.save_checkpoint(is_best=True)                        # Always save at end of epoch            trainer.save_checkpoint()        # Clean up    if world_size > 1:        dist.destroy_process_group() def validate(model, val_loader, device):    """Validate the model on validation dataset"""    model.eval()    total_loss = 0    with torch.no_grad():        for batch in val_loader:            batch = {k: v.to(device) for k, v in batch.items()}            outputs = model(**batch, labels=batch['input_ids'])            total_loss += outputs.loss.item()        avg_loss = total_loss / len(val_loader)    model.train()    return avg_loss def check_for_termination_notice():    """Check if AWS has sent a spot termination notice"""    try:        # On AWS, spot termination notices are available at this URL        response = requests.get(            "http://169.254.169.254/latest/meta-data/spot/instance-action",            timeout=0.1        )        if response.status_code == 200:            # Termination notice received            return True    except:        # Any error means no termination notice or not on AWS        pass    return False def parse_args():    parser = argparse.ArgumentParser(description="Elastic training with spot instances")    parser.add_argument("--model_name", type=str, default="gpt2", help="Model name or path")    parser.add_argument("--batch_size", type=int, default=8, help="Batch size per GPU")    parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")    parser.add_argument("--num_epochs", type=int, default=3, help="Number of epochs")    parser.add_argument("--max_seq_length", type=int, default=512, help="Maximum sequence length")    parser.add_argument("--warmup_steps", type=int, default=500, help="Warmup steps")    parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Gradient clipping norm")    parser.add_argument("--logging_steps", type=int, default=100, help="Log every X steps")    parser.add_argument("--save_steps", type=int, default=1000, help="Save checkpoint every X steps")    parser.add_argument("--termination_check_steps", type=int, default=50, help="Check for spot termination every X steps")    parser.add_argument("--checkpoint_dir", type=str, default="./checkpoints", help="Directory for checkpoints")    parser.add_argument("--s3_bucket", type=str, default=None, help="S3 bucket for checkpoints")    return parser.parse_args() if __name__ == "__main__":    args = parse_args()        # Determine world size and run training    world_size = torch.cuda.device_count()        if world_size > 1:        import torch.multiprocessing as mp        mp.spawn(            train_model,            args=(world_size, args),            nprocs=world_size,            join=True        )    else:        train_model(0, 1, args) 

# Example implementation of memory-efficient optimizersimport torchimport mathfrom torch.optim import Optimizer class Adafactor(Optimizer):    """    Implements Adafactor optimizer from Google Research    (https://arxiv.org/abs/1804.04235)    """    def __init__(self, params, lr=None, beta1=0.9, eps=(1e-30, 1e-3),                 clip_threshold=1.0, decay_rate=-0.8, weight_decay=0.0):        defaults = dict(lr=lr, beta1=beta1, eps=eps,                        clip_threshold=clip_threshold,                        decay_rate=decay_rate, weight_decay=weight_decay)        super(Adafactor, self).__init__(params, defaults)     def _get_lr(self, param_group, param_state):        if param_group['lr'] is None:  # Use adaptive learning rate            return min(1.0, 1.0 / math.sqrt(param_state['step']))        else:            return param_group['lr']     def _factored(self, shape):        """Whether to use factored second moment estimates"""        return len(shape) >= 2     def _compute_factored_second_moment(self, exp_avg_sq_row, exp_avg_sq_col, grad):        """Compute factored second moment statistics"""        row_mean = torch.mean(grad * grad, dim=-1, keepdim=True)        col_mean = torch.mean(grad * grad, dim=-2, keepdim=True)                # Update factored second moment estimates        beta2 = 1.0 - (1.0 / exp_avg_sq_row.shape[0])  # Decreasing beta for larger matrices        exp_avg_sq_row.mul_(beta2).add_(row_mean, alpha=(1.0 - beta2))        exp_avg_sq_col.mul_(beta2).add_(col_mean, alpha=(1.0 - beta2))                # Compute scaling factors        return exp_avg_sq_row, exp_avg_sq_col     def step(self, closure=None):        """Performs a single optimization step"""        loss = None        if closure is not None:            loss = closure()         for group in self.param_groups:            for p in group['params']:                if p.grad is None:                    continue                grad = p.grad.data                                # Handle 16-bit gradients                if grad.dtype == torch.float16:                    grad = grad.float()                 if grad.is_sparse:                    raise RuntimeError("Adafactor does not support sparse gradients")                 state = self.state[p]                                # State initialization                if len(state) == 0:                    state['step'] = 0                    if self._factored(p.shape):                        state['exp_avg_sq_row'] = torch.zeros(p.shape[:-1]).to(p)                        state['exp_avg_sq_col'] = torch.zeros(p.shape[:-2] + p.shape[-1:]).to(p)                    else:                        state['exp_avg_sq'] = torch.zeros_like(p)                    if group['beta1'] > 0.0:                        state['exp_avg'] = torch.zeros_like(p)                                state['step'] += 1                lr = self._get_lr(group, state)                 # Apply weight decay                if group['weight_decay'] != 0:                    grad = grad.add(p, alpha=group['weight_decay'])                                # Compute update                if self._factored(p.shape):                    # Factored second moment estimator for matrix parameters                    exp_avg_sq_row = state['exp_avg_sq_row']                    exp_avg_sq_col = state['exp_avg_sq_col']                                        exp_avg_sq_row, exp_avg_sq_col = self._compute_factored_second_moment(                        exp_avg_sq_row, exp_avg_sq_col, grad                    )                                        # Compute RMS using factored 2nd moment                    rms = torch.rsqrt(                        torch.matmul(exp_avg_sq_row.unsqueeze(-1), exp_avg_sq_col.unsqueeze(-2))                    ).to(grad) + group['eps'][0]                                        update = grad * rms                else:                    # Scalar parameters and vectors use simpler update                    exp_avg_sq = state['exp_avg_sq']                    beta2 = 1.0 - math.pow(state['step'], group['decay_rate'])                    exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)                    update = grad * torch.rsqrt(exp_avg_sq + group['eps'][0])                                # First moment estimate (momentum)                if group['beta1'] > 0.0:                    exp_avg = state['exp_avg']                    exp_avg.mul_(group['beta1']).add_(update, alpha=1 - group['beta1'])                    update = exp_avg                                # Apply update                p.data.add_(update, alpha=-lr)                        return loss # Example: 8-bit Adam (simplified version)class Adam8bit(Optimizer):    """    Implements Adam with 8-bit quantized optimizer states    Memory savings: ~75% compared to standard Adam    """    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8):        defaults = dict(lr=lr, betas=betas, eps=eps)        super(Adam8bit, self).__init__(params, defaults)            def _quantize_to_8bit(self, x):        """Quantize a tensor to 8-bit precision"""        # Compute scale factors per tensor        max_val = torch.max(torch.abs(x)).item()        scale = 127.0 / (max_val + 1e-8)  # Use 127 for int8 range (-127 to 127)                # Quantize by scaling and rounding        x_quant = torch.round(x * scale).to(torch.int8)                return x_quant, scale            def _dequantize_to_float(self, x_quant, scale):        """Dequantize from 8-bit back to float"""        return x_quant.float() / scale        def step(self, closure=None):        """Performs a single optimization step"""        loss = None        if closure is not None:            loss = closure()         for group in self.param_groups:            for p in group['params']:                if p.grad is None:                    continue                grad = p.grad.data                                if grad.is_sparse:                    raise RuntimeError("Adam8bit does not support sparse gradients")                 state = self.state[p]                 # State initialization                if len(state) == 0:                    state['step'] = 0                    # Initialize 8-bit moments and scaling factors                    m_8bit, m_scale = self._quantize_to_8bit(torch.zeros_like(p.data))                    v_8bit, v_scale = self._quantize_to_8bit(torch.zeros_like(p.data))                                        state['m_8bit'] = m_8bit                    state['v_8bit'] = v_8bit                    state['m_scale'] = m_scale                    state['v_scale'] = v_scale                 # Get optimizer parameters                beta1, beta2 = group['betas']                                state['step'] += 1                                # Dequantize 8-bit states to compute updates                m = self._dequantize_to_float(state['m_8bit'], state['m_scale'])                v = self._dequantize_to_float(state['v_8bit'], state['v_scale'])                                # Standard Adam update                m = beta1 * m + (1 - beta1) * grad                v = beta2 * v + (1 - beta2) * (grad * grad)                                # Bias correction                m_hat = m / (1 - beta1 ** state['step'])                v_hat = v / (1 - beta2 ** state['step'])                                # Update parameter                p.data.addcdiv_(m_hat, torch.sqrt(v_hat) + group['eps'], value=-group['lr'])                                # Re-quantize the moments for storage                state['m_8bit'], state['m_scale'] = self._quantize_to_8bit(m)                state['v_8bit'], state['v_scale'] = self._quantize_to_8bit(v)                        return loss # Example usage of the optimizersdef train_with_efficient_optimizers():    # Define a simple model    model = torch.nn.Sequential(        torch.nn.Linear(1024, 1024),        torch.nn.ReLU(),        torch.nn.Linear(1024, 1024),    )        # Total parameters: ~2M    total_params = sum(p.numel() for p in model.parameters())    print(f"Model has {total_params:,} parameters")        # Memory usage comparison    adam_memory = total_params * 3 * 4  # 3x params (weights + two moments), 4 bytes per float32    adafactor_memory = total_params * 4 + 2 * (1024 + 1024)  # Factored representation for matrices    adam8bit_memory = total_params * 4 + 2 * total_params  # 4 bytes for weights, 1 byte each for moments        print(f"Standard Adam memory: {adam_memory/1024/1024:.2f} MB")    print(f"Adafactor memory: {adafactor_memory/1024/1024:.2f} MB")    print(f"8-bit Adam memory: {adam8bit_memory/1024/1024:.2f} MB")        # Create dataset and train    x = torch.randn(100, 1024)    y = torch.randn(100, 1024)        # Choose optimizer    # optimizer = torch.optim.Adam(model.parameters(), lr=0.001)    # optimizer = Adafactor(model.parameters(), lr=0.001)    optimizer = Adam8bit(model.parameters(), lr=0.001)        # Simple training loop    loss_fn = torch.nn.MSELoss()    for epoch in range(3):        optimizer.zero_grad()        output = model(x)        loss = loss_fn(output, y)        loss.backward()        optimizer.step()        print(f"Epoch {epoch}, Loss: {loss.item():.4f}") # Usageif __name__ == "__main__":    train_with_efficient_optimizers() 

# Smart Scheduling and Early Stopping Implementationimport numpy as npimport torchimport torch.nn as nnimport torch.optim as optimimport matplotlib.pyplot as pltfrom torch.utils.data import DataLoader, Subsetfrom sklearn.model_selection import train_test_splitfrom collections import deque class EarlyStopping:    """Early stopping to terminate training when validation loss doesn't improve."""        def __init__(self, patience=5, min_delta=0.0, restore_best_weights=True):        """        Args:            patience (int): How many epochs to wait after last improvement            min_delta (float): Minimum change to qualify as an improvement            restore_best_weights (bool): Whether to restore model weights from the best epoch        """        self.patience = patience        self.min_delta = min_delta        self.restore_best_weights = restore_best_weights        self.best_score = None        self.best_weights = None        self.counter = 0        self.early_stop = False        def __call__(self, val_loss, model):        score = -val_loss  # Higher score is better (less loss)                if self.best_score is None:            self.best_score = score            self.save_checkpoint(model)        elif score < self.best_score + self.min_delta:            self.counter += 1            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')            if self.counter >= self.patience:                self.early_stop = True        else:            self.best_score = score            self.save_checkpoint(model)            self.counter = 0                def save_checkpoint(self, model):        """Save model weights when validation loss decreases."""        if self.restore_best_weights:            self.best_weights = {k: v.cpu().clone() for k, v in model.state_dict().items()}                def restore_checkpoint(self, model):        """Restore model weights to the best observed so far."""        if self.restore_best_weights and self.best_weights is not None:            model.load_state_dict(self.best_weights) class LearningRateScheduler:    """Custom learning rate scheduler with warmup and cosine decay."""        def __init__(self, optimizer, warmup_epochs=5, max_epochs=100,                  min_lr=1e-6, max_lr=1e-3, decay_type='cosine'):        self.optimizer = optimizer        self.warmup_epochs = warmup_epochs        self.max_epochs = max_epochs        self.min_lr = min_lr        self.max_lr = max_lr        self.decay_type = decay_type        self.current_epoch = 0            def step(self):        """Update the learning rate based on the current epoch."""        self.current_epoch += 1        lr = self.calculate_lr()        for param_group in self.optimizer.param_groups:            param_group['lr'] = lr        return lr        def calculate_lr(self):        """Calculate the learning rate based on schedule type."""        if self.current_epoch < self.warmup_epochs:            # Linear warmup            return self.min_lr + (self.max_lr - self.min_lr) * (self.current_epoch / self.warmup_epochs)        else:            # Apply decay after warmup            if self.decay_type == 'cosine':                # Cosine annealing                progress = (self.current_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)                return self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + np.cos(progress * np.pi))            elif self.decay_type == 'linear':                # Linear decay                progress = (self.current_epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)                return self.max_lr - (self.max_lr - self.min_lr) * progress            elif self.decay_type == 'step':                # Step decay                decay_rate = 0.1                step_size = (self.max_epochs - self.warmup_epochs) // 3                factor = decay_rate ** ((self.current_epoch - self.warmup_epochs) // step_size)                return self.max_lr * factor            else:                return self.min_lr class CurriculumSampler:    """Sample data in a curriculum-based manner, from easy to hard examples."""        def __init__(self, dataset, difficulty_scores, num_bins=5, schedule='linear'):        """        Args:            dataset: The dataset to sample from            difficulty_scores: List of scores measuring the difficulty of each example            num_bins: Number of difficulty levels to create            schedule: Type of curriculum schedule ('linear', 'exponential', or 'step')        """        self.dataset = dataset        self.num_bins = num_bins        self.schedule = schedule                # Sort examples by difficulty and divide into bins        sorted_indices = np.argsort(difficulty_scores)        self.bins = []        bin_size = len(sorted_indices) // num_bins                for i in range(num_bins):            start_idx = i * bin_size            end_idx = (i + 1) * bin_size if i < num_bins - 1 else len(sorted_indices)            self.bins.append(sorted_indices[start_idx:end_idx])        def get_sampler_for_epoch(self, epoch, max_epochs):        """Return a sampler for the given epoch that follows the curriculum."""        # Calculate how far through the curriculum we are (0 to 1)        progress = epoch / max_epochs                if self.schedule == 'exponential':            # Exponential schedule focuses more on easier examples early            curriculum_position = 1 - np.exp(-5 * progress)        elif self.schedule == 'step':            # Step schedule increases difficulty in discrete jumps            curriculum_position = min(int(progress * self.num_bins), self.num_bins - 1) / (self.num_bins - 1)        else:            # Linear schedule increases difficulty uniformly            curriculum_position = progress                    # Determine which bins to include based on current position        active_bin_count = max(1, int(np.ceil(curriculum_position * self.num_bins)))        indices = []        for i in range(active_bin_count):            indices.extend(self.bins[i])                # Create a subset dataset with these indices        return Subset(self.dataset, indices) def train_with_smart_scheduling(model, train_dataset, val_dataset,                                 batch_size=32, max_epochs=100,                                 difficulty_fn=None, patience=10,                                 use_curriculum=True, lr_schedule='cosine'):    """Train a model with smart scheduling and early stopping.        Args:        model: PyTorch model to train        train_dataset: Training dataset        val_dataset: Validation dataset        batch_size: Batch size for training        max_epochs: Maximum number of epochs        difficulty_fn: Function to calculate difficulty of each example        patience: Early stopping patience        use_curriculum: Whether to use curriculum learning        lr_schedule: Learning rate schedule type    """    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')    model = model.to(device)        # Define optimizer    optimizer = optim.AdamW(model.parameters(), lr=1e-5, weight_decay=0.01)        # Set up learning rate scheduler    scheduler = LearningRateScheduler(        optimizer, warmup_epochs=5, max_epochs=max_epochs,        min_lr=1e-6, max_lr=1e-3, decay_type=lr_schedule    )        # Set up early stopping    early_stopping = EarlyStopping(patience=patience, min_delta=1e-4)        # Set up curriculum learning if requested    curriculum_sampler = None    if use_curriculum and difficulty_fn is not None:        # Calculate difficulty scores for each example        difficulty_scores = [difficulty_fn(x) for x in train_dataset]        curriculum_sampler = CurriculumSampler(train_dataset, difficulty_scores)        # Training history    history = {        'train_loss': [],        'val_loss': [],        'learning_rates': []    }        # Training loop    for epoch in range(max_epochs):        # Update learning rate        current_lr = scheduler.step()        history['learning_rates'].append(current_lr)                # Get data loader based on curriculum for this epoch        if curriculum_sampler and use_curriculum:            epoch_dataset = curriculum_sampler.get_sampler_for_epoch(epoch, max_epochs)            train_loader = DataLoader(epoch_dataset, batch_size=batch_size, shuffle=True)        else:            train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)                    val_loader = DataLoader(val_dataset, batch_size=batch_size)                # Training phase        model.train()        train_loss = 0.0        for inputs, targets in train_loader:            inputs, targets = inputs.to(device), targets.to(device)                        optimizer.zero_grad()            outputs = model(inputs)            loss = nn.CrossEntropyLoss()(outputs, targets)            loss.backward()            optimizer.step()                        train_loss += loss.item()                    train_loss /= len(train_loader)        history['train_loss'].append(train_loss)                # Validation phase        model.eval()        val_loss = 0.0        with torch.no_grad():            for inputs, targets in val_loader:                inputs, targets = inputs.to(device), targets.to(device)                outputs = model(inputs)                loss = nn.CrossEntropyLoss()(outputs, targets)                val_loss += loss.item()                        val_loss /= len(val_loader)        history['val_loss'].append(val_loss)                print(f'Epoch {epoch+1}/{max_epochs}, LR: {current_lr:.6f}, '              f'Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}')                # Check early stopping        early_stopping(val_loss, model)        if early_stopping.early_stop:            print(f"Early stopping triggered at epoch {epoch+1}")            break        # Restore best model weights    early_stopping.restore_checkpoint(model)        # Plot training history    plt.figure(figsize=(12, 4))        plt.subplot(1, 2, 1)    plt.plot(history['train_loss'], label='Train Loss')    plt.plot(history['val_loss'], label='Validation Loss')    plt.xlabel('Epoch')    plt.ylabel('Loss')    plt.legend()    plt.title('Training and Validation Loss')        plt.subplot(1, 2, 2)    plt.plot(history['learning_rates'])    plt.xlabel('Epoch')    plt.ylabel('Learning Rate')    plt.title('Learning Rate Schedule')        plt.tight_layout()    plt.show()        return model, history # Example difficulty function - sequence length as difficultydef sequence_length_difficulty(example):    """Return the length of a sequence as a measure of difficulty."""    # Replace with actual logic to extract sequence from your data format    sequence = example[0]  # Assuming example is a tuple (input, target)    return len(sequence) # Example usageif __name__ == "__main__":    # Define a simple model    model = nn.Sequential(        nn.Linear(768, 512),        nn.ReLU(),        nn.Dropout(0.1),        nn.Linear(512, 256),        nn.ReLU(),        nn.Linear(256, 10)    )        # Create dummy datasets (replace with your actual data)    X = torch.randn(1000, 768)    y = torch.randint(0, 10, (1000,))    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2)        class DummyDataset(torch.utils.data.Dataset):        def __init__(self, X, y):            self.X = X            self.y = y                def __len__(self):            return len(self.X)                def __getitem__(self, idx):            return self.X[idx], self.y[idx]        train_dataset = DummyDataset(X_train, y_train)    val_dataset = DummyDataset(X_val, y_val)        # Train with smart scheduling    trained_model, history = train_with_smart_scheduling(        model,         train_dataset,         val_dataset,        batch_size=32,        max_epochs=50,        difficulty_fn=sequence_length_difficulty,        patience=7,        use_curriculum=True,        lr_schedule='cosine'    ) 

# Comprehensive energy and carbon footprint estimation for LLM trainingimport pandas as pdimport matplotlib.pyplot as pltfrom datetime import datetime, timedelta class CarbonTracker:    """Track carbon emissions from AI training runs"""        # Energy mix data by region (approximate values)    CARBON_INTENSITY = {        "us-east": 0.38,        # US East Coast        "us-west": 0.22,        # US West Coast (more renewables)        "europe": 0.23,         # European average        "asia-pacific": 0.55,   # Asia Pacific region        "global-average": 0.47  # Global average    }        def __init__(self,                  gpu_model="A100",                  num_gpus=8,                  region="us-east",                  pue=1.1):        """        Initialize a carbon tracker                Args:            gpu_model: GPU model being used (affects power draw)            num_gpus: Number of GPUs in the training cluster            region: Geographic region (affects carbon intensity)            pue: Power Usage Effectiveness of data center (1.1 is excellent, 2.0 is poor)        """        self.gpu_power = self._get_gpu_power(gpu_model)        self.num_gpus = num_gpus        self.region = region        self.carbon_factor = self.CARBON_INTENSITY.get(region, self.CARBON_INTENSITY["global-average"])        self.pue = pue  # Data center efficiency factor                # For tracking        self.start_time = None        self.measurements = []        def _get_gpu_power(self, gpu_model):        """Return typical power draw in watts for common GPU models"""        power_draw = {            "A100": 400,            "H100": 700,            "A6000": 300,            "V100": 300,            "A40": 300,            "A10": 150,        }        return power_draw.get(gpu_model, 400)  # Default to A100 if unknown        def start_tracking(self):        """Start the tracking session"""        self.start_time = datetime.now()        self.measurements = []        print(f"Started carbon tracking at {self.start_time}")        def log_utilization(self, gpu_utilization=1.0):        """Log current GPU utilization (between 0.0-1.0)"""        if self.start_time is None:            raise ValueError("Must call start_tracking first")                    duration = (datetime.now() - self.start_time).total_seconds() / 3600  # hours        self.measurements.append({            "timestamp": datetime.now(),            "duration_hrs": duration,            "utilization": gpu_utilization        })        def estimate_carbon_footprint(self, additional_hours=0, avg_utilization=0.85):        """        Calculate energy usage and carbon emissions                Args:            additional_hours: Future hours to include in projection            avg_utilization: Average GPU utilization for future projection        """        # Calculate duration based on tracking or fixed input        if self.start_time and self.measurements:            # Calculate average utilization from measurements            if len(self.measurements) > 0:                measured_utilization = sum(m["utilization"] for m in self.measurements) / len(self.measurements)            else:                measured_utilization = avg_utilization                            # Measured duration plus projected additional time            total_hours = self.measurements[-1]["duration_hrs"] + additional_hours            avg_util = (measured_utilization * self.measurements[-1]["duration_hrs"] +                        avg_utilization * additional_hours) / total_hours        else:            # If no tracking, just use the provided values            total_hours = additional_hours            avg_util = avg_utilization                # Calculate energy in kWh, accounting for data center PUE        energy_kwh = (self.gpu_power * self.num_gpus * total_hours * avg_util * self.pue) / 1000                # Calculate CO2 emissions in kg        co2_emission = energy_kwh * self.carbon_factor                results = {            "gpu_model": self._get_gpu_model_name(),            "num_gpus": self.num_gpus,            "region": self.region,            "duration_hours": total_hours,            "avg_utilization": avg_util,            "pue": self.pue,            "energy_kwh": energy_kwh,            "carbon_factor": self.carbon_factor,            "co2_emission_kg": co2_emission,            "co2_emission_tons": co2_emission / 1000,            "equivalents": self._get_carbon_equivalents(co2_emission)        }                return results        def _get_gpu_model_name(self):        # Reverse lookup to get model name from power        for model, power in {            "A100": 400,            "H100": 700,            "A6000": 300,            "V100": 300,        }.items():            if power == self.gpu_power:                return model        return "Custom GPU"        def _get_carbon_equivalents(self, co2_kg):        """Convert CO2 emissions to everyday equivalents"""        return {            "flights_ny_to_sf": co2_kg / 1100,  # One-way flight (~1100kg)            "miles_driven": co2_kg / 0.404,     # ~0.404 kg CO2 per mile            "smartphone_charges": co2_kg / 0.005,  # ~5g per full charge            "trees_year_offset": co2_kg / 21,   # One tree absorbs ~21kg/year            "homes_day_energy": co2_kg / 38     # Average US home ~38kg/day        }        def visualize_impact(self, results):        """Create visualizations of the carbon impact"""        # Create figure with two subplots        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))                # Plot 1: Energy and Emissions        data = [results["energy_kwh"], results["co2_emission_kg"]]        labels = ["Energy (kWh)", "CO₂ Emissions (kg)"]        ax1.bar(labels, data, color=["#3498db", "#e74c3c"])        ax1.set_title("Energy Usage and Carbon Emissions")        for i, v in enumerate(data):            ax1.text(i, v + 5, f"{v:.1f}", ha='center')                # Plot 2: Carbon Equivalents        eq = results["equivalents"]        labels = ["Flights\nNY to SF", "Miles\nDriven", "Trees to\nOffset (year)"]        data = [eq["flights_ny_to_sf"], eq["miles_driven"]/1000, eq["trees_year_offset"]]                ax2.bar(labels, data, color=["#2ecc71", "#9b59b6", "#f39c12"])        ax2.set_title("Carbon Emission Equivalents")        for i, v in enumerate(data):            ax2.text(i, v + 0.05*max(data), f"{v:.1f}", ha='center')                plt.tight_layout()        return fig # Example usageif __name__ == "__main__":    # Initialize tracker    tracker = CarbonTracker(        gpu_model="A100",        num_gpus=8,        region="us-east",        pue=1.1  # 1.1 is excellent, industry average is ~1.6    )        # Estimate for a 24-hour training run    results = tracker.estimate_carbon_footprint(additional_hours=24, avg_utilization=0.85)        # Print results    print(f"\nTraining Configuration:")    print(f"- {results['num_gpus']} {results['gpu_model']} GPUs in {results['region']}")    print(f"- {results['duration_hours']:.1f} hours at {results['avg_utilization']*100:.0f}% utilization")    print(f"- Data center PUE: {results['pue']}")        print(f"\nEnvironmental Impact:")    print(f"- Energy used: {results['energy_kwh']:.1f} kWh")    print(f"- CO₂ emitted: {results['co2_emission_kg']:.2f} kg ({results['co2_emission_tons']:.3f} tons)")        print(f"\nThis is equivalent to:")    eq = results["equivalents"]    print(f"- {eq['flights_ny_to_sf']:.2f} one-way flights from NY to SF")    print(f"- {eq['miles_driven']:.0f} miles driven by an average car")    print(f"- {eq['smartphone_charges']:.0f} smartphone charges")    print(f"- {eq['trees_year_offset']:.1f} trees needed for a year to offset")    print(f"- {eq['homes_day_energy']:.1f} days of energy for an average US home")        # Visualize (uncomment to display)    # fig = tracker.visualize_impact(results)    # plt.show()