Add shifted sparse attention

[joecummings]

Summary

Add shifted sparse attention (w/ flash attention) to enable longer context training w/ less memory overhead.

Paper: https://arxiv.org/pdf/2309.12307.pdf
Code: https://github.com/dvlab-research/LongLoRA/tree/main

Testing

• Added test to check for raised ValueError if sample_packing = True and s2_attention = True
pytest tests/utils/test_models.py::ModelsUtilsTest::test_cfg_throws_error_with_s2_attention_and_sample_packing

• Run accelerate launch -m axolotl.cli.train examples/openllama-3b/lora.yml with the following config changes:

datasets:
  - path: Yukang/LongAlpaca-12k # From LongLoRA paper
    type: alpaca
sequence_len:  65536
s2_attention: true

[INSERT WANDB LOG HERE]

Follow-ups

• Add ability to train embed and norm during LoRA, which improves performance according to the above paper (e.g. LoRA+)
• Implement S2 Attn without flash attention (https://github.com/dvlab-research/LongLoRA/blob/be3bb473b4ca2c291c1d26419f6e409d774dd422/llama_attn_replace.py#L225)
• Implement w/ sample packing
• Refactor code

[winglian]

@joecummings do you have time to rebase this onto main? If not, I can take a stab at rebasing later this week.

[joecummings]

@joecummings do you have time to rebase this onto main? If not, I can take a stab at rebasing later this week.

yep I'll do this later today!

[winglian]

Looks good to me. @NanoCode012 xan you take a quick look please to make sure I didn't miss anything? Thanks

[NanoCode012]

Since it applies to llama models, do we need to account for mistral here as well?

[winglian]

@joecummings this should work with mistral and mixtral too, right?

[NanoCode012]

Hm, does it mean, FA won't be enabled for inference mode now?

[winglian]

I don't think FA was ever enabled for flash_attention. here's the original code:

    if cfg.is_llama_derived_model and cfg.flash_attention and cfg.sample_packing:
        if cfg.device not in ["mps", "cpu"] and not inference:
            from axolotl.monkeypatch.llama_attn_hijack_flash import (
                replace_llama_attn_with_flash_attn,
            )

            LOG.info("patching with flash attention for sample packing")
            replace_llama_attn_with_flash_attn(
                packed=cfg.sample_packing,
                cross_entropy=cfg.flash_attn_cross_entropy,
                rms_norm=cfg.flash_attn_rms_norm,
            )

[winglian]

Thanks for all your work on this @joecummings !