테스트 코드를 통한 파일 실행 순서 확인

1. 테스트 코드 레벨 

1.1 import 단계

from aicf_fw.nn import Sequential, Linear, ReLU
from aicf_fw.optim import Adam

실행 파일

• sequential
• linear
• relu
• 공통 베이스
  • module

이 시점에

• 모델 구조만 생성
• IR / CUDA / Graph 아무것도 안 돌아감

 

1.2 모델 생성

model = Sequential(
    Linear(...),
    ReLU(),
    Linear(...),
).to(device)

실행 흐름

• Sequential.__init__
  • add_module("0", Linear)
  • add_module("1", ReLU)
  • add_module("2", Linear)
• Modele.to()
  • 모든 parameter tensor device / dtype 이동

 

1.3 Optimizer 생성

opt = Adam(model, lr=1e-3)

실행 파일

• adam.py
• 상속 : base.py

내부 상태

• m, v, step, bc1_inv, bc2_inv 모두 torch.Tensor
• CUDA Graph 캡처 가능한 상태로 준비됨

 

2. model.compile(...) 호출

model.compile(
    optimizer=opt,
    warmup_inputs={"x": x, "t": t},
)

2-1 fw/module.py

실행 함수

Module.compile()

하는 일

• 내부에서 호출

from aicf_fw.fw.compile import compile_train_step

 

2-2 fw/compile.py

실행 함수

compile_train_step(...)

warmup_inputs 기반 자동 추론

_infer_from_warmup_inputs(...)

체크 내용

• x, t 존재 여부
• rank = 2
• shape 동일
• dtype / device 일치
• 모델 파라미터 전체와 device / dtype cross-check

이 단계에서

• B, D, device, dtype 확정

 

IR Trace 준비

def build():
    ...

 

IR 생성

ir = trace_ir(build, name=name)

실행 파일

• trace
• ir

내부에서 실제로 일어나는 것

• build 실행
• 하지만
  • 실제 tensor 아님
  • 전부 sym_tensor

결과물

• IRGraph
• 노드들
  • Linear
  • ReLU
  • Save
  • LinearBwd
  • ReluBwd
  • MseGrad
  • AdamStep

여기까지는 순수 Python IR 생성

 

Lowering

ir = trace_ir(build, name=name)

실행 파일

• lower.py

여기서 일어나는 것

 

• Linear → gemm + bias_add
• LinearBwd → gemm + gemm + reduce_sum
• AdamStep → adam_step
• kernel_id 부착 (*_v0)

 

결과

List[dict]  # backend primitive ops

 

Binding Plan 생성

plan = build_binding_plan(ir)

실행 파일

• plan.py

역할

• IR value id - runtime tensor 매핑
• params / optimizer state / inputs / outputs 연결

 

Executor 생성

ex = PlannedExecutor(...)

실행 파일

• aicf_fw/core_v2/exec.py
• 내부에서:
  • CUDA kernel dispatch table 준비
  • kernel_id → launcher 매핑

 

CompiledTrainStep 생성

compiled = CompiledTrainStep(...)

실행 파일

• aicf_fw/fw/train_step.py

이 객체가 들고 있는 것

• IR
• lowered ops
• binding plan
• executor
• param tensors
• optimizer state tensors
• CUDA graph 핸들 (아직 없음)

 

compile-time warmup

compiled.warmup(...)

실행 흐름

• CompiledTrainStep.train_step() 반복 호출
• eager 실행
• CUDA Graph 캡처 ❌
• 메모리 안정화 목적

 

⛔ compile 종료

이 시점에:

• 그래프 구조 확정
• 커널 ID 확정
• 메모리 레이아웃 확정

 

3. model.train_step(...)

model.train_step({"x": x, "t": t})

실행 경로

• fw/module.py → Module.train_step
• CompiledTrainStep.train_step

실행 모드

• eager path
• CUDA Graph ❌
• optimizer meta 업데이트 포함

📌 테스트에서 [train_step] |ΔW0| > 0 확인

 

4. model.capture(...)

model.capture({"x": x, "t": t})

실행 경로

• CompiledTrainStep.capture

실행 파일

• aicf_fw/fw/train_step.py
• 내부에서:
  • cudaGraphBegin
  • executor 실행
  • cudaGraphEnd

📌 로그:

[aicf] graph_begin (dedicated stream)

➡️ CUDA Graph 생성 완료

 

5. model.replay(n=3)

model.replay(n=3)

실행 경로

• CompiledTrainStep.replay

내부

• cudaGraphLaunch × 3
• host 코드 거의 없음
• optimizer state + params in-place 업데이트

📌 replay에서도 W 변화 확인됨

 

6. Meta mutation 검증

bc1.fill_(1.0)
bc2.fill_(1.0)
model.replay(1)

의미

• CUDA Graph 외부에서 tensor 값 변경
• Graph는 값만 읽음
• → replay 결과 변화

📌 그래프가 “값에 열려 있음”을 증명

 

7. model.reset()

model.reset()

실행

• CUDA Graph 해제
• executor 상태 정리

 

전체 실행 흐름 요약

python_framework_test.py
 → fw/module.py (compile)
   → fw/compile.py
     → core_v2/trace.py
     → core_v2/lower.py
     → core_v2/plan.py
     → core_v2/exec.py
     → fw/train_step.py
 → train_step (eager)
 → capture (cudaGraph)
 → replay (cudaGraphLaunch)
 → reset