codeparrot-small / codeparrot_training.py

Update codeparrot_training.py

74c33a2 over 2 years ago

No virus

8.81 kB

	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, infinite=False, seq_length=1024,
	num_of_sequences=1024, chars_per_token=3.6):
	self.tokenizer = tokenizer
	self.concat_token_id = tokenizer.bos_token_id
	self.dataset = dataset
	self.seq_length = seq_length
	self.input_characters = seq_length * chars_per_token * num_of_sequences
	self.epoch = 0
	self.infinite = infinite

	def __iter__(self):
	iterator = iter(self.dataset)
	more_examples = True
	while more_examples:
	buffer, buffer_len = [], 0
	while True:
	if buffer_len >= self.input_characters:
	break
	try:
	buffer.append(next(iterator)['content'])
	buffer_len += len(buffer[-1])
	except StopIteration:
	if self.infinite:
	iterator = iter(self.dataset)
	self.epoch += 1
	logger.info(f"Dataset epoch: {self.epoch}")
	else:
	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(dataset_name, args):
	ds_kwargs = {"streaming":True}
	train_data = load_dataset(dataset_name+'-train', split='train', **ds_kwargs)
	train_data = train_data.shuffle(buffer_size=args.shuffle_buffer,
	seed=args.seed)
	valid_data = load_dataset(dataset_name+'-valid', split="train", **ds_kwargs)
	train_dataset = ConstantLengthDataset(tokenizer, train_data, infinite=True,
	seq_length=args.seq_length)
	valid_dataset = ConstantLengthDataset(tokenizer, valid_data, infinite=False,
	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()
	acc_state = {str(k): str(v) for k, v in accelerator.state.__dict__.items()}

	# Hyperparameters
	project_name = 'lvwerra/codeparrot-small'
	dataset_name = '../codeparrot-clean'
	config = {"train_batch_size": 12,
	"valid_batch_size": 12,
	"weight_decay": 0.1,
	"shuffle_buffer": 1_000,
	"learning_rate": 5e-4,
	"lr_scheduler_type": "cosine",
	"num_warmup_steps": 2_000,
	"gradient_accumulation_steps": 1,
	"gradient_checkpointing": False,
	"max_train_steps": 150_000,
	"max_eval_steps": -1,
	"seq_length": 1024,
	"seed": 1,
	"save_checkpoint_steps": 15_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("./")
	if args.gradient_checkpointing:
	model.gradient_checkpointing_enable()
	tokenizer = AutoTokenizer.from_pretrained("./")

	# Load dataset and dataloader
	train_dataloader, eval_dataloader = create_dataloaders(dataset_name, 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)
	unwrapped_model.save_pretrained("./", save_function=accelerator.save)
	if accelerator.is_main_process:
	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)
	unwrapped_model.save_pretrained("./", save_function=accelerator.save)
	if accelerator.is_main_process:
	hf_repo.push_to_hub(commit_message=f'final model')