feat: make only_unmask_final configurable in SFT

examples/configs/sft.yaml

@@ -12,6 +12,9 @@ sft:
   val_at_start: true
   val_at_end: false
   seed: 42
+  # If true, only the final message in each conversation is unmasked for loss
+  # computation. If false, all assistant messages are unmasked.
+  only_unmask_final: false
 
 checkpointing:
   enabled: true

nemo_rl/algorithms/sft.py

@@ -73,6 +73,10 @@ class SFTConfig(TypedDict):
     # final checkpoint has validation metrics, which is required for get_best_checkpoint_path().
     val_at_end: bool
     seed: int
+    # If True, only the final message in each conversation is unmasked for loss
+    # computation, regardless of role. If False, all messages with role in
+    # `roles_to_train_on` (currently hard-coded to ["assistant"]) are unmasked.
+    only_unmask_final: bool
 
 
 class MasterConfig(BaseModel, extra="allow"):
@@ -271,6 +275,7 @@ def validate(
             add_loss_mask_to_message_log(
                 val_batch["message_log"],
                 roles_to_train_on=["assistant"],
+                only_unmask_final=master_config.sft["only_unmask_final"],
             )
 
             cat_and_padded, input_lengths = batched_message_log_to_flat_message(
@@ -434,6 +439,7 @@ def sft_train(
                     add_loss_mask_to_message_log(
                         batch["message_log"],
                         roles_to_train_on=["assistant"],
+                        only_unmask_final=master_config.sft["only_unmask_final"],
                     )
 
                     cat_and_padded, input_lengths = batched_message_log_to_flat_message(

tests/unit/reference_configs/sft.yaml

@@ -12,6 +12,9 @@ sft:
   val_at_start: true
   val_at_end: false
   seed: 42
+  # If true, only the final message in each conversation is unmasked for loss
+  # computation. If false, all assistant messages are unmasked.
+  only_unmask_final: false
 
 checkpointing:
   enabled: true