Preserve heterogeneous inputs in ZORM micro-batches

Summary: Preserve valid heterogeneous model outputs and required inputs when batching ZORM updates. Jagged values retain variable cardinality, optional all-empty fields keep their semantics, and model-output normalization now runs consistently on synchronous and CPU-offloaded paths.

Differential Revision: D115274546

Author: jeffkbkim

torchrec/metrics/cpu_offloaded_metric_module.py

@@ -90,6 +90,7 @@ def log_event(*args: object, **kwargs: object) -> None:
 # an orphaned GLOO all_gather (a peer rank skipped the collective), an unbounded
 # join hangs the post_train_teardown lease and the whole job is StuckJob-killed.
 _COMPUTE_SHUTDOWN_JOIN_TIMEOUT_SEC: float = 300.0
+_JAGGED_LENGTHS_SUFFIX: str = "_lengths"
 
 
 def _format_thread_stack(thread: threading.Thread) -> str:
@@ -102,8 +103,22 @@ def _format_thread_stack(thread: threading.Thread) -> str:
     return "".join(traceback.format_stack(frame))
 
 
+def _jagged_value_keys(tensors: Mapping[str, torch.Tensor]) -> set[str]:
+    return {
+        key[: -len(_JAGGED_LENGTHS_SUFFIX)]
+        for key in tensors
+        if key.endswith(_JAGGED_LENGTHS_SUFFIX)
+        and key[: -len(_JAGGED_LENGTHS_SUFFIX)] in tensors
+    }
+
+
 def _job_batch_size(model_out: Dict[str, torch.Tensor]) -> int:
-    sizes = [t.shape[0] for t in model_out.values() if t.dim() >= 1]
+    jagged_value_keys = _jagged_value_keys(model_out)
+    sizes = [
+        tensor.shape[0]
+        for key, tensor in model_out.items()
+        if key not in jagged_value_keys and tensor.dim() >= 1
+    ]
     return max(sizes) if sizes else 1
 
 
@@ -131,11 +146,54 @@ def _merge_tensors_across_jobs(
     )
 
 
+def _concatenate_variable_cardinality(
+    tensors: list[torch.Tensor],
+) -> torch.Tensor:
+    rank = tensors[0].dim()
+    trailing_shape = tensors[0].shape[1:]
+    if any(
+        tensor.dim() != rank or tensor.shape[1:] != trailing_shape
+        for tensor in tensors[1:]
+    ):
+        raise RecMetricException(
+            "cannot concatenate tensors with shapes "
+            f"{[tuple(tensor.shape) for tensor in tensors]}: dimensions after "
+            "dim 0 must match"
+        )
+    return torch.cat(tensors, dim=0)
+
+
+def _merge_tensor_mappings(
+    mappings: list[Mapping[str, torch.Tensor]],
+    batch_sizes: list[int],
+    mapping_name: str,
+) -> Dict[str, torch.Tensor]:
+    first = mappings[0]
+    jagged_value_keys = _jagged_value_keys(first)
+    merged: Dict[str, torch.Tensor] = {}
+    for key in first:
+        tensors = [mapping[key] for mapping in mappings]
+        try:
+            if key in jagged_value_keys or all(
+                tensor.numel() == 0 for tensor in tensors
+            ):
+                merged[key] = _concatenate_variable_cardinality(tensors)
+            else:
+                merged[key] = _merge_tensors_across_jobs(tensors, batch_sizes)
+        except (RecMetricException, RuntimeError) as error:
+            raise RecMetricException(
+                f"failed to merge {mapping_name} key {key!r}: {error}"
+            ) from error
+    return merged
+
+
 def _merge_update_jobs(jobs: list[MetricUpdateJob]) -> MetricUpdateJob:
     """
     Merge a list of MetricUpdateJobs into a single job by concatenating
     per-key tensors in `model_out` and any tensor entries in
     `kwargs['required_inputs']` along the batch dimension (dim=0).
+    Values with a paired `<key>_lengths` tensor retain their jagged
+    cardinality, and fields that are empty in every job remain empty.
 
     Worker-side input batching: K mini-batches' worth of model outputs
     are passed as a single (concatenated) job, cutting the DtoH transfer,
@@ -152,20 +210,18 @@ def _merge_update_jobs(jobs: list[MetricUpdateJob]) -> MetricUpdateJob:
 
     first = jobs[0]
     batch_sizes = [_job_batch_size(j.model_out) for j in jobs]
-    merged_model_out: Dict[str, torch.Tensor] = {
-        key: _merge_tensors_across_jobs([j.model_out[key] for j in jobs], batch_sizes)
-        for key in first.model_out.keys()
-    }
+    merged_model_out = _merge_tensor_mappings(
+        [job.model_out for job in jobs], batch_sizes, "model_out"
+    )
 
     merged_kwargs: Dict[str, Any] = dict(first.kwargs)
     required_inputs_first = first.kwargs.get("required_inputs")
     if isinstance(required_inputs_first, dict) and required_inputs_first:
-        merged_required: Dict[str, torch.Tensor] = {
-            key: _merge_tensors_across_jobs(
-                [j.kwargs["required_inputs"][key] for j in jobs], batch_sizes
-            )
-            for key in required_inputs_first.keys()
-        }
+        merged_required = _merge_tensor_mappings(
+            [job.kwargs["required_inputs"] for job in jobs],
+            batch_sizes,
+            "required_inputs",
+        )
         merged_kwargs["required_inputs"] = merged_required
 
     return MetricUpdateJob(
@@ -517,6 +573,7 @@ def _process_metric_update_job(self, metric_update_job: MetricUpdateJob) -> None
             if transfer_completed_event is not None:
                 transfer_completed_event.synchronize()
 
+            cpu_model_out = self._prepare_model_out_for_metrics(cpu_model_out)
             labels, predictions, weights, required_inputs = parse_task_model_outputs(
                 self.rec_tasks,
                 cpu_model_out,

torchrec/metrics/metric_module.py

@@ -315,6 +315,17 @@ def load_state_dict_hook(
                 f"Removed key '{key}' from state_dict for backward compatibility"
             )
 
+    def _prepare_model_out_for_metrics(
+        self, model_out: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        """Prepare model outputs before task and required-input parsing.
+
+        Subclasses can override this hook when their model-output format needs
+        normalization. CPU-offloaded modules run it on the worker after transfer
+        to CPU, keeping preparation off the trainer critical path.
+        """
+        return model_out
+
     def _update_rec_metrics(
         self, model_out: Dict[str, torch.Tensor], **kwargs: Any
     ) -> None:
@@ -324,6 +335,7 @@ def _update_rec_metrics(
         """
         # pyrefly: ignore[not-callable]
         if self.rec_metrics and self.rec_tasks:
