mira/unsloth-train-scripts
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: _

Browse Source
This commit is contained in:
Mira
2025-02-15 00:59:47 +06:00
parent 39e69c90b1
commit fc9c04fe9c
2 changed files with 27 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
main.py
View File

@@ -84,72 +84,46 @@ def run_sweep(base_config: TrainingConfig, dataset_path: str):
"method": "bayes",
"metric": {"name": "val_loss", "goal": "minimize"},
"parameters": {
"base_learning_rate": {
"learning_rate": {
"distribution": "log_uniform_values",
"min": 1e-5,
"max": 1e-3,
"max": 1e-4,
},
"lora_r": {"values": [32, 64]},
# "lora_alpha": {"values": [16, 32, 64, 128]},
"per_device_train_batch_size": {"values": [16]},
"gradient_accumulation_steps": {"values": [2, 4, 8]},
"lora_r": {"values": [32]},
"lora_alpha": {"values": [64]},
"per_device_train_batch_size": {"values": [64]},
"gradient_accumulation_steps": {"values": [1]},
"num_train_epochs": {"values": [1]},
},
"early_terminate": {"type": "hyperband", "min_iter": 100},
}
sweep_id = wandb.sweep(sweep_config, project=base_config.wandb.project)
def compute_scaled_config(config_dict: dict) -> dict:
base_batch_size = config_dict["per_device_train_batch_size"] # 16
effective_batch_size = (
config_dict["per_device_train_batch_size"]
* config_dict["gradient_accumulation_steps"]
)
# apply square root scaling for Adam
lr_scale_factor = (effective_batch_size / base_batch_size) ** 0.5
# scale the learning rate
config_dict["learning_rate"] = (
config_dict["base_learning_rate"] * lr_scale_factor
)
config_dict["effective_batch_size"] = effective_batch_size
config_dict["lr_scale_factor"] = lr_scale_factor
config_dict["lora_alpha"] = 2 * wandb.config.lora_r
del config_dict["base_learning_rate"]
return config_dict
def sweep_train():
with wandb.init() as run:
# Convert base config to dict
config_dict = (
asdict(base_config)
if hasattr(base_config, "__dataclass_fields__")
else base_config.__dict__.copy()
)
# Update with sweep parameters
config_dict.update(wandb.config)
config_dict = compute_scaled_config(config_dict)
wandb.log(
{
"effective_batch_size": config_dict["effective_batch_size"],
"lr_scale_factor": config_dict["lr_scale_factor"],
"scaled_learning_rate": config_dict["learning_rate"],
"lora_alpha": config_dict["lora_alpha"],
}
# Set lora_alpha based on lora_r
# config_dict["lora_alpha"] = 2 * config_dict["lora_r"]
# Log effective batch size for monitoring
effective_batch_size = (
config_dict["per_device_train_batch_size"]
* config_dict["gradient_accumulation_steps"]
)
wandb.log({"effective_batch_size": effective_batch_size})
config_dict.pop("effective_batch_size")
config_dict.pop("lr_scale_factor")
run_config = TrainingConfig(**config_dict)
# Create training config and run
run_config = reconstruct_dataclass(TrainingConfig, config_dict)
print(run_config)
train_single_run(run_config, dataset_path)
wandb.agent(sweep_id, function=sweep_train)