??? NVIDIA Megatron-Core ???? ??? AI ??? ? ????? ????

Reading Time: 6 minutes

2019?? ?? ??? NVIDIA Megatron-LM? AI ????? ??? ??? ?????, ???? ???? ? ?? ?? ?????? ??? ?? ?? ??(LLM)? ?? ???????. ??? ?? ?? ?? LLM ??? ????? ???? Megatron-LM ?????? ???? Megatron-LM ?????? ??? ?????, ????? ??? AI ????? ??? ?????? ??? ???? ????. Megatron-LM? ???? ??? ?? LLM ??????? Colossal-AI, Hugging Face Accelerate, NVIDIA NeMo ?? ????.

??????? ???? ???? ???? ?? ???? ?? ?? ?? ??? ???? ? ??? ?? ?? NVIDIA? ?? Megatron-LM? ??????. ??? ??? ?? ??? NVIDIA Megatron-Core? GPU ??? ??, ??? ??? ??? ??, ??? ?? ??? ?? ??? API? ?? PyTorch ?? ?? ?? ????????.

Megatron-Core? ??? ??? ?? ??? ????? ???? ?? ????. ? ?????? ???? ??? ?? ??? LLaVA(Large Language and Vision Assistant) ?????? ??? ?? ??? ? ?? ?????.

NVIDIA Megatron-Core

Megatron-Core?? ??? ??? ??? ??? ??? GPU ??? ??? ???? ????. Megatron-Core? ??? ??? ?? ??? ??? API? ????? ????? ???? ???? NVIDIA ?? ??? ????? ???? ??? ?????? ??? ? ?? ??? ???? ?????.?

Megatron-Core ?????? ??? ????, ????? ?? ? ???, ??? ???, ??? ??? ?? ????? ??? ?? ?? ??? ?????. ??? ??? ? ?? ?????? ??? ?? ??? ? ?????? ????? ???? ????.

Megatron-Core? ???? ??? ??? ??? GPT, Bert, T5, RETRO ? ?? ?? LLM ????? ????? ??? ? ????. ?? Megatron-Core? ?? NVIDIA Tensor ?? GPU? ???? NVIDIA Hopper ????? ???? FP8 ??? ??? ???? ??? ???? ???? ??? ?? ??? ??? ? ????. Megatron-Core ??? Reka AI ? Codeium? ?? ???? ??? ???? ????? ? ?? ?????.

Reka AI? ?? ??? Deyu Fu? “Megatron-Core? ?? ??? ??? ???? ??? ???? LLM ????? ???? ?????. ???? GPU ??? ?? ?? ??? ?? ??? ??? ???? ????? ?? ???? ?? ??? ?? ? ?? ??? ???? ???? ????? ????? ?????.”? ??????.

Codeium? ?? ?? ???? Devin Chotzen-Hartzell? “Megatron-Core? ???? ??????? ???? ???? ???? LLM ?? ?? ???? ??? ??? ? ????. ??? ??? ? ?? ??? ??? ??? ??? ? ?? ?????.”? ?????.

Megatron-Core, ???? ?? ??

??? ?? ??? ????? ??? ???? ??? ???? ?? ??? ??? ???????. ??? ???? ?? ?? ??? ???? ??? ???? ??????, ??? ??? ???? ???? ????? ????? ?? ??? ?????. ??? ???? ??? AI ??? ??? ???? ??? ??? ?? ???????.

?? Megatron-Core v0.7??? ????? ?????. LLaVA? ???? ??? ???? ???? ?????? GitHub? NVIDIA/Megatron-LM? ?????. ?? ???? Megatron?? ?? ?? ???? ??? ??? ???? ???? ?????? ??? ? ???? ?? ??? ? ????. ?? ????? ?? ? ????? ?????.

LLaVA ????? ??? ??? ???? ??? ???? ? ????.

Megatron ? ????? ?? ??? ?? ?? ?? ? ?? ?? ??(SFT) ????? ??
Megatron Core ?? ?? ? ??? ?? ??? ??? Misral ? CLIP?? ???? LLaVA ???? ?? ??
COCO ?? ? VQAv2? ?? ??? ???? ?? ???

Megatron-Core v0.7 ???? LLaVA ?????? ??? ??? ??? ?? ???, ??? ??? ???? ??(MMMU) ??? 38???, ?? 7B ?? ?? LLM ?? LLaVA ????? ?? ??? ????.

?? Megatron-Core(Mcore) ?? ???? ?? ????? PyTorch ?? ???? ?? ??? ?? ???? ? ????. ?? Megatron-Core ?????? ??? ??? ?? ?? ?? ?? ??? ??? ?????. ? ?? ??? ? ?? ?? ??? ?? ????, ????? ???? ??? ?? ?? ??? ?? ? ??? ?? ??? ???? ?????.

