try out fsdp sharding

EleutherAI · May 2, 2023 · 54bee5d · 54bee5d
1 parent ffdef0e
commit 54bee5d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/llama_device_map.py b/llama_device_map.py
@@ -117,9 +117,9 @@ def main(args):
     max_memory = (
         {0: forty_gb, 1: forty_gb}
         if not use_8bit
-        # this is a hack since infer_auto_device_map doesn't detect 8bit
-        # even if we load it in 8bit
-        # for big models, it'll start allocating to disk
+        # this is a hack since infer_auto_device_map can't detect
+        # that we're using 8bit, since we inited an empty model
+        # to analyse.
         else {0: forty_gb * 2, 1: forty_gb * 2}
     )
     autodevice_map = infer_auto_device_map(

diff --git a/llama_fsdp.py b/llama_fsdp.py
@@ -50,7 +50,7 @@ def run_inference(
         sync_module_states=False,
         limit_all_gathers=False,
         forward_prefetch=True,
-        strategy=strategy,
+        sharding_strategy=strategy,
     )
 
     if rank == 0: