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