[benchmarks] Add benchmarks for FSDP (#765)

* add benchmarks for fsdp * fix lint errors * clean up * clean up unused flags * add the benchmarks * remove unused args * fix lint errors * fix lint errors * update command line * add support for multiple devices * try full fp16 mode * try full fp16 mode * lint errors * merge main * lint errors * lint errors * lint error * update intersphinx mapping for numpy * update intersphinx mapping for numpy * skip test * added golden configs * use synthetic benchmarks * fix fn name * fix cuda device id * fix verify * lint fix
facebookresearch · Feb 22, 2022 · f9a125d · f9a125d
1 parent 9090bfd
commit f9a125d
Show file tree

Hide file tree

Showing 4 changed files with 457 additions and 0 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -150,6 +150,12 @@ run_offload_benchmark: &run_offload_benchmark
       command: |
         python benchmarks/experimental/offload.py --checkpoint_activation
 
+run_fsdp_benchmark: &run_fsdp_benchmark
+  - run:
+      name: Run FSDP Benchmark
+      command: |
+        python benchmarks/fsdp.py --use_synthetic_data
+
 run_pipe_benchmark: &run_pipe_benchmark
   - run:
       name: Run Pipe Benchmark
@@ -489,6 +495,8 @@ jobs:
 
       - <<: *run_oss_gloo
 
+      - <<: *run_fsdp_benchmark
+
       - save_cache:
           paths:
             - /tmp/MNIST