Update train.py

train.py

@@ -6,7 +6,6 @@ import trl
 from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM, TrainingArguments, PreTrainedTokenizerFast, AdamW, get_linear_schedule_with_warmup
 from datasets import load_dataset
 from tokenizers import ByteLevelBPETokenizer
-from accelerate import Accelerator
 
 MAX_SEQ_LENGTH = 512
 BATCH_SIZE = 32
@@ -94,7 +93,7 @@ def configure_tokenizer(tokenizer):
     chat_template = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '<|user|>\n' + message['content'] + '<|end|>\n' }}{% elif message['role'] == 'assistant' %}{{ '<|bot|>\n' + message['content'] + '<|end|>\n' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}{{ eos_token }}"
     tokenizer.chat_template = chat_template
 
-def train_model(accelerator, model, tokenizer, dataset, push):
+def train_model(model, tokenizer, dataset, push):
     args = TrainingArguments(
         output_dir="model",
         num_train_epochs=EPOCHS,
@@ -114,8 +113,6 @@ def train_model(accelerator, model, tokenizer, dataset, push):
         num_warmup_steps=args.warmup_steps, 
         num_training_steps=len(dataset) * args.num_train_epochs // args.gradient_accumulation_steps
     )
-
-    model, optimizer = accelerator.prepare(model, optimizer)
     
     dataset = dataset.map(lambda examples: format_prompts(examples, tokenizer), batched=True)
     trainer = trl.SFTTrainer(
@@ -146,8 +143,7 @@ def main(push_to_hub=True):
     tokenizer = create_tokenizer(training_corpus)
     configure_tokenizer(tokenizer)
     model = create_model(tokenizer)
-    accelerator = Accelerator()
-    train_model(accelerator, model, tokenizer, dataset, push_to_hub)
+    train_model(model, tokenizer, dataset, push_to_hub)
 
 if __name__ == "__main__":
     main(PUSH_TO_HUB)