trainer/examples/pytorch/question-answering/fine-tune-distilbert.ipynb

Fine-tuning DistilBERT for question answering

This guide describes fine-tuning DistilBERT model with Stanford Question Answering Dataset (SQuAD) for question-answering using Kubeflow Trainer.

This guide is adapted from HuggingFace question answering task recipe page: https://huggingface.co/docs/transformers/en/tasks/question_answering

Pretrained DistilBERT: https://huggingface.co/docs/transformers/en/model_doc/distilbert

SQuAD dataset: https://huggingface.co/datasets/rajpurkar/squad

Install the Kubeflow SDK

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

In [ ]:

# !pip install -U kubeflow

Install dependencies

In [1]:

!pip install "cloudpathlib[all]" "transformers[torch]"

Requirement already satisfied: cloudpathlib[all] in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (0.21.0)
Requirement already satisfied: transformers[torch] in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (4.51.3)
Requirement already satisfied: filelock in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (3.18.0)
Requirement already satisfied: huggingface-hub<1.0,>=0.30.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (0.30.2)
Requirement already satisfied: numpy>=1.17 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (2.2.5)
Requirement already satisfied: packaging>=20.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (25.0)
Requirement already satisfied: pyyaml>=5.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (6.0.2)
Requirement already satisfied: regex!=2019.12.17 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (2024.11.6)
Requirement already satisfied: requests in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (2.32.3)
Requirement already satisfied: tokenizers<0.22,>=0.21 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (0.21.1)
Requirement already satisfied: safetensors>=0.4.3 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (0.5.3)
Requirement already satisfied: tqdm>=4.27 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (4.67.1)
Requirement already satisfied: torch>=2.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (2.7.0)
Requirement already satisfied: accelerate>=0.26.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from transformers[torch]) (1.6.0)
Requirement already satisfied: psutil in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from accelerate>=0.26.0->transformers[torch]) (7.0.0)
Requirement already satisfied: fsspec>=2023.5.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from huggingface-hub<1.0,>=0.30.0->transformers[torch]) (2025.3.2)
Requirement already satisfied: typing-extensions>=3.7.4.3 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from huggingface-hub<1.0,>=0.30.0->transformers[torch]) (4.13.2)
Requirement already satisfied: setuptools in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from torch>=2.0->transformers[torch]) (80.0.0)
Requirement already satisfied: sympy>=1.13.3 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from torch>=2.0->transformers[torch]) (1.14.0)
Requirement already satisfied: networkx in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from torch>=2.0->transformers[torch]) (3.4.2)
Requirement already satisfied: jinja2 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from torch>=2.0->transformers[torch]) (3.1.6)
Requirement already satisfied: azure-storage-blob>=12 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from cloudpathlib[all]) (12.25.1)
Requirement already satisfied: azure-storage-file-datalake>=12 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from cloudpathlib[all]) (12.20.0)
Requirement already satisfied: google-cloud-storage in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from cloudpathlib[all]) (3.1.0)
Requirement already satisfied: boto3>=1.34.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from cloudpathlib[all]) (1.38.3)
Requirement already satisfied: charset-normalizer<4,>=2 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from requests->transformers[torch]) (3.4.1)
Requirement already satisfied: idna<4,>=2.5 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from requests->transformers[torch]) (3.10)
Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from requests->transformers[torch]) (2.4.0)
Requirement already satisfied: certifi>=2017.4.17 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from requests->transformers[torch]) (2025.4.26)
Requirement already satisfied: azure-core>=1.30.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from azure-storage-blob>=12->cloudpathlib[all]) (1.33.0)
Requirement already satisfied: cryptography>=2.1.4 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from azure-storage-blob>=12->cloudpathlib[all]) (44.0.2)
Requirement already satisfied: isodate>=0.6.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from azure-storage-blob>=12->cloudpathlib[all]) (0.7.2)
Requirement already satisfied: botocore<1.39.0,>=1.38.3 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from boto3>=1.34.0->cloudpathlib[all]) (1.38.3)
Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from boto3>=1.34.0->cloudpathlib[all]) (1.0.1)
Requirement already satisfied: s3transfer<0.13.0,>=0.12.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from boto3>=1.34.0->cloudpathlib[all]) (0.12.0)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from sympy>=1.13.3->torch>=2.0->transformers[torch]) (1.3.0)
Requirement already satisfied: google-auth<3.0dev,>=2.26.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-cloud-storage->cloudpathlib[all]) (2.39.0)
Requirement already satisfied: google-api-core<3.0.0dev,>=2.15.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-cloud-storage->cloudpathlib[all]) (2.25.0rc0)
Requirement already satisfied: google-cloud-core<3.0dev,>=2.4.2 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-cloud-storage->cloudpathlib[all]) (2.4.3)
Requirement already satisfied: google-resumable-media>=2.7.2 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-cloud-storage->cloudpathlib[all]) (2.7.2)
Requirement already satisfied: google-crc32c<2.0dev,>=1.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-cloud-storage->cloudpathlib[all]) (1.7.1)
Requirement already satisfied: MarkupSafe>=2.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from jinja2->torch>=2.0->transformers[torch]) (3.0.2)
Requirement already satisfied: six>=1.11.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from azure-core>=1.30.0->azure-storage-blob>=12->cloudpathlib[all]) (1.17.0)
Requirement already satisfied: python-dateutil<3.0.0,>=2.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from botocore<1.39.0,>=1.38.3->boto3>=1.34.0->cloudpathlib[all]) (2.9.0.post0)
Requirement already satisfied: cffi>=1.12 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from cryptography>=2.1.4->azure-storage-blob>=12->cloudpathlib[all]) (1.17.1)
Requirement already satisfied: googleapis-common-protos<2.0.0,>=1.56.2 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-api-core<3.0.0dev,>=2.15.0->google-cloud-storage->cloudpathlib[all]) (1.70.0)
Requirement already satisfied: protobuf!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<7.0.0,>=3.19.5 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-api-core<3.0.0dev,>=2.15.0->google-cloud-storage->cloudpathlib[all]) (6.30.2)
Requirement already satisfied: proto-plus<2.0.0,>=1.22.3 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-api-core<3.0.0dev,>=2.15.0->google-cloud-storage->cloudpathlib[all]) (1.26.1)
Requirement already satisfied: cachetools<6.0,>=2.0.0 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-auth<3.0dev,>=2.26.1->google-cloud-storage->cloudpathlib[all]) (5.5.2)
Requirement already satisfied: pyasn1-modules>=0.2.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-auth<3.0dev,>=2.26.1->google-cloud-storage->cloudpathlib[all]) (0.4.2)
Requirement already satisfied: rsa<5,>=3.1.4 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from google-auth<3.0dev,>=2.26.1->google-cloud-storage->cloudpathlib[all]) (4.9.1)
Requirement already satisfied: pycparser in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from cffi>=1.12->cryptography>=2.1.4->azure-storage-blob>=12->cloudpathlib[all]) (2.22)
Requirement already satisfied: pyasn1<0.7.0,>=0.6.1 in /Users/andrew/git/trainer/.venv/lib/python3.13/site-packages (from pyasn1-modules>=0.2.1->google-auth<3.0dev,>=2.26.1->google-cloud-storage->cloudpathlib[all]) (0.6.1)

