any plan suport distill in axolotl?

[sankexin]

Now, Distill+SFT+RL is a standardized posttraining method that improves model performance and achieves good results. Hope it can be realized in axolotl .

[NanoCode012]

Hey! This is the first request we had for distill. Could you let us know which tool that you would normally use for distillation?

We'll need to then discuss internally on whether to open a CLI for this option.

[sankexin]

Do I just need to make the YAML file and train like this?

Yes, you can take a regular example config and edit it to add the plugin+dataset which contains your logits. Then you can call axolotl train config.yaml.

Where can I see detailed distill training examples?

there's a linked example here https://huggingface.co/datasets/axolotl-ai-co/evolkit-logprobs-pipeline-75k-v2-sample

The results of the teacher model is offline and requires the use of distanabel to generate infer results. could it be done online in the future? much More convenient to use like this:

'''

Load and preprocess dataset

dataset = load_dataset(config["dataset"]["name"], split=config["dataset"]["split"])
dataset = dataset.shuffle(seed=config["dataset"]["seed"])

Load tokenizers

teacher_tokenizer = AutoTokenizer.from_pretrained(config["models"]["teacher"])
student_tokenizer = AutoTokenizer.from_pretrained(config["models"]["student"])

teacher_model = AutoModelForCausalLM.from_pretrained(config["models"]["teacher"], **model_kwargs)
student_model = AutoModelForCausalLM.from_pretrained(config["models"]["student"], **model_kwargs)

teacher_outputs = teacher_model(**inputs)
student_outputs = student_model(**inputs)

def distillation_loss(self, model, student_logits, teacher_logits, inputs, original_loss):
...
'''

teacher model offline & online infering, Having both ways at the same time is perfect! well control memory!

[NanoCode012]

The results of the teacher model is offline and requires the use of distanabel to generate infer results. could it be done online in the future?

Yes, we had some internal talks on this, but I don't think it's something we're looking into in the near future. It could change if interest in the feature increases however.

[sankexin]

The results of the teacher model is offline and requires the use of distanabel to generate infer results. could it be done online in the future?

Yes, we had some internal talks on this, but I don't think it's something we're looking into in the near future. It could change if interest in the feature increases however.

distilabel need api_key, I have a try with qwen3 to get logits,but I have faild, Its official examples cannot run successfully, looking forward to online release of distilling.

from datasets import load_dataset

from distilabel.models import InferenceEndpointsLLM
from distilabel.pipeline import Pipeline
from distilabel.steps.tasks import TextGeneration

with Pipeline() as pipeline:
    TextGeneration(
        llm=InferenceEndpointsLLM(
            model_id="Qwen/Qwen3-8B",
            generation_kwargs={"temperature": 0.7, "max_new_tokens": 512},
        ),
    )

if __name__ == "__main__":
    dataset = load_dataset("distilabel-internal-testing/instructions", split="test")
    distiset = pipeline.run(dataset=dataset)
    distiset.push_to_hub(repo_id="distilabel-example")

[NanoCode012]

distilabel need api_key

Hey, you can run a vllm serve instance that exposes an OPENAI compatible interface to use for logits generation.

Here's an old example:

        text_generation_vllm = ChatGeneration(
            name='text_generation_vllm',
            llm=OpenAILLM(
                model=f"{name}",
                base_url=f"{base_url}v1",
                api_key=f"{key}",
                max_retries=5,
                generation_kwargs={
                    "temperature": 2.0,
                    "max_new_tokens": 4096,
                    "logprobs": True,
                    "top_logprobs": 64,
                    "skip_special_tokens": False,
                },
                tokenizer=f"{name}",
            ),
            resources=StepResources(replicas=8),
            input_batch_size=2,
        )

[sankexin]

distilabel need api_key

Hey, you can run a vllm serve instance that exposes an OPENAI compatible interface to use for logits generation.

