PyTorch 2.14의 node_creation_hook, 역전파 노드가 만들어지는 순간 기록을 붙이는 방법 | DAKER 커뮤니티

역전파를 살필 때 가장 답답한 지점은, 메모리가 어디서 생겼는지를 나중에 거슬러 올라가며 추정해야 한다는 점입니다. 특히 작은 실험을 재현하거나 참가자 노트를 남길 때는, 결과보다 기록이 언제 붙었는지가 더 중요할 때가 많습니다.

PyTorch 2.14의 torch.autograd.graph.node_creation_hook은 이런 지점을 직접 겨냥합니다. 역전파 그래프의 각 노드가 만들어질 때 바로 실행되기 때문에, 그래프를 다시 훑지 않고도 그 순간 메타데이터를 붙이거나 훅을 등록할 수 있습니다.

핵심은 역전파가 끝난 뒤 추적하는 것이 아니라, 노드가 생기는 시점에 바로 기록을 남길 수 있다는 점입니다.

node_creation_hook이 바꾸는 관찰 시점

예전에는 역전파가 끝난 뒤에야 이 메모리가 어느 앞 연산에서 왔는지를 추측해야 했습니다. 반면 torch.autograd.graph.node_creation_hook을 쓰면, 노드가 생성되는 순간 연산 이름이나 출력 shape 같은 정보를 바로 붙일 수 있습니다. 이 방식의 동기는 역전파 메모리를 그 메모리를 만든 앞 구간에 연결하는 일입니다.

작은 모델에서 훅을 등록한 뒤, 노드가 생길 때마다 연산 이름과 출력 shape를 리스트에 쌓아 보면 됩니다. 참가자 노트에는 훅이 호출된 횟수, 남긴 필드, 훅을 끈 뒤 리스트가 비는지를 적어 두는 것이 좋습니다. 메모리가 줄었는지보다, 노드가 생길 때 기록이 실제로 붙었는지가 재현에 더 도움이 됩니다.

재현해 볼 실험 조건

실험은 단순할수록 좋습니다. 작은 선형층 한 번과 loss.backward()만 있는 경로에서 torch.autograd.graph.node_creation_hook을 등록합니다. 훅 안에서는 노드 타입 또는 이름과 출력 텐서 정보를 리스트에 남기고, backward가 끝난 뒤 리스트 길이가 0이 아닌지 확인하면 됩니다.

이후 같은 코드를 훅 없이 한 번 더 돌려, 기록이 훅 등록에만 의존하는지 비교할 수 있습니다. 필요하면 커스텀 autograd.Function 한 개에도 같은 훅이 도는지 보면 됩니다.

이번 확인 포인트는 메모리 최적화 자체가 아니라, 노드가 생길 때 훅이 실제로 호출되는지입니다.

이번 글에서 다루지 않는 범위

같은 릴리스의 ctx.set_output_grad_dtype나 cdist/pdist 이중 역전파는 다른 주제입니다. 이번 글에서는 노드가 생길 때 훅이 도는지에만 초점을 둡니다.

실험이 끝나면 사용한 PyTorch 버전, 모델 한 줄, 훅 등록 코드를 함께 남겨 두면 다음 참가자가 같은 조건을 바로 따라오기 좋습니다.

실무에서 바로 볼 포인트

이 훅은 그래프가 만들어진 뒤가 아니라 노드가 생기는 순간에 붙습니다. 그래서 등록 시점은 forward 전에 두는 것이 좋습니다. 또 학습 노트에는 훅 호출 횟수와 남긴 필드 이름을 적어 두면 비교가 쉬워집니다.

메모리 귀속, 즉 어느 앞 구간이 역전파 메모리를 만들었는지를 보려면 훅에서 그 구간 식별자를 같이 저장하면 됩니다.

참고 자료

https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/public/learning/materials/pytorch-advanced-2-embedding-self-attention-mini-i
https://daker.ai/community

직접 실험해 보셨다면, 어떤 필드를 남겼을 때 재현에 가장 도움이 됐는지 궁금합니다.