nemo_automodel.components.datasets.zephon.loader

Checkpointable Zephon stream with all-rank aggregation of opaque loader state.

Module Contents

Classes

Name	Description
`ZephonDataLoader`	One owned stream; iteration never reconstructs the work source.

Functions

Name	Description
`convert_text_batch`	Collate text, retaining boundaries independently of the model’s positions.

Data

logger

API

class nemo_automodel.components.datasets.zephon.loader.ZephonDataLoader(
    pipeline: zephon.Pipeline,
    convert_batch: collections.abc.Callable[[SampleBatch], nemo_automodel.components.datasets.zephon.attention.TrainingBatch],
    output_format: nemo_automodel.components.datasets.packing.BatchOutputFormat = 'batched',
    packing_contract: nemo_automodel.components.datasets.packing.PackedSequenceContract
)

Bases: Iterator[TrainingBatch]

One owned stream; iteration never reconstructs the work source.

use_distributed_checkpointing selects BaseRecipe’s all-rank auxiliary state path so TP/PP peers participate in Zephon checkpoint collectives as well as data-owning ranks. StatefulDataLoader checkpoint routing saves only data-owning ranks and cannot provide that participation.

The original SampleBatch keeps Zephon replay/accounting metadata unchanged. The consumer calls to_training and applies lightweight model-input adaptations.

_component_tokens

Counter[int] = Counter()

_iterator

Iterator[SampleBatch] | None = None

_restored_checkpoint

= b''

nemo_automodel.components.datasets.zephon.loader.ZephonDataLoader.__iter__() -> nemo_automodel.components.datasets.zephon.loader.ZephonDataLoader

nemo_automodel.components.datasets.zephon.loader.ZephonDataLoader.__next__() -> nemo_automodel.components.datasets.zephon.attention.TrainingBatch

nemo_automodel.components.datasets.zephon.loader.ZephonDataLoader.close() -> None

Release owned workers at completion, failure, or an explicit training cap.

nemo_automodel.components.datasets.zephon.loader.ZephonDataLoader.load_state_dict(
    state: collections.abc.Mapping[str, typing.Any]
) -> None

Restore opaque bytes before iteration; empty state requires a fresh loader.

Parameters:

state

Mapping[str, Any]

Mapping containing zephon checkpoint bytes. Empty bytes represent a stream that has never started or been restored. Nonempty bytes are validated and restored by Zephon.

nemo_automodel.components.datasets.zephon.loader.ZephonDataLoader.state_dict() -> dict[str, bytes]

Return opaque checkpoint bytes, or empty bytes for an unstarted stream.

DCP calls this method before loading. An unstarted loader returns its initial or restored state without initializing or calibrating a pipeline. Once iteration starts, every rank collects the complete aggregated state, including after exhaustion or close.

Returns: dict[str, bytes]

Mapping containing zephon bytes for the current stream position.

nemo_automodel.components.datasets.zephon.loader.convert_text_batch(
    sample_batch: zephon.SampleBatch,
    packing_contract: nemo_automodel.components.datasets.packing.PackedSequenceContract,
    ignore_index: int,
    mask_segment_starts: bool,
    allow_cross_document_attention: bool,
    output_format: nemo_automodel.components.datasets.packing.BatchOutputFormat = 'batched',
    pad_to_length: int | None = None,
    pad_token_id: int | None = None
) -> nemo_automodel.components.datasets.zephon.attention.TrainingBatch

Collate text, retaining boundaries independently of the model’s positions.

Parameters:

sample_batch

SampleBatch

Prepared documents with unshifted token IDs and scalar subdocument-reset positions [sequence] in each payload. Positions come from online packing, an offline export, or single-document defaults.

packing_contract

PackedSequenceContract

Model-derived return requests; mask and gathered-token requests apply only to batched output.

ignore_index

int

Loss exclusion value for shifted labels.

mask_segment_starts

bool

Mask targets at original subdocument boundaries.

allow_cross_document_attention

bool

Allow earlier subdocuments as context and number each packed document continuously for RoPE.

output_format

BatchOutputFormatDefaults to 'batched'

Batched mask/metadata route or native THD lengths.

pad_to_length

int | NoneDefaults to None

Unshifted target length when online packing is disabled. Shorter supplied documents are padded, with padding excluded from loss.

pad_token_id

int | NoneDefaults to None

Resolved model padding ID, required with pad_to_length.

Returns: TrainingBatch

Token IDs, labels, positions, padding and compact boundary/attention IDs

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