Native diffusion language-model training

Train supported diffusion checkpoints with diffusion_lm: and the usual Axolotl chat datasets. No plugin is required. The model profile selects its noise process, attention layout, logit alignment and reference loss.

Model	Noise and layout	Default native objective
Nemotron-Labs-Diffusion (3B, 8B)	Absorbing mask, full sequence, aligned logits	Inverse-time weighted CE at masked positions
Dream v0	Absorbing mask, full sequence, shifted logits	Source-compatible masked CE with inverse-time weighting
DiffusionGemma	Uniform token noise, causal encoder and bidirectional canvas	Canvas CE, encoder AR CE and two-pass self-conditioning

Other models can still train with the legacy masked-diffusion objective by setting diffusion_lm.from_causal_lm: true with the legacy diffusion plugin; any other diffusion_lm: config on a model without a native profile is rejected rather than trained as causal SFT. See Converting a causal LM.

Installation

Native diffusion training runs on Axolotl’s supported Torch range, using the dependency versions pinned in pyproject.toml. attn_implementation: varlen additionally requires Torch 2.14, so install the torch==2.14.0 wheel for your accelerator from PyTorch’s installation guide before installing Axolotl:

python -m venv .venv-diffusion
source .venv-diffusion/bin/activate
python -m pip install --upgrade pip
python -m pip install torch==2.14.0  # or the accelerator-specific wheel
python -m pip install -e .

Put this environment’s bin directory first on PATH so axolotl train launches the matching Accelerate installation.

Start with Nemotron 3B

The small iteration target is nvidia/Nemotron-Labs-Diffusion-3B. From the repository root:

axolotl preprocess examples/nemotron-diffusion/lora-smoke.yaml
axolotl train examples/nemotron-diffusion/lora-smoke.yaml

This recipe trains two optimizer steps on 32 synthetic chat records. It checks preprocessing, packed training, gradient accumulation and adapter saving; it is not a quality benchmark. It pins the model and remote-code revision with revision_of_model and uses the published bidirectional diffusion mode, aligned logits and the existing mask token 100. Padding reuses the tokenizer EOS token without resizing the vocabulary. When you change models, remove or replace revision_of_model.

Attention backends

Native diffusion defaults to flex_attention when attn_implementation is omitted. eager and sdpa remain explicit fallbacks; the examples use sdpa as a dense reference. FlexAttention requires a compatible GPU, and its first execution includes kernel compilation.

Nemotron also supports PyTorch’s public variable-length attention:

attn_implementation: varlen
sample_packing: true

This path gathers valid Q/K/V tokens into THD format and passes cumulative document offsets to torch.nn.attention.varlen.varlen_attn, without building a dense document mask. PyTorch selects the underlying kernel. varlen requires Torch 2.14 or newer, a native full-sequence diffusion model (Nemotron), uncached bidirectional attention, and zero attention dropout. Dream and DiffusionGemma keep their existing backends: Gemma’s per-layer encoder-prefix/canvas visibility cannot be expressed as one full-sequence varlen call.

DiffusionGemma’s large attention heads can exceed a GPU’s per-kernel shared memory limit with the default Flex tile sizes. Smaller tiles can be configured without changing attention visibility:

attn_implementation: flex_attention
flex_attn_compile_kwargs:
  fwd_BLOCK_M: 16
  fwd_BLOCK_N: 16
  fwd_num_stages: 1
  bwd_BLOCK_M1: 16
  bwd_BLOCK_N1: 16
  bwd_BLOCK_M2: 16
  bwd_BLOCK_N2: 16
  bwd_num_stages: 1

Packed training

Enable sample_packing: true. Logical documents remain isolated even when they share a physical row, and position IDs restart at document boundaries. DiffusionGemma also keeps each canvas associated with its own encoder history.

sample_packing: true
sequence_len: 256
micro_batch_size: 2
diffusion_lm:
  overflow_policy: error
  generate_samples: false

sequence_len limits each example’s logical layout. For packed and batch_flattening rows, Axolotl derives the total allocation as sequence_len * micro_batch_size. Ordinary padded batches have no aggregate row budget. Evaluation uses eval_batch_size when it is set.

