# 2. Moreh의 학습 코드 톺아보기

학습 데이터를 모두 준비하셨다면 다음으로는 실제 fine-tuning 과정을 실행할 train_baichuan2_13b.py 스크립트의 내용을 살펴 보겠습니다. 이번 단계에서는 MoAI Platform은 pytorch와의 완전한 호환성으로 학습 코드가 일반적인 nvidia gpu를 위한 pytorch 코드와 100% 동일하다는 것을 확인하실 수 있습니다. 또한 이를 넘어서 기존의 복잡한 병렬화 기법들을 MoAI Platform에서는 얼마나 효율적으로 구현할 수 있는지도 확인하실 수 있습니다.

우선 제공된 스크립트를 그대로 사용하여 튜토리얼을 끝까지 진행해 보시기를 권장합니다. 이후 스크립트를 원하는 대로 수정하셔서 Baichuan 모델을 다른 방식으로 fine-tuning하는 것도 얼마든지 가능합니다. 필요하시다면 Moreh에서 제공하는 LLM Fine-tuning 파라미터 가이드를 참고하십시오.

# Training Code

모든 코드는 일반적인 pytorch 사용 경험과 완벽하게 동일합니다.

먼저, transformers 라이브러리에서 필요한 모듈을 불러옵니다.

from transformers import AdamW, AutoTokenizer
from modeling_baichuan import BaichuanForCausalLM

HuggingFace에 공개된 모델 체크포인트를 불러옵니다.

모델 최적화를 위해 quickstart repository에 미리 정의된 BaichuanForCausalLM를 사용합니다.

model = BaichuanForCausalLM.from_pretrained('baichuan-inc/Baichuan-13B-Base', trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained('baichuan-inc/Baichuan-13B-Base', trust_remote_code=True)

Hugging Face에 공개된 학습 데이터셋을 불러와 전처리하고, 데이터 로더를 정의합니다.

# Load dataset
dataset = load_dataset("bitext/Bitext-customer-support-llm-chatbot-training-dataset").with_format("torch")
...
dataset = dataset.map(preprocess)

# Create a DataLoader for the training set
train_dataloader = torch.utils.data.DataLoader(
	dataset,
	batch_size=args.batch_size,
	shuffle=True,
	drop_last=True,
)

이후 학습도 일반적인 Pytorch를 사용하여 모델 학습과 동일하게 진행됩니다.

# Mask pad tokens for training
def mask_pads(input_ids, attention_mask, ignore_index = -100):
	idx_mask = attention_mask
	labels = copy.deepcopy(input_ids)
    labels[~idx_mask.bool()] = ignore_index
    return labels

# Define AdamW optimizer
optim = AdamW(model.parameters(), lr=args.lr)

# Start training
for epoch in range(args.epoch):
    for i, batch in enumerate(train_dataloader, 0):
	input_ids = batch["input_ids"]
	attn_mask = batch["attention_mask"]
	labels = mask_pads(input_ids, attn_mask)
	outputs = model(
		input_ids.cuda(),
		attention_mask=attn_mask.cuda(),
		labels=labels.cuda(),
		use_cache=False,
	)

	loss = outputs[0]
	loss.backward()

	optim.step()
	model.zero_grad(set_to_none=True)

위와 같이 MoAI Platform에서는 기존에 사용하시던 PyTorch 스크립트를 수정 없이 동일하게 사용하실 수 있습니다.

# About Advanced Parallelism

본 튜토리얼에 사용되는 학습 스크립트에는 아래와 같은 코드가 추가로 한 줄 존재합니다. 이는 MoAI Platform에서 제공하는 자동 병렬화 기능을 수행하는 코드입니다.

torch.moreh.option.enable_advanced_parallelization()

본 튜토리얼에서 사용하는 Baichuan2 13B 와 같은 거대한 언어 모델의 경우 필연적으로 여러 개의 GPU를 사용하여 학습시켜야만 합니다. 이때, MoAI Platform이 아닌 다른 프레임워크를 사용할 경우, Data Parallel, Pipeline Parallel, Tensor Parallel과 같은 병렬화 기법을 도입하여 학습을 수행해야 합니다.

예를 들어, 사용자가 일반적인 pytorch 코드에서 DDP를 적용하고 싶다면, 다음과 같은 코드 스니펫이 추가되어야 합니다. (https://pytorch.org/tutorials/intermediate/ddp_tutorial.html)

...
def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
...

def main(rank, world_size, args):
	setup(rank, world_size)
...
	sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
	loader = DataLoader(dataset, batch_size=64, sampler=sampler)
...

...
world_size = torch.cuda.device_count()  # Change this if you want a different number of GPUs
rank = int(os.environ['LOCAL_RANK'])
main(rank, world_size, args)
...

# single node 실행
torchrun --standalone --nnodes=1 --nproc_per_node=8 train.py
# multi node 실행
torchrun --nnodes=2 --nproc_per_node=8 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29400 train.py

DDP는 비교적 쉽게 적용할 수 있지만, 파이프라인 병렬 처리나 텐서 병렬 처리를 적용하려면 상당히 복잡한 코드 수정이 필요합니다. 최적화된 병렬화 처리를 적용하려면 학습 스크립트 작성 과정에서 Python 코드가 다중 처리 환경에서 어떻게 동작하는지 이해해야 하며, 특히 다중 노드 설정에서는 학습에 사용되는 각 노드의 환경을 구성해야 합니다. 또한, 모델 종류, 크기, 데이터셋 등을 고려해 최적의 병렬화 방법을 찾기 위해서는 상당히 많은 시간이 필요합니다.

반면, MoAI Platform의 AP 기능을 통해 사용자는 별도의 병렬화 기법을 적용할 필요 없이, 학습 스크립트에 단 한 줄의 코드를 추가하는 것으로도 최적화된 병렬화 학습을 진행할 수 있습니다.

import torch
...
torch.moreh.option.enable_advanced_parallelization()

model = AutoModelForCausalLM.from_pretrained('baichuan-inc/Baichuan-13B-Base', trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained('baichuan-inc/Baichuan-13B-Base', trust_remote_code=True)
...

다른 프레임워크에서는 경험할 수 없는 MoAI Platform만의 Advanced Parallelization(AP) 기능을 통해 최적의 자동화된 분산 병렬처리를 경험해보세요. AP기능을 이용하면 대규모 모델 훈련시 일반적으로 필요한 Pipeline Parallelism, Tensor Parallelism의 최적 매개변수와 환경변수를 아주 간단한 코드 한 줄로 설정할 수 있습니다.

tutorial baichuan2