CUDA Graph 학습 경량 가이드

전체 흐름 (요약)

1. capture_plan.make_plan_for_sequential(model, shape) → 레이어별 버퍼/WS 준비
2. record_step_graph(..., X_buf, y_buf) → fwd→loss→bwd→opt 그래프 캡처
3. TrainGraph.set_batch(X, y); TrainGraph.launch() → 한 스텝 실행
4. 고수준엔 CudaGraphTrainer.compile/one_step만 쓰면 됨

---

모듈별 핵심만

1) graph/capture_plan.py

• 역할: 레이어별 고정 버퍼/워크스페이스 사전할당
• 주요 함수
  • make_plan_for_sequential(model, input_shape, *, loss_kind="softmax_ce", lt_bytes=8<<20) -> CapturePlan
• 생성물
  • CapturePlan.per_layer[i]:
    • y(출력), z(pre-act, 선택), gA/gW/gB(역전파), work(gemm ws 등)
  • CapturePlan.loss.dY(softmax CE일 때)
• 주의: 각 레이어는 compute_output_shape 구현 필요

2) graph/graph_exec.py

• 역할: 한 스텝(fwd→loss→bwd→opt) CUDA Graph 캡처 & 실행
• 주요 함수
  • record_step_graph(model, loss_fn, optimizer_step_fn, plan, *, X_buf, y_buf, stream=None) -> GraphExec | GraphExecLike
    • X_buf, y_buf는 고정 버퍼 (그래프 내부에서 참조)
• 주요 클래스
  • TrainGraph(gexec, io, stream)
    • set_batch(X, y) : 고정 버퍼에 복사
    • launch() : 캡처된 그래프 실행
    • 속성: X_buf, y_buf, logits
• 주의: 캡처/워밍업 모두 X_buf에서 시작 (입력 K mismatch 방지)

3) train/cuda_graph_trainer.py

• 역할: 고수준 트레이너 (한 줄로 캡처/실행)
• API
  • CudaGraphTrainer(model, loss_fn, optimizer)
  • compile((B, D)) : build→plan→grad rebind→record
  • one_step(X, y) -> float : 한 스텝 실행 후 loss 반환
  • tg : TrainGraph 핸들(필요 시 직접 접근)

4) optim/rebind.py

• 역할: plan의 gW/gB 버퍼를 옵티마이저 grad 포인터로 재바인딩
• API
  • try_rebind_grads(model, optimizer, plan)
  • (내부) collect_params_from_plan(...)
• 주의: 옵티마이저가 rebind_grads 구현 시 자동 적용

5) losses/utils.py

• 역할: dY 스케일 감지 → 권장 grad_scale 반환
• API
  • infer_grad_scale(loss_fn, model, X, y) -> (scale: float, "sum"|"mean")

6) utils/streams.py

• 역할: 간단 타이머
• API
  • with Timer(stream) as t: ...; t.ms()

 

import cupy as cp
from graph_executor_v2.layers.sequential import Sequential
from graph_executor_v2.layers.dense_gemm import Dense
from graph_executor_v2.losses.softmax_ce import SoftmaxCrossEntropy
from graph_executor_v2.optim.adamw import AdamWOpt
from graph_executor_v2.train.cuda_graph_trainer import CudaGraphTrainer

# 1) 모델
def make_model(B=64, D=128, H=256, C=11):
    m = Sequential(
        Dense(H, activation="relu", initializer="he", use_native_bwd=True),
        Dense(C, activation="none", initializer="xavier", use_native_bwd=True)
    )
    m.build((B, D)); m.train(True); return m

B,D,H,C = 64,128,256,11
X = cp.random.randn(B, D).astype(cp.float32)
y = cp.random.randint(0, C, size=(B,), dtype=cp.int32)

model = make_model(B,D,H,C)
loss  = SoftmaxCrossEntropy()
opt   = AdamWOpt([], lr=1e-3, wd=1e-4)      # 파라미터 rebind는 내부에서 처리
if hasattr(opt,"ensure_initialized"): opt.ensure_initialized()

# 2) 캡처
trainer = CudaGraphTrainer(model, loss, opt)
trainer.compile((B, D))

# 3) 실행
L0, _ = loss.forward(model(X), y)
trainer.one_step(X, y)
L1, _ = loss.forward(model(X), y)
print(L0, L1)

 

언제 무엇을 손대나?

• 새 레이어 추가 → compute_output_shape, forward_into, backward_into 구현 +
  capture_plan.make_plan_for_sequential에 gW/gB/WS 할당 로직만 추가
• 새 손실 추가 → loss.forward(logits, y, return_scalar=False)에서 (device_loss, dY) 반환 보장
• 새로운 옵티마이저 → rebind_grads(triplets)만 구현하면 그 즉시 그래프 경로에 편입