`bridge.models.diffusion_gemma.data`#

Batch collation and synthetic data for DiffusionGemma block-diffusion SFT.

Module Contents#

Classes#

`_MockDiffusionGemmaDataset`
`MockDiffusionGemmaDatasetConfig`	Synthetic, deterministic DiffusionGemma SFT data for smoke tests.
`DiffusionGemmaJSONLDataset`	Load message/image prompts and completions without fitting them all on a GPU.
`DiffusionGemmaDatasetConfig`	Serializable real-data adapter preserving explicit train/validation/test files.

Functions#

`_as_long`
`collate_diffusion_gemma`	Collate prompt/canvas examples into the DiffusionGemma step contract.

API#

bridge.models.diffusion_gemma.data._as_long(value: Any) → torch.Tensor#

bridge.models.diffusion_gemma.data.collate_diffusion_gemma( examples: Sequence[dict[str, Any]], *, pad_token_id: int, canvas_length: int, pad_to_multiple_of: int = 1, ) → dict[str, torch.Tensor]#

Collate prompt/canvas examples into the DiffusionGemma step contract.

Each example provides input_ids (encoder prompt, including image placeholder tokens) and canvas_ids (EOS-terminated target, at most canvas_length tokens). Optional image tensors are stacked per example. Prompts are left-padded to a common length that is a multiple of pad_to_multiple_of (use the TP size with sequence parallelism).

class bridge.models.diffusion_gemma.data._MockDiffusionGemmaDataset( config: bridge.models.diffusion_gemma.data.MockDiffusionGemmaDatasetConfig, size: int, offset: int, )#

Bases: torch.utils.data.Dataset

Initialization

__len__() → int#

__getitem__(index: int) → dict[str, torch.Tensor]#

collate_fn( examples: Sequence[dict[str, Any]], ) → dict[str, torch.Tensor]#

class bridge.models.diffusion_gemma.data.MockDiffusionGemmaDatasetConfig#

Bases: megatron.bridge.data.base.DatasetProvider

Synthetic, deterministic DiffusionGemma SFT data for smoke tests.

canvas_length: int#: 256

prompt_length: int#: 32

completion_length: int#: 64

num_distinct_answers: int#: 4

token_range: tuple[int, int]#: (1000, 20000)

pad_token_id: int#: 0

eos_token_id: int#: 1

bos_token_id: int#: 2

image_token_id: int | None#: None

boi_token_id: int | None#: None

eoi_token_id: int | None#: None

image_patches_per_side: int#: 0

patch_size: int#: 16

pooling_kernel_size: int#: 3

pad_to_multiple_of: int#: 1

seed: int#: 1234

dataloader_type: str | None#: ‘single’

num_workers: int#: 0

persistent_workers: bool#: False

finalize() → None#

build_datasets(context: megatron.bridge.data.base.DatasetBuildContext)#

class bridge.models.diffusion_gemma.data.DiffusionGemmaJSONLDataset( path: str, *, processor: Any, config: bridge.models.diffusion_gemma.data.DiffusionGemmaDatasetConfig, )#

Bases: torch.utils.data.Dataset

Load message/image prompts and completions without fitting them all on a GPU.

Rows are {"messages": [...], "completion": "..."}. Image items contain local image paths relative to the JSONL file. Each completion is split into fixed-length canvases; previous clean blocks are appended to the encoder context. The current block and future blocks never enter that context.