???? ??? ??? ??? ?? Megatron-Core ???? ?? ??? ? NVIDIA NeMo? ???? NeVa? ?? ??? ?? NeMo? ?? ???? ??? ??? ?????.

??? ?? ??(MoE)? ?? ?? ??? ???

??? ???? ??? AI ???? ??? ??(MoE) ??? ?? ??? ??? ?? ??? ??? ?? ???? ????? ?? ??? ? ?? ???? ????. MoE?? ??? FFN ??? ? ??? ???? ??? ? ?? ????? ????? MoE ???? ?????.?

Megatron-Core v0.7? MoE ??? ???? ??? ?? ?? ? ??? ???? ???? ??? MoE ??? ?? ?? ???? ???? ?????. ???? Megatron-Core? ?? GShard?? ???? ?? ??? ?? MoE ??? ????, ?? CUDA ??? ?? ? ????? ?? ??? ?? ??? GroupedGEMM? ?? ?? ?? ???? ?????.

? 1?? Megatron-Core? ?-?-? ???? ? 4096? ??? ??? BF16 ???? ?? ? GPU? 400TFLOP/s ??? ???? ??? ??? ???? ????. ? ??? ? ?? ???(–moe-router-topk)?? ??????. NVIDIA? MoE? FP8 ???? ????? ????? ??? ?? Megatron-Core ????? ??? ? ??? ??? ?????.

?? Megatron-Core? MoE? ?? ??? ?? ??? ???? Megatron-Core?? ?? ???? Tensor, ???, ??? ? ????? ?? ??? ?? ?? ?? ?? ??? ??? ? ????. ??? ??? ?? ???? ?????.

??	???	GPU ?	MBS	GBS	TP	EP	PP	????? ??	GPU? ???(TFLOP/s/GPU)
Mistral 7B(??? ?? ??)	BF16	128	4	256	2	?? ??	1	1	492
Mixtral 8x7B	BF16	128	1	256	1	8	4	8	402

? 1. Megatron-Core v0.7?? ???? ??? ??? NVIDIA H100 GPU? Mixtral 8x7B? ?? GPU? ???

?? ??? ??? ?? ?? ?? ?????

?? ?????? ??? ???? ???? ???? ? ?? ?????. PyTorch ???? ???? torch.save? ???? ???? ??? ??? ?? ? ???? ???? ???? ??? ?????. ?? ?? Azure Nebula ? AWS Gemini? ??? ?????? ???? PyTorch ?? ?????(DCP)? ???? ???? ??? ?????? ?????. ??? ??? ?? ??? ????? ????? ?? `torch.save`??? ???? ???? ??? ?????? ???? ?? ??? ??? ????.?

?? ??? ???? ?? GPU ?? ??? ?? ?? ?? ?? Python ???(Python Global Interpreter Lock?? ?? ???? ?? ? ??)? ???? ??? ?????? ????? ????? ??? ??? ?? ??? ????. ?? ??? ???? ???? ?????? ???? ? ???? ?? ??? ?? ??(?: PP ? TP ??)?? ?????? ????? ?? ? ??? ???? ?? ?? ?? ??? ?? ????? ?? ?????.

Megatron-Core v0.7??? ??? ?? ? ??? ?? ??? ???? ??? ??? ??????. ?? ?? ??(FPS)? ?? ??? ?? ???? ?? ??? ???? ?? ??? ?? ??? ???? ?? ????? ??? ? ????. ??? ?? ??? ? ?? ????? ?? ??? ?????? ?????? ?????? ???? ????? ???? ?? ?? ?? ?? ??? CPU(?? ?? ?? ????)? ???? ?? ????? ??? ?? ????.

*?? 1.?? ? ?? ??? ?? ??? ?? ???? ???? Megatron-Core? ?? ?? ??*

??? ?? Megatron-Core? ?? ???? ??? Tensor ? ????? ?? ???? ??? ??????? ??? ?? ??? ? ???? ?? ?? ??? ?? ???? ?? ??? ??? ? ??? ????.?

Megatron-Core? ?? ? ?? API? PyTorch ???? API? ?? ???? ???? Megatron-Core ?? ?????? ?? ??? ? ????. ??? ?? ?? ??? ??? ?????? ???? ?? ? ?? PyTorch ???? ?? ????? ????? 26?(Nemotron-4 340B? ??) ?? 50?(Nemotron-4 15B? ??) ?? ? ????. ??? ?????? ???? Nemotron-4 340B? ????? ????? 42? ?? ? ????(?? 2).?

