Add support for GPTNeoX models

[naubull2]

Adds Long-LoRA support for GPTNeoX models.

Tested on a colab A100 40GB x 1 instance, with the scripts

• fine-tune.py
• supervised-fine-tune.py

Using a sample GPTNeoX model

• EleutherAI/pythia-1.4b-deduped

As there was no specific guide on how to contribute, I've tried to make as little modification as possible to the original structure.

Added GPTNeoX support by adding a module gptneox_attn_replace just as the original llama_attn_replace.

How to apply

Application is showcased in the tested scripts fine-tune.py, supervised-fine-tune.py

Add

• model_type argument to switch back and forth between the llama and gpt-neox configuration.
• import

  from gptneox_attn_replace import replace_gpt_neox_attn

• Appropriate changes needed for low rank training

   if training_args.low_rank_training:
        if model_args.model_type == "gpt-neox":
            # added `dense` to match with llama as the vanilla peft config would only target 'query_key_value'
            targets = ["query_key_value", "dense"]
        else:
            targets=["q_proj", "k_proj", "v_proj", "o_proj"],
  
        config = LoraConfig(
            r=8,
            lora_alpha=16,
            target_modules=targets,
            lora_dropout=0,
            bias="none",
            task_type="CAUSAL_LM",
        )
        model = get_peft_model(model, config)

Notes on flash-attention + GPTNeoX

• As the huggingface implementation won't support flash attention off the shelf, I modified some parts from modeling_gpt_neox.py, for the use_flash_attn=True case.
  • transformers == 4.33.3 as of writing.
  • Mainly the part where cached cos/sin rotary embedding is in fp32 where flash-attn requires tensors to be in fp16/bf16 only.
• Some how the original flash-attention2 interface flash_attn_varlen_func would cause a runtime error of "in-place operation" flash-attention code
  • So I've opted for flash_attn_varlen_qkvpacked_func which worked fine.
    • In changing the dimensions to fit in, I referenced codes by Philipp Schmid🤗 ref

[yukang2017]

Hi,

Many thanks for your contribution. These commits are really helpful for this project. I have merged them in to the main branch!

Regards,
Yukang Chen