update

mlfoundations · Dec 15, 2023 · 6196cbc · 6196cbc
1 parent 39124e7
commit 6196cbc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -140,3 +140,4 @@ out*
 tests/assets/*
 .vscode/
 checkpoints/
+experiments/
diff --git a/open_lm/main.py b/open_lm/main.py
@@ -27,7 +27,6 @@
     StateDictType,
     CPUOffload,
 )
-
 from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
 
 from open_lm.data import proc_token
@@ -153,6 +152,7 @@ def load_data_chunks(args):
         )
         return [0 for _ in range(len(args.dataset_manifest))], 0
 
+
 def save_checkpoint(
     args,
     model,
@@ -473,6 +473,7 @@ def main(args):
             if args.ddp_static_graph:
                 # this doesn't exist in older PyTorch, arg only added if enabled
                 ddp_args["static_graph"] = True
+            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device], **ddp_args)
 
     if args.grad_checkpointing:
         model.set_grad_checkpointing()