### Install Dependencies and Setup Directories

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Installs project dependencies using pip and creates necessary output directories for logs and samples. Ensure you are in a Python 3.9 conda environment.

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

mkdir outputs watch_folder logs sample_logs
```

--------------------------------

### Start Unconstrained Generation (Multi-hot Projection)

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Initiates unconstrained generation using the multi-hot projection method. Results are saved to the model directory.

```bash
source loop_eval.sh
```

--------------------------------

### Create Conda Environment and Install Dependencies

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Sets up a new conda environment named 'bd3lm' with Python 3.9 and installs project dependencies from 'requirements.txt'.

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

--------------------------------

### Start Unconstrained Generation (Sampling Projection)

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Initiates unconstrained generation using the sampling projection method (indicated by 'alt' suffix). Results are saved to the model directory.

```bash
source loop_eval_alt.sh
```

--------------------------------

### Submit SSD-LM Model Training Job

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Submits a distributed training job for SSD-LM using Slurm. The job will save checkpoints to 'logging/ssd_dbs25' and can automatically resume from the latest checkpoint if interrupted. Ensure the output directory is empty if starting from scratch.

```bash
sbatch submit_template_ssd_model_train.sbatch
```

--------------------------------

### Create Directories for Outputs and Logs

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Creates necessary directories for storing model outputs, watch files, and logs before running training jobs.

```bash
mkdir outputs watch_folder logs sample_logs
```

--------------------------------

### Train BD3-LMs on OpenWebText

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Use this script to train BD3-LMs with OpenWebText. Adjust BLOCK_SIZE for different context lengths. Set PRETRAIN_CKPT to null to train from scratch.

```bash
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

```

--------------------------------

### Process Data with SSD-LM

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Downloads and processes the OpenWebText dataset. Ensure the year argument is correctly specified.

```bash
bash run_ssd_process_data.sh 2022
```

--------------------------------

### Initialize Diffusion LightningModule

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Instantiate the `Diffusion` module with a configuration and tokenizer. This sets up the backbone model, noise schedule, and EMA. Ensure the configuration and block size are correctly set.

```python
import transformers
import omegaconf
import diffusion as diff_module
import dataloader

# Load config (normally done via Hydra)
config = omegaconf.OmegaConf.load('configs/config.yaml')
config.algo.backbone = 'dit'
config.block_size = 16

tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion(config, tokenizer=tokenizer)

# Check model properties
print(model.block_size)          # 16
print(model.parameterization)    # 'subs'
print(model.sampler)             # 'semi_ar'
print(model.ema.decay)           # 0.9999
```

--------------------------------

### Diffusion Module Initialization

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Instantiates the Diffusion LightningModule with a configuration and tokenizer. It handles training and sampling logic.

```APIDOC
## `Diffusion` — PyTorch Lightning Training Module

`Diffusion` is the central `LightningModule` that wraps the backbone model, noise schedule, EMA, and all training/sampling logic. It is instantiated with a Hydra config and a HuggingFace tokenizer, and handles all Lightning lifecycle hooks.

```python
import transformers
import omegaconf
import diffusion as diff_module
import dataloader

# Load config (normally done via Hydra)
config = omegaconf.OmegaConf.load('configs/config.yaml')
config.algo.backbone = 'dit'
config.block_size = 16

tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion(config, tokenizer=tokenizer)

# Check model properties
print(model.block_size)          # 16
print(model.parameterization)    # 'subs'
print(model.sampler)             # 'semi_ar'
print(model.ema.decay)           # 0.9999
```
```

--------------------------------

### Configure and Instantiate BD3LM Model

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Builds a `BD3LMConfig` object with specified architectural and inference hyperparameters, then instantiates the `BD3LM` model. The config is compatible with HuggingFace's Auto API for saving and loading.

```python
from models.hf import BD3LMConfig, BD3LM
import transformers

# Build config for a 12-layer, 768-dim model with block size 16
config = BD3LMConfig(
    block_size=16,            # tokens denoised in parallel per block
    vocab_size=50258,         # GPT2 vocab + 1 mask token
    model_length=1024,        # training context length
    hidden_dim=768,           # transformer hidden dimension
    cond_dim=128,             # conditioning (sigma) embedding dim
    n_blocks=12,              # number of transformer layers
    n_heads=12,               # attention heads
    dropout=0.1,
    attn_backend='sdpa',      # 'sdpa' for inference, 'flex' for training
    cross_attn=True,          # cross-attention between xt and x0
    adaln=True,               # adaptive LayerNorm conditioning
    time_conditioning=False,  # disable sigma time-conditioning
    var_min=True,             # enable variance-minimizing clipped schedule
    sampling_eps_min=1e-3,    # lower clip bound for noise schedule
    sampling_eps_max=0.999,   # upper clip bound for noise schedule
    return_dict=False
)

