# NVIDIA NeMO Megatron
## 용어 정리
- NVIDIA NeMO Framework: NVIDIA에서 제공하는 엔터프라이즈 프레임워크로 빌드, 커스터마이즈, 배포기능 제공
end-to-end cloud-native enterprise framework for developers to build, customize, and deploy generative AI models with billions of parameters
https://developer.nvidia.com/nemo

- Megatron: NVIDIA에서 개발한 transformer
large, powerful transformer developed by the Applied Deep Learning Research team at NVIDIA.

- Megatron-lm: transformer 언어 모델을 효율적으로 pretrain 할 수 있는 라이브러리
enables training large transformer language models at scale. It provides efficient tensor, pipeline and sequence based model parallelism for pre-training transformer based Language Models such as GPT (Decoder Only), BERT (Encoder Only) and T5 (Encoder-Decoder)
- 3d parallelism figure
- youtube explained 3d parallelism
https://www.youtube.com/watch?v=pTChDs5uD8I

- NeMO Magatron: Megatron 모델을 학습하기 위한 NVIDIA에서 제공하는 end-to-end framework
- 4가지 parallelism
- tensor parallism
Tensor Parallelism (TP)은 노드 내 랭크 간에 추가 통신이 크게 필요하지 않으면서 메모리 풋프린트를 줄이는 기술입니다. 각 텐서는 여러 개의 청크로 분할되며 각 샤드는 별도의 GPU에 위치합니다. 각 단계에서 동일한 미니 배치 데이터가 각 샤드에 의해 독립적으로 병렬로 처리되고 그 후 모든 GPU 간에 동기화됩니다 (all-reduce 연산). 간단한 트랜스포머 레이어의 경우, 이는 순방향 경로에서 2개의 all-reduce 및 역방향 경로에서 2개의 all-reduce를 수행하게 됩니다.
TP를 사용하면 모델의 메모리 요구사항을 줄이고 동시에 더 효율적으로 병렬화할 수 있습니다. 각 GPU는 해당하는 텐서의 일부를 저장하고 처리하여 메모리 요구사항을 감소시킵니다. 미니 배치 데이터는 독립적으로 처리되므로 각 GPU에서 병렬로 계산이 가능합니다. 그리고 모든 GPU 간에 동기화되는 과정을 통해 연산 결과를 통합합니다. 이러한 과정을 통해 TP는 대규모 모델을 효율적으로 훈련시킬 수 있는 기술입니다.
- data parallism
Distributed Optimizer를 통한 Data Parallelism (DP)은 옵티마이저 상태와 그래디언트를 DP 랭크 간에 분산하여 메모리 풋프린트를 줄이는 기술입니다. 이는 기존의 데이터 병렬 랭크 간 옵티마이저 상태 복제 방식과는 다릅니다. 예를 들어, 혼합 정밀도 훈련을 위해 Adam 옵티마이저를 사용하는 경우, 각 매개변수는 12바이트의 메모리를 차지합니다. 이 메모리는 GPU 간에 균등하게 분산되며, 즉, 4개의 GPU가 있다면 각 매개변수는 3바이트 (12/4)를 차지하게 됩니다.
이 방식을 사용하면 GPU 간에 옵티마이저 상태를 복제하지 않고 메모리를 효율적으로 사용할 수 있습니다. 매개변수와 그래디언트는 DP 랭크 간에 분할되어 저장되고 처리됩니다. 이를 통해 각 GPU는 자체적으로 계산을 수행하고 업데이트된 그래디언트를 동기화합니다. 분산 옵티마이저를 통한 Data Parallelism은 대규모 모델을 훈련시키는 데 사용되며 메모리 요구사항을 효과적으로 관리할 수 있는 기술입니다.
- pipeline parallism
파이프라인 병렬성(PP)은 메모리 풋프린트를 줄이고 교차 노드 병렬화를 통해 대규모 훈련을 가능하게 하는 기술입니다. 이 기술은 원활한 PP를 위한 PipeDream-Flush 스케줄 / 1F1B 스케줄 및 교차 1F1B 스케줄을 통해 초기 PP의 비효율성을 줄입니다. 레이어들은 PP 스테이지 간에 균등하게 분산됩니다. 예를 들어, 모델이 24개의 레이어를 가지고 파이프라인 병렬성을 위해 4개의 GPU를 사용한다면, 각 GPU는 6개의 레이어를 가지게 될 것입니다 (24/4).
PP는 여러 GPU 간에 모델의 레이어를 분할하여 병렬적으로 처리함으로써 메모리 사용을 최적화하고 대규모 모델의 훈련을 가능하게 합니다. 각 GPU는 해당하는 레이어를 처리하고 결과를 다음 스테이지로 전달합니다. 이러한 방식으로 각 GPU는 독립적으로 계산을 수행하며 레이어 간에는 데이터 흐름이 순차적으로 이루어집니다. PP를 통해 모델의 훈련 속도와 효율성을 향상시킬 수 있으며, 대규모 모델을 효과적으로 처리할 수 있습니다.
- sequence parallism
시퀀스 병렬성 (SP)은 메모리 풋프린트를 줄이면서 추가 통신 없이 작동합니다. TP를 사용할 때만 적용 가능합니다. 이는 활성화 메모리 요구사항을 줄이기 위해 모든 리듀스 후 텐서 병렬 랭크에서 동일한 복사본이 없도록 reduce-scatter로 교체하고 no-op 연산은 all-gather로 교체함으로써 이루어집니다. all-reduce = reduce-scatter + all-gather이기 때문에 통신 비용을 추가로 지불하지 않고도 많은 양의 활성화 메모리를 절약할 수 있습니다. 간단히 말해, SP는 각 트랜스포머 레이어의 출력을 시퀀스 차원을 따라 분할하여 처리합니다. 예를 들어, 시퀀스 길이가 1024이고 TP 크기가 4인 경우, 각 GPU는 각 샘플에 대해 256개의 토큰 (1024/4)을 가집니다. 이는 훈련에 사용할 수 있는 배치 크기를 증가시킵니다.

## Prerequisite
- pytorch lightning: NeMO framework는 기본적으로 lightning module을 사용함
https://www.youtube.com/playlist?list=PLhhyoLH6IjfyL740PTuXef4TstxAK6nGP
- megatron training: 현재 NeMO Megatron은 GPT-style models (decoder only), T5/BART-style models (encoder-decoder), BERT-style models (encoder only) 학습을 지원한다.
https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/nlp/nemo_megatron/gpt/gpt_training.html
- megatron with adapter: https://github.com/NVIDIA/NeMo/blob/main/tutorials/02_NeMo_Adapters.ipynb