make_evaluator.py

#!/usr/bin/env python3

"""Train a MiniHF evaluator model (instruction tuned LoRA)."""

import argparse
from functools import partial
import os
from pathlib import Path
import sys

os.environ["BITSANDBYTES_NOWELCOME"] = "1"

import accelerate
import datasets
import datasets.distributed
import peft
import torch
from torch import optim
from torch.nn import functional as F
from torch.utils import data
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from tqdm import tqdm

print = tqdm.external_write_mode()(print)


def batch_to_tensors(batch, device="cpu"):
 batch = [item["input_ids"] for item in batch]
 seq_len = max(len(x) for x in batch)
 input_ids = torch.zeros(len(batch), seq_len, dtype=torch.long, device=device)
 attention_mask = torch.zeros(len(batch), seq_len, dtype=torch.long, device=device)
 for i, x in enumerate(batch):
 input_ids[i, : len(x)] = torch.tensor(x, dtype=torch.long, device=device)
 attention_mask[i, : len(x)] = 1
 return input_ids, attention_mask


def weighted_mean(x, w=None, dim=None, keepdim=False, dtype=None):
 w = x.new_tensor(1.0) if w is None else w
 w = w.expand_as(x)
 dim = tuple(range(x.ndim)) if dim is None else dim
 num = torch.sum(x * w, dim=dim, keepdim=keepdim, dtype=dtype)
 denom = torch.sum(w, dim=dim, keepdim=keepdim, dtype=dtype)
 return num / denom


class EndlessHFDataset(data.IterableDataset):
 def __init__(self, dataset):
 super().__init__()
 self.dataset = dataset

 def __iter__(self):
 while True:
 yield from self.dataset
 self.dataset.set_epoch(self.dataset._epoch + 1)