# Instantiate model
model = BD3LM(config)
print(model.config.block_size)  # 16
```

--------------------------------

### Generate Text Samples with Diffusion.restore_model_and_sample

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Generate text samples using the `restore_model_and_sample` method. This method loads a checkpoint, optionally uses EMA weights, and runs the configured sampler. It also computes generative perplexity metrics. Ensure `attn_backend` is set to 'sdpa' for sampling.

```python
import torch
import transformers
import omegaconf
import diffusion as diff_module
import dataloader

config = omegaconf.OmegaConf.load('configs/config.yaml')
config.mode = 'sample_eval'
config.model.attn_backend = 'sdpa'   # required for sampling
config.loader.eval_batch_size = 2
config.sampling.num_sample_batches = 1
config.sampling.nucleus_p = 0.9
config.sampling.first_hitting = True
config.sampling.kv_cache = True
config.sampling.var_length = False

tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion.load_from_checkpoint(
    'checkpoints/last.ckpt',
    tokenizer=tokenizer,
    config=config,
    strict=False
).to('cuda')

# Generate 1024-token sequences using 5000 diffusion steps
samples = model.restore_model_and_sample(num_steps=5000)
for i, text in enumerate(samples):
    print(f"Sample {i}: {text[:100]}...")

# Retrieve generative perplexity (under GPT2-large)
gen_ppl = model.metrics.gen_ppl.compute()
print(f"Generative perplexity: {gen_ppl:.2f}")
```

--------------------------------

### Generate SSD-LM Samples

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Executes a shell script to generate batch samples using the SSD-LM codebase.

```bash
ssd-lm/run_generate_text_batch.sh
```

--------------------------------

### Load Tokenizer with get_tokenizer

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Loads and configures a tokenizer, supporting GPT-2, BERT, and custom tokenizers. It normalizes BOS/EOS/PAD tokens for sequence handling. Demonstrates encoding and decoding a sample string.

```python
import omegaconf
from dataloader import get_tokenizer

config = omegaconf.OmegaConf.load('configs/config.yaml')
# config.data.tokenizer_name_or_path is 'gpt2' by default

tokenizer = get_tokenizer(config)
print(tokenizer.bos_token)   # '<|endoftext|>
print(tokenizer.eos_token)   # '<|endoftext|>
print(tokenizer.pad_token)   # '[PAD]' (added if missing)

# Encode and decode a sample
ids = tokenizer.encode("Hello, world!")
print(ids)
print(tokenizer.decode(ids))
```

--------------------------------

### Apply Forward Diffusion (Noising) with Diffusion.q_xt

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

The `q_xt` method applies the forward diffusion process, masking tokens in a clean sequence `x0` with a given probability `p`. It can optionally enforce constraints on the fraction of masked tokens per block using `sampling_eps_min` and `sampling_eps_max`.

```python
import torch
import omegaconf
import diffusion as diff_module
import dataloader

config = omegaconf.OmegaConf.load('configs/config.yaml')
tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion(config, tokenizer=tokenizer).to('cuda')

batch_size, seq_len = 4, 1024
x0 = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_len), device='cuda')

# p: masking probability per token, shape (batch_size, 1)
p = torch.full((batch_size, 1), 0.3, device='cuda')

xt = model.q_xt(
    x=x0,
    p=p,
    block_size=model.block_size,   # tokens masked per block
    sampling_eps_min=0.1,          # at least 10% of block tokens masked
    sampling_eps_max=0.9           # at most 90% of block tokens masked
)
# xt has the same shape as x0; masked positions contain model.mask_index
num_masked = (xt == model.mask_index).float().mean()
print(f"Fraction masked: {num_masked:.3f}")   # ~0.3
```

--------------------------------

### Distributed BD3LM Training with Slurm

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Provides a bash script for launching distributed BD3LM training on Slurm. It configures hyperparameters as Hydra overrides for `main.py`, including dataset, model size, block size, and pre-training checkpoint.

```bash
# Single-node 8-GPU training on OpenWebText, block size 16,
# fine-tuned from the MDLM pre-train checkpoint
BLOCK_SIZE=16
PRETRAIN_CKPT=kuleshov-group/bd3lm-owt-block_size1024-pretrain

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 \
    data.insert_train_special=False \
    data.insert_valid_special=False \
    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}

