nemo_automodel.recipes.dllm.train_ft

Assemble the [B, 1, R, enc_len + R] block-causal mask from per-example builds.

Each example satisfies build_block_diffusion_training_mask’s prefix + response_length == enc_len contract exactly, so the builder is called once per example (batch is small) and the result is placed into a padded additive mask. Pad query rows (j >= response_length) are kept well-defined by leaving their canvas self-diagonal unmasked, so the shared transformer never sees an all--inf softmax row; those rows are unsupervised and discarded from the loss.

Slice the batch to the response window and build the block-causal masks.

Returns a dict with the canvas-region tensors (all [B, R]) and the forward inputs (canvas ids, masks, position ids). R is the longest response in the batch; shorter responses are right-padded and the pad positions are dropped from the (sliced) loss_mask / noise_mask so they never contribute to the loss.

Count the GLOBAL denominators from the FINAL response-window masks.

The raw pre-mask counts over-count: _build_response_window restricts the diffusion loss to ONE step-selected canvas block (one-canvas scheme) and drops padding, while the AR loss is scored on the padding-masked, next-token-shifted supervised mask. Dividing by num_noise_tokens / num_supervised_tokens would include tokens that cannot contribute, under-scaling diffusion (~by the response-block count) and AR (by padding + the shift). We build each microbatch’s window once here (caching it on the batch so the forward reuses the same step-seeded block selection) and sum the actually-scored tokens.

Per-EXAMPLE two-pass self-conditioning coins -> [B] bool tensor.

Google draws the coin PER EXAMPLE (uniform(B) < p) and mixes conditioned / zero-conditioned examples within a batch. Returns a [B] mask so the model gates the self-cond branch per example, while pass-1 (the no_grad self-cond signal) ALWAYS runs in the forward. Always running pass-1 makes the FSDP collectives identical every step regardless of the coins (no rank desync) and keeps it correct for local_batch_size > 1 — a scalar per-microbatch coin only matched Google’s per-example mix when B == 1. Seeded by (step, microbatch_idx) for reproducibility; rank-correlation of the coins is harmless now that pass-1 no longer branches on them.

Response-window forward/backward: canvas-only logits + canvas-sliced loss.

Inject a dummy mask_token_id (unused by block_diffusion) then build.

Apply token corruption via the configured strategy.

Parameters:

Returns:

Tuple of (noisy_input_ids, noise_mask, p_mask).

Add any model-specific forward inputs derived from the batch.

Default is a no-op; block_diffusion overrides this to attach the block-causal attention masks and canvas position ids.

Return (num_diffusion_tokens, num_ar_tokens) — the GLOBAL token denominators for this step’s diffusion + AR losses.

Base: the pre-mask counts (diffusion per the strategy’s normalization_mode, AR = supervised). Subclasses whose final loss masks differ from these raw counts (e.g. block-diffusion, which restricts the diffusion loss to one selected canvas block and drops padding) override this to count the final masks — otherwise the losses are divided by tokens that cannot contribute and are silently under-scaled.

Override: apply dLLM corruption and compute dLLM loss.

Execute a single training step with dLLM loss.

Follows the parent pattern but uses loss_mask from the collate wrapper instead of labels != -100 for token counting.

Run one validation pass with dLLM corruption and loss.

Computes per-batch loss with proper denominators, then accumulates weighted by noise token count to produce a per-noise-token average across the val set.

Replace dataloader collate functions with the dLLM single-pass collater.

Uses DLLMCollator which goes directly from variable-length sample lists to block-aligned tensors in one pass.

Requires datasets to produce unshifted format (input_ids + loss_mask, via _package_tokenized_example(unshifted=True)).

Log dLLM-specific training metrics.

Build all training components, then apply dLLM-specific overrides.