Add amp_flow_training_single_gpu_full_data.py

src/amp_flow_training_single_gpu_full_data.py

@@ -0,0 +1,561 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR
+import numpy as np
+from tqdm import tqdm
+import json
+import os
+import argparse
+import time
+from torch.cuda.amp import autocast, GradScaler
+import wandb  # For logging (optional)
+
+# Import your existing components
+from compressor_with_embeddings import Compressor, Decompressor, PrecomputedEmbeddingDataset
+from final_flow_model import AMPFlowMatcherCFGConcat, SinusoidalTimeEmbedding
+from cfg_dataset import CFGFlowDataset, create_cfg_dataloader
+
+# ---------------- Optimized Configuration for H100 ----------------
+ESM_DIM = 1280  # ESM-2 hidden dim (esm2_t33_650M_UR50D)
+COMP_RATIO = 16  # compression factor
+COMP_DIM = ESM_DIM // COMP_RATIO
+MAX_SEQ_LEN = 50  # Actual sequence length from final_sequence_encoder.py
+
+# OPTIMIZED H100 hyperparameters - HIGH THROUGHPUT + STABLE TRAINING
+BATCH_SIZE = 512  # PUSH H100 TO LIMITS - using ~70GB memory 
+EPOCHS = 2000  # Slightly more epochs with safer LR for same 5-6 hour target
+BASE_LR = 8e-4  # SAFE but effective LR - 2x original, not 4x
+LR_MIN = 4e-4  # Conservative minimum learning rate
+WARMUP_STEPS = 4000  # Gentler warmup to avoid explosion
+GPU_ID = 0  # Use GPU 0
+
+# Training optimizations
+USE_MIXED_PRECISION = True  # BF16 for H100
+GRADIENT_CLIP_NORM = 0.5  # TIGHTER gradient clipping for flow matching stability
+WEIGHT_DECAY = 0.01  # Weight decay for regularization
+VALIDATION_INTERVAL = 5000  # Validate every 5K steps (more frequent)
+CHECKPOINT_INTERVAL = 300  # Save checkpoint every 300 epochs (more frequent)
+NUM_WORKERS = 32  # MAXIMIZED data loading workers for H100
+
+# CFG training parameters
+CFG_DROPOUT_RATE = 0.15  # 15% of batches as unconditional for CFG
+
+class AMPFlowTrainerSingleGPUFullData:
+    """
+    Optimized Single GPU training pipeline for AMP generation using ProtFlow methodology.
+    Uses ALL available data with H100-optimized settings for overnight training.
+    """
+    
+    def __init__(self, embeddings_path, cfg_data_path, use_wandb=False):
+        self.device = torch.device(f'cuda:{GPU_ID}')
+        self.embeddings_path = embeddings_path
+        self.cfg_data_path = cfg_data_path
+        self.use_wandb = use_wandb
+        
+        # Enable H100 optimizations
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+        
+        print(f"Using GPU {GPU_ID} for optimized H100 training")
+        print(f"Mixed precision: {USE_MIXED_PRECISION}")
+        print(f"Batch size: {BATCH_SIZE}")
+        print(f"Target epochs: {EPOCHS}")
+        print(f"Learning rate: {BASE_LR} -> {LR_MIN}")
+        
+        # Initialize mixed precision training
+        if USE_MIXED_PRECISION:
+            self.scaler = GradScaler()
+            print("✓ Mixed precision training enabled (BF16)")
+        
+        # Initialize wandb if requested
+        if self.use_wandb:
+            wandb.init(
+                project="amp-flow-training",
+                config={
+                    "batch_size": BATCH_SIZE,
+                    "epochs": EPOCHS,
+                    "base_lr": BASE_LR,
+                    "lr_min": LR_MIN,
+                    "warmup_steps": WARMUP_STEPS,
+                    "mixed_precision": USE_MIXED_PRECISION,
+                    "gradient_clip": GRADIENT_CLIP_NORM,
+                    "weight_decay": WEIGHT_DECAY
+                }
+            )
+        
+        print(f"Loading ALL AMP embeddings from {embeddings_path}...")
+        
+        # Load ALL embeddings (use the combined file instead of individual files)
+        self._load_all_embeddings()
+        
+        # Compute normalization statistics
+        print("Computing preprocessing statistics...")
+        self._compute_preprocessing_stats()
+        
+        # Initialize models
+        self._initialize_models()
+        
+        # Initialize datasets and dataloaders
+        self._initialize_data()
+        
+        # Initialize optimizer and scheduler
+        self._initialize_optimizer()
+        
+        print("✓ Optimized Single GPU training setup complete with FULL DATA!")
+    
+    def _load_all_embeddings(self):
+        """Load ALL peptide embeddings from the combined file."""
+        # Try to load the combined embeddings file first
+        combined_path = os.path.join(self.embeddings_path, "all_peptide_embeddings.pt")
+        
+        if os.path.exists(combined_path):
+            print(f"Loading combined embeddings from {combined_path}...")
+            self.embeddings = torch.load(combined_path, map_location=self.device)
+            print(f"✓ Loaded ALL embeddings: {self.embeddings.shape}")
+        else:
+            print("Combined embeddings file not found, loading individual files...")
+            # Fallback to individual files
+            import glob
+            
+            embedding_files = glob.glob(os.path.join(self.embeddings_path, "*.pt"))
+            embedding_files = [f for f in embedding_files if not f.endswith('metadata.json') and not f.endswith('sequence_ids.json') and not f.endswith('all_peptide_embeddings.pt')]
+            
+            print(f"Found {len(embedding_files)} individual embedding files")
+            
+            # Load and stack all embeddings
+            embeddings_list = []
+            for file_path in embedding_files:
+                try:
+                    embedding = torch.load(file_path)
+                    if embedding.dim() == 2:  # (seq_len, hidden_dim)
+                        embeddings_list.append(embedding)
+                    else:
+                        print(f"Warning: Skipping {file_path} - unexpected shape {embedding.shape}")
+                except Exception as e:
+                    print(f"Warning: Could not load {file_path}: {e}")
+            
+            if not embeddings_list:
+                raise ValueError("No valid embeddings found!")
+            
+            self.embeddings = torch.stack(embeddings_list)
+            print(f"Loaded {len(self.embeddings)} embeddings from individual files")
+    
+    def _compute_preprocessing_stats(self):
+        """Compute normalization statistics for embeddings."""
+        # Flatten all embeddings
+        flat_embeddings = self.embeddings.reshape(-1, ESM_DIM)
+        
+        # Compute statistics
+        mean = flat_embeddings.mean(dim=0)
+        std = flat_embeddings.std(dim=0)
+        min_val = flat_embeddings.min()
+        max_val = flat_embeddings.max()
+        
+        self.stats = {
+            'mean': mean,
+            'std': std,
+            'min': min_val,
+            'max': max_val
+        }
+        
+        # Save statistics
+        torch.save(self.stats, 'normalization_stats.pt')
+        print(f"✓ Statistics computed and saved:")
+        print(f"  Total embeddings: {len(self.embeddings):,}")
+        print(f"  Mean: {mean.mean():.4f} ± {mean.std():.4f}")
+        print(f"  Std: {std.mean():.4f} ± {std.std():.4f}")
+        print(f"  Range: [{min_val:.4f}, {max_val:.4f}]")
+    
+    def _initialize_models(self):
+        """Initialize compressor, decompressor, and flow model."""
+        print("Initializing models...")
+        
+        # Load pre-trained compressor and decompressor
+        self.compressor = Compressor().to(self.device)
+        self.decompressor = Decompressor().to(self.device)
+        
+        self.compressor.load_state_dict(torch.load('final_compressor_model.pth', map_location=self.device))
+        self.decompressor.load_state_dict(torch.load('final_decompressor_model.pth', map_location=self.device))
+        
+        # Initialize flow model with CFG
+        self.flow_model = AMPFlowMatcherCFGConcat(
+            hidden_dim=480,
+            compressed_dim=COMP_DIM,
+            n_layers=12,
+            n_heads=16,
+            dim_ff=3072,
+            max_seq_len=25,  # MAX_SEQ_LEN // 2 due to pooling
+            use_cfg=True
+        ).to(self.device)
+        
+        # Compile model for PyTorch 2.x speedup (if available)
+        try:
+            self.flow_model = torch.compile(self.flow_model, mode="reduce-overhead")
+            print("✓ Model compiled with torch.compile for speedup")
+        except Exception as e:
+            print(f"⚠️  Model compilation failed: {e}")
+        
+        # Set models to training mode
+        self.compressor.train()
+        self.decompressor.train()
+        self.flow_model.train()
+        
+        print(f"✓ Models initialized:")
+        print(f"  Compressor parameters: {sum(p.numel() for p in self.compressor.parameters()):,}")
+        print(f"  Decompressor parameters: {sum(p.numel() for p in self.decompressor.parameters()):,}")
+        print(f"  Flow model parameters: {sum(p.numel() for p in self.flow_model.parameters()):,}")
+    
+    def _initialize_data(self):
+        """Initialize datasets and dataloaders with FULL data."""
+        print("Initializing datasets with FULL data...")
+        
+        # Create CFG dataset with FULL UniProt data
+        self.cfg_dataset = CFGFlowDataset(
+            embeddings_path=self.embeddings_path,
+            cfg_data_path=self.cfg_data_path,
+            use_masked_labels=True,
+            max_seq_len=MAX_SEQ_LEN,
+            device=self.device
+        )
+        
+        # Create dataloader with optimized settings
+        self.dataloader = create_cfg_dataloader(
+            self.cfg_dataset,
+            batch_size=BATCH_SIZE,
+            shuffle=True,
+            num_workers=NUM_WORKERS
+        )
+        
+        # Calculate total steps and validation intervals
+        self.total_steps = len(self.dataloader) * EPOCHS
+        self.validation_steps = VALIDATION_INTERVAL
+        
+        print(f"✓ Dataset initialized with FULL data:")
+        print(f"  Total samples: {len(self.cfg_dataset):,}")
+        print(f"  Batch size: {BATCH_SIZE}")
+        print(f"  Batches per epoch: {len(self.dataloader):,}")
+        print(f"  Total training steps: {self.total_steps:,}")
+        print(f"  Validation every: {self.validation_steps:,} steps")
+    
+    def _initialize_optimizer(self):
+        """Initialize optimizer and learning rate scheduler."""
+        print("Initializing optimizer and scheduler...")
+        
+        # Optimizer for flow model only (compressor/decompressor are frozen)
+        self.optimizer = optim.AdamW(
+            self.flow_model.parameters(),
+            lr=BASE_LR,
+            weight_decay=WEIGHT_DECAY,
+            betas=(0.9, 0.98),  # Optimized betas for flow matching
+            eps=1e-6  # Lower epsilon for numerical stability
+        )
+        
+        # Learning rate scheduler with proper warmup and cosine annealing
+        warmup_scheduler = LinearLR(
+            self.optimizer, 
+            start_factor=0.1, 
+            end_factor=1.0, 
+            total_iters=WARMUP_STEPS
+        )
+        
+        main_scheduler = CosineAnnealingLR(
+            self.optimizer,
+            T_max=self.total_steps - WARMUP_STEPS,
+            eta_min=LR_MIN
+        )
+        
+        self.scheduler = SequentialLR(
+            self.optimizer,
+            schedulers=[warmup_scheduler, main_scheduler],
+            milestones=[WARMUP_STEPS]
+        )
+        
+        print(f"✓ Optimizer initialized:")
+        print(f"  Base LR: {BASE_LR}")
+        print(f"  Min LR: {LR_MIN}")
+        print(f"  Warmup steps: {WARMUP_STEPS}")
+        print(f"  Weight decay: {WEIGHT_DECAY}")
+        print(f"  Gradient clip norm: {GRADIENT_CLIP_NORM}")
+    
+    def _preprocess_batch(self, batch):
+        """Preprocess a batch of data for training."""
+        # Extract data
+        embeddings = batch['embeddings'].to(self.device)  # (B, L, ESM_DIM)
+        labels = batch['labels'].to(self.device)  # (B,)
+        
+        # Normalize embeddings
+        m, s = self.stats['mean'].to(self.device), self.stats['std'].to(self.device)
+        mn, mx = self.stats['min'].to(self.device), self.stats['max'].to(self.device)
+        
+        embeddings = (embeddings - m) / (s + 1e-8)
+        embeddings = (embeddings - mn) / (mx - mn + 1e-8)
+        
+        # Compress embeddings
+        with torch.no_grad():
+            compressed = self.compressor(embeddings)  # (B, L, COMP_DIM)
+        
+        return compressed, labels
+    
+    def _compute_validation_metrics(self):
+        """Compute validation metrics on a subset of data."""
+        self.flow_model.eval()
+        val_losses = []
+        
+        # Use a subset of data for validation
+        val_samples = min(1000, len(self.cfg_dataset))
+        val_indices = torch.randperm(len(self.cfg_dataset))[:val_samples]
+        
+        with torch.no_grad():
+            for i in range(0, val_samples, BATCH_SIZE):
+                batch_indices = val_indices[i:i+BATCH_SIZE]
+                batch_data = [self.cfg_dataset[idx] for idx in batch_indices]
+                
+                # Collate batch
+                embeddings = torch.stack([item['embedding'] for item in batch_data])
+                labels = torch.stack([item['label'] for item in batch_data])
+                
+                # Preprocess
+                compressed, labels = self._preprocess_batch({
+                    'embeddings': embeddings,
+                    'labels': labels
+                })
+                
+                B, L, D = compressed.shape
+                
+                # Sample random time
+                t = torch.rand(B, device=self.device)
+                
+                # Sample random noise
+                eps = torch.randn_like(compressed)
+                
+                # Compute target
+                xt = (1 - t.unsqueeze(-1).unsqueeze(-1)) * compressed + t.unsqueeze(-1).unsqueeze(-1) * eps
+                
+                # Predict vector field
+                vt_pred = self.flow_model(xt, t, labels=labels)
+                
+                # Target vector field
+                vt_target = eps - compressed
+                
+                # Compute loss
+                loss = F.mse_loss(vt_pred, vt_target)
+                val_losses.append(loss.item())
+        
+        self.flow_model.train()
+        return np.mean(val_losses)
+    
+    def train_flow_matching(self):
+        """Train the flow matching model with FULL data and optimizations."""
+        print(f"🚀 Starting Optimized Single GPU Flow Matching Training with FULL DATA")
+        print(f"GPU: {GPU_ID}")
+        print(f"Total iterations: {EPOCHS}")
+        print(f"Batch size: {BATCH_SIZE}")
+        print(f"Total samples: {len(self.cfg_dataset):,}")
+        print(f"Mixed precision: {USE_MIXED_PRECISION}")
+        print(f"Estimated time: ~8-10 hours (overnight training with ALL data)")
+        print("=" * 60)
+        
+        # Training loop
+        best_loss = float('inf')
+        losses = []
+        val_losses = []
+        global_step = 0
+        start_time = time.time()
+        
+        for epoch in tqdm(range(EPOCHS), desc="Training Flow Model"):
+            epoch_losses = []
+            epoch_start_time = time.time()
+            
+            for batch_idx, batch in enumerate(self.dataloader):
+                # Preprocess batch
+                compressed, labels = self._preprocess_batch(batch)
+                B, L, D = compressed.shape
+                
+                # CFG training: randomly mask some labels for unconditional training
+                if torch.rand(1).item() < CFG_DROPOUT_RATE:
+                    labels = torch.full_like(labels, fill_value=-1)  # Unconditional
+                
+                # Sample random time
+                t = torch.rand(B, device=self.device)  # (B,)
+                
+                # Sample random noise
+                eps = torch.randn_like(compressed)  # (B, L, D)
+                
+                # Compute target: x_t = (1-t) * x_0 + t * eps
+                xt = (1 - t.unsqueeze(-1).unsqueeze(-1)) * compressed + t.unsqueeze(-1).unsqueeze(-1) * eps
+                
+                # Forward pass with mixed precision
+                if USE_MIXED_PRECISION:
+                    with autocast(dtype=torch.bfloat16):
+                        vt_pred = self.flow_model(xt, t, labels=labels)  # (B, L, D)
+                        vt_target = eps - compressed  # (B, L, D)
+                        loss = F.mse_loss(vt_pred, vt_target)
+                    
+                    # Backward pass with gradient scaling
+                    self.optimizer.zero_grad()
+                    self.scaler.scale(loss).backward()
+                    
+                    # Gradient clipping
+                    self.scaler.unscale_(self.optimizer)
+                    torch.nn.utils.clip_grad_norm_(self.flow_model.parameters(), max_norm=GRADIENT_CLIP_NORM)
+                    
+                    self.scaler.step(self.optimizer)
+                    self.scaler.update()
+                else:
+                    # Standard training
+                    vt_pred = self.flow_model(xt, t, labels=labels)  # (B, L, D)
+                    vt_target = eps - compressed  # (B, L, D)
+                    loss = F.mse_loss(vt_pred, vt_target)
+                    
+                    # Backward pass
+                    self.optimizer.zero_grad()
+                    loss.backward()
+                    
+                    # Gradient clipping
+                    torch.nn.utils.clip_grad_norm_(self.flow_model.parameters(), max_norm=GRADIENT_CLIP_NORM)
+                    
+                    self.optimizer.step()
+                
+                # Update learning rate
+                self.scheduler.step()
+                
+                epoch_losses.append(loss.item())
+                global_step += 1
+                
+                # Logging
+                if batch_idx % 100 == 0:
+                    current_lr = self.scheduler.get_last_lr()[0]
+                    elapsed_time = time.time() - start_time
+                    steps_per_sec = global_step / elapsed_time
+                    eta_hours = (self.total_steps - global_step) / steps_per_sec / 3600
+                    
+                    print(f"Epoch {epoch:4d} | Step {global_step:6d}/{self.total_steps:6d} | "
+                          f"Loss: {loss.item():.6f} | LR: {current_lr:.2e} | "
+                          f"Speed: {steps_per_sec:.1f} steps/s | ETA: {eta_hours:.1f}h")
+                    
+                    # Log to wandb
+                    if self.use_wandb:
+                        wandb.log({
+                            'train/loss': loss.item(),
+                            'train/learning_rate': current_lr,
+                            'train/steps_per_sec': steps_per_sec,
+                            'train/global_step': global_step
+                        })
+                
+                # Validation
+                if global_step % self.validation_steps == 0:
+                    val_loss = self._compute_validation_metrics()
+                    val_losses.append(val_loss)
+                    
+                    print(f"Validation at step {global_step}: Loss = {val_loss:.6f}")
+                    
+                    if self.use_wandb:
+                        wandb.log({
+                            'val/loss': val_loss,
+                            'val/global_step': global_step
+                        })
+                    
+                    # Early stopping check
+                    if val_loss < best_loss:
+                        best_loss = val_loss
+                        self._save_checkpoint(epoch, val_loss, global_step, is_final=False, is_best=True)
+            
+            # Compute epoch statistics
+            avg_loss = np.mean(epoch_losses)
+            losses.append(avg_loss)
+            epoch_time = time.time() - epoch_start_time
+            
+            print(f"Epoch {epoch:4d} | Avg Loss: {avg_loss:.6f} | "
+                  f"LR: {self.scheduler.get_last_lr()[0]:.2e} | "
+                  f"Time: {epoch_time:.1f}s | Samples: {len(self.cfg_dataset):,}")
+            
+            # Save checkpoint
+            if (epoch + 1) % CHECKPOINT_INTERVAL == 0:
+                self._save_checkpoint(epoch, avg_loss, global_step, is_final=True)
+        
+        # Save final model
+        self._save_checkpoint(EPOCHS - 1, losses[-1], global_step, is_final=True)
+        
+        total_time = time.time() - start_time
+        print("=" * 60)
+        print("🎉 Optimized Training Complete with FULL DATA!")
+        print(f"Best validation loss: {best_loss:.6f}")
+        print(f"Total training time: {total_time/3600:.1f} hours")
+        print(f"Total samples used: {len(self.cfg_dataset):,}")
+        print(f"Final model saved as: amp_flow_model_final_optimized.pth")
+        
+        return losses, val_losses
+    
+    def _save_checkpoint(self, step, loss, global_step, is_final=False, is_best=False):
+        """Save model checkpoint."""
+        # Create output directory if it doesn't exist
+        output_dir = '/data2/edwardsun/flow_checkpoints'
+        os.makedirs(output_dir, exist_ok=True)
+        
+        if is_best:
+            filename = os.path.join(output_dir, 'amp_flow_model_best_optimized.pth')
+        elif is_final:
+            filename = os.path.join(output_dir, 'amp_flow_model_final_optimized.pth')
+        else:
+            filename = os.path.join(output_dir, f'amp_flow_checkpoint_optimized_step_{step:04d}.pth')
+        
+        checkpoint = {
+            'step': step,
+            'global_step': global_step,
+            'loss': loss,
+            'flow_model_state_dict': self.flow_model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scheduler_state_dict': self.scheduler.state_dict(),
+            'stats': self.stats,
+            'total_samples': len(self.cfg_dataset),
+            'config': {
+                'batch_size': BATCH_SIZE,
+                'epochs': EPOCHS,
+                'base_lr': BASE_LR,
+                'lr_min': LR_MIN,
+                'warmup_steps': WARMUP_STEPS,
+                'mixed_precision': USE_MIXED_PRECISION,
+                'gradient_clip': GRADIENT_CLIP_NORM,
+                'weight_decay': WEIGHT_DECAY
+            }
+        }
+        
+        torch.save(checkpoint, filename)
+        print(f"✓ Checkpoint saved: {filename} (loss: {loss:.6f}, step: {global_step})")
+
+def main():
+    """Main training function."""
+    global BATCH_SIZE, EPOCHS
+    
+    parser = argparse.ArgumentParser(description='Optimized Single GPU AMP Flow Training with FULL DATA')
+    parser.add_argument('--embeddings', default='/data2/edwardsun/flow_project/peptide_embeddings/', 
+                       help='Path to peptide embeddings directory')
+    parser.add_argument('--cfg_data', default='/data2/edwardsun/flow_project/test_uniprot_processed/uniprot_processed_data.json',
+                       help='Path to FULL CFG data file')
+    parser.add_argument('--use_wandb', action='store_true', help='Use wandb for logging')
+    parser.add_argument('--batch_size', type=int, default=BATCH_SIZE, help='Batch size for training')
+    parser.add_argument('--epochs', type=int, default=EPOCHS, help='Number of training epochs')
+    
+    args = parser.parse_args()
+    
+    # Update global variables if provided
+    if args.batch_size != BATCH_SIZE:
+        BATCH_SIZE = args.batch_size
+    if args.epochs != EPOCHS:
+        EPOCHS = args.epochs
+    
+    print(f"Starting optimized training with batch_size={BATCH_SIZE}, epochs={EPOCHS}")
+    
+    # Initialize trainer
+    trainer = AMPFlowTrainerSingleGPUFullData(args.embeddings, args.cfg_data, args.use_wandb)
+    
+    # Start training
+    losses, val_losses = trainer.train_flow_matching()
+    
+    print("Optimized training completed successfully with FULL DATA!")
+
+if __name__ == "__main__":
+    main()