update training script and requirements

codeparrot_training.py

@@ -12,45 +12,41 @@ from argparse import Namespace
 import torch
 import logging
 import wandb
+import time
 
 
 class ConstantLengthDataset(IterableDataset):
-
-    def __init__(self, tokenizer, dataset, seq_length=1024, batch_size=3,
+    def __init__(self, tokenizer, dataset, seq_length=1024,
                  num_of_sequences=1024, chars_per_token=3.6):
         self.tokenizer = tokenizer
-        self.concatenation_token = tokenizer.bos_token
+        self.concat_token_id = tokenizer.bos_token_id
         self.dataset = dataset
         self.seq_length = seq_length
-        self.batch_size = batch_size
         self.input_characters = seq_length * chars_per_token * num_of_sequences
-
+        self.produced_samples = 0
     def __iter__(self):
         iterator = iter(self.dataset)
         more_examples = True
-        batch = []
         while more_examples:
-            buffer = ''
+            buffer = []
+            buffer_len = 0
             while True:
-                if len(buffer) >= self.input_characters:
+                if buffer_len >= self.input_characters:
                     break
                 try:
-                    next_example = next(iterator)['content']
-                    buffer = buffer + self.concatenation_token + next_example
+                    buffer.append(next(iterator)['content'])
+                    buffer_len += len(buffer[-1])
                 except StopIteration:
                     more_examples = False
                     break
-
-            tokenized_input = tokenizer(buffer, truncation=True,
-                                        max_length=self.seq_length,
-                                        return_overflowing_tokens=True)
-
-            for input_ids in tokenized_input['input_ids']:
+            tokenized_inputs = tokenizer(buffer, truncation=False)['input_ids']
+            all_token_ids = []
+            for tokenized_input in tokenized_inputs:
+                all_token_ids.extend(tokenized_input + [self.concat_token_id])
+            for i in range(0, len(all_token_ids), self.seq_length):
+                input_ids = all_token_ids[i : i + self.seq_length]
                 if len(input_ids) == self.seq_length:
-                    batch.append(input_ids)
-                if len(batch) == self.batch_size:
-                    yield torch.tensor(batch)
-                    batch = []
+                    yield torch.tensor(input_ids)
 
 def setup_logging(project_name):
     logger = logging.getLogger(__name__)
@@ -59,31 +55,31 @@ def setup_logging(project_name):
         datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO,)
     if accelerator.is_main_process: # we only want to setup logging once
         wandb.init(project=project_name, config=args)
+        run_name = wandb.run.name
         tb_writer = SummaryWriter()
         tb_writer.add_hparams(vars(args), {'0': 0})
         logger.setLevel(logging.INFO)
         datasets.utils.logging.set_verbosity_warning()
         transformers.utils.logging.set_verbosity_info()
     else:
+        tb_writer = None
+        run_name = ''
         logger.setLevel(logging.ERROR)
         datasets.utils.logging.set_verbosity_error()
         transformers.utils.logging.set_verbosity_error()
-    return logger, tb_writer
+    return logger, tb_writer, run_name
 
 def create_dataloaders(dataset_name):
     train_data = load_dataset(dataset_name+'-train', split="train",
                               streaming=True)
-    train_data = train_data.shuffle(buffer_size=args.shuffle_buffer)
+    train_data = train_data.shuffle(buffer_size=args.shuffle_buffer,
+                                    seed=args.seed)
     valid_data = load_dataset(dataset_name+'-valid', split="train",
                               streaming=True)
-
     train_dataset = ConstantLengthDataset(tokenizer, train_data,
-                                        seq_length=args.seq_length,
-                                        batch_size=args.train_batch_size)
+                                          seq_length=args.seq_length)
     valid_dataset = ConstantLengthDataset(tokenizer, valid_data,
-                                        seq_length=args.seq_length,
-                                        batch_size=args.valid_batch_size)
-
+                                          seq_length=args.seq_length)
     train_dataloader=DataLoader(train_dataset, batch_size=args.train_batch_size)
     eval_dataloader=DataLoader(valid_dataset, batch_size=args.valid_batch_size)
     return train_dataloader, eval_dataloader
@@ -107,7 +103,7 @@ def evaluate():
     losses = []
     for step, batch in enumerate(eval_dataloader):
         with torch.no_grad():
