trainer/examples/torchtune/llama3_2/alpaca-trainjob-yaml.ipynb

Fine-tune Llama-3.2-1B-Instruct with Alpaca Dataset

This example demonstrates how to fine-tune Llama-3.2-1B-Instruct model with the Alpaca Dataset using TorchTune BuiltinTrainer from Kubeflow Trainer SDK.

This notebooks walks you through the prerequisites of using TorchTune BuiltinTrainer from Kubeflow Trainer SDK, and how to submit TrainJob to bootstrap the fine-tuning workflow.

Llama-3.2-1B-Instruct: https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct

Alpaca Dataset: https://huggingface.co/datasets/tatsu-lab/alpaca

Install the Kubeflow SDK

You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:

In [ ]:

# !pip install -U kubeflow

Prerequisites

Install Official Training Runtimes

You need to make sure that you've installed the Kubeflow Trainer Controller Manager and Kubeflow Training Runtimes mentioned in the installation guide.

In [1]:

# List all available Kubeflow Training Runtimes.
from kubeflow.trainer import *
from kubeflow_trainer_api import models
import os

client = TrainerClient()
for runtime in client.list_runtimes():
    print(runtime)

Runtime(name='deepspeed-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='deepspeed', num_nodes=1, accelerator_count=4), pretrained_model=None)
Runtime(name='mlx-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='mlx', num_nodes=1, accelerator_count=1), pretrained_model=None)
Runtime(name='torch-distributed', trainer=RuntimeTrainer(trainer_type=<TrainerType.CUSTOM_TRAINER: 'CustomTrainer'>, framework='torch', num_nodes=1, accelerator_count='Unknown'), pretrained_model=None)
Runtime(name='torchtune-llama3.2-1b', trainer=RuntimeTrainer(trainer_type=<TrainerType.BUILTIN_TRAINER: 'BuiltinTrainer'>, framework='torchtune', num_nodes=1, accelerator_count='2.0'), pretrained_model=None)
Runtime(name='torchtune-llama3.2-3b', trainer=RuntimeTrainer(trainer_type=<TrainerType.BUILTIN_TRAINER: 'BuiltinTrainer'>, framework='torchtune', num_nodes=1, accelerator_count='2.0'), pretrained_model=None)

Bootstrap LLM Fine-tuning Workflow

Kubeflow TrainJob will train the model in the referenced (Cluster)TrainingRuntime.

In [ ]:

job_name = client.train(
    runtime="torchtune-llama3.2-1b",
    initializer=Initializer(
        dataset=HuggingFaceDatasetInitializer(
            storage_uri="hf://tatsu-lab/alpaca/data"
        ),
        model=HuggingFaceModelInitializer(
            storage_uri="hf://meta-llama/Llama-3.2-1B-Instruct",
            access_token=os.environ["HF_TOKEN"] # Replace with your Hugging Face token,
        )
    ),
    trainer=BuiltinTrainer(
        config=TorchTuneConfig(
            dataset_preprocess_config=TorchTuneInstructDataset(
                source=DataFormat.PARQUET, split="train[:1000]"
            ),
            resources_per_node={
                "memory": "200G",
                "gpu": 1,
            },
            
        )
    )
)

Wait for running status

In [ ]:

# Wait for the running status.
client.wait_for_job_status(name=job_name, status={"Running"})

Watch the TrainJob Logs

We can use the get_job_logs() API to get the TrainJob logs.

Dataset Initializer

In [ ]:

from kubeflow.trainer.constants import constants

for line in client.get_job_logs(job_name, follow=True, step=constants.DATASET_INITIALIZER):
    print(line)

2025-07-02T08:24:01Z INFO     [__main__.py:16] Starting dataset initialization
2025-07-02T08:24:01Z INFO     [huggingface.py:28] Downloading dataset: tatsu-lab/alpaca
2025-07-02T08:24:01Z INFO     [huggingface.py:29] ----------------------------------------
Fetching 3 files: 100%|██████████| 3/3 [00:01<00:00,  1.82it/s]
2025-07-02T08:24:04Z INFO     [huggingface.py:40] Dataset has been downloaded

Model Initializer

In [ ]:

for line in client.get_job_logs(job_name, follow=True, step=constants.MODEL_INITIALIZER):
    print(line)

2025-07-02T08:24:23Z INFO     [__main__.py:16] Starting pre-trained model initialization
2025-07-02T08:24:23Z INFO     [huggingface.py:26] Downloading model: meta-llama/Llama-3.2-1B-Instruct
2025-07-02T08:24:23Z INFO     [huggingface.py:27] ----------------------------------------
Fetching 8 files: 100%|██████████| 8/8 [01:02<00:00,  7.87s/it]
2025-07-02T08:25:27Z INFO     [huggingface.py:43] Model has been downloaded

Trainer Node

In [ ]:

for c in client.get_job(name=job_name).steps:
    print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\n")

for line in client.get_job_logs(job_name, follow=True):
    print(line)

INFO:torchtune.utils._logging:Running FullFinetuneRecipeDistributed with resolved config:

batch_size: 4
checkpointer:
  _component_: torchtune.training.FullModelHFCheckpointer
  checkpoint_dir: /workspace/model
  checkpoint_files:
  - model.safetensors
  model_type: LLAMA3_2
  output_dir: /workspace/output
  recipe_checkpoint: null
clip_grad_norm: null
compile: false
dataset:
  _component_: torchtune.datasets.instruct_dataset
  data_dir: /workspace/dataset/data
  packed: false
  source: parquet
device: cuda
dtype: bf16
enable_activation_checkpointing: false
enable_activation_offloading: false
epochs: 1
gradient_accumulation_steps: 8
log_every_n_steps: 1
log_peak_memory_stats: true
loss:
  _component_: torchtune.modules.loss.CEWithChunkedOutputLoss
max_steps_per_epoch: null
metric_logger:
  _component_: torchtune.training.metric_logging.DiskLogger
  log_dir: /workspace/output/logs
model:
  _component_: torchtune.models.llama3_2.llama3_2_1b
optimizer:
  _component_: torch.optim.AdamW
  fused: true
  lr: 2.0e-05
optimizer_in_bwd: false
output_dir: /workspace/output
profiler:
  _component_: torchtune.training.setup_torch_profiler
  active_steps: 2
  cpu: true
  cuda: true
  enabled: false
  num_cycles: 1
  output_dir: /workspace/output/profiling_outputs
  profile_memory: false
  record_shapes: true
  wait_steps: 5
  warmup_steps: 3
  with_flops: false
  with_stack: false
resume_from_checkpoint: false
seed: null
shuffle: true
tokenizer:
  _component_: torchtune.models.llama3.llama3_tokenizer
  max_seq_len: null
  path: /workspace/model/original/tokenizer.model

DEBUG:torchtune.utils._logging:Setting manual seed to local seed 3686749453. Local seed is seed + rank = 3686749453 + 0
Writing logs to /workspace/output/logs/log_1751444966.txt
INFO:torchtune.utils._logging:Distributed training is enabled. Instantiating model and loading checkpoint on Rank 0 ...
INFO:torchtune.utils._logging:Instantiating model and loading checkpoint took 17.31 secs
INFO:torchtune.utils._logging:Memory stats after model init:
	GPU peak memory allocation: 2.33 GiB
	GPU peak memory reserved: 2.34 GiB
	GPU peak memory active: 2.33 GiB
/opt/conda/lib/python3.11/site-packages/torch/distributed/distributed_c10d.py:4631: UserWarning: No device id is provided via `init_process_group` or `barrier `. Using the current device set by the user. 
  warnings.warn(  # warn only once
INFO:torchtune.utils._logging:Optimizer is initialized.
INFO:torchtune.utils._logging:Loss is initialized.
Generating train split: 52002 examples [00:00, 192755.58 examples/s]
INFO:torchtune.utils._logging:No learning rate scheduler configured. Using constant learning rate.
WARNING:torchtune.utils._logging: Profiling disabled.
INFO:torchtune.utils._logging: Profiler config after instantiation: {'enabled': False}
1|1625|Loss: 1.5839784145355225: 100%|██████████| 1625/1625 [21:54<00:00,  1.23it/s]INFO:torchtune.utils._logging:Saving checkpoint. This may take some time. Retrieving full model state dict...
INFO:torchtune.utils._logging:Getting full model state dict took 1.69 secs
INFO:torchtune.utils._logging:Model checkpoint of size 2.30 GiB saved to /workspace/output/epoch_0/model-00001-of-00001.safetensors
INFO:torchtune.utils._logging:Saving final epoch checkpoint.
INFO:torchtune.utils._logging:The full model checkpoint, including all weights and configurations, has been saved successfully.You can now use this checkpoint for further training or inference.
INFO:torchtune.utils._logging:Saving checkpoint took 6.64 secs
1|1625|Loss: 1.5839784145355225: 100%|██████████| 1625/1625 [22:01<00:00,  1.23it/s]
Running with torchrun...

Verify TrainJob Completion

Check the TrainJob status to ensure it completed successfully.

In [ ]:

client.wait_for_job_status(name=job_name, timeout=20)

Get the Fine-tuned Model

After Trainer node completes the fine-tuning task, the fine-tuned model will be stored into the /workspace/output directory, which can be shared across Pods through PVC mounting. You can find it in another Pod's /<mountDir>/output directory if you mount the PVC under /<mountDir>.