0.8.4

kyegomez/swarms0.8.4Jul 18, 2023by kyegomez

AI Summary

Fixed critical bugs in the Flash Attention implementation and ensured compatibility with PyTorch 2.0.

Key Highlights

  • Fixed Flash Attention bugs (dimensions, scale)
  • Integrated flash attention module
  • CUDA support added
  • PyTorch 2.0 compatibility

New Features

  • Flash Attention integration
  • CUDA configuration for efficient attention
  • PyTorch 2.0 support
  • Error handling improvements for incompatible options

Full Release Notes

# Changelog

## Bugs
1. The flash attention module was missing in the original codebase. This resulted in a module not found error during the execution.
2. Incorrect implementation of flash attention integration with the main attention module. The `forward` method in the `Attend` class wasn't correctly handling flash attention. 
3. The `flash_attn` function within the `Attend` class had incorrect assumptions about the dimensions of the `k` and `v` tensors. This led to dimension mismatch errors during the tensor operations.
4. The original `flash_attn` method was not handling the scale correctly when `qk_norm` was set to `True`.

## Improvements
1. Integrated the flash attention module into the main codebase and ensured the dimensions and operations are correct.
2. Modified the `forward` method in the `Attend` class to handle flash attention correctly. It checks whether flash attention is enabled and accordingly calls the correct attention method.
3. Adjusted the `flash_attn` method to account for possible missing dimensions in `q`, `k`, and `v` tensors, and to correct for possible dimension mismatches.
4. Included a check to determine if the tensor is on a CUDA device and if so, to leverage the appropriate CUDA configuration for efficient attention.
5. Correctly handled the scale in the `flash_attn` method when `qk_norm` was `True`.
6. Added assertions and informative error messages for incompatible options such as 'talking heads' and 'flash attention'.
7. Ensured compatibility with PyTorch version 2.0 and above for using flash attention.