sdlm/trainers/trainer_diffusion.py

import math
import time
import warnings
from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, Dataset
from transformers import Trainer

with warnings.catch_warnings():
    warnings.simplefilter(action="ignore", category=FutureWarning)
    from transformers.deepspeed import deepspeed_init

from transformers.integrations import TensorBoardCallback
from transformers.trainer_pt_utils import (
    IterableDatasetShard,
    find_batch_size,
    nested_concat,
    nested_detach,
    nested_numpify,
)
from transformers.trainer_utils import denumpify_detensorize, has_length, speed_metrics
from transformers.utils import (
    is_apex_available,
    is_datasets_available,
    is_sagemaker_mp_enabled,
    logging,
)

from sdlm.inference.inference_utils import (
    logits_projection,
    predict_conditional_generated,
)
from sdlm.models.utils import is_cdcd_check
from sdlm.pipelines.simplex_ddpm import SimplexDDPMClassifierGuidancePipeline
from sdlm.utils import convert_to_simplex, pad_data, scale, self_condition_preds

if is_apex_available():
    from apex import amp

if is_datasets_available():
    import datasets

GENERATION_RESULTS = "generated"


logger = logging.get_logger(__name__)


class EvalLoopOutput(NamedTuple):
    logits: Union[np.ndarray, Tuple[np.ndarray]]
    simplex: Union[np.ndarray, Tuple[np.ndarray]]
    input_ids: Optional[Union[np.ndarray, Tuple[np.ndarray]]]
    metrics: Optional[Dict[str, float]]
    results: Optional[Dict[str, List[str]]]
    num_samples: Optional[int]


