Upload lightbulb_lm.py

lightbulb_lm.py

@@ -0,0 +1,517 @@
+#!/usr/bin/env python
+
+import argparse
+import math
+import os
+
+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
+from torch.amp import autocast, GradScaler
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+# Set the device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Train Transformer model with advanced features.')
+    parser.add_argument('--model_name', type=str, default='gpt2', help='Pretrained model name or path')
+    parser.add_argument('--dataset_name', type=str, default='wikitext', help='Dataset name from HuggingFace Datasets')
+    parser.add_argument('--dataset_config', type=str, default='wikitext-2-raw-v1', help='Dataset configuration name')
+    parser.add_argument('--batch_size', type=int, default=8, help='Batch size')
+    parser.add_argument('--num_epochs', type=int, default=3, help='Number of epochs')
+    parser.add_argument('--max_length', type=int, default=128, help='Maximum sequence length')
+    parser.add_argument('--accumulation_steps', type=int, default=4, help='Gradient accumulation steps')
+    parser.add_argument('--learning_rate', type=float, default=1e-4, help='Learning rate')
+    parser.add_argument('--weight_decay', type=float, default=1e-2, help='Weight decay')
+    parser.add_argument('--alpha', type=float, default=0.1, help='Entropy regularization weight')
+    parser.add_argument('--beta', type=float, default=0.1, help='Variance regularization weight')
+    parser.add_argument('--max_grad_norm', type=float, default=1.0, help='Max gradient norm for clipping')
+    parser.add_argument('--save_dir', type=str, default='./models', help='Directory to save the models')
+    parser.add_argument('--temperature', type=float, default=1.0, help='Temperature parameter for entropy and variance')
+    args = parser.parse_args()
+    return args
+
+
+def load_data(args, tokenizer):
+    # Load the dataset
+    dataset = load_dataset(args.dataset_name, args.dataset_config)
+    
+    # Ensure the tokenizer has a padding token
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    def tokenize_function(examples):
+        return tokenizer(examples['text'], truncation=True, max_length=args.max_length)
+
+    tokenized_datasets = dataset.map(
+        tokenize_function,
+        batched=True,
+        num_proc=4,
+        remove_columns=dataset['train'].column_names,
+    )
+
+    # Build inputs and labels for language modeling
+    block_size = args.max_length
+
+    def group_texts(examples):
+        # Concatenate all texts
+        concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+        total_length = len(concatenated_examples['input_ids'])
+        # We drop the small remainder
+        total_length = (total_length // block_size) * block_size
+        # Split by chunks of block_size
+        result = {
+            k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+            for k, t in concatenated_examples.items()
+        }
+        result['labels'] = result['input_ids'].copy()
+        return result
+
+    lm_datasets = tokenized_datasets.map(
+        group_texts,
+        batched=True,
+        num_proc=4,
+    )
+
+    # Create DataLoader
+    train_dataset = lm_datasets['train']
+    eval_dataset = lm_datasets['validation'] if 'validation' in lm_datasets else lm_datasets['test']
+
+    data_collator = lambda data: {
+        'input_ids': torch.tensor([f['input_ids'] for f in data], dtype=torch.long),
+        'labels': torch.tensor([f['labels'] for f in data], dtype=torch.long)
+    }
+
+    train_loader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size, collate_fn=data_collator)
+    eval_loader = DataLoader(eval_dataset, shuffle=False, batch_size=args.batch_size, collate_fn=data_collator)
+
+    return train_loader, eval_loader
+
+
+class RotaryPositionalEncoding(nn.Module):
+    def __init__(self, d_model):
+        super(RotaryPositionalEncoding, self).__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))
+        self.register_buffer('inv_freq', inv_freq)
+
+    def forward(self, x):
+        seq_len, batch_size, _ = x.size()
+        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+        sinusoid_inp = torch.einsum("i,j->ij", t, self.inv_freq)
+        sin = sinusoid_inp.sin().unsqueeze(1)  # (seq_len, 1, d_model/2)
+        cos = sinusoid_inp.cos().unsqueeze(1)  # (seq_len, 1, d_model/2)
+
+        x1 = x[..., 0::2]
+        x2 = x[..., 1::2]
+
+        # Apply rotation
+        x_rotated = torch.zeros_like(x)
+        x_rotated[..., 0::2] = x1 * cos - x2 * sin
+        x_rotated[..., 1::2] = x1 * sin + x2 * cos
+
+        return x_rotated
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, num_heads):
+        super(MultiHeadAttention, self).__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        self.d_k = d_model // num_heads
+        self.num_heads = num_heads
+        self.linear_q = nn.Linear(d_model, d_model)
+        self.linear_k = nn.Linear(d_model, d_model)
+        self.linear_v = nn.Linear(d_model, d_model)
+        self.linear_out = nn.Linear(d_model, d_model)
+    
+    def forward(self, query, key, value, mask=None):
+        batch_size = query.size(0)
+        query = self.linear_q(query).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        key = self.linear_k(key).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        value = self.linear_v(value).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        
+        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.d_k)
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, -1e9)
+        attn = F.softmax(scores, dim=-1)
+        output = torch.matmul(attn, value)
+        
+        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k)
+        return self.linear_out(output)
+
+
+class MoE(nn.Module):
+    def __init__(self, d_model, num_experts, d_ff, top_k=2, dropout=0.1):
+        super(MoE, self).__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.experts = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(d_model, d_ff),
+                nn.GELU() if i % 2 == 0 else nn.SiLU(),
+                nn.Linear(d_ff, d_model)
+            )
+            for i in range(num_experts)
+        ])
+        self.gate = nn.Linear(d_model, num_experts)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        batch_size, seq_len, d_model = x.size()
+        # Compute gating scores
+        gate_scores = self.gate(x)  # (batch_size, seq_len, num_experts)
+        top_k_scores, top_k_indices = torch.topk(gate_scores, self.top_k, dim=-1)  # (batch_size, seq_len, top_k)
+        top_k_scores = F.softmax(top_k_scores, dim=-1)  # (batch_size, seq_len, top_k)
+
+        # Initialize output
+        output = torch.zeros_like(x)
+
+        # Flatten batch and sequence dimensions
+        x_flat = x.view(-1, d_model)  # (batch_size * seq_len, d_model)
+        output_flat = output.view(-1, d_model)
+        top_k_indices_flat = top_k_indices.view(-1, self.top_k)  # (batch_size * seq_len, top_k)
+        top_k_scores_flat = top_k_scores.view(-1, self.top_k)  # (batch_size * seq_len, top_k)
+
+        for k in range(self.top_k):
+            expert_idx_flat = top_k_indices_flat[:, k]  # (batch_size * seq_len)
+            expert_scores_flat = top_k_scores_flat[:, k]  # (batch_size * seq_len)
+            for e in range(self.num_experts):
+                mask = (expert_idx_flat == e)  # Boolean mask
+                if mask.any():
+                    x_masked = x_flat[mask]  # Select tokens for expert e
+                    expert_output = self.experts[e](x_masked)  # Apply expert e
+                    output_flat[mask] += expert_scores_flat[mask].unsqueeze(-1) * expert_output
+
+        output = output_flat.view(batch_size, seq_len, d_model)
+        return self.dropout(output)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, d_model, num_heads, d_ff, num_experts, dropout=0.1, top_k=2):
+        super(TransformerBlock, self).__init__()
+        self.self_attention = MultiHeadAttention(d_model, num_heads)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.cross_attention = MultiHeadAttention(d_model, num_heads)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.moe = MoE(d_model, num_experts, d_ff, top_k, dropout)
+        self.norm3 = nn.LayerNorm(d_model)
+
+    def forward(self, x, mask=None, enc_output=None, enc_mask=None):
+        # Self-attention
+        attn_output = self.self_attention(x, x, x, mask)
+        x = self.norm1(x + attn_output)
+        # Cross-attention (only in decoder)
+        if enc_output is not None:
+            cross_attn_output = self.cross_attention(x, enc_output, enc_output, enc_mask)
+            x = self.norm2(x + cross_attn_output)
+        # Feedforward/MoE
+        moe_output = self.moe(x)
+        return self.norm3(x + moe_output)
+
+
+class Transformer(nn.Module):
+    def __init__(self, input_dim, d_model, num_heads, num_layers, d_ff, num_experts, output_dim, dropout=0.1, top_k=2):
+        super(Transformer, self).__init__()
+        self.embedding = nn.Embedding(input_dim, d_model, padding_idx=input_dim - 1)
+        self.rotary_positional_encoding = RotaryPositionalEncoding(d_model)
+        self.encoder_layers = nn.ModuleList(
+            [TransformerBlock(d_model, num_heads, d_ff, num_experts, dropout, top_k) for _ in range(num_layers)]
+        )
+        self.decoder_layers = nn.ModuleList(
+            [TransformerBlock(d_model, num_heads, d_ff, num_experts, dropout, top_k) for _ in range(num_layers)]
+        )
+        self.output_layer = nn.Linear(d_model, output_dim)
+        self.d_model = d_model
+
+    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
+        # Encoder
+        src = self.embedding(src) * math.sqrt(self.d_model)
+        src = src.transpose(0, 1)  # (batch_size, seq_len, d_model) -> (seq_len, batch_size, d_model)
+        src = self.rotary_positional_encoding(src)
+        src = src.transpose(0, 1)  # (seq_len, batch_size, d_model) -> (batch_size, seq_len, d_model)
+        for layer in self.encoder_layers:
+            src = layer(src, src_mask)
+
+        # Decoder
+        tgt = self.embedding(tgt) * math.sqrt(self.d_model)
+        tgt = tgt.transpose(0, 1)
+        tgt = self.rotary_positional_encoding(tgt)
+        tgt = tgt.transpose(0, 1)
+        for layer in self.decoder_layers:
+            tgt = layer(tgt, tgt_mask, src, src_mask)
+        output = self.output_layer(tgt)
+        return output
+
+    def generate(self, src, tokenizer, max_length=20, temperature=1.0):
+        """
+        Generate sequences using differentiable sampling (Gumbel-Softmax).
+
+        Args:
+            src (torch.Tensor): Source input tensor of shape (batch_size, seq_len)
+            tokenizer (transformers.PreTrainedTokenizer): Tokenizer to access special tokens
+            max_length (int): Maximum length of the generated sequence
+            temperature (float): Temperature parameter for Gumbel-Softmax
+
+        Returns:
+            torch.Tensor: Generated sequences of shape (batch_size, max_length)
+            torch.Tensor: Entropy values for each time step
+            torch.Tensor: Variance values for each time step
+        """
+        batch_size = src.size(0)
+
+        # Encode the source
+        src_enc = self.embedding(src) * math.sqrt(self.d_model)
+        src_enc = src_enc.transpose(0, 1)
+        src_enc = self.rotary_positional_encoding(src_enc)
+        src_enc = src_enc.transpose(0, 1)
+        for layer in self.encoder_layers:
+            src_enc = layer(src_enc)
+
+        # Initialize decoder input with <sos> tokens
+        tgt_seq = torch.full((batch_size, 1), tokenizer.bos_token_id, dtype=torch.long, device=src.device)
+        entropies = []
+        variances = []
+
+        for _ in range(max_length):
+            tgt_emb = self.embedding(tgt_seq) * math.sqrt(self.d_model)
+            tgt_emb = tgt_emb.transpose(0, 1)
+            tgt_emb = self.rotary_positional_encoding(tgt_emb)
+            tgt_emb = tgt_emb.transpose(0, 1)
+            tgt_dec = tgt_emb
+            for layer in self.decoder_layers:
+                tgt_dec = layer(tgt_dec, None, src_enc, None)
+            output = self.output_layer(tgt_dec)  # (batch_size, seq_len, vocab_size)
+            logits = output[:, -1, :]  # Get logits for the last time step
+
+            # Compute token probabilities
+            probs = F.softmax(logits / temperature, dim=-1)  # (batch_size, vocab_size)
+
+            # Compute entropy
+            entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)  # (batch_size)
+            entropies.append(entropy)
+
+            # Sample token using Gumbel-Softmax
+            gumbel_noise = -torch.log(-torch.log(torch.rand_like(probs) + 1e-9) + 1e-9)
+            y = (logits + gumbel_noise) / temperature
+            y = F.softmax(y, dim=-1)  # (batch_size, vocab_size)
+
+            # Compute variance
+            variance = torch.var(y, dim=-1)  # (batch_size)
+            variances.append(variance)
+
+            # Get token indices (argmax for hard selection)
+            next_tokens = torch.argmax(y, dim=-1, keepdim=True)  # (batch_size, 1)
+            tgt_seq = torch.cat([tgt_seq, next_tokens], dim=1)
+
+        # Stack entropies and variances
+        entropies = torch.stack(entropies, dim=1)  # (batch_size, max_length)
+        variances = torch.stack(variances, dim=1)  # (batch_size, max_length)
+
+        return tgt_seq[:, 1:], entropies, variances  # Exclude the initial <sos> token
+
+
+def compute_loss(output, target, padding_idx, alpha=0.1, beta=0.1, temperature=1.0):
+    """
+    Compute the loss with entropy and variance regularization.
+
+    Args:
+        output (torch.Tensor): Model output logits of shape (batch_size, seq_len, vocab_size)
+        target (torch.Tensor): Target sequences of shape (batch_size, seq_len)
+        padding_idx (int): Padding index to ignore in the loss
+        alpha (float): Weight for the entropy regularization term
+        beta (float): Weight for the variance regularization term
+        temperature (float): Temperature parameter for computing probabilities
+
+    Returns:
+        torch.Tensor: Scalar loss value
+    """
+    # Cross-entropy loss
+    output_flat = output.contiguous().view(-1, output.size(-1))
+    target_flat = target.contiguous().view(-1)
+    ce_loss = F.cross_entropy(
+        output_flat,
+        target_flat,
+        ignore_index=padding_idx
+    )
+
+    # Compute probabilities
+    probs = F.softmax(output / temperature, dim=-1)  # (batch_size, seq_len, vocab_size)
+
+    # Compute entropy
+    entropy = -torch.sum(probs * torch.log(probs + 1e-9), dim=-1)  # (batch_size, seq_len)
+    entropy_loss = -alpha * torch.mean(entropy)
+
+    # Compute variance
+    variance = torch.var(probs, dim=-1)  # (batch_size, seq_len)
+    variance_loss = -beta * torch.mean(variance)
+
+    # Total loss
+    total_loss = ce_loss + entropy_loss + variance_loss
+    return total_loss
+
+
+def train_epoch(model, train_loader, optimizer, scheduler, scaler, args, padding_idx):
+    model.train()
+    total_loss = 0.0
+    optimizer.zero_grad()
+    print(f"Starting training epoch with {len(train_loader)} batches...")
+    for i, batch in enumerate(train_loader):
+        print(f"Processing batch {i+1}/{len(train_loader)}...")
+        src_batch = batch['input_ids'].to(device)
+        tgt_batch = batch['labels'].to(device)
+
+        with autocast(device_type='cuda'):
+            print("Forward pass...")
+            output = model(src_batch, tgt_batch[:, :-1])
+            print("Computing loss...")
+            loss = compute_loss(
+                output, 
+                tgt_batch[:, 1:], 
+                padding_idx, 
+                alpha=args.alpha, 
+                beta=args.beta, 
+                temperature=args.temperature
+            )
+            loss = loss / args.accumulation_steps
+
+        print("Backward pass...")
+        scaler.scale(loss).backward()
+
+        if (i + 1) % args.accumulation_steps == 0:
+            print("Gradient clipping...")
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+            print("Optimizer step...")
+            scaler.step(optimizer)
+            scaler.update()
+
+            print("Zeroing gradients...")
+            optimizer.zero_grad()
+
+            print("Updating learning rate...")
+            scheduler.step()
+
+        total_loss += loss.item() * args.accumulation_steps
+        print(f"Batch {i+1} completed. Current loss: {loss.item():.4f}")
+
+    avg_loss = total_loss / len(train_loader)
+    print(f"Epoch completed. Average loss: {avg_loss:.4f}")
+    return avg_loss
+
+
+def evaluate(model, eval_loader, args, padding_idx):
+    model.eval()
+    total_loss = 0.0
+    with torch.no_grad():
+        for batch in eval_loader:
+            src_batch = batch['input_ids'].to(device)
+            tgt_batch = batch['labels'].to(device)
+
+            with autocast(device_type='cuda'):
+                # Forward pass
+                output = model(src_batch, tgt_batch[:, :-1])
+                # Compute loss
+                loss = compute_loss(
+                    output, 
+                    tgt_batch[:, 1:], 
+                    padding_idx, 
+                    alpha=args.alpha, 
+                    beta=args.beta, 
+                    temperature=args.temperature
+                )
+
+            total_loss += loss.item()
+
+    avg_loss = total_loss / len(eval_loader)
+    return avg_loss
+
+
+def main():
+    args = parse_args()
+    print("Arguments parsed successfully.")
+
+    # Create save directory
+    if not os.path.exists(args.save_dir):
+        os.makedirs(args.save_dir)
+    print(f"Save directory created: {args.save_dir}")
+
+    # Load tokenizer
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    print("Tokenizer loaded successfully.")
+
+    # Load data
+    print("Loading and preprocessing data...")
+    train_loader, eval_loader = load_data(args, tokenizer)
+    print("Data loaded and preprocessed successfully.")
+
+    # Define model parameters
+    input_dim = len(tokenizer)
+    d_model = 512
+    num_heads = 8
+    num_layers = 6
+    d_ff = 2048
+    num_experts = 4
+    output_dim = input_dim
+    dropout = 0.1
+    top_k = 2
+
+    print("Initializing model...")
+    model = Transformer(
+        input_dim, d_model, num_heads, num_layers, d_ff, num_experts, output_dim, dropout, top_k
+    )
+    model = model.to(device)
+    print(f"Model initialized and moved to device: {device}")
+
+    padding_idx = tokenizer.pad_token_id
+
+    print("Setting up optimizer and scheduler...")
+    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
+    scheduler = CosineAnnealingLR(optimizer, T_max=args.num_epochs)
+    scaler = GradScaler()
+    print("Optimizer and scheduler set up successfully.")
+
+    print("Starting training loop...")
+    for epoch in range(args.num_epochs):
+        print(f"Epoch {epoch + 1}/{args.num_epochs} started.")
+        avg_train_loss = train_epoch(
+            model, 
+            train_loader, 
+            optimizer, 
+            scheduler, 
+            scaler, 
+            args, 
+            padding_idx
+        )
+        print(f"Epoch {epoch + 1}/{args.num_epochs} training completed.")
+        
+        print(f"Starting evaluation for epoch {epoch + 1}...")
+        avg_eval_loss = evaluate(model, eval_loader, args, padding_idx)
+        print(f"Evaluation for epoch {epoch + 1} completed.")
+        
+        print(f"Epoch {epoch + 1}/{args.num_epochs}, Train Loss: {avg_train_loss:.4f}, Eval Loss: {avg_eval_loss:.4f}")
+
+        model_save_path = os.path.join(args.save_dir, f"model_epoch_{epoch + 1}.pt")
+        torch.save(model.state_dict(), model_save_path)
+        print(f"Model saved for epoch {epoch + 1}")
+
+    print("Training completed.")
+
+
+if __name__ == '__main__':
+    main()
+
+
+'''
+Example usage:
+python lightbulb.py --model_name gpt2 --dataset_name wikitext --dataset_config wikitext-2-raw-v1 --batch_size 8 --num_epochs 3
+'''