bridge.models.diffusion_gemma.diffusion_gemma_step

Megatron forward/loss step for DiffusionGemma uniform-state diffusion SFT.

Batch contract (prompt left-padded, canvas right-padded)::

input_ids          [B, P]  encoder prompt, including image placeholder tokens
attention_mask     [B, P]  1 for real prompt tokens, 0 for padding
canvas_ids         [B, L]  clean target tokens x0, EOS-terminated
canvas_mask        [B, L]  1 for scored target tokens, 0 for padding
position_ids       [B, P]  optional
pixel_values, image_position_ids, mm_token_type_ids  optional image inputs

Optional deterministic overrides, mainly for parity tests: diffusion_t [B], noisy_canvas_ids [B, L], self_conditioning_mask [B].

Module Contents

Functions

corrupt_canvas	Sample one time per example and replace valid tokens uniformly with probability t.
diffusion_loss_function	Return Megatron’s (sum_loss, denominator, metrics) contract.
get_batch_from_iterator	Load a DiffusionGemma batch and move tensors to the current CUDA device.
next_microbatch_index	Return this phase’s microbatch index, restarting whenever the step changes.
noise_seed	Derive a deterministic 63-bit RNG seed for one microbatch on one DP rank.
_noise_generator	Return a CUDA RNG keyed by (seed, phase, step, microbatch, DP rank).
_diffusion_setting
forward_step	Corrupt the target canvas, run one DiffusionGemma forward, and return its loss.

Data

_BATCH_KEYS
_MICROBATCH_COUNTERS

API

bridge.models.diffusion_gemma.diffusion_gemma_step._BATCH_KEYS

(‘input_ids’, ‘attention_mask’, ‘position_ids’, ‘pixel_values’, ‘image_position_ids’, ‘mm_token_type…

bridge.models.diffusion_gemma.diffusion_gemma_step.corrupt_canvas(

x0: torch.Tensor,

canvas_mask: torch.Tensor,

*,

vocab_size: int,

noise_min: float = 0.001,

noise_max: float = 1.0 - 0.001,

generator: torch.Generator | None = None,

) → tuple[torch.Tensor, torch.Tensor]

Sample one time per example and replace valid tokens uniformly with probability t.

bridge.models.diffusion_gemma.diffusion_gemma_step.diffusion_loss_function(

output_tensor: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],

*,

variant: str,

vocab_size: int,

diffusion_weight_clip: float | None = None,

check_for_nan_in_loss: bool = True,

check_for_spiky_loss: bool = False,

decoder_loss_weight: float = 1.0,

encoder_loss_weight: float = 1.0,

) → tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]

Return Megatron’s (sum_loss, denominator, metrics) contract.

DiffGemma averages loss per example, so the denominator is the number of examples. Megatron divides by it (or reduces it globally when calculate_per_token_loss is enabled), preserving the reference mean.

bridge.models.diffusion_gemma.diffusion_gemma_step.get_batch_from_iterator(

data_iterator: Iterable,

*,

is_first_pp_stage: bool = True,

is_last_pp_stage: bool = True,

) → dict[str, Any]

Load a DiffusionGemma batch and move tensors to the current CUDA device.

bridge.models.diffusion_gemma.diffusion_gemma_step._MICROBATCH_COUNTERS: dict[str, dict[str, int | None]]

None

bridge.models.diffusion_gemma.diffusion_gemma_step.next_microbatch_index(mode: str, step: int) → int

Return this phase’s microbatch index, restarting whenever the step changes.

Train and eval keep separate counters. A validation pass inside step N must not shift the noise that training step N later draws, otherwise a run resumed from a checkpoint (fresh counters) would diverge from the uninterrupted run.

bridge.models.diffusion_gemma.diffusion_gemma_step.noise_seed(

*,

seed: int,

mode: str,

step: int,

microbatch: int,

dp_rank: int,

) → int

Derive a deterministic 63-bit RNG seed for one microbatch on one DP rank.

bridge.models.diffusion_gemma.diffusion_gemma_step._noise_generator(

state: megatron.bridge.training.state.GlobalState,

model: Any,

) → torch.Generator

Return a CUDA RNG keyed by (seed, phase, step, microbatch, DP rank).

TP/EP peers in one data-parallel replica draw identical noise; DP replicas draw distinct noise. Keying by training step makes resumed runs reproduce the original corruption without checkpointing generator state.

bridge.models.diffusion_gemma.diffusion_gemma_step._diffusion_setting(

config: Any,

name: str,

default: Any,

) → Any

bridge.models.diffusion_gemma.diffusion_gemma_step.forward_step(

state: megatron.bridge.training.state.GlobalState,

data_iterator: Iterable,

model: Any,

) → tuple[tuple[torch.Tensor, ...], functools.partial]

Corrupt the target canvas, run one DiffusionGemma forward, and return its loss.