Update train.py

train.py

@@ -5,11 +5,13 @@ from transformers import (
     AutoTokenizer, LlamaConfig, AutoModelForCausalLM, LlamaForCausalLM,
     TrainingArguments, PreTrainedTokenizerFast, AdamW, get_cosine_schedule_with_warmup
 )
-from datasets import load_dataset
+from datasets import load_dataset, Dataset
 from tokenizers import ByteLevelBPETokenizer
 from torch.utils.data import DataLoader
+from torch.cuda.amp import autocast, GradScaler
+from itertools import islice
 
-BATCH_SIZE = 8
+BATCH_SIZE = 32
 EPOCHS = 1
 LEARNING_RATE = 1e-4
 FACTOR = 768
@@ -19,7 +21,8 @@ INPUT_DATASET = "HuggingFaceTB/smollm-corpus"
 INSTRUCT_DATASET = "nroggendorff/elephant"
 OUTPUT_REPO = "nroggendorff/smallama"
 INSTRUCT_FINETUNE_BOOL = False
-INIT = 0
+INIT = 1#/16
+SHARD_SIZE = int(5e+5)
 FP16 = True
 WARMUP_STEPS = 1000
 WEIGHT_DECAY = 0.01
@@ -30,20 +33,12 @@ NUM_WORKERS = 4
 def load_data():
     if not INSTRUCT_FINETUNE_BOOL:
         dataset = load_dataset(INPUT_DATASET, "cosmopedia-v2", split="train", streaming=True)
-        dataset = custom_shard_stream(dataset)
+        start = INIT * SHARD_SIZE
+        dataset = Dataset.from_dict({'text': [example['text'] for example in islice(dataset, start, start + SHARD_SIZE)]})
     else:
         dataset = load_dataset(INSTRUCT_DATASET, split="train")
     return dataset
 
-def custom_shard_stream(dataset, shard_size=5e5, shard_index=0):
-    def shard_generator():
-        count = 0
-        for example in dataset:
-            if count % shard_size == shard_index:
-                yield example
-            count += 1
-    return shard_generator()
-
 def create_tokenizer(training_corpus):
     tokenizer = ByteLevelBPETokenizer()
     special_tokens = ["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
@@ -59,11 +54,11 @@ def create_tokenizer(training_corpus):
     return fast_tokenizer
 
 def load_tokenizer():
-    return AutoTokenizer.from_pretrained(OUTPUT_REPO)
+    return AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")#OUTPUT_REPO)
 
 def get_training_corpus(dataset):
-    for example in dataset:
-        yield example['text']
+    for i in range(0, len(dataset['text']), 1000):
+        yield dataset['text'][i : i + 1000]
 
 def format_prompts(examples, tokenizer, isinst):
     texts = []
@@ -140,38 +135,44 @@ def train_model(model, tokenizer, dataset, push, isinst):
         save_total_limit=2,
     )
 
+    dataset = dataset.shard(num_shards=len(dataset) // SHARD_SIZE, index=INIT)
+
     optimizer = AdamW(model.parameters(), lr=args.learning_rate, weight_decay=WEIGHT_DECAY)
     scheduler = get_cosine_schedule_with_warmup(
         optimizer,
         num_warmup_steps=args.warmup_steps, 
-        num_training_steps=args.num_train_epochs
+        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,
-        optimizers=(optimizer, scheduler),
-        max_seq_length=MAX_SEQ_LENGTH
+        dataset_text_field='text',
+        max_seq_length=MAX_SEQ_LENGTH,
+        optimizers=(optimizer, scheduler)
     )
-
-    train_result = trainer.train()
-
+    
+    train = 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_result.training_loss:.4f}"
-        trainer.model.push_to_hub(repo_id, commit_message=msg, force=True)
-        trainer.tokenizer.push_to_hub(repo_id, commit_message=msg, force=True)
+        msg = f"Training loss: {train.training_loss:.4f}"
+        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:
-        trainer.model.save_pretrained("model")
-        trainer.tokenizer.save_pretrained("tokenizer")
+        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:
+    if not is_inst_finetune and INIT == 0 and False:
         training_corpus = get_training_corpus(dataset)
         tokenizer = create_tokenizer(training_corpus)
     else: