(fine_tune_id: str, base_url: str)
| 23 | |
| 24 | |
| 25 | def do_train(fine_tune_id: str, base_url: str): |
| 26 | logging.info(f"Beginning training process for model {fine_tune_id}") |
| 27 | |
| 28 | training_info_resp = get_training_info.sync_detailed( |
| 29 | client=AuthenticatedClient( |
| 30 | base_url=base_url, token=os.environ["AUTHENTICATED_SYSTEM_KEY"] |
| 31 | ), |
| 32 | fine_tune_id=fine_tune_id, |
| 33 | ) |
| 34 | |
| 35 | if training_info_resp.status_code != 200: |
| 36 | raise Exception(f"Failed to get training info: {training_info_resp}") |
| 37 | |
| 38 | training_info = training_info_resp.parsed |
| 39 | if not training_info: |
| 40 | raise Exception(f"Failed to get training info: {training_info_resp}") |
| 41 | |
| 42 | logging.info(f"Training info: {training_info.to_dict()}") |
| 43 | |
| 44 | logging.info("Downloading training data") |
| 45 | training_file = "/tmp/train.jsonl" |
| 46 | |
| 47 | urllib.request.urlretrieve(training_info.training_data_url, training_file) |
| 48 | |
| 49 | config_path = "/tmp/training-config.yaml" |
| 50 | lora_model_path = lora_model_cache_dir(fine_tune_id) |
| 51 | |
| 52 | os.makedirs(lora_model_path, exist_ok=True) |
| 53 | |
| 54 | # Clear the lora_model_path and merged_model_path directories |
| 55 | shutil.rmtree(lora_model_path, ignore_errors=True) |
| 56 | |
| 57 | config = training_info.training_config |
| 58 | config.datasets[0].path = training_file |
| 59 | config.output_dir = lora_model_path |
| 60 | |
| 61 | training_yaml = yaml.dump(config.to_dict()) |
| 62 | print(f"Training config:\n{training_yaml}") |
| 63 | with open(config_path, "w") as f: |
| 64 | f.write(training_yaml) |
| 65 | |
| 66 | logging.info("Beginning training") |
| 67 | try: |
| 68 | # We have to run this in a subprocess instead of importing axolotl directly |
| 69 | # because I haven't figured out how to free the GPU memory after training |
| 70 | # and we get OOMs when we reload the peft model to merge it. |
| 71 | subprocess.run( |
| 72 | [ |
| 73 | "accelerate", |
| 74 | "launch", |
| 75 | "-m", |
| 76 | "axolotl.cli.train", |
| 77 | config_path, |
| 78 | ], |
| 79 | check=True, |
| 80 | ) |
| 81 | except subprocess.CalledProcessError as e: |
| 82 | logging.error(f"Training failed: {e}") |
no test coverage detected