6.2 성공 케이스 문서화 : Warmup (no drift) + CUDA Graph Replay == IRExecutor

목표

• Warmup 단계에서 state drift (파라미터 / step / m / v ) 없이 ㅋ커널 / 버퍼 / grad buffer 를 materialize
• CUDA Graph capture / replay 가 실제 학습 업데이트를 포함
• IRExecutor 가 replay 와 동일한 state 변화를 만들고 결정성이 유지

 

핵심 설계 규칙

Capture-safe autograd: leaf grad 포인터 안정성

• capture 중에는 leaf.grad 를 새로 allocate 하면 안 됨
• leaf grad 는 warmup 에서 한 번 materialize 해두고, 이후엔 in-place overwrite 로만 갱싱

요점

• autograd.backward 기본
• leaf grad write 는 copy 로 pointer-stable 유지

 

Tracing 시 backward 는 실행이 아니라 IR 연결

• tracing 동안 실제 autodiff 계산 금지
• Backward IR node 만 emit 하고 forward 에서 만들어진 IRValue 와 동일한 vid 를 재사용해야 링크가 끊기지 않음

요점

• as_ir_value_obj() / cache 기반으로 같은 객체 - 같은 IRValue

 

Lowering 에서 leaf param.grad vid 에 절대 쓰지 말 것

제일 오래 걸린 지점

• IRExecutor 는 lowered 리스트를 그대로 실행
• backward GEMM output 을 leaf grad vid 에 걸어버리면
  • IRExe-only 에서 param.grad 버퍼가 forward / optim 과 alias 나면서 깨짐
  • capture replay 와 달라짐

해결

• backward 출력은 무조건 dedicated grad pool 로만 write
• AdamStep 의 g_in 도 그 grad vid 를 입력으로 사용

 

Warmup 은 커널 / 버퍼 예열임, 학습이 아니다

warmup 단계는 다음만

• leaf grad buffer materialize
• lazy init 방출
• activation buffer 풀 생성

하면 안되는 것

• step 증가
• adam 업데이트

해결법

• warmup 동안만 AICF_WARMUP = 1
• capture 구간에서는 강제 AICF_WARMUP = 0
• 그래서 warmup 은 no-drift, capture 는 업데이트 포함

 

IRExecutor 가 실행 가능한 env 를 완전 바인딩해야 함

lowering 이 grad pool / relu_y_saved 같은 temporaries 를 만들면,

IRExecutor 는 그 vid 에 대응하는 텐서를 반드시 가져야 함

해결

• autobind_env_from_lowered 로
  • lowered 에 등장한 모든 vid 에 대해 torch.empty 를 잡아줌
  • 이건 capture 밖에서만 수행