[LlamaFactory] Qwen3.5 모델의 GDN 선형 어텐션에 대한 시퀀스 병렬 처리 및 Pack 지원 추가

2026년 8월 20일수정: 2026년 8월 20일

PR 링크: hiyouga/LlamaFactory#10727 상태: Merged | 변경: +269 / -16

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기뿐만 아니라 처리할 수 있는 시퀀스 길이의 확장에도 크게 의존하고 있습니다. 특히 Qwen3.5와 같은 모델들은 효율적인 장문 시퀀스 처리를 위해 다양한 어텐션 메커니즘을 활용합니다. 하지만 기존의 시퀀스 병렬 처리(Sequence Parallelism, SP) 또는 컨텍스트 병렬 처리(Context Parallelism, CP) 구현은 특정 어텐션 방식, 예를 들어 표준 Flash Attention에는 잘 적용되었지만, Qwen3.5 모델에서 사용되는 Gated Delta Network (GDN) 선형 어텐션 레이어에서는 제대로 지원되지 않는 문제가 있었습니다. 이로 인해 cp_size > 1 환경에서 GDN 레이어가 CP의 이점을 제대로 활용하지 못하거나, 심지어 잘못된 계산을 수행하는 상황이 발생했습니다. 또한, Pack (cu_seqlens) 지원 부재는 GPU 활용률을 더욱 저하시키는 요인이었습니다.

본 PR은 이러한 문제를 해결하기 위해 Qwen3.5 모델의 GDN 선형 어텐션 레이어에 대한 CP 지원을 추가하고, Pack 기능을 통합하여 GPU 활용도를 극대화하는 것을 목표로 합니다.

코드 분석

이번 PR의 핵심 변경 사항은 src/llamafactory/v1/plugins/model_plugins/parallelization/gdn_attention.py 파일의 신규 추가와 src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py 파일의 수정입니다. 이를 통해 GDN 레이어의 특성을 고려한 CP 및 Pack 기능이 구현되었습니다.

1. `gdn_attention.py` - GDN 레이어를 위한 CP 및 Pack 구현

이 파일은 GDN 레이어에 특화된 CP 로직을 담고 있습니다.

is_gdn_layer(layer): 주어진 모듈이 GDN 레이어인지 식별하는 함수입니다. layer_type == "linear_attention" 또는 block_type == "linear_attention" 속성을 확인합니다.
_get_gdn_module(module): GDN 레이어 또는 GDN을 포함하는 DecoderLayer에서 실제 GDN 모듈을 추출합니다.
get_parameter_local_cp(param, dim, cp_group, split_sections=None): CP 환경에서 각 랭크가 사용할 파라미터(conv1d weight/bias, A_log, dt_bias 등)를 적절히 슬라이싱하는 함수입니다. split_sections 인자를 통해 Q/K/V와 같이 여러 부분으로 나뉘는 파라미터가 각 CP 랭크에 비례적으로 할당되도록 보장합니다.
gdn_forward_with_cp(self, hidden_states, attention_mask=None, **kwargs): CP 환경 (cp_size > 1)에서 GDN 레이어의 forward 메소드를 대체하는 핵심 함수입니다. CP가 아닐 경우 (cp_size <= 1) 원래의 forward 메소드(self.original_forward)로 즉시 폴백(fallback)하여 오버헤드를 최소화합니다. 이 함수는 Pack 모드를 지원하며, position_ids로부터 cu_seqlens를 계산하여 Flash Attention의 가변 길이(varlen) 모드를 활용합니다.

주요 구현 디테일:
- Per-component all_to_all: Q, K, V와 같은 각 컴포넌트를 독립적으로 4D 텐서로 변환하고 SeqAllToAll4D를 사용하여 all_to_all 통신을 수행합니다. 이는 기존의 통합된 hidden-split 방식에서 발생할 수 있는 Q/K/V 서브 그룹 간의 불균등한 할당 버그를 방지합니다.
- Parameter slicing: get_parameter_local_cp 함수를 사용하여 conv1d의 가중치와 편향, A_log, dt_bias 등의 파라미터를 CP 랭크별로 정확하게 분할합니다.
- Pack support: position_ids를 활용하여 cu_seqlens를 계산하고, 이를 causal_conv1d_fn (Flash Linear Attention)에 전달하여 Pack 기능을 활성화합니다. 이는 GPU 메모리 활용률을 크게 향상시킵니다.
- GPU-only Pack: Pack 기능은 causal_conv1d_fn과 chunk_gated_delta_rule에서 Flash Triton 커널에 의존하므로 GPU 환경에서만 작동합니다.

2. `sequence_parallel.py` - GDN 포워드 함수 등록

src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py 파일의 apply_sequence_parallel() 함수는 모델의 각 레이어를 순회하며 CP 적용 대상 모듈을 찾고, GDN 레이어의 경우 gdn_forward_with_cp 함수로 forward 메소드를 교체하는 역할을 합니다. 이전에는 Qwen3.5 모델의 경우 CP 지원이 불가능하다는 예외 처리가 있었으나, 이 PR을 통해 해당 예외 처리가 제거되었습니다.

