Function validate

nemo_rl/algorithms/sft.py:234–349

Run validation on the validation dataset.

(
    policy: PolicyInterface,
    val_dataloader: Optional[StatefulDataLoader],
    tokenizer,
    loss_fn,
    step: int,
    master_config: MasterConfig,
    val_batches: int,
    val_batch_size: int,
    val_mbs: int,
)

# Training & Validation
# =======================================================
def validate(
    policy: PolicyInterface,
    val_dataloader: Optional[StatefulDataLoader],
    tokenizer,
    loss_fn,
    step: int,
    master_config: MasterConfig,
    val_batches: int,
    val_batch_size: int,
    val_mbs: int,
):
    """Run validation on the validation dataset."""
    if val_dataloader is None:
        assert master_config["sft"]["val_period"] <= 0, (
            "val_dataloader is None, so sft.val_period must be <= 0"
        )
        print("  ⚠️ No validation dataloader provided, skipping validation")
        return {}, {}

    timer = Timer()

    with timer.time("total_validation_time"):
        print(f"▶ Starting validation at step {step}...")

        # Show a progress indicator for validation
        # val_total = len(val_dataloader)

        val_metrics = {"val_loss": 0.0}
        sum_num_valid_tokens = 0

        policy.prepare_for_training()
        for batch_idx, val_batch in enumerate(val_dataloader):
            ## add loss mask based on role to every message
            add_loss_mask_to_message_log(
                val_batch["message_log"],
                roles_to_train_on=["assistant"],
            )

            cat_and_padded, input_lengths = batched_message_log_to_flat_message(
                val_batch["message_log"],
                pad_value_dict={"token_ids": tokenizer.pad_token_id},
                make_sequence_length_divisible_by=master_config["policy"][
                    "make_sequence_length_divisible_by"
                ],
            )

            val_data: BatchedDataDict = BatchedDataDict(
                {
                    "input_ids": cat_and_padded["token_ids"],
                    "input_lengths": input_lengths,
                    "token_mask": cat_and_padded["token_loss_mask"],
                    "sample_mask": val_batch["loss_multiplier"],
                }
            )

            # update multimodal data
            val_data.update(cat_and_padded.get_multimodal_dict(as_tensors=False))
            # When running validation with drop_last=False, we might end up with a partial batch.