feat: support drop_last=False during validation (#1029)

Signed-off-by: ashors1 <ashors@nvidia.com>
Signed-off-by: Anna Shors <ashors@nvidia.com>

Author: ashors1

nemo_rl/algorithms/dpo.py

@@ -26,7 +26,7 @@
 from nemo_rl.algorithms.loss_functions import (
     DPOLossFn,
 )
-from nemo_rl.algorithms.utils import set_seed
+from nemo_rl.algorithms.utils import maybe_pad_last_batch, set_seed
 from nemo_rl.data import DataConfig
 from nemo_rl.data.datasets import AllTaskProcessedDataset, preference_collate_fn
 from nemo_rl.distributed.virtual_cluster import ClusterConfig, RayVirtualCluster
@@ -87,7 +87,14 @@ class MasterConfig(TypedDict):
 
 class DPOValMetrics(TypedDict):
     loss: float
+    sft_loss: float
+    preference_loss: float
     accuracy: float
+    rewards_chosen_mean: float
+    rewards_rejected_mean: float
+    num_valid_samples: float
+    global_valid_seqs: float
+    global_valid_toks: float
 
 
 # =======================================================
@@ -187,7 +194,7 @@ def setup(
                 ],
                 add_loss_mask=True,
             ),
-            drop_last=True,
+            drop_last=False,
         )
         for k, v in val_dataset.items()
     }
@@ -255,6 +262,15 @@ def add_ref_logprobs_to_data(dataloader, policy, master_config, is_val=False):
                 else master_config["policy"]["train_micro_batch_size"] * 2
             )
 
