switched to accelerator

train.py

@@ -1,37 +1,29 @@
 import os
 import torch
 import torch.nn as nn
-import torch.distributed as dist
-import torch.multiprocessing as mp
 import bitsandbytes as bnb
 from datasets import load_dataset
 import transformers
 from transformers import AutoTokenizer, AutoConfig, LLaMAForCausalLM, LLaMATokenizer
 from peft import prepare_model_for_int8_training, LoraConfig, get_peft_model
 
-def setup(rank, world_size):
-    os.environ['MASTER_ADDR'] = 'localhost'
-    os.environ['MASTER_PORT'] = '12355'
-
-    # Initialize the process group
-    dist.init_process_group("nccl", rank=rank, world_size=world_size)
-
-def cleanup():
-    dist.destroy_process_group()
-
-def train(rank, world_size):
-    setup(rank, world_size)
-
-    # os.system("nvidia-smi")
-    # os.system("git clone https://github.com/tloen/alpaca-lora.git")
-    # os.chdir("alpaca-lora/")
-    # os.system("pip install -q datasets loralib sentencepiece")
-    # os.system("pip uninstall -y transformers")
-    # os.system("pip install -q git+https://github.com/zphang/transformers@c3dc391")
-    # os.system("pip install -q git+https://github.com/huggingface/peft.git")
-    # os.system("pip install bitsandbytes")
-
-    # os.system("conda install -y -c conda-forge cudatoolkit")
+# Import the necessary Accelerate modules
+from accelerate import Accelerator, DistributedType
+
+def train():
+    # Initialize the Accelerator
+    accelerator = Accelerator(
+        device_placement=True,
+        split_batches=False,
+        mixed_precision="fp16",
+        # distributed_type=DistributedType.MULTI_GPU,
+        gradient_accumulation_steps=1,
+        rng_types=["torch", "cuda"],
+        log_with=["tensorboard", "wandb", "comet_ml"],
+        project_dir="./",
+        even_batches=True,
+        step_scheduler_with_optimizer=True
+    )
 
     MICRO_BATCH_SIZE = 1
     BATCH_SIZE = 16
@@ -42,18 +34,20 @@ def train(rank, world_size):
     LORA_ALPHA = 8
     LORA_DROPOUT = 0.05
 
-    device = torch.device(f"cuda:{rank}")
+    device = accelerator.device
+
     model = LLaMAForCausalLM.from_pretrained(
         "decapoda-research/llama-7b-hf",
         load_in_8bit=True,
         device_map="auto",
     )
-    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank], output_device=rank)
+
+    model = accelerator.prepare(model)
     tokenizer = LLaMATokenizer.from_pretrained(
         "decapoda-research/llama-7b-hf", add_eos_token=True
     )
 
-    model = prepare_model_for_int8_training(model.module)
+    model = prepare_model_for_int8_training(model)
 
     config = LoraConfig(
         r=LORA_R,
@@ -63,7 +57,7 @@ def train(rank, world_size):
         bias="none",
         task_type="CAUSAL_LM",
     )
-    model.module = get_peft_model(model.module, config)
+    model = get_peft_model(model, config)
     tokenizer.pad_token_id = 0
     data = load_dataset("json", data_files="../samples.json")
 
@@ -89,29 +83,28 @@ def train(rank, world_size):
         )
     )
 
+    training_args = transformers.TrainingArguments(
+        per_device_train_batch_size=MICRO_BATCH_SIZE,
+        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
+        warmup_steps=100,
+        num_train_epochs=EPOCHS,
+        learning_rate=LEARNING_RATE,
+        logging_steps=1,
+        output_dir=f"lora-smartscraper-{accelerator.process_index}",
+        save_total_limit=3,
+    )
+    # training_args = accelerator.update_arguments(training_args)
+
     trainer = transformers.Trainer(
         model=model,
         train_dataset=data["train"],
-        args=transformers.TrainingArguments(
-            per_device_train_batch_size=MICRO_BATCH_SIZE,
-            gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
-            warmup_steps=100,
-            num_train_epochs=EPOCHS,
-            learning_rate=LEARNING_RATE,
-            fp16=True,
-            logging_steps=1,
-            output_dir=f"lora-smartscraper-{rank}",
-            save_total_limit=3,
-        ),
+        args=training_args,
         data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
     )
     model.config.use_cache = False
     trainer.train(resume_from_checkpoint=False)
 
-    model.save_pretrained(f"lora-smartscraper-{rank}")
-
-    cleanup()
+    model.save_pretrained(f"lora-smartscraper-{accelerator.process_index}")
 
 if __name__ == "__main__":
-    world_size = torch.cuda.device_count()
-    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)
+    train()