Pinned Loading

1. quad_torch quad_torch Public

   An implementation of PSGD-QUAD optimizer for PyTorch

   Python 5

2. distributed_kron distributed_kron Public

   An implementation of PSGD Kron in JAX for distributed training in JAX or Flax

   Python 10

3. kron_torch kron_torch Public

   An implementation of PSGD Kron second-order optimizer for PyTorch

   Python 97 6

4. psgd_jax psgd_jax Public

   Implementation of PSGD optimizer in JAX

   Python 35 2

5. llm-jax llm-jax Public

   Train a SmolLM-style llm on fineweb-edu in JAX/Flax with an assortment of optimizers.

   Python 18 1