From 1aa3d057ef61ac87e1bc72405b435c6c3cef096e Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sun, 30 Aug 2026 11:15:13 -0700 Subject: [PATCH 1/2] feat(vlm): enable validation under pipeline parallelism Route validation batches through the same _forward_backward_step path as training but via a forward-only AutoPipeline.eval() (which preserves the model-owned kwargs chunk spec that a bare schedule.eval() would drop), broadcast the last-stage loss, and require validation drop_last under PP since AutoPipeline uses a fixed outer batch size. The parity script gains a val_loss metric and the Gemma4 PP2 L2 test now asserts validation parity, so a stale PP validation skip cannot pass. Co-Authored-By: Claude Fable 5 Signed-off-by: HuiyingLi --- .../gemma4/gemma4_31b_tp4_pp2.yaml | 1 + .../gemma4/gemma4_31b_tp4_pp4.yaml | 1 + .../minimax_m3_vl_lora_pp4ep8_8node.yaml | 1 + .../minimax_m3_vl_sft_cp2_medpix_2k.yaml | 1 + .../minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml | 1 + .../mistral3p5/mistral3p5_128b_medpix.yaml | 1 + .../mistral3p5_128b_medpix_lora.yaml | 1 + .../mistral4/mistral4_medpix.yaml | 1 + .../qwen3_5_moe/qwen3_5_35b_neat_packing.yaml | 1 + .../stepfun/step3p7_medpix_200b_ep32pp4.yaml | 1 + ...step3p7_medpix_200b_lora_pp8ep8_8node.yaml | 1 + .../distributed/pipelining/autopipeline.py | 64 ++++++- nemo_automodel/recipes/base_recipe.py | 16 ++ nemo_automodel/recipes/llm/train_ft.py | 7 - nemo_automodel/recipes/vlm/finetune.py | 133 +++++++++++++-- .../L2_Parallelism_VLM_Gemma4_PP2_Parity.sh | 16 +- .../parallelism/compare_parallel_parity.py | 121 +++++++++----- .../pipelining/test_autopipeline.py | 41 +++++ .../recipes/test_finetune_vlm_cp_wiring.py | 129 ++++++++++++++- .../recipes/test_finetune_vlm_helpers.py | 156 +++++++++++++++++- .../test_compare_parallel_parity.py | 113 +++++++++++++ 21 files changed, 737 insertions(+), 70 deletions(-) create mode 100644 tests/unit_tests/test_compare_parallel_parity.py diff --git a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml b/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml index ff3bc0d4cf..b1292a433c 100644 --- a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp2.yaml +++ b/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 diff --git a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml b/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml index 3848cdb375..492f537647 100644 --- a/examples/vlm_finetune/gemma4/gemma4_31b_tp4_pp4.yaml +++ b/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 diff --git a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml b/examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml index 5018fb16ba..4ab3dcd53f 100644 --- a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_lora_pp4ep8_8node.yaml +++ b/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 diff --git a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml b/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml index 3d54946f9c..742f31797d 100644 --- a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_cp2_medpix_2k.yaml +++ b/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 diff --git a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml b/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml index 2aa40543a1..f9103d73d4 100644 --- a/examples/vlm_finetune/minimax_m3/minimax_m3_vl_sft_ep32pp4.yaml +++ b/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 diff --git a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml b/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml index 279c2e8d8e..3cea53752b 100644 --- a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix.yaml +++ b/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 diff --git a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml b/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml index 351a6eac58..952cda63a0 100644 --- a/examples/vlm_finetune/mistral3p5/mistral3p5_128b_medpix_lora.yaml +++ b/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 diff --git a/examples/vlm_finetune/mistral4/mistral4_medpix.yaml b/examples/vlm_finetune/mistral4/mistral4_medpix.yaml index 2176dd709c..dfcdb9f26d 100644 --- a/examples/vlm_finetune/mistral4/mistral4_medpix.yaml +++ b/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 diff --git a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml b/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml index 9ace6728c9..6e696eea7c 100644 --- a/examples/vlm_finetune/qwen3_5_moe/qwen3_5_35b_neat_packing.yaml +++ b/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 diff --git a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml index f979814fe2..8b0b4206f3 100644 --- a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_ep32pp4.yaml +++ b/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 diff --git a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml b/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml index af8a2adaaa..bc849ce7eb 100644 --- a/examples/vlm_finetune/stepfun/step3p7_medpix_200b_lora_pp8ep8_8node.yaml +++ b/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 diff --git a/nemo_automodel/components/distributed/pipelining/autopipeline.py b/nemo_automodel/components/distributed/pipelining/autopipeline.py index b88a509123..1ebcd4ef6e 100644 --- a/nemo_automodel/components/distributed/pipelining/autopipeline.py +++ b/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,6 +283,63 @@ 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. """ @@ -290,15 +347,16 @@ def step( 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 diff --git a/nemo_automodel/recipes/base_recipe.py b/nemo_automodel/recipes/base_recipe.py index 69520f9e2e..d3a82afcc1 100644 --- a/nemo_automodel/recipes/base_recipe.py +++ b/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. diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py index ff396c7f8d..c765a58a1b 100644 --- a/nemo_automodel/recipes/llm/train_ft.py +++ b/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. diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index ac0dd0ce79..8187745252 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -602,8 +602,9 @@ def setup(self): ) from nemo_automodel.components.models.common.packing import configure_packing, get_attn_implementation + model_attn_implementation = get_attn_implementation(self.cfg.model, model=self.model_parts[0]) packing_attn_implementation = dataloader_config.resolve_packing_attn_implementation( - model_attn_implementation=get_attn_implementation(self.cfg.model, model=self.model_parts[0]), + model_attn_implementation=model_attn_implementation, cp_size=self.mesh_context.cp_size, ) if dataloader_config.packing is not None and dataloader_config.packing.packing_format != "thd": @@ -629,6 +630,22 @@ def setup(self): self.val_dataloader = None validation_config = self.cfg.vlm_validation_dataloader if validation_config is not None: + if self.pp_enabled and not validation_config.drop_last: + raise ValueError( + "Pipeline-parallel VLM validation requires validation_dataloader.drop_last=true because " + "AutoPipeline uses a fixed outer batch size. Enable drop_last or remove validation_dataset." + ) + _validate_cp_packing_support( + self.model_parts[0], + packing_enabled=validation_config.packing is not None, + cp_size=self.mesh_context.cp_size, + ) + validation_packing_attn_implementation = validation_config.resolve_packing_attn_implementation( + model_attn_implementation=model_attn_implementation, + cp_size=self.mesh_context.cp_size, + ) + if validation_config.packing is not None and validation_config.packing.packing_format != "thd": + configure_packing(attn_implementation=validation_packing_attn_implementation) validation_build_context = FirstRankPerNode(group=process_group) with ScopedRNG(seed=self.cfg.get("seed", 42), ranked=True): validation_build = validation_config.build( @@ -638,6 +655,8 @@ def setup(self): batch_size=self.cfg.get("step_scheduler.local_batch_size", 1), dataset_build_context=validation_build_context, get_rope_index=get_rope_index, + packing_attn_implementation=validation_packing_attn_implementation, + pp_n_microbatches=pp_n_microbatches, cp_size=self.mesh_context.cp_size, ) self.val_dataloader = validation_build.dataloader @@ -701,12 +720,9 @@ def run_train_validation_loop(self): val_loss = {} if self.step_scheduler.is_val_step and self.val_dataloader is not None: - if self.pp_enabled: - logger.warning("Validation is not supported for pipeline parallelism") - else: - val_log_data = self._run_validation_epoch(self.val_dataloader) - val_loss["val_loss"] = val_log_data.metrics["val_loss"] - self.log_val_metrics(val_log_data) + val_log_data = self._run_validation_epoch(self.val_dataloader) + val_loss["val_loss"] = val_log_data.metrics["val_loss"] + self.log_val_metrics(val_log_data) for mp in self.model_parts: mp.train() @@ -915,10 +931,6 @@ def _forward_backward_step( labels = batch.pop("labels") if self.pp_enabled: - if not is_train: - logging.info("Skipping forward pass for validation because pipeline parallelism is enabled") - return - with self._cp_vision_frame_sharding_context(), train_ctx(): losses = [] if self.pp.info.has_last_stage else None if self.pp.info.has_last_stage: @@ -933,7 +945,12 @@ def _forward_backward_step( self._maybe_set_pp_first_stage_embed_input_meta(model_input) with stage_vlm_media_for_pp(self.pp, self.model_parts, batch): - self.pp.step(model_input, target=targets, losses=losses, **batch) + if is_train: + self.pp.step(model_input, target=targets, losses=losses, **batch) + else: + # Forward-only: validation must not enqueue backward work on + # the schedule, which would also desync the other stages. + self.pp.eval(model_input, target=targets, losses=losses, **batch) if self.pp.info.has_last_stage: local_loss = torch.sum(torch.stack(losses)) @@ -1183,7 +1200,20 @@ def _run_train_optim_step(self, batches, max_grad_norm: float | None = None): @torch.no_grad() def _run_validation_epoch(self, val_dataloader): - """Run one pass over `self.val_dataloader`.""" + """Run one pass over `self.val_dataloader`. + + Args: + val_dataloader: Iterable of collated VLM batches. Each batch maps model + input names to tensors; ``labels`` has shape [batch, sequence] and + marks unsupervised positions with ``-100``. + + Returns: + MetricsSample whose ``val_loss`` is the mean loss per supervised label + token over the whole pass. + """ + if self.pp_enabled: + return self._run_pp_validation_epoch(val_dataloader) + with ScopedRNG(seed=1, ranked=True): for mp in self.model_parts: mp.eval() @@ -1255,6 +1285,83 @@ def _run_validation_epoch(self, val_dataloader): }, ) + @torch.no_grad() + def _run_pp_validation_epoch(self, val_dataloader) -> MetricsSample: + """Run one validation pass under pipeline parallelism. + + Every batch goes through the same ``_forward_backward_step`` path as + training -- media staging, stage-shape updates and CP sharding included -- + but with the schedule's forward-only entry point, so only the last stage + produces a loss. + + Args: + val_dataloader: Iterable of collated VLM batches. Each batch maps model + input names to tensors; ``labels`` has shape [batch, sequence] and + marks unsupervised positions with ``-100``. + + Returns: + MetricsSample whose ``val_loss`` is the mean loss per supervised label + token, matching the non-PP path's metric. + + Raises: + ValueError: If no supervised label token survives DP aggregation. + """ + with ScopedRNG(seed=1, ranked=True): + for mp in self.model_parts: + mp.eval() + + total_loss = torch.tensor(0.0, dtype=torch.float32, device=self.dist_env.device) + total_num_label_tokens = 0 + for batch in val_dataloader: + loss_buffer = [] + # Count on the unsharded batch: `_forward_backward_step` may hand CP + # only a slice of the sequence, but the denominator must stay global. + total_num_label_tokens += int((batch["labels"] != -100).sum().item()) + self._forward_backward_step( + 0, + batch, + loss_buffer=loss_buffer, + num_label_tokens=None, # normalized once below, over the whole pass. + num_batches=1, + is_train=False, + ) + total_loss += torch.sum(torch.stack(loss_buffer)) + + total_loss = self._dp_allreduce(total_loss, include_cp=True) + # Every CP rank counted the full sequence above while `total_loss` is + # reassembled from CP-sharded sums, so the token count must not span CP. + total_num_label_tokens = int( + self._dp_allreduce( + torch.tensor(total_num_label_tokens, dtype=torch.long, device=self.dist_env.device) + ).item() + ) + if total_num_label_tokens <= 0: + raise ValueError( + "VLM validation produced no supervised label tokens after DP aggregation. " + "With pipeline parallelism, validation_dataloader.drop_last=true may have removed every batch " + "because each DP shard is smaller than the local batch size; otherwise verify that labels are not " + "all masked." + ) + + # PP loss microbatches are unnormalized sums, so divide once here to get the + # same mean-per-token metric the non-PP path reports. + val_loss = (total_loss / total_num_label_tokens).float().to(self.dist_env.device) + # Only the last stage owns a loss; the rest pass a receive buffer. The token + # count needs no broadcast: every PP rank sees the same batches and reduces + # over the same DP group, so it already agrees. + val_loss = self._broadcast_from_last_pp_stage(val_loss) + + return MetricsSample( + step=self.step_scheduler.step, + epoch=self.step_scheduler.epoch, + metrics={ + "val_loss": val_loss.item(), + "lr": self.optimizer[0].param_groups[0]["lr"], + "num_label_tokens": total_num_label_tokens, + "mem": torch.cuda.max_memory_allocated() / 1024**3, + }, + ) + def log_val_metrics(self, log_data): """Log metrics to wandb and other loggers Args: diff --git a/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh b/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh index 0f7a8ffa12..a7b97269ea 100644 --- a/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh +++ b/tests/functional_tests/parallelism/L2_Parallelism_VLM_Gemma4_PP2_Parity.sh @@ -17,8 +17,9 @@ # # Runs the Gemma4 31B proxy twice with the same seed and data order -- once on a # single rank, once at pp_size=2 -- and asserts both follow the same loss and -# gradient-norm trajectory. `dp_size` is 1 in both runs, so the dataloader yields -# identical batches and any divergence is attributable to the pipeline split. +# gradient-norm trajectory and validation loss. `dp_size` is 1 in both runs, so +# the dataloader yields identical batches and any divergence is attributable to +# the pipeline split. # # Covers the gap from PR #2983 (commit 00f40419). # @@ -55,8 +56,10 @@ COMMON_ARGS=( --validation_dataset.split validation --validation_dataset.limit_dataset_samples 8 --step_scheduler.max_steps 6 + --step_scheduler.val_every_steps 2 --step_scheduler.global_batch_size 4 --step_scheduler.local_batch_size 2 + --validation_dataloader.drop_last true ) # --- Baseline: single rank, no parallelism --- @@ -102,3 +105,12 @@ python tests/functional_tests/parallelism/compare_parallel_parity.py \ --axis pp \ --loss-tol 0.05 \ --grad-norm-rtol 0.20 + +# Both runs must execute recipe-owned validation. The parity helper also rejects +# empty validation logs, so a stale PP validation skip cannot pass this check. +python tests/functional_tests/parallelism/compare_parallel_parity.py \ + "$RUN_DIR/baseline/validation.jsonl" \ + "$RUN_DIR/pp2/validation.jsonl" \ + --axis pp \ + --metric val_loss \ + --loss-tol 0.05 diff --git a/tests/functional_tests/parallelism/compare_parallel_parity.py b/tests/functional_tests/parallelism/compare_parallel_parity.py index cf2f2588a5..1fa7f939be 100644 --- a/tests/functional_tests/parallelism/compare_parallel_parity.py +++ b/tests/functional_tests/parallelism/compare_parallel_parity.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Parallel-vs-single-rank training parity validator. +"""Parallel-vs-single-rank training and validation parity validator. -Compares two ``training.jsonl`` logs produced by the same recipe and seed: a -single-rank baseline and a run with one parallelism axis enabled (TP, PP, CP, or -EP). Both runs must follow the same loss and gradient-norm trajectory. +Compares two metric logs produced by the same recipe and seed: a single-rank +baseline and a run with one parallelism axis enabled (TP, PP, CP, or EP). +Training logs must follow the same loss and gradient-norm trajectory; +validation logs compare their validation loss. This is the generic net for parallelism correctness. A smoke test only fails on a crash or a hang, but wrong stage metadata, a gradient that syncs over the @@ -32,12 +33,14 @@ Usage: python compare_parallel_parity.py baseline.jsonl pp2.jsonl --axis pp + python compare_parallel_parity.py baseline-validation.jsonl pp2-validation.jsonl --axis pp --metric val_loss """ from __future__ import annotations import argparse import json +import math # Both runs share a seed and a data order, so step 1 differs only by floating- # point reduction order. Later steps accumulate that difference through the @@ -54,15 +57,16 @@ MIN_LOSS_SPREAD = 1e-4 -def read_metrics(jsonl_path: str) -> dict[int, dict[str, float]]: - """Read per-step training metrics from a ``training.jsonl`` log. +def read_metrics(jsonl_path: str, *, metric: str = "loss") -> dict[int, dict[str, float]]: + """Read per-step loss metrics from a JSONL log. Args: - jsonl_path: Path to a ``training.jsonl`` written by ``MetricLogger``. + jsonl_path: Path to a metric JSONL written by ``MetricLogger``. + metric: Loss field to read: ``loss`` for training or ``val_loss`` for validation. Returns: - Mapping of step index to a dict with the ``loss`` key and, when the - recipe reported it, ``grad_norm``. + Mapping of step index to a dict whose ``loss`` key holds the requested + metric and, when the recipe reported it, ``grad_norm``. """ entries: dict[int, dict[str, float]] = {} with open(jsonl_path) as f: @@ -71,9 +75,9 @@ def read_metrics(jsonl_path: str) -> dict[int, dict[str, float]]: if not line: continue record = json.loads(line) - if "step" not in record or "loss" not in record: + if "step" not in record or metric not in record: continue - sample: dict[str, float] = {"loss": float(record["loss"])} + sample: dict[str, float] = {"loss": float(record[metric])} grad_norm = record.get("grad_norm") if grad_norm is not None: sample["grad_norm"] = float(grad_norm) @@ -85,6 +89,8 @@ def _assert_both_runs_trained( baseline: dict[int, dict[str, float]], parallel: dict[int, dict[str, float]], common_steps: list[int], + *, + metric: str = "loss", ) -> None: """Reject a comparison where neither run actually trained. @@ -98,10 +104,11 @@ def _assert_both_runs_trained( baseline: Per-step metrics from the baseline run. parallel: Per-step metrics from the parallel run. common_steps: Steps present in both runs. + metric: Loss field the comparison ran on, used to explain a flat curve. Raises: AssertionError: If either run shows a zero gradient norm, or the loss - does not move across steps that saw different batches. + does not move across steps. """ for name, metrics in (("baseline", baseline), ("parallel", parallel)): dead = [step for step in common_steps if metrics[step].get("grad_norm") == 0.0] @@ -111,26 +118,39 @@ def _assert_both_runs_trained( "above proves nothing. Check that the model's weights were initialized." ) - # Needs at least two steps to have anything to compare against, and the - # steps must have seen different batches -- true for every recipe here, - # which draw from a multi-sample dataset without repeating. + # Needs at least two steps to have anything to compare against. Training + # steps each see a different batch -- true for every recipe here, which draw + # from a multi-sample dataset without repeating -- and validation replays one + # dataset against weights that training has moved in between. if len(common_steps) < 2: return losses = [baseline[step]["loss"] for step in common_steps] spread = max(losses) - min(losses) + why_it_should_move = ( + "Each step sees a different batch, so a loss that never moves means the model's output " + "does not depend on its input" + if metric == "loss" + else "Training updates the weights between validation runs, so a validation loss that " + "never moves means validation is not reading the trained model" + ) assert spread > MIN_LOSS_SPREAD, ( - f"The baseline loss is flat across {len(common_steps)} steps (spread {spread:.3e} <= " - f"{MIN_LOSS_SPREAD}). Each step sees a different batch, so a loss that never moves means " - "the model's output does not depend on its input and the comparison above proves nothing." + f"The baseline {metric} is flat across {len(common_steps)} steps (spread {spread:.3e} <= " + f"{MIN_LOSS_SPREAD}). {why_it_should_move} and the comparison above proves nothing." ) def main() -> None: """Compare a single-rank baseline log against a parallel-run log.""" - parser = argparse.ArgumentParser(description="Compare single-rank vs parallel training parity") - parser.add_argument("baseline_jsonl", help="training.jsonl from the single-rank baseline run") - parser.add_argument("parallel_jsonl", help="training.jsonl from the parallel run") + parser = argparse.ArgumentParser(description="Compare single-rank vs parallel training/validation parity") + parser.add_argument("baseline_jsonl", help="Metric JSONL from the single-rank baseline run") + parser.add_argument("parallel_jsonl", help="Metric JSONL from the parallel run") parser.add_argument("--axis", required=True, help="Parallelism axis under test, e.g. pp/tp/cp/ep") + parser.add_argument( + "--metric", + choices=("loss", "val_loss"), + default="loss", + help="Loss field to compare; val_loss performs validation loss-only parity", + ) parser.add_argument("--loss-tol", type=float, default=DEFAULT_LOSS_TOL, help="Absolute per-step loss tolerance") parser.add_argument( "--grad-norm-rtol", @@ -140,11 +160,18 @@ def main() -> None: ) args = parser.parse_args() - baseline = read_metrics(args.baseline_jsonl) - parallel = read_metrics(args.parallel_jsonl) - - assert len(baseline) > 0, f"No training records in {args.baseline_jsonl}" - assert len(parallel) > 0, f"No training records in {args.parallel_jsonl}" + baseline = read_metrics(args.baseline_jsonl, metric=args.metric) + parallel = read_metrics(args.parallel_jsonl, metric=args.metric) + + assert len(baseline) > 0, f"No {args.metric} records in {args.baseline_jsonl}" + assert len(parallel) > 0, f"No {args.metric} records in {args.parallel_jsonl}" + if args.metric == "val_loss": + # Validation runs on a fixed cadence, so a leg that silently skipped it + # would otherwise pass on whatever handful of steps still overlap. + assert set(baseline) == set(parallel), ( + f"Validation steps differ between {args.baseline_jsonl} (steps {sorted(baseline)}) " + f"and {args.parallel_jsonl} (steps {sorted(parallel)})" + ) common_steps = sorted(set(baseline) & set(parallel)) assert len(common_steps) > 0, ( @@ -156,21 +183,35 @@ def main() -> None: grad_norm_failures: list[str] = [] compared_grad_norms = 0 - print(f"=== {args.axis} parity: {len(common_steps)} common steps ===") + print(f"=== {args.axis} {args.metric} parity: {len(common_steps)} common steps ===") print(f"{'step':>6} {'baseline':>12} {'parallel':>12} {'delta':>12}") for step in common_steps: base_loss = baseline[step]["loss"] par_loss = parallel[step]["loss"] - delta = abs(base_loss - par_loss) - print(f"{step:>6} {base_loss:>12.6f} {par_loss:>12.6f} {delta:>12.6f}") - if delta > args.loss_tol: - loss_failures.append(f"step {step}: baseline={base_loss:.6f} {args.axis}={par_loss:.6f} delta={delta:.6f}") + if not math.isfinite(base_loss) or not math.isfinite(par_loss): + # NaN compares unequal to everything, so the delta check below would + # never fire on a run that diverged into non-finite loss. + loss_failures.append( + f"step {step}: non-finite {args.metric}: baseline={base_loss!r} {args.axis}={par_loss!r}" + ) + else: + delta = abs(base_loss - par_loss) + print(f"{step:>6} {base_loss:>12.6f} {par_loss:>12.6f} {delta:>12.6f}") + if delta > args.loss_tol: + loss_failures.append( + f"step {step}: baseline={base_loss:.6f} {args.axis}={par_loss:.6f} delta={delta:.6f}" + ) base_norm = baseline[step].get("grad_norm") par_norm = parallel[step].get("grad_norm") if base_norm is None or par_norm is None: continue compared_grad_norms += 1 + if not math.isfinite(base_norm) or not math.isfinite(par_norm): + grad_norm_failures.append( + f"step {step}: non-finite gradient norm: baseline={base_norm!r} {args.axis}={par_norm!r}" + ) + continue scale = max(abs(base_norm), 1e-8) norm_delta = abs(base_norm - par_norm) / scale if norm_delta > args.grad_norm_rtol: @@ -187,15 +228,19 @@ def main() -> None: f"relative in gradient norm:\n " + "\n ".join(grad_norm_failures) ) - # A log without grad_norm would silently reduce this to a loss-only check. - assert compared_grad_norms > 0, ( - "Neither log reported grad_norm, so the gradient-sync half of this check did not run. " - "Confirm the recipe logs grad_norm to training.jsonl." - ) + if args.metric == "loss": + # A training log without grad_norm would silently reduce this to a + # loss-only check. Validation logs carry no grad_norm by design. + assert compared_grad_norms > 0, ( + "Neither log reported grad_norm, so the gradient-sync half of this check did not run. " + "Confirm the recipe logs grad_norm to training.jsonl." + ) - _assert_both_runs_trained(baseline, parallel, common_steps) + _assert_both_runs_trained(baseline, parallel, common_steps, metric=args.metric) - print(f"{args.axis} parity OK: {len(common_steps)} steps, {compared_grad_norms} gradient norms compared") + print( + f"{args.axis} {args.metric} parity OK: {len(common_steps)} steps, {compared_grad_norms} gradient norms compared" + ) if __name__ == "__main__": diff --git a/tests/unit_tests/distributed/pipelining/test_autopipeline.py b/tests/unit_tests/distributed/pipelining/test_autopipeline.py index 9bc08cce91..953cffdd7d 100644 --- a/tests/unit_tests/distributed/pipelining/test_autopipeline.py +++ b/tests/unit_tests/distributed/pipelining/test_autopipeline.py @@ -243,6 +243,25 @@ def __init__(self, *, fail_on_step: bool = False): self.args_during_step = None self.kwargs_chunk_spec_during_step = None self.kwargs_split = None + self.entry_points: list[str] = [] + + def eval(self, *args, target=None, losses=None, **kwargs): + """Forward-only entry point, which upstream implements by delegating to ``step``. + + Args: + *args: Positional schedule inputs. Tensor values have arbitrary + model-defined layouts. + target: Optional tensor of shape [batch, sequence] containing loss + targets. + losses: Optional mutable list populated with scalar loss tensors. + **kwargs: Keyword schedule inputs. Tensor values have arbitrary + model-defined layouts. + + Returns: + A sentinel string identifying the schedule result. + """ + self.entry_points.append("eval") + return self.step(*args, target=target, losses=losses, **kwargs) def step(self, *args, target=None, losses=None, **kwargs): """Split schedule inputs using the chunk spec active during the call. @@ -260,6 +279,7 @@ def step(self, *args, target=None, losses=None, **kwargs): A sentinel string identifying the schedule result. """ del target, losses + self.entry_points.append("step") self.args_during_step = args self.kwargs_chunk_spec_during_step = self._kwargs_chunk_spec if self.fail_on_step: @@ -361,6 +381,27 @@ def test_model_hook_cannot_configure_unknown_kwarg(self): with pytest.raises(ValueError, match="unknown kwarg"): ap.step(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) + def test_eval_runs_forward_only_with_the_same_chunk_policy(self): + """Validation must chunk model-owned axes exactly as training does.""" + position_ids = torch.arange(8, dtype=torch.long).view(1, 1, -1).expand(3, 2, -1).clone() + ap = self._pipeline_with_parts(_KwargsChunkHookPart({"position_ids": 1})) + + result = ap.eval(torch.zeros(2, 8, dtype=torch.long), position_ids=position_ids) + + assert result == "schedule-result" + assert ap.info.schedule.entry_points[0] == "eval" + assert ap.info.schedule.kwargs_split[0]["position_ids"].shape == (3, 1, 8) + assert ap.info.schedule._kwargs_chunk_spec is None + + def test_eval_without_model_hook_uses_pytorch_default_chunking(self): + ap = self._pipeline_with_parts(nn.Module()) + + ap.eval(torch.zeros(2, 8), attention_mask=torch.ones(2, 8)) + + assert ap.info.schedule.entry_points[0] == "eval" + assert ap.info.schedule.kwargs_chunk_spec_during_step is None + assert ap.info.schedule.kwargs_split[0]["attention_mask"].shape == (1, 8) + # ----------------------------- # Core build/materialize/step tests diff --git a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py index c1a5917201..33e99ac3d1 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_cp_wiring.py @@ -32,6 +32,7 @@ from contextlib import nullcontext from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -425,7 +426,14 @@ class _StageWithoutCPPrepare: pass -def _patch_pp_setup_minimals(monkeypatch, *, cp_size, stage0, dataloader_calls): +def _patch_pp_setup_minimals( + monkeypatch, + *, + cp_size, + stage0, + dataloader_calls, + validation_loader_config=None, +): monkeypatch.setattr(vlm_finetune, "AutoPipeline", _FakePPModel) monkeypatch.setattr( vlm_finetune, @@ -497,7 +505,7 @@ def _build_dataloader(**kwargs): ) monkeypatch.setattr( "nemo_automodel.recipes._typed_config.RecipeConfig.vlm_validation_dataloader", - property(lambda self: None), + property(lambda self: validation_loader_config), ) monkeypatch.setattr(vlm_finetune, "ScopedRNG", lambda **kwargs: nullcontext()) monkeypatch.setattr( @@ -576,6 +584,123 @@ def test_setup_always_stages_pp_media_under_pp( assert dataloader_calls[0]["cp_size"] == cp_size +def test_setup_stages_pp_validation_media_and_preserves_packing_wiring(monkeypatch): + """The validation loader gets the same PP/packing wiring as the training loader. + + Under PP the recipe now runs validation, so the validation batches must be + pre-chunked per microbatch and built against a resolved packing backend -- + otherwise validation hits the raw-media row-chunking the training path fixed. + """ + dataloader_calls = [] + packing_resolutions = [] + configure_packing_calls = [] + + def _resolve_validation_packing(**kwargs): + packing_resolutions.append(kwargs) + return "sdpa" + + validation_loader_config = SimpleNamespace( + drop_last=True, + packing=SimpleNamespace(packing_format="neat"), + resolve_packing_attn_implementation=_resolve_validation_packing, + build=lambda **kwargs: ( + dataloader_calls.append(kwargs) or SimpleNamespace(dataloader="val_dl", processor="processor") + ), + ) + _patch_pp_setup_minimals( + monkeypatch, + cp_size=1, + stage0=_StageWithoutCPPrepare(), + dataloader_calls=dataloader_calls, + validation_loader_config=validation_loader_config, + ) + monkeypatch.setattr( + "nemo_automodel.components.models.common.packing.configure_packing", + lambda **kwargs: configure_packing_calls.append(kwargs), + ) + trainer = FinetuneRecipeForVLM(_minimal_pp_setup_cfg()) + + trainer.setup() + + assert len(dataloader_calls) == 2 + validation_call = dataloader_calls[1] + assert validation_call["pp_n_microbatches"] == 2 + assert validation_call["packing_attn_implementation"] == "sdpa" + assert validation_call["cp_size"] == 1 + assert packing_resolutions == [{"model_attn_implementation": "sdpa", "cp_size": 1}] + assert configure_packing_calls == [{"attn_implementation": "sdpa"}] + assert trainer.val_dataloader == "val_dl" + + +def test_setup_rejects_incomplete_pp_validation_batches(monkeypatch): + """AutoPipeline runs a fixed outer batch, so a short trailing val batch cannot run.""" + dataloader_calls = [] + validation_loader_config = SimpleNamespace( + drop_last=False, + packing=None, + resolve_packing_attn_implementation=lambda **kwargs: None, + build=lambda **kwargs: pytest.fail("validation loader must not build before drop_last validation"), + ) + _patch_pp_setup_minimals( + monkeypatch, + cp_size=1, + stage0=_StageWithoutCPPrepare(), + dataloader_calls=dataloader_calls, + validation_loader_config=validation_loader_config, + ) + trainer = FinetuneRecipeForVLM(_minimal_pp_setup_cfg()) + + with pytest.raises(ValueError, match=r"validation_dataloader\.drop_last=true"): + trainer.setup() + + assert len(dataloader_calls) == 1 + + +def test_train_loop_runs_validation_when_pipeline_is_enabled(): + """The train loop must no longer skip validation under PP.""" + + class _SingleStepScheduler: + epochs = (0,) + step = 1 + epoch = 0 + is_val_step = True + is_ckpt_step = False + sigterm_flag = False + + def set_epoch(self, epoch): + self.epoch = epoch + + def __iter__(self): + yield [object()] + + recipe = object.__new__(FinetuneRecipeForVLM) + model_part = SimpleNamespace(train=MagicMock()) + recipe.model_parts = [model_part] + recipe.step_scheduler = _SingleStepScheduler() + recipe.val_dataloader = object() + recipe.pp_enabled = True + recipe.max_grad_norm = None + recipe._make_progress_bar = MagicMock(return_value=None) + recipe._run_train_optim_step = MagicMock(return_value=SimpleNamespace(metrics={"loss": 1.0})) + recipe.log_train_metrics = MagicMock() + recipe._update_progress_bar = MagicMock() + validation_metrics = SimpleNamespace(metrics={"val_loss": 0.25}) + recipe._run_validation_epoch = MagicMock(return_value=validation_metrics) + recipe.log_val_metrics = MagicMock() + recipe.save_checkpoint = MagicMock() + recipe._maybe_collect_garbage = MagicMock() + recipe.metric_logger_train = SimpleNamespace(close=MagicMock()) + recipe.metric_logger_valid = SimpleNamespace(close=MagicMock()) + recipe._finalize_and_close_checkpointer = MagicMock() + + recipe.run_train_validation_loop() + + recipe._run_validation_epoch.assert_called_once_with(recipe.val_dataloader) + recipe.log_val_metrics.assert_called_once_with(validation_metrics) + # Once before the loop, once after validation put the parts back in train mode. + assert model_part.train.call_count == 2 + + # ----------------------------------------------------------------------------- # val-side wiring (the bug-fix territory) # ----------------------------------------------------------------------------- diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index 3f4130827b..fc6c36d355 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -1540,6 +1540,7 @@ def step_side_effect(*args, **kwargs): kwargs["losses"].append(torch.tensor(0.5)) self.schedule.step = MagicMock(side_effect=step_side_effect) + self.schedule.eval = MagicMock(side_effect=step_side_effect) class _MockAutoPipeline: @@ -1574,6 +1575,25 @@ def step(self, model_input, *, target=None, losses=None, **kwargs): schedule_args = (model_input,) if self.info.has_first_stage else () return self.info.schedule.step(*schedule_args, target=target, losses=losses, **kwargs) + def eval(self, model_input, *, target=None, losses=None, **kwargs): + """Record and forward a forward-only AutoPipeline step. + + Args: + model_input: Tensor of shape [batch, ...] containing the first + pipeline stage's input. + target: Optional tensor of shape [batch, sequence] containing loss + targets. + losses: Optional mutable list populated with scalar loss tensors. + **kwargs: Keyword schedule inputs. Tensor values have arbitrary + model-defined layouts. + + Returns: + The value returned by the schedule mock. + """ + self.step_batches.append(dict(kwargs)) + schedule_args = (model_input,) if self.info.has_first_stage else () + return self.info.schedule.eval(*schedule_args, target=target, losses=losses, **kwargs) + def _create_pp_recipe(model=None): """Helper to create a PP recipe bypassing BaseRecipe tracking.""" @@ -1610,8 +1630,8 @@ def pp_recipe(self): """Create a recipe configured for PP testing.""" return _create_pp_recipe() - def test_pp_skips_validation_forward(self, pp_recipe, monkeypatch): - """Test that PP mode skips forward pass during validation.""" + def test_pp_validation_runs_forward_only(self, pp_recipe, monkeypatch): + """Validation under PP runs the schedule's forward-only entry point.""" pp_recipe.pp = _MockAutoPipeline() monkeypatch.setattr( @@ -1625,7 +1645,6 @@ def test_pp_skips_validation_forward(self, pp_recipe, monkeypatch): } loss_buffer = [] - # Should return early without error pp_recipe._forward_backward_step( idx=0, batch=batch, @@ -1635,8 +1654,39 @@ def test_pp_skips_validation_forward(self, pp_recipe, monkeypatch): is_train=False, # Validation mode ) - # Loss buffer should be empty (no forward pass) - assert len(loss_buffer) == 0 + # Forward-only: eval() drives the schedule, step() must stay untouched so + # validation never enqueues backward work. + pp_recipe.pp.info.schedule.eval.assert_called_once() + pp_recipe.pp.info.schedule.step.assert_not_called() + # Two mock microbatch losses of 0.5 each, summed on the last stage. + assert len(loss_buffer) == 1 + assert loss_buffer[0].item() == pytest.approx(1.0) + + def test_pp_training_runs_step_not_eval(self, pp_recipe, monkeypatch): + """The training path must keep using the backward-capable entry point.""" + pp_recipe.pp = _MockAutoPipeline() + + monkeypatch.setattr( + "nemo_automodel.components.distributed.context_parallel.utils._make_cp_batch_and_ctx", + lambda device_mesh, batch, *a, **k: (lambda: nullcontext(), batch, None), + ) + + batch = { + "labels": torch.tensor([[1, 2]]), + "input_ids": torch.tensor([[1, 2]]), + } + loss_buffer = [] + + pp_recipe._forward_backward_step( + idx=0, + batch=batch, + loss_buffer=loss_buffer, + num_label_tokens=2, + num_batches=1, + ) + + pp_recipe.pp.info.schedule.step.assert_called_once() + pp_recipe.pp.info.schedule.eval.assert_not_called() def test_pp_vlm_chunking_equal_images_and_batch(self, pp_recipe, monkeypatch): """Test VLM pixel_values chunking when n_images == batch_size.""" @@ -2975,6 +3025,102 @@ def _build_checkpointer(**kwargs): assert build_kwargs["pp_group"] is pp_group +def _make_pp_validation_recipe(monkeypatch, *, allreduce_calls, broadcast_calls, step_calls, losses): + """Build a PP recipe stub whose collectives and forward step are recorded. + + Args: + allreduce_calls: List extended with ``(value, include_cp)`` per DP all-reduce. + broadcast_calls: List extended with each scalar tensor handed to the PP broadcast. + step_calls: List extended with the keyword arguments of each forward step. + losses: Per-batch scalar loss values the stubbed forward step reports. + + Returns: + A ``FinetuneRecipeForVLM`` instance wired for ``_run_validation_epoch``. + """ + recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) + recipe.pp_enabled = True + recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) + recipe.model_parts = [SimpleNamespace(eval=MagicMock())] + recipe.step_scheduler = SimpleNamespace(step=4, epoch=0) + recipe.optimizer = [SimpleNamespace(param_groups=[{"lr": 0.001}])] + + remaining = list(losses) + + def _fake_forward_backward_step(idx, batch, **kwargs): + step_calls.append({"idx": idx, **kwargs}) + kwargs["loss_buffer"].append(torch.tensor(remaining.pop(0))) + + def _fake_allreduce(tensor, include_cp=False): + allreduce_calls.append((tensor.item(), include_cp)) + return tensor + + recipe._forward_backward_step = _fake_forward_backward_step + recipe._dp_allreduce = _fake_allreduce + recipe._broadcast_from_last_pp_stage = lambda tensor: broadcast_calls.append(tensor.item()) or tensor + monkeypatch.setattr("nemo_automodel.recipes.vlm.finetune.ScopedRNG", lambda **kwargs: nullcontext()) + return recipe + + +def test_vlm_pp_validation_reports_mean_loss_per_supervised_token(monkeypatch): + """PP microbatch losses are unnormalized sums, so the epoch divides once at the end.""" + allreduce_calls = [] + broadcast_calls = [] + step_calls = [] + recipe = _make_pp_validation_recipe( + monkeypatch, + allreduce_calls=allreduce_calls, + broadcast_calls=broadcast_calls, + step_calls=step_calls, + losses=[3.0, 2.0], + ) + batches = [ + {"labels": torch.tensor([[1, 2, -100, 4]])}, + {"labels": torch.tensor([[-100, 5, 6]])}, + ] + + metrics = recipe._run_validation_epoch(batches) + + # 5.0 summed loss over 3 + 2 supervised tokens. + assert metrics.metrics["val_loss"] == pytest.approx(1.0) + assert metrics.metrics["num_label_tokens"] == 5 + # Every batch is one whole pipeline outer batch, normalized once outside. + assert [(call["idx"], call["num_batches"], call["num_label_tokens"]) for call in step_calls] == [ + (0, 1, None), + (0, 1, None), + ] + assert all(call["is_train"] is False for call in step_calls) + # Loss spans CP because it is reassembled from CP shards; the token count was + # measured pre-shard on every CP rank, so summing it over CP would inflate it. + assert allreduce_calls == [(5.0, True), (5, False)] + assert broadcast_calls == [pytest.approx(1.0)] + + +@pytest.mark.parametrize( + "batches", + [ + [], + [{"labels": torch.tensor([[-100, -100]])}], + ], +) +def test_vlm_pp_validation_rejects_zero_global_denominator(monkeypatch, batches): + """A validation pass with no supervised token must fail loudly, not divide by zero.""" + allreduce_calls = [] + step_calls = [] + recipe = _make_pp_validation_recipe( + monkeypatch, + allreduce_calls=allreduce_calls, + broadcast_calls=[], + step_calls=step_calls, + losses=[0.0], + ) + + with pytest.raises(ValueError, match="no supervised label tokens.*drop_last=true"): + recipe._run_validation_epoch(batches) + + # Both reductions run before the guard, so every rank raises together. + assert len(allreduce_calls) == 2 + + def test_vlm_rope_fusion_disabled_when_cp_gt_1(monkeypatch): """rope_fusion should be set to False during VLM setup when cp_size > 1.""" cfg = _minimal_vlm_cfg(cp_size=2, rope_fusion=True) diff --git a/tests/unit_tests/test_compare_parallel_parity.py b/tests/unit_tests/test_compare_parallel_parity.py new file mode 100644 index 0000000000..53985e6b9b --- /dev/null +++ b/tests/unit_tests/test_compare_parallel_parity.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# 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 json +import sys + +import pytest + +from tests.functional_tests.parallelism import compare_parallel_parity + + +def _write_metrics(path, records): + path.write_text("".join(json.dumps(record) + "\n" for record in records)) + + +def _run_comparison(monkeypatch, baseline_path, parallel_path, *, metric=None): + argv = ["compare_parallel_parity.py", str(baseline_path), str(parallel_path), "--axis", "pp"] + if metric is not None: + argv += ["--metric", metric] + monkeypatch.setattr(sys, "argv", argv) + compare_parallel_parity.main() + + +def _run_validation_comparison(monkeypatch, baseline_path, parallel_path): + _run_comparison(monkeypatch, baseline_path, parallel_path, metric="val_loss") + + +def _run_training_comparison(monkeypatch, baseline_path, parallel_path): + _run_comparison(monkeypatch, baseline_path, parallel_path) + + +def test_validation_parity_accepts_finite_matching_steps(tmp_path, monkeypatch): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + _write_metrics(baseline_path, [{"step": 2, "val_loss": 1.0}, {"step": 4, "val_loss": 0.9}]) + _write_metrics(parallel_path, [{"step": 2, "val_loss": 1.01}, {"step": 4, "val_loss": 0.91}]) + + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + +def test_grad_norm_guard_applies_to_training_only(tmp_path, monkeypatch): + """Validation logs carry no grad_norm, so only ``--metric loss`` may demand one.""" + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + records = [{"step": 2, "loss": 2.0, "val_loss": 1.0}, {"step": 4, "loss": 1.8, "val_loss": 0.9}] + _write_metrics(baseline_path, records) + _write_metrics(parallel_path, records) + + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + with pytest.raises(AssertionError, match="Neither log reported grad_norm"): + _run_training_comparison(monkeypatch, baseline_path, parallel_path) + + +@pytest.mark.parametrize("nonfinite", [float("nan"), float("inf"), float("-inf")]) +def test_validation_parity_rejects_nonfinite_loss(tmp_path, monkeypatch, nonfinite): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + _write_metrics(baseline_path, [{"step": 2, "val_loss": nonfinite}]) + _write_metrics(parallel_path, [{"step": 2, "val_loss": nonfinite}]) + + with pytest.raises(AssertionError, match="non-finite val_loss"): + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + +def test_validation_parity_requires_identical_step_sets(tmp_path, monkeypatch): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + _write_metrics( + baseline_path, + [ + {"step": 2, "val_loss": 1.0}, + {"step": 4, "val_loss": 0.9}, + ], + ) + _write_metrics(parallel_path, [{"step": 2, "val_loss": 1.0}]) + + with pytest.raises(AssertionError, match="Validation steps differ"): + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + +def test_validation_parity_rejects_flat_validation_curve(tmp_path, monkeypatch): + """Two runs that agree on a frozen val_loss prove nothing about validation.""" + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + records = [{"step": 2, "val_loss": 1.0}, {"step": 4, "val_loss": 1.0}] + _write_metrics(baseline_path, records) + _write_metrics(parallel_path, records) + + with pytest.raises(AssertionError, match="baseline val_loss is flat.*not reading the trained model"): + _run_validation_comparison(monkeypatch, baseline_path, parallel_path) + + +def test_training_parity_rejects_nonfinite_gradient_norm(tmp_path, monkeypatch): + baseline_path = tmp_path / "baseline.jsonl" + parallel_path = tmp_path / "parallel.jsonl" + record = {"step": 1, "loss": 1.0, "grad_norm": float("nan")} + _write_metrics(baseline_path, [record]) + _write_metrics(parallel_path, [record]) + + with pytest.raises(AssertionError, match="non-finite gradient norm"): + _run_training_comparison(monkeypatch, baseline_path, parallel_path) From 38d7e031dc20268776c0848a23e1e8988e8281e6 Mon Sep 17 00:00:00 2001 From: HuiyingLi Date: Sun, 30 Aug 2026 23:50:36 -0700 Subject: [PATCH 2/2] refactor(vlm): align PP validation with LLM Keep the pipeline-validation aggregation and broadcast behavior identical to the LLM recipe so the common path can be extracted later. Retain AutoPipeline.eval() as the VLM-specific chunk-spec adapter. Signed-off-by: HuiyingLi --- nemo_automodel/recipes/vlm/finetune.py | 64 +++++++++---------- .../recipes/test_finetune_vlm_helpers.py | 20 +++--- 2 files changed, 40 insertions(+), 44 deletions(-) diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 8187745252..2bdb9e2991 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -1302,9 +1302,6 @@ def _run_pp_validation_epoch(self, val_dataloader) -> MetricsSample: Returns: MetricsSample whose ``val_loss`` is the mean loss per supervised label token, matching the non-PP path's metric. - - Raises: - ValueError: If no supervised label token survives DP aggregation. """ with ScopedRNG(seed=1, ranked=True): for mp in self.model_parts: @@ -1312,54 +1309,51 @@ def _run_pp_validation_epoch(self, val_dataloader) -> MetricsSample: total_loss = torch.tensor(0.0, dtype=torch.float32, device=self.dist_env.device) total_num_label_tokens = 0 + for batch in val_dataloader: loss_buffer = [] - # Count on the unsharded batch: `_forward_backward_step` may hand CP - # only a slice of the sequence, but the denominator must stay global. - total_num_label_tokens += int((batch["labels"] != -100).sum().item()) + num_label_tokens = (batch["labels"] != -100).sum().item() self._forward_backward_step( 0, batch, loss_buffer=loss_buffer, - num_label_tokens=None, # normalized once below, over the whole pass. + num_label_tokens=None, # we will normalize outside. num_batches=1, is_train=False, ) - total_loss += torch.sum(torch.stack(loss_buffer)) + + total_loss += torch.sum(torch.stack(loss_buffer)).item() + total_num_label_tokens += num_label_tokens total_loss = self._dp_allreduce(total_loss, include_cp=True) - # Every CP rank counted the full sequence above while `total_loss` is - # reassembled from CP-sharded sums, so the token count must not span CP. - total_num_label_tokens = int( - self._dp_allreduce( - torch.tensor(total_num_label_tokens, dtype=torch.long, device=self.dist_env.device) - ).item() - ) - if total_num_label_tokens <= 0: - raise ValueError( - "VLM validation produced no supervised label tokens after DP aggregation. " - "With pipeline parallelism, validation_dataloader.drop_last=true may have removed every batch " - "because each DP shard is smaller than the local batch size; otherwise verify that labels are not " - "all masked." - ) + total_num_label_tokens = self._dp_allreduce( + torch.tensor(total_num_label_tokens, dtype=torch.long, device=self.dist_env.device) + ).item() + val_loss = total_loss / max(total_num_label_tokens, 1e-8) - # PP loss microbatches are unnormalized sums, so divide once here to get the - # same mean-per-token metric the non-PP path reports. - val_loss = (total_loss / total_num_label_tokens).float().to(self.dist_env.device) - # Only the last stage owns a loss; the rest pass a receive buffer. The token - # count needs no broadcast: every PP rank sees the same batches and reduces - # over the same DP group, so it already agrees. - val_loss = self._broadcast_from_last_pp_stage(val_loss) + # For PP, send val_loss and num_label_tokens from last stage to main rank + if self.pp_enabled: + val_loss = val_loss.to(self.dist_env.device) + # On non-last ranks total_num_label_tokens is 0; this tensor is just a recv buffer. + pp_num_tokens = torch.tensor(total_num_label_tokens, dtype=torch.long, device=self.dist_env.device) + val_loss = self._broadcast_from_last_pp_stage(val_loss) + pp_num_tokens = self._broadcast_from_last_pp_stage(pp_num_tokens) + if self.dist_env.is_main: + total_num_label_tokens = pp_num_tokens.item() + + val_loss = val_loss.item() if isinstance(val_loss, torch.Tensor) else val_loss + + metrics = { + "val_loss": val_loss, + "lr": self.optimizer[0].param_groups[0]["lr"], + "num_label_tokens": total_num_label_tokens, + "mem": torch.cuda.max_memory_allocated() / 1024**3, + } return MetricsSample( step=self.step_scheduler.step, epoch=self.step_scheduler.epoch, - metrics={ - "val_loss": val_loss.item(), - "lr": self.optimizer[0].param_groups[0]["lr"], - "num_label_tokens": total_num_label_tokens, - "mem": torch.cuda.max_memory_allocated() / 1024**3, - }, + metrics=metrics, ) def log_val_metrics(self, log_data): diff --git a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py index fc6c36d355..0d5dfa7903 100644 --- a/tests/unit_tests/recipes/test_finetune_vlm_helpers.py +++ b/tests/unit_tests/recipes/test_finetune_vlm_helpers.py @@ -3039,7 +3039,7 @@ def _make_pp_validation_recipe(monkeypatch, *, allreduce_calls, broadcast_calls, """ recipe = FinetuneRecipeForVLM.__new__(FinetuneRecipeForVLM) recipe.pp_enabled = True - recipe.dist_env = SimpleNamespace(device=torch.device("cpu")) + recipe.dist_env = SimpleNamespace(device=torch.device("cpu"), is_main=True) recipe.model_parts = [SimpleNamespace(eval=MagicMock())] recipe.step_scheduler = SimpleNamespace(step=4, epoch=0) recipe.optimizer = [SimpleNamespace(param_groups=[{"lr": 0.001}])] @@ -3092,7 +3092,7 @@ def test_vlm_pp_validation_reports_mean_loss_per_supervised_token(monkeypatch): # Loss spans CP because it is reassembled from CP shards; the token count was # measured pre-shard on every CP rank, so summing it over CP would inflate it. assert allreduce_calls == [(5.0, True), (5, False)] - assert broadcast_calls == [pytest.approx(1.0)] + assert broadcast_calls == [pytest.approx(1.0), 5] @pytest.mark.parametrize( @@ -3102,23 +3102,25 @@ def test_vlm_pp_validation_reports_mean_loss_per_supervised_token(monkeypatch): [{"labels": torch.tensor([[-100, -100]])}], ], ) -def test_vlm_pp_validation_rejects_zero_global_denominator(monkeypatch, batches): - """A validation pass with no supervised token must fail loudly, not divide by zero.""" +def test_vlm_pp_validation_handles_zero_global_denominator_like_llm(monkeypatch, batches): + """PP validation follows LLM's finite zero-loss behavior when no labels are supervised.""" allreduce_calls = [] + broadcast_calls = [] step_calls = [] recipe = _make_pp_validation_recipe( monkeypatch, allreduce_calls=allreduce_calls, - broadcast_calls=[], + broadcast_calls=broadcast_calls, step_calls=step_calls, losses=[0.0], ) - with pytest.raises(ValueError, match="no supervised label tokens.*drop_last=true"): - recipe._run_validation_epoch(batches) + metrics = recipe._run_validation_epoch(batches) - # Both reductions run before the guard, so every rank raises together. - assert len(allreduce_calls) == 2 + assert metrics.metrics["val_loss"] == 0.0 + assert metrics.metrics["num_label_tokens"] == 0 + assert allreduce_calls == [(0.0, True), (0, False)] + assert broadcast_calls == [0.0, 0] def test_vlm_rope_fusion_disabled_when_cp_gt_1(monkeypatch):