102 lines
4.0 KiB
Python
102 lines
4.0 KiB
Python
"""
|
|
LLM Fine-tuning Entry Point
|
|
|
|
Standard RDLoop entry point for LLM fine-tuning, consistent with data science implementation.
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Optional, cast
|
|
|
|
import fire
|
|
|
|
from rdagent.app.finetune.llm.conf import FT_RD_SETTING
|
|
from rdagent.log import rdagent_logger as logger
|
|
from rdagent.scenarios.finetune.loop import LLMFinetuneRDLoop
|
|
|
|
|
|
def main(
|
|
path: Optional[str] = None,
|
|
checkout: bool = True,
|
|
user_target_scenario: Optional[str] = None,
|
|
benchmark: Optional[str] = None,
|
|
benchmark_description: Optional[str] = None,
|
|
dataset: Optional[str] = None,
|
|
base_model: Optional[str] = None,
|
|
upper_data_size_limit: Optional[int] = None,
|
|
step_n: Optional[int] = None,
|
|
loop_n: Optional[int] = None,
|
|
timeout: Optional[str] = None,
|
|
):
|
|
"""
|
|
LLM fine-tuning entry point
|
|
|
|
Parameters
|
|
----------
|
|
path :
|
|
A path like `$LOG_PATH/__session__/1/0_propose`. This indicates that we restore the state after finishing step 0 in loop 1.
|
|
checkout :
|
|
Used to control the log session path. Boolean type, default is True.
|
|
- If True, the new loop will use the existing folder and clear logs for sessions after the one corresponding to the given path.
|
|
- If False, the new loop will use the existing folder but keep the logs for sessions after the one corresponding to the given path.
|
|
dataset : str
|
|
Dataset name for fine-tuning (e.g., 'shibing624/alpaca-zh')
|
|
base_model : str, optional
|
|
Model name for fine-tuning (e.g., 'Qwen/Qwen2.5-1.5B-Instruct').
|
|
If not provided, auto-selects optimal model based on hardware and dataset.
|
|
step_n : int, optional
|
|
Number of steps to run; if None, runs indefinitely until completion or error
|
|
loop_n : int, optional
|
|
Number of loops to run; if None, runs indefinitely until completion or error
|
|
timeout : str, optional
|
|
Maximum duration for the entire process
|
|
|
|
Examples:
|
|
.. code-block:: bash
|
|
dotenv run -- python rdagent/app/finetune/llm/loop.py --dataset shibing624/alpaca-zh --base-model Qwen/Qwen2.5-1.5B-Instruct
|
|
dotenv run -- python rdagent/app/finetune/llm/loop.py --dataset shibing624/alpaca-zh # TODO: not enabled yet
|
|
"""
|
|
|
|
if user_target_scenario:
|
|
FT_RD_SETTING.user_target_scenario = user_target_scenario
|
|
assert (
|
|
FT_RD_SETTING.user_target_scenario is None
|
|
), "user_target_scenario is not yet supported, please specify via benchmark and benchmark_description"
|
|
if upper_data_size_limit:
|
|
FT_RD_SETTING.upper_data_size_limit = upper_data_size_limit
|
|
logger.info(f"Set upper_data_size_limit to {FT_RD_SETTING.upper_data_size_limit}")
|
|
if benchmark:
|
|
FT_RD_SETTING.target_benchmark = benchmark
|
|
if benchmark_description:
|
|
FT_RD_SETTING.benchmark_description = benchmark_description
|
|
assert FT_RD_SETTING.user_target_scenario or (
|
|
FT_RD_SETTING.target_benchmark and FT_RD_SETTING.benchmark_description
|
|
), "Either user_target_scenario or target_benchmark must be specified for LLM fine-tuning."
|
|
|
|
# Update configuration with provided parameters
|
|
if dataset:
|
|
FT_RD_SETTING.dataset = dataset
|
|
if base_model:
|
|
FT_RD_SETTING.base_model = base_model
|
|
|
|
# Create and run LLM fine-tuning loop
|
|
data_set_target = FT_RD_SETTING.dataset if FT_RD_SETTING.dataset else "auto generated dataset"
|
|
model_target = FT_RD_SETTING.base_model if FT_RD_SETTING.base_model else "auto selected model"
|
|
|
|
# Temporary assertion until auto-selection is implemented
|
|
assert (
|
|
FT_RD_SETTING.base_model is not None
|
|
), "Base model auto selection not yet supported, please specify via --base-model"
|
|
|
|
logger.info(f"Starting LLM fine-tuning on dataset='{data_set_target}' with model='{model_target}'")
|
|
|
|
if path is None:
|
|
loop = LLMFinetuneRDLoop(FT_RD_SETTING)
|
|
else:
|
|
loop = cast(LLMFinetuneRDLoop, LLMFinetuneRDLoop.load(str(path), checkout=checkout))
|
|
|
|
asyncio.run(loop.run(step_n=step_n, loop_n=loop_n, all_duration=timeout))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
fire.Fire(main)
|