\begin{equation}
\pi_R(X, i) = X \quad \text{for all } i \in \{0, \ldots, N-1\}
\end{equation}
```*

***셰이딩 ($`S`$):** 텐서는 $`N`$ 개의 연속 셰이드로 분할됩니다. 장치 $`i`$는 셰이드 $`i`$를 저장합니다.
``` math
\begin{equation}
\pi_S(X, i) = X\left[\frac{i \cdot |X|}{N} : \frac{(i+1) \cdot |X|}{N}\right]
\end{equation}
```*

***셰이딩 위드 게더 ($`S^*`$):** $`S`$와 동일하지만 사용 전에 전체 텐서를 All-Gather를 통해 재구성하고, 사용 후에는 비로컬 부분을 삭제합니다. 이는 ZeRO-3/FSDP 파라미터 처리를 캡처합니다.
``` math
\begin{equation}
\pi_{S^*}(X, i) = X\left[\frac{i \cdot |X|}{N} : \frac{(i+1) \cdot |X|}{N}\right] \text{ (지속적인)}, \quad X \text{ (사용 중 일시적)}
\end{equation}
```*

***물리화 ($`M`$):** 어떤 장치도 $`X`$를 지속적으로 저장하지 않습니다. $`X`$가 필요할 때마다 다른 상태에서 재구성하고 사용한 후 삭제합니다. 이는 활성화와 같은 중간 값을 적용합니다.
``` math
\begin{equation}
\pi_M(X, i) = \emptyset \quad \text{(지속적인 저장)}
\end{equation}
```*

***오프로드 ($`O`$):** 텐서는 CPU 메모리 또는 NVMe에 저장되어 있으며 필요할 때 GPU로 전송됩니다.
``` math
\begin{equation}
\pi_O(X, i) = \emptyset \quad \text{(GPU 메모리)}
\end{equation}
```*

</div>

<div class="example">

**예 1** (데이터 병렬화). *데이터 병렬화(DP)에서 모든 모델 상태는 복제됩니다: $`\pi_\Theta = R`$, $`\pi_\Omega = R`$, $`\pi_G = R`$. 각 장치가 전체 파라미터, 전체 옵티마이저 상태 및 전체 그래디언트를 보유합니다. 로컬 그래디언트 계산 후 All-Reduce를 통해 그래디언트를 장치 간 동기화합니다.*

</div>

<div class="example">

**예 2** (ZeRO Stage 3). *ZeRO Stage 3는 모든 것을 셰이딩합니다: $`\pi_\Theta = S^*`$, $`\pi_\Omega = S`$, $`\pi_G = S`$. 파라미터는 장치들 사이에 분산되어 있으며, 각 레이어의 계산 전에는 전체 파라미터를 All-Gather로 재구성하고 사용 후 삭제합니다. 옵티마이저 상태와 그래디언트는 계속 셰이딩됩니다.*

</div>

<div id="tab:placements">

| 전략 | $`\pi_\Theta`$ | $`\pi_\Omega`$ | $`\pi_G`$ | $`\pi_A`$ |
|:---|:--:|:--:|:--:|:--:|
| 데이터 병렬화 (DP) | $`R`$ | $`R`$ | $`R`$ | $`R`$ |
| ZeRO Stage 1 | $`R`$ | $`S`$ | $`R`$ | $`R`$ |
| ZeRO Stage 2 | $`R`$ | $`S`$ | $`S`$ | $`R`$ |
| ZeRO Stage 3 / FSDP | $`S^*`$ | $`S`$ | $`S`$ | $`R`$ |
| ZeRO-오프로드 | $`O`$ | $`O`$ | $`S`$ | $`R`$ |
| 텐서 병렬화 (TP, 레이어 내) | $`S`$ | $`S`$ | $`S`$ | $`S`$ |
| 파이프라인 병렬화 (PP, 레이어 간) | $`S`$ | $`S`$ | $`S`$ | $`R`$ |

공통적인 병렬화 전략에 대한 배치 사양. 텐서 및 파이프라인 병렬화의 경우, 배치는 각 레이어 또는 스테이지별로 적용됩니다. $`S^*`$은 ZeRO-3/FSDP에서 사용되는 계산 전 셰이딩 위드 게더를 나타냅니다. DP = 데이터 병렬화, TP = 텐서 병렬화, PP = 파이프라인 병렬화.

</div>

<figure id="fig:modes" data-latex-placement="t">

<figcaption>다섯 가지 배치 모드. 위: 각 장치의 GPU 메모리 비용 (<span class="math inline"><em>s</em></span> = 텐서 크기, <span
class="math inline"><em>N</em></span> = 장치 수). 아래: 공통적인 전략에서 사용 예. <span
class="math inline"><em>S</em><sup>*</sup></span> (셰이딩 위드 게더)는 ZeRO-3/FSDP의 핵심 혁신입니다: 파라미터는 저장을 위해 셰이딩되지만, 계산을 위해 일시적으로 재조립됩니다.</figcaption>
</figure>

## 배치 사양

<div class="definition">

**정의 2** (배치 사양). *배치 사양은 튜플 $`\Pi = (\pi_\Theta, \pi_\Omega, \pi_G, \pi_A)`$입니다. 여기서 각각의 $`\pi_X \in \{R, S, S^*, M, O\}`$는 상태 $`X`$에 대한 배치 모드를 지정합니다.*

</div>

병렬화 전략은 그 배치 사양으로 완전히 결정됩니다. 표 <a href="#tab:placements" data-reference-type="ref"
data-reference="tab:placements">2</a>는 알려진 전략들의 사양을 보여주며, 도표 <a href="#fig:modes" data-reference-type="ref"
data-reference="fig:modes">1</a>은 다섯 가지 모드와 그 메모리 비용을 설명합니다.

<div class="remark">

**비고 2.** *물리화 모드 ($`M`$)는 표 <a href="#tab:placements" data-reference-type="ref"
data-reference="tab:placements">2</a>에 나타나지 않습니다. 이는 공통적인 전략들이 모든 상태를 지속적으로 보관하기 때문입니다. 그러나 $`M`$은 활성화 체크포인팅을 모델링하는 데 사용할 수 있습니다. 여기서 활성화는 저장 대신 재계산됩니다.*

</div>

# 유도 규칙

현재 우리의 주요 기술적 기여를 제시합니다: 배치 사양으로부터 메모리와 통신 비용을 유도하는 규칙들입니다.

## 준비사항

<div id="def:sunit" class="definition">

**정의 3** (재구성 단위). *$`s_{\text{unit}}`$는 셰이딩 위드 게더 작업 중에 독립적으로 재구성될 수 있는 가장 작은 단위 크기를 나타냅니다. 트랜스포머 모델의 경우 이는 일반적으로 하나의 레이어를 나타냅니다: $`s_{\text{unit}} = 12H^2 \cdot \text{bytes\_per\_param}`$ 표준 트랜스포머 레이어의 숨겨진 차원 $`H`$. $`s_{\text{unit}}`$의 선택은 메모리와 통신 세분화를 교환하는 구현 결정입니다.*

</div>

## 메모리 유도

<div id="thm:memory" class="theorem">

**정리 1** (배치로부터 메모리). *$`\Pi = (\pi_\Theta, \pi_\Omega, \pi_G, \pi_A)`$는 $`N`$ 장치에 대한 배치 사양입니다. 각 장치의 GPU 메모리는 다음과 같습니다:
``` math
\begin{equation}
M(\Pi) = \mu(\pi_\Theta, |\Theta|) + \mu(\pi_\Omega, |\Omega|) + \mu(\pi_G, |G|) + \mu(\pi_A, |A|)
\end{equation}

\begin{align}
\mu(R, s) &= s \\
\mu(S, s) &= s/N \\
\mu(S^*, s) &= s/N + s_{\text{unit}} \\
\mu(M, s) &= s_{\text{unit}} \\
\mu(O, s) &= 0
\end{align}