()
| 48 | return train_dataset |
| 49 | |
| 50 | def main(): |
| 51 | base_model_id = "" |
| 52 | base_model = AutoModelForCausalLM.from_pretrained( |
| 53 | base_model_id, |
| 54 | load_in_4bit=True, |
| 55 | quantization_config=bnb_config, |
| 56 | torch_dtype=torch.bfloat16, |
| 57 | device_map="auto", |
| 58 | trust_remote_code=True, |
| 59 | ) |
| 60 | |
| 61 | base_model.config.use_cache = False # silence the warnings. Please re-enable for inference! |
| 62 | base_model.config.pretraining_tp = 1 |
| 63 | base_model.gradient_checkpointing_enable() |
| 64 | |
| 65 | # Load tokenizer |
| 66 | tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True) |
| 67 | tokenizer.padding_side = 'right' |
| 68 | tokenizer.pad_token = tokenizer.eos_token |
| 69 | tokenizer.add_eos_token = True |
| 70 | |
| 71 | |
| 72 | project = "hyperagent" |
| 73 | # base_model_name = "mistral" |
| 74 | base_model_name = "codellama" |
| 75 | run_name = base_model_name + "_" + project |
| 76 | output_dir = "model/" + run_name |
| 77 | |
| 78 | model = prepare_model_for_kbit_training(base_model) |
| 79 | peft_config = LoraConfig( |
| 80 | lora_alpha=16, |
| 81 | lora_dropout=0.1, |
| 82 | r=64, |
| 83 | bias="none", |
| 84 | task_type="CAUSAL_LM", |
| 85 | target_modules=["q_proj", "k_proj", "v_proj", "o_proj","gate_proj"] |
| 86 | ) |
| 87 | model = get_peft_model(model, peft_config) |
| 88 | |
| 89 | #Hyperparamter |
| 90 | training_arguments = TrainingArguments( |
| 91 | output_dir=output_dir, |
| 92 | num_train_epochs=1, |
| 93 | per_device_train_batch_size=4, |
| 94 | gradient_accumulation_steps=1, |
| 95 | optim="paged_adamw_32bit", |
| 96 | save_steps=25, |
| 97 | logging_steps=25, |
| 98 | learning_rate=2e-4, |
| 99 | weight_decay=0.001, |
| 100 | fp16=False, |
| 101 | bf16=False, |
| 102 | max_grad_norm=0.3, |
| 103 | max_steps=-1, |
| 104 | warmup_ratio=0.03, |
| 105 | group_by_length=True, |
| 106 | lr_scheduler_type="constant", |
| 107 | report_to="wandb" |
no test coverage detected