Adding Training image needed for train api

.github/workflows/publish-core-images.yaml

@@ -10,8 +10,9 @@ jobs:
     uses: ./.github/workflows/build-and-publish-images.yaml
     with:
       component-name: ${{ matrix.component-name }}
-      platforms: linux/amd64,linux/arm64,linux/ppc64le
+      platforms: ${{ matrix.platforms }}
       dockerfile: ${{ matrix.dockerfile }}
+      context: ${{ matrix.context }}
     secrets:
       DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }}
       DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
@@ -22,8 +23,15 @@ jobs:
         include:
           - component-name: training-operator
             dockerfile: build/images/training-operator/Dockerfile
+            platforms: linux/amd64,linux/arm64,linux/ppc64le
           - component-name: kubectl-delivery
             dockerfile: build/images/kubectl-delivery/Dockerfile
+            platforms: linux/amd64,linux/arm64,linux/ppc64le
           - component-name: storage-initializer
-             dockerfile: sdk/python/kubeflow/storage_initializer/Dockerfile
-             context: sdk/python/kubeflow/storage_initializer
+            dockerfile: sdk/python/kubeflow/storage_initializer/Dockerfile
+            context: sdk/python/kubeflow/storage_initializer
+            platforms: linux/amd64,linux/arm64
+          - component-name: trainer-huggingface
+            dockerfile: sdk/python/kubeflow/trainer/hf_dockerfile
+            context: sdk/python/kubeflow/trainer
+            platforms: linux/amd64,linux/arm64

.github/workflows/publish-example-images.yaml

@@ -52,7 +52,6 @@ jobs:
           - component-name: mxnet-auto-tuning
             dockerfile: examples/mxnet/tune/Dockerfile
             context: examples/mxnet/tune
-
 # TODO (tenzen-y): Fix the below broken Dockerfiles
 #          - component-name: pytorch-dist-mnist-mpi
 #            dockerfile: examples/pytorch/mnist/Dockerfile-mpi

examples/sdk/train_api.py

@@ -0,0 +1,57 @@
+from kubeflow.training.api.training_client import TrainingClient
+from kubeflow.storage_initializer.hugging_face import (
+    HuggingFaceModelParams,
+    HuggingFaceTrainParams,
+    HfDatasetParams,
+)
+from peft import LoraConfig
+import transformers
+from transformers import TrainingArguments
+
+client = TrainingClient()
+
+client.train(
+    name="hf-test",
+    num_workers=2,
+    num_procs_per_worker=0,
+    model_provider_parameters=HuggingFaceModelParams(
+        model_uri="hf://Jedalc/codeparrot-gp2-finetune",
+        transformer_type=transformers.AutoModelForCausalLM,
+    ),
+    dataset_provider_parameters=HfDatasetParams(repo_id="imdatta0/ultrachat_10k"),
+    train_parameters=HuggingFaceTrainParams(
+        lora_config=LoraConfig(
+            r=8,
+            lora_alpha=8,
+            target_modules=["c_attn", "c_proj", "w1", "w2"],
+            layers_to_transform=list(range(30, 40)),
+            # layers_pattern=['lm_head'],
+            lora_dropout=0.1,
+            bias="none",
+            task_type="CAUSAL_LM",
+        ),
+        training_parameters=TrainingArguments(
+            num_train_epochs=2,
+            per_device_train_batch_size=1,
+            gradient_accumulation_steps=1,
+            gradient_checkpointing=True,
+            warmup_steps=0.01,
+            # max_steps=50, #20,
+            learning_rate=1,
+            lr_scheduler_type="cosine",
+            bf16=False,
+            logging_steps=0.01,
+            output_dir="",
+            optim=f"paged_adamw_32bit",
+            save_steps=0.01,
+            save_total_limit=3,
+            disable_tqdm=False,
+            resume_from_checkpoint=True,
+            remove_unused_columns=True,
+            evaluation_strategy="steps",
+            eval_steps=0.01,
+            per_device_eval_batch_size=1,
+        ),
+    ),
+    resources_per_worker={"gpu": 0, "cpu": 8, "memory": "8Gi"},
+)

