Skip to content

Commit

Permalink
[KFTO-SDK] Add doc on cutomizing base images for Train API
Browse files Browse the repository at this point in the history
Follow up from kubeflow/training-operator#2261 as
this is a user facing change.

Signed-off-by: Varsha Prasad Narsing <[email protected]>
  • Loading branch information
varshaprasad96 committed Sep 27, 2024
1 parent b83d113 commit f58e58b
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions content/en/docs/components/training/user-guides/fine-tuning.md
Original file line number Diff line number Diff line change
Expand Up @@ -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/<username>/<custom-storage-initiailizer_image>'
os.environ['TRAINER_TRANSFORMER_IMAGE'] = 'docker.io/<username>/<custom-trainer_transformer_image>'

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)
Expand Down

0 comments on commit f58e58b

Please sign in to comment.