Function sft_one_iter

examples/unsloth/sft_algorithm.py:135–329

One iteration of SFT. The idea is to get all trace data from the rollouts, and then use the reward to select the top triplets to train on. Performs (1) rollout - data collection, (2) data conversion, (3) SFT training, and (4) model saving. Args: iteration: The iteration number

(
    *,
    iteration: int,
    store: LightningStore,
    model_path: str,
    train_dataset: Dataset[GsmProblem],
    llm_proxy: LLMProxy,
    data_adapter: TraceToTripletBase,
    triplet_fraction: float,
    vllm_port: int,
)

async def sft_one_iter(
    *,
    iteration: int,
    store: LightningStore,
    model_path: str,
    train_dataset: Dataset[GsmProblem],
    llm_proxy: LLMProxy,
    data_adapter: TraceToTripletBase,
    triplet_fraction: float,
    vllm_port: int,
) -> str:
    """One iteration of SFT.

    The idea is to get all trace data from the rollouts, and then use the reward to select the top triplets to train on.

    Performs (1) rollout - data collection, (2) data conversion, (3) SFT training, and (4) model saving.

    Args:
        iteration: The iteration number.
        store: The LightningStore instance.
        model_path: The path to the model to train. Must be a local path.
        train_dataset: The dataset to train on.
        llm_proxy: The LLM proxy instance. Used to shield between the inference endpoint and the rollout runners.
        data_adapter: The data adapter instance. This is used to convert the trace data recorded by LLM proxy.
        triplet_fraction: The fraction of triplets to use for SFT.
        vllm_port: The port to serve vLLM chat completion endpoint.

    Returns:
        The path to the saved model (next generation).
    """

    console.print(f"\n[bold red][Algo][/bold red] Starting iteration {iteration}")

    # 1. Rollout to get trace data
    if not os.path.exists(model_path):
        raise ValueError(f"Model path {model_path} does not exist.")

    # First launch the vLLM server
    with vllm_server(model_path, vllm_port) as server_address:
        # Update the model list of the LLM proxy and start it
        model_list: List[ModelConfig] = [
            {
                "model_name": "Qwen3-4B-Instruct",
                "litellm_params": {
                    "model": f"hosted_vllm/{model_path}",
                    "api_base": server_address,
                },
            }
        ]
        console.print(f"[bold red][Algo][/bold red] Updating model list and restarting LLM proxy: {model_list}")
        llm_proxy.update_model_list(model_list)
        # Restart the LLM proxy after backend model list update
        # If LLM proxy has never been started, it will be started
        await llm_proxy.restart()

        # Put the LLM proxy address into the store as an address
        resources_update = await store.add_resources(
            {