+            model_out = self._prepare_model_out_for_metrics(model_out)
             labels, predictions, weights, required_inputs = parse_task_model_outputs(
                 self.rec_tasks, model_out, self.get_required_inputs()
             )

torchrec/metrics/tests/test_cpu_offloaded_metric_module.py

@@ -1921,7 +1921,147 @@ def test_unalignable_required_input_raises(self) -> None:
                 kwargs={"required_inputs": {"target_tensor": torch.tensor([7.0, 8.0])}},
             ),
         ]
-        with self.assertRaises(RecMetricException):
+        with self.assertRaisesRegex(
+            RecMetricException,
+            "failed to merge required_inputs key 'target_tensor'",
+        ):
+            _merge_update_jobs(jobs)
+
+    def test_all_empty_unused_model_output_is_preserved(self) -> None:
+        jobs = [
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.1, 0.2]),
+                    "optional_teacher_label": torch.empty(0),
+                },
+                kwargs={},
+            ),
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.3, 0.4]),
+                    "optional_teacher_label": torch.empty(0),
+                },
+                kwargs={},
+            ),
+        ]
+
+        merged = _merge_update_jobs(jobs)
+
+        self.assertEqual(merged.model_out["prediction"].shape, (4,))
+        self.assertEqual(merged.model_out["optional_teacher_label"].shape, (0,))
+
+    def test_all_empty_model_outputs_require_matching_trailing_shapes(self) -> None:
+        jobs = [
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.1, 0.2]),
+                    "optional_teacher_label": torch.empty(0, 5),
+                },
+                kwargs={},
+            ),
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.3, 0.4]),
+                    "optional_teacher_label": torch.empty(0, 3),
+                },
+                kwargs={},
+            ),
+        ]
+
+        with self.assertRaisesRegex(
+            RecMetricException,
+            "failed to merge model_out key 'optional_teacher_label'.*"
+            "dimensions after dim 0 must match",
+        ):
+            _merge_update_jobs(jobs)
+
+    def test_jagged_values_and_lengths_preserve_variable_cardinality(self) -> None:
+        jobs = [
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.1, 0.2]),
+                    "campaign_id": torch.empty(0, dtype=torch.long),
+                    "campaign_id_lengths": torch.tensor([0, 0]),
+                },
+                kwargs={},
+            ),
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.3, 0.4]),
+                    "campaign_id": torch.tensor([7, 7, 8]),
+                    "campaign_id_lengths": torch.tensor([2, 1]),
+                },
+                kwargs={},
+            ),
+        ]
+
+        merged = _merge_update_jobs(jobs)
+
+        torch.testing.assert_close(
+            merged.model_out["prediction"], torch.tensor([0.1, 0.2, 0.3, 0.4])
+        )
+        torch.testing.assert_close(
+            merged.model_out["campaign_id"], torch.tensor([7, 7, 8])
+        )
+        torch.testing.assert_close(
+            merged.model_out["campaign_id_lengths"], torch.tensor([0, 0, 2, 1])
+        )
+
+    def test_jagged_required_inputs_preserve_variable_cardinality(self) -> None:
+        jobs = [
+            MetricUpdateJob(
+                model_out={"prediction": torch.tensor([0.1, 0.2])},
+                kwargs={
+                    "required_inputs": {
+                        "session_id": torch.empty(0, dtype=torch.long),
+                        "session_id_lengths": torch.tensor([0, 0]),
+                    }
+                },
+            ),
+            MetricUpdateJob(
+                model_out={"prediction": torch.tensor([0.3, 0.4])},
+                kwargs={
+                    "required_inputs": {
+                        "session_id": torch.tensor([9, 9, 10]),
+                        "session_id_lengths": torch.tensor([2, 1]),
+                    }
+                },
+            ),
+        ]
+
+        merged = _merge_update_jobs(jobs)
+
+        torch.testing.assert_close(
+            merged.kwargs["required_inputs"]["session_id"],
+            torch.tensor([9, 9, 10]),
+        )
+        torch.testing.assert_close(
+            merged.kwargs["required_inputs"]["session_id_lengths"],
+            torch.tensor([0, 0, 2, 1]),
+        )
+
+    def test_unpaired_mixed_empty_tensor_remains_invalid(self) -> None:
+        jobs = [
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.1, 0.2]),
+                    "ambiguous": torch.empty(0),
+                },
+                kwargs={},
+            ),
+            MetricUpdateJob(
+                model_out={
+                    "prediction": torch.tensor([0.3, 0.4]),
+                    "ambiguous": torch.tensor([1.0]),
+                },
+                kwargs={},
+            ),
+        ]
+
+        with self.assertRaisesRegex(
+            RecMetricException,
+            "failed to merge model_out key 'ambiguous'",
+        ):
             _merge_update_jobs(jobs)
 
     def test_scalar_model_out_expands_and_concats(self) -> None:

