Features:
- Support MaxText model
- Support KerasHub Gemma2 model
- Implement SFT in Kithara
- Support data loading with Ray data
- Support single host SFT workload on TPU and GPU
- Support GCE VM & QR multihost workload on TPU and GPU via Ray
- HF <-> MaxText converter for Gemma2 models
- Support Orbax checkpointing
- Support Xprof and Tensorboard
- Support saving checkpoints to cloud storage
- Support FSDP for Gemma2