Only use meta device if explicitly requested

open_lm/main.py

@@ -249,10 +249,13 @@ def check_args(args):
                 raise ValueError("Sync protocol not supported when using resume latest.")
 
     if args.target_mask_left is not None and args.target_mask_individual == args.target_mask_left:
-        ValueError(f"--target-mask-left and --target-mask-individual set to same value of {args.target_mask_left}.")
+        raise ValueError(f"--target-mask-left and --target-mask-individual set to same value of {args.target_mask_left}.")
 
     if args.lr_scheduler != "cosine":
-        ValueError(f"Unknown scheduler, {args.lr_scheduler}. Available options are: cosine, const, const-cooldown.")
+        raise ValueError(f"Unknown scheduler, {args.lr_scheduler}. Available options are: cosine, const, const-cooldown.")
+
+    if args.init_meta_device and not args.fsdp:
+        raise ValueError("--init-meta-device can only be specified if --fsdp is specified.")
 
 
 def main(args):
@@ -399,10 +402,8 @@ def main(args):
     if args.hf_model is not None:
         model = create_wrapped_hf_model(args)
     else:
-        with torch.device("meta" if args.fsdp else args.device):
+        with torch.device("meta" if args.init_meta_device else args.device):
             model = create_model(args)
-        if not args.fsdp:
-            model.reset_parameters()
 
     args.vocab_size = model.vocab_size
     args.seq_len = model.seq_len

open_lm/params.py

@@ -561,6 +561,12 @@ def parse_args(args):
         default=False,
         help="If true, ignore parse errors in data loading. This should ideally be False, as errors in dataloading can point to bigger issues in your dataset. However, this can be useful when training on a large dataset which has a couple errors.",
     )
+    parser.add_argument(
+        "--init-meta-device",
+        action="store_true",
+        default=False,
+        help="If true, initialize the model on the meta device. This allows creating models larger than CPU memory. Can only be specified if --fsdp is also specified."
+    )
 
     add_model_args(parser)