[notice] A new release of pip is available: 25.0.1 -> 25.1
[notice] To update, run: pip install --upgrade pip

Define the HuggingFace training script

We need to wrap our training script into a function to create the Kubeflow TrainJob.

In [ ]:

def train_distilbert(model_name: str, bucket: str = None):
    import os

    from cloudpathlib import CloudPath
    from datasets import load_dataset
    import torch
    from transformers import AutoTokenizer, DefaultDataCollator, AutoModelForQuestionAnswering, TrainingArguments, Trainer

    import torch.distributed as dist

    # Initialize distributed environment
    _, backend = ("cuda", "nccl") if torch.cuda.is_available() else ("cpu", "gloo")
    dist.init_process_group(backend=backend)

    local_rank = int(os.getenv("LOCAL_RANK", 0))
    print(
        "Distributed Training with WORLD_SIZE: {}, RANK: {}, LOCAL_RANK: {}.".format(
            dist.get_world_size(),
            dist.get_rank(),
            local_rank,
        )
    )

    # Download the dataset and tokenizer
    squad = load_dataset("rajpurkar/squad", split="train[:100]")    

    squad = squad.train_test_split(test_size=0.2, shuffle=False)
    
    tokenizer = AutoTokenizer.from_pretrained(f'distilbert/{model_name}')
    
    # Define the preprocessing function
    def preprocess_function(examples):
        questions = [q.strip() for q in examples["question"]]
        inputs = tokenizer(
            questions,
            examples["context"],
            max_length=384,
            truncation="only_second",
            return_offsets_mapping=True,
            padding="max_length",
        )
    
        offset_mapping = inputs.pop("offset_mapping")
        answers = examples["answers"]
        start_positions = []
        end_positions = []
    
        for i, offset in enumerate(offset_mapping):
            answer = answers[i]
            start_char = answer["answer_start"][0]
            end_char = answer["answer_start"][0] + len(answer["text"][0])
            sequence_ids = inputs.sequence_ids(i)
    
            # Find the start and end of the context
            idx = 0
            while sequence_ids[idx] != 1:
                idx += 1
            context_start = idx
            while sequence_ids[idx] == 1:
                idx += 1
            context_end = idx - 1
    
            # If the answer is not fully inside the context, label it (0, 0)
            if offset[context_start][0] > end_char or offset[context_end][1] < start_char:
                start_positions.append(0)
                end_positions.append(0)
            else:
                # Otherwise it's the start and end token positions
                idx = context_start
                while idx <= context_end and offset[idx][0] <= start_char:
                    idx += 1
                start_positions.append(idx - 1)
    
                idx = context_end
                while idx >= context_start and offset[idx][1] >= end_char:
                    idx -= 1
                end_positions.append(idx + 1)
    
        inputs["start_positions"] = start_positions
        inputs["end_positions"] = end_positions
        return inputs
        
    # Apply the preprocessing function to the dataset
    tokenized_squad = squad.map(preprocess_function, batched=True, remove_columns=squad["train"].column_names)
        
    # Create a batch of examples using DefaultDataCollator
    data_collator = DefaultDataCollator()

    # Load the model
    model = AutoModelForQuestionAnswering.from_pretrained(f'distilbert/{model_name}')

    # Define training hyperparameters
    training_args = TrainingArguments(
        output_dir=model_name,
        eval_strategy="epoch",
        learning_rate=2e-5,
        per_device_train_batch_size=1,
        per_device_eval_batch_size=1,
        num_train_epochs=1,
        weight_decay=0.01,
        push_to_hub=False,
    )
    
    # Prepare trainer with configuration
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_squad["train"],
        eval_dataset=tokenized_squad["test"],
        processing_class=tokenizer,
        data_collator=data_collator,
    )
    
    trainer.train()

    # Upload the fine-tuned model
    if bucket:
        (CloudPath(bucket) / model_name).upload_from(model_name)

