├── .gitignore ├── .gitmodules ├── .readthedocs.yaml ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── constraints.txt ├── docs ├── Makefile ├── README.md ├── images │ ├── kithara_logo.png │ ├── kithara_logo_no_text.png │ └── kithara_logo_with_green_bg.png ├── requirements.txt └── source │ ├── .getting_tpus.rst.swp │ ├── .overview.rst.swp │ ├── api │ ├── kithara.dataset_api.rst │ ├── kithara.model_api.rst │ └── kithara.trainer_api.rst │ ├── basic_examples │ └── full_param_finetuning_example.rst │ ├── checkpointing.rst │ ├── conf.py │ ├── datasets.rst │ ├── ddp.rst │ ├── disk_storage.rst │ ├── dws_flex_start.rst │ ├── finetuning_guide.rst │ ├── getting_tpus.rst │ ├── help_and_feedback.rst │ ├── index.rst │ ├── installation.rst │ ├── installation │ ├── .tpu_gke.rst.swp │ ├── tpu_gke.rst │ ├── tpu_qr.rst │ └── tpu_vm.rst │ ├── lora.rst │ ├── models.rst │ ├── observability.rst │ ├── optimizations.rst │ ├── optimizers.rst │ ├── overview.rst │ ├── overviewxyz.rst │ ├── packing.rst │ ├── pretraining.rst │ ├── quickstart.rst │ ├── scaling_with_ray.rst │ ├── serve_with_vllm.rst │ ├── sft.rst │ ├── temp.rst │ └── troubleshooting.rst ├── examples ├── colab │ └── SFT_with_LoRA_Gemma2-2b.ipynb ├── example_datasets.py ├── multihost │ └── ray │ │ ├── GPU │ │ └── sft_lora_example.py │ │ └── TPU │ │ ├── continued_pretraining_example.py │ │ ├── full_finetuning_example.py │ │ ├── quick_start.py │ │ ├── sft_lora_example.py │ │ └── sft_lora_example.yaml └── singlehost │ ├── continued_pretraining_example.py │ ├── full_finetuning_example.py │ ├── quick_start.py │ └── sft_lora_example.py ├── kithara ├── __init__.py ├── callbacks │ ├── __init__.py │ ├── checkpointer.py │ └── profiler.py ├── config │ ├── __init__.py │ ├── default.yaml │ └── pyconfig.py ├── dataset │ ├── __init__.py │ ├── dataloader.py │ ├── dataset.py │ ├── packed_dataset.py │ ├── sft.py │ ├── text_completion.py │ └── utils.py ├── distributed │ ├── __init__.py │ ├── data │ │ ├── __init__.py │ │ └── split.py │ └── sharding │ │ ├── __init__.py │ │ ├── _data_sharding.py │ │ ├── _layout.py │ │ ├── _mesh.py │ │ ├── maxtext.py │ │ ├── models │ │ ├── __init__.py │ │ ├── gemma.py │ │ └── llama.py │ │ ├── strategy.py │ │ └── utils.py ├── model │ ├── README.md │ ├── __init__.py │ ├── hf_compatibility │ │ ├── __init__.py │ │ ├── model_configs.py │ │ ├── shape_mapping.py │ │ └── to_huggingface.py │ ├── kerashub │ │ ├── __init__.py │ │ ├── ckpt_compatibility │ │ │ ├── param_mapping.py │ │ │ └── to_huggingface.py │ │ └── keras_hub_model.py │ ├── maxtext │ │ ├── __init__.py │ │ ├── ckpt_compatibility │ │ │ ├── __init__.py │ │ │ ├── from_huggingface.py │ │ │ ├── param_mapping.py │ │ │ ├── to_huggingface.py │ │ │ └── utils.py │ │ ├── conversion_utils.py │ │ ├── inference_engine.py │ │ └── maxtext_model.py │ ├── model.py │ └── supported_models.py ├── optimizers │ ├── __init__.py │ ├── optax_optimizer.py │ └── protocol.py ├── trainer │ ├── __init__.py │ └── trainer.py └── utils │ ├── gcs_utils.py │ ├── logging_utils.py │ ├── safetensor_utils.py │ ├── torch_utils.py │ └── tree_utils.py ├── perf ├── kerashub.py ├── kithara_packing.py ├── maxtext_native.py ├── maxtext_via_kithara.py ├── sft_kithara.py ├── sft_unsloth.py └── trl_packing.py ├── pyproject.toml ├── ray ├── GPU │ ├── cluster.yaml │ └── find_devices.py ├── README.md ├── TPU │ ├── GCE │ │ └── cluster.yaml │ ├── GKE │ │ ├── Dockerfile │ │ ├── multi-host.yaml │ │ └── single-host.yaml │ └── QR │ │ ├── cluster.yaml │ │ └── qr_worker_startup_script.sh └── submit_job.py └── tests ├── README.md ├── __init__.py ├── callbacks ├── __init__.py └── test_orbax_checkpointer.py ├── dataset ├── __init__.py ├── test_dataloader_creation.py ├── test_dataset_creation.py ├── test_dataset_packing.py └── utils.py ├── model ├── __init__.py ├── kerashub │ ├── ckpt_compatibility │ │ ├── __init__.py │ │ ├── test_loading_models.py │ │ ├── test_lora_adapters_shape_match.py │ │ ├── test_lora_adapters_value_match.py │ │ └── test_saving_models.py │ ├── test_creation.py │ └── test_inference.py ├── maxtext │ ├── __init__.py │ ├── ckpt_compatibility │ │ ├── __init__.py │ │ ├── test_loading_models.py │ │ └── test_saving_models.py │ ├── test_creation.py │ └── test_inference.py ├── test_prompt.py └── utils.py ├── optimizers ├── __init__.py ├── test_optax_optimizer.py └── test_protocol.py ├── test_utils.py └── trainer ├── __init__.py ├── test_sft_e2e.py └── test_trainer_creation.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/.gitignore -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/.gitmodules -------------------------------------------------------------------------------- /.readthedocs.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/.readthedocs.yaml -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/README.md -------------------------------------------------------------------------------- /constraints.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/constraints.txt -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/README.md -------------------------------------------------------------------------------- /docs/images/kithara_logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/images/kithara_logo.png -------------------------------------------------------------------------------- /docs/images/kithara_logo_no_text.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/images/kithara_logo_no_text.png -------------------------------------------------------------------------------- /docs/images/kithara_logo_with_green_bg.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/images/kithara_logo_with_green_bg.png -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/requirements.txt -------------------------------------------------------------------------------- /docs/source/.getting_tpus.rst.swp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/.getting_tpus.rst.swp -------------------------------------------------------------------------------- /docs/source/.overview.rst.swp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/.overview.rst.swp -------------------------------------------------------------------------------- /docs/source/api/kithara.dataset_api.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/api/kithara.dataset_api.rst -------------------------------------------------------------------------------- /docs/source/api/kithara.model_api.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/api/kithara.model_api.rst -------------------------------------------------------------------------------- /docs/source/api/kithara.trainer_api.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/api/kithara.trainer_api.rst -------------------------------------------------------------------------------- /docs/source/basic_examples/full_param_finetuning_example.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/basic_examples/full_param_finetuning_example.rst -------------------------------------------------------------------------------- /docs/source/checkpointing.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/checkpointing.rst -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/conf.py -------------------------------------------------------------------------------- /docs/source/datasets.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/datasets.rst -------------------------------------------------------------------------------- /docs/source/ddp.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/ddp.rst -------------------------------------------------------------------------------- /docs/source/disk_storage.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/disk_storage.rst -------------------------------------------------------------------------------- /docs/source/dws_flex_start.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/dws_flex_start.rst -------------------------------------------------------------------------------- /docs/source/finetuning_guide.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/finetuning_guide.rst -------------------------------------------------------------------------------- /docs/source/getting_tpus.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/getting_tpus.rst -------------------------------------------------------------------------------- /docs/source/help_and_feedback.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/help_and_feedback.rst -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/index.rst -------------------------------------------------------------------------------- /docs/source/installation.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/installation.rst -------------------------------------------------------------------------------- /docs/source/installation/.tpu_gke.rst.swp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/installation/.tpu_gke.rst.swp -------------------------------------------------------------------------------- /docs/source/installation/tpu_gke.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/installation/tpu_gke.rst -------------------------------------------------------------------------------- /docs/source/installation/tpu_qr.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/installation/tpu_qr.rst -------------------------------------------------------------------------------- /docs/source/installation/tpu_vm.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/installation/tpu_vm.rst -------------------------------------------------------------------------------- /docs/source/lora.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/lora.rst -------------------------------------------------------------------------------- /docs/source/models.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/models.rst -------------------------------------------------------------------------------- /docs/source/observability.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/observability.rst -------------------------------------------------------------------------------- /docs/source/optimizations.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/optimizations.rst -------------------------------------------------------------------------------- /docs/source/optimizers.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/optimizers.rst -------------------------------------------------------------------------------- /docs/source/overview.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/overview.rst -------------------------------------------------------------------------------- /docs/source/overviewxyz.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/overviewxyz.rst -------------------------------------------------------------------------------- /docs/source/packing.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/packing.rst -------------------------------------------------------------------------------- /docs/source/pretraining.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/pretraining.rst -------------------------------------------------------------------------------- /docs/source/quickstart.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/quickstart.rst -------------------------------------------------------------------------------- /docs/source/scaling_with_ray.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/scaling_with_ray.rst -------------------------------------------------------------------------------- /docs/source/serve_with_vllm.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/serve_with_vllm.rst -------------------------------------------------------------------------------- /docs/source/sft.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/sft.rst -------------------------------------------------------------------------------- /docs/source/temp.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/temp.rst -------------------------------------------------------------------------------- /docs/source/troubleshooting.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/docs/source/troubleshooting.rst -------------------------------------------------------------------------------- /examples/colab/SFT_with_LoRA_Gemma2-2b.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/colab/SFT_with_LoRA_Gemma2-2b.ipynb -------------------------------------------------------------------------------- /examples/example_datasets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/example_datasets.py -------------------------------------------------------------------------------- /examples/multihost/ray/GPU/sft_lora_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/multihost/ray/GPU/sft_lora_example.py -------------------------------------------------------------------------------- /examples/multihost/ray/TPU/continued_pretraining_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/multihost/ray/TPU/continued_pretraining_example.py -------------------------------------------------------------------------------- /examples/multihost/ray/TPU/full_finetuning_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/multihost/ray/TPU/full_finetuning_example.py -------------------------------------------------------------------------------- /examples/multihost/ray/TPU/quick_start.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/multihost/ray/TPU/quick_start.py -------------------------------------------------------------------------------- /examples/multihost/ray/TPU/sft_lora_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/multihost/ray/TPU/sft_lora_example.py -------------------------------------------------------------------------------- /examples/multihost/ray/TPU/sft_lora_example.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/multihost/ray/TPU/sft_lora_example.yaml -------------------------------------------------------------------------------- /examples/singlehost/continued_pretraining_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/singlehost/continued_pretraining_example.py -------------------------------------------------------------------------------- /examples/singlehost/full_finetuning_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/singlehost/full_finetuning_example.py -------------------------------------------------------------------------------- /examples/singlehost/quick_start.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/singlehost/quick_start.py -------------------------------------------------------------------------------- /examples/singlehost/sft_lora_example.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/examples/singlehost/sft_lora_example.py -------------------------------------------------------------------------------- /kithara/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/__init__.py -------------------------------------------------------------------------------- /kithara/callbacks/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/callbacks/__init__.py -------------------------------------------------------------------------------- /kithara/callbacks/checkpointer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/callbacks/checkpointer.py -------------------------------------------------------------------------------- /kithara/callbacks/profiler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/callbacks/profiler.py -------------------------------------------------------------------------------- /kithara/config/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /kithara/config/default.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/config/default.yaml -------------------------------------------------------------------------------- /kithara/config/pyconfig.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/config/pyconfig.py -------------------------------------------------------------------------------- /kithara/dataset/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/dataset/__init__.py -------------------------------------------------------------------------------- /kithara/dataset/dataloader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/dataset/dataloader.py -------------------------------------------------------------------------------- /kithara/dataset/dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/dataset/dataset.py -------------------------------------------------------------------------------- /kithara/dataset/packed_dataset.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/dataset/packed_dataset.py -------------------------------------------------------------------------------- /kithara/dataset/sft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/dataset/sft.py -------------------------------------------------------------------------------- /kithara/dataset/text_completion.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/dataset/text_completion.py -------------------------------------------------------------------------------- /kithara/dataset/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/dataset/utils.py -------------------------------------------------------------------------------- /kithara/distributed/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/__init__.py -------------------------------------------------------------------------------- /kithara/distributed/data/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/data/__init__.py -------------------------------------------------------------------------------- /kithara/distributed/data/split.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/data/split.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/__init__.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/_data_sharding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/_data_sharding.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/_layout.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/_layout.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/_mesh.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/_mesh.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/maxtext.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/maxtext.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/models/__init__.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/models/gemma.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/models/gemma.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/models/llama.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/models/llama.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/strategy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/strategy.py -------------------------------------------------------------------------------- /kithara/distributed/sharding/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/distributed/sharding/utils.py -------------------------------------------------------------------------------- /kithara/model/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/README.md -------------------------------------------------------------------------------- /kithara/model/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/__init__.py -------------------------------------------------------------------------------- /kithara/model/hf_compatibility/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/hf_compatibility/__init__.py -------------------------------------------------------------------------------- /kithara/model/hf_compatibility/model_configs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/hf_compatibility/model_configs.py -------------------------------------------------------------------------------- /kithara/model/hf_compatibility/shape_mapping.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/hf_compatibility/shape_mapping.py -------------------------------------------------------------------------------- /kithara/model/hf_compatibility/to_huggingface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/hf_compatibility/to_huggingface.py -------------------------------------------------------------------------------- /kithara/model/kerashub/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/kerashub/__init__.py -------------------------------------------------------------------------------- /kithara/model/kerashub/ckpt_compatibility/param_mapping.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/kerashub/ckpt_compatibility/param_mapping.py -------------------------------------------------------------------------------- /kithara/model/kerashub/ckpt_compatibility/to_huggingface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/kerashub/ckpt_compatibility/to_huggingface.py -------------------------------------------------------------------------------- /kithara/model/kerashub/keras_hub_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/kerashub/keras_hub_model.py -------------------------------------------------------------------------------- /kithara/model/maxtext/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/__init__.py -------------------------------------------------------------------------------- /kithara/model/maxtext/ckpt_compatibility/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/ckpt_compatibility/__init__.py -------------------------------------------------------------------------------- /kithara/model/maxtext/ckpt_compatibility/from_huggingface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/ckpt_compatibility/from_huggingface.py -------------------------------------------------------------------------------- /kithara/model/maxtext/ckpt_compatibility/param_mapping.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/ckpt_compatibility/param_mapping.py -------------------------------------------------------------------------------- /kithara/model/maxtext/ckpt_compatibility/to_huggingface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/ckpt_compatibility/to_huggingface.py -------------------------------------------------------------------------------- /kithara/model/maxtext/ckpt_compatibility/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/ckpt_compatibility/utils.py -------------------------------------------------------------------------------- /kithara/model/maxtext/conversion_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/conversion_utils.py -------------------------------------------------------------------------------- /kithara/model/maxtext/inference_engine.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/inference_engine.py -------------------------------------------------------------------------------- /kithara/model/maxtext/maxtext_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/maxtext/maxtext_model.py -------------------------------------------------------------------------------- /kithara/model/model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/model.py -------------------------------------------------------------------------------- /kithara/model/supported_models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/model/supported_models.py -------------------------------------------------------------------------------- /kithara/optimizers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/optimizers/__init__.py -------------------------------------------------------------------------------- /kithara/optimizers/optax_optimizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/optimizers/optax_optimizer.py -------------------------------------------------------------------------------- /kithara/optimizers/protocol.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/optimizers/protocol.py -------------------------------------------------------------------------------- /kithara/trainer/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/trainer/__init__.py -------------------------------------------------------------------------------- /kithara/trainer/trainer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/trainer/trainer.py -------------------------------------------------------------------------------- /kithara/utils/gcs_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/utils/gcs_utils.py -------------------------------------------------------------------------------- /kithara/utils/logging_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/utils/logging_utils.py -------------------------------------------------------------------------------- /kithara/utils/safetensor_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/utils/safetensor_utils.py -------------------------------------------------------------------------------- /kithara/utils/torch_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/utils/torch_utils.py -------------------------------------------------------------------------------- /kithara/utils/tree_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/kithara/utils/tree_utils.py -------------------------------------------------------------------------------- /perf/kerashub.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/perf/kerashub.py -------------------------------------------------------------------------------- /perf/kithara_packing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/perf/kithara_packing.py -------------------------------------------------------------------------------- /perf/maxtext_native.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/perf/maxtext_native.py -------------------------------------------------------------------------------- /perf/maxtext_via_kithara.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/perf/maxtext_via_kithara.py -------------------------------------------------------------------------------- /perf/sft_kithara.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/perf/sft_kithara.py -------------------------------------------------------------------------------- /perf/sft_unsloth.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/perf/sft_unsloth.py -------------------------------------------------------------------------------- /perf/trl_packing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/perf/trl_packing.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/pyproject.toml -------------------------------------------------------------------------------- /ray/GPU/cluster.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/GPU/cluster.yaml -------------------------------------------------------------------------------- /ray/GPU/find_devices.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/GPU/find_devices.py -------------------------------------------------------------------------------- /ray/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/README.md -------------------------------------------------------------------------------- /ray/TPU/GCE/cluster.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/TPU/GCE/cluster.yaml -------------------------------------------------------------------------------- /ray/TPU/GKE/Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/TPU/GKE/Dockerfile -------------------------------------------------------------------------------- /ray/TPU/GKE/multi-host.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/TPU/GKE/multi-host.yaml -------------------------------------------------------------------------------- /ray/TPU/GKE/single-host.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/TPU/GKE/single-host.yaml -------------------------------------------------------------------------------- /ray/TPU/QR/cluster.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/TPU/QR/cluster.yaml -------------------------------------------------------------------------------- /ray/TPU/QR/qr_worker_startup_script.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/TPU/QR/qr_worker_startup_script.sh -------------------------------------------------------------------------------- /ray/submit_job.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/ray/submit_job.py -------------------------------------------------------------------------------- /tests/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/README.md -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/callbacks/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/callbacks/test_orbax_checkpointer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/callbacks/test_orbax_checkpointer.py -------------------------------------------------------------------------------- /tests/dataset/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/dataset/test_dataloader_creation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/dataset/test_dataloader_creation.py -------------------------------------------------------------------------------- /tests/dataset/test_dataset_creation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/dataset/test_dataset_creation.py -------------------------------------------------------------------------------- /tests/dataset/test_dataset_packing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/dataset/test_dataset_packing.py -------------------------------------------------------------------------------- /tests/dataset/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/dataset/utils.py -------------------------------------------------------------------------------- /tests/model/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/model/kerashub/ckpt_compatibility/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/model/kerashub/ckpt_compatibility/test_loading_models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/kerashub/ckpt_compatibility/test_loading_models.py -------------------------------------------------------------------------------- /tests/model/kerashub/ckpt_compatibility/test_lora_adapters_shape_match.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/kerashub/ckpt_compatibility/test_lora_adapters_shape_match.py -------------------------------------------------------------------------------- /tests/model/kerashub/ckpt_compatibility/test_lora_adapters_value_match.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/kerashub/ckpt_compatibility/test_lora_adapters_value_match.py -------------------------------------------------------------------------------- /tests/model/kerashub/ckpt_compatibility/test_saving_models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/kerashub/ckpt_compatibility/test_saving_models.py -------------------------------------------------------------------------------- /tests/model/kerashub/test_creation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/kerashub/test_creation.py -------------------------------------------------------------------------------- /tests/model/kerashub/test_inference.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/kerashub/test_inference.py -------------------------------------------------------------------------------- /tests/model/maxtext/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/model/maxtext/ckpt_compatibility/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/model/maxtext/ckpt_compatibility/test_loading_models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/maxtext/ckpt_compatibility/test_loading_models.py -------------------------------------------------------------------------------- /tests/model/maxtext/ckpt_compatibility/test_saving_models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/maxtext/ckpt_compatibility/test_saving_models.py -------------------------------------------------------------------------------- /tests/model/maxtext/test_creation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/maxtext/test_creation.py -------------------------------------------------------------------------------- /tests/model/maxtext/test_inference.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/maxtext/test_inference.py -------------------------------------------------------------------------------- /tests/model/test_prompt.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/test_prompt.py -------------------------------------------------------------------------------- /tests/model/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/model/utils.py -------------------------------------------------------------------------------- /tests/optimizers/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/optimizers/test_optax_optimizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/optimizers/test_optax_optimizer.py -------------------------------------------------------------------------------- /tests/optimizers/test_protocol.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/optimizers/test_protocol.py -------------------------------------------------------------------------------- /tests/test_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/test_utils.py -------------------------------------------------------------------------------- /tests/trainer/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/trainer/test_sft_e2e.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/trainer/test_sft_e2e.py -------------------------------------------------------------------------------- /tests/trainer/test_trainer_creation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AI-Hypercomputer/kithara/HEAD/tests/trainer/test_trainer_creation.py --------------------------------------------------------------------------------