Major Improvements
Documentation
- Comprehensive overhaul of documentation for better clarity and usability
- Added more examples and use cases to help users get started quickly
- Improved API references for easier navigation
LoRA (Low-Rank Adaptation)
- Fixed critical issues affecting LoRA functionality
- Introduced a new option to selectively save only A, B, and Alpha matrices in LoRA instead of the entire merged model
- This feature significantly reduces storage requirements and improves flexibility in model management
Checkpoint Managers
- Updated with enhanced functionality for more efficient model state handling
- Improved integration with distributed training workflows
Removed Features
- The FJFormer Linen API has been deprecated and removed from this version
New Features
ImplicitArray
- Using `core.ImplicitArray`, a powerful abstraction for handling large arrays without instantiation
- Enables lazy evaluation and efficient array operations in JAX
Quantization Support
- Added `Array8Bit` for 8-bit quantization
- Reduces model size while maintaining good performance
- Introduced `Array4Bit` for 4-bit quantization (NF4)
- Offers extreme model compression for scenarios where size is critical