wehub-resource-sync
65 lines
1.9 KiB
Python
65 lines
1.9 KiB
Python
"""
|
|
RL Post-training Entry Point
|
|
"""
|
|
|
|
import asyncio
|
|
from typing import Optional
|
|
|
|
import typer
|
|
from typing_extensions import Annotated
|
|
|
|
from rdagent.app.rl.conf import RL_RD_SETTING
|
|
from rdagent.log import rdagent_logger as logger
|
|
from rdagent.scenarios.rl.loop import RLPostTrainingRDLoop
|
|
|
|
|
|
def main(
|
|
base_model: Annotated[Optional[str], typer.Option("--base-model", "-m")] = None,
|
|
benchmark: Annotated[Optional[str], typer.Option("--benchmark", "-b")] = None,
|
|
step_n: Optional[int] = None,
|
|
loop_n: Optional[int] = None,
|
|
timeout: Optional[str] = None,
|
|
):
|
|
"""
|
|
RL post-training entry point
|
|
|
|
Parameters
|
|
----------
|
|
base_model : str
|
|
Model name (e.g., 'Qwen2.5-Coder-0.5B-Instruct')
|
|
Docker path: /models/{base_model}
|
|
benchmark : str
|
|
Benchmark/dataset name (e.g., 'gsm8k')
|
|
Docker path: /data/{benchmark}
|
|
step_n : int, optional
|
|
Number of steps to run; if None, runs all steps per loop
|
|
loop_n : int, optional
|
|
Number of loops to run; if None, runs indefinitely
|
|
timeout : str, optional
|
|
Maximum duration for the entire process
|
|
|
|
Examples
|
|
--------
|
|
.. code-block:: bash
|
|
|
|
export RL_MODELS_DIR=/path/to/models
|
|
export RL_DATA_DIR=/path/to/data
|
|
python rdagent/app/rl/loop.py --base-model Qwen2.5-Coder-0.5B-Instruct --benchmark gsm8k
|
|
"""
|
|
# Update config from CLI
|
|
if base_model:
|
|
RL_RD_SETTING.base_model = base_model
|
|
if benchmark:
|
|
RL_RD_SETTING.benchmark = benchmark
|
|
|
|
logger.info(f"Starting RL post-training: model={RL_RD_SETTING.base_model}, benchmark={RL_RD_SETTING.benchmark}")
|
|
|
|
# RDLoop 会自动根据 RL_RD_SETTING.scen 创建 Scenario
|
|
# Scenario.__init__() 中会自动运行 baseline 评测
|
|
loop = RLPostTrainingRDLoop(RL_RD_SETTING)
|
|
asyncio.run(loop.run(step_n=step_n, loop_n=loop_n, all_duration=timeout))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(main)
|