# Or submit to Slurm:
sbatch scripts/train/train_owt_bd3lm.sh
```

--------------------------------

### get_dataloaders — Dataset and DataLoader Construction

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Creates tokenized, chunked DataLoader instances for training and validation. Supports streaming, data wrapping (concatenating texts into fixed-length blocks with BOS/EOS), and fault-tolerant distributed sampling for Slurm/multi-GPU environments.

```APIDOC
## get_dataloaders

### Description
Creates tokenized, chunked `DataLoader` instances for training and validation. Supports streaming, data wrapping (concatenating texts into fixed-length blocks with BOS/EOS), and fault-tolerant distributed sampling for Slurm/multi-GPU environments.

### Usage
```python
import omegaconf
from dataloader import get_tokenizer, get_dataloaders

config = omegaconf.OmegaConf.load('configs/config.yaml')
# Override to use wikitext2 for a quick test
config.data.train = 'wikitext2'
config.data.valid = 'wikitext2'
config.data.wrap = True
config.model.length = 1024
config.loader.batch_size = 4
config.loader.eval_batch_size = 4
config.loader.num_workers = 2
config.loader.pin_memory = False

tokenizer = get_tokenizer(config)
train_loader, valid_loader = get_dataloaders(config, tokenizer)

batch = next(iter(train_loader))
print(batch['input_ids'].shape)        # torch.Size([4, 1024])
print(batch['attention_mask'].shape)   # torch.Size([4, 1024])
# First token is always BOS, last is always EOS
print(tokenizer.decode([batch['input_ids'][0, 0].item()]))   # <|endoftext|>
```
```

--------------------------------

### Create Datasets and DataLoaders with get_dataloaders

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Constructs tokenized and chunked DataLoaders for training and validation. Supports streaming, data wrapping, and fault-tolerant distributed sampling. Shows batch shape and the first token being BOS.

```python
import omegaconf
from dataloader import get_tokenizer, get_dataloaders

config = omegaconf.OmegaConf.load('configs/config.yaml')
# Override to use wikitext2 for a quick test
config.data.train = 'wikitext2'
config.data.valid = 'wikitext2'
config.data.wrap = True
config.model.length = 1024
config.loader.batch_size = 4
config.loader.eval_batch_size = 4
config.loader.num_workers = 2
config.loader.pin_memory = False

tokenizer = get_tokenizer(config)
train_loader, valid_loader = get_dataloaders(config, tokenizer)

batch = next(iter(train_loader))
print(batch['input_ids'].shape)        # torch.Size([4, 1024])
print(batch['attention_mask'].shape)   # torch.Size([4, 1024])
# First token is always BOS, last is always EOS
print(tokenizer.decode([batch['input_ids'][0, 0].item()]))   # <|endoftext|>
```

--------------------------------

### Generate Arbitrary-Length Sequences (HuggingFace Model)

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Use this script to generate sequences longer than the training context size when using a HuggingFace checkpoint. Ensure BLOCK_SIZE is a multiple of the model's block size and LENGTH is a multiple of BLOCK_SIZE.

```bash
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}
```

--------------------------------

### Train OWT BD3LM Model

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Submits a batch job to train the OWT BD3LM model using the provided script.

```bash
sbatch scripts/train/train_owt_bd3lm.sh
```

--------------------------------

### Likelihood Evaluation on OpenWebText

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

This script computes test perplexity for BD3-LM on the OpenWebText dataset. Ensure BLOCK_SIZE is set appropriately (e.g., 4, 8, 16). The output logs will be saved to `logs/bd3lm_owt_block_size${BLOCK_SIZE}.log`.

```bash
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
```

--------------------------------

### Generate Arbitrary-Length Sequences (Local Checkpoint)

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Use this script to generate sequences longer than the training context size when using a local checkpoint. Ensure BLOCK_SIZE is a multiple of the model's block size and LENGTH is a multiple of BLOCK_SIZE.

```bash
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 \
    data=openwebtext-split \
    model.length=$LENGTH \
    block_size=$BLOCK_SIZE \
    wandb=null \
    mode=sample_eval \
    eval.checkpoint_path=/path/to/checkpoint/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}
```

--------------------------------

### Training — Distributed BD3LM Training with Slurm

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Launch distributed training via Slurm using the provided shell scripts. All hyperparameters are Hydra overrides passed directly to `main.py`.

```APIDOC
## Training

