FJFormer is a powerful and flexible JAX-based package designed to accelerate and simplify machine learning and deep learning workflows. It provides a comprehensive suite of tools and utilities for efficient model development, training, and deployment.

Leverage the power of distributed computing and model parallelism with our advanced JAX sharding utilities. These tools enable efficient splitting and management of large models across multiple devices, enhancing performance and enabling the training of larger models.

Boost your model's performance with our optimized kernels for specific operations. These custom-built kernels, implemented using Pallas and Triton, provide significant speedups for common bottleneck operations in deep learning models.

Jump-start your training with our collection of ready-to-use, efficiently implemented optimization algorithms:

• AdamW: An Adam variant with decoupled weight decay.
• Adafactor: Memory-efficient adaptive optimization algorithm.
• Lion: Recently proposed optimizer combining the benefits of momentum and adaptive methods.
• RMSprop: Adaptive learning rate optimization algorithm.

A rich set of utility functions to streamline your workflow, including:

• Various loss functions (e.g., cross-entropy)
• Metrics calculation
• Data preprocessing tools

Our innovative ImplicitArray class provides a powerful abstraction for representing and manipulating large arrays without instantiation. Benefits include:

• Lazy evaluation for memory efficiency
• Optimized array operations in JAX
• Seamless integration with other FJFormer components

• Implement 4-bit quantization (NF4) effortlessly using our ArrayNF4 class, built on top of ImplicitArray. Reduce model size and increase inference speed without significant loss in accuracy (from QLoRA paper).

• Similar to ArrayNF4, our Array8Lt implementation offers 8-bit quantization via ImplicitArray, providing a balance between model compression and precision.

Efficiently fine-tune large language models with our LoRA implementation, leveraging ImplicitArray for optimal performance and memory usage.

A comprehensive set of tools and utilities for efficient array operations and manipulations in JAX, designed to complement and extend JAX's native capabilities.

Robust utilities for managing model checkpoints, including:

• Efficient saving and loading of model states
• Version control for checkpoints
• Integration with distributed training workflows

You can install FJFormer using pip:

pip install fjformer

For the latest development version, you can install directly from GitHub:

pip install git+https://github.com/yourusername/fjformer.git

For detailed documentation, including API references, please visit:

https://fjformer.readthedocs.org

FJFormer is released under the Apache License 2.0. See the LICENSE file for more details.