FlexAttention rounds capacity down to 128-token buckets. Encoder/canvas packing reserves another 128 tokens because the two streams pad independently. Use a derived allocation of at least 128 for full-sequence models and 256 for DiffusionGemma; increase sequence_len if individual examples do not fit. SDPA and eager use the full derived capacity.

The default overflow_policy: error raises on an oversized example. overflow_policy: drop filters and counts oversized examples during preprocessing; it never truncates them.

Dream can use eos_tail: visible_supervised to extend each logical example to sequence_len with supervised EOS tokens. These are semantic tokens and consume packing capacity. Physical padding is excluded from attention and loss.

Packing mainly saves memory; throughput depends on your lengths and attention backend, so benchmark both settings for your workload.

DiffusionGemma attention boundaries

Prompt prefill is causal. During denoising, each token in the current output canvas attends to the clean prefix and every valid token in that same noisy canvas. Bidirectional visibility applies to the output canvas, not the combined prompt and output sequence.

Training uses the same boundary. The encoder can process the complete clean sequence for its auxiliary autoregressive loss, but the decoder mask exposes only keys before the selected output block. Clean copies of that block and later output tokens stay hidden, and causal encoder attention prevents those future tokens from entering the visible prefix states.

Full-attention decoder layers see the entire eligible prefix. Sliding-attention layers see the retained prefix window and the entire current canvas; they do not apply a query-relative sliding window inside the canvas. Both layer types exclude other packed examples. Dream and Nemotron use their full-sequence layouts.

Objective and recurrence options

Unset options inherit the model’s reference behavior. Dream’s CART variant uses time_weighting: cart and cart_p; its optional focal reweighting uses token_reweighting: true, alpha and gamma.

DiffusionGemma’s default self-conditioning probability (self_conditioning.p) is 0.5. Its optional Rhine objectives use time_weighting: loo or inv_one_minus_t with objective_reduction: example_mean; rhine_weight_clip caps the 1 / (1 - t) weights. encoder_ar_weight scales the separate encoder loss. self_conditioning.train_module: true includes the self-conditioning module in the saved adapter.

diffusion_lm:
  unroll:
    k_max: 2
    grad_through_steps: false
    eval_steps: 1

With unroll.k_max above 1, each step draws a read count uniformly from 1 to k_max. Earlier reads run without retained gradients and the final read supplies the loss. grad_through_steps: true is available only for models with self-conditioning. eval_steps (1 to k_max) fixes the read count for evaluation and prediction instead of drawing it; unset, they draw like training.

Adapters, merging and inference

Native training requires LoRA with explicit text attention or dense-MLP projection targets. DoRA, layer replication, lora_target_linear and lora_target_parameters are rejected, and native vocabulary resizing is disabled. Nemotron additionally supports adapter: qlora (4-bit; 8-bit is rejected), FSDP2 via fsdp_config, and the fused LoRA kernels (lora_qkv_kernel, lora_o_kernel, lora_mlp_kernel) together with lora_fp32_gradients; Dream and DiffusionGemma remain unquantized LoRA without FSDP.

DiffusionGemma’s encoder and decoder text weights are tied, so their LoRA factors are shared. Its targets must be a single expression naming both the encoder language model and the decoder, for example:

lora_target_modules: '^model\.(encoder\.language_model|decoder)\.layers\.\d+\.self_attn\.(q|k|v|o)_proj$'

Vision, router and expert parameters stay frozen.

Nemotron supports opt-in Cut Cross Entropy for its aligned diffusion objective and for typed-decision training. With Axolotl’s CCE fork installed, add:

plugins:
  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
cut_cross_entropy: true

Keep any other plugins the recipe requires in that list. The final training read computes full-vocabulary CE from hidden states without materializing vocabulary logits; decision training still computes its restricted-label and Brier terms over the allowed answers. Targets stay aligned, with no causal shift. This path requires a plain linear output head and does not support a LoRA output head or the loo objective.

axolotl merge-lora config.yaml exports a merged checkpoint. When merge_method is omitted, DiffusionGemma selects legacy so its tied factors merge once; setting merge_method: memory_efficient for DiffusionGemma raises.

For interactive generation, use axolotl inference config.yaml --chat. Each reply is a completion block appended to the conversation and denoised in one piece; see Inference. Native models generate completions only.

Nemotron Diffusion VLM

