PyTorch Lightning DeepSpeed

토치 라이트닝에서 딥스피드를 사용하기 위한 방법 학습 및 정리.

DeepSpeed

딥스피드는 딥러닝 학습 optimization 라이브러리이다. 딥스피드를 통해 트랜스포머 모델의 메모리 를 효율적으로 학습할 수 있게 해, Billion 이상의 파라미터를 가진 모델들을 학습할 수 있게 해준다.

DeepSpeed in PyTorch lightning

딥스피드 패키지를 설치해서 사용하면, 기존의 파이토치 코드에 딥스피드를 적용해서 학습할 수 있다. 적용할때 토치 라이트닝을 이용하면 조금더 편하게 사용할 수 있다고하여, 토치 라이트닝에서 사용해보고자 한다.

토치 라이트닝에서 DeepSpeed Plugin을 이용하면 10 Billion 이상의 파라미터를 가진 모델을 학습할수 있게 해준다.

트랜스포머 모델은 메모리 비용이 크기 때문에, 큰 모델을 학습 하고 이용하기 위해 메모리를 절약하는게 반드시 필요하다. 하지만 학습에서 학습을 위한 모든 상태(Optimizer States, Parameters)들이 GPU 위에 있기 때문에 메모리를 많이 차지 한다.

DeepSpeed의 ZeRO의 경우 ZeRO Redundancy Optimizer의 줄인 말로, 메모리를 최적화 하면서 학습 속도를 향상 시키고자한 옵티마이저 라이브러리이다.

단순하게는 필요한만큼만 메모리에 올려서 메모리의 중복을 제거했다.

1. ZeRO Stage 2

메모리를 줄이기 위해 GPU에서 Optimizer States와 Gradient를 분할하고, 대부분의 경우에서 DDP와 비슷하거나 더 뛰어난 성능을 보인다. 기본 bucket size는 약 3.6GB VRAM을 분산 커뮤니케이션을 위해 사용한다.

from pytorch_lightning import Trainer

model = MyModel()
trainer = Trainer(gpus=4, plugins="deepspeed_stage_2", precision=16)
trainer.fit(model)

2. ZeRO Stage 2 Offload

Offload는 optimizer states와 gradients를 cpu에 내려서 계산하는 방법이다. cpu는 gpu에 비해 느리기 때문에 딥러닝 학습에서 계산 비용이 높은 것들은 cpu에서 수행하면 안된다.

딥러닝에서 시간 복잡도는 보통 $O(model\_size * batch\_size)$를 따르며, GPU에서 대부분의 메모리는 optimizer states에서 차지한다. optimizer states는 계산 비용이 Forward, Backward 연산에 비해 적은 편으로 CPU에서 계산하게 된다면, GPU의 메모리를 절약할수 있다.

Offload optimizer states and gradients to CPU. Increases communication, but significant memory improvement
Basic usage

from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin

model = MyModel()
trainer = Trainer(gpus=4, plugins="deepspeed_stage_2_offload", precision=16)
trainer.fit(model)

More speed benefit
DeepSpeedCPUAdam 을 사용해, 기존에 PyTorch 구현체에서 사용하던것 보다 빠르고, CPU에서 계산 가능한 Optimizer를 제공한다.

import pytorch_lightning
from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin
from deepspeed.ops.adam import DeepSpeedCPUAdam


class MyModel(pl.LightningModule):
    ...

    def configure_optimizers(self):
        # DeepSpeedCPUAdam provides 5x to 7x speedup over torch.optim.adam(w)
        return DeepSpeedCPUAdam(self.parameters())


model = MyModel()
trainer = Trainer(gpus=4, plugins="deepspeed_stage_2_offload", precision=16)
trainer.fit(model)

3. ZeRO Stage 3

ZeRO Stage 3는 optimizer states, gradients, model parameters를 분할하여 단일 GPU에서 부터 다중 GPU에서까지 모델을 크게 증가시킬수 있다.

최적의 메모리 효율과 가장큰 모델을 학습하기 위한 요소:
1. Stage3 DeepSpeed Plugin을 사용.
2. CPU Offloading 사용. (적절한 CPU RAM 필요)
3. Shard Activations에 DeepSpeed Activation Checkpointing 사용

from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin
from deepspeed.ops.adam import FusedAdam


class MyModel(pl.LightningModule):
    ...

    def configure_optimizers(self):
        return FusedAdam(self.parameters())


