Upload 3 files

models/diffusion.py

@@ -1,11 +1,11 @@
 import itertools
 import math
 import torch
-import torch.nn.functional as F
+import numpy as np
 import pytorch_lightning as L
 import torchmetrics
 from dataclasses import dataclass
-from models import dit, ema
+import dit, ema
 import noise_schedule  # Assuming this is part of the MDLM repository
 
 LOG2 = math.log(2)
@@ -22,7 +22,6 @@ class NLL(torchmetrics.MeanMetric):
 class BPD(NLL):
     def compute(self) -> torch.Tensor:
         """Computes the bits per dimension.
-
         Returns:
           bpd
         """
@@ -31,21 +30,24 @@ class BPD(NLL):
 class Perplexity(NLL):
     def compute(self) -> torch.Tensor:
         """Computes the Perplexity.
-
         Returns:
          Perplexity
         """
         return torch.exp(self.mean_value / self.weight)
 
+# Based on MDLM repo
 class Diffusion(L.LightningModule):
-    def __init__(self, config, latent_dim):
+    def __init__(self, config, latent_dim, tokenizer):
         super().__init__()
         self.config = config
         self.latent_dim = latent_dim
+        self.tokenizer = tokenizer
 
-        self.backbone = dit.DIT(config, vocab_size=self.latent_dim)
+        self.backbone = dit.DIT(self.config, vocab_size=self.latent_dim)
         self.T = self.config.T
