본문 바로가기

dev_AI_framework

단순 복붙, 각 부분 수정 필요함, 단 이는 최적화 단계에서, 일단 구현 먼저 다 하자 으쌰으쌰

0) 큰 그림 읽기

  • 커널 타임 59.4% → GPU가 대부분 커널 실행에 시간을 씀. 나머지 ~40%는 런치 오버헤드/메모리 이동/빈 공간.
  • GEMM 계열(총 ~35%+):
    • gemm_bias_act_f32_smoke_ex 13.3% (자체 fused GEMM+bias+act)
    • ampere_sgemm_128x32_nt 12.0% (cublas/lt 내부 커널로 보임)
    • gemmSN_TN_kernel 9.9% (TN/NT 레이아웃 교차)
  • 옵티마이저: adamw_fused_fp32 6.9% (멀티-텐서 fused로 보이지만 비중 큼)
  • 에필로그/활성/전치/덧셈: bwd_epilogue_kernel 6.1%, k_transpose_MN 5.0%, k_add_transpose_into 4.8%, k_apply_act_rows_local 4.2% 등
  • RNG: generate_seed_pseudo 6.0% (꽤 큼)
  • 메모리/유틸: cupy_copy_float32_float32 3.6%, cupy_power_float32_int32_float 2.8%

핵심은 (A) GEMM/에필로그/전치 재배치로 큰 덩어리 최적화, (B) RNG/옵티마이저/elementwise를 더 세게 fuse, (C) 불필요한 D2D copy 제거입니다.


1) GEMM/전치/에필로그: “레이아웃 계약(Layout Contract)” 고정

전치/덧셈 관련 커널이 합쳐서 10%+를 잡아먹고 있어요. 이건 GEMM의 opA/opB/출력 레이아웃을 바꾸면 통째로 사라질 수 있는 비용입니다.

액션

  1. 전치 없애기
    • 가능하면 A op = N, B op = T 같은 cublasLt op 플래그로 해결하고, 명시적 transpose 커널 금지.
    • 당신의 fused GEMM(gemm_bias_act_f32_smoke_ex) 에필로그에서 “전치된 형태로 저장” 옵션을 추가하세요. 그러면 후속 k_transpose_MN/k_add_transpose_into가 사라집니다.
  2. Bias/Act/Dropout 더 강한 에필로그
    • 현재 bwd_epilogue_kernel 6.1%, k_apply_act_rows_local 4.2%가 분리 실행 중.
    • fwd: GEMM + (bias+act+dropout), **bwd: dGEMM + (act’ + bias_grad 축약)**로 각 패스 1커널을 목표.
  3. cublasLt 튜닝(알고리즘/워크스페이스)
    • Ampere에서 TF32 허용(정밀 OK면): cublasSetMathMode(CUBLAS_TF32_TENSOR_OP_MATH) 혹은 Lt의 computeType = CUBLAS_COMPUTE_32F_FAST_TF32.
    • 워크스페이스 크게 잡고 cublasLtMatmulAlgoGetHeuristic로 베스트 algo 선택.
    • 배치/shape가 반복된다면 알고리즘 캐시(shape→algo)로 재사용.

체크리스트

  • 명시적 transpose 커널 호출 0개 만들기
  • opA/opB로 해소 불가 시, 에필로그가 전치된 레이아웃으로 write
  • 활성/드롭아웃 에필로그 일체화
  • TF32/Algo/Workspace 튜닝 및 캐시

2) RNG 6.0%: “카운터 기반(stateless) + 1-pass”

generate_seed_pseudo가 6%는 과합니다. 매 step/레이어마다 별도 RNG 커널을 쏘면 런치도 늘고 L2/DRAM 왕복도 증가합니다.

액션

  • 카운터 기반 Philox(stateless)로 드롭아웃/노이즈 커널 내부에서 즉시 난수 생성:
    • seed = global_base_seed
    • counter = (global_step << 32) ^ layer_id ^ tensor_offset
    • → 별도 generate_seed_* 커널 삭제 가능.
  • CUDA Graph 재생에서도 (global_step, replay_idx) 를 카운터에 반영하면 재현성 확보.
  • 가능하면 드롭아웃도 에필로그에 융합(앞 1)과 결합).

