(
# model/data params
base_model: str = "", # the only required argument
train_file: str="",
eval_file: str="",
output_dir: str = "",
sample: int = -1,
seed: int = 42,
# training hyperparams
batch_size: int = 128,
micro_batch_size: int = 4,
num_epochs: int = 10,
learning_rate: float = 3e-4,
cutoff_len: int = 512,
# llm hyperparams
group_by_length: bool = False, # faster, but produces an odd training loss curve
freeze_LLM: bool = False, # freeze LLM parameters, only train new token embeddings
# wandb params
wandb_project: str = "",
wandb_run_name: str = "",
resume_from_checkpoint: str = None, # either training checkpoint or final adapter
category: str="",
train_from_scratch: bool = False,
sid_index_path: str = "",
item_meta_path: str = "",
)
| 88 | |
| 89 | |
| 90 | def train( |
| 91 | # model/data params |
| 92 | base_model: str = "", # the only required argument |
| 93 | train_file: str="", |
| 94 | eval_file: str="", |
| 95 | output_dir: str = "", |
| 96 | sample: int = -1, |
| 97 | seed: int = 42, |
| 98 | |
| 99 | # training hyperparams |
| 100 | batch_size: int = 128, |
| 101 | micro_batch_size: int = 4, |
| 102 | num_epochs: int = 10, |
| 103 | learning_rate: float = 3e-4, |
| 104 | cutoff_len: int = 512, |
| 105 | # llm hyperparams |
| 106 | group_by_length: bool = False, # faster, but produces an odd training loss curve |
| 107 | freeze_LLM: bool = False, # freeze LLM parameters, only train new token embeddings |
| 108 | # wandb params |
| 109 | wandb_project: str = "", |
| 110 | wandb_run_name: str = "", |
| 111 | resume_from_checkpoint: str = None, # either training checkpoint or final adapter |
| 112 | category: str="", |
| 113 | train_from_scratch: bool = False, |
| 114 | sid_index_path: str = "", |
| 115 | item_meta_path: str = "", |
| 116 | ): |
| 117 | set_seed(seed) |
| 118 | os.environ['WANDB_PROJECT'] = wandb_project |
| 119 | category_dict = {"Industrial_and_Scientific": "industrial and scientific items", "Office_Products": "office products", "Toys_and_Games": "toys and games", "Sports": "sports and outdoors", "Books": "books"} |
| 120 | print(category) |
| 121 | category = category_dict[category] |
| 122 | assert ( |
| 123 | base_model |
| 124 | ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'" |
| 125 | gradient_accumulation_steps = batch_size // micro_batch_size |
| 126 | |
| 127 | device_map = "auto" |
| 128 | world_size = int(os.environ.get("WORLD_SIZE", 1)) |
| 129 | ddp = world_size != 1 |
| 130 | if ddp: |
| 131 | device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)} |
| 132 | gradient_accumulation_steps = gradient_accumulation_steps // world_size |
| 133 | |
| 134 | if not train_from_scratch: |
| 135 | model = AutoModelForCausalLM.from_pretrained( |
| 136 | base_model, |
| 137 | torch_dtype=torch.bfloat16, |
| 138 | ) |
| 139 | else: |
| 140 | config = AutoConfig.from_pretrained(base_model) |
| 141 | model = AutoModelForCausalLM.from_config(config) |
| 142 | print("Training from scratch!") |
| 143 | |
| 144 | tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True) |
| 145 | tokenizer.pad_token = tokenizer.eos_token |
| 146 | tokenizer.pad_token_id = tokenizer.eos_token_id |
| 147 | tokenizer.padding_side = "left" |
nothing calls this directly
no test coverage detected