Skip to content

v0.3.0

Compare
Choose a tag to compare
@christopher-w-murphy christopher-w-murphy released this 05 Sep 15:57
· 10 commits to main since this release
91faff3

Perform "surgery" on existing models. Take a pretrained model with softmax_0 in its attention mechanism and "operate" on it to replace softmax_0 with softmax_n. Based on MosaicML's composer.

Optionally install via:

$ pip install flash-attention-softmax-n[surgery]

New Features:

  • Functional API: add one line of code to your script, flash_attention_n.surgery.apply_attention_softmax_n.
  • Object-oriented API for use with the MosaicML composer trainer, flash_attention_n.surgery.AttentionSoftmaxN.
  • Use flash_attention_n.surgery.surgery_functions. policy_registry to register your model!
    See the README for sample usage.