import itertools from dataclasses import dataclass import hydra.utils import lightning as L import numpy as np import torch import torch.nn.functional as F import transformers import dataloader import metrics import models import utils @dataclass class Loss: loss: torch.FloatTensor nlls: torch.FloatTensor prior_loss: torch.FloatTensor num_tokens: torch.FloatTensor class LogLinear(torch.nn.Module): def __init__(self, eps): super().__init__() self.eps = eps # 1e-3 by default, to be consistent with SEDD: https://github.com/louaaron/Score-Entropy-Discrete-Diffusion/blob/0605786da5ccb5747545e26d66fdf477187598b6/noise_lib.py#L56 def forward(self, t): t = (1 - self.eps) * t alpha_t = 1 - t dalpha_t = - torch.ones_like(alpha_t) * (1 - self.eps) return dalpha_t, alpha_t def get_t_for_alpha(self, alpha_t): return 1 - alpha_t class Cosine(torch.nn.Module): def __init__(self, eps): super().__init__() self.eps = eps self.half_pi = torch.pi / 2 def forward(self, t): t = (1 - self.eps) * t alpha_t = 1 - torch.cos(self.half_pi * (1 - t)) dalpha_t = - torch.sin(self.half_pi * (1 - t)) * self.half_pi return dalpha_t, alpha_t def get_t_for_alpha(self, alpha_t): is_tensor = torch.is_tensor(alpha_t) if not is_tensor: alpha_t = torch.tensor([alpha_t]) t = 1 - 2 / torch.pi * torch.acos(1 - alpha_t) if not is_tensor: t = t.cpu().item() return t 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 _unsqueeze(x, reference): return x.view( * x.shape, * ((1,) * (len(reference.shape) - len(x.shape)))) class TrainerBase(L.LightningModule): def __init__( self, config, tokenizer: transformers.PreTrainedTokenizer, vocab_size=None): super().__init__() self.save_hyperparameters() self.config = config if hasattr(self.config.algo, 'ignore_bos'): self.ignore_bos = config.algo.ignore_bos else: self.ignore_bos = False if hasattr(self.config.algo, 'loss_type'): self.loss_type = config.algo.loss_type self.tokenizer = tokenizer if vocab_size is None: self.vocab_size = len(self.tokenizer) else: self.vocab_size = vocab_size self.sampler = self.config.sampling.predictor self.antithetic_sampling = self.config.training.antithetic_sampling self.parameterization = self.config.algo.parameterization if self.config.algo.backbone == 'dit': self.backbone = models.dit.DIT( self.config, vocab_size=self.vocab_size) elif self.config.algo.backbone == 'unet': self.backbone = models.unet.UNet(self.config, self.vocab_size) elif self.config.algo.backbone == 'dimamba': self.backbone = models.dimamba.DiMamba( self.config, vocab_size=self.vocab_size, pad_token_id=self.tokenizer.pad_token_id) elif self.config.algo.backbone == 'hf_dit': self.backbone = transformers.AutoModelForMaskedLM.from_pretrained( config.eval.checkpoint_path, trust_remote_code=True) self.T = self.config.algo.T self.num_tokens = self.config.model.length self.p_nucleus = self.config.sampling.p_nucleus # Noise Schedule if config.noise.type == 'log-linear': self.noise = LogLinear(config.noise.eps) elif config.noise.type == 'cosine': self.noise = Cosine(config.noise.eps) else: raise ValueError(config.noise.type) # Class-conditional training arguments self.num_classes = config.data.get('num_classes', None) self.class_conditional = self.num_classes is not None self.class_cond_dropout = config.training.class_dropout_p self.metrics = metrics.Metrics( gen_ppl_eval_model_name_or_path=\ self.config.eval.gen_ppl_eval_model_name_or_path, eval_ppl_batch_size=\ self.config.eval.perplexity_batch_size) if self.config.training.ema > 0: self.ema = models.ema.ExponentialMovingAverage( self._get_parameters(), decay=self.config.training.ema) else: self.ema = None self.lr = self.config.optim.lr self.sampling_eps = self.config.training.sampling_eps self.time_conditioning = self.config.algo.time_conditioning self.neg_infinity = -1000000.0 self.fast_forward_epochs = None self.fast_forward_batches = None def _validate_configuration(self): assert self.config.algo.backbone in {'dit', 'hf_dit', 'unet'} if self.config.algo.parameterization == 'ar': assert not self.config.algo.time_conditioning assert self.config.prior.type == 'none' if self.time_conditioning: assert self.config.sampling.predictor != 'ancestral_cache' if self.parameterization in {'score', 'mean'}: assert self.time_conditioning if self.T > 0: assert self.parameterization != 'score' def to(self, *args, **kwargs): self = super().to(*args, **kwargs) self.metrics.to(*args, **kwargs) return self def q_xt(self, x, alpha_t): raise NotImplementedError def _get_parameters(self): return itertools.chain(self.backbone.parameters(), self.noise.parameters()) def _eval_mode(self): if self.ema: self.ema.store(self._get_parameters()) self.ema.copy_to(self._get_parameters()) self.backbone.eval() self.noise.eval() def _train_mode(self): if self.ema: self.ema.restore(self._get_parameters()) self.backbone.train() self.noise.train() def on_load_checkpoint(self, checkpoint): if self.ema: self.ema.load_state_dict(checkpoint['ema']) # Copied from: # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py#L41 self.fast_forward_epochs = checkpoint['loops'][ 'fit_loop']['epoch_progress']['current']['completed'] self.fast_forward_batches = checkpoint['loops'][ 'fit_loop']['epoch_loop.batch_progress'][ 'current']['completed'] def on_save_checkpoint(self, checkpoint): if self.ema: checkpoint['ema'] = self.ema.state_dict() # Copied from: # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/tasks/seq.py # ['epoch_loop.batch_progress']['total']['completed'] # is 1 iteration behind, so we're using the optimizer's progress. checkpoint['loops']['fit_loop'][ 'epoch_loop.batch_progress']['total'][ 'completed'] = checkpoint['loops']['fit_loop'][ 'epoch_loop.automatic_optimization.optim_progress'][ 'optimizer']['step']['total'][ 'completed'] * self.trainer.accumulate_grad_batches checkpoint['loops']['fit_loop'][ 'epoch_loop.batch_progress']['current'][ 'completed'] = checkpoint['loops']['fit_loop'][ 'epoch_loop.automatic_optimization.optim_progress'][ 'optimizer']['step']['current'][ 'completed'] * self.trainer.accumulate_grad_batches # _batches_that_stepped tracks the number of global steps, # not the number of local steps, so we don't multiply with # self.trainer.accumulate_grad_batches here. checkpoint['loops']['fit_loop'][ 'epoch_loop.state_dict'][ '_batches_that_stepped'] = checkpoint['loops']['fit_loop'][ 'epoch_loop.automatic_optimization.optim_progress'][ 'optimizer']['step']['total']['completed'] if 'sampler' not in checkpoint.keys(): checkpoint['sampler'] = {} if hasattr(self.trainer.train_dataloader.sampler, 'state_dict'): sampler_state_dict = self.trainer.\ train_dataloader.sampler.state_dict() checkpoint['sampler'][ 'random_state'] = sampler_state_dict.get( 'random_state', None) else: checkpoint['sampler']['random_state'] = None def on_train_start(self): if self.ema: self.ema.move_shadow_params_to_device(self.device) # Adapted from: # https://github.com/Dao-AILab/flash-attention/blob/main/training/src/datamodules/language_modeling_hf.py distributed = ( self.trainer._accelerator_connector.use_distributed_sampler and self.trainer._accelerator_connector.is_distributed) if distributed: sampler_cls = dataloader.FaultTolerantDistributedSampler else: sampler_cls = dataloader.RandomFaultTolerantSampler updated_dls = [] for dl in self.trainer.fit_loop._combined_loader.flattened: if hasattr(dl.sampler, 'shuffle'): dl_sampler = sampler_cls( dl.dataset, shuffle=dl.sampler.shuffle) else: dl_sampler = sampler_cls(dl.dataset) if (distributed and self.fast_forward_epochs is not None and self.fast_forward_batches is not None): dl_sampler.load_state_dict({ 'epoch': self.fast_forward_epochs, 'counter': (self.fast_forward_batches * self.config.loader.batch_size)}) updated_dls.append( torch.utils.data.DataLoader( dl.dataset, batch_size=self.config.loader.batch_size, num_workers=self.config.loader.num_workers, pin_memory=self.config.loader.pin_memory, sampler=dl_sampler, shuffle=False, persistent_workers=True)) self.trainer.fit_loop._combined_loader.flattened = updated_dls def optimizer_step(self, *args, **kwargs): super().optimizer_step(*args, **kwargs) if self.ema: self.ema.update(self._get_parameters()) def _process_sigma(self, sigma): raise NotImplementedError def _process_model_output(self, model_output, xt, sigma): raise NotImplementedError def forward(self, xt, sigma, labels=None, weights=None, nn_input_idxs=None): if nn_input_idxs is None: nn_input_idxs = xt sigma = self._process_sigma(sigma) with torch.amp.autocast('cuda', dtype=torch.float32): model_output = self.backbone( x=nn_input_idxs, sigma=sigma, class_cond=labels, weights=weights) return self._process_model_output( model_output=model_output, xt=xt, sigma=sigma) def on_train_epoch_start(self): self.metrics.reset() assert self.metrics.train_nlls.nll.mean_value == 0 assert self.metrics.train_nlls.nll.weight == 0 def training_step(self, batch, batch_idx): current_accumulation_step = ( batch_idx % self.trainer.accumulate_grad_batches) losses = self._loss(batch['input_ids'], batch.get('labels', None), batch['attention_mask'], current_accumulation_step, train_mode=True) self.metrics.update_train(losses.nlls, losses.prior_loss, losses.num_tokens) self.log(name='trainer/loss', value=losses.loss.item(), on_step=True, on_epoch=False, sync_dist=True) return losses.loss def on_train_epoch_end(self): for k, v in self.metrics.valid_nlls.items(): self.log(name=k, value=v.compute(), on_step=False, on_epoch=True, sync_dist=True) def on_validation_epoch_start(self): self.metrics.reset() self._eval_mode() assert self.metrics.valid_nlls.nll.mean_value == 0 assert self.metrics.valid_nlls.nll.weight == 0 def validation_step(self, batch, batch_idx): del batch_idx losses = self._loss(batch['input_ids'], batch.get('labels', None), batch['attention_mask']) self.metrics.update_valid(losses.nlls, losses.prior_loss, losses.num_tokens) return losses.loss def on_validation_epoch_end(self): for k, v in self.metrics.valid_nlls.items(): self.log(name=k, value=v.compute(), on_step=False, on_epoch=True, sync_dist=True) if ((self.config.eval.compute_perplexity_on_sanity or not self.trainer.sanity_checking) and self.config.eval.generate_samples): samples, text_samples = None, None for _ in range( self.config.sampling.num_sample_batches): samples = self.generate_samples( num_samples=self.config.loader.eval_batch_size) self.metrics.record_entropy(samples) # Decode the samples to be re-tokenized by eval model text_samples = self.tokenizer.batch_decode(samples) if self.config.eval.compute_generative_perplexity: self.metrics.record_generative_perplexity( text_samples, self.num_tokens, self.device) if text_samples is not None: if self.trainer.global_rank == 0 and hasattr( self.trainer.logger, 'log_table'): # Log the last generated samples text_samples = text_samples[ : self.config.sampling.num_sample_log] self.trainer.logger.log_table( key=f'samples@global_step{self.global_step}', columns=['Generated Samples'], data=[[s] for s in text_samples]) if self.config.eval.compute_generative_perplexity: self.log('val/gen_ppl', self.metrics.gen_ppl.compute(), on_epoch=True, on_step=False, sync_dist=True) self.log('val/sample_entropy', self.metrics.sample_entropy.compute(), on_epoch=True, on_step=False, sync_dist=True) self._train_mode() def configure_optimizers(self): optimizer = torch.optim.AdamW( self._get_parameters(), lr=self.config.optim.lr, betas=(self.config.optim.beta1, self.config.optim.beta2), eps=self.config.optim.eps, weight_decay=self.config.optim.weight_decay) scheduler = hydra.utils.instantiate( self.config.lr_scheduler, optimizer=optimizer) scheduler_dict = {'scheduler': scheduler, 'interval': 'step', 'monitor': 'val/loss', 'name': 'trainer/lr'} return [optimizer], [scheduler_dict] def generate_samples(self, num_samples, num_steps, eps): raise NotImplementedError def restore_model_and_sample(self, num_steps, eps=1e-5): """Generate samples from the model.""" # Lightning auto-casting is not working in this method for some reason self._eval_mode() samples = self.generate_samples( num_samples=self.config.loader.eval_batch_size, num_steps=num_steps, eps=eps) self._train_mode() return samples def _process_model_input(self, x0, valid_tokens): raise NotImplementedError def nll(self, input_tokens, labels, output_tokens, current_accumulation_step=None, train_mode=False): raise NotImplementedError def _loss(self, x0, labels, valid_tokens, current_accumulation_step=None, train_mode=False): (input_tokens, output_tokens, valid_tokens) = self._process_model_input( x0, valid_tokens) loss = self.nll(input_tokens, labels, output_tokens, current_accumulation_step, train_mode) assert loss.ndim == 2 if self.ignore_bos: loss[:, 1:] = loss[:, 1:] valid_tokens[:, 1:] = valid_tokens[:, 1:] nlls = (loss * valid_tokens).sum() num_tokens = valid_tokens.sum() token_nll = nlls / num_tokens return Loss(loss=token_nll, nlls=nlls, prior_loss=0.0, num_tokens=num_tokens) class Diffusion(TrainerBase): def _validate_configuration(self): super()._validate_configuration() assert self.config.sampling.noise_removal in { 'none', 'ancestral', 'greedy'} assert self.loss_type in {'elbo', 'low_var'} if self.config.sampling.noise_removal == 'greedy': assert self.sampler != 'analytic' assert self.parameterization in {'mean', 'subs'} def _process_model_input(self, x0, valid_tokens): return x0, None, valid_tokens def _process_sigma(self, sigma): assert sigma.ndim == 2 sigma = sigma.mean(-1).squeeze() if sigma.ndim == 0: sigma = sigma.unsqueeze(0) if not self.time_conditioning: sigma = torch.zeros_like(sigma) assert sigma.ndim == 1, sigma.shape return sigma def _sample_t(self, n, accum_step): if accum_step is not None: # During training batch_dim = n n = self.config.loader.global_batch_size _eps_t = torch.rand(n, device=self.device) if self.antithetic_sampling: offset = torch.arange(n, device=self.device) / n _eps_t = (_eps_t / n + offset) % 1 t = (1 - self.sampling_eps) * _eps_t + self.sampling_eps if accum_step is not None: t = t.chunk(self.trainer.num_nodes)[self.trainer.node_rank] t = t.chunk(self.trainer.num_devices)[self.trainer.local_rank] t = t.chunk(self.trainer.accumulate_grad_batches)[ accum_step] # corner case for the last datapoint t = t[:batch_dim] return t def _sigma_from_alphat(self, alpha_t): return -torch.log(alpha_t) def _reconstruction_loss(self, x0): t0 = torch.zeros(1, x0.shape[0], dtype=self.dtype, device=self.device) sigma_t0 = self._sigma_from_alphat(self.noise(t0)[1]) model_output_t0 = self.forward(x0, sigma_t0) return - torch.gather(input=model_output_t0, dim=-1, index=x0[:, :, None]).squeeze(-1) def nll_per_token(self, model_output, xt, x0, alpha_t, dalpha_t, low_var): raise NotImplementedError def nll(self, x0, labels, output_tokens, current_accumulation_step=None, train_mode=False): del output_tokens t = self._sample_t(x0.shape[0], current_accumulation_step) assert t.shape[0] == x0.shape[0] 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) dalpha_t, alpha_t = self.noise(t) alpha_t = alpha_t.unsqueeze(-1) dalpha_t = dalpha_t.unsqueeze(-1) assert alpha_t.ndim == 2 assert dalpha_t.ndim == 2 sigma = self._sigma_from_alphat(alpha_t) xt = self.q_xt(x0, alpha_t) # Handle class-conditional training, with class dropout if self.class_conditional: assert labels is not None rand = torch.rand(size=labels.shape, dtype=torch.float32, device=self.device) # num_classes represent the absence of class-conditioning labels = torch.where(rand < self.class_cond_dropout, self.num_classes, labels) else: assert labels is None log_x_theta = self.forward(xt, sigma=sigma, labels=labels) utils.print_nans(log_x_theta, 'model_output') return self.nll_per_token( log_x_theta=log_x_theta, xt=xt, x0=x0, alpha_t=alpha_t, dalpha_t=dalpha_t, low_var=train_mode and self.loss_type == 'low_var') def _get_score(self, **kwargs): del kwargs raise NotImplementedError def _denoiser_update(self, x, t): raise NotImplementedError def _analytic_update(self, x, t, dt): raise NotImplementedError def _posterior_from_x0(self, x0, xt, alpha_s, alpha_t): """From the clean x0, or denoiser predictions, implement the mathematical expression for q(z_s | z_t, x)""" raise NotImplementedError def _forward_process(self, q_x0, alpha_s): """Apply forward noising: q(x_s | x_0), where x_0 is a K-dimensional vector""" raise NotImplementedError def _get_ancestral_posterior(self, xt, sigma, labels, alpha_s, alpha_t, p_x0): """Top-level call from generate_samples. Compute the standard posterior sampling distribution from the noisy sequence xt""" gamma = self.config.sampling.guid_weight # If class cond but gamma is None or 0, same as sampling # from the class unconditional case. if not self.class_conditional or gamma in (None, 0.0): if labels is not None: labels = torch.full_like(labels, self.num_classes) return self._get_posterior_from_xt(xt, sigma, labels, alpha_s, alpha_t, p_x0) elif gamma == 1.0: # Simply sample from the cond. posterior only. assert labels is not None return self._get_posterior_from_xt(xt, sigma, labels, alpha_s, alpha_t, p_x0) else: # Case gamma not in {0, 1}, and class conditional # -> mix conditional and unconditional predictions. return self._get_guided_posterior_from_xt(self, xt, sigma, labels, gamma, alpha_s, alpha_t, p_x0) def _get_posterior_from_xt(self, xt, sigma, labels, alpha_s, alpha_t, p_x0=None): """From xt, compute a single denoiser predictions, cast to correct dtype, and call _posterior_from_x0.""" if p_x0 is None: log_x0_pred = self.forward(xt, sigma, labels) if self.config.sampling.use_float64: log_x0_pred = log_x0_pred.to(torch.float64) if self.p_nucleus < 1: log_x0_pred = utils.top_k_top_p_filtering( log_x0_pred, top_p=self.p_nucleus) p_x0 = log_x0_pred.exp() return p_x0, self._posterior_from_x0(x0=p_x0,xt=xt, alpha_s=alpha_s, alpha_t=alpha_t) def _get_guided_posterior_from_xt(self, xt, sigma, labels, gamma, alpha_s, alpha_t, p_x0=None): """From xt, combine the class cond / uncond predictions of the denoiser. Call _get_posterior_from_xt.""" # unpack the cache if p_x0 is None: p_x0_cond = p_x0_uncond = None else: p_x0_cond, p_x0_uncond = p_x0 p_x0_cond, cond_posterior = self._get_posterior_from_xt( xt, sigma, labels, alpha_s, alpha_t, p_x0_cond) log_cond_posterior = cond_posterior.log() # NOTE: conditioning on self.num_classes represents the # class-unconditional predictions. p_x0_uncond, uncond_posterior = self._get_posterior_from_xt( xt, sigma, torch.full_like(labels, self.num_classes), alpha_s, alpha_t, p_x0_uncond) log_uncond_posterior = uncond_posterior.log() un_normalized_posterior = gamma * log_cond_posterior \ + (1 - gamma) * log_uncond_posterior # Handle cases where the posterior is zero (eg after # nucleus sampling) is_inf_mask = torch.logical_or( log_cond_posterior.isinf(), log_uncond_posterior.isinf()) un_normalized_posterior[is_inf_mask] = self.neg_infinity return ((p_x0_cond, p_x0_uncond), un_normalized_posterior.softmax(-1)) def _ancestral_update(self, x, t, labels, dt, p_x0=None, noise_removal_step=False): _, alpha_t = self.noise(t) if noise_removal_step: alpha_s = torch.ones_like(alpha_t) else: _, alpha_s = self.noise(t - dt) assert alpha_t.ndim == 2 sigma = self._sigma_from_alphat(alpha_t) assert alpha_t.ndim == 2, f'{alpha_t.ndim=}' p_x0, q_xs = self._get_ancestral_posterior(x, sigma, labels, alpha_s, alpha_t, p_x0) return p_x0, sample_categorical(q_xs) def _psi_update(self, x, t, labels, dt, kappa, p_x0=None, noise_removal_step=False): _, alpha_t = self.noise(t) if noise_removal_step: alpha_s = torch.ones_like(alpha_t) else: _, alpha_s = self.noise(t - dt) alpha_0 = torch.ones_like(alpha_t) sigma = self._sigma_from_alphat(alpha_t) # Standard posterior q(x_s | x_t) p_x0, q_xs = self._get_ancestral_posterior( x, sigma, labels, alpha_s, alpha_t, p_x0) # Posterior targeting t=0, reuse predictions p_x0 _, q_x0 = self._get_ancestral_posterior( x, sigma, labels, alpha_0, alpha_t, p_x0) # PC: forward-noise q_x0 back to time s pc_q_xs = self._forward_process(q_x0, alpha_s) q_sample = kappa * q_xs + (1 - kappa) * pc_q_xs return p_x0, sample_categorical(q_sample) def _get_sampling_time_profile(self, eps, num_steps): profile = self.config.sampling.psi.time_profile num_steps += 1 if profile == 'linear' \ or self.config.sampling.predictor != 'psi': # Default: linearly decrease return torch.linspace(1, eps, num_steps) if not profile.startswith('linear-constant-linear'): raise ValueError(profile) c = float(profile.split('-')[3]) if 'inv' in profile: c = self.noise.get_t_for_alpha(c) psi_cfg = self.config.sampling.psi n_hi = round(psi_cfg.high_frac * num_steps) n_mid = round(psi_cfg.middle_frac * num_steps) return torch.cat([ torch.linspace(1, c, n_hi), torch.full((n_mid,), c), torch.linspace(c, eps, num_steps - n_hi - n_mid)]) def _mode_to_psi_kappas(self, mode, timesteps): n = len(timesteps) if mode == 'pure-posterior': return torch.ones(n) if mode == 'pure-pc': return torch.zeros(n) eta = float(mode.split('-')[-1]) if (mode.startswith('constant-') and not mode.startswith('constant-remdm')): return torch.full((n,), eta) # Noise-schedule-dependent modes (ReMDM-like) _, all_alphas = self.noise(timesteps) alpha_t, alpha_s = all_alphas[:-1], all_alphas[1:] eta_t = torch.tensor([eta]) if mode.startswith('max-capped-'): sigma = torch.minimum( eta_t.expand_as(alpha_t), (1 - alpha_s) / alpha_t) sigma = torch.where(alpha_t == 0, eta, sigma) elif mode.startswith('max-rescale-'): sigma_max = torch.minimum( eta_t.expand_as(alpha_t), (1 - alpha_s) / alpha_t) sigma = torch.where(alpha_t > 0, sigma_max, 1) * eta elif mode.startswith('constant-remdm'): sigma = eta_t else: raise ValueError(mode) kappas = torch.clip(1 - sigma / (1 - alpha_s), 0, 1) if len(kappas) > 0: kappas = torch.cat([kappas, torch.ones(1)]) return kappas def _get_kappas(self, timesteps): cfg = self.config.sampling.psi n = len(timesteps) n_hi = round(cfg.high_frac * n) n_mid = round(cfg.middle_frac * n) kappas = torch.cat([ self._mode_to_psi_kappas(cfg.high_mode, timesteps[:n_hi]), self._mode_to_psi_kappas(cfg.middle_mode, timesteps[n_hi:n_hi + n_mid]), self._mode_to_psi_kappas(cfg.low_mode, timesteps[n_hi + n_mid:])]) assert (kappas >= 0).all() and (kappas <= 1).all() assert len(kappas) == n return kappas @torch.no_grad() def generate_samples(self, num_samples, labels=None, num_steps=None, eps=1e-5): """Generate samples from the model.""" # Lightning auto-casting is not working in this method for some reason if num_steps is None: num_steps = self.config.sampling.steps x = self.prior_sample(num_samples, self.num_tokens) use_psi_sampler = self.config.sampling.predictor == 'psi' timesteps = self._get_sampling_time_profile(eps, num_steps) if use_psi_sampler: kappas = self._get_kappas(timesteps).to(self.device) p_x0_cache = None if labels is not None: labels = labels.to(self.device) for i in range(num_steps): t = timesteps[i] * torch.ones( x.shape[0], 1, device=self.device) dt = timesteps[i] - timesteps[i+1] if self.sampler == 'ancestral': _, x = self._ancestral_update( x=x, t=t, labels=labels, dt=dt, p_x0=None) elif self.sampler == 'ancestral_cache': p_x0_cache, x_next = self._ancestral_update( x=x, t=t, labels=labels, dt=dt, p_x0=p_x0_cache) if (not torch.allclose(x_next, x) or self.time_conditioning): # Disable caching p_x0_cache = None x = x_next elif self.sampler == 'psi': _, x = self._psi_update(x=x, t=t, kappa=kappas[i], labels=labels, dt=dt, p_x0=None) elif self.sampler == 'analytic': assert labels is None, 'class-conditional sampling ' \ 'is not implemented with the analytic sampler' x = self._analytic_update(x=x,t=t, dt=dt) else: raise ValueError(self.sampler) t0 = timesteps[-1] * torch.ones(x.shape[0], 1, device=self.device) if self.config.sampling.noise_removal == 'ancestral': if self.sampler == 'analytic': x = self._denoiser_update(x=x, t=t0) else: _, x = self._ancestral_update(x=x, t=t0, labels=labels, dt=None, p_x0=p_x0_cache, noise_removal_step=True) elif self.config.sampling.noise_removal == 'greedy': sigma = self._sigma_from_alphat(self.noise(t0)[1]) x = self.forward(xt=x, sigma=sigma).argmax(dim=-1) return x @torch.no_grad def _semi_ar_sampler( self, n_samples, stride_length, num_strides, dt=0.001): # TODO(subham): Test this method after refactoring. 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.prior_sample(n_samples, self.num_tokens) if target is not None: x[:, : -stride_length] = target for i in range(num_steps + 1): p_x0_cache, x_next = self._ancestral_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.batch_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 # TODO(subham): Test this method after refactoring. self._eval_mode() (sampling_steps, samples, sequence_lengths) = self._semi_ar_sampler( n_samples=self.config.loader.eval_batch_size, stride_length=stride_length, num_strides=num_strides, dt=dt) self._train_mode() return sampling_steps, samples, sequence_lengths class AbsorbingState(Diffusion): def __init__(self, config, tokenizer): # NOTE: Ideally, we should do # vocab_size = len(tokenizer), so that we account # for the special tokens added in dataloader.py. # But we use tokenizer.vocab_size so as to to be # consistent with the prior checkpoints. vocab_size = tokenizer.vocab_size if (not hasattr(tokenizer, 'mask_token') or tokenizer.mask_token is None): self.mask_index = vocab_size vocab_size += 1 else: self.mask_index = tokenizer.mask_token_id self.subs_masking = config.algo.subs_masking super().__init__(config, tokenizer, vocab_size=vocab_size) self.save_hyperparameters() def _validate_configuration(self): super()._validate_configuration() if self.parameterization in {'score', 'mean'}: assert self.time_conditioning assert not (self.parameterization == 'mean' and self.T == 0) if self.T > 0: assert self.parameterization in {'mean', 'subs'} if self.subs_masking: assert self.parameterization == 'mean' def q_xt(self, x, alpha_t): """Computes the noisy sample xt. Args: x: int torch.Tensor with shape (batch_size, diffusion_model_input_length), input. alpha_t: float torch.Tensor with shape (batch_size, 1). """ move_indices = torch.rand( * x.shape, device=x.device) < 1 - alpha_t xt = torch.where(move_indices, self.mask_index, x) if self.ignore_bos: xt[:, 0] = x[:, 0] return xt def prior_sample(self, *batch_dims): return self.mask_index * torch.ones( * batch_dims, dtype=torch.int64, device=self.device) def _posterior_from_x0(self, x0, xt, alpha_s, alpha_t): """From the clean x0, or denoiser predictions, implement the mathematical expression for q(z_s | z_t, x)""" assert x0.dtype == torch.float64, 'Requires float64 prec.' # should be one-hot on clean tokens orig_mask = xt[:, :, None] != self.mask_index orig_mask = orig_mask.expand(-1, -1, x0.shape[-1]) orig_output_on_clean = x0[orig_mask] q_xs = ((alpha_s - alpha_t) / (1 - alpha_t))[..., None] * x0 q_xs[..., self.mask_index] = (1 - alpha_s) / (1 - alpha_t) q_xs[orig_mask] = orig_output_on_clean return q_xs def _forward_process(self, x0, alpha_s): out = alpha_s[..., None] * x0 out[..., self.mask_index] = 1 - alpha_s return out def _staggered_score(self, score, dsigma): score = score.clone() extra_const = (1 - dsigma.exp()) * score.sum(dim=-1) score *= dsigma.exp()[:, None] score[..., self.mask_index] += extra_const return score def _analytic_update(self, x, t, dt): sigma_t = self._sigma_from_alphat(self.noise(t)[1]) sigma_s = self._sigma_from_alphat(self.noise(t - dt)[1]) dsigma = sigma_t - sigma_s score = self._get_score(x, sigma_t) if self.config.sampling.use_float64: score = score.to(torch.float64) stag_score = self._staggered_score(score, dsigma) probs = stag_score * self._transp_transition(x, dsigma) return sample_categorical(probs) def _denoiser_update(self, x, t): sigma = self._sigma_from_alphat(self.noise(t)[1]) score = self._get_score(x, sigma) if self.config.sampling.use_float64: score = score.to(torch.float64) stag_score = self._staggered_score(score, sigma) probs = stag_score * self._transp_transition(x, sigma) probs[..., self.mask_index] = 0 samples = sample_categorical(probs) return samples def _transp_transition(self, i, sigma): sigma = _unsqueeze(sigma, reference=i[..., None]) edge = torch.exp(-sigma) * F.one_hot( i, num_classes=self.vocab_size) edge += torch.where(i == self.mask_index, 1 - torch.exp(-sigma).squeeze(-1), 0)[..., None] return edge class UniformState(Diffusion): def _validate_configuration(self): super()._validate_configuration() assert self.time_conditioning assert self.parameterization == 'mean' if self.config.algo.name != 'distillation': assert self.T == 0 def _forward_process(self, x0, alpha_s): return (alpha_s[..., None] * x0 + (1 - alpha_s[..., None]) / self.vocab_size) def q_xt(self, x, alpha_t): """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). """ move_indices = torch.rand( *x.shape, device=x.device) < 1 - alpha_t uniform_tensor = torch.randint( 0, self.vocab_size, x.shape, device=x.device) xt = torch.where(move_indices, uniform_tensor, x) if self.ignore_bos: xt[:, 0] = x[:, 0] return xt def prior_sample(self, *batch_dims): return torch.randint( 0, self.vocab_size, batch_dims, dtype=torch.int64, device=self.device)