class DiffusionTrainer(Trainer):
    def __init__(
        self,
        noise_scheduler,
        inference_noise_schedulers,
        diffusion_args,
        data_args,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.original_data_collator = self.data_collator
        self.noise_scheduler = noise_scheduler
        self.diffusion_args = diffusion_args
        self.data_args = data_args
        self.vocab_size = self.model.config.vocab_size
        self.inference_noise_schedulers = inference_noise_schedulers
        self.inference_timesteps = diffusion_args.num_inference_diffusion_steps
        self.tb_writer = self.get_tb_writer()
        self.eos_token_id = self.tokenizer.eos_token_id
        self.classifier_free_guidance = (
            diffusion_args.guidance_scale > 1.0
            and data_args.conditional_generation is not None
        )
        self.counter = 0
        # TODO: control seed.
        self.self_cond_generator = np.random.default_rng(42)

    def annotated_split(self, split):
        return f"{split}_top_p_{self.diffusion_args.top_p}_temperature_{self.diffusion_args.temperature}_seed_{self.args.seed}_guidance_scale_{self.diffusion_args.guidance_scale}"

    def save_metrics(self, split, metrics, combined=True):
        super().save_metrics(self.annotated_split(split), metrics, combined)

    def log_metrics(self, split, metrics):
        super().log_metrics(self.annotated_split(split), metrics)

    def get_tb_writer(self):
        for cb in self.callback_handler.callbacks:
            if isinstance(cb, TensorBoardCallback):
                return cb
        return None

    def training_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]
    ) -> torch.Tensor:
        """
        Perform a training step on a batch of inputs.
        Subclass and override to inject custom behavior.
        Args:
            model (`nn.Module`):
                The model to train.
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.
                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.
        Return:
            `torch.Tensor`: The tensor with training loss on this batch.
        """
        model.train()
        inputs = self._prepare_inputs(inputs)

        # Truncate the length if needed.
        if self.data_args.truncation_length > 0:
            inputs["input_ids"] = inputs["input_ids"][
                :, : -self.data_args.truncation_length
            ]
            inputs["span_mask"] = inputs["span_mask"][
                :, : -self.data_args.truncation_length
            ]

        # Creates the noisy simplex and timesteps.
        simplex = convert_to_simplex(
            inputs["input_ids"], self.diffusion_args.simplex_value, self.vocab_size
        )
        noise = self.diffusion_args.simplex_value * torch.randn(
            simplex.shape, device=simplex.device, dtype=simplex.dtype
        )
        bsz = simplex.shape[0]
        # Sample a random timestep for each simplex token representation.
        # testing just sampling the same place. This better matches reality.
        if True:  # np.random.rand(1) > 0.5:
            timesteps = torch.randint(
                0,
                len(self.noise_scheduler),
                (bsz, inputs["input_ids"].shape[1])
                if False  # is_tokenwise_cdcd_check(self.model)
                else (bsz,),
                device=simplex.device,
                dtype=torch.int64,
            )
            timesteps = timesteps[:, None].expand(-1, inputs["input_ids"].shape[1])
        else:
            timesteps = torch.randint(
                0,
                len(self.noise_scheduler),
                (bsz, inputs["input_ids"].shape[1])
                if True  # is_tokenwise_cdcd_check(self.model)
                else (bsz,),
                device=simplex.device,
                dtype=torch.int64,
            )
        # expand out timesteps to match tokenwise setup
        # if True:  # not is_tokenwise_cdcd_check(self.model):
        #     timesteps = timesteps[:, None].expand(-1, inputs["input_ids"].shape[1])

        # save original timesteps for warping
        original_timesteps = timesteps
        # warp timesteps according to cdf
        # we re-scale the timesteps to the correct range.
        # the -1 is due to the timestep should be in range [0, 5000)
        if is_cdcd_check(self.model):
            input_ids = inputs["input_ids"]
            span_mask = inputs["span_mask"]
            token_input = torch.where(
                (input_ids * span_mask) > 1, self.tokenizer.pad_token_id, input_ids
            )
            timesteps = self.model.warp_timesteps(
                timesteps,
                token_input=token_input,
                span_mask=span_mask,
                t_max=len(self.noise_scheduler) - 1,
            )
        # Adds noise to each simplex representation (Forward diffusion process).
        noisy_simplex = self.noise_scheduler.add_noise(simplex, noise, timesteps)
        # the warper model will scale the timesteps to the correct range.
        timesteps = scale(timesteps, len(self.noise_scheduler))
        # original_timesteps_scaled = scale(original_timesteps, len(self.noise_scheduler))
        # inputs.update(
        #     {"original_timesteps": scale(original_timesteps, len(self.noise_scheduler))}
        # )

        inputs.update(
            {
                "timesteps": timesteps,
                "simplex": noisy_simplex,
            }
        )
        # inputs.update({"max_timestep": len(self.noise_scheduler)})
        if self.diffusion_args.self_condition is not None:
            previous_pred = None
            # previous_hidden = None
            if self.self_cond_generator.random(1) > 0.5:
                next_timestep = inputs.pop("timesteps")
                next_simplex = inputs.pop("simplex")
                timesteps = torch.clamp(
                    (next_timestep * len(self.noise_scheduler)) + 1,
                    max=len(self.noise_scheduler) - 1,
                )
                if is_cdcd_check(self.model):
                    input_ids = inputs["input_ids"]
                    span_mask = inputs["span_mask"]
                    token_input = torch.where(
                        (input_ids * span_mask) > 1,
                        self.tokenizer.pad_token_id,
                        input_ids,
                    )
                    timesteps = self.model.warp_timesteps(
                        timesteps,
                        token_input=token_input,
                        span_mask=span_mask,
                        t_max=len(self.noise_scheduler) - 1,
                    )
                noisy_simplex = self.noise_scheduler.add_noise(
                    simplex, noise, timesteps
                )
                timesteps = scale(timesteps, len(self.noise_scheduler))
                inputs.update(
                    {
                        "timesteps": timesteps,
                        "simplex": noisy_simplex,
                    }
                )
                # we don't backprop through this.
                with torch.no_grad():
                    outputs = model(**inputs, previous_pred=previous_pred)
                logits_projection_fct = lambda x: logits_projection(  # noqa: E731
                    x,
                    self.diffusion_args.sampling_type,
                    self.diffusion_args.top_p,
                    self.diffusion_args.simplex_value,
                    self.diffusion_args.temperature,
                )
                previous_pred = self_condition_preds(
                    self.diffusion_args.self_condition,
                    outputs.logits,
                    logits_projection_fct,
                ).detach()
                # following rest of self-conditioning, don't backprop through.
                # previous_hidden = outputs.hidden_states.detach()
                # pop timestep/simplex and put the old ones back.
                inputs.update(
                    {
                        "timesteps": next_timestep,
                        "simplex": next_simplex,
                    }
                )
            inputs.update({"previous_pred": previous_pred})
            # inputs.update({"previous_hidden": previous_hidden})
        else:
            inputs.update({"previous_pred": None})
            # inputs.update({"previous_hidden": None})
            # previous_hidden = None
        # NOTE: we do this after computation of self-conditioning to not affect that one.
        # inputs.update(
        #     {"classifier_free_guidance_in_train": self.classifier_free_guidance}
        # )
        # re-warp based on previous hidden state
        if is_cdcd_check(self.model):
            # replace masked tokens with <mask> token.
            input_ids = inputs["input_ids"]
            span_mask = inputs["span_mask"]
            token_input = torch.where(
                (input_ids * span_mask) > 1, self.tokenizer.pad_token_id, input_ids
            )
            timesteps = self.model.warp_timesteps(
                original_timesteps,
                t_max=len(self.noise_scheduler) - 1,
                token_input=token_input,
                span_mask=span_mask,
            )
            noisy_simplex = self.noise_scheduler.add_noise(simplex, noise, timesteps)
            timesteps = scale(timesteps, len(self.noise_scheduler))
            inputs.update(
                {
                    "timesteps": timesteps,
                    "simplex": noisy_simplex,
                }
            )
        with self.compute_loss_context_manager():
            loss = self.compute_loss(model, inputs)

        if self.args.n_gpu > 1:
            loss = loss.mean()  # mean() to average on multi-gpu parallel training
        # HACK: transformer update
        # if self.do_grad_scaling:
        #     self.scaler.scale(loss).backward()
        elif self.use_apex:
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            self.accelerator.backward(loss)
        return loss.detach() / self.args.gradient_accumulation_steps

    def light_prediction_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
        with torch.no_grad():
            inputs = self._prepare_inputs(inputs)
            # Truncate the length if needed.
            if self.data_args.truncation_length > 0:
                inputs["input_ids"] = inputs["input_ids"][
                    :, : -self.data_args.truncation_length
                ]
                inputs["span_mask"] = inputs["span_mask"][
                    :, : -self.data_args.truncation_length
                ]
            # Creates the noisy simplex and timesteps.
            simplex = convert_to_simplex(
                inputs["input_ids"], self.diffusion_args.simplex_value, self.vocab_size
            )
            noise = self.diffusion_args.simplex_value * torch.randn(
                simplex.shape, device=simplex.device, dtype=simplex.dtype
            )
            bsz = simplex.shape[0]
            # Sample a random timestep for each simplex token representation.
            # we use the train timesteps to be consistent with the training process.
            # randomly flip between random batchwise and tokenwise timesteps.
            if True:
                timesteps = torch.randint(
                    0,
                    len(self.noise_scheduler),
                    (bsz, inputs["input_ids"].shape[1])
                    if False  # is_tokenwise_cdcd_check(self.model)
                    else (bsz,),
                    device=simplex.device,
                    dtype=torch.int64,
                )
                timesteps = timesteps[:, None].expand(-1, inputs["input_ids"].shape[1])
            else:
                timesteps = torch.randint(
                    0,
                    len(self.noise_scheduler),
                    (bsz, inputs["input_ids"].shape[1])
                    if True  # is_tokenwise_cdcd_check(self.model)
                    else (bsz,),
                    device=simplex.device,
                    dtype=torch.int64,
                )
            # original_timesteps = timesteps

            # if cdcd, we need to wrap the timesteps in a cdf.
            # make sure we scale the timesteps to the correct range!
            if is_cdcd_check(self.model):
                input_ids = inputs["input_ids"]
                span_mask = inputs["span_mask"]
                token_input = torch.where(
                    (input_ids * span_mask) > 1, self.tokenizer.pad_token_id, input_ids
                )
                timesteps = self.model.warp_timesteps(
                    timesteps,
                    t_max=len(self.noise_scheduler) - 1,
                    token_input=token_input,
                    span_mask=span_mask,
                )

            # Adds noise to each simplex representation (Forward diffusion process).
            noisy_simplex = self.noise_scheduler.add_noise(simplex, noise, timesteps)

            timesteps = scale(timesteps, len(self.noise_scheduler))
            # original_timesteps_scaled = scale(
            #     original_timesteps, len(self.noise_scheduler)
            # )
            # inputs.update({"original_timesteps": original_timesteps_scaled})

            inputs.update(
                {
                    "timesteps": timesteps,
                    "simplex": noisy_simplex,
                }
            )
            # inputs.update({"max_timestep": len(self.noise_scheduler)})
            if self.diffusion_args.self_condition is not None:
                previous_pred = None
                # last_hidden_state = None
                if np.random.rand(1) > 0.5:
                    outputs = model(**inputs, previous_pred=previous_pred)
                    logits_projection_fct = lambda x: logits_projection(  # noqa: E731
                        x,
                        self.diffusion_args.sampling_type,
                        self.diffusion_args.top_p,
                        self.diffusion_args.simplex_value,
                        self.diffusion_args.temperature,
                    )
                    previous_pred = self_condition_preds(
                        self.diffusion_args.self_condition,
                        outputs.logits,
                        logits_projection_fct,
                    )
                    # last_hidden_state = outputs.hidden_states
                inputs.update(
                    {
                        "previous_pred": previous_pred,
                        # "previous_hidden": last_hidden_state,
                    }
                )
            # NOTE: we do this after computation of self-conditioning to not affect that one.
            # inputs.update(
            #     {"classifier_free_guidance_in_train": self.classifier_free_guidance}
            # )
            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs)

            if self.args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training

            return (
                loss.detach()
            )  # no division by gradient accumulation steps for eval. we want per-sample avg loss.

    # TODO: argument for doing one step.
    def prediction_step(
        self,
        inputs: Dict[str, Union[torch.Tensor, Any]],
        model: nn.Module,
        pipeline: List[SimplexDDPMClassifierGuidancePipeline],
    ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
        inputs = self._prepare_inputs(inputs)
        # full inference.
        with torch.no_grad():
            with self.compute_loss_context_manager():
                for i, x in enumerate(
                    pipeline(
                        seq_length=self.data_args.max_seq_length
                        - self.data_args.truncation_length,
                        batch=inputs,
                        guidance_scale=self.diffusion_args.guidance_scale,
                        generator=torch.Generator(device=self.args.device).manual_seed(
                            self.args.seed
                        )
                        if self.diffusion_args.generate_with_seed
                        else None,
                        is_generator=False,
                        use_gumbel_softmax=self.diffusion_args.use_gumbel_softmax,
                        do_hard_sample=self.diffusion_args.do_hard_sample,
                        softmax_temperature=self.diffusion_args.softmax_temperature,
                        num_guidance_steps=self.diffusion_args.num_guidance_steps,
                    )
                ):
                    outputs = x
        logits = nested_detach(outputs.logits)
        simplex = nested_detach(outputs.simplex)

        return (simplex, logits)

    def evaluation_loop(
        self,
        dataloader: DataLoader,
        description: str,
        prediction_loss_only: Optional[bool] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
        noise_scheduler=None,
        light_eval_dataloader=None,
        do_light_eval=False,
    ) -> EvalLoopOutput:
        """
        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.
        Works both with or without labels.
        """
        args = self.args
        is_conditional_generation = self.data_args.conditional_generation is not None
        save_prefixes = is_conditional_generation

        prediction_loss_only = (
            prediction_loss_only
            if prediction_loss_only is not None
            else args.prediction_loss_only
        )
        # if eval is called w/o train handle model prep here
        if self.is_deepspeed_enabled and self.model_wrapped is self.model:
            _, _ = deepspeed_init(self, num_training_steps=0, inference=True)

        model = self._wrap_model(self.model, training=False, dataloader=dataloader)

        if len(self.accelerator._models) == 0 and model is self.model:
            model = (
                self.accelerator.prepare(model)
                if self.is_deepspeed_enabled
                else self.accelerator.prepare_model(model, evaluation_mode=True)
            )

            if self.is_fsdp_enabled:
                self.model = model

            # for the rest of this function `model` is the outside model, whether it was wrapped or not
            if model is not self.model:
                self.model_wrapped = model

            # backward compatibility
            if self.is_deepspeed_enabled:
                self.deepspeed = self.model_wrapped

        # if full fp16 or bf16 eval is wanted and this ``evaluation`` or ``predict`` isn't called
        # while ``train`` is running, cast it to the right dtype first and then put on device
        if not self.is_in_train:
            if args.fp16_full_eval:
                model = model.to(dtype=torch.float16, device=args.device)
            elif args.bf16_full_eval:
                model = model.to(dtype=torch.bfloat16, device=args.device)

        batch_size = self.args.eval_batch_size

        logger.info(f"***** Running {description} *****")
        if has_length(dataloader):
            logger.info(f"  Num examples = {self.num_examples(dataloader)}")
        else:
            logger.info("  Num examples: Unknown")
        logger.info(f"  Batch size = {batch_size}")

        model.eval()

        pipeline = SimplexDDPMClassifierGuidancePipeline(
            model=model,
            scheduler=noise_scheduler,
            simplex_value=self.diffusion_args.simplex_value,
            top_p=self.diffusion_args.top_p,
            sampling_type=self.diffusion_args.sampling_type,
            is_conditional_generation=is_conditional_generation,
            tokenizer=self.tokenizer,
            classifier_free_uncond_input=self.diffusion_args.classifier_free_uncond_input,
            temperature=self.diffusion_args.temperature,
            guidance_softmax_combination=self.diffusion_args.guidance_softmax_combination,
            classifier_model_name_or_path=self.diffusion_args.classifier_model_name_or_path,
        )

        self.callback_handler.eval_dataloader = dataloader
        # Do this before wrapping.
        eval_dataset = getattr(dataloader, "dataset", None)

        # Initialize containers
        # logits/simplex/labels on GPU/TPU (accumulated for eval_accumulation_steps)
        losses_host = None
        logits_host = None
        simplex_host = None
        inputs_host = None
        masks_host = None
        prefixes_host = None

        # logits/simplex/labels on CPU (final containers)
        all_losses = None
        all_logits = None
        all_simplex = None
        all_inputs = None
        all_masks = None
        all_prefixes = None
        observed_num_examples = 0

        # light evaluation loop.
        if light_eval_dataloader is not None and do_light_eval:
            for step, inputs in enumerate(light_eval_dataloader):
                # Truncate the length if needed.
                if self.data_args.truncation_length > 0:
                    inputs["input_ids"] = inputs["input_ids"][
                        :, : -self.data_args.truncation_length
                    ]
                    inputs["span_mask"] = inputs["span_mask"][
                        :, : -self.data_args.truncation_length
                    ]
                    max_seq_length = (
                        self.data_args.max_seq_length - self.data_args.truncation_length
                    )
                    assert self.data_args.eval_context_size < max_seq_length
                # predict loss mimicking training.
                loss = self.light_prediction_step(model, inputs)

                if loss is not None:
                    losses = self._nested_gather(loss.repeat(batch_size))
                    losses_host = (
                        losses
                        if losses_host is None
                        else torch.cat((losses_host, losses), dim=0)
                    )
                if (
                    args.eval_accumulation_steps is not None
                    and (step + 1) % args.eval_accumulation_steps == 0
                ):
                    if losses_host is not None:
                        losses = nested_numpify(losses_host)
                        all_losses = (
                            losses
                            if all_losses is None
                            else np.concatenate((all_losses, losses), axis=0)
                        )
                    losses_host = None

        # Main evaluation loop
        for step, inputs in enumerate(dataloader):
            has_mask = True if "span_mask" in inputs else False

            # Truncate the length if needed.
            if self.data_args.truncation_length > 0:
                inputs["input_ids"] = inputs["input_ids"][
                    :, : -self.data_args.truncation_length
                ]
                inputs["span_mask"] = inputs["span_mask"][
                    :, : -self.data_args.truncation_length
                ]
                max_seq_length = (
                    self.data_args.max_seq_length - self.data_args.truncation_length
                )
                assert self.data_args.eval_context_size < max_seq_length

            # Update the observed num examples
            observed_batch_size = find_batch_size(inputs)
            if observed_batch_size is not None:
                observed_num_examples += observed_batch_size
                # For batch samplers, batch_size is not known by the dataloader in advance.
                if batch_size is None:
                    batch_size = observed_batch_size

            # Prediction step
            simplex, logits = self.prediction_step(inputs, model, pipeline=pipeline)
            inputs_decode = self._prepare_input(inputs["input_ids"])
            masks = self._prepare_input(inputs["span_mask"]) if has_mask else None
            if save_prefixes:
                prefixes = (
                    pad_data(
                        [input[~mask] for input, mask in zip(inputs_decode, masks)],
                        self.tokenizer,
                    )
                    if has_mask
                    else None
                )
                prefixes = self._prepare_input(prefixes)
            else:
                prefixes = None
            # Update containers on host
            if prefixes is not None:
                prefixes = self.accelerator.pad_across_processes(
                    prefixes, dim=1, pad_index=self.eos_token_id
                )
                prefixes = self._nested_gather(prefixes)
                prefixes_host = (
                    prefixes
                    if prefixes_host is None
                    else nested_concat(
                        prefixes_host, prefixes, padding_index=self.eos_token_id
                    )
                )
            if inputs_decode is not None:
                inputs_decode = self.accelerator.pad_across_processes(
                    inputs_decode, dim=1, pad_index=self.eos_token_id
                )
                inputs_decode = self._nested_gather(inputs_decode)
                inputs_host = (
                    inputs_decode
                    if inputs_host is None
                    else nested_concat(
                        inputs_host, inputs_decode, padding_index=self.eos_token_id
                    )
                )
            # Note that this block should be before masks block, since we need masks here.
            if simplex is not None:
                # In case of having a mask softmax is applied over the simplex non-masked values.
                if has_mask:
                    mask_value = torch.finfo(simplex.dtype).min
                    mask_value = torch.tensor(
                        mask_value, dtype=simplex.dtype, device=simplex.device
                    )
                    simplex = torch.where(masks[:, :, None], simplex, mask_value)
                simplex = F.softmax(simplex, dim=-1)
                if self.preprocess_logits_for_metrics is not None:
                    simplex = self.preprocess_logits_for_metrics(simplex)
                simplex = self.accelerator.pad_across_processes(
                    simplex, dim=1, pad_index=self.eos_token_id
                )
                simplex = self._nested_gather(simplex)
                # TODO: note that this is no more a simplex, but the processed one.
                simplex_host = (
                    simplex
                    if simplex_host is None
                    else nested_concat(
                        simplex_host, simplex, padding_index=self.eos_token_id
                    )
                )
            if masks is not None:
                masks = self.accelerator.pad_across_processes(masks, dim=1, pad_index=0)
                masks = self._nested_gather(masks)
                # We pad masks with False tokens.
                masks_host = (
                    masks
                    if masks_host is None
                    else nested_concat(masks_host, masks, padding_index=0)
                )
            if logits is not None:
                if self.preprocess_logits_for_metrics is not None:
                    logits = self.preprocess_logits_for_metrics(logits)
                logits = self.accelerator.pad_across_processes(
                    logits, dim=1, pad_index=self.eos_token_id
                )
                logits = self._nested_gather(logits)
                logits_host = (
                    logits
                    if logits_host is None
                    else nested_concat(
                        logits_host, logits, padding_index=self.eos_token_id
                    )
                )

            self.control = self.callback_handler.on_prediction_step(
                args, self.state, self.control
            )

        # Gather all tensors and put them back on the CPU if we have done enough accumulation steps.
        if (
            args.eval_accumulation_steps is not None
            and (step + 1) % args.eval_accumulation_steps == 0
        ):
            if logits_host is not None:
                logits = nested_numpify(logits_host)
                all_logits = (
                    logits
                    if all_logits is None
                    else nested_concat(
                        all_logits, logits, padding_index=self.eos_token_id
                    )
                )
            if simplex_host is not None:
                simplex = nested_numpify(simplex_host)
                all_simplex = (
                    simplex
                    if all_simplex is None
                    else nested_concat(
                        all_simplex, simplex, padding_index=self.eos_token_id
                    )
                )
            if inputs_host is not None:
                inputs_decode = nested_numpify(inputs_host)
                all_inputs = (
                    inputs_decode
                    if all_inputs is None
                    else nested_concat(
                        all_inputs, inputs_decode, padding_index=self.eos_token_id
                    )
                )
            if masks_host is not None:
                masks = nested_numpify(masks_host)
                all_masks = (
                    masks
                    if all_masks is None
                    else nested_concat(all_masks, masks, padding_index=0)
                )
            if prefixes_host is not None:
                prefixes = nested_numpify(prefixes_host)
                all_prefixes = (
                    prefixes
                    if all_prefixes is None
                    else nested_concat(
                        all_prefixes, prefixes, padding_index=self.eos_token_id
                    )
                )

            # Set back to None to begin a new accumulation
            logits_host, simplex_host, inputs_host, masks_host, prefixes_host = (
                None,
                None,
                None,
                None,
                None,
            )

        # Gather all remaining tensors and put them back on the CPU
        if losses_host is not None:
            all_losses = nested_numpify(losses_host)
        if logits_host is not None:
            all_logits = nested_numpify(logits_host)
        if simplex_host is not None:
            all_simplex = nested_numpify(simplex_host)
        if inputs_host is not None:
            all_inputs = nested_numpify(inputs_host)
        if masks_host is not None:
            all_masks = nested_numpify(masks_host)
        if prefixes_host is not None:
            all_prefixes = nested_numpify(prefixes_host)

        if args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of the evaluation loop
            delattr(self, "_past")

        # Number of samples
        if has_length(eval_dataset):
            num_samples = len(eval_dataset)
        # The instance check is weird and does not actually check for the type, but whether the dataset has the right
        # methods. Therefore we need to make sure it also has the attribute.
        elif (
            isinstance(eval_dataset, IterableDatasetShard)
            and getattr(eval_dataset, "num_examples", 0) > 0
        ):
            num_samples = eval_dataset.num_examples
        else:
            if has_length(dataloader):
                num_samples = self.num_examples(dataloader)
            else:  # both len(dataloader.dataset) and len(dataloader) fail
                num_samples = observed_num_examples
        if num_samples == 0 and observed_num_examples > 0:
            num_samples = observed_num_examples

        # Generates the texts.
        results = {}
        if is_conditional_generation:
            # We predict the masked tokens only. Here, we compute the masked tokens.
            results.update(
                predict_conditional_generated(
                    all_masks,
                    all_inputs,
                    self.tokenizer,
                    all_simplex,
                    "pred_texts_from_simplex",
                    self.data_args.skip_special_tokens,
                )
            )
            results.update(
                predict_conditional_generated(
                    all_masks,
                    all_inputs,
                    self.tokenizer,
                    all_logits,
                    "pred_texts_from_logits",
                    self.data_args.skip_special_tokens,
                )
            )
        else:
            results.update(
                {
                    "pred_texts_from_simplex": self.tokenizer.batch_decode(
                        all_simplex,
                        skip_special_tokens=self.data_args.skip_special_tokens,
                    )
                }
            )
            results.update(
                {
                    "pred_texts_from_logits": self.tokenizer.batch_decode(
                        all_logits,
                        skip_special_tokens=self.data_args.skip_special_tokens,
                    )
                }
            )
        if is_conditional_generation:
            results.update(
                {
                    "gold_texts_masked": [
                        self.tokenizer.decode(
                            input[mask],
                            skip_special_tokens=self.data_args.skip_special_tokens,
                        )
                        for mask, input in zip(all_masks, all_inputs)
                    ]
                }
            )
            if save_prefixes:
                results.update(
                    {
                        "prefixes": [
                            self.tokenizer.decode(
                                x, skip_special_tokens=True
                            )  # self.data_args.skip_special_tokens)
                            for x in all_prefixes
                        ]
                    }
                )

        # Metrics.
        if self.compute_metrics is not None:
            metrics = self.compute_metrics(results)
        else:
            metrics = {}

        # To be JSON-serializable, we need to remove numpy types or zero-d tensors
        metrics = denumpify_detensorize(metrics)

        if all_losses is not None:
            metrics[f"{metric_key_prefix}_loss"] = all_losses.mean().item()

        # Prefix all keys with metric_key_prefix + '_'
        for key in list(metrics.keys()):
            if not key.startswith(f"{metric_key_prefix}_"):
                metrics[f"{metric_key_prefix}_{key}"] = metrics.pop(key)

        return EvalLoopOutput(
            logits=all_logits,
            simplex=all_simplex,
            input_ids=all_inputs,
            metrics=metrics,
            num_samples=num_samples,
            results=results,
        )

    def evaluate(
        self,
        eval_dataset: Optional[Dataset] = None,
        ignore_keys: Optional[List[str]] = None,
        metric_key_prefix: str = "eval",
    ) -> Dict[str, float]:
        """
        Run evaluation and returns metrics.
        The calling script will be responsible for providing a method to compute metrics, as they are task-dependent
        (pass it to the init `compute_metrics` argument).
        You can also subclass and override this method to inject custom behavior.
        Args:
            eval_dataset (`Dataset`, *optional*):
                Pass a dataset if you wish to override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns
                not accepted by the `model.forward()` method are automatically removed. It must implement the `__len__`
                method.
            ignore_keys (`Lst[str]`, *optional*):
                A list of keys in the output of your model (if it is a dictionary) that should be ignored when
                gathering predictions.
            metric_key_prefix (`str`, *optional*, defaults to `"eval"`):
                An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named
                "eval_bleu" if the prefix is "eval" (default)
        Returns:
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The
            dictionary also contains the epoch number which comes from the training state.
        """
        # memory metrics - must set up as early as possible
        self._memory_tracker.start()
        eval_dataloader = self.get_eval_dataloader(eval_dataset)
        light_eval_dataloader = self.get_light_eval_dataloader(eval_dataset)
        start_time = time.time()

        outputs = []
        timesteps = self.inference_timesteps
        for timestep, noise_scheduler in zip(
            timesteps, self.inference_noise_schedulers
        ):
            output = self.evaluation_loop(
                eval_dataloader,
                description="Evaluation",
                # No point gathering the predictions if there are no metrics, otherwise we defer to
                # self.args.prediction_loss_only
                prediction_loss_only=True if self.compute_metrics is None else None,
                ignore_keys=ignore_keys,
                metric_key_prefix=metric_key_prefix,
                noise_scheduler=noise_scheduler,
                light_eval_dataloader=light_eval_dataloader,
                do_light_eval=timestep
                == timesteps[
                    0
                ],  # we only need the loss once, since it is the same for all timesteps
            )
            outputs.append(output)
            key_prefix = f"inference_{timestep}_"
            metrics = {key_prefix + k: v for k, v in output.metrics.items()}
            results = {key_prefix + k: v for k, v in output.results.items()}
            # reset output with new metrics / results
            output = EvalLoopOutput(
                logits=output.logits,
                simplex=output.simplex,
                input_ids=output.input_ids,
                metrics=metrics,
                num_samples=output.num_samples,
                results=results,
            )

            total_batch_size = self.args.eval_batch_size * self.args.world_size
            output.metrics.update(
                speed_metrics(
                    metric_key_prefix,
                    start_time,
                    num_samples=output.num_samples,
                    num_steps=math.ceil(output.num_samples / total_batch_size),
                )
            )
            self.log(output.metrics)
            self.control = self.callback_handler.on_evaluate(
                self.args, self.state, self.control, output.metrics
            )
            self._memory_tracker.stop_and_update_metrics(output.metrics)

            # Save the results
            self.save_metrics(
                GENERATION_RESULTS + "_" + key_prefix + metric_key_prefix,
                output.results,
            )
            logger.info("Results are saved now")

        # log outside so we can group generations together
        if self.args.log_generated_texts:
            length = len(outputs[0].logits)
            results = {
                f"{k}_inference_{i}": v
                for o, i in zip(outputs, timesteps)
                for k, v in o.results.items()
            }
            self.log_results_to_tensorboard(self.state, length, results)

        return output.metrics

    def log_results_to_tensorboard(self, state, length, results):
        # TODO: we need to fix this which happens during the only eval option.
        if self.tb_writer.tb_writer is None:
            return
        for i in range(length):
            total_text = ""
            for k, v in results.items():
                total_text += f"*** {k} ***: {v[i]}" + "  \n"
            self.tb_writer.tb_writer.add_text(
                f"sample_{i}", total_text, state.global_step
            )

    def get_train_dataloader(self) -> DataLoader:
        self.data_collator = self.original_data_collator("train")
        return super().get_train_dataloader()

    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
        self.data_collator = self.original_data_collator("eval")
        return super().get_eval_dataloader(eval_dataset)

    def get_light_eval_dataloader(
        self, eval_dataset: Optional[Dataset] = None
    ) -> DataLoader:
        """
        Returns the evaluation [`~torch.utils.data.DataLoader`].
        Used for the light evaluation, which matches masking with training.
        Args:
            eval_dataset (`torch.utils.data.Dataset`, *optional*):
                If provided, will override `self.eval_dataset`. If it is a [`~datasets.Dataset`], columns not accepted
                by the `model.forward()` method are automatically removed. It must implement `__len__`.
        """
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        data_collator = self.original_data_collator("train")

        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
            eval_dataset = self._remove_unused_columns(
                eval_dataset, description="evaluation"
            )
        else:
            data_collator = self._get_collator_with_removed_columns(
                data_collator, description="evaluation"
            )

        dataloader_params = {
            "batch_size": self.args.eval_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(eval_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset)
            dataloader_params["drop_last"] = self.args.dataloader_drop_last

        return self.accelerator.prepare(DataLoader(eval_dataset, **dataloader_params))

    def create_optimizer(self):
        from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
        from transformers.trainer_pt_utils import get_parameter_names

        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model

        if self.optimizer is not None:
            return self.optimizer

        decay_parameters = get_parameter_names(opt_model, ALL_LAYERNORM_LAYERS)
        decay_parameters = [name for name in decay_parameters if "bias" not in name]
        optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
            self.args
        )

        # override to apply higher lr to timestep_embed and cdcd cdf
        optimizer_grouped_parameters = [
            {
                "params": [
                    p
                    for n, p in opt_model.named_parameters()
                    if (
                        n in decay_parameters
                        and p.requires_grad
                        and not ("timestep_embed" in n or "cdf" in n)
                    )
                ],
                "weight_decay": self.args.weight_decay,
                "lr": optimizer_kwargs["lr"],
            },
            {
                "params": [
                    p
                    for n, p in opt_model.named_parameters()
                    if (
                        n not in decay_parameters
                        and p.requires_grad
                        and not ("timestep_embed" in n or "cdf" in n)
                    )
                ],
                "weight_decay": 0.0,
                "lr": optimizer_kwargs["lr"],
            },
            {
                "params": [
                    p
                    for n, p in opt_model.named_parameters()
                    if (("timestep_embed" in n) and p.requires_grad)
                ],
                "weight_decay": 0.0,
                "lr": self.args.timestep_embed_lr or self.args.learning_rate,
            },
        ]
        # check cdcd
        cdf_params = [
            p
            for n, p in opt_model.named_parameters()
            if (("cdf" in n) and p.requires_grad)
        ]
        if cdf_params:
            optimizer_grouped_parameters.append(
                {
                    "params": cdf_params,
                    "weight_decay": 0.0,
                    "lr": 1e-3,
                }
            )

        optimizer_kwargs.pop("lr")

        self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)

        return self.optimizer