### Description
Launch distributed training via Slurm using the provided shell scripts. All hyperparameters are Hydra overrides passed directly to `main.py`.

### Usage
```bash
# Single-node 8-GPU training on OpenWebText, block size 16,
# fine-tuned from the MDLM pre-train checkpoint
BLOCK_SIZE=16
PRETRAIN_CKPT=kuleshov-group/bd3lm-owt-block_size1024-pretrain

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 \
    data.insert_train_special=False \
    data.insert_valid_special=False \
    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}

# Or submit to Slurm:
sbatch scripts/train/train_owt_bd3lm.sh
```
```

--------------------------------

### Perplexity Evaluation with `mode=ppl_eval`

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Evaluates the NELBO as a perplexity estimate on the validation set using EMA parameters and the Lightning Trainer. Redirects output to a log file.

```bash
BLOCK_SIZE=4

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

# Slurm variant:
sbatch scripts/ppl/ppl_owt_bd3lm.sh
```

--------------------------------

### Sample Generation with `mode=sample_eval`

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Generates text samples and computes generative perplexity and token entropy. Supports fixed-length and variable-length generation with optional nucleus sampling and KV caching.

```bash
# Fixed-length generation (1024 tokens, block size 4)
BLOCK_SIZE=4
python -u main.py \
    loader.eval_batch_size=8 \
    model=small \
    algo=bd3lm \
    algo.T=5000 \
    algo.backbone=hf_dit \
    data=openwebtext-split \
    model.length=1024 \
    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=./sample_logs/bd3lm_block${BLOCK_SIZE}
```

```bash
# Variable-length generation (up to 2048 tokens)
LENGTH=2048
BLOCK_SIZE=16
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.var_length=true \
    sampling.logdir=./varlen_sample_logs/bd3lm_block${BLOCK_SIZE}
```

--------------------------------

### Publishing Models to HuggingFace Hub

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Loads a Lightning checkpoint, extracts EMA parameters, and pushes the BD3LM model with its custom config to the HuggingFace Hub. The model can then be loaded via `AutoModelForMaskedLM`.

```python
import torch
import transformers
from models.hf import BD3LMConfig, BD3LM

BLOCK_SIZE = 16
BD3LMConfig.register_for_auto_class()
BD3LM.register_for_auto_class('AutoModelForMaskedLM')

tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
name_or_path = f'your-org/bd3lm-owt-block_size{BLOCK_SIZE}'

config = BD3LMConfig(
    block_size=BLOCK_SIZE,
    vocab_size=tokenizer.vocab_size + 1,  # +1 for mask token
    model_length=1024,
    hidden_dim=768,
    cond_dim=128,
    n_blocks=12,
    n_heads=12,
    dropout=0.1,
    time_conditioning=False,
    return_dict=False
)

model = BD3LM(config).to('cuda')
ckpt = torch.load('checkpoints/last.ckpt', weights_only=False)
model.load_state_dict(ckpt['state_dict'])

# Copy EMA shadow params into model weights before pushing
for s_param, param in zip(ckpt['ema']['shadow_params'], model.parameters()):
    if param.requires_grad:
        param.data.copy_(s_param.data)

model.push_to_hub(name_or_path, private=True)
# Model is now loadable via:
# transformers.AutoModelForMaskedLM.from_pretrained(name_or_path, trust_remote_code=True)
```

--------------------------------

### Build Default Log-Linear Schedule

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Initializes and uses a LogLinearNoise schedule to determine loss scaling and masking probabilities based on the diffusion process time `t`. The `sigma_max` value represents the maximum noise level.

```python
noise = LogLinearNoise(eps=1e-3)

# t in [0, 1]: fraction through the diffusion process
t = torch.linspace(0, 1, steps=11)

# Returns (loss_scaling, move_chance = masking probability)
loss_scale, p = noise(t)
print("t:           ", t.tolist())
print("move_chance: ", p.round(decimals=3).tolist())
# move_chance goes from ~0.001 (barely noised) to ~1.0 (fully masked)

# sigma = -log(1 - p); sigma_max at t=1
sigma_max = noise.sigma_max
print(f"sigma_max: {sigma_max:.4f}")   # ~6.908

# Using the factory function driven by a config
import omegaconf
config = omegaconf.OmegaConf.load('configs/config.yaml')
noise_from_config = get_noise(config)   # returns LogLinearNoise by default
```

--------------------------------

### Run GPT-2 Baseline Evaluation

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Executes the evaluation script to obtain baseline results using GPT-2.

```bash
source loop_baseline_gpt2.sh
```

--------------------------------

### `BD3LMConfig` — HuggingFace Configuration Class

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

The `BD3LMConfig` class extends HuggingFace's `PretrainedConfig` and manages all hyperparameters for the BD3LM model architecture and inference. It is compatible with HuggingFace's Auto API for saving and loading configurations from the Hub.

```APIDOC
## `BD3LMConfig` — HuggingFace Configuration Class

