Polyaxon & Unsloth Fine-tuning

How to use Polyaxon and Unsloth Fine-tuning together

Polyaxon schedules Unsloth fine-tuning as a GPU-backed job. Unsloth remains normal Python code inside your container; Polyaxon handles scheduling, parameters, logs, resource tracking, and artifacts.

Use this integration when you want to:

run Unsloth fine-tuning on Kubernetes GPU nodes,
parameterize model, dataset, and training settings,
track metrics through Polyaxon's Hugging Face trainer callback,
save LoRA adapters and optional merged model outputs under the run artifacts.

For the tracking-only API details, see the Unsloth tracking integration.

Container image

For a quick proof of concept, you can install dependencies in the job. For actual training, build an image with compatible CUDA, PyTorch, Unsloth, TRL, datasets, and Polyaxon versions:

FROM pytorch/pytorch:<cuda-runtime-tag>

RUN pip install --no-cache-dir \
    polyaxon \
    unsloth \
    datasets \
    trl \
    accelerate

Use the CUDA and PyTorch versions that match your GPU nodes. Runtime installs make runs slower and less reproducible; bake the image once and reuse it.

Polyaxon component

The following component creates an inline Unsloth training script and passes fine-tuning settings as typed Polyaxon inputs.

version: 1.1
kind: component
name: unsloth-finetune
tags: ["unsloth", "llm", "finetuning"]
inputs:
- name: model_name
  type: str
  isOptional: true
  value: unsloth/Llama-3.2-1B-Instruct
- name: dataset_name
  type: str
  isOptional: true
  value: yahma/alpaca-cleaned
- name: dataset_split
  type: str
  isOptional: true
  value: train
- name: dataset_text_field
  type: str
  isOptional: true
  value: text
- name: max_seq_length
  type: int
  isOptional: true
  value: 2048
- name: max_steps
  type: int
  isOptional: true
  value: 60
- name: learning_rate
  type: float
  isOptional: true
  value: 0.0002
- name: save_merged
  type: bool
  isOptional: true
  value: false
run:
  kind: job
  init:
  - file:
      filename: train_unsloth.py
      content: |
        import argparse

        from unsloth import FastLanguageModel, is_bfloat16_supported
        from datasets import load_dataset
        from transformers import TrainingArguments
        from trl import SFTTrainer

        from polyaxon import tracking
        from polyaxon.tracking.contrib.hugging_face import PolyaxonCallback


        def parse_args():
            parser = argparse.ArgumentParser()
            parser.add_argument("--model_name", type=str, required=True)
            parser.add_argument("--dataset_name", type=str, required=True)
            parser.add_argument("--dataset_split", type=str, default="train")
            parser.add_argument("--dataset_text_field", type=str, default="text")
            parser.add_argument("--max_seq_length", type=int, default=2048)
            parser.add_argument("--max_steps", type=int, default=60)
            parser.add_argument("--learning_rate", type=float, default=2e-4)
            parser.add_argument("--save_merged", action="store_true")
            return parser.parse_args()


        def main():
            args = parse_args()

            tracking.init()
            tracking.log_inputs(
                model_name=args.model_name,
                dataset_name=args.dataset_name,
                dataset_split=args.dataset_split,
                dataset_text_field=args.dataset_text_field,
                max_seq_length=args.max_seq_length,
                max_steps=args.max_steps,
                learning_rate=args.learning_rate,
            )

            model, tokenizer = FastLanguageModel.from_pretrained(
                model_name=args.model_name,
                max_seq_length=args.max_seq_length,
                dtype=None,
                load_in_4bit=True,
            )

            model = FastLanguageModel.get_peft_model(
                model,
                r=16,
                target_modules=[
                    "q_proj",
                    "k_proj",
                    "v_proj",
                    "o_proj",
                    "gate_proj",
                    "up_proj",
                    "down_proj",
                ],
                lora_alpha=16,
                lora_dropout=0,
                bias="none",
                use_gradient_checkpointing="unsloth",
                random_state=3407,
            )

            dataset = load_dataset(args.dataset_name, split=args.dataset_split)

            training_args = TrainingArguments(
                per_device_train_batch_size=2,
                gradient_accumulation_steps=4,
                warmup_steps=5,
                max_steps=args.max_steps,
                learning_rate=args.learning_rate,
                fp16=not is_bfloat16_supported(),
                bf16=is_bfloat16_supported(),
                logging_steps=1,
                optim="adamw_8bit",
                weight_decay=0.01,
                lr_scheduler_type="linear",
                seed=3407,
                output_dir=tracking.get_outputs_path("trainer"),
                report_to="none",
            )

            trainer = SFTTrainer(
                model=model,
                tokenizer=tokenizer,
                train_dataset=dataset,
                dataset_text_field=args.dataset_text_field,
                max_seq_length=args.max_seq_length,
                packing=False,
                args=training_args,
                callbacks=[PolyaxonCallback()],
            )
            trainer.train()

            adapter_dir = tracking.get_outputs_path("lora-adapter")
            model.save_pretrained(adapter_dir)
            tokenizer.save_pretrained(adapter_dir)
            tracking.log_model_ref(
                path=adapter_dir,
                name="lora-adapter",
                framework="unsloth",
            )

            if args.save_merged:
                merged_dir = tracking.get_outputs_path("merged-model")
                model.save_pretrained_merged(
                    merged_dir,
                    tokenizer,
                    save_method="merged_16bit",
                )
                tracking.log_model_ref(
                    path=merged_dir,
                    name="merged-model",
                    framework="unsloth",
                )


        if __name__ == "__main__":
            main()
  container:
    image: your-registry/unsloth:cuda
    workingDir: "{{ globals.artifacts_path }}"
    command: [python, -u, train_unsloth.py]
    args:
    - "{{ params.model_name.as_arg }}"
    - "{{ params.dataset_name.as_arg }}"
    - "{{ params.dataset_split.as_arg }}"
    - "{{ params.dataset_text_field.as_arg }}"
    - "{{ params.max_seq_length.as_arg }}"
    - "{{ params.max_steps.as_arg }}"
    - "{{ params.learning_rate.as_arg }}"
    - "{{ params.save_merged.as_arg }}"
    resources:
      limits:
        nvidia.com/gpu: "1"

Replace your-registry/unsloth:cuda with the image you build for your cluster. If you need private Hugging Face models or datasets, inject the token through a Polyaxon connection or Kubernetes secret; do not hard-code it in the component.

Run the job

Save the component as unsloth-finetune.yaml and run it with the defaults:

polyaxon run -f unsloth-finetune.yaml

Override the model, dataset, or training budget with -P:

polyaxon run -f unsloth-finetune.yaml \
  -P model_name=unsloth/Llama-3.2-3B-Instruct \
  -P dataset_name=your-org/your-dataset \
  -P max_steps=500

To also export a merged model:

polyaxon run -f unsloth-finetune.yaml -P save_merged=true

Operational notes

Use one GPU for small QLoRA jobs and increase resources only after measuring memory headroom. Larger base models, longer context, and merged exports can change storage and memory requirements quickly.

For repeated jobs, mount or configure a Hugging Face cache through your cluster storage. Re-downloading base models on every run wastes startup time and stresses shared network paths.

Keep LoRA adapters as the default logged output. Export merged checkpoints only for serving, registry, or handoff workflows that require a standalone model.

Polyaxon & Unsloth Fine-tuning

Container image

Polyaxon component

Run the job

Operational notes

On this page