def main():
 parser = argparse.ArgumentParser(
 description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
 )
 parser.add_argument("--batch-size", type=int, default=4, help="batch size per process")
 parser.add_argument("--examples", type=int, default=100000, help="train for n examples")
 parser.add_argument("--output-dir", type=Path, default="evaluator", help="output directory")
 parser.add_argument("--save-every", type=int, default=10000, help="save every n examples")
 args = parser.parse_args()

 dataset_seed = 100
 lora_rank = 32
 lr = 1e-4
 max_len = 2048
 model_name = "openlm-research/open_llama_7b"

 # Initialize Accelerate
 accelerator = accelerate.Accelerator(mixed_precision="bf16", dispatch_batches=False)
 device = accelerator.device
 print0 = accelerator.on_local_main_process(print)

 # Load tokenizer
 print0(f"### Loading tokenizer: {model_name}", file=sys.stderr)
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
 tokenizer.pad_token = tokenizer.eos_token

 # Load model
 print0(f"### Loading model: {model_name}", file=sys.stderr)
 bnb_config = BitsAndBytesConfig(
 load_in_4bit=True,
 bnb_4bit_compute_dtype=torch.bfloat16,
 bnb_4bit_quant_type="nf4",
 bnb_4bit_use_double_quant=True,
 )
 with accelerator.main_process_first():
 model = AutoModelForCausalLM.from_pretrained(
 model_name,
 device_map="auto" if accelerator.num_processes == 1 else {"": device},
 quantization_config=bnb_config,
 torch_dtype=torch.bfloat16,
 trust_remote_code=True,
 )
 accelerator.wait_for_everyone()

 # Set up the LoRA
 print0("### Setting up the LoRA", file=sys.stderr)
 peft_config = peft.LoraConfig(
 peft.TaskType.CAUSAL_LM,
 inference_mode=False,
 r=lora_rank,
 lora_alpha=8,
 lora_dropout=0.0,
 target_modules=[
 "self_attn.q_proj",
 "self_attn.k_proj",
 "self_attn.v_proj",
 "self_attn.o_proj",
 "mlp.gate_proj",
 "mlp.up_proj",
 "mlp.down_proj",
 "lm_head",
 ],
 )
 model = peft.get_peft_model(model, peft_config)
 accelerator.wait_for_everyone()

 # Set up the model
 model.train()
 model.gradient_checkpointing_enable()
 model.enable_input_require_grads()
 if accelerator.is_local_main_process:
 model.print_trainable_parameters()

 # Dataset helper functions
 def combine_flan(row):
 return row["inputs"] + "<|end|>" + row["targets"] + tokenizer.eos_token

 def combine_dolly(row):
 return (
 row["context"]
 + "\n\n"
 + row["instruction"]
 + "<|end|>"
 + row["response"]
 + tokenizer.eos_token
 )

 def to_tokens(combine_fn, row):
 return tokenizer(combine_fn(row))

 def exclude_too_long(row):
 return len(row["input_ids"]) <= max_len

 # Load dataset
 print0("### Loading datasets", file=sys.stderr)
 with accelerator.main_process_first():
 dataset_1 = datasets.load_dataset("Muennighoff/flan", streaming=True)
 dataset_2 = datasets.load_dataset("databricks/databricks-dolly-15k", streaming=True)
 accelerator.wait_for_everyone()
 dataset_1 = dataset_1["train"].map(partial(to_tokens, combine_flan))
 dataset_2 = dataset_2["train"].map(partial(to_tokens, combine_dolly))
 dataset = (
 datasets.interleave_datasets([dataset_1, dataset_2], probabilities=[0.9, 0.1])
 .filter(exclude_too_long)
 .shuffle(seed=dataset_seed)
 .select_columns(["input_ids"])
 )
 dataset = datasets.distributed.split_dataset_by_node(
 dataset, accelerator.process_index, accelerator.num_processes
 )
 dataloader = data.DataLoader(
 EndlessHFDataset(dataset),
 batch_size=args.batch_size,
 collate_fn=batch_to_tensors,
 drop_last=True,
 )

 # Set up optimizer
 opt = optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.99))

 # Wrap objects
 model, opt, dataloader = accelerator.prepare(model, opt, dataloader)

 # Test max sequence length
 print0("### Testing max sequence length", file=sys.stderr)
 input_ids = torch.zeros([args.batch_size, max_len], dtype=torch.long, device=device)
 attention_mask = torch.ones([args.batch_size, max_len], dtype=torch.long, device=device)
 outputs = model(input_ids, attention_mask=attention_mask, use_cache=False)
 accelerator.backward(outputs.logits.sum() * 0)
 opt.zero_grad()
 torch.cuda.empty_cache()

 def save_model():
 print0("### Saving model", file=sys.stderr)
 accelerator.wait_for_everyone()
 if accelerator.is_main_process:
 unwrapped_model = accelerator.unwrap_model(model)
 unwrapped_model.save_pretrained(args.output_dir, safe_serialization=True)
 tokenizer.save_pretrained(args.output_dir)

 # Train
 print0("### Training", file=sys.stderr)
 examples = 0
 last_save = 0
 pbar = tqdm(
 disable=not accelerator.is_local_main_process,
 total=args.examples,
 unit="ex",
 smoothing=0.01,
 )

 try:
 for batch in dataloader:
 input_ids, attention_mask = batch
 with accelerator.accumulate(model):
 # Forward pass
 outputs = model(
 input_ids[:, :-1],
 attention_mask=attention_mask[:, :-1],
 use_cache=False,
 )
 losses = F.cross_entropy(
 outputs.logits.transpose(-1, -2),
 input_ids[:, 1:],
 reduction="none",
 )
 mask = attention_mask[:, :-1] * attention_mask[:, 1:]
 loss = weighted_mean(losses, mask, dtype=torch.float32)

 # Backward pass and optimizer step
 accelerator.backward(loss)
 opt.step()
 opt.zero_grad()

 global_batch_size = args.batch_size * accelerator.num_processes
 examples += global_batch_size
 pbar.update(global_batch_size)

 global_loss = accelerator.reduce(loss, "mean")
 print0(f"examples: {examples}, loss: {global_loss.item():g}")

 if examples >= args.examples:
 save_model()
 break

 if examples - last_save >= args.save_every:
 save_model()
 last_save += args.save_every

 except KeyboardInterrupt:
 pass

 finally:
 pbar.close()


if __name__ == "__main__":
 main()