[Liger-Kernel] LinkedIn Liger-Kernel GRPO Loss Backward 최적화: 불필요한 연산 제거로 성능 50% 향상

2026년 10월 22일수정: 2026년 10월 22일

PR 링크: linkedin/Liger-Kernel#1550 상태: Merged | 변경: +10 / -0

들어가며

딥러닝 모델, 특히 강화학습 기반의 언어 모델을 학습시킬 때 손실 함수(Loss Function)의 계산은 전체 학습 시간에서 상당한 부분을 차지합니다. PPO(Proximal Policy Optimization)와 같은 알고리즘에서 사용되는 GRPO(Gated Reward Policy Optimization) 손실 함수의 backward 연산은 모델의 성능과 안정성에 직접적인 영향을 미치지만, 때로는 불필요한 계산으로 인해 성능 저하를 야기할 수 있습니다.

이번 분석 대상인 LinkedIn의 Liger-Kernel 레포지토리의 PR은 GRPO 손실 함수의 backward 커널에서 특정 조건 하에 불필요한 연산을 조기에 건너뛰도록 하여 성능을 크게 향상시킨 사례입니다. 구체적으로, dlogp 값이 0.0인 경우, 즉 PPO 클리핑 토큰(beta == 0.0)이거나 제로 어드밴티지(zero-advantage) 시퀀스에 해당하는 경우, HBM(High Bandwidth Memory)에서 불필요한 로짓(logits)을 로드하고 tl.exp 연산을 수행하는 것을 방지합니다. 이 PR을 통해 해당 연산이 약 51.2%의 성능 향상을 가져왔다고 합니다.

본 글에서는 이 PR의 코드 변경 사항을 자세히 살펴보고, 왜 이러한 변경이 성능 향상으로 이어졌는지, 그리고 이 최적화가 주는 일반적인 교훈은 무엇인지 기술 블로그 형식으로 풀어보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 src/liger_kernel/ops/grpo_loss.py 파일 내 두 개의 함수, _grpo_loss_bwd_kernel_seq와 _grpo_loss_bwd_kernel에 집중되어 있습니다. 두 함수 모두 GRPO 손실 함수의 backward pass를 처리하며, 특정 조건에서 조기 반환(early return) 로직을 추가하여 불필요한 계산을 건너뛰도록 수정되었습니다.

`_grpo_loss_bwd_kernel_seq` 함수

이 함수는 시퀀스 레벨에서의 GRPO 손실 backward를 처리합니다. 변경 전후 코드를 비교해보겠습니다.

Before:

# ... (이전 코드) ...

    dlogp = dlogp / TEMPERATURE
    tl.debug_barrier()
    for start_n in tl.range(0, N, BLOCK_N):
        cols = start_n + tl.arange(0, BLOCK_N)
        tl.store(DLOGITS + cols, 0.0, mask=cols < N)

# ... (이후 코드) ...

After:

