정적, 동적 학습 루프로 분리한, sequential 및 graph_executor 관련 파일들의 구조 / 관계 / 실행 순서

• layers/sequential.py 모델 컨테이너 + 오케스트레이터
  • 빌드/요약/파라미터 유틸 ( build() / parameters() / zero_grad() 등)
  • 정적 경로 : compile() - (고정 입력/모델 전체) 캡처 준비 - one_step() 로 재생
  • 동적 경로 : one_step_dynamic() - 분기/반복/EaryExit 경로 평탄화 - 경로별 키 생성, 캐시 조회 - 미스 시 캡처 후 풀에 저장 - 이후 재생
  • 풀 : graph_executor.py 의 graph_pool + 내부 풀백
• graph/captrue_plan.py 캡처 사전계획 생성기
  • 레이어 타입별 플래너 registry 로 forward 출력, pre-act, backward, workspace 를 사전할당 - 캡처 구간의 동적 malloc 제거
  • API : make_plan_for_sequential(model, in_shape) / make_plan_for_path(path_layers, in_shape)
  • 드롭아웃 마스크 롤링 : advance_dropout(plan, seed_bump)
• graph/graph_exec.py 한 스텝 ( fwd - loss - bwd - opt) 캡처/재생
  • 러너 : _run_fwd(layers, plan, X, stream_ptr), _run_bwd(layers, plan, g_in, stream_ptr)
  • 캡처 : record_step_graph() - 워밍업 1회 - cuda.graph.capture_stream - instantiate() 또는 폴백 실행자
  • 실행 핸들 : TrainGraph(gexec, io, stream) - set_batch() / launch()
• graph/graph_executor.py 키/폴만 유지하는 전역 캐시
  • GraphSignature(shape, dtype, layout) + GraphKey(signature, branch_id, variant)
  • MultiGraphPool ( LRU, 스레드 세이프 ) & 전역 graph_pool 인스턴스
  • sequential 이 동적 경로 그래프를 키로 조회/저장해 재사용

[Sequential]  ──uses──>  [capture_plan] ──produces──> CapturePlan(per_layer, loss.dY, ...)
     │                         │
     │                         └─(buffers/workspaces for each layer)
     │
     ├─(static/dynamic)──calls──> [graph_exec.record_step_graph]
     │                                  │
     │                                  └─captures fwd→loss→bwd→opt as CUDA Graph
     │                                               │
     │                                               └─returns TrainGraph(gexec, io, stream)
     │
     └─(dynamic only)──uses──> [graph_executor.graph_pool]
                               └─keyed by GraphKey(signature+branch+variant)

 

실행 순서: 정적 캡처/재생

전제 : 모델 구조/입력 shape 가 고정, 모든 레이어가 capture-ready

1. Sequential.compile(input_shape, loss, optimier, ... )
   • (필요 시) build()
   • CapturePlan 생성 : make_plan_for_sequential(model, in_shape) - 레이어별 PerLayerBufs(y, z, gA, gW, gB, work) 사전 할당, loss.dY 준비
   • 옵티마이저 바인딩(필요 시) : try_rebind_grads(self, optimizer, plan)
   • 그래프 캡처 record_step_graph(self, loss, optimizer.step_into, plan, X_buf, y_buf, stream, loss_out)
     • 워밍업 1회 (fwd - loss - bwd - opt) 로 버퍼/WS 고정
     • capture_stream 내 동일 시쿠너스 캡처 - instantiate()
   • TrainGraph 생성/보관 : self._tg = TrainGraph(gexec, io, stream)
2. Sequential.one_step(X, y)
   • self._tg.set_batch(X, y)
   • self._tg.launch() ( CUDA Graph 재생 )
   • loss_out.get() 반환

핵심 포인트

• Plan 단계에서 모든 메모리가 준비돼 캡처 내 malloc 금지
• _run_fwd/_run_bwd 가 model.layers 기준으로 순회

 

 

주요 데이터 구조 & 불변식

• CapturePlan
  • per_layer[i] <-> layers[i] 1:1 정렬 ( 정적은 model.layers, 동적은 path_layers )
  • PerLayerBufs
    • y : 다음 레이어 입력, z : pre_activation
    • gA/gW/gB : backward 버퍼
    • work : Conv/GEMM/RNN 등 WS 묶음
  • loss.dY : CE 등 일부 손실의 dY 고정 버퍼
• TrainGraph
  • io : 고정 X, y, logits 버퍼 핸들
  • set_batch() 는 그래프와 같은 스트림에서 H2D/D2D 복사
  • launch() 는 내부 gexec.launch(stream_ptr) 호출
• GraphKey
  • signature + batch_id + variant
  • 키가 같으면 가은 경로 모드 입력서명을 가리키므로 재사용 가능