`BD3LMConfig` extends `transformers.PretrainedConfig` and holds all architectural and inference hyperparameters for the BD3LM model. It is registered with the HuggingFace Auto API and can be used for saving/loading from the Hub.

```python
from models.hf import BD3LMConfig, BD3LM
import transformers

# Build config for a 12-layer, 768-dim model with block size 16
config = BD3LMConfig(
    block_size=16,            # tokens denoised in parallel per block
    vocab_size=50258,         # GPT2 vocab + 1 mask token
    model_length=1024,        # training context length
    hidden_dim=768,           # transformer hidden dimension
    cond_dim=128,             # conditioning (sigma) embedding dim
    n_blocks=12,              # number of transformer layers
    n_heads=12,               # attention heads
    dropout=0.1,
    attn_backend='sdpa',      # 'sdpa' for inference, 'flex' for training
    cross_attn=True,          # cross-attention between xt and x0
    adaln=True,               # adaptive LayerNorm conditioning
    time_conditioning=False,  # disable sigma time-conditioning
    var_min=True,             # enable variance-minimizing clipped schedule
    sampling_eps_min=1e-3,    # lower clip bound for noise schedule
    sampling_eps_max=0.999,   # upper clip bound for noise schedule
    return_dict=False
)

# Instantiate model
model = BD3LM(config)
print(model.config.block_size)  # 16
```
```

--------------------------------

### Noise Schedule Factory and Classes

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

The `get_noise` factory function and schedule classes (e.g., `LogLinearNoise`, `CosineNoise`) control the forward diffusion process parameters like masking probability and loss weighting. The `LogLinearNoise` schedule is recommended.

```python
import torch
from noise_schedule import get_noise, LogLinearNoise
```

--------------------------------

### Generate Controlled Continuations (Sentiment Analysis)

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Generates controlled continuations based on off-the-shelf sentiment classifiers. This script is for the 'ctrsa' experiment.

```bash
source loop_eval_ctrsa.sh
```

--------------------------------

### get_tokenizer — Tokenizer Loader

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Loads and configures the tokenizer for the chosen dataset. Handles GPT-2, BERT, and custom text8 tokenizers, normalizing BOS/EOS/PAD tokens for correct sequence wrapping.

```APIDOC
## get_tokenizer

### Description
Loads and configures the tokenizer for the chosen dataset. Handles GPT-2, BERT, and custom text8 tokenizers, normalizing BOS/EOS/PAD tokens for correct sequence wrapping.

### Usage
```python
import omegaconf
from dataloader import get_tokenizer

config = omegaconf.OmegaConf.load('configs/config.yaml')
# config.data.tokenizer_name_or_path is 'gpt2' by default

tokenizer = get_tokenizer(config)
print(tokenizer.bos_token)   # '<|endoftext|>' 
print(tokenizer.eos_token)   # '<|endoftext|>' 
print(tokenizer.pad_token)   # '[PAD]' (added if missing)

# Encode and decode a sample
ids = tokenizer.encode("Hello, world!")
print(ids)
print(tokenizer.decode(ids))
```
```

--------------------------------

### Initialize KV-Cache for BD3LM Sampling

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Initializes the KV cache in all transformer blocks of a BD3LM model to enable efficient semi-autoregressive sampling. This reduces memory access overhead during inference.

```python
import transformers

BLOCK_SIZE = 4
model = transformers.AutoModelForMaskedLM.from_pretrained(
    f'kuleshov-group/bd3lm-owt-block_size{BLOCK_SIZE}',
    trust_remote_code=True
).eval().to('cuda')

# Initialize KV cache for a batch size of 2
model.reset_kv_cache(eval_batch_size=2)

# After initialization, each DDiTBlock has a pre-allocated cache tensor:
# shape: (eval_batch_size, model_length, hidden_dim * 3)
block = model.backbone.blocks[0]
print(block.kv_cache.shape)   # torch.Size([2, 1024, 2304])
print(block.cache_idx)        # 0
```

--------------------------------

### Compute Log-Probabilities with Diffusion.forward

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Use the `forward` method to run the backbone network on noisy token IDs and noise levels to obtain normalized log-probabilities. This is useful for calculating scores during training or inference. Ensure the model is in evaluation mode and tensors are on the correct device.

```python
import torch
import transformers
import omegaconf
import diffusion as diff_module
import dataloader