*?? 2. PyTorch ???? ? ?? Megatron-Core ????? Nemotron-4 340B ??? ????? ???? ??*

??? ???

Megatron-Core? v0.5 ????? ?? ??? ?? ??? ?? ?? ????? all-reduce? ???? ??? ?????. ?? ??? ???? ????? ??? ?? ??? ?? ?? ?????? ???? ??? ?? ??? ?? ??? ?????. ??? ??? ??? ??? ?? ??? ?? ?? Megatron-Core ???? ????, GPU? ?? ??? ?? ????? ??? ?? ??? ??? ? ?? ?????.?

?? 3? ??? ?? ??? ???? ?? ?? ??? ???? ?? ???? ??(??? ?? ??? 3*data_parallel_size), ?? ?? ??? 8? Nemotron-4 15B ??? GPU? ???? ?????. ??? ??? ?? ??? ?? ??? 32?? ?? ??? 96? ? ???? 34% ???? ?? ??? ? ?????. --overlap-grad-reduce ???? ??? ?? ??? ??? ? ? ?? ??? ???? ? ????. ??? ??? Megatron-Core ??? ?????.

?? 3. NVIDIA H100 GPU ? BF16 ???? ???? Nemotron-4 15B? ?? --overlap-grad-reduce ???? ??

Megatron-Core v0.6 ?????? ????? ??? ?? ??? ?? ???? ???? ?? ??? ?? ??? ??? ??? ????? ??? ??????. ?? ??? ?????? ??? ???? ????? all-reduce? ????? RS(reduce-scatter) ? ?? ?? AG(all-gather)? ?????. Megatron-Core? ?? ?? ??? reduce-scatter? ???? ?? ?? ??? all-gather? ?????. ??? ???? ?? ??? ??? Nemotron-4 15B? ??? ? ????. ?? 4? ?? 3? ??? ?? ???? ??? ???? ???? 15B ??? GPU? ???? ?????.

*?? 4. NVIDIA H100 GPU? ???? Nemotron-4 15B ??? ?? ??? RS(reduce-scatter) ??? ???? ?? ??? AG(all-gather) ??? ??? ?? ??*

?? v0.7 ???? ?? ?? ?? ??? ?? ? ?? ?? ????? ???? ??? ??? ?? ???? Megatron-Core? ??? ?? ?????. ?? ?? 300(GPU ? 3,000 ??)? ???? DP ????? ?? ??? ?? ???? ???? ??? ??? ? ????. ?? 5? ??? ?? ?? ?? ??? ?? ???? ???? ??(reduce-scatter ? all-gather ??? ?? ???? ??? ?????)?? ??? ?? ??? ?? 384? ? 15B ??? GPU? ???? ??? ????.

?? 5. ?? ??? ?? ??? 384? NVIDIA H100 GPU? ???? Nemotron-4 15B?? Megatron-Core 0.6 ? 0.7 ??? ??

Megatron-Core ???? ????? ?? ??? ???? ?? ?? ?? ??? ??? ? ? ??? ???? ?? ??? ? ????. Nemotron-4 340B ??? Megatron-Core?? ??? ???? ???? BF16? ???? ??? GPU?? ?? ?? ???? ?????. ? 2? Nemotron-4 340B ?? ???? ??? ?? ??? ??? Megatron-Core? GPU? ???? ?????. ?? TP ??? 8, PP ??? 12, ?? ????? ?? ?? 8, ??? ??? 4096???. ??? ??? Nemotron-4 340B ?? ???? ?????. ?? ?? ?? ??? ?? ???? ?? ? ?? ? ????.

???	GPU ?(H100)	??? ?? ??	?? ??	GPU? ???(TFLOP/s/GPU)
BF16	1536	16	768	419.3
BF16	3072	32	1536	418.3
BF16	6144	64	2304	405.0

? 2. Nemotron-4 340B ?? ???? ??? ?? ??? ??? Megatron-Core? GPU? ???

????

Megatron-Core? GitHub? NVIDIA/Megatron-LM ??????? ?? ??? ??? ? ??? Megatron-LM ?? NVIDIA NeMo? ?? ??? ? ????. ?? ?? ?????? Megatron-LM? ??? ?? PyTorch ?? ??? ????? ? ?? ?? ??? ??? ???? ????? ?????. Megatron-LM? Megatron-Core? ???? ?? ??????. ??? ??? Megatron-Core ??? ?????.

Megatron-Core? ??, ???, ?? ?? ? ??? ???? ??????? AI ????? ???? NVIDIA NeMo? ???? ???? ??? LLM? Megatron-Core ??? ???? ???? ? ?? AI? ?? ?? ???? ?? ??? ?????. ??? ??? NVIDIA NeMo ????? ??? ?????.