@@ -413,6 +413,11 @@ def _grpo_loss_bwd_kernel_seq(
             dlogp += BETA * (1 - tl.exp(ref_logp - logp)) * dloss
 
     dlogp = dlogp / TEMPERATURE
+    if dlogp == 0.0:
+        for start_n in tl.range(0, N, BLOCK_N):
+            cols = start_n + tl.arange(0, BLOCK_N)
+            tl.store(DLOGITS + cols, 0.0, mask=cols < N)
+        return
     tl.debug_barrier()
     for start_n in tl.range(0, N, BLOCK_N):
         cols = start_n + tl.arange(0, BLOCK_N)

변경 사항을 보면, dlogp = dlogp / TEMPERATURE 라인 직후에 if dlogp == 0.0: 조건문이 추가되었습니다. 이 조건이 참일 경우, DLOGITS 텐서에 0.0을 저장하고 즉시 함수를 반환합니다. 이전 코드에서는 dlogp 값이 0이더라도 이후의 tl.debug_barrier()와 tl.store 연산이 계속 수행되었습니다.

`_grpo_loss_bwd_kernel` 함수

이 함수는 더 일반적인 GRPO 손실 backward를 처리하며, _grpo_loss_bwd_kernel_seq와 유사한 구조를 가집니다. 변경 전후 코드는 다음과 같습니다.

Before:

# ... (이전 코드) ...

    dlogp = dlogp * dloss / TEMPERATURE
    tl.debug_barrier()
    for start_n in tl.range(0, N, BLOCK_N):
        cols = start_n + tl.arange(0, BLOCK_N)
        tl.store(DLOGITS + cols, 0.0, mask=cols < N)

# ... (이후 코드) ...

After:

@@ -543,6 +548,11 @@ def _grpo_loss_bwd_kernel(
             dlogp += BETA * (1 - tl.exp(ref_logp - logp))
 
     dlogp = dlogp * dloss / TEMPERATURE
+    if dlogp == 0.0:
+        for start_n in tl.range(0, N, BLOCK_N):
+            cols = start_n + tl.arange(0, BLOCK_N)
+            tl.store(DLOGITS + cols, 0.0, mask=cols < N)
+        return
     tl.debug_barrier()
     for start_n in tl.range(0, N, BLOCK_N):
         cols = start_n + tl.arange(0, BLOCK_N)

마찬가지로, dlogp = dlogp * dloss / TEMPERATURE 계산 직후 if dlogp == 0.0: 조건이 추가되어, 조건이 참일 경우 DLOGITS에 0.0을 저장하고 함수를 종료합니다. 이전에는 dlogp가 0이더라도 이후의 tl.debug_barrier()와 tl.store 연산이 실행되었습니다.

왜 이게 좋은가?

이 PR의 변경 사항은 다음과 같은 이유로 매우 훌륭한 최적화라고 할 수 있습니다.

불필요한 연산 제거: dlogp == 0.0인 경우는 PPO 클리핑 토큰(beta == 0.0)이거나 제로 어드밴티지 시퀀스에 해당하는 경우입니다. 이 경우, 손실 함수 업데이트에 기여하는 dlogp의 값이 0이므로, 해당 토큰에 대한 그래디언트(DLOGITS)는 0이 됩니다. 이전 코드에서는 dlogp가 0이더라도 HBM에서 관련 로짓을 로드하고 tl.exp와 같은 연산을 수행했습니다. 이 PR은 dlogp == 0.0 조건을 확인하여 이러한 불필요한 메모리 로드 및 계산을 건너뛰도록 합니다. 이는 특히 GPU 메모리 대역폭과 연산 능력을 절약하는 데 크게 기여합니다.
명확한 성능 향상: PR 설명에 따르면, 특정 워크로드((B, L, V, zero-grad %) = (4, 512, 32000, 50%))에서 main 대비 PR 버전의 실행 시간이 2.014ms에서 0.983ms로 51.2% 단축되었습니다. 이는 절반 가까이 되는 시간을 절약한 것으로, 대규모 모델 학습 시 상당한 시간 절약 효과를 기대할 수 있습니다.
정확성 유지: PR에서는 Outputs and gradients are bitwise identical (torch.equal).라고 명시하며, 이 최적화가 모델의 정확성에 영향을 미치지 않음을 검증했습니다. 이는 성능 최적화에서 가장 중요한 부분 중 하나입니다. 즉, 성능을 높이면서도 모델의 학습 결과는 동일하게 유지된다는 것을 보장합니다.
코드 가독성 및 유지보수성 향상: 조기 반환(early return) 로직은 코드의 흐름을 더 명확하게 만들어 줍니다. 특정 조건에서 함수가 즉시 종료된다는 것을 명시함으로써, 코드의 의도를 파악하기 쉬워지고 잠재적인 버그를 줄이는 데 도움이 됩니다.

일반적인 교훈

이 PR은 다음과 같은 일반적인 최적화 교훈을 제공합니다:

'Nothing is faster than not doing it': 불필요한 연산을 아예 실행하지 않는 것이 가장 효과적인 최적화 방법입니다. 조건문을 활용하여 특정 상황에서만 실행되어야 하는 코드 블록을 명확히 구분하는 것이 중요합니다.
종료 조건(Exit Condition) 최적화: 함수의 시작 부분이나 계산 중간에 명확한 종료 조건을 두어, 더 이상 계산이 필요 없거나 결과가 결정된 경우 즉시 반환하는 패턴은 성능 향상에 큰 도움이 됩니다.
벤치마킹의 중요성: PR에서는 구체적인 성능 수치(51.2% 향상)를 제시하며 변경의 효과를 입증했습니다. 실제 환경에서의 성능 측정은 최적화의 가치를 증명하는 데 필수적입니다.
정확성 검증: 성능 최적화 과정에서 모델의 정확성이나 수렴성에 영향을 미치지 않는다는 것을 철저히 검증하는 것이 중요합니다. torch.equal과 같은 도구를 사용하여 변경 전후의 결과가 동일함을 보장해야 합니다.

리뷰 피드백 분석

이 PR에는 별도의 리뷰 댓글이 제공되지 않았지만, PR 설명 자체에서 테스트 결과와 검증 과정을 상세히 기술하고 있습니다. 특히, make test를 통해 91개의 테스트 케이스를 통과했고, make test-convergence를 통해 수렴성까지 확인했다는 점은 이 변경이 안정적이고 신뢰할 수 있음을 보여줍니다. 또한, 하드웨어 정보(NVIDIA GeForce RTX 4050 Laptop GPU), 실험 설정(locked GPU clocks, 5 interleaved ABBA rounds, bf16) 등을 명시하여 결과의 재현성을 높였습니다.

References

Triton Language Documentation - Triton 언어의 기본 연산 및 저장(store) 관련 정보
PyTorch Documentation - PyTorch 텐서 연산 관련 일반 정보 (이 PR에서 직접적으로 사용된 특정 함수는 없으나, Triton은 PyTorch 백엔드와 함께 사용되는 경우가 많음)

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer SM107(Rubin) 아키텍처를 위한 MXFP8 GEMM 최적화 분석
현재글 : [Liger-Kernel] LinkedIn Liger-Kernel GRPO Loss Backward 최적화: 불필요한 연산 제거로 성능 50% 향상
다음글 없음

댓글

관련 포스트

PR Analysis 의 다른글