(
# 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 = "",
)
| 131 | |
| 132 | |
| 133 | def train( |
| 134 | # model/data params |
| 135 | base_model: str = "", # the only required argument |
| 136 | train_file: str="", |
| 137 | eval_file: str="", |
| 138 | output_dir: str = "", |
| 139 | sample: int = -1, |
| 140 | seed: int = 42, |
| 141 | |
| 142 | # training hyperparams |
| 143 | batch_size: int = 128, |
| 144 | micro_batch_size: int = 4, |
| 145 | num_epochs: int = 10, |
| 146 | learning_rate: float = 3e-4, |
| 147 | cutoff_len: int = 512, |
| 148 | # llm hyperparams |
| 149 | group_by_length: bool = False, # faster, but produces an odd training loss curve |
| 150 | freeze_LLM: bool = False, # freeze LLM parameters, only train new token embeddings |
| 151 | # wandb params |
| 152 | wandb_project: str = "", |
| 153 | wandb_run_name: str = "", |
| 154 | resume_from_checkpoint: str = None, # either training checkpoint or final adapter |
| 155 | category: str="", |
| 156 | train_from_scratch: bool = False, |
| 157 | sid_index_path: str = "", |
| 158 | item_meta_path: str = "", |
| 159 | ): |
| 160 | set_seed(seed) |
| 161 | os.environ['WANDB_PROJECT'] = wandb_project |
| 162 | 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"} |
| 163 | print(category) |
| 164 | category = category_dict[category] |
| 165 | assert ( |
| 166 | base_model |
| 167 | ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'" |
| 168 | gradient_accumulation_steps = batch_size // micro_batch_size |
| 169 | |
| 170 | device_map = "auto" |
| 171 | world_size = int(os.environ.get("WORLD_SIZE", 1)) |
| 172 | ddp = world_size != 1 |
| 173 | if ddp: |
| 174 | device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)} |
| 175 | gradient_accumulation_steps = gradient_accumulation_steps // world_size |
| 176 | |
| 177 | if not train_from_scratch: |
| 178 | model = AutoModelForCausalLM.from_pretrained( |
| 179 | base_model, |
| 180 | torch_dtype=torch.bfloat16, |
| 181 | ) |
| 182 | else: |
| 183 | config = AutoConfig.from_pretrained(base_model) |
| 184 | model = AutoModelForCausalLM.from_config(config) |
| 185 | print("Training from scratch!") |
| 186 | |
| 187 | tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True) |
| 188 | original_vocab_size = len(tokenizer) |
| 189 | |
| 190 | # Add Special Tokens |
nothing calls this directly
no test coverage detected