import argparse import math import os from functools import partial from collections import Counter import torch import torch.nn as nn import torch.nn.functional as F from datasets import load_from_disk from torch.optim import AdamW from torch.optim.lr_scheduler import LambdaLR from torch.utils.data import DataLoader import pytorch_lightning as pl from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor from pytorch_lightning.loggers import WandbLogger from pytorch_lightning.strategies import DDPStrategy from rdkit import Chem from smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer from peptide_analyzer import PeptideAnalyzer import dataloading_for_dynamic_batching as dynamic_dataloader class RotaryPositionalEmbedding(nn.Module): def __init__(self, dim, max_position_embeddings=2048, base=10000): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq) def forward(self, x, seq_len=None): if seq_len is None: seq_len = x.shape[1] t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype) freqs = torch.einsum("i,j->ij", t, self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1) cos_emb = emb.cos()[None, :, :] sin_emb = emb.sin()[None, :, :] return cos_emb, sin_emb def rotate_half(x): x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb(q, k, cos, sin): q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed # --- Model Architecture with RoPE --- def modulate(x, shift, scale): return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) class TimestepEmbedder(nn.Module): def __init__(self, hidden_size): super().__init__() self.mlp = nn.Sequential( nn.Linear(1, hidden_size, bias=True), nn.SiLU(), nn.Linear(hidden_size, hidden_size, bias=True), ) def forward(self, t): return self.mlp(t.unsqueeze(-1)) class MultiHeadAttentionWithRoPE(nn.Module): def __init__(self, hidden_size, n_heads): super().__init__() self.hidden_size = hidden_size self.n_heads = n_heads self.head_dim = hidden_size // n_heads assert self.head_dim * n_heads == hidden_size, "hidden_size must be divisible by n_heads" self.q_proj = nn.Linear(hidden_size, hidden_size, bias=False) self.k_proj = nn.Linear(hidden_size, hidden_size, bias=False) self.v_proj = nn.Linear(hidden_size, hidden_size, bias=False) self.out_proj = nn.Linear(hidden_size, hidden_size) self.rope = RotaryPositionalEmbedding(self.head_dim) def forward(self, x): batch_size, seq_len, hidden_size = x.shape q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2) cos, sin = self.rope(q, seq_len) q, k = apply_rotary_pos_emb(q, k, cos, sin) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn_weights = F.softmax(scores, dim=-1) attn_output = torch.matmul(attn_weights, v) attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size) output = self.out_proj(attn_output) return output class DiTBlock(nn.Module): def __init__(self, hidden_size, n_heads): super().__init__() self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.attn = MultiHeadAttentionWithRoPE(hidden_size, n_heads) self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.mlp = nn.Sequential( nn.Linear(hidden_size, 4 * hidden_size), nn.GELU(), nn.Linear(4 * hidden_size, hidden_size) ) self.adaLN_modulation = nn.Sequential( nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True) ) def forward(self, x, c): shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) x_norm1 = modulate(self.norm1(x), shift_msa, scale_msa) attn_output = self.attn(x_norm1) x = x + gate_msa.unsqueeze(1) * attn_output x_norm2 = modulate(self.norm2(x), shift_mlp, scale_mlp) mlp_output = self.mlp(x_norm2) x = x + gate_mlp.unsqueeze(1) * mlp_output return x class MDLM(nn.Module): def __init__(self, vocab_size, model_dim, n_heads, n_layers): super().__init__() self.vocab_size = vocab_size self.model_dim = model_dim self.mask_token_id = vocab_size self.token_embedder = nn.Embedding(vocab_size, model_dim) self.time_embedder = TimestepEmbedder(model_dim) self.transformer_blocks = nn.ModuleList([ DiTBlock(model_dim, n_heads) for _ in range(n_layers) ]) self.final_norm = nn.LayerNorm(model_dim) self.lm_head = nn.Linear(model_dim, vocab_size) self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, (nn.Linear, nn.Embedding)): module.weight.data.normal_(mean=0.0, std=0.02) if isinstance(module, nn.Linear) and module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.LayerNorm): if module.bias is not None: module.bias.data.zero_() if module.weight is not None: module.weight.data.fill_(1.0) def forward(self, x, t): x_embed = self.token_embedder(x) t_embed = self.time_embedder(t) for block in self.transformer_blocks: x_embed = block(x_embed, t_embed) x_embed = self.final_norm(x_embed) logits = self.lm_head(x_embed) return logits # --- PyTorch Lightning Module --- class MDLMLightningModule(pl.LightningModule): def __init__(self, args, tokenizer): super().__init__() self.save_hyperparameters(ignore=['tokenizer']) self.args = args self.tokenizer = tokenizer self.peptide_analyzer = PeptideAnalyzer() # Initialize model self.model = MDLM( vocab_size=tokenizer.vocab_size, model_dim=args.model_dim, n_heads=args.n_heads, n_layers=args.n_layers ) # For tracking steps self.automatic_optimization = True self.validation_step_outputs = [] def forward(self, x, t): return self.model(x, t) def _compute_invalid_loss(self, logits): batch_token_ids = torch.argmax(logits, dim=-1) # (batch_size, seq_length) sampled_sequences = self.tokenizer.batch_decode(batch_token_ids) penalties = torch.tensor( [1 if not self.peptide_analyzer.is_peptide(seq) else 0 for seq in sampled_sequences], dtype=torch.float32, device=self.device ) sampled_probs = torch.softmax(logits, dim=-1).gather(dim=-1, index=batch_token_ids.unsqueeze(-1)).squeeze(-1).to(self.device) scaled_penalty = penalties[:, None] * sampled_probs # (batch_size, seq_length) return scaled_penalty def _loss(self, logits, x_1, attn_mask): # Standard cross-entropy loss ce_loss = F.cross_entropy( logits.view(-1, self.model.vocab_size), x_1.view(-1), reduction='none' ).view(x_1.shape[0], -1) # ce_loss = (ce_loss * attn_mask).sum() / attn_mask.sum() # validity_weight = self.args.validity_weight * min(1.0, (self.current_epoch + 1) / self.trainer.max_epochs) invalid_loss = self._compute_invalid_loss(logits) # (batch_size, seq_length) loss = ce_loss + self.args.validity_weight * invalid_loss nlls = loss * attn_mask num_tokens = attn_mask.sum() batch_nll = nlls.sum() token_nll = batch_nll / num_tokens return token_nll, (ce_loss*attn_mask).sum() / num_tokens, (invalid_loss*attn_mask).sum() / num_tokens def training_step(self, batch, batch_idx): x_1 = batch['input_ids'].clone().detach().to(self.device) attn_mask = batch['attention_mask'].clone().detach().to(self.device) bond_mask = batch['bond_mask'].clone().detach().to(self.device).bool() batch_size, _ = x_1.shape # ReDi approach: random start -> target x_0 = torch.randint(0, self.model.vocab_size, x_1.shape, device=self.device) t_continuous = torch.rand(batch_size, device=self.device) # mask = torch.rand(x_1.shape, device=self.device) < t_continuous.view(-1, 1) # x_t = torch.where(mask, x_1, x_0) peptide_bond_prob = t_continuous.view(-1, 1) ** self.args.gamma # slower increase non_peptide_prob = t_continuous.view(-1, 1) # linear increase masking_prob = torch.where(bond_mask, peptide_bond_prob, non_peptide_prob) mask = torch.rand(x_1.shape, device=self.device) < masking_prob x_t = torch.where(mask, x_1, x_0) logits = self.model(x_t, t_continuous) token_nll, ce_loss, invalid_loss = self._loss(logits, x_1, attn_mask) # Logging self.log('train/token_nll', token_nll.item(), on_step=True, on_epoch=True, prog_bar=True, batch_size=x_1.size(0), sync_dist=True) self.log('train/ce_loss', ce_loss.item(), on_step=True, on_epoch=True, batch_size=x_1.size(0), sync_dist=True) self.log('train/invalid_loss', invalid_loss.item(), on_step=True, on_epoch=True, batch_size=x_1.size(0), sync_dist=True) # self.log('train/validity_weight', validity_weight, on_step=False, on_epoch=True, batch_size=x_1.size(0)) return token_nll def validation_step(self, batch, batch_idx): x_1 = batch['input_ids'].clone().detach().to(self.device) attn_mask = batch['attention_mask'].clone().detach().to(self.device) bond_mask = batch['bond_mask'].clone().detach().to(self.device).bool() batch_size, _ = x_1.shape # ReDi approach: random start -> target x_0 = torch.randint(0, self.model.vocab_size, x_1.shape, device=self.device) t_continuous = torch.rand(batch_size, device=self.device) peptide_bond_prob = t_continuous.view(-1, 1) ** self.args.gamma # slower increase non_peptide_prob = t_continuous.view(-1, 1) # linear increase masking_prob = torch.where(bond_mask, peptide_bond_prob, non_peptide_prob) mask = torch.rand(x_1.shape, device=self.device) < masking_prob x_t = torch.where(mask, x_1, x_0) logits = self.model(x_t, t_continuous) token_nll, ce_loss, invalid_loss = self._loss(logits, x_1, attn_mask) self.log('val/token_nll', token_nll.item(), on_step=True, on_epoch=True, prog_bar=True, batch_size=x_1.size(0), sync_dist=True) self.log('val/ce_loss', ce_loss.item(), on_step=True, on_epoch=True, batch_size=x_1.size(0), sync_dist=True) self.log('val/invalid_loss', invalid_loss.item(), on_step=True, on_epoch=True, batch_size=x_1.size(0), sync_dist=True) def configure_optimizers(self): optimizer = AdamW( self.parameters(), lr=self.args.learning_rate, weight_decay=self.args.weight_decay ) # Calculate total steps if hasattr(self.trainer, 'estimated_stepping_batches'): num_training_steps = self.trainer.estimated_stepping_batches else: # Fallback calculation num_training_steps = len(self.trainer.datamodule.train_dataloader()) * self.trainer.max_epochs warmup_steps = int(num_training_steps * 0.1) def lr_lambda(current_step): if current_step < warmup_steps: lr_range = self.args.learning_rate - (self.args.learning_rate * 0.1) lr = (self.args.learning_rate * 0.1) + lr_range * (current_step / warmup_steps) return lr / self.args.learning_rate else: progress = (current_step - warmup_steps) / (num_training_steps - warmup_steps) cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress)) lr_range = self.args.learning_rate - (self.args.learning_rate * 0.1) lr = (self.args.learning_rate * 0.1) + lr_range * cosine_decay return lr / self.args.learning_rate scheduler = LambdaLR(optimizer, lr_lambda) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": scheduler, "interval": "step", "frequency": 1, }, } # --- Main Execution --- def main(args): # Set up checkpoint directory checkpoint_dir = (args.checkpoint_dir + f"correct_lr{args.learning_rate}_wd{args.weight_decay}_layer{args.n_layers}_" f"head{args.n_heads}_valweight{args.validity_weight}") print(f"Saving to {checkpoint_dir}") os.makedirs(checkpoint_dir, exist_ok=True) print("Loading tokenizer...") tokenizer = SMILES_SPE_Tokenizer('/scratch/pranamlab/tong/ReDi_discrete/smiles/smiles_tokenizer/old_vocab.txt', '/scratch/pranamlab/tong/ReDi_discrete/smiles/smiles_tokenizer/old_splits.txt') print(f"Tokenizer loaded. Vocab size: {tokenizer.vocab_size}") # Initialize data module data_module = dynamic_dataloader.CustomDataModule('./data/11M_smiles_old_tokenizer_no_limit/', tokenizer) # Initialize model model = MDLMLightningModule(args, tokenizer) # Set up logger logger = WandbLogger( project="smiles-redi-training", # or your preferred project name entity="programmablebio", name=f"lr{args.learning_rate}_dim{args.model_dim}_head{args.n_heads}_layer{args.n_layers}", save_dir=checkpoint_dir ) # Set up callbacks callbacks = [ ModelCheckpoint( dirpath=checkpoint_dir, filename='best', monitor='val/token_nll', mode='min', save_top_k=1, save_last=True, # every_n_train_steps=10000 # This will save every 1000 steps AND when val/nll improves ), LearningRateMonitor(logging_interval='step') ] # Initialize trainer trainer = pl.Trainer( max_epochs=args.epochs, devices=torch.cuda.device_count(), accelerator='gpu', strategy=DDPStrategy(find_unused_parameters=False), num_nodes=int(os.environ.get("SLURM_NNODES", 1)), precision="bf16", gradient_clip_val=args.grad_clip if args.grad_clip > 0 else None, callbacks=callbacks, logger=logger, log_every_n_steps=100, check_val_every_n_epoch=True, # val_check_interval=10000, accumulate_grad_batches=8, enable_progress_bar=True, enable_model_summary=True ) print(f"Model initialized with {sum(p.numel() for p in model.parameters()):,} parameters.") print("Starting training...") # Train the model trainer.fit(model, data_module) print("Training complete.") print(f"Best checkpoint saved at: {trainer.checkpoint_callback.best_model_path}") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Train ReDi model for SMILES generation with RoPE using PyTorch Lightning") # Model arguments parser.add_argument("--model_dim", type=int, default=1024) parser.add_argument("--n_heads", type=int, default=8) parser.add_argument("--n_layers", type=int, default=6) # Training arguments parser.add_argument("--epochs", type=int, default=50) parser.add_argument("--learning_rate", type=float, default=1e-4) parser.add_argument("--weight_decay", type=float, default=1e-5) parser.add_argument("--label_smoothing", type=float, default=0) parser.add_argument("--grad_clip", type=float, default=1.0) parser.add_argument("--gamma", type=float, default=2.0) # Validity arguments parser.add_argument("--validity_weight", type=float, default=0.1) parser.add_argument("--validity_check_freq", type=int, default=10) parser.add_argument("--validity_eval_batches", type=int, default=20) # Logging arguments parser.add_argument("--checkpoint_dir", type=str, default="./checkpoints_smiles") args = parser.parse_args() main(args)