본문 바로가기

AI Compiler framework

accumulata = False / True, 먼뜻인지 True 환경에서 deterministic 하게

역전파가 leaf.grad 에 어떻게 쓰느냐의 차이, 메모리 / 상태 모델을 결정

 

accumulate = False

이번 backward 에서 나온 grad 를 이번 스텝용으로 새로 만든 값으로 취급, leaf grad 는 덮어쓰기 모델

  • leaf 에 grad 버퍼가 있으면 copy / set / write 로 채움
  • 또는 backward 결과를 바로 optimizer 로 넘기고 grad 저장을 최소화

특징

  • 결정적이고 단순
  • 한 번의 루프에서 가장 일반적인 모드
  • capture-safe 만들기 쉬움

언제

  • 일반적인 단일 loss 학습 루프
  • microbatch / gradient accumulation 안 할 때
  • 최소 프레임워크 구현

 

accumulate = True

이번 backward 에서 나온 grad 를 기존 leaf.grad 에 더한다 즉, leaf grad 는 누적 모델

  • leaf.grad 가 있으면 add
  • 없으면 새로 생성

ㅌ특징

  • gradient accumulation / microbatch / multi-loss 로 바로 이어짐
  • 하지만 capture-safe 관점에선 까다로움
    • 매 replay 마다 leaf.grad 상태가 이전 replay 의 잔재를 먹기 쉬움
    • 몇 번 누적했는지가 상태로 들어가면 deterministic 계약이 흔들림
    • leaf.grad 를 동적으로 만들거나 분기하면 capture 중 위험

언제

  • microbatch : 여러 번 forwad / backward 해서 grad 를 모은 뒤 한 번 step
  • multi-loss : loss1.backward(), loss2.backward() 처럼 합산
  • 메모리 제한으로 큰 batch 를 쪼개는 경우

 

같은 메모리 주소에 같은 커널 시퀀스를 반복하는 것이 핵심인데 

  • accumulate = True 는 grad 가 반드시 초기화 되어 있어야 이번 스텝 grad 가 됨

 

직접 in-place 누적이 아닌, pool 이 고나리하는 grad_accum 버퍼를 둬서 사용