config = omegaconf.OmegaConf.load('configs/config.yaml')
tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion(config, tokenizer=tokenizer).to('cuda')
model.eval()

batch_size, seq_len = 2, 1024
# Noisy token sequence (some tokens replaced with mask_index)
x = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_len), device='cuda')
x[:, 100:110] = model.mask_index   # simulate masked tokens

# Sigma: noise level per token (shape: batch x seq_len)
sigma = torch.full((batch_size, seq_len), 0.5, device='cuda')

with torch.no_grad():
    log_probs = model.forward(x, sigma)
# log_probs shape: (batch_size, seq_len, vocab_size)
print(log_probs.shape)       # torch.Size([2, 1024, 50258])
print(log_probs.exp().sum(-1).mean().item())  # ~1.0 (probability simplex)
```

--------------------------------

### `BD3LM.reset_kv_cache` — KV-Cache Initialization for Sampling

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

This method initializes the KV cache within all transformer blocks of the BD3LM model. This is crucial for enabling autoregressive-style KV caching during semi-autoregressive sampling, which helps reduce memory access overhead during inference.

```APIDOC
## `BD3LM.reset_kv_cache` — KV-Cache Initialization for Sampling

Initializes the KV cache in all transformer blocks to enable autoregressive-style KV caching during BD3LM semi-autoregressive sampling, reducing memory access overhead at inference time.

```python
import transformers

BLOCK_SIZE = 4
model = transformers.AutoModelForMaskedLM.from_pretrained(
    f'kuleshov-group/bd3lm-owt-block_size{BLOCK_SIZE}',
    trust_remote_code=True
).eval().to('cuda')

# Initialize KV cache for a batch size of 2
model.reset_kv_cache(eval_batch_size=2)

# After initialization, each DDiTBlock has a pre-allocated cache tensor:
# shape: (eval_batch_size, model_length, hidden_dim * 3)
block = model.backbone.blocks[0]
print(block.kv_cache.shape)   # torch.Size([2, 1024, 2304])
print(block.cache_idx)        # 0
```
```

--------------------------------

### Perform Forward Pass with BD3LM Model

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Loads a pre-trained BD3LM model from HuggingFace and performs a forward pass with noisy token IDs and optional timesteps. The model returns per-token logits, which can be used to predict tokens.

```python
import torch
import transformers
from models.hf import BD3LM, BD3LMConfig

BLOCK_SIZE = 4
tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')

# Load a pre-trained model from HuggingFace
model = transformers.AutoModelForMaskedLM.from_pretrained(
    f'kuleshov-group/bd3lm-owt-block_size{BLOCK_SIZE}',
    trust_remote_code=True
).eval().to('cuda')

# Prepare a batch of token IDs (e.g., noisy tokens at some diffusion timestep)
tokens = tokenizer("The quick brown fox", return_tensors='pt')['input_ids'].to('cuda')
timesteps = torch.zeros(tokens.shape[0], device='cuda')  # sigma=0 → clean

with torch.no_grad():
    # Returns logits of shape (batch, seq_len, vocab_size)
    logits = model(input_ids=tokens, timesteps=timesteps)
    predicted_tokens = logits.argmax(-1)
    print(tokenizer.decode(predicted_tokens[0]))
```

--------------------------------

### Score Generated Continuations

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Scores the generated continuations from unconstrained generation. Results are saved in files ending with '_ssd_eval.txt' and '_ssd_eval_sampling.txt'.

```bash
source loop_scoring_eval.sh
```

--------------------------------

### Diffusion.restore_model_and_sample — Sample Generation

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Generates text samples from the model using EMA weights and a configured sampler, optionally computing generative perplexity metrics.

```APIDOC
## `Diffusion.restore_model_and_sample` — Sample Generation

Generates text samples from the model by optionally swapping in EMA weights, setting the backbone to eval mode, and running the configured sampler (`semi_ar` or `analytic`). Automatically computes and stores generative perplexity metrics.

