From 3c6a39139b5cb3e7557d0ea684e21b6e4da95a1d Mon Sep 17 00:00:00 2001 From: Sanket Jayant Purandare Date: Tue, 25 Jun 2024 13:58:59 -0700 Subject: [PATCH 1/2] Adding integration test for FSDP Memory Tracking and Estimation [ghstack-poisoned] --- test_runner.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test_runner.py b/test_runner.py index b26f1947..63377edd 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 From 79ddab246f2e099d3c301bb68515b61a779e352b Mon Sep 17 00:00:00 2001 From: Sanket Jayant Purandare Date: Tue, 25 Jun 2024 14:30:22 -0700 Subject: [PATCH 2/2] Update on "Adding integration test for FSDP Memory Tracking and Estimation" Adds an integration test for `FSDPMemTracker` which will help keep `estimation.py` in sync with `train.py`. `python test_runner.py test_outputs --test fsdp2_mem_tracker` cc: gnadathur [ghstack-poisoned] --- .ci/docker/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/.ci/docker/requirements.txt b/.ci/docker/requirements.txt index 520eb8d4..71669e74 100644 --- a/.ci/docker/requirements.txt +++ b/.ci/docker/requirements.txt @@ -5,3 +5,4 @@ tensorboard sentencepiece tiktoken blobfile +tabulate