nemo_automodel.components.datasets.vlm.dspark_collate

Multimodal (image+text) data pipeline for DSpark speculative-decoding training.

DSpark’s own anchor-sampling / label-gathering logic (dspark.common, and every draft model’s forward()) expects a single, unshifted input_ids of length T plus a loss_mask of the same length T, and derives its own future-token targets internally (see label_offsets/safe_label_indices in each draft’s forward()). This is different from :func:~.collate_fns.default_collate_fn, which additionally shifts (labels = labels[:, 1:]) and truncates (input_ids = input_ids[:, :-1]) for the standard “predict next token” causal-LM loss — feeding DSpark that already-shifted, truncated output would silently double-shift / misalign labels. :func:dspark_vlm_collate_fn reuses the same tokenization + label-marking building blocks as :func:~.collate_fns.default_collate_fn but stops before that shift-and-truncate step.

Module Contents

Functions

Name	Description
`build_dspark_vlm_dataloader`	Build a multimodal DataLoader for DSpark training.
`dspark_vlm_collate_fn`	Collate multimodal conversations into a DSpark-ready batch.

Data

__all__

API

nemo_automodel.components.datasets.vlm.dspark_collate.build_dspark_vlm_dataloader(
    dataset_cfg,
    processor,
    batch_size: int,
    max_length: int,
    shuffle: bool,
    num_workers: int = 0,
    distributed: bool = False,
    dp_mesh = None
) -> torch.utils.data.DataLoader

Build a multimodal DataLoader for DSpark training.

dataset_cfg is any config node exposing .instantiate() (the repo’s standard _target_ convention, e.g. nemo_automodel.components.datasets .vlm.datasets.make_medpix_dataset or any other function under that module returning {"conversation": [...]} examples) — the same convention the VLM finetune recipe uses for its own dataset: config block. Deliberately simpler than that recipe’s build_dataloader: no packing, no pipeline- parallel microbatch chunking, no mRoPE position-id generation, since none of that applies to DSpark’s non-packed, non-pipelined target-capture call (MiniMax M3’s text decoder defaults to plain sequential position ids regardless of whether images are present).

nemo_automodel.components.datasets.vlm.dspark_collate.dspark_vlm_collate_fn(
    examples: typing.Sequence[typing.Dict[str, typing.Any]],
    processor,
    max_length: int
) -> typing.Dict[str, torch.Tensor]

Collate multimodal conversations into a DSpark-ready batch.

Unlike :func:~.collate_fns.default_collate_fn, this does not shift labels or truncate input_ids/attention_mask by one token: it returns a loss_mask the same length as input_ids, matching what DSpark’s anchor sampling and every draft’s forward() expect. labels itself is dropped — DSpark re-derives its own future-token targets from input_ids + loss_mask, so carrying a separate labels tensor would be dead weight moved to device every batch for no consumer.

max_length is required (unlike default_collate_fn’s optional one): DSpark needs a fixed shape across every batch/rank/step for its DFlash attention mask and the FSDP-sharded target’s forward to stay consistent.

Every conversation without an image/video gets the same fake-image injection default_collate_fn uses (:func:~.fake_image .inject_fake_image_into_conversation): MiniMax M3’s vision_tower is its own FSDP2-sharded unit, so a batch mixing text-only and image-containing samples across data-parallel ranks would have some ranks skip the vision_tower’s all-gather collective while others don’t, hanging training. The fake image’s vision tokens get attention_mask = 0 (:func:~.fake_image.mask_fake_vision_tokens_batch) so they never influence the captured hidden states.

nemo_automodel.components.datasets.vlm.dspark_collate.__all__ = ['dspark_vlm_collate_fn', 'build_dspark_vlm_dataloader']