Function sft_algorithm

examples/unsloth/sft_algorithm.py:332–379

Run the complete SFT algorithm with multiple iterations. This is the main entry point for running the SFT training pipeline. It sets up the LLM proxy, data adapter, and runs multiple iterations of model training. The function performs these steps for each iteration: 1. Serves the c

(*, store: LightningStore)

async def sft_algorithm(*, store: LightningStore) -> None:
    """Run the complete SFT algorithm with multiple iterations.

    This is the main entry point for running the SFT training pipeline. It sets up
    the LLM proxy, data adapter, and runs multiple iterations of model training.

    The function performs these steps for each iteration:
    1. Serves the current model via vLLM
    2. Collects rollout data using the model
    3. Converts trace data to training triplets
    4. Trains the model on top-performing examples
    5. Saves the improved model for the next iteration

    Args:
        store: The LightningStore instance for managing rollouts and trace data.
    """
    train_dataset = load_math_dataset()

    # Constants for the SFT algorithm
    MAX_ITERATIONS = 2
    VLLM_PORT = 12316
    LLM_PROXY_PORT = 12358
    TRAIN_TRIPLET_FRACTION = 0.5

    # Download the model before starting the script:
    # hf download unsloth/Qwen3-4B-Instruct-2507 --local-dir models/version_0
    model_path = "models/version_0"

    # Create the LLM proxy for rollout worker access and trace data collection
    llm_proxy = LLMProxy(port=LLM_PROXY_PORT, store=store)

    # This data adapter util is used to convert the trace data recorded by LLM proxy
    # into a format suitable for SFT
    data_adapter = LlmProxyTraceToTriplet()

    for iteration in range(MAX_ITERATIONS):
        model_path = await sft_one_iter(
            iteration=iteration,
            store=store,
            model_path=model_path,
            train_dataset=train_dataset,
            llm_proxy=llm_proxy,
            data_adapter=data_adapter,
            triplet_fraction=TRAIN_TRIPLET_FRACTION,
            vllm_port=VLLM_PORT,
        )

    console.print(f"[bold red][Algo][/bold red] Final model path: {model_path}")


if __name__ == "__main__":