fine_tune.py

import time
import logging
import pickle
import os
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForCausalLM
from accelerate import Accelerator

class TokenizedDataset(Dataset):
    def __init__(self, tokenized_data):
        self.tokenized_data = tokenized_data

    def __len__(self):
        return len(self.tokenized_data)

    def __getitem__(self, idx):
        return torch.tensor(self.tokenized_data[idx])


def fine_tune(
    pretrained_model_file_path="zephyr-7b-beta",
    lr=5e-5,
    gradient_clip=1.0,
    num_epochs=3,
    out_dir=None,
):
    """
    Fine-tune a pretrained model on a dataset using pytorch.
    """
    if out_dir is None:
        raise ValueError("out_dir must be specified.")
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    logging.basicConfig(
        filename=os.path.join(out_dir, "fine_tune.log"),
        filemode="w",
        encoding="utf-8",
        level=logging.INFO,
    )
    start_time = time.time()
    logging.info(f"start of log. initial time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}")
    logging.info(f"pretrained_model_file_path: {pretrained_model_file_path}")
    logging.info(f"lr: {lr}")
    logging.info(f"gradient_clip: {gradient_clip}")
    logging.info(f"out_dir: {out_dir}")
    logging.info(f"num_epochs: {num_epochs}")

    logging.info("load tokenized dataset")
    with open(os.path.join(out_dir, 'tokenized_dataset.pkl'), "rb") as file:
        dataset = TokenizedDataset(pickle.load(file))

    logging.info("initialize dataloader")
    batch_size = 4
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    del dataset # free memory
    
    logging.info("initialize model")
    model = AutoModelForCausalLM.from_pretrained(pretrained_model_file_path, torch_dtype=torch.bfloat16)
    logging.info(f"model dtype: {model.dtype}")

    logging.info("initialize optimizer")
    # optimizer = AdamW(model.parameters(), lr=lr)
    optimizer = torch.optim.SGD(model.parameters(), lr=lr)

    last_save_time = start_time
    save_id = 0
    loss_history = []
    accelerator = Accelerator()
    dataloader, model, optimizer = accelerator.prepare(dataloader, model, optimizer)
    model.train() # set model to training mode
    logging.info("start training loop")
    for epoch in range(num_epochs):
        total_loss = 0
        for batch_idx, input_ids in enumerate(dataloader):
            # Forward pass
            with accelerator.autocast():
                outputs = model(input_ids, labels=input_ids)
                loss = outputs.loss
            loss_history.append(loss.item())

            # log stats
            elapsed_time = time.time() - start_time
            total_batches_processed = epoch * len(dataloader) + batch_idx + 1
            estimated_total_time = elapsed_time * num_epochs * len(dataloader) / total_batches_processed
            estimated_time_remaining = estimated_total_time - elapsed_time
            remaining_hours, remaining_rem = divmod(estimated_time_remaining, 3600)
            remaining_minutes, remaining_seconds = divmod(remaining_rem, 60)
            info_string = (f"Epoch {epoch + 1}/{num_epochs}, Batch {batch_idx + 1}/{len(dataloader)}, "
                           f"Loss: {loss.item():.4f}, GPU Usage: {torch.cuda.memory_allocated() / 1024**3:.2f} GB, "
                           f"Time Remaining: {int(remaining_hours)}h {int(remaining_minutes)}m {int(remaining_seconds)}s")
            logging.info(info_string)

            # Backward pass and optimization
            optimizer.zero_grad()
            accelerator.backward(loss)
            if gradient_clip:
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=gradient_clip)
            optimizer.step()

            total_loss += loss.item()

            # for every 10 hours that pass, save one temporary checkpoint
            if time.time() - last_save_time > 36000:
                save_id += 1
                cp_dir = f"{out_dir}/cp_{save_id}"
                logging.info(f"Checkpoint {save_id}")
                logging.info("Saving checkpoint weights")
                model.save_pretrained(cp_dir)
                logging.info("Saving loss history so far")
                with open(os.path.join(cp_dir, "loss_history.pkl"), "wb") as file:
                    pickle.dump(loss_history, file)
                last_save_time = time.time()

    model.save_pretrained(out_dir)
    with open(os.path.join(out_dir, "loss_history.pkl"), "wb") as file:
        pickle.dump(loss_history, file)

if __name__ == "__main__":
    fine_tune(
        pretrained_model_file_path="Yi-6B",
        lr=1e-3,
        gradient_clip=1.0,
        num_epochs=1,
        out_dir="Yi-6B_textbooks_v1"
    )