Merge pull request #3 from sanketpurandare/run_est

Merge Andrew's updates

.gitattributes

@@ -5,3 +5,11 @@
 .github/scripts/gql_mocks.json linguist-generated=true
 third_party/LICENSES_BUNDLED.txt linguist-generated=true
 tools/build/bazel/requirements.txt linguist-generated=true
+torch/distributed/_tools/a100_models/bmm.joblib filter=lfs diff=lfs merge=lfs -text
+torch/distributed/_tools/a100_models/mm.joblib filter=lfs diff=lfs merge=lfs -text
+torch/distributed/_tools/a100_models/sdpa.joblib filter=lfs diff=lfs merge=lfs -text
+torch/distributed/_tools/a100_models/sdpa_backward.joblib filter=lfs diff=lfs merge=lfs -text
+torch/distributed/_tools/h100_models/sdpa.joblib filter=lfs diff=lfs merge=lfs -text
+torch/distributed/_tools/h100_models/sdpa_backward.joblib filter=lfs diff=lfs merge=lfs -text
+torch/distributed/_tools/h100_models/bmm.joblib filter=lfs diff=lfs merge=lfs -text
+torch/distributed/_tools/h100_models/mm.joblib filter=lfs diff=lfs merge=lfs -text

test/distributed/_tools/test_runtime_estimator.py

@@ -133,6 +133,7 @@ def _init_model_and_args(
     def test_transformer_runtime(
         self,
     ):
+        print("Transformer Test")
         """Runs a basic GPT-2 model"""
         vocab_size = 8192
         bsz, seq_len = 8, 1024
@@ -155,22 +156,30 @@ def test_transformer_runtime(
             roofline_estimate = self._runtime_estimate(
                 "operator-level-cost-model", self._train_step, fake_args
             )
+            learned_estimate = self._runtime_estimate(
+                "operator-level-learned-model", self._train_step, fake_args
+            )
         benchmark_accuracy = actual_runtime / benchmark_estimate
         roofline_accuracy = actual_runtime / roofline_estimate
+        learned_accuracy = actual_runtime / learned_estimate
         print(
             f"Actual: {actual_runtime} Benchmark Estimate: {benchmark_estimate} Accuracy: {benchmark_accuracy}"
-            f"\n Actual: {actual_runtime} Roofline Estimatee: {roofline_estimate} Accuracy: {roofline_accuracy}"
+            f"\nActual: {actual_runtime} Roofline Estimate: {roofline_estimate} Accuracy: {roofline_accuracy}"
+            f"\nActual: {actual_runtime} Learned Estimate: {learned_estimate} Accuracy: {learned_accuracy}"
         )
+
         # No accuracy check for benchmark in CI as it is highly variable
         # self.assertAlmostEqual(benchmark_accuracy, 1.0, delta=0.2)
         # self.assertAlmostEqual(roofline_accuracy, 1.0, delta=0.3)
 
+
     @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/115653")
     @unittest.skipIf(not TEST_CUDA, "CUDA not available")
     def test_conv_model_runtime(
         self,
     ):
         """Runs a simple CNN model"""
+        print("CNN Test")
         num_classes = 100
         bsz, img_sz = 256, 128
         model_args = ConvArgs(img_sz, num_classes)
@@ -184,11 +193,16 @@ def test_conv_model_runtime(
             roofline_estimate = self._runtime_estimate(
                 "operator-level-cost-model", self._train_step, fake_args
             )
+            learned_estimate = self._runtime_estimate(
+                "operator-level-learned-model", self._train_step, fake_args
+            )
         benchmark_accuracy = actual_runtime / benchmark_estimate
         roofline_accuracy = actual_runtime / roofline_estimate
+        learned_accuracy = actual_runtime / learned_estimate
         print(
-            f"Actual: {actual_runtime} Benchmark Estimate: {benchmark_estimate} Accuracy: {benchmark_accuracy}\n"
-            f"Actual: {actual_runtime} Roofline Estimatee: {roofline_estimate} Accuracy: {roofline_accuracy}"
+            f"Actual: {actual_runtime} Benchmark Estimate: {benchmark_estimate} Accuracy: {benchmark_accuracy}"
+            f"\nActual: {actual_runtime} Roofline Estimate: {roofline_estimate} Accuracy: {roofline_accuracy}"
+            f"\nActual: {actual_runtime} Learned Estimate: {learned_estimate} Accuracy: {learned_accuracy}"
         )
         # No accuracy check for benchmark in CI as it is highly variable
         # self.assertAlmostEqual(benchmark_accuracy, 1.0, delta=0.2)