DOCSEARCHbyAlgolia
Submit Docs

BD3-LMs (Block Discrete Denoising Diffusion Language Models)

repository·main·Indexed 21 days ago

https://github.com/kuleshov-group/bd3lms

A framework that interpolates between autoregressive and diffusion language models by performing discrete diffusion within blocks of tokens. BD3-LMs enables state-of-the-art likelihoods and arbitrary-length sequence generation. The repository includes implementations for training, likelihood (perplexity) evaluation, and sequence generation, as well as a codebase for the SSD-LM baseline.

Tokens
3K
Snippets
13
Records
14
Agent score
77%

What's inside BD3-LMs

  1. BD3-LM Code Organization

    main

    The repository is structured as follows:

    • main.py: Primary entry point for training and evaluation.
    • noise_schedule.py: Implementation of noise schedules.
    • diffusion.py: Forward and reverse diffusion logic.
    • dataloader.py: Data loading routines.
    • utils.py: Learning rate schedulers, logging, and fsspec handling.
    • models/: Network architectures (supports DiT and AR transformer).
    • configs/: Configuration files for datasets, models, noise schedules, and LR schedules.
    • scripts/: Shell scripts for various tasks:
      • train/: Training scripts (LM1B, OWT).
      • ppl/: Likelihood evaluation.
      • zs_ppl/: Zero-shot likelihood evaluation.
      • gen_ppl/: Sample quality evaluation.
      • var_len/: Arbitrary-length sequence generation.
    • ssd-lm/: Codebase for the SSD-LM baseline.

  2. Train the SSD-LM model

    main

    Distributed training is managed via Slurm. Use sbatch to submit the training template.

    • Resuming Training: If a job is interrupted, resubmitting the same command will automatically resume from the most recent recoverable checkpoint.
    • Starting from Scratch: To restart training completely, manually delete the existing output directory (logging/ssd_dbs25) before resubmitting.
    • Manual Checkpoints: You can manually download checkpoints from Hugging Face and place them in logging/ssd_dbs25.
    sbatch submit_template_ssd_model_train.sbatch

  3. Perform controlled generation evaluation

    main

    Controlled generation is performed using off-the-shelf sentiment classifiers.

    1. Generate controlled continuations: source loop_eval_ctrsa.sh
    2. Score the generations: source loop_scoring_eval_ctrsa.sh

    Results are saved to files ending with _ssd_ctrsa_eval.txt.

    source loop_eval_ctrsa.sh
source loop_scoring_eval_ctrsa.sh

  4. Set up the SSD-LM environment

    main

    The project uses Conda for environment management. An environment.yml file is provided to create an environment named ssdlm.

    Note on Slurm and Paths: If using a Slurm system, you must update the shell (.sh) and Slurm (.sbatch) scripts to replace any hardcoded home directory paths (e.g., /private/home/xhan77) with your own. For Slurm files, ensure you provide your specific partition name, device constraints, and job time limits. The default scripts assume 8 GPUs per compute node.

    conda env create -f environment.yml

  6. Install and set up BD3-LMs

    main

    To use the BD3-LMs framework, create a Conda environment with Python 3.9 and install the required dependencies via pip. Note that while BD3-LMs do not require FlashAttention, evaluating the MDLM baseline requires flash-attn==2.5.6.

    You must also create specific directories for outputs, logs, and samples before running training or evaluation scripts.

    conda create --name bd3lm python=3.9
conda activate bd3lm
pip install -r requirements.txt

mkdir outputs watch_folder logs sample_logs

  7. Generate arbitrary-length sequences with BD3-LM

    main

    You can generate sequences of arbitrary length by setting mode=sample_eval. The model.length must be a multiple of the block_size.

    When using a model from HuggingFace, provide the model ID to eval.checkpoint_path. For local models, provide the absolute path. Use sampling.nucleus_p for nucleus sampling and sampling.kv_cache=true for efficiency.

    BLOCK_SIZE=4 # 4, 8, 16
LENGTH=2048 # arbitrary; needs to be a multiple of the block size

