nemo_automodel.components.datasets.datum

Typed input contract for training: :class:Datum and :func:collate_datums.

A Datum is the single-example input boundary between user/algorithm code (SFT, RL post-training) and the training loop. It lives in components.datasets because feeding and collating examples is a data concern — and, crucially, because that lets :func:collate_datums reuse the canonical collaters (default_collater for padded [B, T] and packed_sequence_thd_collater for THD) instead of forking a second padding/packing implementation that could drift from them.

The companion output contract (ModelOutput and the per-token extraction helpers) lives in components.training — that side touches model logits, so it is a forward concern, not a dataset one.

Conventions

A Datum holds one sequence. input_ids is 1-D, shape [T].
loss_fn_inputs carries everything the loss needs, aligned to input_ids token positions (length T) for per-token entries:

=============== ======================================================= key meaning =============== ======================================================= target_tokens next-token targets, shape [T] (becomes labels) weights per-token loss mask / weight (0 disables a position) logprobs old/behavior-policy logprobs (importance sampling) advantages advantage signal (PPO/GRPO), per-token or per-sample =============== =======================================================
Masking convention matches the codebase: a target position with weights == 0 becomes ignore_index (default -100) in labels.

Module Contents

Classes

Name	Description
`Datum`	A single training example.

Functions

Name	Description
`collate_datums`	Collate a list of :class:`Datum` into a model-ready batch dict.

Data

CROSS_ENTROPY_IGNORE_IDX

__all__

API

class nemo_automodel.components.datasets.datum.Datum(
    input_ids: torch.Tensor,
    loss_fn_inputs: dict[str, torch.Tensor] = dict()
)

Dataclass

A single training example.

Parameters:

input_ids

torch.Tensor

1-D LongTensor of token ids, shape [T].

loss_fn_inputs

dict[str, torch.Tensor]Defaults to dict()

per-key tensors the loss consumes. Per-token entries are 1-D and length T; per-sample entries are scalar or shape [1]. See the module docstring for the well-known keys.

input_ids

Tensor

loss_fn_inputs

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

seq_len

int

Number of tokens in this example.

nemo_automodel.components.datasets.datum.Datum.__post_init__() -> None

nemo_automodel.components.datasets.datum.Datum.to_features(
    ignore_index: int = CROSS_ENTROPY_IGNORE_IDX
) -> dict[str, list[int]]

Emit the per-example dict the canonical collaters expect.

Every position of a Datum is a real token, so attention_mask is all ones: it tells the padded collater exactly which positions it added, instead of leaving it to infer them from the pad token value — which misreads a real token that happens to equal the pad id (commonly pad_token_id == eos_token_id) as padding.

labels is included only when loss_fn_inputs["target_tokens"] is present, with positions where loss_fn_inputs["weights"] == 0 set to ignore_index. Only integer token fields are emitted here — the collaters cast to LongTensor; float side-inputs are batched separately by :func:collate_datums.

Returns: dict[str, list[int]]

{"input_ids": [...], "attention_mask": [...], "labels": [...]}

nemo_automodel.components.datasets.datum.collate_datums(
    datums: list[nemo_automodel.components.datasets.datum.Datum],
    packed: bool = False,
    pad_seq_len_divisible: int | None = None,
    ignore_index: int = CROSS_ENTROPY_IGNORE_IDX
) -> dict[str, torch.Tensor]

Collate a list of :class:Datum into a model-ready batch dict.

Token fields are delegated to the existing canonical collaters so the padded / THD schema (attention_mask / qkv_format / seq_lens) is produced by the same code paths the dataset pipeline uses — no fork:

packed=False → default_collater (padded [B, T]).
packed=True → :func:pack_features_for_thd concatenates all datums into one pre-packed record, then packed_sequence_thd_collater emits the flat [1, total_tokens] THD schema (qkv_format="thd", per-sequence seq_lens for splitting outputs back per datum).

Float per-token side-inputs (every loss_fn_inputs key shared by all datums except target_tokens, e.g. weights / logprobs / advantages) are batched under their own key — this is the part the token collaters cannot carry (they cast to LongTensor). Padded mode right-pads them to the collated width and stacks to [B, T]; packed mode concatenates them in datum order to [1, total_tokens], aligned with input_ids. Per-sample (scalar / length-1) entries are stacked into a [num_datums] tensor without padding in both modes. A length-1 entry on a single-token sequence matches both shapes; it is read as per-token.

Parameters:

datums

list[Datum]

examples for this microbatch. Must be non-empty. One Datum is treated as one sequence.

packed

boolDefaults to False

pack all datums into one flat [1, total_tokens] THD row instead of the padded [B, T] layout.

pad_seq_len_divisible

int | NoneDefaults to None

pad sequence length to a multiple of this value (padded mode only; TP/CP/FP8 alignment).

ignore_index

intDefaults to CROSS_ENTROPY_IGNORE_IDX

label value for masked positions.

Returns: dict[str, torch.Tensor]

The collater output dict, augmented with the float side-input tensors.

nemo_automodel.components.datasets.datum.CROSS_ENTROPY_IGNORE_IDX = -100

nemo_automodel.components.datasets.datum.__all__ = ['Datum', 'collate_datums']