torchrec/metrics/tests/test_metric_module.py

@@ -15,7 +15,7 @@
 import tempfile
 import unittest
 from collections.abc import Mapping
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, cast, Dict, List, Optional
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -101,6 +101,15 @@ def _update_rec_metrics(
         self.rec_metrics.update(predictions=predictions, labels=labels, weights=weights)
 
 
+class PreparingMetricModule(RecMetricModule):
+    def _prepare_model_out_for_metrics(
+        self, model_out: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        prepared_model_out = dict(model_out)
+        prepared_model_out["task-prediction"] = prepared_model_out["raw-prediction"]
+        return prepared_model_out
+
+
 class MetricModuleTest(unittest.TestCase):
     @seed_and_log
     def setUp(self) -> None:
@@ -189,6 +198,41 @@ def test_metric_module(self) -> None:
                 self.assertTrue("optimizers-optimizers|learning_rate" in ret)
             dist.destroy_process_group()
 
+    def test_prepares_model_output_before_parsing(self) -> None:
+        tasks = gen_test_tasks(["task"])
+        rec_metric = MockRecMetric(
+            world_size=1,
+            my_rank=0,
+            batch_size=2,
+            tasks=tasks,
+        )
+        metric_module = PreparingMetricModule(
+            batch_size=2,
+            world_size=1,
+            rec_tasks=tasks,
+            rec_metrics=RecMetricList([rec_metric]),
+        )
+        raw_predictions = torch.tensor([0.25, 0.75])
+        model_out = gen_test_batch(
+            batch_size=2,
+            label_name="task-label",
+            prediction_name="unused-prediction",
+            weight_name="task-weight",
+        )
+        model_out["raw-prediction"] = raw_predictions
+
+        metric_module.update(model_out)
+
+        self.assertEqual(1, rec_metric.update_called_count)
+        actual_predictions = rec_metric.predictions_update_calls[0]
+        self.assertIsInstance(actual_predictions, dict)
+        self.assertTrue(
+            torch.equal(
+                raw_predictions,
+                cast(Dict[str, torch.Tensor], actual_predictions)["task"],
+            )
+        )
+
     def test_rectask_info(self) -> None:
         mock_optimizer = MockOptimizer()
         config = DefaultMetricsConfig