new config

train.py

@@ -1,30 +1,14 @@
 import os
 import torch
 import torch.nn as nn
-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
-
-# Import the necessary Accelerate modules
-from accelerate import Accelerator, DistributedType
+from accelerate import Accelerator
+from torch.utils.data import DataLoader
 
 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
     GRADIENT_ACCUMULATION_STEPS = BATCH_SIZE // MICRO_BATCH_SIZE
@@ -34,15 +18,11 @@ def train():
     LORA_ALPHA = 8
     LORA_DROPOUT = 0.05
 
-    device = accelerator.device
+    accelerator = Accelerator()
 
     model = LLaMAForCausalLM.from_pretrained(
-        "decapoda-research/llama-7b-hf",
-        load_in_8bit=True,
-        device_map="auto",
+        "decapoda-research/llama-7b-hf"
     )
-
-    model = accelerator.prepare(model)
     tokenizer = LLaMATokenizer.from_pretrained(
         "decapoda-research/llama-7b-hf", add_eos_token=True
     )
@@ -59,52 +39,48 @@ def train():
     )
     model = get_peft_model(model, config)
     tokenizer.pad_token_id = 0
-    data = load_dataset("json", data_files="../samples.json")
+    data = load_dataset("json", data_files="samples.json")
 
     def generate_prompt(data_point):
         if data_point["input"]:
-            return f"""### Instruction:
-    {data_point["instruction"]}
-    ### Input:
-    {data_point["input"]}
-    ### Response:
-    {data_point["output"]}"""
+            prompt = f"""### Instruction:
+        {data_point["instruction"]}
+        ### Input:
+        {data_point["input"]}
+        ### Response:
+        {data_point["output"]}"""
         else:
-            return f"""### Instruction:
-    {data_point["instruction"]}
-    ### Response:
-    {data_point["output"]}"""
+            prompt = f"""### Instruction:
+        {data_point["instruction"]}
+        ### Response:
+        {data_point["output"]}"""
 
-    data = data.shuffle().map(
-        lambda data_point: tokenizer(
-            generate_prompt(data_point),
-            truncation=False,
-            padding='longest',
-        )
-    )
+        input_tokens = tokenizer(prompt, truncation=False, padding='longest', return_tensors='pt')
+        output_tokens = tokenizer(data_point["output"], truncation=False, padding='longest', return_tensors='pt')
 
-    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)
+        return input_tokens, output_tokens["input_ids"].squeeze()
 
-    trainer = transformers.Trainer(
-        model=model,
-        train_dataset=data["train"],
-        args=training_args,
-        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
-    )
-    model.config.use_cache = False
-    trainer.train(resume_from_checkpoint=False)
+    data = data.shuffle().map(generate_prompt)
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)
+    model, optimizer = accelerator.prepare(model, optimizer)
+
+    train_dataloader = DataLoader(data["train"], batch_size=MICRO_BATCH_SIZE, shuffle=True)
+    train_dataloader = accelerator.prepare(train_dataloader)
+
+    for epoch in range(EPOCHS):
+        for step, batch in enumerate(train_dataloader):
+            inputs, labels = batch
+            inputs_tensor = torch.tensor(inputs["input_ids"], dtype=torch.long).unsqueeze(0).to(accelerator.device)
+            outputs = model(inputs_tensor)
+            labels_tensor = torch.tensor(labels, dtype=torch.long).to(accelerator.device)
+            loss = nn.CrossEntropyLoss()(outputs.logits.view(-1, outputs.logits.size(-1)), labels_tensor.view(-1))
+
+            accelerator.backward(loss)
+            optimizer.step()
+            optimizer.zero_grad()
 
     model.save_pretrained(f"lora-smartscraper-{accelerator.process_index}")
 
-if __name__ == "__main__":
-    train()
+    if __name__ == "__main__":
+        train()