Inference mode configuration support for benchmarking

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -91,6 +91,9 @@ class RunOptions:
             Default is "EXACT_ADAGRAD".
         sparse_lr (float): Learning rate for sparse parameters.
             Default is 0.1.
+        training_mode (bool): Whether to run the model in training mode.
+            If True, model remains in training mode. If False, model.eval() is called.
+            Default is True.
     """
 
     world_size: int = 2
@@ -110,6 +113,7 @@ class RunOptions:
     sparse_lr: float = 0.1
     sparse_momentum: Optional[float] = None
     sparse_weight_decay: Optional[float] = None
+    training_mode: bool = True
 
 
 @dataclass
@@ -343,6 +347,9 @@ def runner(
             planner=planner,
         )
 
+        if not run_option.training_mode:
+            sharded_model.eval()
+
         def _func_to_benchmark(
             bench_inputs: List[ModelInput],
             model: nn.Module,

torchrec/distributed/benchmark/benchmark_utils.py

@@ -524,6 +524,9 @@ def wrapper() -> Any:
                 if origin in (list, List):
                     elem_type = get_args(ftype)[0]
                     arg_kwargs.update(nargs="*", type=elem_type)
+                elif ftype is bool:
+                    # Special handling for boolean arguments
+                    arg_kwargs.update(type=lambda x: x.lower() in ["true", "1", "yes"])
                 else:
                     arg_kwargs.update(type=ftype)
 

torchrec/distributed/train_pipeline/train_pipelines.py

@@ -95,6 +95,20 @@
     has_2d_support = False
 
 
+# Returns (losses, output) from model forward pass. Losses is None if model is in eval() mode
+# pyre-ignore[3]
+def unpack_model_fwd(
+    model_fwd_fn: Callable[[Any], Any], batch: Any, training: bool  # pyre-ignore[2]
+) -> Tuple[torch.Tensor, Any]:
+    result = model_fwd_fn(batch)
+    if training:
+        # result expected to be (losses, output)
+        return result
+    else:
+        # result expected to be output only
+        return None, result  # pyre-ignore[7]
+
+
 class ModelDetachedException(Exception):
     pass
 
@@ -667,7 +681,9 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
 
         # forward
         with record_function("## forward ##"):
-            losses, output = self._model_fwd(self.batches[0])
+            losses, output = unpack_model_fwd(
+                self._model_fwd, self.batches[0], self._model.training
+            )
 
         if self._enqueue_batch_after_forward:
             # batch i+2: load data and copy to gpu, the dataload iter will first exhaust here.
@@ -1030,7 +1046,9 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
 
         # forward
         with record_function("## forward ##"):
-            losses, output = self._model_fwd(self.batches[0])
+            losses, output = unpack_model_fwd(
+                self._model_fwd, self.batches[0], self._model.training
+            )
 
         if len(self.batches) >= 2:
             # invoke data (values, lengths, etc.) all_to_all comms (second part of input_dist)
@@ -1254,7 +1272,10 @@ def _mlp_forward(
             _wait_for_events(
                 batch, context, torch.get_device_module(self._device).current_stream()
             )
-            return self._model_fwd(batch)
+            losses, output = unpack_model_fwd(
+                self._model_fwd, batch, self._model.training
+            )
+            return losses, output
 
     def embedding_backward(self, context: EmbeddingTrainPipelineContext) -> None:
         assert len(context.embedding_features) == len(context.embedding_tensors)
@@ -1466,7 +1487,9 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
         self._wait_sparse_data_dist()
         # forward
         with record_function("## forward ##"):
-            losses, output = self._model_fwd(self._batch_i)
+            losses, output = unpack_model_fwd(
+                self._model_fwd, self._batch_i, self._model.training
+            )
 
         self._prefetch(self._batch_ip1)
 
@@ -2024,7 +2047,9 @@ def progress(self, dataloader_iter: Iterator[In]) -> Out:
         # forward
         ctx = self.get_compiled_autograd_ctx()
         with ctx, torchrec_use_sync_collectives(), record_function("## forward ##"):
-            losses, output = self._model_fwd(self.batches[0])
+            losses, output = unpack_model_fwd(
+                self._model_fwd, self.batches[0], self._model.training
+            )
 
         if len(self.batches) >= 2:
             self.wait_sparse_data_dist(self.contexts[1])