본문으로 건너뛰기

[sglang] ROCm 환경에서 BF16 All-Reduce의 수치 안정성 확보하기: QuickReduce의 FP16 Saturation 이슈 해결

PR 링크: sgl-project/sglang#34484 상태: Merged | 변경: +184 / -5

들어가며

대규모 언어 모델(LLM) 학습 및 추론에서 분산 컴퓨팅의 핵심은 GPU 간의 효율적인 데이터 교환, 즉 All-Reduce 연산입니다. SGLang 프로젝트의 QuickReduce는 성능 최적화를 위해 BF16(Bfloat16) 입력을 FP16(Half-precision)으로 변환하여 연산을 수행합니다. FP16은 BF16보다 연산 속도가 빠르지만, 표현 가능한 수의 범위(Dynamic Range)가 훨씬 좁다는 치명적인 단점이 있습니다.

최근 GPT-OSS와 같은 모델에서 BF16 값이 FP16의 최대 유한 값인 65,504를 초과하는 현상이 발생했습니다. 이로 인해 FP16 변환 과정에서 inf(무한대)가 발생하고, 이것이 전체 텐서로 전파되어 모델의 출력이 망가지는 문제가 보고되었습니다. 이번 PR은 이러한 'FP16 Saturation' 문제를 해결하면서도 성능 저하를 최소화한 기법들을 담고 있습니다.

코드 분석: 핵심 변경 사항

1. Range Guard (Scaling) 도입

가장 직접적인 해결책은 FP16으로 변환하기 전에 값의 크기를 줄이는 것입니다. CodecFP 구조체에 kCastScaleLog2를 도입하여 데이터를 로드할 때 16($2^4$)으로 나누고, 저장할 때 다시 16을 곱하는 방식을 채택했습니다.

[Before]

// 기존에는 별도의 스케일링 없이 바로 변환
for (int j = 0; j < 4; ++j) {
  float2 f = __bfloat1622float2(bf_buf[j]);
  half2_buf[j] = __float22half2_rn(f);
}

[After]

// quick_all_reduce_base.h
static constexpr int kQRFp16CastScaleLog2Fp = 4;     // S = 16

// quick_all_reduce.cuh
if constexpr (Codec::kCastScaleLog2 == 0) {
    // 기존 경로 (Quantized Codec 등)
    f.x *= kCastInvScale;
    f.y *= kCastInvScale;
    half_buf[j] = __float22half2_rn(f);
} else {
    // 새로운 Range Guard 경로
    half_buf[j] = scaled_bfloat162_to_half2(bf_buf[j], kCastInvScale);
}

여기서 $S=16$을 선택한 이유는 2의 거듭제곱으로 나누고 곱하는 연산이 부동소수점 연산에서 정밀도 손실 없이 정확하게 수행되기 때문입니다.

2. 컴파일러 최적화 방지 (Inline Assembly)

단순히 (value * scale)을 수행한 뒤 FP16으로 캐스팅하면, LLVM 컴파일러가 최적화 과정에서 순서를 바꿔 (value를 FP16으로 캐스팅) * scale로 실행할 위험이 있습니다. 이 경우 캐스팅 시점에 이미 inf가 되어버려 스케일링이 무의미해집니다. 이를 방지하기 위해 AMD GPU의 ISA 명령어를 직접 호출하는 인라인 어셈블리를 사용했습니다.

[New Function]

__quickreduce_device_inline__ half2 scaled_bfloat162_to_half2(nv_bfloat162 value, float scale) {
  float2 scaled = __bfloat1622float2(value);
  scaled.x *= scale;
  scaled.y *= scale;

  int packed;
  // v_cvt_pk_f16_f32 명령어를 명시적으로 사용하여 컴파일러의 재배치를 차단
  asm volatile("v_cvt_pk_f16_f32 %0, %1, %2" : "=v"(packed) : "v"(scaled.x), "v"(scaled.y));
  return *reinterpret_cast<half2*>(&packed);
}

3. 하드웨어 레벨의 Clamping 설정

ROCm 환경(gfx942, gfx950)에서 FP16 연산 결과가 범위를 벗어날 때 inf 대신 최대값(65,504)으로 고정(Clamping)되도록 하드웨어 레지스터를 설정했습니다. 또한 memory 클로버(clobber)를 추가하여 컴파일러가 이 설정 이전에 연산을 수행하지 못하도록 강제했습니다.

[Before]

__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
#if defined(__gfx942__)
  if (value) {
    asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::);
  } // ...
#endif
}

[After]

__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
#if defined(__gfx942__) || defined(__gfx950__)
  if (value) {
    // "memory" 제약 조건을 추가하여 컴파일러 배리어 역할 수행
    asm volatile("s_setreg_imm32_b32 0xdc1, 1;" ::: "memory");
  } // ...
#endif
}

왜 이게 좋은가?

1. 성능과 정확도의 균형

BF16 네이티브 커널을 사용하는 대안은 수치적으로 완벽하지만, 기존 FP16 기반 QuickReduce보다 22~46% 더 느립니다. 이번 PR은 스케일링과 하드웨어 클램핑을 조합하여 성능 저하 없이(오히려 일부 케이스에서 9.6% 향상) inf 발생 문제를 완벽히 해결했습니다.

2. 비트 정확도(Bit-exact) 유지

테스트 결과, 일반적인 수치 범위 내에서는 기존 방식과 완전히 동일한 비트 결과를 생성함을 확인했습니다. 이는 최적화가 모델의 추론 결과에 부정적인 영향을 주지 않음을 보장합니다.

3. 컴파일러에 대한 깊은 이해

단순한 C++ 코드 작성을 넘어, asm volatile과 memory 배리어를 통해 컴파일러의 최적화 동작을 제어한 점은 시니어 엔지니어링의 정수를 보여줍니다. 특히 리뷰어 yctseng0211이 지적한 gfx942 호환성 이슈(특정 명령어가 지원되지 않음)는 실제 하드웨어 타겟팅 시 발생할 수 있는 미묘한 차이를 잘 보여줍니다.

일반적인 교훈

  1. 수치적 한계 인지: 고성능을 위해 데이터 타입을 하향 변환(Downcasting)할 때는 항상 타겟 타입의 Dynamic Range를 고려해야 합니다.
  2. 컴파일러는 믿되 검증하라: 고도로 최적화된 코드를 작성할 때는 컴파일러가 의도와 다르게 코드를 재배치할 수 있음을 인지하고, 적절한 Barrier나 Inline Assembly를 사용해야 합니다.
  3. 하드웨어 특성 활용: GPU 아키텍처가 제공하는 특수 레지스터(예: Clamping 모드)를 활용하면 소프트웨어 레벨의 복잡한 예외 처리를 하드웨어 가속으로 대체할 수 있습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글