Generate text via iterative unmasking. Args: model: DiffusionTransformer prompt: tensor of shape (B, L) — the prompt token i

Tokenize full_text, find the split point using prompt_text, handle tokenizer boundary issues and context window truncation. Returns (toke

Dataloader for masked diffusion pretraining. Unlike autoregressive training, we don't need (x, y) pairs with offset. We just need packed

Randomly mask tokens in a batch. Each sequence gets a random mask ratio. Position 0 (BOS) is never masked. Returns: (noisy_batch, masked

Precompute how many tokens to unmask at each step (uniform schedule).

Gumbel-max sampling for categorical distributions.

Compute SFT diffusion loss. Same as pretraining loss but: 1. Only actual response tokens are masked (prompt and padding stay visible) 2.

Estimate negative log-likelihood of the continuation (tokens after prompt_len) via Monte Carlo sampling of the ELBO, following LLaDA eval_lla

Yields (input_ids [B, T], prompt_lengths [B]) batches. Pads shorter sequences with BOS tokens (which are in the prompt region, so never maske

Build a few-shot prefix string from examples, excluding the current item.

Delete all but the most recent `keep_last` checkpoints.

Language modeling: check if greedy unmasking produces the exact continuation.

Multiple choice: vary the answer, pick lowest ELBO.

Schema: vary the context, same continuation, pick lowest ELBO.

Compute the masked diffusion training loss (ELBO). 1% of the time, randomly truncates the sequence (LLaDA paper guideline).

Download and extract the CORE eval bundle.

Evaluate on the CORE benchmark using ELBO-based multiple-choice scoring. Same task data as nanochat, but scored via diffusion likelihood inst

Evaluate a single example. Dispatches to the right logic based on task_type. Returns True if correct.

Quick SFT loss on a subset of data.

Evaluate a task across all examples, distributed across ranks. Returns mean accuracy.

Compute average diffusion loss on validation data.

Auto-detect the first checkpoint directory.

Check if greedy one-step unmasking of the continuation matches the target. Masks all continuation tokens, runs one forward pass, checks if ar

Load model + config from a checkpoint directory.

Load SFT data mix following nanochat's recipe: - SmolTalk (general conversations) - MMLU x3 (teaches multiple choice answering) - GSM

Render a conversation to token ids and compute prompt length. Format: <|bos|> User: {msg}\nAssistant: {response}<|eos|> Only the last ass

Forward pass. Returns logits of shape (B, T, vocab_size). idx: (B, T) token ids (may contain mask tokens)

Same as generate(), but yields (step, total_steps, x) at each unmasking step for visualization. Use with torch.no_grad() externally. Yie