# train.py # Description: A complete script to train a ReDi model with a continuous time variable t in [0, 1]. 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 from tqdm import tqdm import wandb # --- Model Architecture --- # Based on the DiT (Diffusion Transformer) architecture, adapted for discrete data (MDLM). def modulate(x, shift, scale): """ Modulates the input tensor x with a shift and scale. This is a key component of the DiT architecture, allowing conditioning on the timestep embedding. """ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) class TimestepEmbedder(nn.Module): """ Embeds a continuous scalar timestep t in [0, 1] into a vector representation. """ 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): # t is shape (batch_size,), needs to be (batch_size, 1) for the Linear layer. return self.mlp(t.unsqueeze(-1)) class DiTBlock(nn.Module): """ A single block of the Diffusion Transformer. """ def __init__(self, hidden_size, n_heads): super().__init__() self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) self.attn = nn.MultiheadAttention(hidden_size, n_heads, batch_first=True) 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_norm1, 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): """ Masked Diffusion Language Model (MDLM) using a DiT backbone. """ def __init__(self, vocab_size, seq_len, model_dim, n_heads, n_layers): super().__init__() self.vocab_size = vocab_size self.seq_len = seq_len self.model_dim = model_dim self.mask_token_id = vocab_size # Use vocab_size as the ID for the mask token self.token_embedder = nn.Embedding(vocab_size + 1, model_dim) # +1 for the mask token self.pos_embedder = nn.Parameter(torch.randn(1, seq_len, 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): seq_len = x.shape[1] x_embed = self.token_embedder(x) + self.pos_embedder[:, :seq_len, :] 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 # --- Learning Rate Scheduler --- def get_lr_scheduler(optimizer, warmup_steps, total_steps, lr_min, lr_max): """ Creates a step-based learning rate scheduler with a linear warmup phase from lr_min to lr_max, followed by a cosine annealing phase from lr_max back down to lr_min. """ def lr_lambda(current_step): # Linear warmup phase if current_step < warmup_steps: lr_range = lr_max - lr_min lr = lr_min + lr_range * (current_step / warmup_steps) return lr / lr_max # Cosine annealing phase else: progress = (current_step - warmup_steps) / (total_steps - warmup_steps) cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress)) lr_range = lr_max - lr_min lr = lr_min + lr_range * cosine_decay return lr / lr_max return LambdaLR(optimizer, lr_lambda) # --- Training and Validation Functions --- def train_one_epoch(model, dataloader, optimizer, scheduler, device, epoch, args): model.train() total_loss = 0.0 progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1} [Train]") for batch in progress_bar: optimizer.zero_grad() x_1 = torch.tensor(batch['input_ids']).to(device) batch_size, _ = x_1.shape x_0 = torch.randint(0, model.vocab_size, x_1.shape, device=device) t_continuous = torch.rand(batch_size, device=device) mask = torch.rand(x_1.shape, device=device) < t_continuous.view(-1, 1) x_t = torch.where(mask, x_1, x_0) logits = model(x_t, t_continuous) loss = F.cross_entropy(logits.view(-1, model.vocab_size), x_1.view(-1), label_smoothing=args.label_smoothing) loss.backward() optimizer.step() scheduler.step() total_loss += loss.item() progress_bar.set_postfix(loss=loss.item(), lr=scheduler.get_last_lr()[0]) # wandb.log({"train_loss_step": loss.item(), "learning_rate": scheduler.get_last_lr()[0]}) return total_loss / len(dataloader) def validate(model, val_dataloader, device, epoch, args): """ Performs validation, calculating NLL, Perplexity, and TC error. """ model.eval() total_val_nll = 0.0 total_tc = 0.0 tc_batches = 0 progress_bar = tqdm(val_dataloader, desc=f"Epoch {epoch+1} [Val]") with torch.no_grad(): for i, batch in enumerate(progress_bar): x_1 = torch.tensor(batch['input_ids']).to(device) batch_size, seq_len = x_1.shape x_0 = torch.randint(0, model.vocab_size, x_1.shape, device=device) t_continuous = torch.rand(batch_size, device=device) mask = torch.rand(x_1.shape, device=device) < t_continuous.view(-1, 1) x_t = torch.where(mask, x_1, x_0) logits = model(x_t, t_continuous) val_nll = F.cross_entropy(logits.view(-1, model.vocab_size), x_1.view(-1)) total_val_nll += val_nll.item() if i < args.tc_batches: k = args.tc_k_samples p_marginal = F.softmax(logits, dim=-1) sampled_x1 = torch.multinomial(p_marginal.view(-1, model.vocab_size), k, replacement=True).view(batch_size, seq_len, k) kl_divs = [] for b in range(batch_size): sample_tuples = [tuple(s.tolist()) for s in sampled_x1[b].T] joint_counts = Counter(sample_tuples) p_joint_est = {k: v / len(sample_tuples) for k, v in joint_counts.items()} kl_sum = 0 for seq_tuple, p_j in p_joint_est.items(): log_p_marginal_prod = 0 for pos, token_id in enumerate(seq_tuple): log_p_marginal_prod += torch.log(p_marginal[b, pos, token_id] + 1e-9) kl_sum += p_j * (math.log(p_j + 1e-9) - log_p_marginal_prod) kl_divs.append(kl_sum) total_tc += sum(kl_divs) / len(kl_divs) tc_batches += 1 avg_val_nll = total_val_nll / len(val_dataloader) perplexity = math.exp(avg_val_nll) avg_tc = total_tc / tc_batches if tc_batches > 0 else 0 return avg_val_nll, perplexity, avg_tc # --- Main Execution --- def main(args): # try: # wandb.login(key="811c943b63ebdf9409a9365602a39da3cfcf0062") # except Exception as e: # print(f"Could not log in to wandb: {e}") # return # wandb.init(project=args.wandb_project, name=f"lr{args.learning_rate}_wd{args.weight_decay}_layer{args.n_layers}_head{args.n_heads}_labelsmoothing{args.label_smoothing}", entity="programmablebio", config=args) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") args.checkpoint_dir = args.checkpoint_dir + f"lr{args.learning_rate}_wd{args.weight_decay}_layer{args.n_layers}_head{args.n_heads}_labelsmoothing{args.label_smoothing}" print(f"Saving to {args.checkpoint_dir}") os.makedirs(args.checkpoint_dir, exist_ok=True) print("Loading datasets...") train_dataset = load_from_disk(args.train_dataset_path) val_dataset = load_from_disk(args.val_dataset_path) train_dataloader = DataLoader(train_dataset, batch_size=None, shuffle=False) val_dataloader = DataLoader(val_dataset, batch_size=None, shuffle=False) print("Initializing model...") model = MDLM(args.vocab_size, args.seq_len, args.model_dim, args.n_heads, args.n_layers).to(device) print(f"Model initialized with {sum(p.numel() for p in model.parameters()):,} parameters.") optimizer = AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) num_training_steps = args.epochs * len(train_dataloader) warmup_steps = int(num_training_steps * 0.1) scheduler = get_lr_scheduler(optimizer, warmup_steps, num_training_steps, args.learning_rate * 0.1, args.learning_rate) best_val_nll = float('inf') print("Starting training...") for epoch in range(args.epochs): train_loss = train_one_epoch(model, train_dataloader, optimizer, scheduler, device, epoch, args) val_nll, perplexity, tc_error = validate(model, val_dataloader, device, epoch, args) print(f"Epoch {epoch+1}/{args.epochs} -> Train Loss: {train_loss:.4f}, Val NLL: {val_nll:.4f}, Val PPL: {perplexity:.2f}, TC: {tc_error:.4f}") # wandb.log({ # "epoch": epoch + 1, # "train_loss_epoch": train_loss, # "val_nll_epoch": val_nll, # "val_perplexity": perplexity, # "conditional_total_correlation": tc_error, # }) # checkpoint_path = os.path.join(args.checkpoint_dir, f"epoch_{epoch+1}.pt") # torch.save({ # 'epoch': epoch + 1, 'model_state_dict': model.state_dict(), # 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), # 'val_nll': val_nll, 'args': args # }, checkpoint_path) # print(f"Checkpoint saved to {checkpoint_path}") if val_nll < best_val_nll: best_val_nll = val_nll best_checkpoint_path = os.path.join(args.checkpoint_dir, "best_checkpoint.pt") torch.save({ 'epoch': epoch + 1, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict(), 'val_nll': val_nll, 'tc_error': tc_error, 'args': args }, best_checkpoint_path) print(f"New best checkpoint saved to {best_checkpoint_path} (Val NLL: {val_nll:.4f})") # wandb.finish() print("Training complete.") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Train a ReDi (MDLM) model with self-contained evaluation.") parser.add_argument("--train_dataset_path", type=str, required=True) parser.add_argument("--val_dataset_path", type=str, required=True) 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) parser.add_argument("--vocab_size", type=int, default=24) parser.add_argument("--seq_len", type=int, default=100) 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("--tc_batches", type=int, default=20, help="Number of validation batches to use for TC calculation.") parser.add_argument("--tc_k_samples", type=int, default=50, help="Number of samples (k) per data point for TC approximation.") parser.add_argument("--wandb_project", type=str, default="redi-training") parser.add_argument("--checkpoint_dir", type=str, default="./checkpoints") args = parser.parse_args() main(args)