feat(vlm): enable validation under pipeline parallelism

examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml

@@ -93,6 +93,7 @@ validation_dataset:
 
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.gemma4_prefix_collate_fn
 

examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml

@@ -98,6 +98,7 @@ validation_dataset:
 
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.gemma4_prefix_collate_fn
 

examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml

@@ -139,6 +139,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml

@@ -117,6 +117,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml

@@ -121,6 +121,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml

@@ -108,6 +108,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml

@@ -106,6 +106,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/mistral4/mistral4_medpix.yaml

@@ -100,6 +100,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
 

examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml

@@ -118,6 +118,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml

@@ -114,6 +114,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml

@@ -123,6 +123,7 @@ validation_dataset:
 validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   num_workers: 1
+  drop_last: true
   collate_fn:
     _target_: nemo_automodel.components.datasets.vlm.collate_fns.default_collate_fn
     max_length: 2048

nemo_automodel/components/distributed/pipelining/autopipeline.py

@@ -270,7 +270,7 @@ def step(
         losses: list[torch.Tensor] | None = None,
         **kwargs: Any,
     ) -> Any:
-        """Run one pipeline schedule step with model-owned input chunking.
+        """Run one forward-and-backward pipeline schedule step with model-owned input chunking.
 
         Args:
             model_input: Tensor of shape [batch, ...] containing the first
@@ -283,22 +283,80 @@ def step(
                 model-defined layouts; model-owned metadata identifies any
                 nonstandard batch axis.
 
+        Returns:
+            The value returned by the underlying PyTorch pipeline schedule.
+        """
+        return self._run_schedule(model_input, forward_only=False, target=target, losses=losses, **kwargs)
+
+    def eval(
+        self,
+        model_input: torch.Tensor,
+        *,
+        target: torch.Tensor | None = None,
+        losses: list[torch.Tensor] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Run one forward-only pipeline schedule step with model-owned input chunking.
+
+        Same inputs as :meth:`step`, but the schedule runs no backward, so callers
+        that only need a loss (validation) do not build or free a backward graph.
+
+        Args:
+            model_input: Tensor of shape [batch, ...] containing the first
+                pipeline stage's input. Ignored on ranks without the first stage.
+            target: Tensor with a model-defined target layout, or ``None`` on
+                ranks without the last pipeline stage.
+            losses: Mutable list populated with scalar loss tensors, or ``None``
+                on ranks without the last pipeline stage.
+            **kwargs: Keyword schedule inputs. Tensor values may have arbitrary
+                model-defined layouts; model-owned metadata identifies any
+                nonstandard batch axis.
+
+        Returns:
+            The value returned by the underlying PyTorch pipeline schedule.
+        """
+        return self._run_schedule(model_input, forward_only=True, target=target, losses=losses, **kwargs)
+
+    def _run_schedule(
+        self,
+        model_input: torch.Tensor,
+        *,
+        forward_only: bool,
+        target: torch.Tensor | None,
+        losses: list[torch.Tensor] | None,
+        **kwargs: Any,
+    ) -> Any:
+        """Drive the pipeline schedule with the model-owned kwargs chunk spec installed.
+
+        Args:
+            model_input: Tensor of shape [batch, ...] containing the first
+                pipeline stage's input. Ignored on ranks without the first stage.
+            forward_only: Whether to run the schedule's forward-only ``eval`` entry
+                point instead of ``step``.
+            target: Tensor with a model-defined target layout, or ``None`` on
+                ranks without the last pipeline stage.
+            losses: Mutable list populated with scalar loss tensors, or ``None``
+                on ranks without the last pipeline stage.
+            **kwargs: Keyword schedule inputs. Tensor values may have arbitrary
+                model-defined layouts.
+
         Returns:
             The value returned by the underlying PyTorch pipeline schedule.
         """
         schedule = self._info.schedule
         if schedule is None:
             raise RuntimeError("AutoPipeline.build() must be called before running a PP schedule step")
 
+        run = schedule.eval if forward_only else schedule.step
         schedule_args = (model_input,) if self._info.has_first_stage else ()
         kwargs_chunk_spec = self._get_schedule_kwargs_chunk_spec(kwargs)
         if kwargs_chunk_spec is None:
-            return schedule.step(*schedule_args, target=target, losses=losses, **kwargs)
+            return run(*schedule_args, target=target, losses=losses, **kwargs)
 
         previous_kwargs_chunk_spec = schedule._kwargs_chunk_spec
         schedule._kwargs_chunk_spec = kwargs_chunk_spec
         try:
-            return schedule.step(*schedule_args, target=target, losses=losses, **kwargs)
+            return run(*schedule_args, target=target, losses=losses, **kwargs)
         finally:
             schedule._kwargs_chunk_spec = previous_kwargs_chunk_spec
 

nemo_automodel/recipes/base_recipe.py

@@ -833,6 +833,22 @@ def _dp_allreduce(self, tensor, op=dist.ReduceOp.SUM, include_cp: bool = False):
             tensor = tensor.cpu()
         return tensor
 
+    def _broadcast_from_last_pp_stage(self, tensor: torch.Tensor) -> torch.Tensor:
+        """Broadcast a PP last-stage scalar to the other ranks in its pipeline group.
+
+        Args:
+            tensor: Scalar tensor on the current device. On the last pipeline stage
+                it holds the value to publish; on every other stage it is only a
+                receive buffer and its contents are overwritten in place.
+
+        Returns:
+            The same tensor, now holding the last stage's value on every rank.
+        """
+        pp_group = self.device_mesh["pp"].get_group()
+        pp_src_rank = dist.get_global_rank(pp_group, dist.get_world_size(pp_group) - 1)
+        dist.broadcast(tensor, src=pp_src_rank, group=pp_group)
+        return tensor
+
     def _make_progress_bar(self, total: int | None = None, initial: int = 0):
         """Create a tqdm progress bar on rank 0; returns None on other ranks.
 

nemo_automodel/recipes/llm/train_ft.py

@@ -1192,13 +1192,6 @@ def _forward_backward_step(
                 if is_train:
                     (local_loss * self._get_dp_group_size(include_cp=True)).backward()
 
-    def _broadcast_from_last_pp_stage(self, tensor: torch.Tensor) -> torch.Tensor:
-        """Broadcast a PP last-stage scalar to the other ranks in its pipeline group."""
-        pp_group = self.device_mesh["pp"].get_group()
-        pp_src_rank = torch.distributed.get_global_rank(pp_group, torch.distributed.get_world_size(pp_group) - 1)
-        torch.distributed.broadcast(tensor, src=pp_src_rank, group=pp_group)
-        return tensor
-
     def _run_train_optim_step(self, batches, max_grad_norm: float | None = None):
         """Execute a single training step.