SIGN IN SIGN UP

feat(datasets): Datum record and collate-time THD sequence packing (#3514)

* feat(datasets): Datum record and collate-time THD sequence packing

Add Datum, the typed single-example input record (input_ids + loss_inputs
side-channel for weights / logprobs / advantages / target_tokens), and
collate_datums, which delegates to the existing canonical collaters:
padded [B, T] via default_collater, or — with packed=True — one flat
[1, total_tokens] THD pack.

The packing itself is a small concatenation helper,
pack_features_for_thd, that emits the pre-packed record schema
packed_sequence_thd_collater already accepts (flat tokens, per-sequence
position_id resets, seq_lens); no second packing implementation. The
same helper powers thd_packing_collater, a YAML-selectable collate_fn
that gives unpacked datasets (e.g. ChatDataset) cross-sample sequence
packing at collate time, without an offline pack_dataset pass. Deciding
which examples share a pack stays with the caller (sampler or an RL
framework's microbatcher).

Datum mirrors the equivalent record in the tinker training API and
NeMo-RL's DatumSpec vocabulary; it is the input contract the upcoming
Engine training API (#2556) consumes.

Tests: 76 new/updated across test_datum.py and test_utils.py; full
tests/unit_tests/datasets suite green (1302 passed).

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(datasets): drop the unused packing collater and Datum serialization

Review of this PR turned up API with no consumer and one silent failure:

- thd_packing_collater had no caller outside its own tests, and could not
  be selected safely anyway: DataloaderConfig.emits_thd is an identity check
  against packed_sequence_thd_collater, so a YAML choosing it would report
  emits_thd=False while emitting qkv_format=thd batches, silently taking the
  wrong branch in validation packing and the CP+PP microbatch override.
  ThdPackingConfig already covers this job. pack_features_for_thd stays --
  it has real callers (collate_datums here, molt externally).
- Datum.to / to_dict / from_dict had zero callers; the Engine moves the
  collated dict to device, not the Datum.
- collate_datums intersected the datums' loss_inputs keys, so one datum
  missing 'weights' silently dropped the loss mask for the whole batch. It
  now raises.

Also documents which way a length-1 loss_inputs entry is read on a
single-token sequence, where per-token and per-sample shapes coincide.

Tests: covers the new error, the weights-absent labels path, and the
single-token tie-break; drops the tests of the removed API.

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* fix(datasets): emit attention_mask from Datum.to_features

A Datum holds only real tokens, but to_features did not say so, so the
padded collate had to fall back to inferring padding from the pad token
value: padding_mask = (input_ids == pad_id). That misreads any real token
whose id equals the pad id as padding -- pad_token_id == eos_token_id is a
common config, and default_collater's own comment records the consequence
(real eos/separator tokens masked out of the MoE experts).

The padded batch also carried no attention_mask at all, so a model forward
over it attended to the padding it had just added.

Emitting a ones mask fixes both: the collater now derives padding_mask from
the mask it padded, and the batch carries the mask the model needs. Packed
batches are unaffected -- pack_features_for_thd builds its own record and
THD describes boundaries with seq_lens.

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

* refactor(datasets): rename Datum.loss_inputs to loss_fn_inputs

Matches the field name in the tinker API this contract mirrors
(thinking-machines-lab/tinker-cookbook, recipes/rl_loop.py:233 and
recipes/sdft/sdft_test.py:215), so an algorithm author moving between the
two reads the same key.

Only the name changes. tinker wraps its payloads in ModelInput/TensorData
for transport across a hosted service; this API is in-process, so plain
tensors and input_ids stay.

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>

---------

Signed-off-by: HuiyingLi <willwin.lee@gmail.com>
H
Huiying committed
2afe8b7aeb54181bc2b1067bd77d2766c4801f49
Parent: 265b454
Committed by GitHub <noreply@github.com> on 8/14/2026, 6:44:49 PM