Automate benchmarking - changed the scripts for backward pass. (#759)

## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
The scripts are updated so that we get benchmarking results on backward
pass as well.
All benchmark scripts were run on a H100, it took around 55 minutes to
run.

There are total of 26 scripts, so each generates a plot for speed -
forward, backward and full pass
for memory - full pass. Total of 104.
And the output (csv file) was of 460kB.
<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Manan Shah <[email protected]>
Co-authored-by: Vaibhav Jindal <[email protected]>

Author: 3 people

benchmark/scripts/benchmark_cpo_loss.py

@@ -147,7 +147,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_cpo_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

benchmark/scripts/benchmark_cross_entropy.py

@@ -109,7 +109,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_cross_entropy,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

benchmark/scripts/benchmark_distill_jsd_loss.py

@@ -248,7 +248,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_jsd_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

benchmark/scripts/benchmark_dpo_loss.py

@@ -166,7 +166,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_dpo_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

benchmark/scripts/benchmark_embedding.py

@@ -112,7 +112,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_embedding,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

benchmark/scripts/benchmark_fused_linear_cross_entropy.py

@@ -156,7 +156,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_cross_entropy,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

benchmark/scripts/benchmark_fused_linear_jsd.py

@@ -246,7 +246,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_fused_linear_jsd,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,

benchmark/scripts/benchmark_jsd.py

@@ -143,7 +143,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_jsd,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_args,

benchmark/scripts/benchmark_kl_div.py

@@ -110,7 +110,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_kldiv,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_args,

benchmark/scripts/benchmark_kto_loss.py

@@ -299,7 +299,7 @@ def full():
 
     run_benchmarks(
         bench_test_fn=bench_speed_kto_loss,
-        kernel_operation_modes=["forward", "full"],
+        kernel_operation_modes=["forward", "backward", "full"],
         metric_name="speed",
         metric_unit="ms",
         **common_configs,