nvidia/Nemotron-Labs-Diffusion-VLM-8B (model_type: nemotron_labs_diffusion_vlm) pairs the Nemotron diffusion backbone with a Pixtral vision tower. Axolotl resolves it to its own profile: the vision tower and projector are frozen on load, and 4-bit loading and the QLoRA merge leave the vision stack and diffusion_head unquantized. In a lora_target_modules list, the names the Pixtral tower shares with the decoder (q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj) are anchored to encoder.layers.N.self_attn/mlp, so they cannot attach to the tower’s projections of the same name. Other entries (for example diffusion_head) keep PEFT’s suffix matching, peft_layers_to_transform is folded into the anchored pattern, and projector-only names (linear_1, linear_2, merging_layer, patch_conv) are rejected. A string (regex) target is used as written; LoRA that lands on a vision module still fails after the adapter loads. The profile also patches two transformers-version mismatches in the published remote code. FSDP2, QLoRA, varlen attention and the fused LoRA kernels (lora_qkv_kernel, lora_o_kernel, lora_mlp_kernel) are supported and were measured on text-only batches; image batches are unmeasured. Under FSDP2 with auto_wrap_policy: TRANSFORMER_BASED_WRAP, fsdp_config.transformer_layer_cls_to_wrap defaults to Ministral3DecoderLayer and may name only that class: a microbatch without images skips the frozen tower, so a vision, projector or Ministral3RMSNorm unit would all-gather on some ranks only and hang. Text-only training on this model matches the remote forward exactly in fp32. The profile sets complementary_mask: false on the loaded config so the remote forward does not duplicate the batch, and saved or merged config.json files keep that value; restore the published true before reusing the remote training code.

Its decoder weights are not the LM’s: the VLM checkpoint was trained further, and its tokenizer carries one extra added token, so step losses differ from Nemotron-Labs-Diffusion-8B on the same data.