```python
import torch
import transformers
import omegaconf
import diffusion as diff_module
import dataloader

config = omegaconf.OmegaConf.load('configs/config.yaml')
config.mode = 'sample_eval'
config.model.attn_backend = 'sdpa'   # required for sampling
config.loader.eval_batch_size = 2
config.sampling.num_sample_batches = 1
config.sampling.nucleus_p = 0.9
config.sampling.first_hitting = True
config.sampling.kv_cache = True
config.sampling.var_length = False

tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion.load_from_checkpoint(
    'checkpoints/last.ckpt',
    tokenizer=tokenizer,
    config=config,
    strict=False
).to('cuda')

# Generate 1024-token sequences using 5000 diffusion steps
samples = model.restore_model_and_sample(num_steps=5000)
for i, text in enumerate(samples):
    print(f"Sample {i}: {text[:100]}...")

# Retrieve generative perplexity (under GPT2-large)
gen_ppl = model.metrics.gen_ppl.compute()
print(f"Generative perplexity: {gen_ppl:.2f}")
```
```

--------------------------------

### Exponential Moving Average for Model Parameters

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Implements Exponential Moving Average (EMA) to track model parameters for stable evaluation. It allows storing, copying, and restoring weights, ensuring EMA is used only during evaluation.

```python
import torch
import torch.nn as nn
from models.ema import ExponentialMovingAverage

model = nn.Linear(128, 64)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

# Initialize EMA with decay=0.9999
ema = ExponentialMovingAverage(model.parameters(), decay=0.9999)

# Training loop
for step in range(100):
    loss = model(torch.randn(8, 128)).sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    ema.update(model.parameters())  # update shadow params after each step

# Evaluation: swap in EMA weights, validate, then restore
ema.store(model.parameters())       # save current weights
ema.copy_to(model.parameters())     # overwrite with EMA weights

with torch.no_grad():
    val_out = model(torch.randn(4, 128))
    print("Val output shape:", val_out.shape)

ema.restore(model.parameters())     # restore original weights for continued training
print("Restored original weights.")
```

--------------------------------

### Diffusion.q_xt — Forward Diffusion (Noising)

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Applies the forward diffusion process to clean token IDs by masking tokens independently with a given probability.

```APIDOC
## `Diffusion.q_xt` — Forward Diffusion (Noising)

Applies the forward diffusion process to clean token IDs `x0` by masking tokens independently with probability `p`. Optionally resamples to enforce that the fraction of masked tokens per block stays within `[sampling_eps_min, sampling_eps_max]`.

```python
import torch
import omegaconf
import diffusion as diff_module
import dataloader

config = omegaconf.OmegaConf.load('configs/config.yaml')
tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion(config, tokenizer=tokenizer).to('cuda')

batch_size, seq_len = 4, 1024
x0 = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_len), device='cuda')

# p: masking probability per token, shape (batch_size, 1)
p = torch.full((batch_size, 1), 0.3, device='cuda')

xt = model.q_xt(
    x=x0,
    p=p,
    block_size=model.block_size,   # tokens masked per block
    sampling_eps_min=0.1,          # at least 10% of block tokens masked
    sampling_eps_max=0.9           # at most 90% of block tokens masked
)
# xt has the same shape as x0; masked positions contain model.mask_index
num_masked = (xt == model.mask_index).float().mean()
print(f"Fraction masked: {num_masked:.3f}")   # ~0.3
```
```

--------------------------------

### `BD3LM.forward` — HuggingFace Model Forward Pass

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

The `BD3LM` model, extending `transformers.PreTrainedModel`, is compatible with `AutoModelForMaskedLM`. Its `forward` method takes noisy token IDs and optional diffusion timesteps (sigma) to produce per-token logits over the vocabulary.

```APIDOC
## `BD3LM.forward` — HuggingFace Model Forward Pass

`BD3LM` extends `transformers.PreTrainedModel` and is compatible with `AutoModelForMaskedLM`. Its `forward` method accepts noisy token IDs and optional diffusion timestep (sigma) values, returning per-token logits over the vocabulary.

```python
import torch
import transformers
from models.hf import BD3LM, BD3LMConfig

BLOCK_SIZE = 4
tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')

# Load a pre-trained model from HuggingFace
model = transformers.AutoModelForMaskedLM.from_pretrained(
    f'kuleshov-group/bd3lm-owt-block_size{BLOCK_SIZE}',
    trust_remote_code=True
).eval().to('cuda')

# Prepare a batch of token IDs (e.g., noisy tokens at some diffusion timestep)
tokens = tokenizer("The quick brown fox", return_tensors='pt')['input_ids'].to('cuda')
timesteps = torch.zeros(tokens.shape[0], device='cuda')  # sigma=0 → clean

with torch.no_grad():
    # Returns logits of shape (batch, seq_len, vocab_size)
    logits = model(input_ids=tokens, timesteps=timesteps)
    predicted_tokens = logits.argmax(-1)
    print(tokenizer.decode(predicted_tokens[0]))
```
```

--------------------------------

### Noise Schedule Classes

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Factory function and classes for controlling the noise schedule during the diffusion process, including masking probability and loss-weighting.

```APIDOC
## `get_noise` / Noise Schedule Classes — Forward Process Schedules

