Fix benchmark pipeline runner (#3236)

Summary:
Pull Request resolved: #3236

Changed the return type of benchmark pipeline runner from `List[BenchmarkResult]` to single `BenchmarkResult` since runner should return only one result for the given configuration.
Remove with/without`apply_jit` comparison as both configurations should call pipeline benchmarking separately.

Reviewed By: jd7-tr

Differential Revision: D78941101

fbshipit-source-id: cf25de1132d6d4a8365b279d83e4b13de1cb1410

Author: SSYernar

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -356,48 +356,32 @@ def _func_to_benchmark(
                 except StopIteration:
                     break
 
-        # Run comparison if apply_jit is True, otherwise run single benchmark
-        jit_configs = (
-            [(True, "WithJIT"), (False, "WithoutJIT")]
-            if pipeline_config.apply_jit
-            else [(False, "")]
+        pipeline = generate_pipeline(
+            pipeline_type=pipeline_config.pipeline,
+            emb_lookup_stream=pipeline_config.emb_lookup_stream,
+            model=sharded_model,
+            opt=optimizer,
+            device=ctx.device,
+            apply_jit=pipeline_config.apply_jit,
+        )
+        pipeline.progress(iter(bench_inputs))
+
+        result = benchmark_func(
+            name=type(pipeline).__name__,
+            bench_inputs=bench_inputs,  # pyre-ignore
+            prof_inputs=bench_inputs,  # pyre-ignore
+            num_benchmarks=5,
+            num_profiles=2,
+            profile_dir=run_option.profile,
+            world_size=run_option.world_size,
+            func_to_benchmark=_func_to_benchmark,
+            benchmark_func_kwargs={"model": sharded_model, "pipeline": pipeline},
+            rank=rank,
+            export_stacks=run_option.export_stacks,
         )
-        results = []
-
-        for apply_jit, jit_suffix in jit_configs:
-            pipeline = generate_pipeline(
-                pipeline_type=pipeline_config.pipeline,
-                emb_lookup_stream=pipeline_config.emb_lookup_stream,
-                model=sharded_model,
-                opt=optimizer,
-                device=ctx.device,
-                apply_jit=apply_jit,
-            )
-            pipeline.progress(iter(bench_inputs))
-
-            name = (
-                f"{type(pipeline).__name__}{jit_suffix}"
-                if jit_suffix
-                else type(pipeline).__name__
-            )
-            result = benchmark_func(
-                name=name,
-                bench_inputs=bench_inputs,  # pyre-ignore
-                prof_inputs=bench_inputs,  # pyre-ignore
-                num_benchmarks=5,
-                num_profiles=2,
-                profile_dir=run_option.profile,
-                world_size=run_option.world_size,
-                func_to_benchmark=_func_to_benchmark,
-                benchmark_func_kwargs={"model": sharded_model, "pipeline": pipeline},
-                rank=rank,
-                export_stacks=run_option.export_stacks,
-            )
-            results.append(result)
 
         if rank == 0:
-            for result in results:
-                print(result)
+            print(result)
 
 
 if __name__ == "__main__":