From f58e58b7c2e2a369a7fccadfeab98d87d46df7b3 Mon Sep 17 00:00:00 2001 From: Varsha Prasad Narsing Date: Tue, 24 Sep 2024 08:45:44 -0700 Subject: [PATCH] [KFTO-SDK] Add doc on cutomizing base images for Train API Follow up from https://github.com/kubeflow/training-operator/pull/2261 as this is a user facing change. Signed-off-by: Varsha Prasad Narsing --- .../components/training/user-guides/fine-tuning.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/content/en/docs/components/training/user-guides/fine-tuning.md b/content/en/docs/components/training/user-guides/fine-tuning.md index b190291130..5313ed296d 100644 --- a/content/en/docs/components/training/user-guides/fine-tuning.md +++ b/content/en/docs/components/training/user-guides/fine-tuning.md @@ -87,6 +87,19 @@ TrainingClient().train( After you execute `train`, the Training Operator will orchestrate the appropriate PyTorchJob resources to fine-tune the LLM. +## Using custom images with Fine-Tuning API + +Platform engineers can customize the storage initializer and trainer images by setting the `STORAGE_INITIALIZER_IMAGE` and `TRAINER_TRANSFORMER_IMAGE` environment variables before executing the `train` command. + +For example: In your python code, set the env vars before executing `train`: +```python +... +os.environ['STORAGE_INITIALIZER_IMAGE'] = 'docker.io//' +os.environ['TRAINER_TRANSFORMER_IMAGE'] = 'docker.io//' + +TrainingClient().train(...) +``` + ## Next Steps - Run the example to [fine-tune the TinyLlama LLM](https://github.com/kubeflow/training-operator/blob/6ce4d57d699a76c3d043917bd0902c931f14080f/examples/pytorch/language-modeling/train_api_hf_dataset.ipynb)