()
| 62 | return run_dir |
| 63 | |
| 64 | def finetune(): |
| 65 | run_dir = create_run_directory(args) |
| 66 | |
| 67 | # Initialize wandb with the run directory |
| 68 | wandb_run_name = os.path.basename(run_dir) |
| 69 | wandb_run = wandb.init( |
| 70 | project="project-name", |
| 71 | config=args, |
| 72 | dir=os.path.join(run_dir, "logs") |
| 73 | ) |
| 74 | |
| 75 | # Save wandb run ID to a file |
| 76 | with open(os.path.join(run_dir, "wandb_run_id.txt"), "w") as f: |
| 77 | f.write(wandb_run.id) |
| 78 | |
| 79 | # Create model and tokenizer |
| 80 | model, tokenizer = create_model_tokenizer_it(args) |
| 81 | |
| 82 | # Data handling |
| 83 | train_dataset = load_and_preprocess_it(tokenizer=tokenizer, args=args) |
| 84 | data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer) |
| 85 | train_loader = DataLoader( |
| 86 | train_dataset, |
| 87 | batch_size=args.eg_bs, |
| 88 | shuffle=True, |
| 89 | collate_fn=data_collator |
| 90 | ) |
| 91 | data_module = dict(train_dataset=train_dataset, data_collator=data_collator) |
| 92 | |
| 93 | named_grads = None |
| 94 | |
| 95 | total_training_steps = len(train_loader) * args.epochs |
| 96 | |
| 97 | eff_lr = args.lr/(args.warmup_ratio * total_training_steps) |
| 98 | |
| 99 | |
| 100 | named_grads = estimate_and_process_grads_torch( |
| 101 | model=model, |
| 102 | dataloader=train_loader, |
| 103 | lr=eff_lr, |
| 104 | num_samples=50, |
| 105 | ) |
| 106 | |
| 107 | # Create peft model |
| 108 | model, lora_config = create_peft_model_it(model, args) |
| 109 | |
| 110 | |
| 111 | # Convert model to xs |
| 112 | |
| 113 | reconstr_config_path = os.path.join(run_dir, "reconstruct_config.yaml") |
| 114 | # Copy reconstruct config to run directory |
| 115 | with open("config/reconstruct_config.yaml", 'r') as src, open(reconstr_config_path, 'w') as dst: |
| 116 | reconstr_config = yaml.load(src, Loader=yaml.FullLoader) |
| 117 | reconstr_config['svd']['rank'] = args.lora_r |
| 118 | yaml.dump(reconstr_config, dst) |
| 119 | |
| 120 | # Save the required JSON file with the correct name |
| 121 | json_path = os.path.join(run_dir, "reconstr_config.json") # Note: reconstr not reconstruct |
no test coverage detected