Update train.py

train.py

@@ -5,12 +5,12 @@ from transformers import (
     AutoTokenizer, LlamaConfig, AutoModelForCausalLM, LlamaForCausalLM,
     TrainingArguments, PreTrainedTokenizerFast, AdamW, get_cosine_schedule_with_warmup
 )
-from datasets import load_dataset, Dataset
+from datasets import load_dataset
 from tokenizers import ByteLevelBPETokenizer
 from torch.utils.data import DataLoader
 from torch.cuda.amp import autocast, GradScaler
 
-BATCH_SIZE = 8
+BATCH_SIZE = 32
 EPOCHS = 1
 LEARNING_RATE = 1e-4
 FACTOR = 768
@@ -32,9 +32,9 @@ NUM_WORKERS = 4
 def load_data():
     if not INSTRUCT_FINETUNE_BOOL:
         dataset = load_dataset(INPUT_DATASET, "cosmopedia-v2", split="train", streaming=True)
-        dataset = Dataset.from_generator(lambda: dataset.take(int(8e+6)))
+        dataset = dataset.take(int(8e+6))  # Keep streaming, no conversion to in-memory dataset
     else:
-        dataset = load_dataset(INSTRUCT_DATASET, split="train")
+        dataset = load_dataset(INSTRUCT_DATASET, split="train", streaming=True)
     return dataset
 
 def create_tokenizer(training_corpus):
@@ -54,10 +54,6 @@ def create_tokenizer(training_corpus):
 def load_tokenizer():
     return AutoTokenizer.from_pretrained(OUTPUT_REPO)
 
-def get_training_corpus(dataset):
-    for i in range(0, len(dataset['text']), 1000):
-        yield dataset['text'][i : i + 1000]
-
 def format_prompts(examples, tokenizer, isinst):
     texts = []
     for text in examples['text']:
@@ -139,53 +135,36 @@ def train_model(model, tokenizer, dataset, push, isinst):
     scheduler = get_cosine_schedule_with_warmup(
         optimizer,
         num_warmup_steps=args.warmup_steps, 
-        num_training_steps=(len(dataset) // args.per_device_train_batch_size) * args.num_train_epochs
-    )
-    
-    dataset = dataset.map(lambda examples: format_prompts(examples, tokenizer, isinst), batched=True, remove_columns=dataset.column_names)
-   
-    trainer = trl.SFTTrainer(
-        model=model,
-        tokenizer=tokenizer,
-        args=args,
-        train_dataset=dataset,
-        dataset_text_field='text',
-        max_seq_length=MAX_SEQ_LENGTH,
-        optimizers=(optimizer, scheduler)
+        num_training_steps=(SHARD_SIZE // args.per_device_train_batch_size) * args.num_train_epochs
     )
     
-    train = trainer.train()
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)
+
+    for batch in dataloader:
+        batch = format_prompts(batch, tokenizer, isinst)
+        trainer = trl.SFTTrainer(
+            model=model,
+            tokenizer=tokenizer,
+            args=args,
+            train_dataset=batch,
+            dataset_text_field='text',
+            max_seq_length=MAX_SEQ_LENGTH,
+            optimizers=(optimizer, scheduler)
+        )
+        trainer.train()
     
     trained_model = trainer.model
     trained_tokenizer = trainer.tokenizer
     
     if push:
         repo_id = OUTPUT_REPO + "-it" if INSTRUCT_FINETUNE_BOOL else OUTPUT_REPO
-        msg = f"Training loss: {train.training_loss:.4f}"
+        msg = "Training completed."
         trained_model.push_to_hub(repo_id, commit_message=msg, force=True)
         trained_tokenizer.push_to_hub(repo_id, commit_message=msg, force=True)
     else:
         trained_model.save_pretrained("model")
         trained_tokenizer.save_pretrained("tokenizer")
 
-def main(push_to_hub=True, is_inst_finetune=False):
-    dataset = load_data()
-    if not is_inst_finetune and INIT == 0:
-        training_corpus = get_training_corpus(dataset)
-        tokenizer = create_tokenizer(training_corpus)
-    else:
-        tokenizer = load_tokenizer()
-    
-    configure_tokenizer(tokenizer)
-    
-    if is_inst_finetune:
-        model = load_model()
-        model.resize_token_embeddings(len(tokenizer))
-    else:
-        model = create_model(tokenizer) if INIT == 0 else load_model()
-    
-    train_model(model, tokenizer, dataset, push_to_hub, is_inst_finetune)
-
 if __name__ == "__main__":
     main(PUSH_TO_HUB, INSTRUCT_FINETUNE_BOOL)
     raise Exception("Done baking!")