Add argument to offload model from memory to disk

environment.yml

@@ -20,6 +20,7 @@ dependencies:
   - torchaudio=0.13.1
   - torchvision=0.14.1
   - pip:
+      - accelerate==0.17.1
       - datasets==2.10.1
       - loguru==0.6.0
       - netifaces==0.11.0

pretrained/GPT-NeoX-20B/prepare.py

@@ -12,6 +12,8 @@
                         help='model-name')
     parser.add_argument('--save-dir', type=str, default=DIR, 
                         help='model-name')
+    parser.add_argument('--offload-dir', type=str, default=None,
+                        help='directory to offload from memory')
     args = parser.parse_args()
 
     if not os.path.exists(args.save_dir):
@@ -24,7 +26,14 @@
     config.save_pretrained(save_path)
     tokenizer = AutoTokenizer.from_pretrained(args.model_name)
     tokenizer.save_pretrained(save_path)
-    model = AutoModelForCausalLM.from_pretrained(args.model_name, torch_dtype=torch.float16)
+
+    # offload model from memory to disk if offload-dir is specified
+    if args.offload_dir is not None:
+        if not os.path.exists(args.offload_dir):
+            os.mkdir(args.offload_dir)
+        model = AutoModelForCausalLM.from_pretrained(args.model_name, torch_dtype=torch.float16, device_map="auto", offload_folder=args.offload_dir)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(args.model_name, torch_dtype=torch.float16)
 
     item = {}
     item['embed_in.weight'] = model.gpt_neox.embed_in.weight