nemo_automodel.components.datasets.vlm.loader

Typed construction for VLM processors, datasets, packing, and dataloaders.

Module Contents

Classes

Name	Description
`RankedVlmDatasetConfig`	Dataset config whose build requires the runtime data-parallel rank.
`VlmCollatorConfig`	Declarative VLM collator factory.
`VlmDataloaderBuild`	Materialized VLM dataloader and its processor.
`VlmDataloaderConfig`	Typed construction config for the complete VLM input pipeline.
`VlmProcessorConfig`	Declarative processor factory and keyword arguments.

Data

VlmCollateFn

__all__

logger

API

class nemo_automodel.components.datasets.vlm.loader.RankedVlmDatasetConfig()

Protocol

Dataset config whose build requires the runtime data-parallel rank.

builds_with_data_parallel_rank

bool

nemo_automodel.components.datasets.vlm.loader.RankedVlmDatasetConfig.build(
    rank: int,
    world_size: int
) -> object

Build the per-rank dataset shard.

class nemo_automodel.components.datasets.vlm.loader.VlmCollatorConfig(
    factory: collections.abc.Callable[..., object],
    kwargs: dict[str, object] = dict()
)

Dataclass

Declarative VLM collator factory.

factory

Callable[..., object]

Function that accepts examples and the runtime processor.

kwargs

dict[str, object] = field(default_factory=dict)

Declarative keyword arguments bound once while building the dataloader.

nemo_automodel.components.datasets.vlm.loader.VlmCollatorConfig.build(
    processor: transformers.ProcessorMixin | None
) -> nemo_automodel.components.datasets.vlm.loader.VlmCollateFn

Bind the runtime processor to the configured collator.

class nemo_automodel.components.datasets.vlm.loader.VlmDataloaderBuild(
    dataloader: torch.utils.data.DataLoader,
    processor: transformers.ProcessorMixin | None
)

Dataclass

Materialized VLM dataloader and its processor.

dataloader

DataLoader

processor

ProcessorMixin | None

class nemo_automodel.components.datasets.vlm.loader.VlmDataloaderConfig(
    dataset_config: nemo_automodel.components.datasets.loader.DatasetConfig | nemo_automodel.components.datasets.vlm.loader.RankedVlmDatasetConfig,
    processor_config: nemo_automodel.components.datasets.vlm.loader.VlmProcessorConfig = VlmProcessorConfig(),
    pretokenization: nemo_automodel.components.datasets.vlm.datasets.PreTokenizedDatasetWrapperConfig | None = None,
    packing: nemo_automodel.components.datasets.vlm.neat_packing_vlm.NeatPackConfig | None = None,
    collator: nemo_automodel.components.datasets.vlm.loader.VlmCollatorConfig | None = None,
    chat_template: str | None = None,
    shuffle: bool = True,
    num_workers: int = 0,
    pin_memory: bool = False,
    persistent_workers: bool = False,
    prefetch_factor: int | None = None,
    drop_last: bool = False
)

Dataclass

Typed construction config for the complete VLM input pipeline.

chat_template

str | None = None

collator

VlmCollatorConfig | None = None

dataset_config

DatasetConfig | RankedVlmDatasetConfig

drop_last

bool = False

num_workers

int = 0

packing

NeatPackConfig | None = None

persistent_workers

bool = False

pin_memory

bool = False

prefetch_factor

int | None = None

pretokenization

PreTokenizedDatasetWrapperConfig | None = None

processor_config

VlmProcessorConfig = field(default_factory=VlmProcessorConfig)

shuffle

bool = True

nemo_automodel.components.datasets.vlm.loader.VlmDataloaderConfig._build_source(
    pretrained_model_name_or_path: str,
    dp_rank: int,
    dp_world_size: int,
    dataset_build_context: contextlib.AbstractContextManager[object] | None
) -> tuple[object, transformers.ProcessorMixin | None]

Build the processor and raw dataset under the caller-owned ordering context.

nemo_automodel.components.datasets.vlm.loader.VlmDataloaderConfig.build(
    pretrained_model_name_or_path: str,
    dp_rank: int,
    dp_world_size: int,
    batch_size: int,
    dataset_build_context: contextlib.AbstractContextManager[object] | None = None,
    get_rope_index: collections.abc.Callable[..., object] | None = None,
    packing_attn_implementation: str | None = None,
    pp_n_microbatches: int | None = None
) -> nemo_automodel.components.datasets.vlm.loader.VlmDataloaderBuild

Build the processor, dataset wrappers, sampler, collator, and dataloader.

Parameters:

pretrained_model_name_or_path

str

Runtime model identifier used to build the processor.

dp_rank

int

Rank within the data-parallel group.

dp_world_size

int

Size of the data-parallel group.

batch_size

int

Runtime local training batch size.

dataset_build_context

AbstractContextManager[object] | NoneDefaults to None

Optional rank-ordering context used only for processor and source-dataset build.

get_rope_index

Callable[..., object] | NoneDefaults to None

Optional model callback used to create packed multimodal position IDs.

packing_attn_implementation

str | NoneDefaults to None

Resolved attention backend for packed-mask construction.

pp_n_microbatches

int | NoneDefaults to None

Optional pipeline microbatch count used to pre-chunk media tensors.

Returns: VlmDataloaderBuild

Named result containing the stateful dataloader and runtime processor.

nemo_automodel.components.datasets.vlm.loader.VlmDataloaderConfig.resolve_packing_attn_implementation(
    model_attn_implementation: str | None,
    cp_size: int
) -> str | None

Resolve the packed-collator mask backend.

Parameters:

model_attn_implementation

str | None

Attention implementation selected by the model config.

cp_size

int

Runtime context-parallel world size.

Returns: str | None

Attention implementation used to choose the packed attention-mask representation.

class nemo_automodel.components.datasets.vlm.loader.VlmProcessorConfig(
    factory: collections.abc.Callable[..., transformers.ProcessorMixin] | None = None,
    kwargs: dict[str, object] = dict()
)

Dataclass

Declarative processor factory and keyword arguments.

factory

Callable[..., ProcessorMixin] | None = None

Configured processor factory; None selects AutoProcessor.from_pretrained.

kwargs

dict[str, object] = field(default_factory=dict)

Declarative keyword arguments for the configured processor factory.

nemo_automodel.components.datasets.vlm.loader.VlmProcessorConfig.build(
    pretrained_model_name_or_path: str
) -> transformers.ProcessorMixin | None

Build the configured processor.

Parameters:

pretrained_model_name_or_path

str

Runtime model identifier used by the default AutoProcessor factory.

Returns: ProcessorMixin | None

Processor instance, or None when the model has no compatible AutoProcessor.

nemo_automodel.components.datasets.vlm.loader.VlmCollateFn = Callable[[list[object]], object]

nemo_automodel.components.datasets.vlm.loader.__all__ = ['RankedVlmDatasetConfig', 'VlmCollatorConfig', 'VlmDataloaderBuild', 'VlmDatalo...

nemo_automodel.components.datasets.vlm.loader.logger = logging.getLogger(__name__)