reference/code/flm/main.py

import json
import os
import itertools
import functools
import argparse
import fsspec
import hydra
import lightning as L
import omegaconf
import rich.syntax
import rich.tree
import torch
torch.load = functools.partial(torch.load, weights_only=False)
from torch.distributed import init_process_group, destroy_process_group
import wandb
import algo
import dataloader
import utils

import numpy as np
from datetime import datetime

import uuid

# Allow torch.load(weights_only=True) to safely unpickle Hydra configs stored in checkpoints
torch.serialization.add_safe_globals([omegaconf.dictconfig.DictConfig, omegaconf.base.ContainerMetadata, omegaconf.base.Metadata])

omegaconf.OmegaConf.register_new_resolver(
    'cwd', os.getcwd)
omegaconf.OmegaConf.register_new_resolver(
    'device_count', torch.cuda.device_count)
omegaconf.OmegaConf.register_new_resolver(
    'eval', eval)
omegaconf.OmegaConf.register_new_resolver(
    'div_up', lambda x, y: (x + y - 1) // y)


def _load_from_checkpoint(diffusion_model, config, tokenizer):
    if 'hf' in config.algo.backbone:
        return diffusion_model(
            config, tokenizer=tokenizer).to('cuda')

    return diffusion_model.load_from_checkpoint(
        config.eval.checkpoint_path,
        tokenizer=tokenizer,
        config=config,
        weights_only=False)


@L.pytorch.utilities.rank_zero_only
def _print_config(
        config: omegaconf.DictConfig,
        resolve: bool = True,
        save_cfg: bool = True) -> None:
    """Prints content of DictConfig using Rich library and its tree structure.

    Args:
      config (DictConfig): Configuration composed by Hydra.
      resolve (bool): Whether to resolve reference fields of DictConfig.
      save_cfg (bool): Whether to save the configuration tree to a file.
    """

    style = 'dim'
    tree = rich.tree.Tree('CONFIG', style=style, guide_style=style)

    fields = config.keys()
    for field in fields:
        branch = tree.add(field, style=style, guide_style=style)

        config_section = config.get(field)
        branch_content = str(config_section)
        if isinstance(config_section, omegaconf.DictConfig):
            branch_content = omegaconf.OmegaConf.to_yaml(
                config_section, resolve=resolve)

        branch.add(rich.syntax.Syntax(branch_content, 'yaml'))
    rich.print(tree)
    if save_cfg:
        with fsspec.open(
            '{}/config_tree.txt'.format(
                config.checkpointing.save_dir), 'w') as fp:
            rich.print(tree, file=fp)


@L.pytorch.utilities.rank_zero_only
def _print_batch(train_ds, valid_ds, tokenizer, k=64):
    for dl_type, dl in [
            ('train', train_ds), ('valid', valid_ds)]:
        print(f'Printing {dl_type} dataloader batch.')
        batch = next(iter(dl))
        print('Batch input_ids.shape', batch['input_ids'].shape)
        first = batch['input_ids'][0, :k]
        last = batch['input_ids'][0, -k:]
        print(f'First {k} tokens:', tokenizer.decode(first))
        print('ids:', first)
        print(f'Last {k} tokens:', tokenizer.decode(last))
        print('ids:', last)


def _generate_samples(diffusion_model, config, logger,
                      tokenizer):
    logger.info('Starting Sample Eval.')
    model = _load_from_checkpoint(
        diffusion_model=diffusion_model,
        config=config,
        tokenizer=tokenizer)
    model.metrics.gen_ppl.reset()
    model.metrics.sample_entropy.reset()
    if config.eval.disable_ema:
        logger.info('Disabling EMA.')
        model.ema = None
    stride_length = config.sampling.stride_length
    num_strides = config.sampling.num_strides
    all_samples = []

    print("generation start: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    for _ in range(config.sampling.num_sample_batches):
        if config.sampling.semi_ar:
            _, intermediate_samples, _ = model.restore_model_and_semi_ar_sample(
                stride_length=stride_length,
                num_strides=num_strides,
                dt=1 / config.sampling.steps)
            text_samples = intermediate_samples[-1]
            # Note: Samples generated using semi-ar method
            # need to to be processed before computing generative perplexity
            # since these samples contain numerous <|endoftext|> tokens
            # and diffusion.compute_generative_perplexity() discards
            # any text after the first EOS token.
        else:
            samples = model.restore_model_and_sample(
                num_steps=config.sampling.steps)
            model.metrics.record_entropy(samples)
            text_samples = model.tokenizer.batch_decode(samples)
            model.metrics.record_generative_perplexity(
                text_samples, config.model.length, model.device)
            all_samples.extend(list(text_samples))

    print("generation end: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    generative_ppl = 0.
    entropy = 0.
    if not config.sampling.semi_ar:
        generative_ppl = model.metrics.gen_ppl.compute().item()
        entropy = model.metrics.sample_entropy.compute().item()
        print('Generative perplexity:', generative_ppl)
        print('Sample entropy:', entropy)
    samples_path = config.eval.generated_samples_path
    with fsspec.open(samples_path, 'w') as f:
        json.dump({'generative_ppl': generative_ppl,
                   'entropy': entropy,
                   'generated_seqs': all_samples}, f, indent=4)
    print('Samples saved at:', samples_path)


def _generate_samples_with_tc(diffusion_model, config, logger,
                              tokenizer):
    logger.info('Starting Sample Eval.')
    model = _load_from_checkpoint(
        diffusion_model=diffusion_model,
        config=config,
        tokenizer=tokenizer)
    model.metrics.gen_ppl.reset()
    model.metrics.sample_entropy.reset()
    if config.eval.disable_ema:
        logger.info('Disabling EMA.')
        model.ema = None
    stride_length = config.sampling.stride_length
    num_strides = config.sampling.num_strides
    all_samples = []

    print("generation start: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    for i in range(config.sampling.num_sample_batches):
        if config.sampling.semi_ar:
            _, intermediate_samples, _ = model.restore_model_and_semi_ar_sample(
                stride_length=stride_length,
                num_strides=num_strides,
                dt=1 / config.sampling.steps)
            text_samples = intermediate_samples[-1]
            # Note: Samples generated using semi-ar method
            # need to to be processed before computing generative perplexity
            # since these samples contain numerous <|endoftext|> tokens
            # and diffusion.compute_generative_perplexity() discards
            # any text after the first EOS token.
        else:
            assert config.loader.eval_batch_size % config.sampling.duplicate == 0
            different_in_batch = config.loader.eval_batch_size // config.sampling.duplicate
            samples = model.restore_model_and_sample(
                num_steps=config.sampling.steps, duplicate=config.sampling.duplicate)
            model.metrics.record_entropy(samples)
            text_samples = model.tokenizer.batch_decode(samples)
            model.metrics.record_generative_perplexity(
                text_samples, config.model.length, model.device)
            model.metrics.record_tc([i*different_in_batch + j for _ in range(
                config.sampling.duplicate) for j in range(different_in_batch)], samples)
            all_samples.extend(list(text_samples))

    print("generation end: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    generative_ppl = 0.
    entropy = 0.
    if not config.sampling.semi_ar:
        generative_ppl = model.metrics.gen_ppl.compute().item()
        entropy = model.metrics.sample_entropy.compute().item()
        avg_tc, avg_joints, avg_marginals = model.metrics.tc.compute()
        print('Generative perplexity:', generative_ppl)
        print('Sample entropy:', entropy)
        print('Total average correlation:', avg_tc)
        print('Average joint entropy:', avg_joints)
        print('Average marginal entropy:', avg_marginals)
    samples_path = config.eval.generated_samples_path
    with fsspec.open(samples_path, 'w') as f:
        json.dump({'generative_ppl': generative_ppl,
                   'entropy': entropy,
                   'avg_tc': avg_tc,
                   'avg_joints': avg_joints,
                   'avg_marginals': avg_marginals,
                   'generated_seqs': all_samples}, f, indent=4)
    print('Samples saved at:', samples_path)


@torch.inference_mode()
def generate_reflow_dataset(diffusion_model, config, logger, tokenizer):
    # TODO: implement with lightning_module.test with pseudo-data
    logger.info('Generating samples.')
    model = _load_from_checkpoint(diffusion_model=diffusion_model,
                                  config=config,
                                  tokenizer=tokenizer)
    if config.eval.disable_ema:
        logger.info('Disabling EMA.')
        model.ema = None
    # if model.ema:
    #   model.ema.store(itertools.chain(
    #       model.backbone.parameters(),
    #       model.noise.parameters()))
    #   model.ema.copy_to(itertools.chain(
    #       model.backbone.parameters(),
    #       model.noise.parameters()))
    #   model.backbone.eval()
    #   model.noise.eval()

    test_ds = dataloader.get_pseudo_dataloader(config, tokenizer, model)
    trainer = hydra.utils.instantiate(
        config.trainer,
        default_root_dir=os.getcwd(),
        callbacks=None,
        strategy=hydra.utils.instantiate(config.strategy),
        logger=None)
    trainer.test(model, test_ds)
    return


def _eval_ppl(diffusion_model, config, logger, tokenizer):
    logger.info('Starting Perplexity Eval.')

    model = _load_from_checkpoint(
        diffusion_model=diffusion_model,
        config=config,
        tokenizer=tokenizer)
    if config.eval.disable_ema:
        logger.info('Disabling EMA.')
        model.ema = None

    wandb_logger = None
    if config.get('wandb', None) is not None:
        wandb_logger = L.pytorch.loggers.WandbLogger(
            config=omegaconf.OmegaConf.to_object(config),
            ** config.wandb)
    callbacks = []
    if 'callbacks' in config:
        for _, callback in config.callbacks.items():
            callbacks.append(hydra.utils.instantiate(callback))
    trainer = hydra.utils.instantiate(
        config.trainer,
        default_root_dir=os.getcwd(),
        callbacks=callbacks,
        strategy=hydra.utils.instantiate(config.strategy),
        logger=wandb_logger)
    _, valid_ds = dataloader.get_dataloaders(
        config, tokenizer, skip_train=True, valid_seed=config.seed)
    trainer.validate(model, valid_ds)


@torch.inference_mode()
def generate_reflow_dataset_with_perturbed_rect(diffusion_model, config, logger, tokenizer):
    logger.info('Generating samples.')
    model = _load_from_checkpoint(
        diffusion_model=diffusion_model,
        config=config,
        tokenizer=tokenizer)
    if config.eval.disable_ema:
        logger.info('Disabling EMA.')
        model.ema = None

    train_ds, _ = dataloader.get_dataloaders(
        config, tokenizer, skip_valid=True)

    # i is given by random sequence of train_ds's N
    shuffled_indices = np.random.permutation(len(train_ds.dataset))

    eval_batch_size = config.loader.eval_batch_size
    generate_samples = config.sampling.num_reflow_samples

    x0s = []
    xTs = []
    ts = []

    print("generation start: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
    for j in range(generate_samples // eval_batch_size):
        if config.sampling.semi_ar:
            raise NotImplementedError(
                "Semi-AR sampling is not implemented. Please use standard sampling.")
        else:
            assert eval_batch_size == 1
            x0 = train_ds.dataset[shuffled_indices[j *
                                                   eval_batch_size:(j+1)*eval_batch_size]]['input_ids']
            x0 = torch.from_numpy(x0).to(model.device)
            x1 = torch.randint(0, 50258, x0.shape,
                               device=model.device, dtype=x0.dtype)
            rand_t = torch.randint(
                0, x0.shape[1], (1, ), device=model.device).float().item() / x0.shape[1]
            num_step = max(int(config.sampling.steps * (1 - rand_t)), 1)
            # random interpolate between x1 and x0
            # =y_given_t where y0=noise, y1=data
            xt = torch.where(rand_t > torch.rand(
                x1.shape, device=model.device), x0, x1)

            samples = model.restore_model_and_sample(
                num_steps=num_step, xT=xt.clone(), given_t=rand_t)
            x0s.append(samples.clone())
            xTs.append(xt.clone())
            ts.append(rand_t)
        if j % 500 == 0:
            print(f"Generated {(j+1) * eval_batch_size} samples")
    x0s = torch.cat(x0s, dim=0)
    xTs = torch.cat(xTs, dim=0)
    ts = torch.tensor(ts, device=model.device)

    print("generation end: " + datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

    x0s = x0s.cpu().numpy()
    xTs = xTs.cpu().numpy()
    ts = ts.cpu().numpy()

    save_path = config.data.save_dir
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    xT_path = os.path.join(save_path, 'xT.npy')
    x0_path = os.path.join(save_path, 'x0.npy')
    ts_path = os.path.join(save_path, 'ts.npy')

    np.save(x0_path, x0s)
    np.save(xT_path, xTs)
    np.save(ts_path, ts)


def _train(diffusion_model, config, logger, tokenizer):
    logger.info('Starting Training.')
    wandb_logger = None
    if config.get('wandb', None) is not None:
        wid = config.wandb.get('id')
        if not wid or len(str(wid)) > 16:
            wid = str(uuid.uuid4().hex[:8])
        config.wandb.id = wid
        if config.wandb.get('name'):
            config.wandb.name = f"{config.wandb.name}_{wid}"
        wandb_logger = L.pytorch.loggers.WandbLogger(
            config=omegaconf.OmegaConf.to_object(config),
            ** config.wandb)

    if (config.checkpointing.resume_from_ckpt
        and config.checkpointing.resume_ckpt_path is not None
        and utils.fsspec_exists(
            config.checkpointing.resume_ckpt_path)):
        ckpt_path = config.checkpointing.resume_ckpt_path
    else:
        ckpt_path = None

    # Lightning callbacks
    callbacks = []
    if 'callbacks' in config:
        for _, callback in config.callbacks.items():
            callbacks.append(hydra.utils.instantiate(callback))

    train_ds, valid_ds = dataloader.get_dataloaders(
        config, tokenizer)
    _print_batch(train_ds, valid_ds, tokenizer)

    if config.training.finetune_path != '':
        assert utils.fsspec_exists(config.training.finetune_path)
        model = diffusion_model.load_from_checkpoint(
            config.training.finetune_path,
            tokenizer=tokenizer,
            config=config,
            weights_only=False)
    else:
        model = diffusion_model(config, tokenizer=valid_ds.tokenizer)

    trainer = hydra.utils.instantiate(
        config.trainer,
        default_root_dir=os.getcwd(),
        callbacks=callbacks,
        strategy=hydra.utils.instantiate(config.strategy),
        logger=wandb_logger)
    # Force weights_only=False to allow full checkpoint restore (PyTorch 2.6 defaults torch.load to weights_only=True)
    trainer.fit(model, train_ds, valid_ds, ckpt_path=ckpt_path)


@hydra.main(version_base=None, config_path='configs',
            config_name='config')
def main(config):
    """Main entry point for training."""
    L.seed_everything(config.seed)
    _print_config(config, resolve=True, save_cfg=True)

    logger = utils.get_logger(__name__)
    tokenizer = dataloader.get_tokenizer(config)
    if config.algo.name == 'ar':
        diffusion_model = algo.AR
    elif config.algo.name == 'mdlm':
        diffusion_model = algo.MDLM
    elif config.algo.name == 'duo_base':
        diffusion_model = algo.DUO_BASE
    elif config.algo.name == 'duo':
        diffusion_model = algo.DUO
    elif config.algo.name == 'flm':
        diffusion_model = algo.FLM
    elif config.algo.name == 'fmlm_twomodel':
        diffusion_model = algo.FMLM_TwoModel
    elif config.algo.name == 'fmlm_twostage':
        diffusion_model = algo.FMLM_TwoStage
    elif config.algo.name == 'fmlm':
        diffusion_model = algo.FMLM
    elif config.algo.name == 'd3pm':
        diffusion_model = algo.D3PMAbsorb
    elif config.algo.name == 'sedd':
        diffusion_model = algo.SEDDAbsorb
    elif config.algo.name == 'distillation':
        diffusion_model = algo.Distillation
    elif config.algo.name == 'rectification':
        diffusion_model = algo.Rectification
    else:
        raise ValueError(
            f'Invalid algorithm name: {config.algo.name}')
    kwargs = {'diffusion_model': diffusion_model,
              'config': config,
              'tokenizer': tokenizer,
              'logger': logger}
    if config.mode == 'sample_eval':
        _generate_samples(**kwargs)
    elif config.mode == 'sample_eval_recon':
        _generate_samples(**kwargs)
    elif config.mode == 'sample_eval_with_tc':
        _generate_samples_with_tc(**kwargs)
    elif config.mode == 'ppl_eval':
        _eval_ppl(**kwargs)
    elif config.mode == 'generate_reflow_data':
        generate_reflow_dataset(diffusion_model, config, logger, tokenizer)
    elif config.mode == 'generate_reflow_data_with_perturbed_rect':
        generate_reflow_dataset_with_perturbed_rect(**kwargs)
    else:
        _train(**kwargs)


if __name__ == '__main__':
    # allow_tf32
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    main()