[KoE5-v1.0] 최초의 한국어 특화 임베딩 모델 (Multilingual E5 Finetune) 2024.10.15 updated
이번 연구실 과제로 한국어 임베딩 모델을 만들게 되었다. 먼저 작은 모델을 finetuning 하는 것부터 시작하여, 나중에는 연구실에서 오래 두고 쓸 수 있는 성능 좋은 임베딩모델을 만들기로 계획했다.
그중, 이번에는 Microsoft의 Multilingual-e5-large 모델을 한국어로 fine-tuning한 방법, 결과, 트러블슈팅 과정 등을 남겨보고자 한다.
[24.10.15 updated]
KoE5-v1.0 모델 및 ko-triplet-v1.0 데이터셋, 논문이 출시되었습니다 !!
다음 링크에서 확인 가능합니다. 많은 사용 후기 기다리겠습니다 !
1. E5의 구조
먼저, 시작하기에 앞서, E5라는 임베딩모델은 어떻게 만들어졌는지에 대해 간략하게 살펴보자. 논문을 보면 E5는 BERT를 기반으로, 다음과 같이 2단계를 걸쳐 만들어졌다.
- Unsupervised Large-Scale Text Dataset 을 구축하여 pre-train
- MS-Marco(Retrieval), NLI (Similarity), NQ 데이터 셋으로 fine-tune
각각의 단계에 관해 간단하게 살펴보자.
1–1. Pre-train
- 본 논문에서는 CCPairs라는 데이터셋을 자체 구축한다. CCPairs는 Reddit, Stackexchange, Scientific papers, Common Crawl 등으로부터 가져온 query-passage 쌍으로 이루어진 1.3B 크기의 데이터셋이다.
- 논문에서는 먼저 1.3B의 노이즈가 많은 텍스트 쌍을 가지고 모델을 pre-train하고, 학습된 모델을 사용하여 각 텍스트 쌍을 100만 개의 random passage와 비교하여 랭킹을 매긴다. 이후, top 2 랭크 passage를 남기고, 모델의 예측이 label과 일치하는 경우에만 passage를 남긴다.
- 쉽게 풀어 설명하면, 1.3B의 query-passage 쌍이 있을 때, 모델을 사용하여 하나의 query와 100만개의 다른 passage를 비교한다. 이후, 가장 유사도가 높은 top 2 passage를 가져와, 만약 처음 이루어진 query-passage쌍의 passage가 그 2개 안에 포함돼 있으면, query와 높은 유사도를 보이는 passage로 판단하고, 해당 데이터셋을 퀄리티가 높은 데이터셋으로 선정한다.
- 이렇게 필터링 과정을 거치면, 1.3B의 noise 쌍이 270M까지 줄게 된다.
- 이후, 정제된 270M 크기의 데이터셋으로 모델을 contrastive pre-training 한다. 이때, 다음과 같은 InfoNCE contrastive loss를 사용한다.
- 식에서 negative 쌍을 mining 하기 위해서는 in-batch negative를 사용하는데, 32,768이라는 충분히 큰 batch size를 사용하여 negative를 선정한다고 한다.
1–2. Fine-tune
- pre-training 이후 fine-tuning을 진행하는데, 이때 MS-MARCO와 NQ dataset을 사용한다.
- 여기서는 mined hard negatives과 knowledge distillation from a cross-encoder (CE) teacher model을 사용하여 loss function을 구축하는데, 식은 다음과 같다.
1–3. Results
이러한 방식을 이용해서 small, base, large의 3가지 크기 model을 만들었다.
이 모델들은 링크에서 확인해 볼 수 있다.
2. Finetune Multilingual-E5
이제 모델에 대해서 알아 보았으니, fine-tuning 과정을 살펴보자.
우리는 한국어 임베딩 모델을 만들 것이니, 다국어 임베딩 모델인 Multilingual-E5를 사용할 예정이다. (multilingual-e5은 BERT 기반이 아닌, xlm-roberta 기반 모델이다.)
2–1. Dataset: hard negative mining & filtering
기업 과제로 해당 프로젝트를 진행했기 때문에, 어떤 데이터셋을 사용했는지는 밝힐 수 없지만, 10종의 오픈된 데이터와, 2종의 비공개 데이터를 사용했다.
본 프로젝트에서는 데이터를 모아, 유사한 query-passage 쌍을 만든 후, hard negative를 mining하고 filtering하는 과정을 거쳤다.
여기서 hard negative이란, 주어진 기준 데이터와 유사하지는 않지만, 표면적으로는 그럴듯하게 보이는 데이터를 뜻한다. 예를 들어, query가 “고양이는 포유류이다.”의 random negative는 “나는 어제 바다에서 수영을 했다." 이고, hard negative는 “고양이는 수염이 있으며, 유연한 동물이다." 정도가 될 수 있다.
Hard negative mining을 위해서는
- query와 documents를 형태소 단위로 분절하여, BM25 점수가 높은 documents를 1차적으로 mining 했다.
- 그 중 Multilingual-E5-large 모델을 사용하여 각 query, documents를 임베딩한 후, 유사도가 높은 top3 데이터를 mining하였다 (gold 제외). 가능한 random negative가 아닌 hard negative가 mine되게끔 하기 위해 gold가 아닌 document 중 query와의 유사도가 가장 높은 document를 mining하여 해당 데이터를 query의 hard negative으로 지정하였다.
이후, 데이터를 정제하기 위해 Filtering 과정을 거쳤는데, 본 과정은
“query와 유사하지 않은 document(positive)를 필터링하고, negatives 중 random negative와 positive에 해당하는 경우(위음성)를 필터링” 하기 위함이다.
Filtering을 위해서는
- 모든 query-positive, query-negative에 대한 유사도를 구하여,
- 점수의 분포를 4개의 구간으로 나누었다.
점수의 1/4 되는 곳을 q1, 1/2 되는 곳을 q2, 3/4 되는 곳을 q3로 지정하여, - positive의 경우, 유사도가 낮은 데이터를 없애기 위해 [0, q1] 구간에 해당하는 데이터들을 잘라냈고,
- negative의 경우, 유사도가 너무 낮은 random negative 또는 유사도가 너무 높은 false negative 데이터를 없애기 위해 [0, q1], [q3, 1] 구간에 해당하는 데이터셋을 잘라냈다.
최종적으로, 데이터셋의 구성은 다음과 같다.
{
"query": "질문",
"document": "query와 관련된, 유사도가 어느 정도 높은 document",
"hard_negative": "query와 적당히 무관한, mined negative"
}최종 데이터의 크기는 800M 정도이다. (파일 크기 기준)
2–2. Train 파이프라인 구현
지난 글에서 다룬 대로, Huggingface의 Trainer를 사용하여 구현해 보고자 하였다.
구현체들이 어떻게 되어 있는지 찾아보고자, ClusterLLM에서 다룬 구현체를 참고했다. 본 코드에서는 forward 과정에서 query, positive, negative의 input_ids, attention_mask를 먼저 받아와 forward를 진행한다.
그럼 이제 코드를 살펴보자. 먼저 폴더 구조는 다음과 같고,
train.py에서 준비한 데이터셋을 로드한다.
# train.py
tokenizer = AutoTokenizer.from_pretrained(
model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
cache_dir=cache_dir,
use_fast=model_args.use_fast_tokenizer,
revision=model_args.model_revision,
use_auth_token=True if model_args.use_auth_token else None,
)
model = AutoModel.from_pretrained(model_name_or_path)
logger.info("Loading train_dataset...")
train_dataset = KoE5Dataset(
args=data_args,
tokenizer=tokenizer,
mode='train',
cache_dir=cache_dir,
test=test,
)
logger.info("Loading eval_dataset...")
eval_dataset = KoE5Dataset(
args=data_args,
tokenizer=tokenizer,
mode='dev',
cache_dir=cache_dir,
test=test
)
data_collator = DataCollatorForKoE5(
tokenizer=tokenizer,
padding=True,
max_length=None,
pad_to_multiple_of=None,
label_pad_token_id=-100,
return_tensors="pt",
)- 처음에 tokenizer와 model를 정의하고, KoE5Dataset이라는 class를 사용하여 데이터셋을 불러온다. KoE5Dataset에 대해서 살펴보자.
# dataset.py
class KoE5Dataset(Dataset):
args: KoE5DataTrainingArguments
features: List[InputFeatures]
def __init__(
self,
args: KoE5DataTrainingArguments,
tokenizer: PreTrainedTokenizerBase,
limit_length: Optional[int] = None,
mode: Union[str, Split] = Split.train,
cache_dir: Optional[str] = None,
test: Optional[bool] = False,
):
self.args = args
self.processor = KoE5MRCProcessor()
if isinstance(mode, str):
try:
mode = Split[mode]
except KeyError:
raise KeyError("mode is not a valid split name")
# Load data features from cache or dataset file
cached_features_file = os.path.join(
cache_dir if cache_dir is not None else args.data_dir,
f"cached_{mode.value}_{tokenizer.__class__.__name__}_{args.max_seq_length}",
)
logger.debug(f"cache file exits: {os.path.exists(cached_features_file)}")
if os.path.exists(cached_features_file) and not args.overwrite_cache:
logger.info("Start loading features...")
start = time.time()
self.features = torch.load(cached_features_file)
logger.info(
f"Loading features from cached file {cached_features_file} [took %.3f s]", time.time() - start
)
else:
lock_path = cached_features_file + ".lock"
with FileLock(lock_path):
logger.info(f"No cache files. Creating features from dataset file at {args.data_dir}")
if mode == Split.dev:
examples = self.processor.get_dev_examples(args.data_dir)
elif mode == Split.test:
examples = self.processor.get_test_examples(args.data_dir)
else:
examples = self.processor.get_train_examples(args.data_dir)
if test:
examples = examples[:100]
logger.info("Test mode activated: Got 100 examples!")
self.features = convert_examples_to_features(examples, tokenizer, args.max_seq_length)
logger.info("Converted examples to features!")
start = time.time()
torch.save(self.features, cached_features_file)
logger.info(
f"Saving features into cached file {cached_features_file} [took {time.time() - start:.3f} s]"
)
def __len__(self):
return len(self.features)
def __getitem__(self, i) -> InputFeatures:
return self.features[i]
def get_labels(self):
return [1, 0]- 이 코드에서는 cache 파일이 존재할 경우, 해당 파일을 불러와 features라는 변수에 저장하고, 없으면 processor.get_train_examples를 통해 example을 불러온다.
- 이때 processor은 KoE5MRCProcessor라는 클래스를 정의했는데, 이를 살펴보자.
# processor.py
class E5InputExample(InputExample):
def __init__(self, query: str, positive_passage: str, negative_passage: str):
super().__init__(
f"query: {query}",
f"passage: {positive_passage}",
f"passage: {negative_passage}"
)
@dataclass(frozen=True)
class InputFeatures:
# question
input_ids: List[int]
token_type_ids: List[int]
attention_mask: Optional[List[int]] = None
def to_json_string(self):
"""Serializes this instance to a JSON string."""
return json.dumps(dataclasses.asdict(self)) + "\n"
class KoE5MRCProcessor(DataProcessor):
"""Processor for the KoE5 dataset."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# warnings.warn(DEPRECATION_WARNING.format("processor"), FutureWarning)
def get_example_from_tensor_dict(self, tensor_dict):
"""See base class."""
return InputExample(
tensor_dict["idx"].numpy(),
tensor_dict["sentence"].numpy().decode("utf-8"),
None,
str(tensor_dict["label"].numpy()),
)
def get_train_examples(self, data_dir):
"""See base class."""
if "train.json" in os.listdir(data_dir):
return self._create_examples(self._read_json(os.path.join(data_dir, "train.json")), "train")
else:
return self._create_examples(self._read_jsonl(os.path.join(data_dir, "train.jsonl")), "train")
def _create_examples(self, datas, set_type):
"""Creates examples for the training, dev and test sets."""
examples = []
for i, data in enumerate(datas):
if isinstance(data["query"], list):
query = data["query"][0]
else:
query = data["query"]
if isinstance(data["document"], list):
document = data["document"][0]
else:
document = data["document"]
if "hard_negative" in data:
if isinstance(data["hard_negative"], list):
hard_negative = data["hard_negative"][0]
else:
hard_negative = data["hard_negative"]
else:
hard_negative = None
examples.append(E5InputExample(
query=query,
positive_passage=document,
negative_passage=hard_negative,
))
return examples
@classmethod
def _read_json(cls, input_file):
"""Reads a tab separated value file."""
with open(input_file, "r", encoding="utf-8-sig") as f:
return json.load(f)
def _read_jsonl(file_path, input_file):
"""Read a JSONL file and return a list of dictionaries."""
data = []
with open(input_file, 'r', encoding='utf-8-sig') as file:
for line in file:
data.append(json.loads(line))
return data- get_train_examples 부분을 살펴보면, train.json 또는 train.jsonl 이라는 명칭의 파일을 불러와, _create_examples 함수로 전달한다.
- _create_examples에서는 데이터로부터 query, document, hard_negative를 가져와 mE5-large에서 지정한 형태로 저장한다. (참고)
- 그렇게 examples가 반환되면, convert_examples_to_features() 함수로 전달되어 해당 값이 feature로 저장된다. convert_examples_to_features를 살펴보자.
# processor.py
def convert_examples_to_features(
examples: List[InputExample],
tokenizer: PreTrainedTokenizer,
max_length: Optional[int] = None,
):
if max_length is None:
max_length = tokenizer.model_max_length
query_batch_encoding = tokenizer(
[example.query for example in examples],
max_length=max_length,
padding="max_length",
truncation=True,
)
document_batch_encoding = tokenizer(
[example.positive_passage for example in examples],
max_length=max_length,
padding="max_length",
truncation=True,
)
negative_batch_encoding = tokenizer(
[example.negative_passage for example in examples],
max_length=max_length,
padding="max_length",
truncation=True,
)
features = []
from tqdm import tqdm
for i in tqdm(range(len(examples)), desc="Converting examples to features..."):
query_inputs = {k: query_batch_encoding[k][i] for k in query_batch_encoding}
doc_inputs = {k: document_batch_encoding[k][i] for k in document_batch_encoding}
neg_inputs = {k: negative_batch_encoding[k][i] for k in negative_batch_encoding}
features.append({
'query_input_ids': query_inputs['input_ids'],
'query_attention_mask': query_inputs['attention_mask'],
'document_input_ids': doc_inputs['input_ids'],
'document_attention_mask': doc_inputs['attention_mask'],
'hard_negative_input_ids': neg_inputs['input_ids'],
'hard_negative_attention_mask': neg_inputs['attention_mask'],
})
for i, example in enumerate(examples[:3]):
logger.info("*** Example ***")
logger.info(f"features: {features[i]}")
return features- convert_examples_to_features 함수에서는 각각의 query, document, hard_negative를 tokenize하여 input_ids와 attention_mask를 따로 features에 저장한다. (앞서 첨부한 ClusterLLM에서 진행한 방식을 차용했다.)
- 이후, 이 features는 DataCollatorForKoE5로 전달된다.
# data_collator.py
class DataCollatorForKoE5(DataCollatorMixin):
tokenizer: PreTrainedTokenizerBase
padding: Union[bool, str, PaddingStrategy] = True
max_length: Optional[int] = None
pad_to_multiple_of: Optional[int] = None
label_pad_token_id: int = -100
return_tensors: str = "pt"
def torch_call(self, features):
import torch
label_name = "label" if "label" in features[0].keys() else "labels"
labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
# 정규 형식 (input_ids, attention_mask 로 처리한 뒤 다시 prefix (query, document, hard_negative) 붙여줌
cat_batch = {}
for k in ['query', 'document', 'hard_negative']:
batch_features = [
{
'input_ids': feature[f'{k}_input_ids'],
'attention_mask': feature[f'{k}_attention_mask']
} for feature in features if f'{k}_input_ids' in feature and f'{k}_attention_mask' in feature
]
if batch_features:
batch = pad_without_fast_tokenizer_warning(
self.tokenizer,
batch_features,
padding=self.padding,
max_length=self.max_length,
pad_to_multiple_of=self.pad_to_multiple_of,
return_tensors="pt",
)
cat_batch[f'{k}_input_ids'] = batch['input_ids']
cat_batch[f'{k}_attention_mask'] = batch['attention_mask']
if labels is None:
return cat_batch- 이전 단계에서 query-docuement-negative의 input_ids와 attention_mask를 따로 저장했다면, 이 단계에서는 batch마다의 {query, document, hard_negative}의 {input_ids, attention_mask}를 따로 저장한다.
- 각각의 key값 이름은 query_input_ids, query_attention_mask, document_input_ids, document_attention_mask, hard_negative_input_ids, hard_negative_attention_mask 가 된다.
- 이렇게 batch별로 저장된 cat_batch는 train.py에서 data_collator로 반환된다.
2–3. CustomTrainer
이제 데이터는 모두 준비됐고,최종적으로 학습 전에 loss를 정의한다.
# train.py 이어서
trainer = CustomTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
tokenizer=tokenizer,
data_collator=data_collator
)
checkpoint = None
if training_args.resume_from_checkpoint is not None:
checkpoint = training_args.resume_from_checkpoint
elif last_checkpoint is not None:
checkpoint = last_checkpoint
if model_args.init_checkpoint is not None:
print(f"Loading from {model_args.init_checkpoint} ...")
state_dict = torch.load(os.path.join(model_args.init_checkpoint, 'pytorch_model.bin'))
model.load_state_dict(state_dict)
train_result = trainer.train(resume_from_checkpoint=checkpoint)
trainer.save_model(output_dir=training_args.output_dir)- 여기서는 fine-tuning할 loss함수를 따로 지정해주기 위해 Trainer를 custom한 CustomTrainer를 정의한다.
- 우리의 목적은 위 Loss 식을 구현하는 것이므로, Trainer의 compute_loss 부분을 다음과 같이 override한다.
class CustomTrainer(Trainer):
def __init__(self, model, args, train_dataset, eval_dataset, tokenizer, data_collator):
super().__init__(model=model, args=args, train_dataset=train_dataset, eval_dataset=eval_dataset,
tokenizer=tokenizer, data_collator=data_collator)
def compute_loss(self, model, inputs, return_outputs=False):
embeddings = {}
for k in ['query', 'document', 'hard_negative']:
input_ids = inputs[f'{k}_input_ids']
attention_mask = inputs[f'{k}_attention_mask']
input_dict = {
'input_ids': input_ids,
'attention_mask': attention_mask,
'token_type_ids': torch.zeros_like(input_ids)
}
output: BaseModelOutput = model(**input_dict)
pooled_output = average_pool(output.last_hidden_state, input_dict['attention_mask'])
embeddings[k] = F.normalize(pooled_output, p=2, dim=1)
# InfoNCE Loss
query_embeddings = embeddings['query']
positive_embeddings = embeddings['document']
negative_embeddings = embeddings['hard_negative']
similarity_fct = nn.CosineSimilarity(dim=-1)
tau = self.args.cl_temperature
positive_scores = similarity_fct(query_embeddings, positive_embeddings) / tau
negative_scores = similarity_fct(query_embeddings.unsqueeze(1), negative_embeddings) / tau
max_positive_scores = torch.max(positive_scores, dim=0, keepdim=True)[0]
max_negative_scores = torch.max(negative_scores, dim=1, keepdim=True)[0]
max_scores = torch.max(max_positive_scores, max_negative_scores)
stable_positive_scores = positive_scores - max_scores
stable_negative_scores = negative_scores - max_scores.unsqueeze(1)
exp_positive_scores = torch.exp(stable_positive_scores) # 분자
exp_negative_scores = torch.exp(stable_negative_scores)
total_scores_sum = exp_positive_scores + exp_negative_scores.sum(dim=1) # 분모
log_prob = torch.log(exp_positive_scores / total_scores_sum)
loss = -log_prob.mean()
return loss2–4. Train
이제 효율적인 train 자동화를 위해 script를 작성한다.
본 논문의 후속 논문인 Multilingual E5 Text Embeddings: A Technical Report를 보면, fine-tuning hyper parameters에 관해 다음과 같이 설명돼 있다.
For fine-tuning, we use batch size 512 and learning rate {3, 2, 1}×10^(−5) for the {small, base, large} models. All models are fine-tuned for 2 epochs.
여기서 얻은 하이퍼파라미터를 기반으로 script파일을 작성한다.
(메모리 이슈로 batch size를 512까지 높이지는 못했다.)
# scripts/finetune.sh
EPOCH=2
LR=1e-5
BATCH_SIZE=32
DATE=
export WANDB_PROJECT="KoE5"
export WANDB_NAME="KoE5-large-v1.2-InfoNCE-bs=${BATCH_SIZE}-ep=${EPOCH}-lr=${LR}-${DATE}"
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 torchrun --nproc_per_node=8 train.py \
--model_name_or_path intfloat/multilingual-e5-large \
--output_dir /data/KoE5/MODELS/${WANDB_NAME} \
--data_dir /data/KoE5/DATA/datasets/ \
--cache_dir projects/KoE5/cache \
--num_train_epochs $EPOCH \
--learning_rate $LR \
--per_device_train_batch_size $BATCH_SIZE \
--per_device_eval_batch_size $BATCH_SIZE \
--warmup_steps 100 \
--logging_steps 2 \
--save_steps 100 \
--cl_temperature 0.02 \
--test False- 이렇게 코드를 완성한 후, “bash scripts/finetune.sh”를 실행하면 학습이 시작된다 !!
3. 결과
3-1. Ko-StrategyQA와 자체 구축 벤치마크
학습을 마치고, 완료된 모델의 평가를 위해 MTEB의 Ko-StrategyQA와 계약 기업에서 자체적으로 구축한 private Benchmark 2개를 사용했다.
Ko-StrategyQA는 위키피디아 기반의, 한국어 multi-hop 질문들을 해결하기 위해 복수 개의 단락을 자동으로 단락 뭉치에서 검색하고, 찾아낼 수 있는 성능을 측정하는 데이터셋이다. 평가 메트릭으로는 ndcg@5을 사용했다.
ndcg란 검색 시스템의 성능을 측정하는 지표로, 검색 결과의 순서와 관련성을 반영한다. 연관성이 높은 문서가 상위에 있을 때 더 높은 점수를 부여하는데, ndcg@5는 top-5개의 문서를 가지고 ndcg 점수를 측정한 것이다.
3가지 벤치마크 데이터셋에 대해 평균을 산출한 결과는 다음과 같았다.
베이스라인 모델인 Multilingual-E5-large에 비해 약 3점 정도 향상된 성능을 보였다 !!
3-2. AutoRAG-example-korean-embedding-benchmark
추가) velog에 올라온 글을 보던 중, Markr 라는 AutoRAG를 만든 회사에서 임베딩모델 벤치마크를 만들었다는 사실을 알게 되었다. 블로그
블로그를 보자마자 바로 테스트해보고 싶어, repository를 fork해서 바로 모델을 평가해보았다. 결과는 다음과 같았다.
결과적으로 top-k 1,3 (그리고 첨부는 안했지만 5, 10, 50)에서 모두 매우 좋은 성능을 보여주었다 !!
(다시 한번 좋은 벤치마크를 공유해주신 Markr, Allganize 팀께 감사드립니다 🙇♂️🙇♂️)
3–3. 결과 분석
사실 다른 결과들은 다 잘 나와주어서 분석할 것이 없지만, KoE5-v1.1이 Ko-Strategy에서 베이스라인 모델보다 낮은 결과를 보여, 왜 그런지 생각해 보았다.
정확하지 않을 수 있겠지만, 아마도 Ko-StrategyQA도 위키피디아를 기반으로 만들어진 벤치마크 데이터이고, 기존 베이스라인 모델 또한 대부분 위키피디아 문서를 기반으로 pre-train 되었으니, 오히려 많은 위키피디아와 연관되지 않은 데이터로 finetuning한 것이 악영향을 준 것 아닌가 싶었다.
이 부분에 대해서 어떻게 개선할지는 조금 더 고민해 보아야 할 것 같다.
4. Troubleshooting
마지막으로 학습 코드를 작성하고 실행하며 겪은 어려움에 대해 공유한다.
혹시나 비슷한 문제를 겪으시는 분들은 보시고 잘 해결하시길..
4–1. 유사도 계산 시 값이 너무 커져 backward가 안되는 이슈
정확한 error는 다음과 같았다:
RuntimeError: Function ‘LogBackward0’ returned nan values in its 0th output.
왜 이런 에러가 나는지 모르겠어서 compute_loss의 모든 부분을 logging 해보니, 다음과 같았다.
자세히 보니, score가 inf가 되어, loss 계산 시 probability가 NAN이 되는 경우가 발생했다. 그리고 total_scores의 다른 값들도 보니, 2.31133e+38 등 스코어가 너무 크다는 것을 알게 되었다.
이를 해결하기 위해 리서치 해본 결과, LogSumExp라는 기법을 사용하면 될 것 같았다. 변형한 식은 다음과 같다.
지수에서의 뺄셈은 진수에서의 나눗셈에 해당하므로, 분자와 분모에서 같은 값을 빼주면 된다. 이에 기존 코드
def compute_loss(self, model, inputs, return_outputs=False):
embeddings = {}
for k in ['query', 'document', 'hard_negative']:
input_dict = {
'input_ids': inputs[f'{k}_input_ids'],
'attention_mask': inputs[f'{k}_attention_mask']
}
output: BaseModelOutput = model(**input_dict)
embeddings[k] = average_pool(output.last_hidden_state, input_dict['attention_mask'])
query_embeddings = embeddings['query']
positive_embeddings = embeddings['document']
negative_embeddings = embeddings['hard_negative']
similarity_fct = nn.CosineSimilarity(dim=-1)
tau = self.args.cl_temperature
positive_scores = similarity_fct(query_embeddings, positive_embeddings) / tau
positive_scores = torch.exp(positive_scores) # InfoNCE loss의 분자
negative_scores = similarity_fct(query_embeddings.unsqueeze(1), negative_embeddings.unsqueeze(0)) / tau
negative_scores = torch.exp(negative_scores)
total_scores = torch.cat((positive_scores.unsqueeze(1), negative_scores), dim=1)
total_scores_sum = total_scores.sum(dim=1) # InfoNCE loss의 분모
infoNCE_loss = -torch.log(positive_scores / total_scores_sum)
loss = infoNCE_loss.mean()
return loss에서 LogSumExp를 적용한 코드로 바꾸었다.
def compute_loss(self, model, inputs, return_outputs=False):
...
query_embeddings = embeddings['query']
positive_embeddings = embeddings['document']
negative_embeddings = embeddings['hard_negative']
similarity_fct = nn.CosineSimilarity(dim=-1)
tau = self.args.cl_temperature
positive_scores = similarity_fct(query_embeddings, positive_embeddings) / tau
negative_scores = similarity_fct(query_embeddings.unsqueeze(1), negative_embeddings) / tau
max_positive_scores = torch.max(positive_scores, dim=0, keepdim=True)[0]
max_negative_scores = torch.max(negative_scores, dim=1, keepdim=True)[0]
max_scores = torch.max(max_positive_scores, max_negative_scores) # max_score을 구하여
stable_positive_scores = positive_scores - max_scores # max_score을 빼줌
stable_negative_scores = negative_scores - max_scores.unsqueeze(1) # 여기서도 max_score을 빼줌
exp_positive_scores = torch.exp(stable_positive_scores)
exp_negative_scores = torch.exp(stable_negative_scores)
total_scores_sum = exp_positive_scores + exp_negative_scores.sum(dim=1)
log_prob = torch.log(exp_positive_scores / total_scores_sum)
loss = -log_prob.mean()
return loss4–2. Token_type_ids 관련 에러
정확한 error는 다음과 같았다:
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.cuda.LongTensor [8, 512]] is at version 4; expected version 3 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
에러에는 inplace 연산자 관련한 에러라고 되어 있어, 눈을 씻고 찾아보았지만 inplace 연산자는 전혀 볼 수 없었다. 그래서 정말 오래동안 검색해보니…
xlm-roberta에 token_type_ids를 전달해주지 않아서 그런 것이었다.
분명 intfloat의 mE5는 token_type_ids를 전달해주지 않아도 된다고 써 있었는데…backward 시에는 다른가보다.
여하튼 그리하여 trainer의 compute_loss 부분을
def compute_loss(self, model, inputs, return_outputs=False):
embeddings = {}
for k in ['query', 'document', 'hard_negative']:
input_ids = inputs[f'{k}_input_ids']
attention_mask = inputs[f'{k}_attention_mask']
input_dict = {
'input_ids': input_ids,
'attention_mask': attention_mask
}
output: BaseModelOutput = model(**input_dict)
pooled_output = average_pool(output.last_hidden_state, input_dict['attention_mask'])
embeddings[k] = pooled_output이 코드에서
def compute_loss(self, model, inputs, return_outputs=False):
embeddings = {}
for k in ['query', 'document', 'hard_negative']:
input_ids = inputs[f'{k}_input_ids']
attention_mask = inputs[f'{k}_attention_mask']
input_dict = {
'input_ids': input_ids,
'attention_mask': attention_mask,
'token_type_ids': torch.zeros_like(input_ids)
}
output: BaseModelOutput = model(**input_dict)
pooled_output = average_pool(output.last_hidden_state, input_dict['attention_mask'])
embeddings[k] = pooled_output이렇게 변경하였더니 잘 해결되었다.
token_type_ids는 각 토큰이 어떤 문장에 속하는지 구분하는 데 사용한다. 우리의 경우, query, document, hard_negative 하나씩 추출하기 때문에 각 원소는 모두 한 문장이라고 본다. 그러므로 torch.zeros_like(input_ids)로 input_ids 모두 같은 token을 가지도록 하였다.
4–3. DDP 적용 시 torchrun을 사용하는데도 GPU에 병렬적으로 데이터가 로드되지 않는 이슈
상황은 이러했다: GPU 8개를 사용하고, torchrun을 사용하여 ddp를 구현하는데, 계속 gpu에 데이터가 병렬적으로 올라가는 것이 아닌, 순차적으로 올라가는 문제가 있었다. 그래서 각 프로세스마다 데이터를 불러오는데 5분이 걸려, 데이터를 불러오는데만 거의 40분을 차지했다.
처음에는 torch.load의 문제인가 싶었다. 왜냐하면 로드하려는 데이터가 크지도 않은데도 6분이나 걸렸기 때문.. ⇒ 근데 이건 그냥 torch.load가 자체적으로 느린 것이었다. (참고)
그래서 뭐가 문제인지 계속 찾아봤는데, 모든 부분에서 로그를 찍어보니, 데이터 로드하는 부분에서만 계속 병렬적으로 로드가 안되고 있었다. (다른 로그들은 다 병렬적으로 잘 찍히더라) 그래서 코드를 살펴보니
with FileLock(lock_path):
if os.path.exists(cached_features_file) and not args.overwrite_cache:
logger.info("Start loading features...")
start = time.time()
self.features = torch.load(cached_features_file)
logger.info(
f"Loading features from cached file {cached_features_file} [took %.3f s]", time.time() - start
)처음에 데이터를 cache에 저장하면서 mutual exclusion을 위해 넣었던 FileLock이 DDP 환경에서 데이터를 병렬적으로 불러오지 못하는 원인이었다. (내 덫에 내가 걸림)
이를 해결하니, 위 사진과 다르게 병렬적으로 잘 불러오는 모습을 볼 수 있었다 ㅎ
그리고 DDP로 코드를 구현하니, 원래 24시간 걸리던 학습 시간이 3시간으로 줄어드는 놀라운 현상을 볼 수 있었다..(약 8배 단축) 이를 계기로 분산처리를 더욱 더 공부하고 싶은 마음이 생겼다.
5. 후기
이렇게 E5 모델을 리뷰하고, 한국어로 fine-tuning하여 평가, 트러블슈팅한 과정까지 모두 작성해 보았다. 오랜 시간 공을 들였는데 결과가 좋게 나와 뿌듯했다.
이번에 임베딩 모델에 관한 리서치를 해보면서 알게 되었는데, 생각보다 한국어 특화된 임베딩 모델을 원하고, 필요해하시는 분들이 많은 것 같다.
KoE5 모델은 결과가 괜찮았지만, multilingual-e5 특성상 context size가 작아, 효과적으로 사용하기에는 어렵다는 단점이 있다.
또, 계약연구센터와의 협업으로 인해 모델을 비공개로 만들었다는 것이 좀 속상했다.
(오픈소스로 만들어야 사람들이 써보고 피드백도 받아보고 하기에..)
그래서 조만간 pretraining부터 진행해 “한국어에 특화된" 임베딩모델을 만들고, 가능하다면 일부를 오픈소스로 공개해보려고 한다.
글 읽어주셔서 감사하고, 궁금한 점 있으시면 댓글 남겨주시면 감사하겠습니다 !!
