What the new Nemotron SFT path adds, next to the generic one that already existed

Companion illustration for NVIDIA-NeMo/RL #4108 — all permalinks pinned to d044f40.

Before
One way to turn a multimodal SFT dataset into training batches: the generic cooker plus the generic task encoder.
This PR
Adds a second family for Nemotron models — ten cookers and one task encoder, built as a subclass of the generic one.
Status
Added by this PR. This page describes the new design, not a defect. The problems found in review are on the other two pages.

Three pages on this PR, one question each — this one: what the new path is and where it rejoins the old one.

NeMo-RL already had one way to turn a multimodal SFT dataset into training batches: a cooker that reads the raw sample, and a task encoder that tokenizes it. This PR adds a second family for Nemotron models — ten new cookers and one new task encoder. The interesting part is how little of the pipeline it actually replaces. NemotronMultiModalTaskEncoder subclasses the generic encoder and overrides exactly two methods, then sets one flag — and that flag is what decides who builds the training loss mask.

Both paths, and where they become one

One sample → one packed batch. Blue = new here, green = inherited. differs shared generic path (already existed) cook_conversation reads messages[] + media manifest preencode_sample adapter.encode() postencode_sample returns it unchanged nemotron path (new in this PR) cook_nemotron_conversation reads conversation[].fragments[] preencode_sample omni.preencode() postencode_sample omni.postencode() both continue into the same code shared — the Nemotron encoder inherits all of this select_samples_to_pack() groups samples into packs — not overridden prepare_packed_sft_batch() builds the tensors — one branch differs, below loss_mask_mode is None → generic add_loss_mask_to_message_log() builds it = "precomputed" → nemotron only checks the mask the encoder wrote

They merge at the task encoder itself. Packing, grouping and batching are inherited, not rewritten. After the merge exactly one thing still differs: the flag loss_mask_mode, which the Nemotron encoder sets to "precomputed" at construction and the generic one leaves as None (default).

The same tiny sample down both paths

One conversation: 2 turns, 1 image, 7 tokens after tokenizing, of which 3 are the assistant's answer and so are the only ones training should learn from.

generic nemotron after preencode token_ids = 7 tokens token_ids = 7 tokens (no mask yet) token_loss_mask = 7 values after postencode unchanged mask filled in by the encoder (still no mask) [0,0,0,0,1,1,1] ← 3 ones in prepare_packed_ reads the roles and checks the mask is 7 long sft_batch builds [0,0,0,0,1,1,1] and only 0s and 1s, then uses it from "assistant"

Both paths hand the trainer the same 7 tokens and the same 3 unmasked positions. The difference is only who computed the mask. The generic path derives it at batch time from the message roles; the Nemotron path carries it from the encoder, because its source format marks trainable spans per fragment rather than per whole turn.

Why keep both

genericnemotron
Source formatmessages[] plus a media manifest, in one webdataset sampleconversation[].fragments[], with separate media stores and a cache pool
Loss maskderived from roles at batch timecomputed in the encoder, checked at batch time
Mediaimagesimages, video and audio
Packingsame code for both

They are not two versions of one thing. They read datasets produced by different teams in different shapes, so they change for different reasons — which is why the PR subclasses rather than adding if model_family == … branches to the generic encoder. The cost is that a third family would need edits in more than one place, because the typed options block (EnergonTaskEncoderOptions) is shared by every encoder but documents itself as Nemotron's, and the loader only forwards it when the encoder is named nemotron_multimodal (sft_dataloader.py:381).

What to take away. The new Nemotron path is a two-method subclass, not a parallel pipeline — packing and batching are shared, so a change there still only has to be made once. The one behaviour that genuinely forks is the loss mask, and it forks on a single flag. If you are reviewing this PR, that flag and the two overridden methods are where the real new logic lives; everything downstream is code that was already there.