KoSimCSE Training on Amazon SageMaker

Usage

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoConfig, PretrainedConfig, PreTrainedModel
from transformers import AutoModel, AutoTokenizer, logging

class SimCSEConfig(PretrainedConfig):
    def __init__(self, version=1.0, **kwargs):
        self.version = version
        super().__init__(**kwargs)

class SimCSEModel(PreTrainedModel):
    config_class = SimCSEConfig

    def __init__(self, config):
        super().__init__(config)
        self.backbone = AutoModel.from_pretrained(config.base_model)
        self.hidden_size: int = self.backbone.config.hidden_size
        self.dense = nn.Linear(self.hidden_size, self.hidden_size)
        self.activation = nn.Tanh()

    def forward(
        self,
        input_ids: Tensor,
        attention_mask: Tensor = None,
        # RoBERTa variants don't have token_type_ids, so this argument is optional
        token_type_ids: Tensor = None,
    ) -> Tensor:
        # shape of input_ids: (batch_size, seq_len)
        # shape of attention_mask: (batch_size, seq_len)
        outputs: BaseModelOutputWithPoolingAndCrossAttentions = self.backbone(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )

        emb = outputs.last_hidden_state[:, 0]

        if self.training:
            emb = self.dense(emb)
            emb = self.activation(emb)

        return emb
    
def show_embedding_score(tokenizer, model, sentences):
    inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
    embeddings = model(**inputs)
    score01 = cal_score(embeddings[0,:], embeddings[1,:])
    score02 = cal_score(embeddings[0,:], embeddings[2,:])
    print(score01, score02)

def cal_score(a, b):
    if len(a.shape) == 1: a = a.unsqueeze(0)
    if len(b.shape) == 1: b = b.unsqueeze(0)
    a_norm = a / a.norm(dim=1)[:, None]
    b_norm = b / b.norm(dim=1)[:, None]
    return torch.mm(a_norm, b_norm.transpose(0, 1)) * 100 

# Load pre-trained model
model = SimCSEModel.from_pretrained("daekeun-ml/KoSimCSE-supervised-kobigbird-roberta-large") 
tokenizer = AutoTokenizer.from_pretrained("daekeun-ml/KoSimCSE-supervised-kobigbird-roberta-large")

# Inference example

s1 = """
대규모 언어 모델(LLM; Large Language Models)을 파인튜닝(Fine-tuning) 하면 오픈 소스 파운데이션 모델(Foundation model)을 개선하여 도메인별 작업에서 더욱 향상된 성능을 끌어낼 수 있습니다. 
이 게시물에서는 최신 오픈 소스 모델을 파인튜닝 하기 위해 Amazon SageMaker 노트북을 사용할 때의 이점에 관해 설명합니다. 
단일 노트북 인스턴스에서 대규모 모델의 대화형 파인튜닝 작업을 하기 위하여 허깅페이스(Hugging Face)의 효율적인 파라미터 파인튜닝(PEFT; Parameter-efficient Fine-tuning) 라이브러리와 bitsandbytes 패키지를 통한 양자화 기법을 활용합니다. 
본 게시물에서는 NVIDIA A10G GPU 4장이 탑재된 단일 ml.g5.12xlarge 인스턴스로 Falcon-40B 모델을 파인튜닝 하는 방법을 보여드리지만, 같은 전략으로 p4d/p4de 노트북 인스턴스에서 훨씬 더 큰 모델을 파인튜닝 할 수 있습니다.

보통 이러한 대규모 모델의 전체 정밀도 표현의 모델 파라미터는 한 장의 GPU나 여러 장의 GPU의 메모리에 모두 들어가지 않습니다. 
이러한 크기의 모델에 대한 파인튜닝하고 추론을 수행하는 대화형 노트북 환경을 위해 저희는 Quantized LLMs with Low-Rank Adapters (QLoRA) 라는 신규 기법을 사용합니다. 
QLoRA는 LLM의 메모리 사용량을 줄이면서도 견고한 성능을 유지하는 효율적인 파인튜닝 방식입니다. 
허깅페이스와 앞서 언급된 논문의 저자들은 기본 개념과 트랜스포머 및 PEFT 라이브러리와의 연동 내용을 다루는 자세한 블로그 게시물을 게시했습니다.
"""

s2 = """
대규모 언어 모델(LLM; Large Language Models) 분야의 발전과 LLM이 가치 있는 인사이트를 제공하는 문제 세트 수가 계속 증가하고 있다는 소식을 들어보셨을 겁니다. 
대규모 데이터 세트와 여러 작업을 통해 훈련된 대규모 모델은 훈련되지 않은 특정 작업에도 잘 일반화됩니다. 
이러한 모델을 파운데이션 모델(foundation model)이라고 하며, 스탠포드 HAI 연구소(Stanford Institute for Human-Centered Artificial Intelligence)에서 처음 대중화한 용어입니다. 
이러한 파운데이션 모델은 프롬프트 엔지니어링(prompt engineering) 기법의 도움으로 잘 활용할 수 있지만, 유스케이스가 도메인에 매우 특화되어 있거나 작업의 종류가 매우 다양하여 모델을 추가로 커스터마이징해야 하는 경우가 많습니다. 
특정 도메인이나 작업에 대한 대규모 모델의 성능을 개선하기 위한 한 가지 접근 방식은 더 작은 작업별 데이터 세트로 모델을 추가로 훈련하는 것입니다. 
파인 튜닝(fine-tuning)으로 알려진 이 접근 방식은 LLM의 정확도를 성공적으로 개선하지만, 모든 모델 가중치를 수정해야 합니다. 
파인 튜닝 데이터 세트 크기가 훨씬 작기 때문에 사전 훈련(pre-training)하는 것 보다 훨씬 빠르지만 여전히 상당한 컴퓨팅 성능과 메모리가 필요합니다. 
파인 튜닝은 원본 모델의 모든 파라미터 가중치를 수정하므로 비용이 많이 들고 원본 모델과 동일한 크기의 모델을 생성하므로 스토리지 용량에도 부담이 됩니다.

이러한 문제를 해결하기 위해 허깅페이스는 효율적인 파라미터 파인 튜닝(PEFT; Parameter-Efficient Fine-Tuning) 라이브러리를 도입했습니다. 
이 라이브러리를 사용하면 대부분의 원래 모델 가중치를 동결하고 훨씬 작은 추가 파라미터 세트만 훈련하여 모델 레이어를 교체하거나 확장할 수 있습니다. 
따라서 필요한 컴퓨팅 및 메모리 측면에서 훨씬 적은 비용으로 훈련을 수행합니다.
"""