+            # when running validation with drop_last=False, we might end up with a partial batch.
+            # In this case, we pad the batch to the next multiple of micro_batch_size * dp_size.
+            dp_size = policy.sharding_annotations.get_axis_size("data_parallel")
+            if batch.size % (dp_size * micro_batch_size) != 0:
+                assert is_val, (
+                    "Partial batches should only happen during validation, but got a partial batch during training."
+                )
+                batch = maybe_pad_last_batch(batch, dp_size, micro_batch_size)
+
             ## append ref policy logprobs to batch
             logprobs = policy.get_reference_policy_logprobs(
                 batch,
@@ -342,7 +358,7 @@ def validate_one_dataset(
     with timer.time("total_validation_time"):
         print(f"▶ Starting validation at step {step} for `{dataset_name}` set..")
 
-        val_metrics = defaultdict(lambda: 0.0)
+        val_metrics = defaultdict(list)
         num_valid_batches = 0
         for batch_idx, val_batch in enumerate(
             add_ref_logprobs_to_data(val_dataloader, policy, master_config, is_val=True)
@@ -352,7 +368,7 @@ def validate_one_dataset(
                 val_batch,
                 loss_fn,
                 eval_mode=True,
-                gbs=val_batch_size * 2,
+                gbs=val_batch.size,
                 mbs=val_mbs * 2,
             )
 
@@ -361,22 +377,61 @@ def validate_one_dataset(
                     "No validation metrics were collected for this batch."
                     " This is likely because there were no valid samples."
                 )
-
             else:
-                for k, v in val_results["all_mb_metrics"].items():
-                    if k in {"lr", "wd", "global_valid_seqs", "global_valid_toks"}:
-                        val_metrics[k] += np.mean(v).item()
-                    else:
-                        val_metrics[k] += np.sum(v).item()
+                for metric_name in DPOValMetrics.__annotations__.keys():
+                    reduction = (
+                        np.mean
+                        if metric_name in {"global_valid_seqs", "global_valid_toks"}
+                        else sum
+                    )
+                    val_metrics[metric_name] += [
+                        reduction(val_results["all_mb_metrics"][metric_name])
+                    ]
+
                 num_valid_batches += 1
 
             if val_batches > 0 and batch_idx >= val_batches - 1:
                 break
 
-        for k, v in val_metrics.items():
-            if k == "num_valid_samples":
-                continue
-            val_metrics[k] /= num_valid_batches
+        if num_valid_batches > 0:
+            sum_num_valid_samples = sum(val_metrics["num_valid_samples"])
+            global_valid_toks = sum(val_metrics["global_valid_toks"])
+            global_valid_seqs = sum(val_metrics["global_valid_seqs"])
+            val_metrics = DPOValMetrics(
+                num_valid_samples=sum_num_valid_samples,
+                global_valid_seqs=global_valid_seqs,
+                global_valid_toks=global_valid_toks,
+                **{
+                    metric_name: sum(
+                        [
+                            value * weight
+                            for value, weight in zip(
+                                val_metrics[metric_name],
+                                val_metrics["num_valid_samples"],
+                            )
+                        ]
+                    )
+                    / sum_num_valid_samples
+                    for metric_name in DPOValMetrics.__annotations__.keys()
+                    if metric_name
+                    not in {
+                        "num_valid_samples",
+                        "global_valid_seqs",
+                        "global_valid_toks",
+                    }
+                },
+            )
+        else:
+            warnings.warn(
+                "No validation metrics were collected."
+                " This is likely because there were no valid samples in the validation set."
+            )
+            val_metrics = DPOValMetrics(
+                **{
+                    metric_name: 0.0
+                    for metric_name in DPOValMetrics.__annotations__.keys()
+                }
+            )
 
         # Calculate validation metrics
         policy.prepare_for_training()

nemo_rl/algorithms/rm.py

@@ -26,7 +26,7 @@
 from nemo_rl.algorithms.loss_functions import (
     PreferenceLoss,
 )
-from nemo_rl.algorithms.utils import set_seed
+from nemo_rl.algorithms.utils import maybe_pad_last_batch, set_seed
 from nemo_rl.data import DataConfig
 from nemo_rl.data.datasets import (
     AllTaskProcessedDataset,
@@ -172,7 +172,7 @@ def setup(
                 ],
                 add_loss_mask=False,
             ),
-            drop_last=True,
+            drop_last=False,
         )
         for k, v in val_dataset.items()
     }
@@ -307,14 +307,20 @@ def validate_one_dataset(
         dict_val_metrics = defaultdict(list)
         num_valid_batches = 0
         for batch_idx, val_batch in enumerate(val_dataloader):
+            # When running validation with drop_last=False, we might end up with a partial batch.
+            # In this case, we pad the batch to the next multiple of micro_batch_size * dp_size.
+            if val_batch.size < val_batch_size * 2:
+                dp_size = policy.sharding_annotations.get_axis_size("data_parallel")
+                val_batch = maybe_pad_last_batch(val_batch, dp_size, val_mbs * 2)
+
             ## just run model fwd
             val_results = policy.train(
                 val_batch,
                 loss_fn,
                 eval_mode=True,
-                ## NOTE: we double the batch size here because each preference example corresponds to a pair of
-                ## examples, chosen and rejected, and the pair needs to be processed as part of the same microbatch.
-                gbs=val_batch_size * 2,
+                gbs=val_batch.size,
+                # NOTE: we double the batch size because each preference example corresponds to a pair of
+                # examples, chosen and rejected, and the pair needs to be processed as part of the same microbatch.
                 mbs=val_mbs * 2,
             )
 

nemo_rl/algorithms/sft.py

@@ -24,7 +24,7 @@
 from nemo_rl.algorithms.loss_functions import (
     NLLLoss,
 )
-from nemo_rl.algorithms.utils import set_seed
+from nemo_rl.algorithms.utils import maybe_pad_last_batch, set_seed
 from nemo_rl.data import DataConfig
 from nemo_rl.data.datasets import AllTaskProcessedDataset, rl_collate_fn
 from nemo_rl.data.interfaces import TaskDataSpec
@@ -150,7 +150,7 @@ def setup(
         batch_size=sft_config["val_global_batch_size"],
         shuffle=False,
         collate_fn=rl_collate_fn,
-        drop_last=True,
+        drop_last=False,
     )
 
     # ==========================
@@ -240,7 +240,7 @@ def validate(
         # val_total = len(val_dataloader)
 
         val_metrics = {"val_loss": 0.0}
-        num_valid_batches = 0
+        sum_num_valid_tokens = 0
 
         policy.prepare_for_training()
         for batch_idx, val_batch in enumerate(val_dataloader):
@@ -269,13 +269,18 @@ def validate(
 
             # 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.
+            # Check if we need to pad the final batch to make it divisible by micro_batch_size * dp_size.
+            if val_data.size < val_batch_size:
+                dp_size = policy.sharding_annotations.get_axis_size("data_parallel")
+                val_data = maybe_pad_last_batch(val_data, dp_size, val_mbs)
 
             ## just run model fwd
             val_results = policy.train(
                 val_data,
                 loss_fn,
                 eval_mode=True,
-                gbs=val_batch_size,
+                gbs=val_data.size,
                 mbs=val_mbs,
             )
 
@@ -285,14 +290,17 @@ def validate(
                     " This is likely because there were no valid samples."
                 )
             else:
-                val_metrics["val_loss"] += float(val_results["loss"])
-                num_valid_batches += 1
+                num_valid_tokens = (
+                    val_data["sample_mask"].unsqueeze(-1) * val_data["token_mask"]
+                ).sum()
+                val_metrics["val_loss"] += float(val_results["loss"]) * num_valid_tokens
+                sum_num_valid_tokens += num_valid_tokens
 
             if val_batches > 0 and batch_idx >= val_batches - 1:
                 break
 
-        if num_valid_batches > 0:
-            val_metrics["val_loss"] /= num_valid_batches
+        if sum_num_valid_tokens > 0:
+            val_metrics["val_loss"] /= sum_num_valid_tokens
         else:
             warnings.warn(
                 "No validation metrics were collected."
@@ -306,7 +314,7 @@ def validate(
     timing_metrics = timer.get_timing_metrics(reduction_op="sum")
     validation_time = timing_metrics.get("total_validation_time", 0)
 
-    if num_valid_batches > 0:
+    if sum_num_valid_tokens > 0:
         # Print summary of validation results
         print("\n📊 Validation Results:")
         print(f"    • Validation loss: {val_metrics['val_loss']:.4f}")

nemo_rl/algorithms/utils.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+import math
 import random
 import warnings
 from functools import wraps
@@ -265,3 +267,62 @@ def get_tokenizer(
         processor.name_or_path = tokenizer.name_or_path
 
     return tokenizer if processor is None else processor
+
+
+def maybe_pad_last_batch(batch: dict, dp_size: int, mbs: int) -> dict:
+    """Pads the given batch so that its size is divisible by (mbs * dp_size).
+
+    Args:
+        batch (dict): The batch to pad.
+        dp_size (int): Data parallel size.
+        mbs (int): Micro batch size.
+
+    Returns:
+        dict: The padded batch.
+    """
+    min_padding = (math.ceil(batch.size / (mbs * dp_size)) * mbs * dp_size) - batch.size
+    if min_padding > 0:
+        print(f"Padding last validation batch with {min_padding} padding samples")
+        # Pad input_ids
+        batch["input_ids"] = torch.cat(
+            [
+                batch["input_ids"],
+                batch["input_ids"][-1].unsqueeze(0).repeat(min_padding, 1),
+            ]
+        )
+        # Pad input_lengths
+        batch["input_lengths"] = torch.cat(
+            [
+                batch["input_lengths"],
+                batch["input_lengths"][-1].unsqueeze(0).repeat(min_padding),
+            ]
+        )
+        if "token_mask" in batch:
+            # Pad token_mask
+            batch["token_mask"] = torch.cat(
+                [
+                    batch["token_mask"],
+                    batch["token_mask"][-1].unsqueeze(0).repeat(min_padding, 1),
+                ]
+            )
+        # Pad sample_mask
+        batch["sample_mask"] = torch.cat(
+            [
+                batch["sample_mask"],
+                torch.zeros_like(batch["sample_mask"][-1])
+                .unsqueeze(0)
+                .repeat(min_padding),
+            ]
+        )
+
+        if "reference_policy_logprobs" in batch:
+            # Pad reference_policy_logprobs
+            batch["reference_policy_logprobs"] = torch.cat(
+                [
+                    batch["reference_policy_logprobs"],
+                    batch["reference_policy_logprobs"][-1]
+                    .unsqueeze(0)
+                    .repeat(min_padding, 1),
+                ]
+            )
+    return batch

tests/unit/algorithms/test_dpo.py

@@ -14,14 +14,31 @@
 
 from unittest.mock import MagicMock
 
+import numpy as np
 import torch
 
 from nemo_rl.algorithms.dpo import add_ref_logprobs_to_data
+from nemo_rl.distributed.batched_data_dict import BatchedDataDict
+from nemo_rl.distributed.named_sharding import NamedSharding
 
 
 class MockPolicy:
     def __init__(self, logprobs):
         self.logprobs = logprobs
+        self.sharding_annotations = NamedSharding(
+            layout=np.arange(2).reshape(
+                1,  # PP
+                -1,  # DP
+                1,  # CP
+                1,  # TP
+            ),
+            names=[
+                "pipeline_parallel",
+                "data_parallel",
+                "context_parallel",
+                "tensor_parallel",
+            ],
+        )
 
     def get_reference_policy_logprobs(self, batch, micro_batch_size):
         return {"reference_logprobs": self.logprobs}
@@ -30,7 +47,7 @@ def get_reference_policy_logprobs(self, batch, micro_batch_size):
 def test_add_logprobs_to_batch():
     """Test that add_ref_logprobs_to_data correctly adds reference policy logprobs to batches."""
     # Create mock data
-    batch_size = 2
+    batch_size = 8
     seq_len = 4
     vocab_size = 16
 
@@ -45,7 +62,7 @@ def test_add_logprobs_to_batch():
 
     # Create a mock dataloader that yields our mock batch
     mock_dataloader = MagicMock()
-    mock_dataloader.__iter__.return_value = iter([mock_batch])
+    mock_dataloader.__iter__.return_value = iter([BatchedDataDict(mock_batch)])
 
     # Create a mock policy that returns our mock logprobs
     mock_policy = MockPolicy(mock_logprobs)