# Kithara

A LLM Post-training Library for TPUs and GPUs. 

# Set up

Kithara requires `Python>=3.11`.

### On CPU 

``` 
pip install kithara[cpu] 
```

### On TPU 

``` 
pip install kithara[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html --extra-index-url https://download.pytorch.org/whl/cpu 
```
### On GPU 

``` 
pip install kithara[gpu]
```

# Examples

## SFT with LoRA 

Example of LoRA finetuning gemma2-2b. This script runs on single-host and multi-host environments, on both TPUs and GPUs. For multi-host set up, we included a Ray guide in the next section. 

```
python kithara/examples/singlehost/sft_lora_example.py
```

## Full parameter finetuning

Example of training a MaxText model. 

```
python kithara/examples/singlehost/full_finetuning_example.py
```

## Multi-host examples

Following instructions in `ray/README.md` to set up a Ray Cluster for running multi-host workloads. Here are example of how to  run tuning tasks once your cluster has been set up.

First copy the example script in the `examples/multihost` folder to a new folder on your your local machine, let's call it `ray_workdir`.
Then, use the `kithara multihost` CLI` to run the script on your Ray Cluster. 


```
cd ray_workdir
kithara multihost sft_lora_example.py --hf-token your_token
```

Similarly, you can run the full parameter finetuning example using the following command

```
cd ray_workdir
kithara multihost full_finetuning_example.py --hf-token your_token
```

You can early-stop your job using 

```ray job stop ray_job_id```

# Troubleshooting

1. Disk OOM when loading HF model checkpoint 

    First try emptying your cache by running the following code on your Ray Cluster.

    ```
    import shutil
    shutil.rmtree("/home/ubuntu/.cache/huggingface/hub/", ignore_errors=True)
    shutil.rmtree("/home/ubuntu/.keras/models", ignore_errors=True)
   ```

    If you are using a single VM, the path may be different.

    ```
    import shutil
    shutil.rmtree("~/.cache/huggingface/hub/", ignore_errors=True)
    shutil.rmtree("~/.keras/models", ignore_errors=True)
    ```

    If emptying the cache still doesn't help, try attaching a disk to your VM and change HF cache directory using the environment variable `export HF_HOME=<your_new_cache_dir>`. 
    
    You may have to copy your HF token to this new cache directory with `cp .cache/huggingface/token <your_new_cache_dir>/token`. 

2. Permission denied error when uploading checkpoint to GCS 

    First verify your current authentication :

    ```
    gcloud auth list
    gsutil ls gs://your_bucket
    ```

    For your Python code, you likely need to ensure you're using the same credentials.

    ```
    gcloud auth application-default login
    ```

3. jaxlib.xla_extension.XlaRuntimeError errors

    Try uninstall and reinstalling `jax` and `jaxlib`

    ```
    pip uninstall jax jaxlib
    pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    ```