s3 = """
2023년 6월부터 AWS 서울 리전에서 EC2 G5인스턴스를 사용할 수 있게 되었습니다. 
여기서는 Falcon Foundation Model을 Amazon SageMaker JumpStart를 이용해 AWS 서울 리전의 EC2 G5에 설치하고, 웹 브라우저 기반의 Chatbot을 생성하는 방법에 대해 설명합니다. 
Falcon FM은 HuggingFace의 Open LLM Leaderboard에서 상위권(2023년 7월 기준)에 위치할 만큼 우수한 성능을 가지고 있으면서도, 아파치 2.0 라이선스 정책에 따라 상용을 포함하여 누구나 자유롭게 사용할 수 있습니다.

Falcon FM은 SageMaker JumpStart를 이용하여 편리하게 설치하고, EC2 G5 인스턴스와 같은 우수한 인스턴스를 통해 AWS 서울 리전에서 개발 및 상용이 가능합니다. 
EC2 G5인스턴스는 EC2 G4dn 인스턴스보다 그래픽 집약적 애플리케이션 및 기계 학습 추론에 대해 최대 3배 더 높은 성능을 제공할 수 있어서, 
고성능을 요구하는 자연어 처리, 컴퓨터 비전 및 추천 엔진과 같은 분야에 적합합니다. 또한, LLM (Large Language Models) 기반의 Chatbot은 기존 Rule-based의 Chatbot에 비하여 훨씬 우수한 대화능력을 보여주며, 
AWS Lambda, Amazon CloudFront, AWS API Gateway, Amazon S3와 같은 AWS 인프라를 이용하여 쉽게 구현할 수 있습니다.

본 게시글에서는 SageMaker JumpStart를 통해 Falcon FM 기반의 Endpoint를 설치하고, AWS CDK를 이용하여 Chatbot를 위한 제공하기 위한 인프라를 준비합니다. 
생성된 Chatbot은 REST API를 통하여 텍스트로 요청을 하고 결과를 화면에 표시할 수 있습니다. 상세한 Architecture는 아래와 같습니다.
"""

sentences = [s1, s2, s3]
show_embedding_score(tokenizer, model.cpu(), sentences)

Introduction

SimCSE is a highly efficient and innovative embedding technique based on the concept of contrastive learning. Unsupervised learning can be performed without the need to prepare ground-truth labels, and high-performance supervised learning can be performed if a good NLI (Natural Language Inference) dataset is prepared. The concept is very simple and the psudeo-code is intuitive, so the implementation is not difficult, but I have seen many people still struggle to train this model.

The official implementation code from the authors of the paper is publicly available, but it is not suitable for a step-by-step implementation. Therefore, we have reorganized the code based on Simple-SIMCSE's GitHub so that even ML beginners can train the model from the scratch with a step-by-step implementation. It's minimalist code for beginners, but data scientists and ML engineers can also make good use of it.

Added over Simple-SimCSE

Requirements

We recommend preparing an Amazon SageMaker instance with the specifications below to perform this hands-on.

SageMaker Notebook instance

SageMaker Training instance

Datasets

For supervised learning, you need an NLI dataset that specifies the relationship between the two sentences. For unsupervised learning, we recommend using wikipedia raw data separated into sentences. This hands-on uses the dataset registered with huggingface, but you can also configure your own dataset.

The datasets used in this hands-on are as follows

Supervised

Unsupervised

How to train

Performance

We trained with parameters similar to those in the paper and did not perform any parameter tuning. Higher max sequence length does not guarantee higher performance; building a good NLI dataset is more important

{
  "batch_size": 64,
  "num_epochs": 1 (for unsupervised training), 3 (for supervised training)
  "lr": 3e-05,
  "num_warmup_steps": 0,
  "temperature": 0.05,
  "lr_scheduler_type": "linear",
  "max_seq_len": 64,
  "use_fp16": "True",
}

KLUE-STS

Model Avg Cosine Pearson Cosine Spearman Euclidean Pearson Euclidean Spearman Manhattan Pearson Manhattan Spearman Dot Pearson Dot Spearman
KoSimCSE-KoBigBird-RoBERTa-large (Supervised) 85.23 84.55 85.66 85.82 85.34 85.87 85.41 84.22 84.94

Kor-STS

Model Avg Cosine Pearson Cosine Spearman Euclidean Pearson Euclidean Spearman Manhattan Pearson Manhattan Spearman Dot Pearson Dot Spearman
KoSimCSE-KoBigBird-RoBERTa-large (Supervised) 85.30 85.02 85.69 84.95 85.31 84.94 85.34 85.37 85.77

References