arg_parser.py

import os
import argparse
from Finetune4bConfig import Finetune4bConfig


def parse_commandline():
    parser = argparse.ArgumentParser(
        prog=__file__.split(os.path.sep)[-1],
        description="Produce LoRA in 4bit training",
        usage="%(prog)s [config] [training]\n\nAll arguments are optional",
    )

    parser.add_argument(
        "dataset",
        nargs="?",
        default="./dataset.json",
        help="Path to dataset file. Default: %(default)s",
    )

    parser_config = parser.add_argument_group("config")
    parser_training = parser.add_argument_group("training")

    # Config args group
    parser_config.add_argument(
        "--ds_type",
        choices=["shot", "txt", "alpaca", "gpt4all", "bluemoon"],
        default="alpaca",
        required=False,
        help="Dataset structure format. Default: %(default)s",
    )
    parser_config.add_argument(
        "--lora_out_dir",
        default="alpaca_lora",
        required=False,
        help="Directory to place new LoRA. Default: %(default)s",
    )
    parser_config.add_argument(
        "--lora_apply_dir",
        default=None,
        required=False,
        help="Path to directory from which LoRA has to be applied before training. Default: %(default)s",
    )
    parser_training.add_argument(
        "--resume_checkpoint",
        default=None,
        required=False,
        help="Resume training from specified checkpoint. Default: %(default)s",
    )
    parser_config.add_argument(
        "--llama_q4_config_dir",
        default="./llama-13b-4bit/",
        required=False,
        help="Path to the config.json, tokenizer_config.json, etc. Default: %(default)s",
    )
    parser_config.add_argument(
        "--llama_q4_model",
        default="./llama-13b-4bit.pt",
        required=False,
        help="Path to the quantized model in huggingface format. Default: %(default)s",
    )

    # Training args group
    parser_training.add_argument(
        "--mbatch_size",
        default=1,
        type=int,
        help="Micro-batch size. Default: %(default)s",
    )
    parser_training.add_argument(
        "--batch_size", default=2, type=int, help="Batch size. Default: %(default)s"
    )
    parser_training.add_argument(
        "--epochs", default=3, type=int, help="Epochs. Default: %(default)s"
    )
    parser_training.add_argument(
        "--lr", default=2e-4, type=float, help="Learning rate. Default: %(default)s"
    )
    parser_training.add_argument(
        "--cutoff_len", default=256, type=int, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "--lora_r", default=8, type=int, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "--lora_alpha", default=16, type=int, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "--lora_dropout", default=0.05, type=float, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "--grad_chckpt",
        action="store_true",
        required=False,
        help="Use gradient checkpoint. For 30B model. Default: %(default)s",
    )
    parser_training.add_argument(
        "--grad_chckpt_ratio",
        default=1,
        type=float,
        help="Gradient checkpoint ratio. Default: %(default)s",
    )
    parser_training.add_argument(
        "--val_set_size",
        default=0.2,
        type=float,
        help="Validation set size. Default: %(default)s",
    )
    parser_training.add_argument(
        "--warmup_steps", default=50, type=int, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "--save_steps", default=50, type=int, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "--save_total_limit", default=3, type=int, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "--logging_steps", default=10, type=int, help="Default: %(default)s"
    )
    parser_training.add_argument(
        "-c",
        "--checkpoint",
        action="store_true",
        help="Produce checkpoint instead of LoRA. Default: %(default)s",
    )
    parser_training.add_argument(
        "--skip",
        action="store_true",
        help="Don't train model. Can be useful to produce checkpoint from existing LoRA. Default: %(default)s",
    )
    parser_training.add_argument(
        "--verbose",
        action="store_true",
        help="If output log of training. Default: %(default)s",
    )

    #
    parser_config.add_argument(
        "--ppo-dataset",
        default=None,
        required=False,
        help="Path to dataset for PPO training. Default: %(default)s",
    )
    parser_training.add_argument(
        "--ppo-train",
        action="store_true",
        required=False,
        help="Train a PPO model after training the LoRA. Default: %(default)s",
    )
    parser_training.add_argument(
        "--train",
        default=True,
        action="store_true",
        required=True,
        help="Train the LoRA. Default: %(default)s",
    )

    # Data args
    parser_training.add_argument(
        "--txt_row_thd", default=-1, type=int, help="Custom thd for txt rows."
    )
    parser_training.add_argument(
        "--use_eos_token",
        default=1,
        type=int,
        help="Use eos token instead if padding with 0. enable with 1, disable with 0.",
    )

    # V2 model support
    parser_training.add_argument(
        "--groupsize", type=int, default=-1, help="Groupsize of v2 model"
    )
    parser_training.add_argument("--v1", action="store_true", help="Use V1 model")

    # Multi GPU Support
    parser_training.add_argument(
        "--local_rank",
        type=int,
        default=0,
        help="local rank if using torch.distributed.launch",
    )

    # Flash Attention
    parser_training.add_argument(
        "--flash_attention",
        action="store_true",
        help="enables flash attention, can improve performance and reduce VRAM use",
    )
    parser_training.add_argument(
        "--xformers",
        action="store_true",
        help="enables xformers memory efficient attention, can improve performance and reduce VRAM use",
    )

    # Train Backend
    parser_training.add_argument(
        "--backend", type=str, default="cuda", help="Backend to use. Triton or Cuda."
    )

    # Positional embeddings
    parser_training.add_argument(
        "--xpos",
        action="store_true",
        help="Use xPos rotary embedding to extrapolate over longer seq len than the model is trained on",
    )
    parser_training.add_argument(
        "--nope",
        action="store_true",
        help="Don't use any positional embeddings (https://arxiv.org/abs/2305.19466)",
    )

    return vars(parser.parse_args())


def get_config() -> Finetune4bConfig:
    args = parse_commandline()
    return Finetune4bConfig(
        dataset=args["dataset"],
        ds_type=args["ds_type"],
        lora_out_dir=args["lora_out_dir"],
        lora_apply_dir=args["lora_apply_dir"],
        resume_checkpoint=args["resume_checkpoint"],
        llama_q4_config_dir=args["llama_q4_config_dir"],
        llama_q4_model=args["llama_q4_model"],
        mbatch_size=args["mbatch_size"],
        batch_size=args["batch_size"],
        epochs=args["epochs"],
        lr=args["lr"],
        cutoff_len=args["cutoff_len"],
        lora_r=args["lora_r"],
        lora_alpha=args["lora_alpha"],
        lora_dropout=args["lora_dropout"],
        val_set_size=args["val_set_size"],
        gradient_checkpointing=args["grad_chckpt"],
        gradient_checkpointing_ratio=args["grad_chckpt_ratio"],
        warmup_steps=args["warmup_steps"],
        save_steps=args["save_steps"],
        save_total_limit=args["save_total_limit"],
        logging_steps=args["logging_steps"],
        checkpoint=args["checkpoint"],
        skip=args["skip"],
        verbose=args["verbose"],
        txt_row_thd=args["txt_row_thd"],
        use_eos_token=args["use_eos_token"] != 0,
        groupsize=args["groupsize"],
        v1=args["v1"],
        local_rank=args["local_rank"],
        flash_attention=args["flash_attention"],
        xformers=args["xformers"],
        backend=args["backend"],
        xpos=args["xpos"],
        ppo_train=args["ppo_train"],
        train=args["train"],
        ppo_dataset=args["ppo_dataset"],
        nope=args["nope"],
    )