Hugging Face와 Axolotl이 Accelerate 라이브러리에 ND-Parallel 통합으로 여러 GPU 간 복합 병렬화 전략 조합을 설정 파일 몇 줄로 구현 가능하게 구축

Context

대규모 모델을 여러 GPU에서 학습할 때 서로 다른 병렬화 전략의 복잡성으로 인해 메모리 제약과 통신 오버헤드가 발생하며, 단일 디바이스에 맞지 않는 모델의 경우 기존 데이터 병렬화만으로는 학습이 불가능했다.

Technical Solution

• ParallelismConfig 클래스를 통한 병렬화 전략 조합: dp_shard_size, dp_replicate_size, cp_size, tp_size 파라미터로 Data Parallelism, Fully Sharded Data Parallelism, Context Parallelism, Tensor Parallelism을 자유롭게 구성
• FSDP2(Fully Sharded Data Parallelism)를 모델 가중치, 그래디언트, 옵티마이저 상태 전체에 적용: 각 GPU에 전체 모델을 복제하지 않고 파라미터를 균등하게 분산하여 메모리 사용량 감소
• 레이어 단위 동적 가중치 수집: 정방향/역방향 패스 전에만 필요한 레이어의 가중치를 모아 사용 후 다시 분산하여 메모리와 통신 오버헤드 트레이드오프 조정
• Axolotl 프레임워크에 ND-Parallel 기능 통합: dp_shard_size, dp_replicate_size, context_parallel_size, tensor_parallel_size 설정값만 YAML 파일에 추가하여 복잡한 코드 작성 없이 병렬화 구성
• 그래디언트 체크포인팅(Gradient Checkpointing) 지원: 정방향 패스에서 활성화값의 일부만 메모리에 유지하고 역방향 패스에서 재계산하여 활성화 메모리를 60~80% 감소
• 학습률 스케일링 적용: 유효 배치 크기 증가 시 선형 스케일링(scaled_lr = base_lr × (effective_batch_size / base_batch_size)) 또는 제곱근 스케일링(scaled_lr = base_lr × √(effective_batch_size / base_batch_size))으로 학습 안정성 유지

Impact

그래디언트 체크포인팅 적용 시 활성화 메모리 60~80% 감소, 학습 시간 약 20~30% 증가.

Key Takeaway

수십억에서 수백억 파라미터를 가진 대규모 모델 학습 시 단일 병렬화 기법이 아닌 여러 병렬화 전략을 조합하되, 메모리 효율성과 통신 오버헤드 사이의 트레이드오프를 명확히 이해하고 각 디바이스의 제약 조건에 맞춰 선택해야 한다.