Factory function and schedule classes controlling the masking probability and loss-weighting at each diffusion timestep. The default `LogLinearNoise` schedule is strongly recommended; alternatives include `CosineNoise`, `ExpNoise`, and `LogarithmicNoise`.

```python
import torch
from noise_schedule import get_noise, LogLinearNoise
```
```

--------------------------------

### Diffusion.forward — Log-Score Computation

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Computes the normalized log-probabilities over the vocabulary for each token position by running the backbone network on noisy token IDs.

```APIDOC
## `Diffusion.forward` — Log-Score Computation

Runs the backbone network on noisy token IDs `x` with noise level `sigma` and applies the appropriate parameterization (`subs`, `sedd`, or `ar`) to return normalized log-probabilities over the vocabulary for each token position.

```python
import torch
import transformers
import omegaconf
import diffusion as diff_module
import dataloader

config = omegaconf.OmegaConf.load('configs/config.yaml')
tokenizer = dataloader.get_tokenizer(config)
model = diff_module.Diffusion(config, tokenizer=tokenizer).to('cuda')
model.eval()

batch_size, seq_len = 2, 1024
# Noisy token sequence (some tokens replaced with mask_index)
x = torch.randint(0, tokenizer.vocab_size, (batch_size, seq_len), device='cuda')
x[:, 100:110] = model.mask_index   # simulate masked tokens

# Sigma: noise level per token (shape: batch x seq_len)
sigma = torch.full((batch_size, seq_len), 0.5, device='cuda')

with torch.no_grad():
    log_probs = model.forward(x, sigma)
# log_probs shape: (batch_size, seq_len, vocab_size)
print(log_probs.shape)       # torch.Size([2, 1024, 50258])
print(log_probs.exp().sum(-1).mean().item())  # ~1.0 (probability simplex)
```
```

--------------------------------

### ExponentialMovingAverage — EMA Parameter Tracking

Source: https://context7.com/kuleshov-group/bd3lms/llms.txt

Maintains exponential moving averages of model parameters for improved evaluation stability. Supports storing and restoring original weights so that EMA parameters are used only during validation/sampling without polluting the optimization state.

```APIDOC
## ExponentialMovingAverage

### Description
Maintains exponential moving averages of model parameters for improved evaluation stability. Supports storing and restoring original weights so that EMA parameters are used only during validation/sampling without polluting the optimization state.

### Usage
```python
import torch
import torch.nn as nn
from models.ema import ExponentialMovingAverage

model = nn.Linear(128, 64)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

# Initialize EMA with decay=0.9999	ema = ExponentialMovingAverage(model.parameters(), decay=0.9999)

# Training loop
for step in range(100):
    loss = model(torch.randn(8, 128)).sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    ema.update(model.parameters())  # update shadow params after each step

# Evaluation: swap in EMA weights, validate, then restore	ema.store(model.parameters())       # save current weights	ema.copy_to(model.parameters())     # overwrite with EMA weights

with torch.no_grad():
    val_out = model(torch.randn(4, 128))
    print("Val output shape:", val_out.shape)
	ema.restore(model.parameters())     # restore original weights for continued training
print("Restored original weights.")
```
```

--------------------------------

### Score Controlled Generations (Sentiment Analysis)

Source: https://github.com/kuleshov-group/bd3lms/blob/main/ssd-lm/README.md

Scores the controlled generations produced by the 'ctrsa' experiment. Results are saved to files ending with '_ssd_ctrsa_eval.txt'.

```bash
source loop_scoring_eval_ctrsa.sh
```

--------------------------------

### BD3-LMs Citation

Source: https://github.com/kuleshov-group/bd3lms/blob/main/README.md

Citation details for the Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models paper.

```bibtex
@inproceedings{
  arriola2025block,
  title={Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models},
  author={Marianne Arriola and Aaron Gokaslan and Justin T Chiu and Zhihan Yang and Zhixuan Qi and Jiaqi Han and Subham Sekhar Sahoo and Volodymyr Kuleshov},
  booktitle={The Thirteenth International Conference on Learning Representations},
  year={2025},
  url={https://arxiv.org/abs/2503.09573}
}

```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.