Machine Learning/DL - NLP

PEFT 기법 (LoRA, IA3)

IP_DataScientist 2023. 7. 7.
반응형

PEFT (Parameter Effcient Fine-Tuning)

LLM 모델 튜닝, 하나의 GPU로 가능할까? Parameter Efficient Fine-Tuning(PEFT)을 소개합니다!

PEFT 기법

모델의 모든 파라미터를 튜닝하는 것이 아닌 일부 파라미터만을 튜닝함으로써 모델의 성능을 적은 자원으로도 높게 유지하는 방법론

그중 가장 많이 알려진 방법은 LoRA 기법

  • 적은양 파라미터(예를 들면 0.01%)를 학습함으로써 빠른 시간 내에 새로운 문제를 LLM의 In Context-Learning과 거의 비슷한 성능으로 풀 수 있게 하자는게 주된 목표

LoRA

개념

LoRA(Low-Rank Adaptation)의 개념을 간단하게 설명하자면, 고정된 weights를 갖는 pretrained model에 학습이 가능한 rank decomposition 행렬을 삽입한것으로

중간중간 학습이 가능한 파라미터를 삽입했다는 점에서는 어댑터와 비슷하지만 구조적으로 조금 다르다고 할 수 있습니다.

적은 양의 파라미터로 모델을 튜닝하는 방법론이기 때문에 적은수의 GPU로 빠르게 튜닝할 수 있다는 장점이 있습니다.

LoRA에서 나온 rank decomposition이라는 말이 처음에는 어렵게 느껴졌었는데요.

아래 보이는 그림에서 처럼 행렬의 차원을 r 만큼 줄이는 행렬과 다시 원래 크기로 키워주는 행렬의 곱으로 나타내는 것을 의미합니다.

 

💡 LoRA를 중학생이 이해할 수 있도록 설명하면, 이건 마치 당신이 학교에서 시험을 볼 때, 전체 교과서를 외울 필요 없이 중요한 부분만 외워서 시험 점수를 올리는 방법에 비유할 수 있어요.
일반적으로, 기계 학습 모델은 많은 양의 데이터를 '공부'하고 이를 바탕으로 '질문'에 대답하게 됩니다. 그런데 이 모든 데이터를 '공부'하는데는 많은 시간과 연산력이 필요하죠. 이것이 마치 우리가 모든 교과서를 다 외우려고 할 때처럼 말이에요.
그런데, LoRA라는 기술을 이용하면, 모든 데이터를 '공부'하지 않고도 모델이 '질문'에 잘 대답할 수 있게 됩니다. 이건 우리가 시험을 위해 교과서의 모든 내용을 외우지 않고, 중요한 부분만 외우는 것과 비슷해요. 적은 양의 데이터만 가지고도 좋은 성능을 낼 수 있게 되는 것이죠.
이렇게 하면, 적은 양의 계산만으로도 모델을 학습시킬 수 있어, 시간과 연산력을 많이 절약할 수 있습니다.
여기서 'rank decomposition'이라는 개념이 나오는데, 이건 마치 큰 그림을 작게 그린 다음 다시 크게 확대하는 것과 같아요. 원래의 큰 그림을 작게 그리면서 필요한 정보만을 남겨두는 과정을 통해, 이를 다시 확대했을 때에도 원래의 그림과 비슷하게 만들 수 있습니다. 이런 방식을 통해, 적은 양의 정보만을 사용하여 원래의 모델과 비슷한 성능을 낼 수 있게 되는 것이죠. 이것이 바로 LoRA에서 말하는 'rank decomposition'의 개념입니다.

위 그림처럼 레이어 중간중간마다 존재하는 hidden states h에 값을 더해줄수 있는 파라미터를 추가해줘서

모델의 출력 값을 원하는 타겟 레이블에 맞게 튜닝하는 것이 LoRA의 핵심 개념이라고 할 수 있습니다.

💡 레이어 중간마다 파라미터를 추가해주는 것을 이해하려면, 우리가 케이크를 만든다고 상상해봅시다.

이 케이크 만드는 과정에서 각 단계별로 재료를 추가하거나 바꾸어서 원하는 맛을 만들어내는 것처럼, 이 LoRA 기술에서는 모델이 정보를 처리하는 각 단계(레이어)에서 값을 조금씩 더해주거나 바꾸어서, 원하는 결과(레이블)를 만들어내는 것입니다. 이 과정에서 사용되는 추가적인 파라미터들이 마치 케이크 만드는 과정에서 각 단계별로 사용하는 재료와 비슷한 역할을 합니다.

즉, LoRA의 핵심은 이 추가 파라미터를 활용해서 모델의 출력 값을 우리가 원하는 방향으로 잘 조절(튜닝)하는 것입니다. 이렇게 해서 모델이 우리가 원하는 대답을 잘 할 수 있도록 도와주는 것이죠. 이러한 방식으로, 적은 계산만으로도 원하는 결과를 얻을 수 있게 됩니다.

코드상으로는 아래와 같이 구현할 수 있는데요. 기존에 모델에서 사용하던 Linear Layer를 LoRA의 로직이 적용된 커스텀 클래스로 교체해주면 적용할 수 있습니다.

if self.r > 0: 라는 if 문이 추가된 부분이 LoRA가 적용된 부분입니다.

class Linear(nn.Linear, LoRALayer):
    # LoRA implemented in a dense layer
    def __init__(
        self, 
        in_features: int, 
        out_features: int, 
        r: int = 0, 
        lora_alpha: int = 1, 
        lora_dropout: float = 0.,
        fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
        merge_weights: bool = True,
        **kwargs
    ):
        nn.Linear.__init__(self, in_features, out_features, **kwargs)
        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout,
                           merge_weights=merge_weights)

        self.fan_in_fan_out = fan_in_fan_out
        # Actual trainable parameters
        if r > 0:
            self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)))
            self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)))
            self.scaling = self.lora_alpha / self.r
            # Freezing the pre-trained weight matrix
            self.weight.requires_grad = False
        self.reset_parameters()
        if fan_in_fan_out:
            self.weight.data = self.weight.data.T

    def reset_parameters(self):
        nn.Linear.reset_parameters(self)
        if hasattr(self, 'lora_A'):
            # initialize A the same way as the default for nn.Linear and B to zero
            nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
            nn.init.zeros_(self.lora_B)

    def train(self, mode: bool = True):
        def T(w):
            return w.T if self.fan_in_fan_out else w
        nn.Linear.train(self, mode)
        if self.merge_weights and self.merged:
            # Make sure that the weights are not merged
            if self.r > 0:
                self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling
            self.merged = False
    
    def eval(self):
        def T(w):
            return w.T if self.fan_in_fan_out else w
        nn.Linear.eval(self)
        if self.merge_weights and not self.merged:
            # Merge the weights and mark it
            if self.r > 0:
                self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling
            self.merged = True

    def forward(self, x: torch.Tensor):
        def T(w):
            return w.T if self.fan_in_fan_out else w
        if self.r > 0 and not self.merged:
            result = F.linear(x, T(self.weight), bias=self.bias)
            if self.r > 0:
                result += (self.lora_dropout(x) @ self.lora_A.T @ self.lora_B.T) * self.scaling
            return result
        else:
            return F.linear(x, T(self.weight), bias=self.bias)

이 코드는 PyTorch를 사용한 Linear 레이어에 LoRA를 적용한 예제입니다. 주요 기능과 부분에 대해 설명해보겠습니다.

  1. 초기화 (init): 이 부분에서는 입력 특성(in_features), 출력 특성(out_features), 그리고 LoRA 특성(r, lora_alpha, lora_dropout 등)을 설정합니다.
  2. if r > 0:: 이 코드 라인에서 r이 0보다 크면 (즉, LoRA가 적용되어야 할 경우), 우리는 두 개의 새로운 파라미터 lora_A와 lora_B를 생성하고, 원래의 weight는 더 이상 학습이 일어나지 않도록 고정합니다. 이 두 파라미터는 LoRA에서 언급한 low-rank 행렬의 부분이며, 이 행렬은 모델의 출력을 우리가 원하는 대로 조정하는데 사용됩니다.
  3. reset_parameters(): 이 메소드는 원래의 파라미터와 lora_A, lora_B 파라미터를 초기화합니다.
  4. train()과 eval(): 이 메소드들은 모델이 학습 모드인지 평가 모드인지를 결정합니다. 이 둘 사이의 차이는 lora_A와 lora_B 파라미터가 원래의 weight에 병합되는지 여부에 있습니다. 학습 모드에서는 이들이 분리되어 있으며, 평가 모드에서는 병합됩니다.
  5. forward(): 이것은 모델의 출력을 계산하는 메소드입니다. 여기서는 원래의 Linear 레이어의 출력에 LoRA의 출력을 더해줍니다. 이 부분이 바로 LoRA의 핵심 동작을 나타냅니다.

요약하면, 이 코드는 Linear 레이어에 LoRA를 적용하여, 레이어의 출력을 조정하고 모델의 성능을 향상시키는 방법을 보여주고 있습니다. 여기서 가장 중요한 부분은 if r > 0: 부분으로, 이 부분이 LoRA가 적용되는 핵심 구간입니다.

위 코드에서 주목할 부분은 eval 함수쪽인데요.

LoRA가 행렬 연산을 기반으로 하기 때문에 기존 행렬 W_0를 LoRA에서 사용하는 A, B 행렬을 기반으로 다음과 같이 재구성할 수 있습니다.

W = W_0 + BA

이렇게 함으로써 얻을 수 있는 이점은 새롭게 학습한 파라미터를 기존에 학습된 pretrained model에 합쳐줌으로써

추가적인 연산이 필요하지 않게 되어 속도도 그대로 유지하면서 아키텍쳐의 변경도 필요없어지게 됩니다.

💡 LoRA를 적용하면 모델의 가중치 W는 원래의 가중치 W_0와 LoRA 행렬 A와 B의 곱의 합, 즉 W = W_0 + BA 로 표현이 가능합니다.

이렇게 되면 모델은 두 부분으로 나누어집니다. 하나는 원래의 pretrained 모델 (W_0)이고, 다른 하나는 LoRA를 통해 학습된 파라미터 (BA)입니다. 이렇게 나누면, 새로운 파라미터를 학습하면서도 원래의 모델을 유지할 수 있습니다.

이런 방식의 장점은 여러 가지가 있지만, 가장 큰 것은 아마도 '효율성'일 것입니다. 첫째, 추가적인 연산이 필요하지 않습니다. 즉, 모델의 복잡성을 늘리지 않고도 새로운 파라미터를 학습할 수 있습니다. 둘째, 모델의 구조를 변경할 필요가 없습니다. 즉, 원래의 모델 구조를 그대로 유지하면서도 새로운 파라미터를 추가할 수 있습니다.

따라서, LoRA는 기존의 모델을 유지하면서도 새로운 학습을 가능하게 하므로, 효율적인 학습과 빠른 속도를 제공합니다.

Alpaca-LoRA

최근에 유행하는 LLaMA의 변형인 Alpaca에도 LoRA가 적용된 오픈소스 프로젝트들이 공개

Huggingface에서 공개한 PEFT 라이브러리를 이용하면 아래와 같이 적용도 매우 간단하게 할 수 있습니다.

from peft import (
    LoraConfig,
    get_peft_model,
    get_peft_model_state_dict,
    prepare_model_for_int8_training,
    set_peft_model_state_dict,
)
from transformers import LlamaForCausalLM, LlamaTokenizer

def train(...):
    model = prepare_model_for_int8_training(model)
    config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        target_modules=lora_target_modules,
        lora_dropout=lora_dropout,
        bias="none",
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, config)

IA3

IA3 (Infused Adapter by Inhibiting and Amplifying Inner Activations)

  • 내부 활성화 억제 및 증폭을 통한 인퓨즈드 어댑터

LoRA와 비슷한 방법으로 적은 파라미터만을 추가해서 모델을 튜닝할 수 있는 방법론

이름에도 나와있듯이 뉴럴네트워크의 Inner Activation을 줄이기도하고 늘리기도하는 어댑터를 중간에 삽입하는 방법론인데요.

 

LoRA의 경우에는 hidden state에 새로운 값을 더해주는 기법이었다면

IA3의 경우에는 Self-Attention, Cross-Attention에서의 Key, Value 값을 rescale해주는 벡터와

position-wise feed-forward network의 값에 rescale을 해주는 벡터를 추가해서 모델을 튜닝해주는 기법 입니다.

💡 우리가 음악을 듣는다고 상상해보세요.

각각의 악기 소리가 믹서를 통해 조절되어서 우리 귀에 들려집니다. 어떤 악기는 크게, 어떤 악기는 작게 들려지게 믹서에서 소리를 조절하는 것이죠.

그런데 만약 우리가 악기 소리의 크기를 바꿔서 새로운 음악을 만들고 싶다면, 믹서에서 각각의 악기 볼륨을 조절해야 할 것입니다.

IA3 기술도 비슷합니다. 뉴럴네트워크(이 경우엔 음악)에서 각 레이어의 출력(악기 소리)을 조절하는 믹서 역할을 하는 것이 IA3입니다. 여기서 "Inner Activation"이란 각 레이어의 출력을 의미합니다.

IA3는 기존 모델에 추가되는 "어댑터"로, 뉴럴 네트워크의 각 레이어에서 출력되는 값을 '높이거나' '낮추는' 기능을 합니다. 즉, 모델의 출력을 우리가 원하는 방향으로 조절해줍니다. 이를 통해 원래 모델이 가지고 있던 정보를 유지하면서도 새로운 학습을 가능하게 하며, 이렇게 조절된 모델을 사용하면 원하는 결과를 더 잘 얻을 수 있게 됩니다.

LoRA가 정보를 더하는 방식으로 모델을 튜닝했다면, IA3는 정보의 크기를 조절하는 방식으로 모델을 튜닝한다고 보시면 됩니다.

💡 IA3, 또는 T-Few는 효율성과 성능에서 매우 뛰어난 방법론으로 알려져 있습니다.

기존에 사용되던 LoRA에 비해 IA3는 더 적은 수의 파라미터를 사용하면서도 높은 성능을 보입니다.

즉, IA3는 더 적은 계산 복잡성과 더 적은 메모리 사용량으로 높은 성능을 내는 것이 가능하다는 뜻입니다. 그리고 이런 특성 덕분에 IA3는 빠르게 모델을 튜닝할 수 있습니다.

실제로, A100 GPU 하나를 사용해서 30분만에 모델을 튜닝할 수 있었다는 사실은 이 방법론의 효율성을 더욱 부각시킵니다.

또한 IA3는 기존의 GPT-3가 'in-context learning' 방식을 사용했을 때보다도 더 좋은 성능을 보였다고 알려져 있습니다.

'In-context learning'은 주어진 context 내에서만 학습하는 방법으로, 한정된 context 안에서만 학습이 가능하다는 단점이 있습니다. 그러나 IA3는 이런 단점을 극복하면서도 더 높은 성능을 보여주었습니다.

이런 이유로 IA3는 뉴럴 네트워크를 더 효율적으로, 더 빠르게, 그리고 더 성능 좋게 튜닝하는 데 있어서 매우 유용한 도구로 간주되고 있습니다.

PEFT 성능비교 그래프
In Context Learning 성능비교 그래프

IA3도 LoRA와 마찬가지로 Linear Layer를 커스텀 구현체로 변경함으로써 구현이 가능하며,

LoRA Layer에 대한 configuration을 수정해서 구현할 수 있습니다. 다음은 IA3에 대한 구현체입니다.

  • configs/ia3.json
{
    "lora_scaling_rank": 1,
    "lora_rank": 0,
    "lora_init_scale": 0.0,
    "lora_modules": ".*SelfAttention|.*EncDecAttention|.*DenseReluDense",
    "lora_layers": "k|v|wi_1.*",
    "trainable_param_names": ".*lora_b.*",
    "model_modifier": "lora",
    "lr": 3e-3,
    "num_steps": 1000
}

LoRA기반 IA3 구현체

def modify_with_lora(transformer, config):
    for m_name, module in dict(transformer.named_modules()).items():
        if re.fullmatch(config.lora_modules, m_name):
            for c_name, layer in dict(module.named_children()).items():
                if re.fullmatch(config.lora_layers, c_name):
                    assert isinstance(
                        layer, nn.Linear
                    ), f"LoRA can only be applied to torch.nn.Linear, but {layer} is {type(layer)}."
                    setattr(
                        module,
                        c_name,
                        LoRALinear(layer, config.lora_rank, config.lora_scaling_rank, config.lora_init_scale),
                    )
    return transformer

class LoRALinear(nn.Module):
    def __init__(self, linear_layer, rank, scaling_rank, init_scale):
        super().__init__()
        self.in_features = linear_layer.in_features
        self.out_features = linear_layer.out_features
        self.rank = rank
        self.scaling_rank = scaling_rank
        self.weight = linear_layer.weight
        self.bias = linear_layer.bias
        if self.rank > 0:
            self.lora_a = nn.Parameter(torch.randn(rank, linear_layer.in_features) * init_scale)
            if init_scale < 0:
                self.lora_b = nn.Parameter(torch.randn(linear_layer.out_features, rank) * init_scale)
            else:
                self.lora_b = nn.Parameter(torch.zeros(linear_layer.out_features, rank))
        if self.scaling_rank:
            self.multi_lora_a = nn.Parameter(
                torch.ones(self.scaling_rank, linear_layer.in_features)
                + torch.randn(self.scaling_rank, linear_layer.in_features) * init_scale
            )
            if init_scale < 0:
                self.multi_lora_b = nn.Parameter(
                    torch.ones(linear_layer.out_features, self.scaling_rank)
                    + torch.randn(linear_layer.out_features, self.scaling_rank) * init_scale
                )
            else:
                self.multi_lora_b = nn.Parameter(torch.ones(linear_layer.out_features, self.scaling_rank))

    def forward(self, input):
        if self.scaling_rank == 1 and self.rank == 0:
            # parsimonious implementation for ia3 and lora scaling
            if self.multi_lora_a.requires_grad:
                # 이 부분에서 IA3의 Key, Value 값에 대한 rescaling이 이루어집니다.
                hidden = F.linear((input * self.multi_lora_a.flatten()), self.weight, self.bias)
            else:
                hidden = F.linear(input, self.weight, self.bias)
            if self.multi_lora_b.requires_grad:
								# 이 부분에서 position-wise feed-forward network의 값에 대한 rescaling이 이루어집니다.
                hidden = hidden * self.multi_lora_b.flatten()
            return hidden
        else:
            # general implementation for lora (adding and scaling)
            weight = self.weight
            if self.scaling_rank:
                weight = weight * torch.matmul(self.multi_lora_b, self.multi_lora_a) / self.scaling_rank
            if self.rank:
                weight = weight + torch.matmul(self.lora_b, self.lora_a) / self.rank
            return F.linear(input, weight, self.bias)

위 코드에서 **self.multi_lora_a**와 **self.multi_lora_b**는 IA3에서 사용되는 스케일링 벡터입니다. input * self.multi_lora_a.flatten() 연산에서 IA3의 Key, Value 값에 대한 rescaling이 이루어지고, hidden * self.multi_lora_b.flatten() 연산에서는 position-wise feed-forward network의 값에 대한 rescaling이 이루어집니다. 이러한 스케일링 작업을 통해 모델이 더 좋은 성능을 내도록 튜닝됩니다.

 

전반적으로 LoRA의 구현안에서 변형된 형태임을 확인할 수 있습니다.

 

코드

키 라이브러리

loralib

LoRA는 순위 분해 행렬 쌍을 학습하고 원래 가중치를 동결하여 학습 가능한 파라미터의 수를 줄입니다. 이를 통해 특정 작업에 맞게 조정된 대규모 언어 모델에 필요한 스토리지 요구량을 크게 줄이고 추론 지연 없이 배포 중에 효율적으로 작업을 전환할 수 있습니다. 또한 LoRA는 어댑터, 접두사 튜닝, 미세 튜닝 등 다른 여러 가지 적응 방법보다 성능이 뛰어납니다.

일부 파라미터만 학습하고 저장하면서 RoBERTa(Liu et al., 2019) 베이스 및 대형과 DeBERTa(He et al., 2020) XXL 1.5B를 사용하여 GLUE 벤치마크에서 전체 미세 조정과 비슷하거나 더 우수한 결과를 얻었습니다. 아래 숫자를 클릭하면 DeBERTa LoRA 체크포인트를 다운로드할 수 있습니다.

https://github.com/microsoft/LoRA

https://github.com/huggingface/peft

https://github.com/TimDettmers/bitsandbytes

논문

LoRA: Low-Rank Adaptation of Large Language Models

자연어 처리의 중요한 패러다임은 일반 도메인 데이터에 대한 대규모 사전 학습과 특정 작업 또는 도메인에 대한 적응으로 구성됩니다. 대규모 모델을 사전 학습할수록 모든 모델 파라미터를 재학습하는 전체 미세 조정은 실현 가능성이 낮아집니다. GPT-3 175B를 예로 들면, 각각 175B의 파라미터를 가진 미세 조정된 모델의 독립적인 인스턴스를 배포하는 것은 엄청난 비용이 소요됩니다. 유니티는 사전 학습된 모델 가중치를 동결하고 학습 가능한 순위 분해 행렬을 트랜스포머 아키텍처의 각 레이어에 주입하여 다운스트림 작업에서 학습 가능한 매개변수의 수를 크게 줄이는 저순위 적응(Low-Rank Adaptation, LoRA)을 제안합니다. Adam으로 미세 조정된 GPT-3 175B와 비교했을 때 LoRA는 훈련 가능한 파라미터 수를 10,000배, GPU 메모리 요구량을 3배까지 줄일 수 있습니다. LoRA는 훈련 가능한 파라미터 수가 더 적고, 훈련 처리량이 더 많으며, 어댑터와 달리 추가적인 추론 지연 시간이 없음에도 불구하고 RoBERTa, DeBERTa, GPT-2 및 GPT-3에서 모델 품질이 미세 조정과 동등하거나 더 나은 성능을 발휘합니다. 또한 언어 모델 적응의 순위 결핍에 대한 실증적 조사를 통해 LoRA의 효능을 조명합니다. 또한, LoRA를 PyTorch 모델과 쉽게 통합할 수 있는 패키지를 출시하고, https://github.com/microsoft/LoRA 에서 RoBERTa, DeBERTa, GPT-2에 대한 구현과 모델 체크포인트를 제공합니다.

Quickstart

  1. Installing loralib is simply
pip install loralib
# Alternatively
# pip install git+https://github.com/microsoft/LoRA

  1. You can choose to adapt some layers by replacing them with counterparts implemented in loralib. We only support nn.Linear for now. We also support a MergedLinear for cases where a single nn.Linear represents more than one layers, such as in some implementations of the attention qkv projection (see Additional Notes for more).
# ===== Before =====
# layer = nn.Linear(in_features, out_features)

# ===== After ======
import loralib as lora
# Add a pair of low-rank adaptation matrices with rank r=16
layer = lora.Linear(in_features, out_features, r=16)

  1. Before the training loop begins, mark only LoRA parameters as trainable.
model = BigModel()
import loralib as lora
# This sets requires_grad to False for all parameters without the string "lora_" in their name
lora.mark_only_lora_as_trainable(model)
# Training loop
for batch in dataloader:
   ...

  1. When saving a checkpoint, generate a state_dict that only contains LoRA parameters.
# ===== Before =====
# torch.save(model.state_dict(), checkpoint_path)
# ===== After =====
torch.save(lora.lora_state_dict(model), checkpoint_path)

  1. When loading a checkpoint using load_state_dict, be sure to set strict=False.
# Load the pretrained checkpoint first
model.load_state_dict(torch.load('ckpt_pretrained.pt'), strict=False)
# Then load the LoRA checkpoint
model.load_state_dict(torch.load('ckpt_lora.pt'), strict=False)

Now training can proceed as usual.

Reference

QLoRA: 48GB GPU로 65B 모델의 미세조정(파인튜닝)이 가능하다고요?

반응형

댓글

💲 Google Ads.