import os import collections import copy import pickle import time from dataclasses import dataclass import fsspec import numpy as np import torch import torch.nn.functional as F import transformers import math import trainer_base import utils from samplers import _get_esolm_sampler, safe_probability class AR(trainer_base.TrainerBase): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self.save_hyperparameters() self._validate_configuration() def _validate_configuration(self): super()._validate_configuration() assert not self.config.algo.time_conditioning assert self.config.prior.type == 'none' def _process_model_input(self, x0, valid_tokens): input_tokens = x0[:, :-1] output_tokens = x0[:, 1:] valid_tokens = valid_tokens[:, 1:] return input_tokens, output_tokens, valid_tokens def nll(self, input_tokens, output_tokens, current_accumulation_step, train_mode): del train_mode, current_accumulation_step dummy_t0 = torch.zeros(input_tokens.shape[0], dtype=self.dtype, device=self.device) output = self.backbone(input_tokens, dummy_t0) output[:, :, self.mask_index] = self.neg_infinity output = output.log_softmax(-1) return - output.gather( -1, output_tokens[:, :, None])[:, :, 0] def _process_sigma(self, sigma): del sigma return None @torch.no_grad() def augment_batch_cfg(self, batch, prompt_index, sigma=None, sort_idx=None): raise NotImplementedError("AR does not support CFG") class MDLM(trainer_base.AbsorbingState): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self._validate_configuration() def _process_model_output(self, model_output, xt, sigma): del xt, sigma # zero-masking probabilities model_output[:, :, self.mask_index] = self.neg_infinity # Normalize the model_output such that x.exp() is # a probability distribution over vocab_size. model_output = model_output.log_softmax(-1) return model_output def nll_per_token(self, log_x_theta, xt, x0, alpha_t, dalpha_t, low_var=False): log_p_theta = log_x_theta.gather( dim=-1, index=x0[:, :, None])[:, :, 0] # carry-over unmasking loss_mask = xt == self.mask_index log_p_theta = log_p_theta * loss_mask if low_var: return -log_p_theta else: return dalpha_t / (1 - alpha_t) * log_p_theta @torch.no_grad() def forward_sample(self, x, sigma, latent=None, attn_mask=None, prompt_index=None): return self.forward(x, sigma, latent=latent, attn_mask=attn_mask, prompt_index=prompt_index) class EsoLM(MDLM): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self.alpha_0 = config.algo.alpha_0 self.noise = trainer_base.LogLinear(self.alpha_0) def _sort_indices( self, indices, shuffle, keep_masks_unshuffled=False): masked = (indices == self.mask_index) if shuffle: offsets = torch.rand( indices.shape).to(indices.device) * 0.9 if keep_masks_unshuffled: # induce left-to-right order within masked tokens # only for sequential part offsets[masked] = torch.linspace( 0, 1, torch.sum(masked)).to(indices.device) else: offsets = torch.linspace( 0, 0.9, indices.shape[1]).to(indices.device) sort_idx = (masked + offsets).argsort(descending=False) return sort_idx def _loss(self, x0, valid_tokens, current_accumulation_step=None, train_mode=False, dynamic=False): batch_size = x0.shape[0] # batch size used for diffusion loss split_batch = int( self.config.algo.batch_split * batch_size) x0_reconstruction = x0[split_batch:] x0_diffusion = x0[:split_batch] valid_tokens_reconstruction = valid_tokens[split_batch:] valid_tokens_diffusion = valid_tokens[:split_batch] num_recons = valid_tokens_reconstruction.sum() num_diffusion = valid_tokens_diffusion.sum() do_sequential = self.config.algo.alpha_0 != 1 do_diffusion = self.config.algo.alpha_0 != 0 if do_sequential: assert num_recons > 0 alpha_start = self.config.algo.alpha_0 z0 = self.q_xt(x0_reconstruction, alpha_start) reconstruction_loss, sort_idx = ( self._reconstruction_loss(x0_reconstruction, z0, dynamic=dynamic)) valid_tokens_reconstruction = torch.gather( valid_tokens_reconstruction, dim=1, index=sort_idx) reconstruction_loss = ( reconstruction_loss * valid_tokens_reconstruction).sum() # artificially scale the reconstruction loss so that the # NLL is computed correctly. recons_loss_per_token = reconstruction_loss / num_recons else: recons_loss_per_token = torch.tensor( [0.0]).to(x0.device) if do_diffusion: assert num_diffusion > 0 diffusion_loss, sort_idx = self.nll( x0_diffusion, None, current_accumulation_step, train_mode, dynamic=dynamic) valid_tokens_diffusion = torch.gather( valid_tokens_diffusion, dim=1, index=sort_idx) diffusion_loss = ( diffusion_loss * valid_tokens_diffusion).sum() diffusion_loss_per_token = diffusion_loss / num_diffusion else: diffusion_loss_per_token = torch.tensor( [0.0]).to(x0.device) loss_per_token = (recons_loss_per_token + diffusion_loss_per_token) if num_recons == 0: num_tokens = num_diffusion elif num_diffusion == 0: num_tokens = num_recons else: num_tokens = num_diffusion return trainer_base.Loss( loss=loss_per_token, nlls=loss_per_token * num_tokens, reconstruction_loss=recons_loss_per_token * num_tokens, num_tokens=num_tokens) def _reconstruction_loss(self, x0, z0, dynamic=False): dummy_t0 = torch.zeros(1, z0.shape[0], dtype=self.dtype, device=self.device) # sort inputs and targets before passing to the model sort_idx = self._sort_indices( z0, shuffle=self.config.algo.sequential_shuffle, keep_masks_unshuffled=True) z0 = torch.gather(z0, dim=1, index=sort_idx) x0 = torch.gather(x0, dim=1, index=sort_idx) # pass sort_idx into the model to also sort pos. embeddings # _process_model_output performs zero-masking trick model_output_t0 = self.forward( z0, dummy_t0, sort_idx, x0=x0, dynamic=dynamic) reconstruction_loss = - torch.gather( input=model_output_t0, dim=-1, index=x0[:, :, None]).squeeze(-1) # carry-over loss masking loss_mask = z0 == self.mask_index reconstruction_loss = reconstruction_loss * loss_mask return reconstruction_loss, sort_idx def nll(self, x0, output_tokens, current_accumulation_step=None, train_mode=False, dynamic=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) assert alpha_t.ndim == 2 sigma = self._sigma_from_alphat(alpha_t) xt = self.q_xt(x0, alpha_t) # sort inputs and targets before passing to the model sort_idx = self._sort_indices( xt, shuffle=self.config.algo.diffusion_shuffle) xt = torch.gather(xt, dim=1, index=sort_idx) x0 = torch.gather(x0, dim=1, index=sort_idx) # pass sort_idx into the model to also sort pos. embeddings # _process_model_output performs zero-masking trick log_x_theta = self.forward( xt, sigma=sigma, sort_idx=sort_idx, dynamic=dynamic) # nll_per_token performs carry-over loss masking 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'), sort_idx def _sample_t(self, n, accum_step): if accum_step is not None: # During training batch_dim = n n = int(self.config.loader.global_batch_size * self.config.algo.batch_split) _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 @torch.no_grad() def forward_sample(self, x, sort_idx, attn_mode, cutoffs, kv_cache, last_k_start, unmasked_tokens, k, prompt_index): x, _, sort_idx = self.augment_batch_cfg(x, prompt_index, sigma=None, sort_idx=sort_idx) logits = self.backbone.forward_sample( zt=x, # shape[1] is model.length sort_idx=sort_idx, # shape[1] is model.length attn_mode=attn_mode, cutoffs=cutoffs, kv_cache=self.kv_cache, last_k_start=last_k_start, curr_k_start=unmasked_tokens, # also last_k_end curr_k_end=unmasked_tokens + k, prompt_index=self.condition_mask) logits = self.reduce_batch_cfg(x, logits) logits = self._process_model_output(model_output=logits, xt=x, sigma=None) return logits @torch.no_grad() def generate_samples(self, num_samples, eps=1e-5, condition=None): """Generate samples from the model.""" # Lightning auto-casting is not working in this method for some reason sampler = _get_esolm_sampler( self.config, self, self.tokenizer) return sampler( num_samples=num_samples, eps=eps, condition=condition) class DUO_BASE(trainer_base.UniformState): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self._validate_configuration() def on_save_checkpoint(self, checkpoint): checkpoint['state_dict'] = collections.OrderedDict( (k, v) for k, v in checkpoint['state_dict'].items() if not k.startswith('teacher')) super().on_save_checkpoint(checkpoint) def on_load_checkpoint(self, checkpoint): checkpoint['state_dict'] = collections.OrderedDict( (k, v) for k, v in checkpoint['state_dict'].items() if not k.startswith('teacher')) super().on_load_checkpoint(checkpoint) def _process_model_output(self, model_output, xt, sigma): del xt, sigma model_output[:, :, self.mask_index] = self.neg_infinity return model_output.log_softmax(dim=-1) def nll_per_token(self, log_x_theta, xt, x0, alpha_t, dalpha_t, low_var=False): if log_x_theta.shape[1] == self.num_tokens: return self._nll_per_token_torch_compile( log_x_theta, xt, x0, alpha_t, dalpha_t, low_var) return self._nll_per_token(log_x_theta, xt, x0, alpha_t, dalpha_t, low_var) @torch.compile def _nll_per_token_torch_compile(self, *args, **kwargs): return self._nll_per_token(*args, **kwargs) def _nll_per_token(self, log_x_theta, xt, x0, alpha_t, dalpha_t, low_var=False): del low_var assert alpha_t.ndim == 2 assert x0.ndim == 2 assert xt.ndim == 2 assert not torch.is_tensor(dalpha_t) or dalpha_t.ndim == 2 x_reconst = log_x_theta.exp() x_bar_theta = self.vocab_size * alpha_t[ :, :, None] * x_reconst + 1 - alpha_t[:, :, None] coeff = dalpha_t / (self.vocab_size * alpha_t) x_eq_xt = (x0 == xt).float() x_neq_xt = 1 - x_eq_xt xbar_xt = (1 - alpha_t) + self.vocab_size * alpha_t * x_eq_xt xbar_theta_xt = torch.gather( x_bar_theta, -1, xt.unsqueeze(-1)).squeeze(-1) xbar_theta_x = torch.gather( x_bar_theta, -1, x0.unsqueeze(-1)).squeeze(-1) term1 = self.vocab_size * (1 / xbar_xt - 1 / xbar_theta_xt) const = (1 - alpha_t) / (self.vocab_size * alpha_t + 1 - alpha_t) term2_coefs = x_eq_xt * const + x_neq_xt term2_offset = ((self.vocab_size - 1) * const * x_eq_xt - (1 / const) * x_neq_xt) * const.log() term2_theta = - term2_coefs * ( x_bar_theta.log().sum(-1) - self.vocab_size * xbar_theta_xt.log()) term2_theta = ( term2_theta - self.vocab_size * alpha_t / (1 - alpha_t) * ( xbar_theta_x.log() - xbar_theta_xt.log()) * x_neq_xt) term2 = term2_theta + term2_offset diffusion_loss = coeff * (term1 - term2) assert diffusion_loss.ndim == 2 return diffusion_loss @torch.no_grad() def forward_sample(self, x, sigma, latent=None, prompt_index=None): return self.forward(x, sigma, prompt_index=prompt_index) class LLaDaWrapper(MDLM): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self._validate_configuration() self.mask_index = 126336 def _validate_configuration(self): super()._validate_configuration() assert self.config.algo.backbone == "llada_dit" @torch.no_grad() def get_logits(self, batch, prompt_index): if self.config.sampling.cfg > 0.: assert len(prompt_index) == batch.shape[1] prompt_index = prompt_index.unsqueeze( 0).repeat(batch.shape[0], 1) un_batch = batch.clone() un_batch[prompt_index] = self.mask_index batch = torch.cat([batch, un_batch]) logits = self.backbone(batch).logits if self.config.sampling.cfg > 0.: logits, un_logits = torch.chunk(logits, 2, dim=0) logits = un_logits + \ (self.config.sampling.cfg + 1) * \ (logits - un_logits) return logits[:, :batch.shape[1]] def forward(self, xt, sigma=None, sort_idx=None, x0=None): # with torch.amp.autocast('cuda', dtype=torch.float32): ### Fails to reproduce results, because they keep the bfloat16 precision at inference # (probably matching the training precision is super important??) model_output = self.get_logits( xt, prompt_index=xt.shape[1]) #### Don't use that: log_softmax does approximation errors (because we use bfloat16 probably) # model_output = self._process_model_output( # model_output=model_output, xt=xt, sigma=sigma) return model_output @torch.no_grad() def forward_sample(self, x, sigma=None, latent=None): return self.forward(x, sigma) @dataclass class VDLMLoss: loss: torch.FloatTensor nlls: torch.FloatTensor reconstruction_loss: torch.FloatTensor auxiliary_loss: dict surrogate_posterior_loss: torch.FloatTensor num_tokens: torch.FloatTensor def freeze_parameters(module: torch.nn.Module): if module is not None: for param in module.parameters(): param.requires_grad = False def unfreeze_parameters(module: torch.nn.Module): if module is not None: for param in module.parameters(): param.requires_grad = True class VDLMBase(MDLM): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self._validate_configuration() self.training_mode = self.config.algo.training_mode if self.training_mode == "diffusion_training": print("Freezing surrogate posterior parameters") freeze_parameters(self.backbone.surrogate_posterior) elif self.training_mode == "posterior_training": print("Freezing encoder and decoder parameters") freeze_parameters(self.backbone.encoder) freeze_parameters(self.backbone.decoder) unfreeze_parameters(self.backbone.surrogate_posterior) else: raise ValueError(f"Invalid training mode: {self.training_mode}") def forward(self, xt, sigma, sort_idx=None, x0=None, training_mode="diffusion"): raise NotImplementedError("Not implemented") @torch.no_grad() def forward_sample(self, x, sigma, latent): return self.backbone(x=x, sigma=sigma, latent=latent, mode="sampling") def _auxiliary_loss(self, latent_mean, latent_logvar, latent): raise NotImplementedError("Not implemented") def _loss(self, x0, valid_tokens, current_accumulation_step=None, train_mode=False, mode="diffusion_training"): (input_tokens, output_tokens, valid_tokens) = self._process_model_input( x0, valid_tokens) if mode == "diffusion_training": loss, nlls, auxiliary_loss = self.diffusion_loss( input_tokens, output_tokens, valid_tokens, current_accumulation_step, train_mode) surrogate_posterior_loss = torch.tensor(0) elif mode == "posterior_training": loss = self.surrogate_posterior_loss( x0, output_tokens, valid_tokens, current_accumulation_step, train_mode) surrogate_posterior_loss = loss nlls, auxiliary_loss = torch.tensor( 0), torch.tensor(0) else: raise ValueError( f"Invalid training phase: {training_phase}") return VDLMLoss( loss=loss, nlls=nlls, auxiliary_loss=auxiliary_loss, reconstruction_loss=torch.tensor(0), surrogate_posterior_loss=surrogate_posterior_loss, num_tokens=valid_tokens.sum()) def sample_latent(self, xt): if self.backbone.surrogate_posterior is not None: return self.sample_latent_posterior(xt) else: return self.sample_latent_prior(xt) def sample_latent_prior(self, xt): raise NotImplementedError("Not implemented") def sample_latent_posterior(self, xt): raise NotImplementedError("Not implemented") def diffusion_loss(self, x0, output_tokens, valid_tokens, current_accumulation_step, train_mode): raise NotImplementedError("Not implemented") def surrogate_posterior_loss(self, x0, output_tokens, valid_tokens, current_accumulation_step, train_mode): raise NotImplementedError("Not implemented") @dataclass class LossParams: name: str weight: float params: dict class VDLMContinuousLatent(VDLMBase): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self._validate_configuration() self.auxiliary_loss_rampup = self.config.algo.auxiliary_loss.rampup self.auxiliary_loss_startup = self.config.algo.auxiliary_loss.startup self.auxiliary_loss_annealing_strategy = self.config.algo.auxiliary_loss.annealing_strategy self.auxiliary_losses = {} self.auxiliary_loss_weights = {} for k, v in self.config.algo.auxiliary_loss.items(): if "loss_" in k: self.auxiliary_losses[v["name"]] = LossParams(name=v["name"], weight=v["weight"], params=v["params"]) def _validate_configuration(self): super()._validate_configuration() assert self.config.algo.latent_type == "continuous" assert self.config.algo.surrogate_posterior.name in [ "dm", "none"] # assert self.config.algo.backbone == "vdlm_dit" def _set_loss_weight(self, training_step): if self.auxiliary_loss_annealing_strategy == "constant": self.auxiliary_loss_weights = {k: v.weight for k, v in self.auxiliary_losses.items()} elif self.auxiliary_loss_annealing_strategy == "linear": self.auxiliary_loss_weights = {k: np.where(training_step < self.auxiliary_loss_startup, 0., np.where(training_step < self.auxiliary_loss_rampup and training_step > self.auxiliary_loss_startup, v.weight * (training_step - self.auxiliary_loss_startup) / (self.auxiliary_loss_rampup - self.auxiliary_loss_startup), v.weight )).item() for k, v in self.auxiliary_losses.items()} elif self.auxiliary_loss_annealing_strategy == "cosine": self.auxiliary_loss_weights = {k: np.where(training_step < self.auxiliary_loss_startup, 0., np.where(training_step < self.auxiliary_loss_rampup and training_step > self.auxiliary_loss_startup, v.weight * (1 + np.cos(np.pi * (1 - (training_step - self.auxiliary_loss_startup) / (self.auxiliary_loss_rampup - self.auxiliary_loss_startup)))) / 2, v.weight )).item() for k, v in self.auxiliary_losses.items()} else: raise ValueError(f"Invalid auxiliary loss weight annealing strategy: {self.auxiliary_loss_annealing_strategy}") def empirical_kernel_estimator(self, x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: x1 = x1.unsqueeze(-2) # Make it into a column tensor x2 = x2.unsqueeze(-3) # Make it into a row tensor sigma = 2. * x2.size(-1) * self.auxiliary_losses["mmd"].params["kernel_std"] result = torch.exp(-((x1 - x2).pow(2).mean(-1) / sigma)) return result def _mmd_loss(self, latent): # Sample from prior (Gaussian) distribution prior = torch.randn_like(latent) prior_kernel = self.empirical_kernel_estimator(prior, prior) latent_kernel = self.empirical_kernel_estimator(latent, latent) cross_kernel = self.empirical_kernel_estimator(prior, latent) mmd = prior_kernel.mean() + \ latent_kernel.mean() - \ 2 * cross_kernel.mean() return mmd def _kld_loss(self, latent_mean, latent_logvar): return torch.mean(-0.5 * torch.mean(1 + latent_logvar - latent_mean.pow(2) - latent_logvar.exp(), dim=1), dim=0) def _auxiliary_loss(self, latent_mean, latent_logvar, latent): auxiliary_loss = {"value": torch.tensor(0.).to(self.device)} if self.auxiliary_losses["kld"].weight > 0.: kld_loss = self._kld_loss(latent_mean, latent_logvar) auxiliary_loss["value"] += self.auxiliary_loss_weights["kld"] * kld_loss auxiliary_loss["kld"] = kld_loss if self.auxiliary_losses["mmd"].weight > 0.: mmd_loss = self._mmd_loss(latent) bias_corr = latent.shape[0] * (latent.shape[0] - 1) auxiliary_loss["value"] += self.auxiliary_loss_weights["mmd"] * mmd_loss / bias_corr auxiliary_loss["mmd"] = mmd_loss return auxiliary_loss def training_step(self, batch, batch_idx): self._set_loss_weight(self.global_step) for k, v in self.auxiliary_losses.items(): self.log(name=f'trainer/{k}_weight', value=v.weight, on_step=True, on_epoch=False, sync_dist=True) return super().training_step(batch, batch_idx) def validation_step(self, batch, batch_idx): self._set_loss_weight(self.global_step) return super().validation_step(batch, batch_idx) def forward(self, xt, sigma, sort_idx=None, x0=None, latent=None, mode="diffusion_training"): sigma = self._process_sigma(sigma) if mode == "diffusion_training": with torch.amp.autocast('cuda', dtype=torch.float32): log_x_theta, latent_mean, latent_logvar, latent = self.backbone( x=xt, sigma=sigma, x0=x0, latent=latent, mode="diffusion_training") log_x_theta = self._process_model_output( model_output=log_x_theta, xt=xt, sigma=sigma) return log_x_theta, latent_mean, latent_logvar, latent elif mode == "posterior_training": assert self.config.algo.surrogate_posterior.name != "none", "Surrogate posterior must be used for posterior training" with torch.amp.autocast('cuda', dtype=torch.float32): latent = self.backbone( x=xt, sigma=sigma, x0=x0, latent=None, mode="posterior_training") return surrogate_latent, latent else: raise ValueError( f"Invalid mode: {mode}") def diffusion_loss(self, x0, output_tokens, valid_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) assert alpha_t.ndim == 2 sigma = self._sigma_from_alphat(alpha_t) xt = self.q_xt(x0, alpha_t) log_x_theta, latent_mean, latent_logvar, latent = self.forward( xt=xt, sigma=sigma, x0=x0, latent=None, mode="diffusion_training") # utils.print_nans(log_x_theta, 'model_output') nll_per_token = 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') nlls = (nll_per_token * valid_tokens).sum() num_tokens = valid_tokens.sum() token_nll = nlls / num_tokens auxiliary_loss = self._auxiliary_loss(latent_mean, latent_logvar, latent) loss = token_nll + auxiliary_loss["value"] return loss, nlls, auxiliary_loss def sample_latent_prior(self, xt): return torch.randn(xt.shape[0], self.config.model.latent_dim, device=xt.device) def sample_latent_posterior(self, xt): # Classical ancestral continuous ODE sampling here (based on EDM) num_steps = self.config.sampling.surrogate_posterior.num_steps rho = self.config.sampling.surrogate_posterior.rho sigma_max = self.config.sampling.surrogate_posterior.sigma_max sigma_min = self.config.sampling.surrogate_posterior.sigma_min latent = sigma_max * torch.randn(xt.shape[0], self.config.model.latent_dim, device=xt.device) step_indices = torch.arange( num_steps, dtype=torch.float64, device=xt.device) t_steps = (sigma_max ** (1 / rho) + step_indices / (num_steps - 1) * (sigma_min ** (1 / rho) - sigma_max ** (1 / rho))) ** rho t_steps = torch.cat( [t_steps, torch.zeros_like(t_steps[:1])]) for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])): latent0_est = self.backbone.surrogate_posterior( latent, sigma=t_cur, cond=xt).to(torch.float64) d_cur = (latent - latent0_est) / t_cur latent = latent + (t_next - t_cur) * d_cur return latent def surrogate_posterior_loss(self, x0, output_tokens, valid_tokens, current_accumulation_step=None, train_mode=False): del output_tokens del valid_tokens t = self._sample_t(x0.shape[0], current_accumulation_step) assert t.shape[0] == x0.shape[0] dalpha_t, alpha_t = self.noise(t) alpha_t = alpha_t.unsqueeze(-1) assert alpha_t.ndim == 2 sigma = self._sigma_from_alphat(alpha_t) xt = self.q_xt(x0, alpha_t) latent = self.backbone.forward(xt, sigma, mode="posterior_training") P_std = self.config.algo.surrogate_posterior.P_std P_mean = self.config.algo.surrogate_posterior.P_mean sigma_data = self.config.algo.surrogate_posterior.sigma_data rnd_normal = torch.randn( [xt.shape[0], 1, 1, 1], device=xt.device) latent_sigma = (rnd_normal * P_std + P_mean).exp() weight = (latent_sigma ** 2 + sigma_data ** 2) / (latent_sigma * sigma_data) ** 2 perturbed_latent = latent + \ torch.randn_like(latent) * latent_sigma D_yn = self.backbone.surrogate_posterior( perturbed_latent, latent_sigma, sigma=sigma, cond=xt) loss = (weight * ((D_yn - latent) ** 2)).sum() return loss class VDLMDiscreteLatent(VDLMBase): def __init__(self, config, tokenizer): super().__init__(config, tokenizer) self._validate_configuration() def _validate_configuration(self): super()._validate_configuration() assert self.config.algo.latent_type == "discrete" assert self.config.algo.surrogate_posterior.name in [ "ar", "none"] assert self.config.algo.backbone == "vdlm_vq_dit" def forward(self, xt, sigma, sort_idx=None, x0=None, mode="diffusion_training"): sigma = self._process_sigma(sigma) if mode == "diffusion_training": with torch.amp.autocast('cuda', dtype=torch.float32): log_x_theta, vq_loss = self.backbone( x=xt, sigma=sigma, x0=x0, latent=latent, mode="diffusion_training") log_x_theta = self._process_model_output( model_output=log_x_theta, xt=xt, sigma=sigma) return log_x_theta, vq_loss elif mode == "posterior_training": assert self.config.algo.surrogate_posterior.name != "none", "Surrogate posterior must be used for posterior training" with torch.amp.autocast('cuda', dtype=torch.float32): surrogate_codes, codes = self.backbone( x=xt, sigma=sigma, x0=x0, latent=None, mode="posterior_training") return surrogate_codes, codes else: raise ValueError( f"Invalid training mode: {mode}") def diffusion_loss(self, x0, output_tokens, valid_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) assert alpha_t.ndim == 2 sigma = self._sigma_from_alphat(alpha_t) xt = self.q_xt(x0, alpha_t) log_x_theta, vq_loss = self.forward( xt=xt, sigma=sigma, x0=x0, latent=None, mode="diffusion_training") # utils.print_nans(log_x_theta, 'model_output') nll_per_token = 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') nlls = (nll_per_token * valid_tokens).sum() num_tokens = valid_tokens.sum() token_nll = nlls / num_tokens token_loss = token_nll + \ self.config.algo.vq_weight * vq_loss / num_tokens # Weird but easier to compare to other runs return token_loss, nlls, vq_loss def surrogate_posterior_loss(self, x0, output_tokens, valid_tokens, current_accumulation_step=None, train_mode=False): del output_tokens del valid_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) assert alpha_t.ndim == 2 sigma = self._sigma_from_alphat(alpha_t) xt = self.q_xt(x0, alpha_t) latent = self.forward( xt, sigma, x0, mode="posterior_training") input_tokens = latent[:, :-1] output_tokens = latent[:, 1:] output = self.backbone.surrogate_posterior( input_tokens, cond=xt) output = output.log_softmax(-1) return - output.gather( -1, output_tokens[:, :, None])[:, :, 0] def sample_latent_prior(self, xt): """ Samples from an uniform distribution over the latent vocabulary, which is definitely not the real prior.""" return torch.randint(0, self.config.algo.surrogate_posterior.vocab_size, (xt.shape[0], self.config.model.latent_dim), device=xt.device) def sample_latent_posterior(self, xt): # precompute token buffer num_pred_tokens = self.config.model.latent_dim - 1 num_samples = xt.shape[0] latent = torch.zeros( (num_samples, num_pred_tokens + 1), dtype=torch.long, device=self.device) latent[:, 0] = self.config.algo.surrogate_posterior.bos_token_id # precompute noise noise = (torch.distributions.Gumbel(0, 1) .sample((num_samples, num_pred_tokens, self.config.algo.surrogate_posterior.vocab_size)) .to(self.device)) if self.config.sampling.use_float64: noise = noise.to(torch.float64) for i in range(num_pred_tokens): output = self.backbone.surrogate_posterior( latent[:, :i + 1], cond=xt) output = output.log_softmax(-1) y = (output[:, -1, :] + noise[:, i, :]).argmax(-1) latent[:, i + 1] = y return latent