분산 LLM 학습 10 - Sequence Parallel과 긴 컨텍스트의 비용
모델 크기만 커지는 것이 아니라 컨텍스트 길이도 길어지면 activation 메모리와 통신 패턴이 다시 달라진다
긴 컨텍스트는 다른 종류의 압박을 만든다
LLM 학습에서 모델 파라미터만 커지는 것이 아니다. context length가 길어지면 activation과 attention 관련 텐서가 빠르게 커진다. 이때 기존 tensor parallel만으로는 메모리와 통신 문제를 충분히 다루기 어려울 수 있다.
이 지점에서 sequence parallel 같은 아이디어가 나온다.
sequence 방향으로 나눈다는 뜻
tensor parallel이 hidden dimension이나 weight 축을 나누는 데 집중한다면, sequence parallel은 토큰 차원 쪽으로 일부 작업을 분산하는 관점에 가깝다.
왜 이런 방식이 필요할까?
- normalization 같은 연산은 hidden 방향으로 분할만으로는 메모리 이점이 제한적일 수 있다
- 긴 sequence에서는 activation 저장량이 급격히 커진다
- attention 관련 중간 결과가 통신과 메모리 모두를 압박한다
즉 긴 컨텍스트에서는 "모델이 크다"와는 다른 형태의 병목이 나타난다.
좋은 점과 까다로운 점
sequence parallel의 장점은 activation 메모리를 줄이고 일부 연산을 더 잘 분산할 수 있다는 점이다. 하지만 대가도 있다.
- rank 간 shape 정렬이 더 복잡해진다
- 특정 연산 전후로 gather/scatter가 필요할 수 있다
- 디버깅이 훨씬 어려워진다
특히 어떤 연산은 sequence 분할과 잘 맞지만, 어떤 연산은 결국 전체 정보가 다시 필요해진다. 따라서 이 전략은 단독으로 보기보다 tensor parallel과 조합해서 이해하는 편이 맞다.
긴 컨텍스트에서 보게 되는 질문
- attention 메모리 사용량의 주범이 무엇인가
- activation recomputation으로 충분한가
- sequence parallel을 도입하면 통신 증가를 감당할 수 있는가
- FlashAttention류 kernel 최적화와 병렬화 전략을 어떻게 조합할 것인가
이 질문은 점점 kernel과 framework, distributed runtime의 경계가 흐려지는 지점이다.
왜 여기서 pipeline parallel로 넘어가야 하는가
지금까지는 하나의 레이어 내부를 쪼개는 관점이 강했다. 하지만 모델이 더 커지면 레이어 그룹 자체를 여러 stage로 나누는 방식도 필요해진다. 그게 pipeline parallel이다.
다음 글에서는 pipeline parallel의 기본 구조와 왜 micro-batch scheduling이 필수인지 본다.