PyTorch 2.0

PyTorch : Pytorch란 META에서 개발한 python 머신 러닝 프레임워크이다

머신 러닝 프레임워크로는 대표적으로 구글의 Tensorflow와 메타의 Pytorch가 존재합니다.

하지만 22년 HuggingFace 모델 기준으로 약 85%가 Pytorch 전용으로 공개되었으며, 가장 인기 있는 모델 30가지 기준으로도 약 66%(2/3) 이상이 Pytorch 전용으로 공개되었습니다.
21년 기준 75%의 논문에서 Pytorch를 택할만큼 이제는 대부분의 사람들이 Pytorch를 사용합니다.

그렇다면 PyTorch 2.0에서는 어떠한 점이 달라졌을까요?
발췌 : PyTorch Get Started

fundamentally changing and supercharging how PyTorch operates at compiler level under the hood.

Pytorch 2.0의 주된 변화는 Compiler level의 연산이라고 합니다.
이는 torch.compile의 이야기이며, "new direction for PyTorch" 라고 칭하며 강조하는 기능입니다.
다음과 같은 모델을 통해 해당 기능의 성능을 검증합니다.

• HuggingFace Transformers의 46개의 모델 -> 52% speed up
• TIMM의 61개의 모델 -> 38% speed up
• TorchBench의 56개의 모델 -> 76% speed up

open-source 모델을 변경하지 않고 torch.compile을 추가하기만 했을 때
A100 GPU와 Float32 precision 기준으로 43% 속도 향상이 있었으며,
Automatic Mixed Precision(AMP) 기준으로 51%의 속도 향상이 있었습니다.

해당 기능은 간단히 코드 한줄을 추가함으로써 사용 가능합니다.

# 모델을 넣으면 컴파일된 모델을 반환해준다.
compiled_model = torch.compile(model)

이를 실제 HuggingFace 모델에 접목시키려면 다음과 같이 사용하면 됩니다.

import torch
from transformers import BertTokenizer, BertModel
# Copy pasted from here https://huggingface.co/bert-base-uncased
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")
model = torch.compile(model) # This is the only line of code that we changed
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
output = model(**encoded_input)

이 외에도 Modes라는 기능이 존재합니다.

# API NOT FINAL
# default: optimizes for large models, low compile-time
#          and no extra memory usage
torch.compile(model)

# reduce-overhead: optimizes to reduce the framework overhead
#                and uses some extra memory. Helps speed up small models
torch.compile(model, mode="reduce-overhead")

# max-autotune: optimizes to produce the fastest model,
#               but takes a very long time to compile
torch.compile(model, mode="max-autotune")

이는 small model에 적합한 mode인 "reduce-overhead"를 사용하여 framework overhead를 줄여 속도를 향상시킬 수 있는 기능을 제공합니다.

PyTorch Lightning 2.0

PyTorch Lightning : Pytorch의 Boilerplate를 제거하여 가독성이 좋으며 간결한 코드를 만드는 프레임 워크

버전이 변경되면서 다음과 같이 변경된 것을 볼 수 있었습니다.

import pytorch_lightning as pl
# ->
import lightning as L

해당 프로젝트에서는 lightning을 사용하는 것이 아닌 FABRIC을 사용하고자 합니다.
FABRIC을 사용하는 이유는 기존 Pytorch 유저에게 PL은 너무나 많은 코드 변화를 유발합니다.
하지만 FABRIC을 사용할 경우 적은 수준의 변화만을 사용하고도 Lightning에서 제공하는 기능을 사용할 수 있기에 기존 Torch 유저가 편하게 접근할 수 있다고 생각했기 때문입니다.

다음은 FABRIC 코드 스타일입니다.

import lightning as L


class LitModel(L.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = ...

    def training_step(self, batch, batch_idx):
        # Main forward, loss computation, and metrics goes here
        x, y = batch
        y_hat = self.model(x)
        loss = self.loss_fn(y, y_hat)
        acc = self.accuracy(y, y_hat)
        ...
        return loss

    def configure_optimizers(self):
        # Return one or several optimizers
        return torch.optim.Adam(self.parameters(), ...)

    def train_dataloader(self):
        # Return your dataloader for training
        return DataLoader(...)

    def on_train_start(self):
        # Do something at the beginning of training
        ...

    def any_hook_you_like(self, *args, **kwargs):
        ...

기존 PL과 비슷하게 LightningModule을 사용해서 hooks를 정의 후

import lightning as L

fabric = L.Fabric(...)

# Instantiate the LightningModule
model = LitModel()

# Get the optimizer(s) from the LightningModule
optimizer = model.configure_optimizers()

# Get the training data loader from the LightningModule
train_dataloader = model.train_dataloader()

# Set up objects
model, optimizer = fabric.setup(model, optimizer)
train_dataloader = fabric.setup_dataloaders(train_dataloader)

# Call the hooks at the right time
model.on_train_start()

model.train()
for epoch in range(num_epochs):
    for i, batch in enumerate(dataloader):
        optimizer.zero_grad()
        loss = model.training_step(batch, i)
        fabric.backward(loss)
        optimizer.step()

        # Control when hooks are called
        if condition:
            model.any_hook_you_like()

이와 같이 hooks를 부르면서 실행하면 됩니다.
여기서 기존 PL과의 차이점은 Trainer를 사용하지 않는다는 점입니다.

model.train()
for epoch in range(num_epochs):
    for i, batch in enumerate(dataloader):
        optimizer.zero_grad()
        loss = model.training_step(batch, i)
        fabric.backward(loss)
        optimizer.step()

        # Control when hooks are called
        if condition:
            model.any_hook_you_like()

기존 PL의 경우 해당 부분이 Trainer로 감싸져 PL에 익숙하지 않은 분들에게는 가독성이 좋지 않았지만 FABRIC의 경우 해당 부분의 코드가 기존 Pytorch와 함수를 제외하면 동일하기에 가독성이 좋습니다.

하지만 FABRIC에도 단점은 존재합니다.
PL 혹은 HuggingFace Trainer가 제공하는 Mixed Precision 혹은 Multi GPU를 쉽게 적용하는 기능은 제공하지만 Gradient Accumulation의 경우 코드를 다음과 같이 변경해서 적용해야합니다.

for iteration, batch in enumerate(dataloader):
    # Accumulate gradient 8 batches at a time
    is_accumulating = iteration % 8 != 0

    output = model(input)
    loss = ...

    # .backward() accumulates when .zero_grad() wasn't called
    fabric.backward(loss)
    ...

    if not is_accumulating:
        # Step the optimizer after the accumulation phase is over
        optimizer.step()
        optimizer.zero_grad()

이는 템플릿을 만들 때 여러움으로 작용할 수 있기에 해결방안을 모색해야할 것 같습니다.