python -u main.py \
    loader.eval_batch_size=1 \
    model=small \
    algo=bd3lm \
    algo.T=5000 \
    algo.backbone=hf_dit \
    data=openwebtext-split \
    model.length=$LENGTH \
    block_size=$BLOCK_SIZE \
    wandb=null \
    mode=sample_eval \
    eval.checkpoint_path=kuleshov-group/bd3lm-owt-block_size${BLOCK_SIZE} \
    model.attn_backend=sdpa \
    sampling.nucleus_p=0.9 \
    sampling.kv_cache=true \
    sampling.logdir=$PWD/sample_logs/samples_genlen_bd3lm_blocksize${BLOCK_SIZE}

  8. Train BD3-LM models

    main

    Training is performed using mode=train (the default).

    Key configuration details:

    • block_size: Recommended values are 4, 8, or 16. It must be a factor of the context length.
    • training.from_pretrained: Set this to a checkpoint path (e.g., a pretraining checkpoint) to continue training. Set to null to train from scratch.
    • loader.batch_size and loader.eval_batch_size: Controls batch size per GPU. If loader.batch_size * num_gpus is less than loader.global_batch_size, PyTorch Lightning will use gradient accumulation.
    • algo.clip_search_widths: Used for the training algorithm.
    BLOCK_SIZE=4 # we recommend 4, 8, or 16. must be a factor of the context length
PRETRAIN_CKPT=kuleshov-group/bd3lm-owt-block_size1024-pretrain # to train from scratch, set to null

python -u main.py \
    loader.global_batch_size=512 \
    loader.eval_global_batch_size=512 \
    loader.batch_size=16 \
    loader.eval_batch_size=16 \
    model=small \
    algo=bd3lm \
    algo.clip_search_widths=[0.5,0.6,0.7,0.8,0.9] \
    data=openwebtext-split \
    model.length=1024 \
    block_size=$BLOCK_SIZE \
    wandb.name=bd3lm-owt-block_size${BLOCK_SIZE} \
    mode=train \
    model.attn_backend=flex \
    training.resample=True \
    training.from_pretrained=$PRETRAIN_CKPT

  9. Perform unconstrained generation evaluation

    main

    To evaluate unconstrained (prompted) generation, use the following scripts. Results are saved to the model directory logging/ssd_dbs25.

    • Multi-hot projection: source loop_eval.sh
    • Sampling projection: source loop_eval_alt.sh

    To score the generated continuations, run: source loop_scoring_eval.sh

    Results are output to files ending in _ssd_eval.txt or _ssd_eval_sampling.txt. To run a GPT-2 baseline, use source loop_baseline_gpt2.sh.

    source loop_eval.sh
source loop_scoring_eval.sh

  10. Evaluate likelihood (perplexity) with BD3-LM

    main

    To compute test perplexity, use mode=ppl_eval. This mode evaluates the model on a dataset (e.g., openwebtext-split) using a specified model.length and block_size.

    BLOCK_SIZE=4 # 4, 8, 16

python -u main.py \
    loader.eval_batch_size=16 \
    model=small \
    algo=bd3lm \
    algo.backbone=hf_dit \
    data=openwebtext-split \
    data.insert_valid_special=False \
    model.length=1024 \
    model.attn_backend=flex \
    block_size=${BLOCK_SIZE} \
    eval.checkpoint_path=kuleshov-group/bd3lm-owt-block_size${BLOCK_SIZE} \
    wandb=null \
    mode=ppl_eval > logs/bd3lm_owt_block_size${BLOCK_SIZE}.log

  11. Run the Block Diffusion pipeline via main.py

    main

    The main.py file serves as the primary entrypoint for the project, using Hydra for configuration management. Depending on the config.mode value, it executes one of three primary workflows:

    1. Training: The default mode. Initializes a diffusion.Diffusion model and uses PyTorch Lightning to run the training loop.
    2. sample_eval: Loads a model from a checkpoint and generates text samples, calculating generative perplexity, entropy, and saving results to a CSV.
    3. ppl_eval: Loads a model from a checkpoint and performs perplexity evaluation on a validation dataset.

    To run the script, you must provide a Hydra configuration. The script automatically registers several OmegaConf resolvers (cwd, device_count, eval, div_up) to allow dynamic configuration values.

    # Example execution (assuming hydra configs are set up)
python main.py mode=sample_eval eval.checkpoint_path=path/to/model.ckpt

  12. Generate text samples with generate_samples()

    main

    Use generate_samples to perform inference and evaluate generative metrics. It loads the model from the path specified in config.eval.checkpoint_path and uses the model.restore_model_and_sample method.

    Key behaviors:

    • If config.eval.disable_ema is set, Exponential Moving Average weights are ignored.
    • It calculates and prints Generative Perplexity and Entropy.
    • It saves a CSV containing gen_ppl, gen_nfes, gen_entropy, gen_lengths, and the samples themselves to the directory specified in config.sampling.logdir.
    • If config.sampling.var_length is enabled, the actual text samples are omitted from the CSV to save space/handle variable lengths.
    # Conceptual usage within the pipeline
samples = generate_samples(config, logger, tokenizer)