Image inputs are supported only through the typed-decision plugin (diffusion_decision); plain diffusion_lm SFT on this model is text-only and rejects multimodal datasets. The model ships no HF processor; the plugin uses NemotronDiffusionVLMProcessingStrategy from axolotl.processing_strategies, which ports the remote image math (PNG, JPEG, WEBP, GIF and BMP only; URL fetches refuse redirects and non-image/* responses and stop at 32 MB or 60 s). An image is resized so its longest edge is at most max_image_size and each 28x28 block (patch 14, spatial merge 2) becomes one <|image_pad|> token: a W x H grid costs W * H pad tokens plus H - 1 row breaks and the start/end markers (1400 -> 50x50 = 2551 prompt tokens; 784 -> 28x28 = 813). Pixel inputs reach the forward as pixel_values (N, 3, H, W) padded to the largest image in the batch and image_sizes (N, 2), in prompt order; the model scatters them into the pad tokens in sequence order. Image inputs through the typed-decision plugin are described under Typed-decision training; examples/nemotron-diffusion/decision-lora-vlm-8b.yaml is the recipe.

Typed-decision training

The diffusion_decision plugin trains a native diffusion model to answer typed multiple-choice, score and yes/no questions. Each record renders a fixed decision system prompt from its question schema, sends its state as the user message, and places one answer label per question in the output canvas using the canvas template of the public mmastrac/djev project.

examples/nemotron-diffusion/decision-lora-8b.yaml is the starting recipe for nvidia/Nemotron-Labs-Diffusion-8B. It uses BF16, attn_implementation: varlen, and rank-64/alpha-128 LoRA over all seven attention and MLP projections with loraplus_lr_ratio: 8. It trains one epoch with AdamW at 6.5e-6, microbatch 16 and eight accumulation steps (effective batch 128). The decision loss combines candidate-restricted CE, full-vocabulary CE and a 0.1 Brier term, with unroll.k_max: 2, no latent slots and non-reentrant gradient checkpointing. The YAML leaves the model revision unpinned; add revision_of_model to a local copy for an exact reproduction.

axolotl preprocess examples/nemotron-diffusion/decision-lora-8b.yaml
axolotl train examples/nemotron-diffusion/decision-lora-8b.yaml

Set the data_files of the datasets and test_datasets entries to separate normalized train and development JSONL files. The recipe evaluates every 50 steps and keeps the adapter with the best eval_loss. Keep held-out test data out of both entries.

Dataset format

type: diffusion_decision.jsonl reads already-normalized JSONL, one record per line. diffusion_decision.typed_decisions and diffusion_decision.procedural normalize rows from the typed_decisions and procedural typed-decision source schemas into the same form.

{"source":"support_demo","family":"returns","id":"case-0042","group":"case-0042","instructions":"Decide from the supplied case only.","state":{"customer":"A customer says their unopened order arrived yesterday and asks for a refund.","policy":"Unopened items delivered within 30 days may be refunded."},"questions":{"refund":{"type":"choice","instructions":"Choose the next action.","options":[{"name":"approve","description":"Approve a refund under the stated policy."},{"name":"deny","description":"Deny the request because the policy does not apply."},{"name":"escalate","description":"Send the case to a human reviewer."}]},"urgency":{"type":"score","instructions":"Rate how urgently this needs review.","levels":["routine","soon","urgent"]},"eligible":{"type":"noul","instructions":"Is the customer eligible for a refund?","criteria":{"true":"The stated policy permits a refund.","false":"The stated policy does not permit a refund."}}},"labels":{"refund":{"kind":"dist","probs":[0.78,0.02,0.20]},"urgency":{"kind":"hard","gold_idx":0},"eligible":{"kind":"set","allowed_set":[0]}},"source_metadata":{"provenance":{"dataset":"my-org/support-decisions","revision":"immutable-revision","split":"train","source_id":"case-0042"}}}

source, id and group identify the row; family is optional and is used for split isolation. Records with the same source, family, state and top-level instructions that share a group are rendered on one canvas, up to max_questions_per_canvas questions. There is no per-row system field: put shared context in state, global instructions in instructions, and per-question instructions on each question. Keep question IDs stable and nonempty.

Type	Required schema field	Canonical alternative order
`choice`	`options`	List order. An option is a string or an object with `name` and `description`; descriptions are rendered.
`score`	`levels`	List order, lowest to highest as defined by the source. Do not sort levels lexically.
`noul`	none	`yes`, then `no`. Optional `criteria` documents the true/false conditions and is rendered.

Label indexes always refer to this canonical order, never to option names or token IDs. Answer tokens come from labels.codebook: vendored26 uses the djev letters A-Z (up to 26 alternatives); expanded52 allows up to 52 (A-Z, a-z); expanded128 adds selected Greek and Cyrillic letters for up to 128. Every label must be one token in the model’s tokenizer, which the canvas template check enforces at preprocessing.

On nvidia/Nemotron-Labs-Diffusion-VLM-8B, a record may carry a top-level images list of http(s) URLs, absolute paths, or paths relative to the directory of the local data file (relative paths are rejected for Hub sources). Images are opened as stored, without EXIF rotation. The images are placed before state in the user turn and expanded to their token grids in the prompt, so they count against sequence_len (diffusion_decision.max_image_size, default 1400, caps the longest edge; see Nemotron Diffusion VLM for the token accounting). Records with different images are never grouped onto one canvas, and decontamination fingerprints an image record by its state plus the sha256 of each image file. The prepared-row cache stores resolved references and sizes only and hashes each image’s bytes, so an edited image or a changed max_image_size invalidates it. Only the jsonl adapter reads images; URL-referenced images bypass the cache. Each collator, and each dataloader worker’s copy of it, keeps up to diffusion_decision.image_cache_size (default 64; 0 disables) normalized float32 image tensors, up to about 23.5 MB each at 1400 px.

Labels

Every question has exactly one label:

{"kind":"hard","gold_idx":1}
{"kind":"dist","probs":[0.10,0.70,0.20]}
{"kind":"set","allowed_set":[0,2]}

hard names one zero-based alternative. set is a nonempty set of unique valid indexes. dist has one finite, nonnegative probability per alternative that sums to one within 1e-4; zero entries are valid and must not be omitted. Partial distributions, ranking targets, per-row candidate token IDs and hard/soft blending are not accepted.

labels.label_softmax selects candidate-restricted CE (restricted), full-vocabulary CE over the allowed label tokens (full), or both (the default). The full-vocabulary term still penalizes probability outside the valid answers. full_ce_weighting: dft applies DFT weighting to hard-label full CE. hard_label_smoothing spreads that fraction of a one-hot CE target uniformly across the question’s valid answer tokens; Brier targets are unchanged.

Splits and provenance

Keep train, development and test data in separate files or dataset entries, and use split: train only for training entries. Preserve source_metadata.provenance with an immutable source location and revision, the split, and the source row ID. typed_decisions rows train only from a declared train split and reject a row split that disagrees with it. After normalization, the loader rejects states or grouped contexts shared between splits.

Mixture and weighting

mixture.weights sets per-source sampling weights by source. Without weights, each source is sampled in proportion to its row count raised to mixture.temperature (default 0.5). max_examples_per_source caps rows per source. loss_weight multiplies each source’s loss. Loss is averaged over the questions on a canvas and then over canvases, so a canvas with many questions does not outweigh a single-question canvas.

Every epoch is drawn afresh with the same number of microbatches as the first: each capped source contributes its next max_examples_per_source window of one seeded shuffle (so later epochs see records the first did not, before any repeat), the mixture draw and the stratified microbatches are redrawn, and question and option order is seeded by the record, the epoch and how many times the epoch has already drawn that record. A record’s first draw in epoch 0 uses the order prepared at startup; any other draw builds its canvas when the batch is loaded by concatenating the record’s token spans pre-tokenized at startup (prompt preamble, each question’s header and option lines, the format instruction, the chat turn around them, and each question’s answer piece), so re-ordering questions and options or reassigning labels needs no tokenizer (a few milliseconds per canvas, most of it seeding the permutation and validating the canvas, against about 25 ms to re-tokenize one on a 16k-token context). This is exact only where no pre-token crosses a span boundary; the plugin proves that for the Nemotron-Labs-Diffusion tokenizers, checks each record’s first draw against full tokenization when the spans are built, and re-tokenizes records or tokenizers it cannot prove. A drawn record that was not prepared at startup and does not fit the token budget is replaced by a prepared record of the same source. Epochs are a function of the seed, the batch settings and the epoch number, so ranks agree and a resumed run replays the same batches. With training sample_packing (core multipack) or premixed rows, every epoch reuses the first epoch’s draws.

To train on a prebuilt draw sequence, set mixture.premixed: true and per_batch_stratified: false. Each training row then needs a unique draw ID and an origin matching its own identity; copies of one origin must be identical apart from this metadata:

"source_metadata": {
  "premix": {
    "draw_id": "run-000001",
    "origin": {"source": "support_demo", "id": "case-0042", "group": "case-0042"}
  }
}

Batching

batch_flattening: true concatenates the logical microbatch and keeps decisions isolated through the varlen attention path. For variable-length decisions, set batch_flattening: false, sample_packing: true and eval_sample_packing: false instead. Packing uses sequence_len as the per-decision limit and sequence_len * micro_batch_size as the physical token budget. It requires mixture.per_batch_stratified: false and does not support sampled latent-slot counts. Reduce micro_batch_size and raise gradient_accumulation_steps if the effective batch does not fit.

Evaluation batches hold a fixed eval_batch_size rows by default. With sample_packing: false, set eval_sample_packing: true to pack evaluation instead: consecutive rows are concatenated, in dataset order, into one varlen row of at most sequence_len * eval_batch_size tokens, so short rows share a forward and a micro-batch never exceeds the fixed-batch worst case. This works with or without batch_flattening for training, applies to full-sequence layouts (encoder/canvas layouts keep fixed batches), and repeats no row across ranks. eval_loss and the eval_decision/* metrics remain per-example means, equal to the fixed-batch values. With sample_packing: true, evaluation is packed by the same multipack sampler as training unless eval_sample_packing: false; only one packing mechanism applies.

On a full-sequence layout without free latent slots, training and evaluation run only the final pass of a k_max unroll: label slots keep their noisy tokens between passes, nothing self-conditions and no gradient flows through earlier passes, so those passes would re-read the same canvas and their logits would be discarded. Loss and gradients are identical to running them, and the read count is still drawn so the random stream is unchanged. The trainer logs this once. Earlier passes only feed forward through latent.mode: free slots, which commit their argmax predictions between passes. The final pass computes logits only at label slots, so it never materializes full-vocabulary logits over the sequence.

Prepared-row cache

With dataset_prepared_path set and local JSONL sources, prepared canvases are cached under <dataset_prepared_path>/diffusion_decision_cache/ and reused when the cache identity matches: the decision config and the other preparation fields (seed, sequence_len, model config), the tokenizer files and vocabulary, the plugin’s code, and the source files. A source file is matched by its resolved path, size and modification time, so touching it without changing its content still re-prepares once.

The cache holds the canvas layer only: every decontaminated, grouped training record per source, with canvases built for the first epoch’s capped subset (the rest of a capped source is kept compressed for later epochs), the pre-tokenized canvas spans of every training record, and the evaluation canvases. Batch settings are not part of its identity: micro_batch_size, gradient_accumulation_steps, eval_batch_size, sample_packing, eval_sample_packing, batch_flattening and the mixture’s weights, temperature and per_batch_stratified only shape the batch plan (token-budget filter over the cached lengths, mixture draws and microbatch grouping), which is rebuilt at startup in seconds.

diffusion_decision.prepared_cache_verify: true trades load time for stronger checks: sources are matched by content sha256 instead, and a cache hit hashes the stored rows against the hash recorded when they were written. Without it a damaged or truncated cache file is still rejected when its framing or contents do not decode, and the rows are prepared again. The two modes key sources differently, so switching the option prepares and caches once more.

Outputs

Each saved adapter includes diffusion_decision_manifest.json, which records the base model and revision, diffusion spec, canvas width, label codebook, latent slots, read protocol and tokenizer special IDs that an external inference integration must reproduce. When any training or eval canvas carries images it also records image_inputs: max_image_size and the placement (images_before_state). Serving must resize with the same limit; the remote image_processing.py defaults to DEFAULT_MAX_IMAGE_SIZE = 1400, so set it to the recorded value. Axolotl does not ship a decision-serving endpoint. Use a separate config for a frozen test split, and do not use test data for checkpoint selection.

The manifest also records max_questions_per_canvas, canvas_prefix_ids and, under noise_read_protocol:

read_passes: model forward passes a decision read depends on. Reads mask every label slot in one pass; on a full-sequence layout without free latent slots the later steps of a k_max unroll re-read the same canvas, so this is 1. null means the count is drawn per batch and changes the result.
unroll_k_max and unroll_eval_steps: the configured unroll.
generation_diffusion_steps: diffusion_lm.num_diffusion_steps, used only by free-text generation, never by decision training or reads. Older manifests call this field diffusion_steps; it is read as generation_diffusion_steps.

Loading weights through lora_model_dir, or a merged checkpoint directory (base_model pointing at a local directory with the manifest), refuses a config whose codebook, canvas width, canvas prefix, decision layout, questions per canvas, latent slots, read passes or mask token disagree with the manifest. Fields an older manifest does not record are not compared. To continue training an adapter under a deliberately different canvas or codebook, set diffusion_decision.allow_manifest_mismatch: true; the mismatch is then logged as a warning. Reads should leave it off.

diffusion_decision.read_overflow_canvas_width is a read-time option for full-sequence layouts: a read canvas whose answer template does not fit canvas_width is rebuilt at this wider width instead of being dropped. Canvases that fit are unchanged, and training rows never widen. The wider canvases are outside the training distribution, so validate their accuracy before relying on them.

Public procedural decision mix

scripts/diffusion_lm/build_public_decision_mix.py materializes a reproducible train/dev mix from tasksource/procedural-typed-decisions at revision 916e6cce365a65c37d58db70c9e369795817651e:

python scripts/diffusion_lm/build_public_decision_mix.py ./data/public-procedural
axolotl train examples/nemotron-diffusion/decision-lora-8b-public-procedural.yaml

The output contains only rows from that dataset plus a contract.json recording the pinned revision and licenses. The dataset card declares Apache-2.0; the separate generator repository publishes its code under CC-BY-4.0. Keep contract.json with any materialized output you share. The model is covered by its own model-card license.

--protected-normalized PATH (repeatable) excludes any public row whose canonical state matches a row in your normalized evaluation JSONL; that file is never copied to the output. This is canonical-state decontamination only, not cross-source group or family disjointness. The train and dev outputs are checked against each other for state, family and group overlap.

decision-lora-8b-public-procedural.yaml pins the model revision and sets mixture.temperature: 1.0, which schedules each usable row once. At the pinned revisions, preprocessing keeps 30,080 train rows at the 2,048-token budget, which is 235 effective batches of 128, hence max_steps: 235. This mix is a reproducible starting point for the format, not a quality result.

Public image decision mix

scripts/diffusion_lm/build_public_image_decision_mix.py materializes a reproducible image typed-decision train/dev/test mix for nvidia/Nemotron-Labs-Diffusion-VLM-8B:

python scripts/diffusion_lm/build_public_image_decision_mix.py ./data/public-image-decisions
axolotl train examples/nemotron-diffusion/decision-lora-vlm-8b-public-image.yaml

Train and dev rows come from HuggingFaceM4/A-OKVQA at revision d1b0efa3a436e9101dfbde3752db7607da696c35 (train and validation splits) and the image-bearing rows of derek-thomas/ScienceQA at revision f18b0a70359ebfb41f658fd564208d0355b013f4 (train and validation splits). The test file is nyu-visionx/CV-Bench at revision bc284db50d036958861cb60cdd7b77612052ce0d (default config, test split), a held-out benchmark that is never used for training or checkpoint selection.

Each row is one choice question named answer with the source’s answer options in their original order and a hard label. state.question holds the question text; ScienceQA rows also carry state.context from the source hint. family is the A-OKVQA question_id, the ScienceQA row id (its subject and topic stay in source_metadata.attributes), or the CV-Bench task. Every row has images: ["images/<sha256>.<png|jpg>"], a path relative to the JSONL file. Image bytes are written once per content hash; PNG and JPEG sources keep their original encoding and anything else is re-encoded as PNG. Rows without an image, with duplicate or empty options, or with an answer outside the option list are dropped.

--train-per-source (default 4096), --dev-per-source (default 256) and --test-cap (default 0) set the rows taken from each source. Each nonzero value is an exact per-source count, and the build fails when a source has fewer usable rows; 0 keeps every usable row for all three flags; negative values are rejected. Selection sorts rows by sha256(seed\0source\0split\0id), so the same caps always yield the same rows. CV-Bench states and images are excluded from train and dev, train states and images are excluded from dev (images by sha256 of their bytes; a state is its fingerprint plus its image digests, as in the plugin’s decontamination, so a templated question over a different image is kept), and source/group pairs, families and image digests are disjoint across splits, so the plugin’s family decontamination keeps every train row. contract.json records the pinned revisions, license evidence, per-split counts and drop reasons (no_image and invalid-option drops count the whole split; state_overlap and image_overlap count only the rank-ordered rows read until the count filled, as recorded under scanned), the sha256 of each JSONL file, and the image count and total bytes.

Licenses: CV-Bench declares Apache-2.0, but its images are repurposed from COCO, ADE20K and Omni3D and keep those sources’ terms. The ScienceQA card metadata says cc-by-sa-4.0 while its Licensing Information section links CC BY-NC-SA 4.0; treat the NonCommercial reading as binding until the card is resolved. Its images were collected with the questions from K-12 science curricula, and the card states no image-specific licence. The A-OKVQA card declares no license; the A-OKVQA project repository is Apache-2.0. Its images are COCO 2017 Flickr photos under the Flickr Terms of Use with per-image licences; COCO’s CC BY 4.0 covers the annotations, not the images. These flags are repeated in contract.json["license_flags"]. Keep contract.json and the images/ directory with any materialized output you share.

decision-lora-vlm-8b-public-image.yaml sets diffusion_decision.max_image_size: 784, so each image costs at most 813 prompt tokens (28 x 28 = 784 pads, 27 row breaks, start and end) inside the 2,048-token budget, and mixture.temperature: 1.0, which schedules each usable row once: the 8,192 train rows all survive decontamination, which is 64 effective batches of 128, hence max_steps: 64. Lower max_steps if preprocessing drops rows at the budget. It evaluates and saves every 32 steps; keep eval_steps and save_steps dividing max_steps so load_best_model_at_end also scores the final step. This mix is a reproducible starting point for the image record format, not a quality result.

Converting a causal LM

diffusion_lm.from_causal_lm: true trains an ordinary causal checkpoint with the legacy masked-diffusion objective, controlled by noise_schedule, min_mask_ratio, max_mask_ratio, importance_weighting and the mask-token options. It requires plugins: [axolotl.integrations.diffusion.DiffusionPlugin], which supplies the diffusion trainer; without it the config is rejected. Existing configurations with a deprecated diffusion: block still load through the same plugin, which translates them to this form.