model = MyModel()
trainer = Trainer(gpus=4, plugins="deepspeed_stage_3", precision=16)
trainer.fit(model)

trainer.test()
trainer.predict()

4. ZeRO Stage 3 Offload

Offload 기능을 사용해 메모리 사용량을 낮춤.

from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin

# Enable CPU Offloading
model = MyModel()
trainer = Trainer(gpus=4, plugins="deepspeed_stage_3_offload", precision=16)
trainer.fit(model)

# Enable CPU Offloading, and offload parameters to CPU
model = MyModel()
trainer = Trainer(
    gpus=4,
    plugins=DeepSpeedPlugin(
        stage=3,
        offload_optimizer=True,
        offload_parameters=True,
    ),
    precision=16,
)
trainer.fit(model)

DeepSpeed Infinity(NVMe Offloading)

딥스피드는 NVMe 드라이브에 Offloading 기능을 제공한다.

NVMe 란?
PCI Express 인터페이스로 연결된 비휘발성 저장장치들을 위한 새로운 통신 프로토콜. AHCI의 후계 프로토콜로 취급하고 있으며, SATA + AHCI 조합의 대역폭을 뛰어넘는 I/O 성능을 가진 고성능 SSD를 위한 규격이다.

from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin

# Enable CPU Offloading
model = MyModel()
trainer = Trainer(gpus=4, plugins="deepspeed_stage_3_offload", precision=16)
trainer.fit(model)

# Enable CPU Offloading, and offload parameters to CPU
model = MyModel()
trainer = Trainer(
    gpus=4,
    plugins=DeepSpeedPlugin(
        stage=3,
        offload_optimizer=True,
        offload_parameters=True,
        remote_device="nvme",
        offload_params_device="nvme",
        offload_optimizer_device="nvme",
        nvme_path="/local_nvme",
    ),
    precision=16,
)
trainer.fit(model)

DeepSpeed Activation Checkpointing

Activation Checkpoint는 forward 되는 동안 필요하지 않은 activation의 메모리를 해제한다. 그리고 backward 과정에서 필요할때 다시 계산하여 사용한다.

사용하기 위해서는 아래와 같이 checkpoint 함수를 이용해야 한다.

from pytorch_lightning import Trainer
from pytorch_lightning.plugins import DeepSpeedPlugin
import deepspeed


class MyModel(pl.LightningModule):
    ...

    def configure_sharded_model(self):
        self.block = nn.Sequential(nn.Linear(32, 32), nn.ReLU())

    def forward(self, x):
        # Use the DeepSpeed checkpointing function instead of calling the module directly
        output = deepspeed.checkpointing.checkpoint(self.block, x)
        return output


model = MyModel()


trainer = Trainer(gpus=4, plugins="deepspeed_stage_3_offload", precision=16)

# Enable CPU Activation Checkpointing
trainer = Trainer(
    gpus=4,
    plugins=DeepSpeedPlugin(
        stage=3,
        cpu_offload=True,  # Enable CPU Offloading
        cpu_checkpointing=True,  # (Optional) offload activations to CPU
    ),
    precision=16,
)
trainer.fit(model)

DeepSpeed ZeRO Stage 3 Tip

토치 라이트닝으로 DeepSpeed ZeRO Stage 3를 사용할때 팁.

• Adam 또는 AdamW를 사용하는 경우 기본 토치 옵티마이저 대신 FusedAdam 또는 DeepSpeedCPUAdam(CPU 오프로딩용)을 사용.

• GPU/CPU 메모리를 하나의 큰 풀로 다룰것. 어떤 경우에는 모델 매개변수를 오프로드하기 위해 더 많은 공간을 제공하기 위해 특정 항목(예: Activations)을 오프로드하고 싶지 않을 수 있다.

• CPU로 오프로드할 때 GPU 메모리가 해제되므로 배치 크기를 늘릴것.

• 분할된 체크포인트를 지원. save_full_weights=False하여 DeepSpeedPlugin을 통해 매우 큰 모델을 저장할 수 있는 모델의 조각을 저장할수 있다. 그러나 모델을 로드하고 테스트/검증/예측을 실행하려면 Trainer 개체를 사용해야 한다.

Custom DeepSpeed Config

원하는 config를 사용해서 학습할 수 있음. 링크참고.

References

• Advanced GPU Optimized Training | DeepSpeed
• 집현전 시즌2| Season2-1조-Large-sclae LM에 대한 얕고 넓은 지식들 (part 1).pdf