from transformers import GPT2LMHeadModel, AutoTokenizer from transformers import AdamW, get_scheduler, set_seed from datasets import load_dataset from accelerate import Accelerator import datasets, transformers from huggingface_hub import Repository from torch.utils.data import IterableDataset from torch.utils.data.dataloader import DataLoader from torch.utils.tensorboard import SummaryWriter from argparse import Namespace import torch import logging import wandb class ConstantLengthDataset(IterableDataset): def __init__( self, tokenizer, dataset, da_type, seq_length=1024, num_of_sequences=1024, chars_per_token=5.2, ): self.tokenizer = tokenizer self.concat_token_id = tokenizer.bos_token_id self.dataset = dataset self.da_type = da_type self.seq_length = seq_length self.input_characters = seq_length * chars_per_token * num_of_sequences def __iter__(self): iterator = iter(self.dataset[f"{self.da_type}"]) more_examples = True while more_examples: buffer, buffer_len = [], 0 while True: if buffer_len >= self.input_characters: break try: buffer.append(next(iterator)["text"]) buffer_len += len(buffer[-1]) except StopIteration: more_examples = False break 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: yield torch.tensor(input_ids) def setup_logging(project_name): logger = logging.getLogger(__name__) logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, handlers=[ logging.FileHandler(f"log/debug_{accelerator.process_index}.log"), logging.StreamHandler(), ], ) 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_info() 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, run_name def create_dataloaders(args): ds_kwargs = {"streaming": True} train_data = load_dataset( "text", data_files={"train": ["train_raw.txt"]}, **ds_kwargs ) train_data = train_data.shuffle(buffer_size=args.shuffle_buffer, seed=args.seed) valid_data = load_dataset( "text", data_files={"valid": ["valid_raw.txt"]}, **ds_kwargs ) train_dataset = ConstantLengthDataset( tokenizer, train_data, da_type="train", seq_length=args.seq_length ) valid_dataset = ConstantLengthDataset( tokenizer, valid_data, da_type="valid", 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 def get_grouped_params(model, args, no_decay=["bias", "LayerNorm.weight"]): params_with_wd, params_without_wd = [], [] for n, p in model.named_parameters(): if any(nd in n for nd in no_decay): params_without_wd.append(p) else: params_with_wd.append(p) return [ {"params": params_with_wd, "weight_decay": args.weight_decay}, {"params": params_without_wd, "weight_decay": 0.0}, ] def log_metrics(step, metrics): logger.info(f"Step {step}: {metrics}") if accelerator.is_main_process: wandb.log(metrics) [tb_writer.add_scalar(k, v, step) for k, v in metrics.items()] def evaluate(args): model.eval() losses = [] for step, batch in enumerate(eval_dataloader): with torch.no_grad(): 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 loss = torch.mean(torch.cat(losses)) try: perplexity = torch.exp(loss) except OverflowError: perplexity = float("inf") return loss.item(), perplexity.item() # Accelerator accelerator = Accelerator(dispatch_batches=True) acc_state = {str(k): str(v) for k, v in accelerator.state.__dict__.items()} # Hyperparameters project_name = "krupalkp/custom_llm-small" dataset_name = "../codeparrot" config = { "train_batch_size": 2, "valid_batch_size": 2, "weight_decay": 0.1, "shuffle_buffer": 1_000, "learning_rate": 2e-4, "lr_scheduler_type": "cosine", "num_warmup_steps": 750, "gradient_accumulation_steps": 16, "max_train_steps": 50_000, "max_eval_steps": -1, "seq_length": 1024, "seed": 1, "save_checkpoint_steps": 50_000, } args = Namespace(**config, **acc_state) samples_per_step = accelerator.state.num_processes * args.train_batch_size set_seed(args.seed) # Logging logger, tb_writer, run_name = setup_logging(project_name.split("/")[1]) logger.info(accelerator.state) # Load model and tokenizer if accelerator.is_main_process: hf_repo = Repository("./", clone_from=project_name, revision=run_name) model = GPT2LMHeadModel.from_pretrained("./") tokenizer = AutoTokenizer.from_pretrained("./") # Load dataset and dataloader train_dataloader, eval_dataloader = create_dataloaders(args) # Prepare the optimizer and learning rate scheduler optimizer = AdamW(get_grouped_params(model, args), lr=args.learning_rate) lr_scheduler = get_scheduler( name=args.lr_scheduler_type, optimizer=optimizer, num_warmup_steps=args.num_warmup_steps, num_training_steps=args.max_train_steps, ) def get_lr(): return optimizer.param_groups[0]["lr"] # Prepare everything with our `accelerator`. model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare( model, optimizer, train_dataloader, eval_dataloader ) # Train model model.train() completed_steps = 0 for step, batch in enumerate(train_dataloader, start=1): loss = model(batch, labels=batch, use_cache=False).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: accelerator.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() lr_scheduler.step() optimizer.zero_grad() completed_steps += 1 if step % args.save_checkpoint_steps == 0: logger.info("Evaluating and saving model checkpoint") eval_loss, perplexity = evaluate(args) log_metrics(step, {"loss/eval": eval_loss, "perplexity": perplexity}) accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) 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: break # Evaluate and save the last checkpoint logger.info("Evaluating and saving model after training") eval_loss, perplexity = evaluate(args) log_metrics(step, {"loss/eval": eval_loss, "perplexity": perplexity}) accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) if accelerator.is_main_process: unwrapped_model.save_pretrained("./") hf_repo.push_to_hub(commit_message=f"final model")