Metadata-Version: 2.5
Name: adjaxt
Version: 0.0.2.10
Summary: Distributed training with JAX on free compute, for humans
License-Expression: MIT
License-File: LICENSE
Requires-Python: >=3.11
Requires-Dist: datasets>=2.19
Requires-Dist: huggingface-hub>=0.23
Requires-Dist: jax>=0.4.30
Requires-Dist: numpy
Requires-Dist: optax>=0.2.2
Requires-Dist: orbax-checkpoint>=0.5
Requires-Dist: safetensors>=0.4
Provides-Extra: cuda
Requires-Dist: jax[cuda12]>=0.4.30; extra == 'cuda'
Provides-Extra: test
Requires-Dist: pytest; extra == 'test'
Requires-Dist: torch; extra == 'test'
Requires-Dist: transformers; extra == 'test'
Description-Content-Type: text/markdown

Hi, this package is in development, not ready yet.

It is being tested and most stuff already works.

Needs more layers and archictectures.

Here is how it works:


```mermaid
flowchart LR
    A[You] -->|Want convenience| B(Pytorch)
    B -->|Works slowly because adds python overhead| C[CUDA]

    D[Me] --> |Want to do math| E(Rewrite popular hf models in JAX)
    E --> |Models convert to pytorch through a simple parser and prebuilt layers| F[huggingface]
```

JAX is fast (because of XLA compilation), works on CUDA and TPU.
This repo implements DiLoCo to use any number of heterogenous compute (kaggle, molab, colab, lightning) without direct connection (only huggingface chekpointing).
