Skip to content

Commit 841bb53

Browse files
Adding integration test for FSDP Memory Tracking and Estimation
ghstack-source-id: cc224db Pull Request resolved: #426
1 parent 236d2ff commit 841bb53

File tree

2 files changed

+9
-0
lines changed

2 files changed

+9
-0
lines changed

.ci/docker/requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,3 +5,4 @@ tensorboard
55
sentencepiece
66
tiktoken
77
blobfile
8+
tabulate

test_runner.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,14 @@ def build_test_list():
263263
],
264264
"Fused Optimizer Test",
265265
),
266+
OverrideDefinitions(
267+
[
268+
["--estimate.memory=True", "--estimate.mode=real"],
269+
],
270+
"FSDP2 Memory Tracking and Estimation",
271+
"fsdp2_mem_tracker",
272+
ngpu=4,
273+
),
266274
]
267275
return integration_tests_flavors
268276

0 commit comments

Comments
 (0)