In [2]:

from kubeflow.trainer import TrainerClient, CustomTrainer

for r in TrainerClient().list_runtimes():
    print(f"Name: {r.name}, Framework: {r.trainer.framework}, Trainer Type: {r.trainer.trainer_type.value}")

Runtime mpi-distributed must have trainer.kubeflow.org/framework label.

Name: deepspeed-distributed, Framework: deepspeed, Trainer Type: CustomTrainer
Name: mlx-distributed, Framework: mlx, Trainer Type: CustomTrainer
Name: torch-distributed, Framework: torch, Trainer Type: CustomTrainer
Name: torchtune-llama3.2-1b, Framework: torchtune, Trainer Type: BuiltinTrainer
Name: torchtune-llama3.2-3b, Framework: torchtune, Trainer Type: BuiltinTrainer

In [ ]:

# To upload to object storage (S3, GCS or Azure Blob Storage), set the bucket with protocol, e.g., "s3://my-bucket/folder"
BUCKET = None

MODEL_NAME = "distilbert-base-uncased"
args = {
    "model_name": MODEL_NAME,
    "bucket": BUCKET,
}

job_id = TrainerClient().train(
    trainer=CustomTrainer(
        func=train_distilbert,
        func_args=args,
        num_nodes=1,
        packages_to_install=["datasets", "transformers[torch]", "cloudpathlib[all]"],
        resources_per_node={
            "cpu": "2",
            "memory": "12Gi",
            # Uncomment this to distribute the TrainJob using GPU nodes
            # "nvidia.com/gpu": 1,
        },
    ),
)

In [5]:

# Train API generates a random TrainJob id.
job_id

Out [5]:

'rafd89de924b'

Check the TrainJob details

Use list_jobs() and get_job() APIs to get details about the created TrainJob and its steps.

In [6]:

for job in TrainerClient().list_jobs():
    print(f"TrainJob: {job.name}, Status: {job.status}, Created at: {job.creation_timestamp}")

TrainJob: rafd89de924b, Status: Unknown, Created at: 2025-04-29 01:22:14+00:00

In [ ]:

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

In [8]:

# We execute mpirun command on node-0, which functions as the MPI Launcher node.
for c in TrainerClient().get_job(name=job_id).steps:
    print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}")

Step: node-0, Status: Running, Devices: cpu x 2

Show the TrainJob logs

Use get_job_logs() API to retrieve the TrainJob logs.

In [ ]:

for logline in TrainerClient().get_job_logs(job_id, follow=True):
    print(logline)

[node-0]: WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager, possibly rendering your system unusable.It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv. Use the --root-user-action option if you know what you are doing and want to suppress this warning.
[node-0]: W0429 01:22:31.907000 1 site-packages/torch/distributed/run.py:793] 
[node-0]: W0429 01:22:31.907000 1 site-packages/torch/distributed/run.py:793] *****************************************
[node-0]: W0429 01:22:31.907000 1 site-packages/torch/distributed/run.py:793] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. 
[node-0]: W0429 01:22:31.907000 1 site-packages/torch/distributed/run.py:793] *****************************************
[node-0]: Distributed Training with WORLD_SIZE: 2, RANK: 1, LOCAL_RANK: 1.
[node-0]: Distributed Training with WORLD_SIZE: 2, RANK: 0, LOCAL_RANK: 0.
Generating train split: 100%|██████████| 87599/87599 [00:00<00:00, 627416.25 examples/s]
Generating validation split: 100%|██████████| 10570/10570 [00:00<00:00, 1041555.11 examples/s]
Map: 100%|██████████| 80/80 [00:00<00:00, 1361.10 examples/s]
Map: 100%|██████████| 20/20 [00:00<00:00, 2123.27 examples/s]
Map: 100%|██████████| 80/80 [00:00<00:00, 1030.86 examples/s]
Map: 100%|██████████| 20/20 [00:00<00:00, 1877.53 examples/s]
[node-0]: Xet Storage is enabled for this repo, but the 'hf_xet' package is not installed. Falling back to regular HTTP download. For better performance, install the package with: `pip install huggingface_hub[hf_xet]` or `pip install hf_xet`
[node-0]: Some weights of DistilBertForQuestionAnswering were not initialized from the model checkpoint at distilbert/distilbert-base-uncased and are newly initialized: ['qa_outputs.bias', 'qa_outputs.weight']
[node-0]: You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
[node-0]: Some weights of DistilBertForQuestionAnswering were not initialized from the model checkpoint at distilbert/distilbert-base-uncased and are newly initialized: ['qa_outputs.bias', 'qa_outputs.weight']
[node-0]: You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
  0%|          | 0/40 [00:00<?, ?it/s][rank0]:[W429 01:22:58.895439547 reducer.cpp:1400] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration,  which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())
[node-0]: [rank1]:[W429 01:22:58.895689005 reducer.cpp:1400] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration,  which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())
100%|██████████| 40/40 [02:36<00:00,  4.10s/it]
  0%|          | 0/10 [00:00<?, ?it/s]
 20%|██        | 2/10 [00:00<00:03,  2.54it/s]
 30%|███       | 3/10 [00:01<00:03,  1.80it/s]
 40%|████      | 4/10 [00:02<00:03,  1.55it/s]
 50%|█████     | 5/10 [00:03<00:03,  1.43it/s]
 60%|██████    | 6/10 [00:03<00:02,  1.37it/s]
 70%|███████   | 7/10 [00:04<00:02,  1.32it/s]
 80%|████████  | 8/10 [00:05<00:01,  1.30it/s]
 90%|█████████ | 9/10 [00:06<00:00,  1.29it/s]
                                               
[node-0]: {'eval_loss': 5.543211936950684, 'eval_runtime': 7.9713, 'eval_samples_per_second': 2.509, 'eval_steps_per_second': 1.254, 'epoch': 1.0}
100%|██████████| 40/40 [02:45<00:00,  4.10s/it]
100%|██████████| 10/10 [00:07<00:00,  1.28it/s]
[node-0]: {'train_runtime': 165.12, 'train_samples_per_second': 0.484, 'train_steps_per_second': 0.242, 'train_loss': 5.764264678955078, 'epoch': 1.0}
100%|██████████| 40/40 [02:45<00:00,  4.13s/it]

Verify TrainJob Completion

Check the TrainJob status to ensure it completed successfully.

In [ ]:

TrainerClient().wait_for_job_status(name=job_id, timeout=20)

Inference

Download the model and run inference on some examples.

In [10]:

from cloudpathlib import CloudPath
from transformers import pipeline

if BUCKET:
    (CloudPath(BUCKET) / MODEL_NAME).download_to(MODEL_NAME)

    question = "How many programming languages does BLOOM support?"
    context = "BLOOM has 176 billion parameters and can generate text in 46 languages natural languages and 13 programming languages."

    question_answerer = pipeline("question-answering", model=f"./{MODEL_NAME}/checkpoint-375")
    question_answerer(question=question, context=context)

Clean up

To delete the TrainJob you can use the delete_job() API and pass the generated job_id.

In [11]:

# _ = TrainerClient().delete_job(job_id)