본문 바로가기

AI Compiler framework

Trainer 의 역할 정의 ( AICF Framework - Execution Orchestrator ) 학습을 위해서가 아닌 실행을 고정하고 성능을 관찰하기 위해 등장한 구조

Trainer 는 학습 step 을 고정된 실행 단위로 정의

그 실행을 eager / capture / replay 어느 모드에서도 동일한 의미로 수행하게 만드는 조정자

Trainer 는 커널, IR 도 직접 다루지 않음, 어떤 실행이 언제, 어떤 입력으로, 어떤 규칙 아래 반복되는가를 책임진다.

 

1. Trainer 가 보장해야 하는 핵심 계약 ( Non-Negotiable Rules)

학습 step 의 의미는 항상 동일하다

Trainer 가 정의하는 1 step 은 다음 순서를 반드시 포함

ZeroGrad
- Forward
- Loss
- Backward
- Optimizer Step
  • eager / capture / replay 모드에서도 의미적 차이가 없어야 함
  • capture 는 step 을 기록하는 방식, step 의 의미를 바꾸지 않음

 

Step은 결정적 deterministic 이어야 한다. ( 동일한 입력 메모리 주소 )

Trainer 는 다음 조건에서 완전히 동일하게 실행되도록 강제

  • 동일한 op sequence
  • 동일한 tensor shape / dtype / device
  • 동일한 optimizer 상태 전이
  • 동일한 입력 메모리 주소

이를 위해 Trainer 는 

  • shape 변화 시 재캡처 요구
  • capture 시점 이후에는 입력 포인터가 변하지 않도록 설계

 

Trainer 는 데이터 흐름과 실행 흐름을 분리한다

  • 데이터 흐름 : batch - input tensors
  • 실행 흐름 : step 실행 방식 ( eager vs replay ) 

capture 모드에서

  • 데이터는 매 step 마다 바뀐다.
  • 실행 흘므은 고정된 CUDA Graph 를 재사용한다.

Trainer 는 이 분리를 유지하기 위해 고정된 입력 버퍼 + copy 갱신 규칙을 강제한다.

 

2. Trainer 의 책임 범위 ( What Trainer Owns )

2.1 Execution Orchestration

실행의 언제 / 어떻게를 결정한다

  • warpup 단계
  • capture 시점
  • replay 반복
  • 재캡처 조건

 

3. Trainer 가 책임지지 않는 것 ( Intentional Non-Responsibilities)

  • 커널 선택 로직
  • IR 생성 / lowering
  • CUDA stream 관리 상세
  • 개별 op 구현

Trainer 는 실행 계획을 소비하는 쪽

 

Capture 는 최적화가 아닌 고정화, 한 번 정의된 step 을 다시 정의하지 않기 위해서

  • 실행 경로를 고정
  • 커널 선택을 고정
  • 메모리 주소를 고정

CUDA Graph 는 캡처 시점의 포인터 주소를 박제, Trainer 는 이를 규칙으로 강제한다.

 

4. Trainer 내부 구조 요약

Trainer.fit
 ├─ Warmup (eager semantics)
 ├─ Capture (fixed buffers)
 │   └─ train_step_eager(cap_batch)
 └─ Replay Loop
     ├─ bind_batch (copy_)
     ├─ replay()
     └─ read loss tensor

 

커널 성능 개선을 확인하려면 바뀐 것이 커널 뿐이어야 함

일반적인 eager 실행에서는 

  • 커널 선택의 변경 가능성
  • 메모리 주소 변경
  • 실행 경로 변경
  • stream 타이밍 문제

 

고정된 실행 환경의 필요성

 

뿐만 아니라!!

Graph Capture 와 Trainer 는 고정된 학습 실행 환경을 만드는 수단일 뿐만 아니라 그 환경을 기준점으로 삼아 일부만 바꾼 실행을반복 실험하기 위하 장치로도 해석 가능