만약 당신이 하나 이상의 모달리티 (Modality)에 걸쳐 모델을 학습시킨다면, 아마 이런 경험을 해보았을 것입니다. Megatron-LM은 LLM (대규모 언어 모델) 학습에는 검증되었지만, 시각-언어-행동 (Vision-Language-Action, VLA) 로봇을 위해 구축되지는 않았습니다. LeRobot과 OpenPI는 체화된 모델 (Embodied models)의 연구 반복을 위해 구축되었을 뿐, 프로덕션 처리량 (Production throughput)을 위한 것이 아닙니다. 확산 모델 (Diffusion) 학습은 또 다른 격리된 영역에 존재합니다. 결국 당신은 각각 고유한 병렬화 전략 (Parallelism strategy), 체크포인트 형식 (Checkpoint format), 성능 한계 (Performance ceiling)를 가진 세 개 또는 네 개의 코드베이스를 유지 관리하게 됩니다.

우리는 이 모든 것을 하나의 프레임워크로 통합하기 위해 LoongForge를 방금 오픈 소스로 공개했습니다. 이것은 Baidu Baige의 학습 가속 스택인 AIAK-Training-LLM의 오픈 소스 버전으로, 교육, 컴퓨터 비전 (Computer vision), 그리고 체화된 AI (Embodied AI) 분야의 기업 고객을 위해 프로덕션 학습을 수행해 왔습니다. 일반적으로 고객의 베이스라인 (Baseline) 대비 30~50%의 속도 향상을 제공하며, 가장 큰 규모의 실행에서는 **5,000개 이상의 가속기 (Accelerators)**에 도달했습니다.

Repo: https://github.com/baidu-baige/LoongForge (Apache-2.0)

무엇이 다른가

1. 체화된 모델 (Embodied models) 학습 가속.
이 부분은 대부분의 학습 프레임워크가 다루지 않는 영역입니다. LoongForge는 VLA 및 월드-액션 (World-Action, WAM) 모델 — Pi0.5, GR00T N1.6, X-VLA, FastWAM, DreamZero, LingBot-VA — 을 위해 Megatron 코어와 분리된 전용 torch-native DDP/FSDP 서브시스템을 제공하며, 유연한 DDP / ZeRO-1 / FSDP / HSDP 전략을 지원합니다. 현재 로봇 정책 (Robot policies)을 학습시킨다면, 프로덕션급 처리량 엔지니어링 (Throughput engineering)은 매우 드뭅니다. LoongForge에서는 이것이 내장되어 있습니다.

2. 하나의 프레임워크, 네 가지 모달리티.
LLM, VLM (시각-언어 모델), 확산 모델 (Diffusion), 그리고 체화된 모델 (Embodied) — 이 모든 것이 통일된 모델 추상화 (Model abstraction) 아래에 있습니다. VLM은 교체 가능한 ViT + LLM 구성 요소로부터 설정 기반 (Config-driven)으로 조립되므로, 커스텀 ViT + LLM 백본 (Backbone)을 만드는 것은 포크 (Fork)가 아닌 YAML 변경만으로 가능합니다.

3. 두 가지 하드웨어 백엔드 지원, 네이티브 방식.
최소한의 개입을 특징으로 하는 플러그인 설계를 통해 NVIDIA GPU와 Kunlun XPU를 지원합니다. 동일한 학습 코드가 두 하드웨어 모두에서 실행됩니다.

성능

각 모델 유형이 통상적으로 학습되는 특정 오픈 소스 베이스라인 (Baseline)을 능가하도록 튜닝되었습니다.

내부 엔지니어링 (The engineering underneath)

• 유연한 멀티모달 구성 (Flexible Multi-Modal Composition) — 교체 가능한 ViT 및 LLM 구성 요소를 사용하여 설정(Configuration) 기반으로 VLM을 조립합니다.

• 이종 병렬성 (Heterogeneous Parallelism) — 최적의 처리량(Throughput)과 메모리 활용을 위해 모델 구성 요소별(예: ViT vs. LLM)로 독립적인 TP / DP / 재계산(recompute)을 적용합니다. oxed{blog}