-        self.subs_masking = self.config.subs_masking
+        self.subs_masking = self.config.SUBS_MASKING
+        self.antithetic_sampling = self.config.Training.ANTITHETIC_SAMPLING
+        self.mask_index = self.tokenizer.mask_token_id
 
         self.softplus = torch.nn.Softplus()
         metrics = torchmetrics.MetricCollection({
@@ -59,30 +61,252 @@ class Diffusion(L.LightningModule):
         self.test_metrics = metrics.clone(prefix='test/')
 
         self.noise = noise_schedule.get_noise(self.config, dtype=self.dtype)
-        self.lr = self.config.optim["lr"]
-        self.sampling_eps = self.config.training.get("sampling_eps", 1e-5)
-        self.time_conditioning = self.config.get("time_conditioning", True)
+        self.lr = self.config.Optim.LR
+        self.sampling_eps = self.config.Training.SAMPLING_EPS
+        self.time_conditioning = self.config.TIME_CONDITIONING
         self.neg_infinity = -1000000.0
 
+
+    ############ FORWARD DIFFUSION #########
+    def subs_parameterization(self, logits, noised_latents):
+        # log prob at the mask index = - infinity
+        logits[:, :, self.mask_index] += self.neg_infinity
+        
+        # Normalize the logits such that x.exp() is
+        # a probability distribution over vocab_size.
+        logits = logits - torch.logsumexp(logits, dim=-1,
+                                        keepdim=True)
+
+        # Apply updates directly in the logits matrix.
+        # For the logits of the unmasked tokens, set all values
+        # to -infinity except for the indices corresponding to
+        # the unmasked tokens.
+        unmasked_indices = (noised_latents != self.mask_index)
+        logits[unmasked_indices] = self.neg_infinity
+        logits[unmasked_indices, noised_latents[unmasked_indices]] = 0
+        return logits
+
     def forward(self, latents, sigma):
+        latents = latents.long()
+        with torch.cuda.amp.autocast(dtype=torch.float32):
+            logits = self.backbone(latents, sigma)
+        print(logits)
+        optimized_logits = self.subs_parameterization(logits, latents)
+        return optimized_logits
+    
+    def q_xt(self, latents, move_chance):
+        """
+        Computes the noisy sample xt.
+        Args:
+            x: int torch.Tensor with shape (batch_size, diffusion_model_input_length), input. 
+            move_chance: float torch.Tensor with shape (batch_size, 1).
+        """
+        latents = latents.mean(dim=1) # [bsz x seq_len x 1280] --> [bsz x 1280] as per args
+        move_indices = torch.rand(* latents.shape, device=latents.device) < move_chance
+        noised_latents = torch.where(move_indices, self.mask_index, latents)
+        return noised_latents
+
+    def sample_timestep(self, n, device):
+        _eps_t = torch.rand(n, device=device)
+        if self.antithetic_sampling:
+            offset = torch.arange(n, device=device) / n
+            _eps_t = (_eps_t / n + offset) % 1
+        t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps
+        # if self.importance_sampling:
+        #     return self.noise.importance_sampling_transformation(t)
+        return t
+
+
+    def d3pm_loss(self, model_output, xt, x0, t):
+        """Computes the D3PM loss between noisy latents and the original input at a given time step."""
+        dt = 1 / self.T
+
+        if torch.is_tensor(t):
+            t = t[:, None]
+            assert t.ndim == 2
+            t = t.clamp(0., 1. - 1e-4)
+        alpha_t = 1 - t + torch.zeros_like(xt)
+        alpha_s = 1 - (t - dt) + torch.zeros_like(xt)
+
+        x0 = x0.to(torch.int64)
+        log_x_theta_at_x0 = torch.gather(model_output, -1, x0[:, :, None]).squeeze(-1)
+        log_x_theta_at_m = model_output[:, :, self.mask_index]
+        x_theta_at_m = log_x_theta_at_m.exp()
+        
+        term_1_coef = dt / t
+        term_1_log_nr = torch.log(alpha_t * x_theta_at_m / t + 1)
+        term_1_log_dr = log_x_theta_at_x0
+        
+        term_2_coef = 1 - dt / t
+        term_2_log_nr = term_1_log_nr
+        term_2_log_dr = torch.log(alpha_s * x_theta_at_m / (t - dt) + 1)
+
+        L_vb_masked = (
+            term_1_coef * (term_1_log_nr - term_1_log_dr)
+            + term_2_coef * (term_2_log_nr - term_2_log_dr))
+
+        L_vb = L_vb_masked * (xt == self.mask_index)
+
+        return self.T * L_vb
+
+    def forward_diffusion(self, latents):
         """Forward diffusion process, adds noise to the latents."""
-        noise = sigma * torch.randn_like(latents)
-        noisy_latents = latents + noise
-        return noisy_latents
 
-    def reverse_diffusion(self, noisy_latents, sigma):
-        """Reverse diffusion process, denoises the latents."""
-        denoised_latents = self.backbone(noisy_latents, sigma)
-        return denoised_latents
+        t = self.sample_timestep(latents.shape[0], latents.device)
+        if self.T > 0:
+            t = (t * self.T).to(torch.int)
+            t = t / self.T
+            # t \in {1/T, 2/T, ..., 1}
+            t += (1 / self.T)
+
+        sigma, dsigma = self.noise(t)
+        unet_conditioning = sigma[:, None]
+        move_chance = 1 - torch.exp(-sigma[:, None])
+        
+        noised_latents = self.q_xt(latents, move_chance)
+        model_output = self.forward(noised_latents, unet_conditioning)
+    
+        if self.T > 0:
+            diffusion_loss = self.d3pm_loss(model_output=model_output, xt=noised_latents, x0=latents, t=t)
+            return diffusion_loss
+        # SUBS parameterization, continuous time.
+        else:
+            log_p_theta = torch.gather(input=model_output, dim=-1, index=latents[:, :, None]).squeeze(-1)
+            return - log_p_theta * (dsigma / torch.expm1(sigma))[:, None]
+    
 
+    ######### LOSS CALCULATIONS #########
+    def maybe_sub_sample(self, x0, attention_mask):
+        # seqlen = x0.shape[1]
+        # print(seqlen)
+        # if seqlen > self.config.model.length:
+        #     assert seqlen == 2 * self.config.model.length
+        #     # cropping is needed for text8-crop dataset
+        #     # try the same starting point for now
+        #     start = np.random.choice(self.config.model.length)
+        #     end = start + self.config.model.length
+        #     input_tokens = x0[:, start: end]
+        #     output_tokens = x0[:, start + 1: end + 1]
+        #     new_attention_mask = attention_mask[:, start: end]
+
+        #     # Helps with validation PPL, since the val
+        #     # examples will all start and end with BOS/EOS
+        #     input_tokens[:, 0] = self.tokenizer.bos_token_id
+        #     output_tokens[:, -1] = self.tokenizer.eos_token_id
+    
+        # elif self.parameterization == 'ar':
+        #     input_tokens = x0[:, :-1]
+        #     output_tokens = x0[:, 1:]
+        #     new_attention_mask = attention_mask[:, 1:]
+        # else:
+        input_tokens = x0
+        output_tokens = None
+        new_attention_mask = attention_mask
+    
+        return input_tokens, output_tokens, new_attention_mask
+
+    def compute_loss(self, latents, attention_mask):
+        """"Average of MLM losses to stabilize training"""
+        (input_tokens, output_tokens, attention_mask) = self.maybe_sub_sample(latents, attention_mask)
+        loss = self.forward_diffusion(input_tokens)
+
+        nlls = loss * attention_mask
+        count = attention_mask.sum()
+        batch_nll = nlls.sum()
+        token_nll = batch_nll / count
+
+        return Loss(loss=token_nll, nlls=nlls, token_mask=attention_mask)
+
+
+    ######### TRAINING #########
     def training_step(self, batch, batch_idx):
-        sigma = torch.rand(batch.size(0), device=self.device)
-        noisy_latents = self.forward(batch, sigma)
-        denoised_latents = self.reverse_diffusion(noisy_latents, sigma)
-        loss = F.mse_loss(denoised_latents, batch)
-        self.log("train_loss", loss)
+        latents, attention_mask = batch
+        loss = self.compute_loss(latents, attention_mask)
         return loss
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
         return optimizer
+
+    def validation_step(self, batch):
+        latents, attention_mask = batch
+        loss = self.compute_loss(latents, attention_mask)
+        return loss
+    
+
+    ######### GENERATION #########
+    def sample_prior(self, *batch_dims):
+        return self.mask_index * torch.ones(* batch_dims, dtype=torch.int64)
+
+    def sample_categorical(categorical_probs):
+        gumbel_norm = (1e-10 - (torch.rand_like(categorical_probs) + 1e-10).log())
+        return (categorical_probs / gumbel_norm).argmax(dim=-1)
+
+    def ddpm_caching_update(self, x, t, dt, p_x0=None):
+        assert self.config.noise.type == 'loglinear'
+        sigma_t, _ = self.noise(t)
+        if t.ndim > 1:
+            t = t.squeeze(-1)
+        assert t.ndim == 1
+        move_chance_t = t[:, None, None]
+        move_chance_s = (t - dt)[:, None, None]
+        assert move_chance_t.ndim == 3, move_chance_t.shape
+        if p_x0 is None:
+            p_x0 = self.forward(x, sigma_t).exp()
+    
+        assert move_chance_t.ndim == p_x0.ndim
+        q_xs = p_x0 * (move_chance_t - move_chance_s)
+        q_xs[:, :, self.mask_index] = move_chance_s[:, :, 0]
+        _x = self.sample_categorical(q_xs)
+        
+        copy_flag = (x != self.mask_index).to(x.dtype)
+        return p_x0, copy_flag * x + (1 - copy_flag) * _x
+
+
+    @torch.no_grad()
+    def sample_subs_guidance(self, n_samples, stride_length, num_strides, dt=0.001):
+        ones = torch.ones(n_samples, dtype=self.dtype,device=self.device)
+        num_steps = int(1 / dt)
+        sampling_steps = 0
+        intermediate_tokens = []
+        target = None
+
+        for _ in range(num_strides + 1):
+            p_x0_cache = None
+            x = self._sample_prior(n_samples,self.config.model.length).to(self.device)
+            
+            if target is not None:
+                x[:, : -stride_length] = target
+            
+            for i in range(num_steps + 1):
+                p_x0_cache, x_next = self.ddpm_caching_update(x=x, t=(1 - i * dt) * ones, dt=dt, p_x0=p_x0_cache)
+                if (not torch.allclose(x_next, x) or self.time_conditioning):
+                    p_x0_cache = None
+                    sampling_steps += 1
+                x = x_next
+            x = self.forward(x, 0 * ones).argmax(dim=-1)
+            intermediate_tokens.append(x[:, :stride_length].cpu().numpy())
+            target = x[:, stride_length:]
+    
+        intermediate_tokens.append(target.cpu().numpy())
+        intermediate_text_samples = []
+        sequence_lengths = ((np.concatenate(intermediate_tokens, axis=1)[:, 1:]
+                                 == self.tokenizer.eos_token_id).cumsum(-1) == 0).sum(-1)
+    
+        for i in range(2, len(intermediate_tokens) + 1):
+            intermediate_text_samples.append(self.tokenizer.decode(np.concatenate(intermediate_tokens[:i], axis=1)))
+        
+        return (sampling_steps, intermediate_text_samples,
+            sequence_lengths)
+
+    def restore_model_and_semi_ar_sample(self, stride_length, num_strides, dt=0.001):
+        """Generate samples from the model."""
+        # Lightning auto-casting is not working in this method for some reason
+        self.backbone.eval()
+        self.noise.eval()
+    
+        (sampling_steps, samples, sequence_lengths) = self.sample_subs_guidance(n_samples=self.config.Loader.BATCH_SIZE,stride_length=stride_length,num_strides=num_strides,dt=dt)
+
+        self.backbone.train()
+        self.noise.train()
+        return sampling_steps, samples, sequence_lengths

models/dit.py

@@ -246,8 +246,7 @@ class DDiTBlock(nn.Module):
 
     bias_dropout_scale_fn = self._get_bias_dropout_scale()
 
-    (shift_msa, scale_msa, gate_msa, shift_mlp,
-     scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None].chunk(6, dim=2)
+    (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp) = self.adaLN_modulation(c)[:, None][0].chunk(6, dim=2)
 
     # attention operation
     x_skip = x
@@ -315,7 +314,7 @@ class DDitFinalLayer(nn.Module):
 
 
   def forward(self, x, c):
-    shift, scale = self.adaLN_modulation(c)[:, None].chunk(2, dim=2)
+    shift, scale = self.adaLN_modulation(c)[:, None][0].chunk(2, dim=2)
     x = modulate_fused(self.norm_final(x), shift, scale)
     x = self.linear(x)
     return x
@@ -348,7 +347,7 @@ class DIT(nn.Module, huggingface_hub.PyTorchModelHubMixin):
       config.model.hidden_size,
       vocab_size,
       config.model.cond_dim)
-    self.scale_by_sigma = config.model.scale_by_sigma
+    #self.scale_by_sigma = config.model.scale_by_sigma
 
   def _get_bias_dropout_scale(self):
     if self.training: