Backward 구현 및 Graph_excutor_v2 에 이식

1) 파일 단위 변경 사항

(A) CUDA 커널 & 런처

✅ src/regemm_backward.cu

• 역전파 메인 엔트리: regemm::gemm_bias_act_bwd_f32(const GemmBiasActBwdParams&, cudaStream_t)
  • 1. bwd_epilogue_kernel에서 gZ = gY * act'(Z) 계산
    • 옵션: gC = β*gZ를 동시에 기록 (FUSE_GC)
    • 옵션: gBias를 원자 누적(Scalar/PerM/PerN)
    • 주의: 원자 누적이라 실행 전 gBias 버퍼를 0으로 memset
  • 1. cuBLAS로 gA/gB 계산
    • 행메이저 맞춤 래퍼 sgemm_rm() 사용(내부 col-major → op 바꿔 호출)

이 파일이 “수학적 정의”에 가장 가깝고, 성능·정확도의 중심입니다.

 

(B) CUDA 백엔드: 파라미터 어댑터

✅ backends/cuda/ops/gemm/backward.cu

• 상위 프레임워크 텐서(ai::Tensor) → regemm::GemmBiasActBwdParams로 매핑
• 체크 내용:
  • 디바이스/레이아웃/타입(f32 row-major CUDA) 제약
  • 치수 일치(A[M,K], B[K,N], gY[M,N], Z[M,N])
  • LD(leading dim) 추론: row-major이면 stride[0] 또는 cols
  • bias_kind 추론: gBias 텐서가 있으면 **크기(1/M/N)**로 매핑
  • alpha=1.0, beta=(C&&gC)?1.0:0.0 (현 스키마에 scale 노출 안 되어 기본값 유지)
• 마지막에 regemm::gemm_bias_act_bwd_f32() 호출

이 파일은 런타임 텐서 → 커널 파라미터 변환에만 집중합니다.

 

(C) 오퍼레이터 디스패치(호출 구심점)

✅ src/ops/gemm_backward.cpp (간단 디스패치 엔트리)

• ai::ops::gemm_bwd_run(...) 정의
  • 복잡한 레지스트리 경유 없이 직접 CUDA 백엔드 함수(위의 GemmCudaBackward) 호출
  • 에러코드만 일관되게 반환
• 목적: Python 바인딩이 이 엔트리를 간단히 부르도록 통일

Forward는 레지스트리 테이블을 쓰지만, Backward는 단일 구현이라 직접 경유가 단순합니다.

 

(E) Forward 디스패치 정리(참고)

✅ src/ops/gemm.cpp

• 검증 강화:
  • 널 포인터(-1), 전치 미지원(-11), 타입/레이아웃/디바이스 체크
  • Bias를 nullptr 정규화(effBias) → 레지스트리 키 일관성 유지
• OpRegistry로 등록된 CUDA Gemm 런처를 찾아 호출

Backward는 레지스트리로 안 돌리고 직접 호출로 간소화.

 

3) 호출 경로(요약 다이어그램)

 

Python
 ├─ gemm_bias_act_fwd_with_z(A,B,bias, act)  ---> regemm.gemm_bias_act_f32_ex (Z 저장)
 │             └─ returns (Y, Z)
 │
 └─ gemm_bias_act_bwd(A,B,gY,Z, act, bias_kind)
    └─ ai::ops::gemm_bwd_run(...)
         └─ backends/cuda/ops/gemm/backward.cu (GemmCudaBackward)
              └─ regemm.gemm_bias_act_bwd_f32(...)
                  ├─ bwd_epilogue_kernel: gZ(+gC,+gBias)
                  ├─ cuBLAS: gA, gB
                  └─ free temporaries

4) CMake 연결 포인트(요약)

• ai_backend_cuda 타겟에 backward.cu/regemm_backward.cu를 추가 컴파일
• ai_core/_core (pybind 모듈) 링크에 cuBLAS/cudart 포함
• 바인딩 파일에 의존하는 외부 등록 함수 ai_backend_cuda_register_all()는 기존 FWD 등록용(Backward는 직접 호출이라 등록 불필요)

---

5) 디버깅/정확도 체크 포인트

1. Z stash 일치성
   • gemm_bias_act_fwd_with_z로 받은 Y,Z가 반드시 Y == act(Z)인지 확인
   • 스크립트에서 np.max(|Y - act(Z)|)가 tol 이내인지 체크
2. gBias 초기화
   • 원자 연산(atomicAdd)로 누적 → 실행 전에 cudaMemsetAsync(0) 필수
   • 이 초기화가 빠지면 잡음이 껴서 오차 증가
3. row-major vs cuBLAS
   • cuBLAS는 col-major 기본 → sgemm_rm()로 op 순서 스왑하여 행메이저 수학과 동일하게 보정
   • 호출 시 M/N/K와 transA/B 정확히 매핑했는지 반드시 체크
4. LD(leading dim)
   • fwd/bwd 모두 row-major이면 ld = cols가 기본
   • 바인딩/백엔드에서 stride 기반으로 추론해 방어
5. Tolerance 설정
   • Forward: 5e-5
   • Backward(수치미분 비교): abs<=5e-3 or rel<=1e-2 (fp32 + 중심차분 오차 고려)

---

6) 테스트 시나리오

• 단일 케이스 빠른 확인
  1. (Y,Z)=gemm_bias_act_fwd_with_z(..)
  2. out=gemm_bias_act_bwd(A,B,gY,Z,..)
  3. 수치미분으로 gA/gB/gBias 비교
• 자동 스위트
  • python/test/test_autograd_gemm.py
    • 여러 act(none/relu/leaky_relu/gelu/sigmoid/tanh)
    • bias kinds(none/scalar/perm/pern)
    • 작은 행렬(8×7×5)로 수치미분도 빠르게
• 성능 벤치
  • python/test/bench_gemm.py
    • --mode end2end(H2D+D2H 포함) vs --mode kernel(커널 타임만)
    • GemmPlan을 통해 재업로드 없이 반복 실행 → 안정적 타이밍

---

7) 자주 맞닥뜨리는 이슈 & 해결

• “gBias가 이상하게 나옴”
  • 초기화(0) 안 했을 가능성 높음 → bwd 에필로그 앞에서 cudaMemsetAsync(gBias, 0, ...) 확인
• “오차가 케이스에 따라 튐”
  • 수치미분 eps 조정(기본 1e-3, 필요시 5e-4~2e-3로 탐색)
  • LeakyReLU leaky_slope 불일치 여부 확인
• “cuBLAS 호출에서 shape/ld 오류”
  • sgemm_rm() M/N/K, lda/ldb/ldc 재확인 (특히 row-major에서 ld=cols)

---

8) 요약

• regemm에 이미 있는 bwd 커널을 사용하고,
• CUDA 백엔드에서 파라미터를 안전하게 어댑트,
• 디스패치 엔트리(gemm_bwd_run)로 바인딩과 연결,
• Python에서 Z를 항상 stash하여 정확한 역전파를 구현,
• 수치미분 기반 테스트로 신뢰성 확인.