-            outputs = model(batch[0], labels=batch[0])
+            outputs = model(batch, labels=batch)
         loss = outputs.loss.repeat(args.valid_batch_size)
         losses.append(accelerator.gather(loss))
         if args.max_eval_steps > 0 and step >= args.max_eval_steps: break
@@ -119,19 +115,19 @@ def evaluate():
 # Hyperparameters
 project_name = 'transformersbook/codeparrot-small'
 dataset_name = 'transformersbook/codeparrot'
-config = {"train_batch_size": 4,
-          "valid_batch_size": 4,
+config = {"train_batch_size": 12,
+          "valid_batch_size": 12,
           "weight_decay": 0.1,
           "shuffle_buffer": 1000,
           "learning_rate": 5e-4,
           "lr_scheduler_type": "cosine",
-          "num_warmup_steps": 1000,
-          "gradient_accumulation_steps": 8,
-          "max_train_steps": 4096,
-          "max_eval_steps": 1024,
+          "num_warmup_steps": 2000,
+          "gradient_accumulation_steps": 1,
+          "max_train_steps": 8192,
+          "max_eval_steps": 512,
           "seq_length": 1024,
           "seed": 1,
-          "save_checkpoint_steps":4096,}
+          "save_checkpoint_steps":512,}
 args = Namespace(**config)
 set_seed(args.seed)
 
@@ -140,12 +136,12 @@ accelerator = Accelerator()
 samples_per_step = accelerator.state.num_processes * args.train_batch_size
 
 # Logging
-logger, tb_writer = setup_logging(project_name.split("/")[1])
+logger, tb_writer, run_name = setup_logging(project_name.split("/")[1])
 logger.info(accelerator.state)
-run_name = wandb.run.name
 
 # Load model and tokenizer
-hf_repo = Repository("./", clone_from=project_name, revision=run_name)
+if accelerator.is_main_process: # we only want to setup logging once
+    hf_repo = Repository("./", clone_from=project_name, revision=run_name)
 model = GPT2LMHeadModel.from_pretrained("./")
 tokenizer = AutoTokenizer.from_pretrained("./")
 
@@ -167,12 +163,12 @@ model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
 model.train()
 completed_steps = 0
 for step, batch in enumerate(train_dataloader, start=1):
-    loss = model(batch[0], labels=batch[0]).loss
+    loss = model(batch, labels=batch).loss
     log_metrics(step, {'lr': get_lr(), 'samples': step*samples_per_step,
                        'steps': completed_steps, 'loss/train': loss.item()})
     loss = loss / args.gradient_accumulation_steps
     accelerator.backward(loss)
-    if step % args.gradient_accumulation_steps == 0:     
+    if step % args.gradient_accumulation_steps == 0:
         optimizer.step()
         lr_scheduler.step()
         optimizer.zero_grad()
@@ -183,8 +179,8 @@ for step, batch in enumerate(train_dataloader, start=1):
         log_metrics(step, {'loss/eval': eval_loss, 'perplexity': perplexity})
         accelerator.wait_for_everyone()
         unwrapped_model = accelerator.unwrap_model(model)
-        unwrapped_model.save_pretrained("./")
         if accelerator.is_main_process:
+            unwrapped_model.save_pretrained("./")
             hf_repo.push_to_hub(commit_message=f'step {step}')
         model.train()
     if completed_steps >= args.max_train_steps:
@@ -196,7 +192,6 @@ eval_loss, perplexity = evaluate()
 log_metrics(step, {'loss/eval': eval_loss, 'perplexity': perplexity})
 accelerator.wait_for_everyone()
 unwrapped_model = accelerator.unwrap_model(model)
-unwrapped_model.save_pretrained("./")
 if accelerator.is_main_process:
-    try: hf_repo.push_to_hub(commit_message=f'final model')
-    except: logger.info('No changes to previously saved model.')
+    unwrapped_model.save_pretrained("./")
+    hf_repo.push_to_hub(commit_message=f'final model')

requirements.txt

@@ -1,7 +1,6 @@
-torch==1.9.0
 wandb
 tensorboard
-git+https://github.com/huggingface/huggingface_hub.git@push-branching
+git+https://github.com/huggingface/huggingface_hub.git
 git+https://github.com/huggingface/transformers.git
 git+https://github.com/huggingface/datasets.git@load_dataset-no-dataset-script
 git+https://github.com/huggingface/accelerate.git