[NLP] PEFT-IA3로 LLM 학습하기

잉송·2023년 8월 17일
0

NLP

목록 보기
8/9
post-thumbnail

▶ 발단

기존의 LLM을 peft 라이브러리의 lora를 사용하여 학습했다.
이번에 peft 라이브러리에 IA3가 추가되면서 IA3로 LLM을 학습시키기를 시도했다.

▶ 위기

IA3를 작동시키기 위해서 여러 오류들이 발생했다.
아직 dev 버전 패키지라서 오류가 많다.

패키지 버전

  • transformers : 4.31.0
  • peft : 0.5.0.dev0

▶ 절정

1. 시작 코드

우선 huggingface에 가면 시작 코드를 아래와 같은 방식으로 작성하라고 한다.

from transformers import AutoModelForCausalLM, ia3Config
from peft import IA3Model, IA3Config

config = IA3Config(
        peft_type="IA3",
        task_type="SEQ_2_SEQ_LM",
        target_modules=["k", "v", "w0"],
        feedforward_modules=["w0"],
        )

model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
ia3_model = IA3Model(config, model)

나는 t5-base가 아닌 kovicuna를 사용하기 위해 아래와 같이 수정했다. 그리고 ia3Config가 transformers에 없다고 에러가 뜨기도 하고, 지금 당장 사용하는 곳이 없어 삭제했다. 하지만 아래 코드를 실행하면 'adapter_name'이 없다는 에러가 뜬다.

from transformers import AutoModelForSeq2SeqLM
from peft import IA3Model, IA3Config

model = AutoModelForCausalLM.from_pretrained(
        'junelee/ko_vicuna_7b',
        load_in_8bit=True,
        torch_dtype=torch.float16,
        device_map=device_map,
    )
    
config = IA3Config(
            peft_type="IA3",
            task_type="CAUSAL_LM",
            target_modules=['q_proj', 'k_proj', 'down_proj'],
            feedforward_modules=["down_proj"],
        )
        
model = IA3Model(config, model)   

아래는 모델별 적정 target_module과 feedforward_modules 관련 refernece이다.

target module reference

TRANSFORMERS_MODELS_TO_IA3_TARGET_MODULES_MAPPING = {
    "t5": ["k", "v", "wo"],
    "mt5": ["k", "v", "wi_1"],
    "gpt2": ["c_attn", "mlp.c_proj"],
    "bloom": ["query_key_value", "mlp.dense_4h_to_h"],
    "roberta": ["key", "value", "output.dense"],
    "opt": ["q_proj", "k_proj", "fc2"],
    "gptj": ["q_proj", "v_proj", "fc_out"],
    "gpt_neox": ["query_key_value", "dense_4h_to_h"],
    "gpt_neo": ["q_proj", "v_proj", "c_proj"],
    "bart": ["q_proj", "v_proj", "fc2"],
    "gpt_bigcode": ["c_attn", "mlp.c_proj"],
    "llama": ["k_proj", "v_proj", "down_proj"],
    "bert": ["key", "value", "output.dense"],
    "deberta-v2": ["key_proj", "value_proj", "output.dense"],
    "deberta": ["in_proj", "output.dense"],
}
TRANSFORMERS_MODELS_TO_IA3_FEEDFORWARD_MODULES_MAPPING = {
    "t5": ["wo"],
    "mt5": [],
    "gpt2": ["mlp.c_proj"],
    "bloom": ["mlp.dense_4h_to_h"],
    "roberta": ["output.dense"],
    "opt": ["fc2"],
    "gptj": ["fc_out"],
    "gpt_neox": ["dense_4h_to_h"],
    "gpt_neo": ["c_proj"],
    "bart": ["fc2"],
    "gpt_bigcode": ["mlp.c_proj"],
    "llama": ["down_proj"],
    "bert": ["output.dense"],
    "deberta-v2": ["output.dense"],
    "deberta": ["output.dense"],
}


2. adapter_name

'adapter_name' 에러를 수정하기 위해 구글을 검색했다.

# before
model = IA3Model(config, model)  

# after
model = IA3Model(config, model, adapter_name="ia3_adapter")  

하지만 그 이후에도 IA3Model에서는 에러가 난무했다.
위 링크를 통해 더 알게된 사실로 IA3Model가 아닌 기존 lora에서 사용했던 get_peft_model로 코드를 수정하면 작동한다는 것이다.

# before
model = IA3Model(config, model, adapter_name="ia3_adapter")  

# after
model = get_peft_model(model, config)

그렇게 2번까지 완성된 코드는 아래와 같다.

from transformers import AutoModelForSeq2SeqLM
from peft import IA3Model, IA3Config

model = AutoModelForCausalLM.from_pretrained(
        'junelee/ko_vicuna_7b',
        load_in_8bit=True,
        torch_dtype=torch.float16,
        device_map=device_map,
    )
    
config = IA3Config(
            peft_type="IA3",
            task_type="CAUSAL_LM",
            target_modules=['q_proj', 'k_proj', 'down_proj'],
            feedforward_modules=["down_proj"],
        )
        
model = get_peft_model(model, config)

하지만 위 코드도 에러가 떴다.



3. forward - result

에러는 다음과 같이 ia3.py 에서 result 부분 즉 forward 함수 부분에서의 에러이다.

UnboundLocalError: local variable 'result' referenced before assignment
result = result * self.ia3_l[self.active_adapter].flatten()
  • source: https://github.com/huggingface/peft/pull/794
    위 링크를 통해 기존의 forward 함수를 아래와 같이 변경하면 작동한 다는 것이다. ia3.py를 열어 forward함수를 아래 함수로 변경했다. 그 이후로는 잘 작동했다!!
        # current forward (incl. bugfix)
        def forward(self, x: torch.Tensor):
            if self.disable_adapters or self.active_adapter not in self.ia3_l.keys():
                return super().forward(x)
            else:
                if not torch.is_autocast_enabled():
                    if x.dtype != torch.float32:
                        x = x.float()
                    if self.is_feedforward:
                        result = super().forward(x * self.ia3_l[self.active_adapter].flatten())
                    else:
                        result = super().forward(x)
                        expected_dtype = result.dtype
                        result = (result * self.ia3_l[self.active_adapter].flatten()).to(expected_dtype)
                else:
                    if self.is_feedforward:
                        result = super().forward(x * self.ia3_l[self.active_adapter].flatten())
                    else:
                        result = super().forward(x)
                        result = result * self.ia3_l[self.active_adapter].flatten()
            return result

        # proposed new forward (incl. bugfix)
        def forward(self, x: torch.Tensor):
            if self.disable_adapters or (self.active_adapter not in self.ia3_l.keys()):
                return super().forward(x)

            requires_conversion = (not torch.is_autocast_enabled()) and (x.dtype != torch.float32)
            if requires_conversion:
                x = x.float()

            ia3_scaling = self.ia3_l[self.active_adapter].flatten()
            if self.is_feedforward:
                result = super().forward(x * ia3_scaling)
            else:
                result = super().forward(x)
                expected_dtype = result.dtype
                result = result * ia3_scaling

            if requires_conversion:
                result = result.to(expected_dtype)
            return result



▶ 결말

  • huggingface에서 알려준 시작방식과 다르다! 아래와 같이 시작해야한다!
from transformers import AutoModelForSeq2SeqLM
from peft import IA3Config

model = AutoModelForCausalLM.from_pretrained(
        'junelee/ko_vicuna_7b',
        load_in_8bit=True,
        torch_dtype=torch.float16,
        device_map=device_map,
    )
    
config = IA3Config(
            peft_type="IA3",
            task_type="CAUSAL_LM",
            target_modules=['q_proj', 'k_proj', 'down_proj'],
            feedforward_modules=["down_proj"],
        )
        
model = get_peft_model(model, config)
  • 기존의 peft 0.5.0.dev0에 패키지에 있는 peft>tuners>ia3.py에서 forward 함수 부분을 아래와 같이 변경해야한다!
        def forward(self, x: torch.Tensor):
            if self.disable_adapters or (self.active_adapter not in self.ia3_l.keys()):
                return super().forward(x)

            requires_conversion = (not torch.is_autocast_enabled()) and (x.dtype != torch.float32)
            if requires_conversion:
                x = x.float()

            ia3_scaling = self.ia3_l[self.active_adapter].flatten()
            if self.is_feedforward:
                result = super().forward(x * ia3_scaling)
            else:
                result = super().forward(x)
                expected_dtype = result.dtype
                result = result * ia3_scaling

            if requires_conversion:
                result = result.to(expected_dtype)
            return result
  • 이 모든것을 끝내면 잘 작동한다!!!!!
profile
NLP 공부하는 사람

1개의 댓글

comment-user-thumbnail
2023년 8월 17일

좋은 글이네요. 공유해주셔서 감사합니다.

답글 달기