From 95dd0ef4a98edb540797e119c620796208b32842 Mon Sep 17 00:00:00 2001
From: Sanket Jayant Purandare <sanketpurandare@meta.com>
Date: Tue, 25 Jun 2024 13:59:03 -0700
Subject: [PATCH] Adding integration test for FSDP Memory Tracking and
 Estimation

ghstack-source-id: 9d25d1fa8303db03dd036cfa115bef1374c4853f
Pull Request resolved: https://github.com/pytorch/torchtitan/pull/426
---
 test_runner.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/test_runner.py b/test_runner.py
index b26f19478..63377edd0 100755
--- a/test_runner.py
+++ b/test_runner.py
@@ -263,6 +263,14 @@ def build_test_list():
             ],
             "Fused Optimizer Test",
         ),
+        OverrideDefinitions(
+            [
+                ["--estimate.memory=True", "--estimate.mode=real"],
+            ],
+            "FSDP2 Memory Tracking and Estimation",
+            "fsdp2_mem_tracker",
+            ngpu=4,
+        ),
     ]
     return integration_tests_flavors