sdk/python/kubeflow/trainer/hf_dockerfile

@@ -0,0 +1,18 @@
+# Use an official Pytorch runtime as a parent image
+FROM nvcr.io/nvidia/pytorch:23.12-py3
+
+ # Set the working directory in the container
+ WORKDIR /app
+
+ # Copy the Python package and its source code into the container
+ COPY . /app
+
+ # Copy the requirements.txt file into the container
+ COPY requirements.txt /app/requirements.txt
+
+ # Install any needed packages specified in requirements.txt
+ RUN pip install --no-cache-dir -r requirements.txt
+
+ # Run storage.py when the container launches
+ ENTRYPOINT ["torchrun", "hf_llm_training.py"]
+

sdk/python/kubeflow/trainer/hf_llm_training.py

@@ -0,0 +1,118 @@
+import argparse
+import transformers
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    AutoConfig,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+    Trainer,
+)
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model
+from urllib.parse import urlparse
+import os
+import json
+
+
+def setup_model_and_tokenizer(model_uri, transformer_type, model_dir):
+    # Set up the model and tokenizer
+
+    parsed_uri = urlparse(model_uri)
+    model_name = parsed_uri.netloc + parsed_uri.path
+    transformer_type_class = getattr(transformers, transformer_type)
+
+    model = transformer_type_class.from_pretrained(
+        pretrained_model_name_or_path=model_name,
+        cache_dir=model_dir,
+        local_files_only=True,
+        device_map="auto",
+    )
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path=model_name,
+        cache_dir=model_dir,
+        local_files_only=True,
+        device_map="auto",
+    )
+
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.add_pad_token = True
+
+    # Freeze model parameters
+    for param in model.parameters():
+        param.requires_grad = False
+
+    return model, tokenizer
+
+
+def load_and_preprocess_data(dataset_name, dataset_dir):
+    # Load and preprocess the dataset
+    print("loading dataset")
+    dataset = load_dataset(dataset_name, cache_dir=dataset_dir)
+    train_data = dataset["train"]
+
+    try:
+        eval_data = dataset["eval"]
+    except Exception as err:
+        eval_data = None
+
+    return train_data, eval_data
+
+
+def setup_peft_model(model, lora_config):
+    # Set up the PEFT model
+    lora_config = LoraConfig(**json.loads(lora_config))
+    print(lora_config)
+    model = get_peft_model(model, lora_config)
+    return model
+
+
+def train_model(model, train_data, eval_data, tokenizer, train_params):
+    # Train the model
+    trainer = Trainer(
+        model=model,
+        train_dataset=train_data,
+        eval_dataset=eval_data,
+        tokenizer=tokenizer,
+        args=TrainingArguments(
+            **train_params,
+            data_collator=DataCollatorForLanguageModeling(
+                tokenizer, pad_to_multiple_of=8, return_tensors="pt", mlm=False
+            )
+        ),
+    )
+
+    trainer.train()
+    print("training done")
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="Script for training a model with PEFT configuration."
+    )
+
+    parser.add_argument("--model_uri", help="model uri")
+    parser.add_argument("--transformer_type", help="model transformer type")
+    parser.add_argument("--model_dir", help="directory containing model")
+    parser.add_argument("--dataset_dir", help="directory contaning dataset")
+    parser.add_argument("--dataset_name", help="dataset name")
+    parser.add_argument("--lora_config", help="lora_config")
+    parser.add_argument(
+        "--training_parameters", help="hugging face training parameters"
+    )
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_arguments()
+    model, tokenizer = setup_model_and_tokenizer(
+        args.model_uri, args.transformer_type, args.model_dir
+    )
+    train_data, eval_data = load_and_preprocess_data(
+        args.dataset_name, args.dataset_dir
+    )
+    model = setup_peft_model(model, args.lora_config)
+    train_model(model, train_data, eval_data, tokenizer, args.training_parameters)

sdk/python/kubeflow/trainer/requirements.txt

@@ -0,0 +1,3 @@
+peft==0.7.0
+datasets==2.15.0
+transformers==4.35.2