Use accelerator pin memory

Author: tjruwase

deepnvme/ds_io/ds_io_sweep.sh

@@ -23,7 +23,7 @@ fi
 if [[ ${XFER} == "cpu" ]]; then
     xfer_opt=""
 elif [[ ${XFER} == "gpu" ]]; then
-    xfer_opt="--gpu"
+    xfer_opt="--gpu --use_accelerator_pin_memory"
 elif [[ ${XFER} == "gds" ]]; then 
     xfer_opt="--gpu --use_gds"
 else

deepnvme/model_checkpoint/save_model_utils.py

@@ -1,6 +1,6 @@
 import argparse
 import os
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import T5ForConditionalGeneration
 from torch_save_utils import PINNED_BUFFER_MB
 
@@ -23,10 +23,13 @@ def _get_hf_model(tag):
     model_name = HF_MODELS_DICT[tag]
     if tag == TINY_T5:
         model = T5ForConditionalGeneration.from_pretrained(model_name)
+
     else:
         model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
 
-    return model, model_name, tag
+    return model, tokenizer, model_name, tag
 
 def get_model(model_tag):
     return _get_hf_model(model_tag)
@@ -108,6 +111,13 @@ def parse_arguments():
                         action='store_true',
                         help='Disable double buffering of i/o buffer.')
 
+    parser.add_argument('--safetensors',
+                        action='store_true',
+                        help='Use safetensors load/save.')
+
+    parser.add_argument('--regular_torch_save',
+                        action='store_true',
+                        help='Use vanilla torch.save.')
 
     #parser.add_argument('--single_writer', action='store_true', help='Disable parallel rank writes of data parallel (replicated) state')
 

deepnvme/model_checkpoint/torch_save_model.py

@@ -57,7 +57,7 @@ def main():
     if not validate_arguments(args):
         quit()
     load_io_ops(args)
-    model, model_name, ckpt_name = get_model(args.model)
+    model, tokenizer, model_name, ckpt_name = get_model(args.model)
     if args.half:
         model = model.half()
     if args.gpu: