728x90

reference
https://huggingface.co/blog/stackllama

하기 내용은 hugging face stackllama tech blog posting 내용을 학습 하고 정리한 것임

개요

#️⃣ LLaMA model

  • Meta AI에서 개발한 large language model
  • 7B ~ 65B 개의 parameter 사이즈, 1T ~ 1.4T 개 token으로 학습됨

StackLLaMA에서는 7B 사용

데이터 처리

#️⃣ Data

수집 데이터
https://huggingface.co/datasets/HuggingFaceH4/stack-exchange-preferences

Pasted image 20240106190055.png

전처리 데이터
https://huggingface.co/datasets/lvwerra/stack-exchange-paired

Pasted image 20240106190137.png

StackExchange 사이트에서 질문과 답변 데이터셋 사용
답변에 추천 수와 채택 답변의 라벨까지 포함되어있음

점수 계산방법

score = log2 (1 + upvotes) rounded to the nearest integer, plus 1 if the questioner accepted the answer (we assign a score of −1 if the number of upvotes is negative)

  • reward model에서 질문에 대해 2개의 답변을 비교
  • 답변이 많은 경우 -> 최대 10쌍의 답변으로 제한
  • HTML을 Markdown으로 바꿈

학습

#️⃣ Efficient training starategies

가장 작은 LLaMA 모델 학습시키는데에도 큰 메모리가 필요함

AI에서 weight을 위한 data type

Datatype Remark
bf16 bfloat16 16-bit 부동 소수점 형식. fp32 대비 메모리 절반만 사용
fp16 half-precision 16-bit 부동 소수점 형식. predict에 사용
fp32 training에 사용. 연산이 오래 걸리고 메모리를 많이 사용

bf16에서

  • Adam optimizer에서 사용되는 8바이트 외에 추가로
  • 파라미터 당 2 bytes (fp32에서는 4 bytes)가 필요

7B 파라미터 모델은 (2+8)*7B=70GB 메모리가 필요하고, 어텐션 스코어 연산에 메모리가 더 필요함

  • 80GB A100 1개로도 학습이 불가능함
  • efficient optimizers of half-precision training 로 메모리를 일부 절약할 수 있지만 충분하지 않음

다른 대안으로는
Parameter-Efficient Fine-Tuning (PEFT)
https://github.com/huggingface/peft

=> Low-Rank Adaptation (LoRA)로 8-bit에 모델을 올리는 방법이 있음

오렌지색의 Low-Rank Adaptation of linear layers 는 파란색의 frozen layer 옆에 붙고
결과가 인코딩된 hidden state들은 frozen layer의 hidden states랑 합쳐지게 됨

모델을 8bit에 로딩하는 것은 메모리 공간을 절약해서
가중치에 파라미터당 1 byte만 필요하게 됨 (7B LLaMa는 7GB 메모리)

LoRA는 원본 가중치를 직접 업데이트하지 않고, 특정 레이어의 위에 작은 어댑터 레이어를 추가함
frozen layer만 빼고 학습시키게되니까 -> 학습 파라미터 수가 많이 감소

배치사이즈와 시퀀스 길이에 따라 다르지만 1B(10억)개 파라미터당 ~1.2-1.4GB를 할당해 전체 파인튜닝 세팅을 함
=> 이렇게 되면 50-60B 크기의 모델을 NVIDIA A100 80GB 1개로 파인튜닝이 가능해짐

PEFT 기법으로 아래와 같은 학습이 가능함
fine-tuning facebook/opt-6.7b (13GB in float16 )
openai/whisper-large on Google Colab (15GB GPU RAM)

학습속도는 data parallelism(같은 트레이닝 세팅을 여러대의 GPU에 복제하고 각 GPU에 다른 배치를 보냄)으로 개선 가능
이 방법으로 모델의 forward, backward pass를 병렬화하고 스케일링할 수 있음

  • transformers.Trainer 나 accelerate로 병렬화해서 torchrun 혹은 accelerate launch를 arguments로 전달하면 병렬화가 됨
  • 하기 명령어는 머신 1대에 8개의 GPU로 학습시키는 것임
  • accelerate: accelerate launch --multi_gpu --num_machines 1 --num_processes 8 my_accelerate_script.py
  • torchrun: torchrun --nnodes 1 --nproc_per_node 8 my_torch_script.py

#️⃣ Supervised fine-tuning

RLHF 이전의 fine-tuning은 일반적인 언어 모델링 과정임

Packing (효율적 데이터 사용을 위한 기술)
배치에서 샘플당 하나의 텍스트를 가지고 가장 긴 텍스트에 맞춰 패딩을 하는 대신
사이 사이에 EOS 토큰을 끼워 많은 텍스트를 concat하고 context 사이즈 크기만큼으로 잘라서 사용하는 것 (패딩 X)

  • padding 토큰이 loss에서 마스킹되는것 대비 효율적인 학습 가능
  • 데이터가 적거나 overflowing 토큰이 잘리는 것에 회의적이라면 기존 dataloader 사용 권장

packing은 ConstantLengthDataset 으로 사용 가능
peft로 모델 로딩 후에 Trainer 사용 가능
모델을 int8로 로딩하고, training을 위해 준비한다음에 LoRA 어댑터를 추가

# load model in 8bit
model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        load_in_8bit=True,
        device_map={"": Accelerator().local_process_index}
    )
model = prepare_model_for_int8_training(model)

# add LoRA to model
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, config)

LLAMA base model weights 사용 요청 한 다음에
https://docs.google.com/forms/d/e/1FAIpQLSfqNECQnMkycAp2jP4Z9TFX0cGR4uf7b_fBxjY_OjhJILlKGA/viewform

허깅페이스 포맷으로 변환하여 파인튜닝 가능
https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/convert_llama_weights_to_hf.py

#️⃣ Reward modeling and human preferences

human annotation이 있는 RLHF를 사용해 모델을 fine-tuning 하는 것은 고비용이 듦
이러한 다이렉트 피드백을 대체하는 트릭으로 강화학습 전에 수집된 human annotation에 기반한 reward model을 학습시키는 방법이 있음

reward model의 목적은 인간이 text에 대해 점수를 어떻게 매길지에 대해 모방하는 것
가장 직접적인 방법은 점수나, "좋음"과 "나쁨"의 이진 분류로 annotation을 예측하는 것임

두 개 예시의 랭킹을 예측하는 것이 잘 동작할 것인데,
reward model이 주어진 프롬프트 에 대해 () 두개의 후보를 제시하면 human annotator에 의해 어떤 것이 더 좋게 평가될 것인지 예측하는 방식임

loss function 수식

  • : 모델의 스코어
  • : 바람직한 대답(후보)

StackExchange 데이터셋에서 스코어에 따라 유저가 어떤 답변을 더 선호하는지 추론할 수 있음
이 정보와 loss function으로 transformers.Trainer에 커스텀 loss function을 추가해서 아래와 같이 수정

class RewardTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        rewards_j = model(input_ids=inputs["input_ids_j"],  attention_mask=inputs["attention_mask_j"])[0]
        rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
        if return_outputs:
            return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
        return loss
  • StackLLaMA 팀에서는 10만 쌍의 후보 subset을 이용해 5만쌍의 보류 집합을 평가 했음
  • 배치사이즈는 4
  • LLaMA 모델을 LoRA peft 어댑터로 단일 에폭에 대해 Adam 옵티마이저(BF16 정확도)로 학습시켰음

LoRA 설정

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)
  • A100 GPU 8대로 몇 시간동안 학습 시켰고, 최종 accuracy는 67%
  • 어댑터의 결과는 frozen model에 merge되어 향후에 사용됨

#️⃣ Reinforcement Learning from Human Feedback

강화학습 반복 과정

  1. 프롬프트로부터 결과 생성
  2. reward model로 결과를 평가
  3. 평가 점수로 policy-optimization을 학습시키는 강화학습
Question: <Query>

Answer: <Response>
  • 쿼리와 결과 프롬프트는 tokenized되고 모델로 전달 되기 전 위와 같이 템플릿화
  • SFT, RM RLHF 단계에서도 같은 템플릿이 사용됨
  • 언어모델을 강화학습으로 학습시킬 때 흔한 문제는 reward model이 횡설수설하는 대답을 생성하는 것에 높은 점수를 매기는 사태임
  • 이러한 문제로부터 밸런스를 잡기 위해 reward에 페널티를 적용함
  • train하지 않는 모델의 레퍼런스를 유지하고 KL-divergence를 계산하여 새 모델의 generation을 레퍼런스 모델과 비교하는 방식
  • : reward model의 reward
  • : 현재 policy와 레퍼런스 모델간의 KL-Divergence

레퍼런스 모델과 policy는 베이스 SFT 모델을 공유함
8-bit로 로딩해서 training 동안 freeze하고, PPO를 사용해 policy의 LoRA weight만 최적화함

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    question_tensors = batch["input_ids"]
        
    # sample from the policy and generate responses
    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

    # Compute sentiment score
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
    rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]

    # Run PPO step
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
    # Log stats to WandB
    ppo_trainer.log_stats(stats, batch, rewards)

평가

#️⃣ 높은 reward가 더 좋은 성능인지 여부

  • 강화학습에서는 일반적으로 가장 높은 보상을 달성하는 것이 목표
  • RLHF에서는 reward model을 사용하는데, 보상모델은 불완전하고 PPO 알고리즘이 이 불완전을 악용할 수 있음
  • 그래서 reward가 갑자기 올랐지만 확인했을 때 ‘’’ 같은 것이 반복되는 현상이 나오기도 함
  • 이러한 문제는 KL 페널티를 통해 보완 가능

#️⃣ KL은 항상 양수인지 여부

  • KL 페널티는 모델 출력이 base policy의 출력에 가깝도록 유지하기 위함
  • KL divergence는 이 두 개의 분포 거리를 계산하고 항상 양수임
  • 하지만 trl에서는 실제 KL divergence와 같은 KL estimate를 사용함

  • SFT 모델보다 낮은 확률의 policy에서 토큰이 샘플링되면 마이너스 KL 페널티가 발생함

    • 평균적으로 KL 페널티는 양수임 (그렇지 않다면 policy에서 샘플링이 적절히 되기 어려움)

  • 전략으로 어떤 토큰은 생성되고, 어떤 토큰은 suppress되게 할 수 있음

    • 배치생성시에는 시퀀스가 패딩되고
    • 최소 길이를 설정했을 때는 EOS 토큰이 억제됨

  • 모델은 이런 경우에 토큰에 매우 높거나 낮은 확률을 하게 되어 마이너스 KL의 결과가 나올 수 있음

  • PPO 알고리즘이 reward를 최적화하기때문에 이러한 마이너스 페널티로 불안정성이 유발되기도 함

따라서 샘플링시 단순한 샘플링을 먼저 시도해볼 것을 권장

그 외에 학습시 loss의 spike가 종종 발생하는데 불안정성을 유발할 수 있기 때문에 해결해야 할 것임

#️⃣ 결론

  1. peft를 사용해 단일 GPU에서도 느리지만 학습이 가능하고, GPU를 추가한한다면 스케일링해서 코드 변경 없이 데이터 병렬 학습이 가능
  2. 실제 유즈케이스로, 학습 모델이 준비되면 다른 모델과 버전별로 순위를 매기는 등 비교 평가 해야함.
  3. evaluation 과정을 추가하는 경우, training 세팅을 수정하는 것을 반복해 모델이 개선될 방법을 찾으면 됨. 다른 데이터를 합치거나, 기존에 필터를 추가하는 방식 적용 가능. 반대로 reward model 아키텍처나 모델 사이즈가 다른 것을 시도해 볼 수 있음
반응형

'AI > LLM' 카테고리의 다른 글

LLM의 Reinforcement Learning & RLHR & DPO  (0) 2024.01.05

티스토리툴바