3) AdamW 6.9%: 멀티-텐서/벡터화/스케줄 융합

이미 “fused”지만 여전히 비중이 큼. 개선 여지 있습니다.

액션

  • 벡터화 로드/스토어(float4/float2), 128B 정렬 보장.
  • 모든 파라미터를 큰 “multi-tensor list”로 묶어 grid를 크게 유지(작은 텐서 다수 → occupancy/런치 오버헤드↑).
  • LR/WD 스케줄이 cupy_power_float32_int32_float 같이 따로 돌고 있다면 옵티마이저 커널 안에서 적용(별도 elementwise 제거).
  • 혼합정밀 사용 시: m/v는 FP32 유지, grad는 FP16/FP32 입력 모두 처리.

4) Elementwise/유틸: copy + power 줄이기

  • cupy_copy_float32_float32 3.6%: D2D copy 감소
    • astype('float32', copy=True) 남용, 혹은 불필요한 ascontiguousarray가 원인일 수 있어요.
    • dtype/stride 규약을 잡고 래퍼에서 강제 변환을 없애세요(또는 첫 입구에서 1회 정규화).
  • cupy_power_float32_int32_float 2.8%:
    • LR decay나 normalize류라면 옵티마이저/에필로그로 융합하거나, 최소한 여러 power/op를 1커널로 합치기.

5) 런치/동시성: 스트림/그래프 활용 점검

59.4%가 커널이면 나머지 40.6%에 런치 오버헤드/빈틈이 섞여 있을 확률이 큽니다.

액션

  • 이미 CUDA Graph를 쓰고 있으니, 그래프 내에서 커널 수를 줄이는 방향(상기 fuse) 이 1순위.
  • 스트림 배치: 독립 파이프(예: data transform vs GEMM)를 서로 다른 스트림으로 보내 동시성 확보. 단, 암묵 sync(디폴트 스트림/이벤트 누락) 주의.
  • 메모리 풀/워크스페이스를 미리 할당(그래프 외부) → 캡처 안전.
  • NVTX로 forward/loss/backward/optimizer 레인지 명확화 후 nsys “CUDA Kernel Statistics”에서 range별 커널 수/시간 비교.

6) 빠른 검증 플랜 (실험 순서)

  1. 전치 제거: opA/opB/에필로그 write-layout 수정 → k_transpose_*/k_add_transpose_into 사라지는지 확인
  2. RNG stateless화: generate_seed_pseudo 제거 → 드롭아웃 커널 내부난수
  3. AdamW 튜닝: 멀티-텐서 일괄/벡터화 + 스케줄 융합 → 커널 1~2개로 수렴
  4. Elementwise 합치기: power/copy 감소
  5. cublasLt 알고리즘/TF32/WS 튜닝: gemm_* 총합 시간 최소화

각 스텝마다 nsys stats --report gpukernsum,summary run_trace.nsys-rep(또는 새 리포트)에 비중 변화를 체크하세요.


7) (원하면 바로 써먹을) 구현 힌트

cublasLt 레이아웃 + 에필로그 전치 write

  • 출력 텐서의 strides를 “전치된 형태”로 설정하거나, fused GEMM 에필로그에서 (m,n) 좌표계로 쓰되 ld/stride를 바꿔 “전치 결과가 되도록” 저장.
  • 후속 연산이 N-major을 기대한다면 처음부터 N-major로 계약.

Philox stateless 드롭아웃 스케치

  • key = hash(global_seed, layer_id)
  • counter = global_step * N + element_linear_idx
  • 1개 커널에서 uniform(counter,key) → mask 생성 → 곧바로 적용/저장. 별도 seed 커널 없음.

AdamW 멀티-텐서

  • 파라미터 포인터/크기 테이블을 큰 배열로 만들어 한 런치에 소화.
  • 내부에서 각 텐서를 루프 돌리되, 워프가 같은 크기 버킷을 처리하도록 bucketing by size.