Repository: kaito-project/kaito

Issue:
Testing coverage for distributed inference
**Is your feature request related to a problem? Please describe.**

Add more test for distributed inference.

**Describe the solution you'd like**

Development workflow for adding a new runtime feature.
Build and push base image to acr -> run preset test against the new base image with distributed inference -> build mcr image -> bump image tag in controller side -> run e2e test. 

- [x] Add test [here](https://github.com/kaito-project/kaito/blob/main/.github/workflows/e2e-preset-test.yml) to test image before pushing to mcr    #1138 
- [x] e2e test against distributed inference
  - [x] ensure adapter works
  - [x] ensure health check works for both leader and worker pods
  - [x] ensure configfile works. (adjust model length and see if it takes effect)

Hints from issue comments:
@chewong: Currently, we have no preset model that uses distributed inference. To test changes in the kaito-base image, we need to create a new preset model with the following requirements:

- `downloadAtRuntime: true`
- Large enough that it cannot fit on a single node, but not too large that it requires a lot of time waiting to finish downloading. Llama 3.3 70B might be a good choice

Additionally, we need to update [`determine_models.py`](https://github.com/kaito-project/kaito/blob/039437a8ccc2158666a421fe55a9d6d8be7f77b2/.github/workflows/kind-cluster/determine_models.py) to output the list of models affected when the base image changes. All models with `downloadAtRuntime: true` will be affected and should be tested during e2e test.

Solution patch:
diff --git a/api/v1alpha1/workspace_validation.go b/api/v1alpha1/workspace_validation.go
index 1910fa5..aaf2083 100644
--- a/api/v1alpha1/workspace_validation.go
+++ b/api/v1alpha1/workspace_validation.go
@@ -209,10 +209,6 @@ func (r *DataSource) validateCreate() (errs *apis.FieldError) {
 	if len(r.URLs) > 0 {
 		sourcesSpecified++
 	}
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-		sourcesSpecified++
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse input image reference: %s", err), "Image"))
@@ -233,9 +229,6 @@ func (r *DataSource) validateUpdate(old *DataSource, isTuning bool) (errs *apis.
 	if isTuning && !reflect.DeepEqual(old.Name, r.Name) {
 		errs = errs.Also(apis.ErrInvalidValue("During tuning Name field cannot be changed once set", "Name"))
 	}
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse input image reference: %s", err), "Image"))
@@ -247,11 +240,6 @@ func (r *DataSource) validateUpdate(old *DataSource, isTuning bool) (errs *apis.
 
 func (r *DataDestination) validateCreate() (errs *apis.FieldError) {
 	destinationsSpecified := 0
-	// TODO: Implement Volumes
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-		destinationsSpecified++
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse output image reference: %s", err), "Image"))
@@ -273,10 +261,6 @@ func (r *DataDestination) validateCreate() (errs *apis.FieldError) {
 }
 
 func (r *DataDestination) validateUpdate() (errs *apis.FieldError) {
-	// TODO: Implement Volumes
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse output image reference: %s", err), "Image"))
diff --git a/api/v1beta1/workspace_validation.go b/api/v1beta1/workspace_validation.go
index 1b0d182..ab530a6 100644
--- a/api/v1beta1/workspace_validation.go
+++ b/api/v1beta1/workspace_validation.go
@@ -209,10 +209,6 @@ func (r *DataSource) validateCreate() (errs *apis.FieldError) {
 	if len(r.URLs) > 0 {
 		sourcesSpecified++
 	}
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-		sourcesSpecified++
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse input image reference: %s", err), "Image"))
@@ -233,9 +229,6 @@ func (r *DataSource) validateUpdate(old *DataSource, isTuning bool) (errs *apis.
 	if isTuning && !reflect.DeepEqual(old.Name, r.Name) {
 		errs = errs.Also(apis.ErrInvalidValue("During tuning Name field cannot be changed once set", "Name"))
 	}
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse input image reference: %s", err), "Image"))
@@ -247,11 +240,6 @@ func (r *DataSource) validateUpdate(old *DataSource, isTuning bool) (errs *apis.
 
 func (r *DataDestination) validateCreate() (errs *apis.FieldError) {
 	destinationsSpecified := 0
-	// TODO: Implement Volumes
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-		destinationsSpecified++
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse output image reference: %s", err), "Image"))
@@ -273,10 +261,6 @@ func (r *DataDestination) validateCreate() (errs *apis.FieldError) {
 }
 
 func (r *DataDestination) validateUpdate() (errs *apis.FieldError) {
-	// TODO: Implement Volumes
-	if r.Volume != nil {
-		errs = errs.Also(apis.ErrInvalidValue("Volume support is not implemented yet", "Volume"))
-	}
 	if image := r.Image; image != "" {
 		if _, err := reference.ParseDockerRef(image); err != nil {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unable to parse output image reference: %s", err), "Image"))
diff --git a/charts/kaito/workspace/templates/lora-params.yaml b/charts/kaito/workspace/templates/lora-params.yaml
index 97f55de..8492381 100644
--- a/charts/kaito/workspace/templates/lora-params.yaml
+++ b/charts/kaito/workspace/templates/lora-params.yaml
@@ -24,7 +24,7 @@ data:
         output_dir: "/mnt/results"
         # num_train_epochs: <Defaults to 3, adjustable>
         ddp_find_unused_parameters: false # Default to false to prevent errors during distributed training
-        save_strategy: "epoch" # Default to save at end of each epoch
+        save_strategy: "no" # Default to save at end of each epoch
         per_device_train_batch_size: 1
     
       DataCollator: # Configurable Parameters: https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/data_collator#transformers.DataCollatorForLanguageModeling
diff --git a/charts/kaito/workspace/templates/qlora-params.yaml b/charts/kaito/workspace/templates/qlora-params.yaml
index a7604dc..8f8f0ac 100644
--- a/charts/kaito/workspace/templates/qlora-params.yaml
+++ b/charts/kaito/workspace/templates/qlora-params.yaml
@@ -27,7 +27,7 @@ data:
         output_dir: "/mnt/results"
         # num_train_epochs: <Defaults to 3, adjustable>
         ddp_find_unused_parameters: false # Default to false to prevent errors during distributed training
-        save_strategy: "epoch" # Default to save at end of each epoch
+        save_strategy: "no" # Default to save at end of each epoch
         per_device_train_batch_size: 1
     
       DataCollator: # Configurable Parameters: https://huggingface.co/docs/transformers/v4.40.2/en/main_classes/data_collator#transformers.DataCollatorForLanguageModeling
diff --git a/docs/img/llama-3.3-70b-instruct.png b/docs/img/llama-3.3-70b-instruct.png
new file mode 100644
index 0000000..9599743
Binary files /dev/null and b/docs/img/llama-3.3-70b-instruct.png differ
diff --git a/docs/proposals/20250529-llama-3.3-70b-instruct.md b/docs/proposals/20250529-llama-3.3-70b-instruct.md
new file mode 100644
index 0000000..fcd2708
--- /dev/null
+++ b/docs/proposals/20250529-llama-3.3-70b-instruct.md
@@ -0,0 +1,87 @@
+---
+title: Proposal for new model support
+authors:
+  - @chewong
+reviewers:
+  - "Kaito contributor"
+creation-date: 2025-05-29
+last-updated: 2025-05-29
+status: provisional
+---
+
+# Title
+
+Add [meta-llama/Llama-3.3-70B-Instruct](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct) to Kaito supported model list
+
+## Glossary
+
+N/A
+
+## Summary
+
+- **Model description**: The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out). The Llama 3.3 instruction tuned text only model is optimized for multilingual dialogue use cases and outperforms many of the available open source and closed chat models on common industry benchmarks.
+- **Model usage statistics**: 700,643 downloads in the last 30 days as of 2025-05-29, according to [Hugging Face](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct).
+- **Model license**: The model is licensed under the [LLAMA 3.3 COMMUNITY LICENSE](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct/blob/main/LICENSE). Therefore, users must agree to the terms of the license and supply Hugging Face credentials to download the model files.
+
+```yaml
+apiVersion: v1
+kind: Secret
+metadata:
+  name: hf-token
+  namespace: default
+data:
+  HF_TOKEN: <base64-encoded-huggingface-token>
+type: Opaque
+```
+
+In the Workspace configuration, you can specify the secret name to use the Hugging Face token:
+
+```yaml
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-llama-3-3-70b-instruct
+resource:
+  instanceType: Standard_ND96isr_H100_v5
+  labelSelector:
+    matchLabels:
+      apps: llama-3-3-70b-instruct
+inference:
+  preset:
+    name: llama-3.3-70b-instruct
+    presetOptions:
+      modelAccessSecret: hf-token
+```
+
+## Requirements
+
+The following table describes the basic model characteristics and the resource requirements of running it.
+
+| Field            | Notes                                                                                                                                                |
+| ---------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------- |
+| Family name      | Llama 3                                                                                                                                              |
+| Type             | `conversational`                                                                                                                                     |
+| Download site    | https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct/tree/main                                                                                   |
+| Version          | [6f6073b423013f6a7d4d9f39144961bfbfbc386b](https://huggingface.co/meta-llama/Llama-3.3-70B-Instruct/commit/6f6073b423013f6a7d4d9f39144961bfbfbc386b) |
+| Storage size     | 140GB                                                                                                                                                |
+| GPU count        | 4                                                                                                                                                    |
+| Total GPU memory | 320GB (~16k context length, more GPUs required for longer context lengths)                                                                           |
+| Per GPU memory   | 80GB                                                                                                                                                 |
+
+![Requirements](../img/llama-3.3-70b-instruct.png)
+
+## Runtimes
+
+This section describes how to configure the runtime framework to support the inference calls.
+
+| Options               | Notes                                                                                                                                                   |
+| --------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| Runtime               | vLLM                                                                                                                                                    |
+| Distributed Inference | True                                                                                                                                                    |
+| Custom configurations | `--max_model_len=16384` by default but users with more GPUs can override it in the Workspace configuration if they want to use a longer context length. |
+
+# History
+
+- [x] 2025-05-29: Open proposal PR.
+- [ ] YYYY-MM-DD: Start model integration.
+- [ ] YYYY-MM-DD: Complete model support.
diff --git a/docs/tuning/README.md b/docs/tuning/README.md
index b976ef2..b444d34 100644
--- a/docs/tuning/README.md
+++ b/docs/tuning/README.md
@@ -2,7 +2,8 @@
 This document presents how to use the Kaito `workspace` Custom Resource Definition (CRD) for parameter-efficient fine-tuning (PEFT) of models, how a Kubernetes job is designed to automate the tuning workflow, and several best practices for troubleshooting.
 
 ## Usage
-Kaito tuning APIs allow users to specify supported tuning methods like [LoRA or QLoRA](https://huggingface.co/docs/peft/main/en/conceptual_guides/lora), the input dataset and configuration settings, and the output destination for saving the tuning results. Currently, Kaito supports both URL and image as the types of tuning input sources. It only supports image as the type of output destination. In the future, Kaito will additionally support the Kubernetes `v1.Volume` API for both the input source and the output destination.
+Kaito tuning APIs allow users to specify supported tuning methods like [LoRA or QLoRA](https://huggingface.co/docs/peft/main/en/conceptual_guides/lora), the input dataset and configuration settings, and the output destination for saving the tuning results. Currently, Kaito supports URL, image 
+and Kubernetes volume as the types of tuning input sources, and image, Kubernetes volume as the types of tuning output destination. 
 
 
 ### Tuning workspace
@@ -35,6 +36,8 @@ tuning:
 
 ```
 
+Example 3: Tuning [`phi-3-mini`](../../examples/fine-tuning/kaito_workspace_tuning_phi_3_with_pvc_volume.yaml). This example shows how to use a Kubernetes volume as the source of input dataset and output destination. We use AzureFile as an example, but any other supported volume type can be used. You should save your input dataset in the volume before creating the workspace, and the output adapter will be saved in the output volume after the tuning job is completed.
+
 The detailed `TuningSpec` API definitions can be found [here](https://github.com/kaito-project/kaito/blob/2ccc93daf9d5385649f3f219ff131ee7c9c47f3e/api/v1alpha1/workspace_types.go#L145).
 
 ### Tuning configurations
@@ -106,7 +109,10 @@ If your dataset is not in one of these formats, it will be passed directly to th
 Note: if you build a container image for the input dataset, please copy the dataset to the **`/data`** directory inside the container.
 
 # Tuning Job
-Kaito uses the Kubernetes **batchv1.job** workload to manage the tuning Pod. When a tuning workspace custom resource is created, the Kaito controller will create a job with the same name as the workspace in the same namespace. To streamline the tuning workflow, Kaito adds two containers in addition to the main container that runs the tuning process. The pod structure is illustrated in Figure 1.
+Kaito uses the Kubernetes **batchv1.job** workload to manage the tuning Pod. When a tuning workspace custom resource is created, the Kaito controller will create a job with the same name as the workspace in the same namespace. The pod structure will be different depending on whether the input and output are specified as URLs, images, or Kubernetes volumes.
+
+## Pod structure when input and output are specified as URLs or images
+To streamline the tuning workflow, Kaito adds two containers in addition to the main container that runs the tuning process. The pod structure is illustrated in Figure 1.
 <div align="left">
   <img src="../img/kaito-fine-tuning.png" width=40% title="Kaito fine tuning" alt="Kaito fine tuning">
 </div>
@@ -120,6 +126,11 @@ Figure 1. Kaito tuning pod structure.
 
 All three containers use shared local volumes (by mounting the same `EmptyDir` volumes), hence file copies between containers are avoided.
 
+## Pod structure when input and output are specified as Kubernetes volumes
+When the input and output are specified as Kubernetes volumes, the initcontainer and sidecar container are removed and only the main container is used. Since the input dataset and ouput destination are already available in the specified volumes, the main container can directly access them. 
+
+Other than the absence of the init and sidecar containers, the main container is the same as described in the previous section. 
+
 # Troubleshooting
 
 ### Job pod failures
diff --git a/examples/fine-tuning/kaito_workspace_tuning_phi_3_with_pvc_volume.yaml b/examples/fine-tuning/kaito_workspace_tuning_phi_3_with_pvc_volume.yaml
new file mode 100644
index 0000000..2ae2b5d
--- /dev/null
+++ b/examples/fine-tuning/kaito_workspace_tuning_phi_3_with_pvc_volume.yaml
@@ -0,0 +1,45 @@
+kind: PersistentVolumeClaim
+apiVersion: v1
+metadata:
+  name: pvc-azurefile-input
+spec:
+  accessModes:
+    - ReadWriteMany
+  resources:
+    requests:
+      storage: 100Gi
+  storageClassName: azurefile-csi
+---
+kind: PersistentVolumeClaim
+apiVersion: v1
+metadata:
+  name: pvc-azurefile-output
+spec:
+  accessModes:
+    - ReadWriteMany
+  resources:
+    requests:
+      storage: 100Gi
+  storageClassName: azurefile-csi
+---
+apiVersion: kaito.sh/v1beta1
+kind: Workspace
+metadata:
+  name: workspace-tuning-phi-3
+resource:
+  instanceType: "Standard_NC6s_v3"
+  labelSelector:
+    matchLabels:
+      app: tuning-phi-3
+tuning:
+  preset:
+    name: phi-3-mini-128k-instruct
+  method: qlora
+  input:
+    volumeSource:
+      persistentVolumeClaim:
+        claimName: pvc-azurefile-input
+  output:
+    volumeSource:
+      persistentVolumeClaim:
+        claimName: pvc-azurefile-output
\ No newline at end of file
diff --git a/pkg/utils/common-preset.go b/pkg/utils/common-preset.go
index f3af3ab..9b76f43 100644
--- a/pkg/utils/common-preset.go
+++ b/pkg/utils/common-preset.go
@@ -15,13 +15,16 @@ const (
 	DefaultAdapterVolumePath  = "/mnt/adapter"
 )
 
-func ConfigResultsVolume(outputPath string) (corev1.Volume, corev1.VolumeMount) {
+func ConfigResultsVolume(outputPath string, outputVolume *corev1.VolumeSource) (corev1.Volume, corev1.VolumeMount) {
 	sharedWorkspaceVolume := corev1.Volume{
 		Name: "results-volume",
 		VolumeSource: corev1.VolumeSource{
 			EmptyDir: &corev1.EmptyDirVolumeSource{},
 		},
 	}
+	if outputVolume != nil {
+		sharedWorkspaceVolume.VolumeSource = *outputVolume
+	}
 	sharedVolumeMount := corev1.VolumeMount{
 		Name:      "results-volume",
 		MountPath: outputPath,
@@ -137,16 +140,12 @@ func ConfigCMVolume(cmName string) (corev1.Volume, corev1.VolumeMount) {
 	return volume, volumeMount
 }
 
-func ConfigDataVolume(hostPath *string) (corev1.Volume, corev1.VolumeMount) {
+func ConfigDataVolume(inputVolumeSource *corev1.VolumeSource) (corev1.Volume, corev1.VolumeMount) {
 	var volume corev1.Volume
 	var volumeMount corev1.VolumeMount
 	var volumeSource corev1.VolumeSource
-	if hostPath != nil {
-		volumeSource = corev1.VolumeSource{
-			HostPath: &corev1.HostPathVolumeSource{
-				Path: *hostPath,
-			},
-		}
+	if inputVolumeSource != nil {
+		volumeSource = *inputVolumeSource
 	} else {
 		volumeSource = corev1.VolumeSource{
 			EmptyDir: &corev1.EmptyDirVolumeSource{},
diff --git a/pkg/workspace/manifests/manifests.go b/pkg/workspace/manifests/manifests.go
index 765742c..b0f35c7 100644
--- a/pkg/workspace/manifests/manifests.go
+++ b/pkg/workspace/manifests/manifests.go
@@ -208,6 +208,11 @@ func GenerateTuningJobManifest(wObj *kaitov1beta1.Workspace, revisionNum string,
 		},
 	}, sidecarContainers...)
 
+	shouldShareProcessNamespace := ptr.To(true)
+	if len(sidecarContainers) == 0 {
+		shouldShareProcessNamespace = ptr.To(false)
+	}
+
 	return &batchv1.Job{
 		TypeMeta: v1.TypeMeta{
 			APIVersion: "batch/v1",
@@ -233,7 +238,7 @@ func GenerateTuningJobManifest(wObj *kaitov1beta1.Workspace, revisionNum string,
 					InitContainers:        initContainers,
 					Containers:            containers,
 					RestartPolicy:         corev1.RestartPolicyNever,
-					ShareProcessNamespace: ptr.To(true),
+					ShareProcessNamespace: shouldShareProcessNamespace,
 					Volumes:               volumes,
 					Tolerations:           tolerations,
 					ImagePullSecrets:      imagePullSecretRefs,
diff --git a/pkg/workspace/tuning/preset-tuning.go b/pkg/workspace/tuning/preset-tuning.go
index 9c5071d..b5a8831 100644
--- a/pkg/workspace/tuning/preset-tuning.go
+++ b/pkg/workspace/tuning/preset-tuning.go
@@ -134,9 +134,9 @@ func GetTrainingOutputDir(ctx context.Context, configMap *corev1.ConfigMap) (str
 }
 
 // SetupTrainingOutputVolume adds shared volume for results dir
-func SetupTrainingOutputVolume(ctx context.Context, configMap *corev1.ConfigMap) (corev1.Volume, corev1.VolumeMount, string) {
+func SetupTrainingOutputVolume(ctx context.Context, configMap *corev1.ConfigMap, outputVolume *corev1.VolumeSource) (corev1.Volume, corev1.VolumeMount, string) {
 	outputDir, _ := GetTrainingOutputDir(ctx, configMap)
-	resultsVolume, resultsVolumeMount := utils.ConfigResultsVolume(outputDir)
+	resultsVolume, resultsVolumeMount := utils.ConfigResultsVolume(outputDir, outputVolume)
 	return resultsVolume, resultsVolumeMount, outputDir
 }
 
@@ -182,35 +182,34 @@ func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspac
 
 	// TODO: make containers only mount the volumes they need
 
-	// Add shared volume for training output
-	trainingOutputVolume, trainingOutputVolumeMount, outputDir := SetupTrainingOutputVolume(ctx, configVolume)
+	// Add volume for training output
+	outputVolume := workspaceObj.Tuning.Output.Volume
+	trainingOutputVolume, trainingOutputVolumeMount, outputDir := SetupTrainingOutputVolume(ctx, configVolume, outputVolume)
 	volumes = append(volumes, trainingOutputVolume)
 	volumeMounts = append(volumeMounts, trainingOutputVolumeMount)
 
 	initContainer, dataSourceVolumes, dataSourceVolumeMounts := prepareDataSource(ctx, workspaceObj)
 	volumes = append(volumes, dataSourceVolumes...)
 	volumeMounts = append(volumeMounts, dataSourceVolumeMounts...)
-	if initContainer.Name != "" {
+	if initContainer != nil && initContainer.Name != "" {
 		initContainers = append(initContainers, *initContainer)
 	}
 
-	sidecarContainer, imagePushSecret, dataDestVolume, dataDestVolumeMount := prepareDataDestination(ctx, workspaceObj, outputDir)
-	volumes = append(volumes, dataDestVolume)
-	volumeMounts = append(volumeMounts, dataDestVolumeMount)
+	sidecarContainer, imagePushSecret, imagePushSecretVolume, imagePushSecretVolumeMount := prepareDataDestination(ctx, workspaceObj, outputDir)
+	volumes = append(volumes, imagePushSecretVolume...)
+	volumeMounts = append(volumeMounts, imagePushSecretVolumeMount...)
 	if sidecarContainer != nil {
-		sidecarContainers = append(sidecarContainers, *sidecarContainer)
+		pauseContainer := corev1.Container{
+			Name:            "pause",
+			Image:           "registry.k8s.io/pause:latest",
+			ImagePullPolicy: corev1.PullAlways,
+		}
+		sidecarContainers = append(sidecarContainers, pauseContainer, *sidecarContainer)
 	}
 	if imagePushSecret != nil {
 		imagePullSecrets = append(imagePullSecrets, *imagePushSecret)
 	}
 
-	pauseContainer := corev1.Container{
-		Name:            "pause",
-		Image:           "registry.k8s.io/pause:latest",
-		ImagePullPolicy: corev1.PullAlways,
-	}
-	sidecarContainers = append(sidecarContainers, pauseContainer)
-
 	modelCommand, err := prepareModelRunParameters(ctx, tuningObj)
 	if err != nil {
 		return nil, err
@@ -246,6 +245,18 @@ func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspac
 		Name:  "PYTORCH_CUDA_ALLOC_CONF",
 		Value: "expandable_segments:True",
 	})
+
+	// remove duplicate volumes
+	seen := make(map[string]bool)
+	var deduplicatedVolumes []corev1.Volume
+	for _, vol := range volumes {
+		if !seen[vol.Name] {
+			deduplicatedVolumes = append(deduplicatedVolumes, vol)
+			seen[vol.Name] = true
+		}
+	}
+	volumes = deduplicatedVolumes
+
 	jobObj := manifests.GenerateTuningJobManifest(workspaceObj, revisionNum, tuningImage, imagePullSecrets, *workspaceObj.Resource.Count, commands,
 		containerPorts, nil, nil, resourceReq, tolerations, initContainers, sidecarContainers, volumes, volumeMounts, envVars)
 
@@ -256,14 +267,18 @@ func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspac
 	return jobObj, nil
 }
 
-// Now there are two options for data destination 1. HostPath - 2. Image
-func prepareDataDestination(ctx context.Context, workspaceObj *kaitov1beta1.Workspace, inputDirectory string) (*corev1.Container, *corev1.LocalObjectReference, corev1.Volume, corev1.VolumeMount) {
+// Now there are two options for data destination 1. Volume - 2. Image
+func prepareDataDestination(ctx context.Context, workspaceObj *kaitov1beta1.Workspace, inputDirectory string) (*corev1.Container, *corev1.LocalObjectReference, []corev1.Volume, []corev1.VolumeMount) {
 	tuning := workspaceObj.Tuning
 	output := tuning.Output
 
+	if output.Volume != nil {
+		return nil, nil, nil, nil
+	}
+
 	outputImage := output.Image
 	if outputImage == "" {
-		return nil, nil, corev1.Volume{}, corev1.VolumeMount{}
+		return nil, nil, nil, nil
 	}
 
 	var annotationsData map[string]map[string]string
@@ -283,10 +298,10 @@ func prepareDataDestination(ctx context.Context, workspaceObj *kaitov1beta1.Work
 
 	volume, volumeMount := utils.ConfigImagePushSecretVolume(imagePushSecretRef.Name)
 
-	return pusherContainer, &imagePushSecretRef, volume, volumeMount
+	return pusherContainer, &imagePushSecretRef, []corev1.Volume{volume}, []corev1.VolumeMount{volumeMount}
 }
 
-// Now there are three options for DataSource: 1. URL - 2. HostPath - 3. Image
+// Now there are three options for DataSource: 1. URL - 2. Volume - 3. Image
 func prepareDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspace) (*corev1.Container, []corev1.Volume, []corev1.VolumeMount) {
 	input := workspaceObj.Tuning.Input
 
@@ -301,6 +316,10 @@ func prepareDataSource(ctx context.Context, workspaceObj *kaitov1beta1.Workspace
 		initContainer, volume, volumeMount := handleURLDataSource(ctx, workspaceObj)
 		return initContainer, []corev1.Volume{volume}, []corev1.VolumeMount{volumeMount}
 
+	case input.Volume != nil:
+		dataVolume, dataVolumeMount := utils.ConfigDataVolume(input.Volume)
+		return nil, []corev1.Volume{dataVolume}, []corev1.VolumeMount{dataVolumeMount}
+
 	default:
 		return nil, nil, nil
 	}
diff --git a/presets/workspace/tuning/text-generation/fine_tuning.py b/presets/workspace/tuning/text-generation/fine_tuning.py
index 213670e..829b17a 100644
--- a/presets/workspace/tuning/text-generation/fine_tuning.py
+++ b/presets/workspace/tuning/text-generation/fine_tuning.py
@@ -5,13 +5,14 @@ import os
 from dataclasses import asdict
 from datetime import datetime
 from parser import parse_configs, load_chat_template
+from cli import ModelConfig, ExtDataCollator, ExtLoraConfig, DatasetConfig
 
 import torch
 from accelerate import Accelerator
 from dataset import DatasetManager
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 from transformers import (AutoModelForCausalLM, AutoTokenizer,
-                          BitsAndBytesConfig,
+                          BitsAndBytesConfig, TrainingArguments,
                           TrainerCallback, TrainerControl, TrainerState)
 from trl import SFTTrainer
 
@@ -26,12 +27,12 @@ logging.basicConfig(
 CONFIG_YAML = os.environ.get('YAML_FILE_PATH', '/mnt/config/training_config.yaml')
 parsed_configs = parse_configs(CONFIG_YAML)
 
-model_config = parsed_configs.get('ModelConfig')
-bnb_config = parsed_configs.get('QuantizationConfig')
-ext_lora_config = parsed_configs.get('LoraConfig')
-ta_args = parsed_configs.get('TrainingArguments')
-ds_config = parsed_configs.get('DatasetConfig')
-dc_args = parsed_configs.get('DataCollator')
+model_config: ModelConfig = parsed_configs.get('ModelConfig')
+bnb_config: BitsAndBytesConfig = parsed_configs.get('QuantizationConfig') 
+ext_lora_config: ExtLoraConfig = parsed_configs.get('LoraConfig')
+ta_args: TrainingArguments = parsed_configs.get('TrainingArguments')
+ds_config: DatasetConfig = parsed_configs.get('DatasetConfig')
+dc_args: ExtDataCollator = parsed_configs.get('DataCollator')
 
 accelerator = Accelerator()
 
@@ -111,7 +112,7 @@ empty_cache_callback = EmptyCacheCallback()
 torch.cuda.set_device(accelerator.process_index)
 torch.cuda.empty_cache()
 # Training the Model
-trainer = accelerator.prepare(SFTTrainer(
+trainer: SFTTrainer = accelerator.prepare(SFTTrainer(
     model=model,
     tokenizer=tokenizer,
     train_dataset=train_dataset,
@@ -124,7 +125,8 @@ trainer = accelerator.prepare(SFTTrainer(
 ))
 trainer.train()
 os.makedirs(ta_args.output_dir, exist_ok=True)
-trainer.save_model(ta_args.output_dir)
+# only save the adapter weights
+trainer.model.save_pretrained(ta_args.output_dir)
 
 # Write file to signify training completion
 timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")



Test patch:
diff --git a/pkg/workspace/tuning/preset-tuning_test.go b/pkg/workspace/tuning/preset-tuning_test.go
index ba9902a..be5c896 100644
--- a/pkg/workspace/tuning/preset-tuning_test.go
+++ b/pkg/workspace/tuning/preset-tuning_test.go
@@ -218,7 +218,7 @@ training_config:
 
 	for name, tc := range testcases {
 		t.Run(name, func(t *testing.T) {
-			_, _, resultOutputDir := SetupTrainingOutputVolume(context.Background(), tc.configMap)
+			_, _, resultOutputDir := SetupTrainingOutputVolume(context.Background(), tc.configMap, nil)
 			assert.Equal(t, tc.expectedOutputDir, resultOutputDir)
 		})
 	}
@@ -341,8 +341,8 @@ func TestPrepareDataDestination_ImageDestination(t *testing.T) {
 
 	// Assertions
 	assert.Equal(t, expectedSidecarContainer, sidecarContainer)
-	assert.Equal(t, expectedVolume, volume)
-	assert.Equal(t, expectedVolumeMount, volumeMount)
+	assert.Equal(t, []corev1.Volume{expectedVolume}, volume)
+	assert.Equal(t, []corev1.VolumeMount{expectedVolumeMount}, volumeMount)
 	assert.Equal(t, expectedImagePushSecret, *imagePushSecret)
 }