• 디커플링된 인코더-디코더 학습 (Decoupled Encoder-Decoder Training) — ViT와 LLM을 독립적인 작업으로 분리하여, 인코더로 인해 발생하는 파이프라인 버블(pipeline bubbles)을 제거합니다.

• DP 부하 분산 (DP Load Balancing) — 부하 인지형(Load-aware) 데이터 재분배를 통해 시퀀스 패킹(sequence-packing) 불균형을 완화하고, 멀티 노드 확장 효율성을 향상시킵니다. oxed{blog}

• MoE 네이티브 최적화 (MoE-Native Optimization) — All2All / 활성화 오프로드(activation offload) / 연산(compute)을 중첩(Overlapped) 처리하며, DeepSeek-V3, Qwen3-MoE 등에서 상위 단계인 Megatron-LM보다 추가적인 메모리 절감을 실현합니다.

• MoE 전문가 부하 분산 (MoE Expert Load Balancing) — 토폴로지 인지(topology-aware) 알고리즘을 사용하여 빈번하게 사용되는 전문가(hot experts)를 동적으로 복제함으로써, 전문가 병렬(Expert Parallel, EP) 워크로드를 균형 있게 맞추고 학습 효율을 높입니다.

• 적응형 FP8 학습 (Adaptive FP8 Training) — 표준 **블록 단위 FP8 (blockwise FP8)**을 사용하여 LLM 및 VLM에 대해 엔드 투 엔드(End-to-end) FP8을 지원합니다. 선택 가능한 적응형 (adaptive) 모드는 GEMM 형상(shape)과 효율성에 따라 연산자별 정밀도를 결정합니다.

• 커스텀 융합 연산자 (Custom Fused Operators) — DSA 스타일 모델을 위한 FusedDSA와 같은 융합 커널(Fused kernels)을 제공합니다.

• 유연한 체크포인팅 (Flexible Checkpointing) — 오프라인 양방향 Megatron ↔ HuggingFace 변환 기능과 네이티브 온라인 HF 로드/저장 기능을 지원하여, 워크플로 전반에 걸쳐 포맷 장벽이 없습니다.

• 다재다능한 파이프라인 및 데이터 도구 (Versatile Pipelines & Data Tools) — 사전 학습 (Pretrain) / 중간 학습 (MidTrain) / SFT / LoRA를 즉시 사용할 수 있으며, 데이터셋 형식 변환 및 시퀀스 패킹(sequence packing) 기능이 내장되어 있습니다.

• Embodied Model Training (체화된 모델 학습) — Megatron 코어와 분리되어 있으며, VLA 및 world-action (WAM) 모델(예: Pi0.5, GR00T N1.6, FastWAM)을 위한 전용 torch-native DDP/FSDP 서브시스템을 제공하며, 유연한 DDP / ZeRO-1 / FSDP / HSDP 전략을 지원합니다.

• Heterogeneous Hardware (이기종 하드웨어) — 최소한의 개입만 필요한 플러그인 설계를 통해 NVIDIA GPU 및 Kunlun XPU를 네이티브로 지원합니다.

대상 사용자

언어, 시각-언어(vision-language), 확산(diffusion), 로보틱스를 위해 별도의 스택을 유지 관리하는 데 지쳤고, 처리량(throughput)을 중요하게 생각하며 대규모로 사전 학습(pre-training) 또는 미세 조정(fine-tuning)을 수행하는 팀에게 적합합니다. Apache-2.0 라이선스를 따르며, Pretrain / MidTrain / SFT / LoRA를 즉시 사용할 수 있고, NVIDIA 및 Kunlun에서 실행됩니다.

사용해 보기 / 기여하기

• ⭐ Repo: https://github.com/baidu-baige/LoongForge
• Docs: https://loongforge.readthedocs.io

이 분야에서 모델을 학습시키고 계신다면, 여러분의 피드백 — 특히 체화된(embodied) 서브시스템에 대한 의견 — 을 진심으로 환영합니다. 다음에 어떤 모델이 지원되기를 원하시나요?