Here's an old example:

        text_generation_vllm = ChatGeneration(
            name='text_generation_vllm',
            llm=OpenAILLM(
                model=f"{name}",
                base_url=f"{base_url}v1",
                api_key=f"{key}",
                max_retries=5,
                generation_kwargs={
                    "temperature": 2.0,
                    "max_new_tokens": 4096,
                    "logprobs": True,
                    "top_logprobs": 64,
                    "skip_special_tokens": False,
                },
                tokenizer=f"{name}",
            ),
            resources=StepResources(replicas=8),
            input_batch_size=2,
        )

from distilabel.pipeline import Pipeline
from distilabel.steps import LoadDataFromDicts, LoadDataFromHub
from distilabel.steps.tasks import TextGeneration
from distilabel.models import OpenAILLM

from transformers import AutoTokenizer

def extract_token_ids_and_logprobs(response, tokenizer_name="Qwen/Qwen3-8B"):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) 
    result = []

    if response == [[]]:
        return [result]


    for item in response[0]:
        encoded = tokenizer.encode(item["token"], add_special_tokens=False)
        token_id = encoded[0] if len(encoded) > 0 else None

        if token_id is None:
            continue

        result.append({
            "logprob": item["logprob"],
            "token": f"token_id:{token_id}"
        })

    return [result]


VLLM_BASE_URL = "http://0.0.0.0:8000/v1" 
MODEL_NAME = "Qwen/Qwen3-8B" 

with Pipeline() as pipeline:
    # Step 1: 
    '''
    load_data = LoadDataFromDicts(
        data=[{"prompt": "Write a poem about the sun and moon."}]
    )
    '''
    load_data = LoadDataFromHub(
        output_mappings={"instruction": "instruction"}  
    )
    '''
    load_data = LoadDataFromHub(
        repo_id="tatsu-lab/alpaca",
        split="train",
        batch_size=8,
        output_mappings={"instruction": "instruction"}
    )
    load_data.load()
    result = next(load_data.process())
    print(result)
    '''

    # Step 2:
    llm = OpenAILLM(
        model=MODEL_NAME,
        base_url=VLLM_BASE_URL,
        api_key="EMPTY",  
        max_retries=5,
        generation_kwargs={
            "temperature": 2.0,
            "max_new_tokens": 4096,
            "top_p": 0.9,
            "logprobs": True,       
            "top_logprobs": 20,     
            # "skip_special_tokens": True,
        }
    )

    # Step 3: 
    text_generation = TextGeneration(
        name="text_generation",
        llm=llm,
        input_batch_size=4,  
        # input_mappings={"instruction": "instruction"},
    )

    # Step 
    load_data >> text_generation

    '''
    # 
    llm.load()
    output = llm.generate_outputs(inputs=[[{"role": "user", "content": "Give three tips for staying healthy"}]])
    print(output)
    '''

if __name__ == "__main__":
    # pipeline
    distiset = pipeline.run(
        parameters={
            load_data.name: {
                "repo_id": "tatsu-lab/alpaca",  
                "split": "train", 
                "batch_size": 8,
            },
        }
    )

    # 
    distiset.save_to_disk("distilabel-internal-testing/distilabel-example-instructions")
    # distiset.push_to_hub(repo_id="xxx/distilabel-example-instructions-Qwen3-8B")
    # print(distiset)
    import pandas as pd
    for step_name in list(distiset.keys()):
        formatted_data = []
        step_output = distiset[step_name]["train"]
        formatted_data = []
        for item in step_output:
            item["llm_text_generation_logprobs"] = extract_token_ids_and_logprobs(item.get("llm_text_generation_logprobs", ""))
            messages_combined = [
                {"role": "user", "content": item.get("instruction", "")},
                {"role": "assistant", "content": item.get("generation", "")}
                ]
            item["messages"] = {"role": "user", "content": item.get("instruction", "")}
            item["messages_combined"] =  messages_combined
            formatted_data.append(item)

        df = pd.DataFrame(formatted_data)
        # print(df)
        df.to_parquet(f"tatsu-lab/{step_name}.parquet", engine="pyarrow")
        print(f"✅ {step_name}.parquet")


    '''
    from datasets import load_dataset
    ds = load_dataset("distilabel-internal-testing/distilabel-example-instructions")
    print(ds['train'][0])
    '''

my result of train is not good, anything wrong with my code ?