--- a/src/llamafactory/v1/core/base_trainer.py
+++ b/src/llamafactory/v1/core/base_trainer.py
@@ -147,11 +147,6 @@
         self.state.epoch = self._resume_epoch
 
         if self.args.cp_size > 1:
-            # qwen3.5 is not supported because of the different attention implementation, which will be supported in the future.
-            if model.config.model_type == "qwen3_5":
-                raise RuntimeError(
-                    "Sequence parallel is not supported for qwen3.5 model due to its different attention implementation, which will be supported in the future."
-                )
             from ..plugins.model_plugins.parallelization.sequence_parallel import SequenceParallelModelPlugin
 
             if model.config._attn_implementation != "flash_attention_2":

위 diff는 기존의 Qwen3.5 모델에 대한 CP 지원 불가 예외 처리가 제거되었음을 보여줍니다.

왜 이게 좋은가?

1. 정확성 보장 및 버그 수정

기존 CP 구현은 GDN 레이어의 내부 구조를 고려하지 않아 cp_size > 1 환경에서 잘못된 계산을 유발할 수 있었습니다. 이 PR은 GDN 레이어의 각 컴포넌트(Q, K, V 등)를 독립적으로 처리하고 파라미터를 비례적으로 슬라이싱함으로써 이러한 계산 오류를 수정하고 CP 환경에서의 모델 정확도를 보장합니다.

2. GPU 활용률 극대화 (Pack 지원)

GDN 레이어에 Pack (cu_seqlens) 지원을 추가함으로써, 특히 배치 크기가 1이고 시퀀스 길이가 가변적인 경우 GPU의 연산 능력을 더욱 효율적으로 사용할 수 있게 되었습니다. 이는 Flash Attention의 가변 길이(varlen) 모드를 활용하여 불필요한 패딩 연산을 줄이고 메모리 접근을 최적화하는 효과를 가져옵니다. 리뷰어 frozenleaves가 언급했듯이, GDN과 Flash Attention(FLA)의 통합은 성능 향상에 매우 중요합니다.

3. 확장성 향상

GDN 레이어에 대한 CP 지원은 Qwen3.5 모델을 더 큰 배치 크기 또는 더 긴 시퀀스 길이로 확장할 수 있는 기반을 마련합니다. 이는 대규모 모델 학습 및 추론 시 필수적인 요소입니다.

4. 성능 수치 (테스트 결과)

PR 설명에 제시된 테스트 결과는 이 변경의 효과를 명확히 보여줍니다:

CP=1 vs CP=2 손실 정렬: CP=2 환경에서도 CP=1 환경 대비 상대 정밀도 2% 이내를 유지하며, CP 적용 시에도 모델 성능 저하가 미미함을 입증합니다.
CP+Pack vs CP+non-Pack 손실 정렬: Pack 기능 적용 시에도 상대 편차가 2% 미만으로 유지되어, Pack 기능이 정확도를 해치지 않으면서 성능 향상에 기여함을 보여줍니다.
메모리 사용량 감소: CP=2 환경에서의 단일 GPU 메모리 사용량이 CP=1보다 감소하는 것을 확인하여, CP의 메모리 절감 효과를 입증합니다.

5. 일반적인 교훈

모델별 특성 고려: 범용적인 병렬 처리 기법을 특정 모델 아키텍처(예: GDN 어텐션)에 적용할 때는 해당 아키텍처의 고유한 특성(예: 컴포넌트 분리, 파라미터 구조)을 면밀히 분석하고 반영해야 합니다.
최신 기술 통합: Flash Attention과 같은 최신 고성능 커널을 기존 병렬 처리 프레임워크와 통합하는 것은 GPU 활용률을 극대화하는 데 매우 효과적입니다.
점진적 검증: 복잡한 병렬 처리 로직은 작은 규모의 모델이나 레이어 수를 줄인 모델로 빠르게 검증하고, 점진적으로 실제 환경에 적용하는 것이 안정성을 높이는 방법입니다.

결론

이번 PR은 Qwen3.5 모델의 GDN 선형 어텐션 레이어에 대한 시퀀스 병렬 처리 및 Pack 지원을 성공적으로 추가함으로써, 모델의 정확성을 유지하면서 GPU 활용률과 확장성을 크게 향상시켰습니다. 이는 LLM의 효율적인 학습 및 추론을 위한 중요한 진전이며, 향후 유사한 아키텍처를 가진 모델에도 적용될 수 있는 좋은 선례를 남겼습니다.

참고 자료

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.

PR Analysis 의 다른글

이전글 [flashinfer] Blackwell 아키텍처를 위한 MXFP8/MXFP4 기반 고성능 MoE 추론 최적화
현재글 : [LlamaFactory] Qwen3.5 모델의 GDN 선형 어텐션에 대한 시퀀스 병렬 처리 및 Pack 지원 추가
다음글 [Liger-Kernel] Liger-Kernel의 RMSNorm 최적화: cuTile 도입과 CuTe DSL 성능 개선