nisbenz의 TensorLib: 순수 C로 구현한 CPU 전용 텐서와 역전파 자동미분 라이브러리

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

TensorLib은 순수 C로 작성된 CPU 전용 텐서·자동미분 라이브러리로서 NumPy 스타일의 브로드캐스팅과 제로 카피 스트라이드 뷰를 제공하고 동적 reverse-mode autograd로 그래디언트를 계산한다. 핵심 구현은 Linear·Embedding·LayerNorm·causal multi-head self-attention·MLP로 구성된 프리-노름 GPT 스타일 디코더 스택과 SGD·AdamW 옵티마이저, 블록/타일 방식의 AVX2+FMA hand-rolled matmul을 포함하며 저장소 버전 관리와 원자적 체크포인팅으로 재현성을 확보한다. 약 6,100줄 코드와 33개 유닛 테스트, 1.9M 파라미터 tiny_lm 학습 사례는 CPU 환경에서의 실용성과 성능 최적화에 초점을 맞춘 설계 선택을 보여주며 외부 ML 의존성을 배제한 경량 구현의 가능성을 제시한다.

섹션별 상세

TensorLib의 텐서 코어는 N차원 float32 배열을 NumPy 스타일 브로드캐스팅 규칙으로 처리하고 reshape·transpose·slice·expand 연산이 모두 동일한 저장소를 공유하는 제로 카피 스트라이드 뷰로 동작한다. 참조 카운팅이 메모리 관리를 담당하며 저장소 버전 카운터로 stale-graph 버그를 탐지해 일관성을 확보한다. 이러한 설계는 CPU 환경에서 메모리 복사 비용을 줄이고 연산 파이프라인의 효율을 높여 대규모 중간 버퍼 없이도 다차원 연산을 수행할 수 있게 한다.

자동미분 엔진은 순전파 동안 연산을 즉시 기록하는 동적 역전파 방식(define-by-run)을 채택해 런타임 제어 흐름을 자연스럽게 지원한다. 연산이 실행될 때 역전파 정보가 구성되며 23개의 미분 가능 연산을 기본으로 제공하고 저장소 버전 카운터로 그래디언트 무결성을 검증한다. 이 구조는 사용자 정의 연산과 복잡한 그래프 변형을 처리하면서도 역전파 시 정확한 기울기 누적을 보장한다.

신경망 모듈 집합은 Linear, Embedding, LayerNorm, Dropout, multi-head causal self-attention, MLP와 GPT 스타일 프리-노름 디코더 스택으로 구성되며 간단한 C 스타일 OOP 모듈 시스템으로 조합된다. 모듈 호출은 함수 포인터 디스패치와 부모/자식 모듈 트리를 통해 입력을 받고 파라미터와 출력을 연결하는 방식으로 동작한다. 이 구조는 모델 구성 요소를 직접 연결해 CPU 상에서 완전한 디코더 기반 트랜스포머를 학습할 수 있게 한다.

성능 최적화는 핸드롤 SIMD matmul 커널과 체크포인팅, 광범위한 단위 테스트로 뒷받침된다. matmul은 블록/타일 기반의 AVX2+FMA 마이크로커널과 RHS 패킹으로 캐시 지역성을 높여 OpenBLAS와 비교하는 벤치마크를 수행했으며 저장소 버전 관리와 원자적 바이너리 저장 방식으로 옵티마이저 상태와 RNG 상태를 보존한다. 약 6,100줄의 C 코드와 33개의 유닛 테스트, 그리고 1.9M 파라미터 규모의 tiny_lm 학습 사례는 CPU 전용 구현의 재현성과 실용성을 입증한다.

용어 해설

역방향 자동미분(Reverse-mode AD): — 역방향 자동미분은 순전파에서 수행된 연산의 연결을 기록하고 역전파에서 체인룰로 기울기를 계산하는 방식이다. 이 구현은 런타임에 그래프를 구성하는 define-by-run 접근을 사용하므로 동적 제어 흐름을 처리할 수 있다. 학습 과정에서 파라미터별 정확한 그래디언트를 제공하므로 미분 가능 연산의 구현과 그래프 무결성이 중요하다.
제로 카피 스트라이드 뷰(Zero-copy strided views): — 제로 카피 스트라이드 뷰는 동일한 저장소를 공유하면서 reshape, transpose, slice, expand 같은 뷰 연산을 메모리 복사 없이 구현하는 방식이다. 스트라이드와 오프셋만 조작해 다차원 배열 형태를 바꾸므로 메모리 사용을 최소화한다. 참조 카운팅과 함께 사용하면 뷰가 원본 버퍼의 라이프사이클을 안전하게 관리한다.
브로드캐스팅(Broadcasting): — 브로드캐스팅은 서로 다른 차원의 배열을 원소 연산을 위해 가상으로 확장해 정렬하는 규칙 세트이다. 실제 데이터 복사 없이 연산 시점에 인덱싱 규칙을 적용해 연산을 수행하므로 메모리와 성능 면에서 유리하다. NumPy 스타일 연산의 호환성과 코드 단순화를 위해 중요하다.
AVX2+FMA SIMD(AVX2+FMA): — AVX2와 FMA는 x86 계열의 벡터 명령 세트로, 동시에 여러 부동소수점 연산을 처리해 행렬 연산 성능을 높인다. 마이크로커널은 블록/타일 방식과 RHS(우측 행렬) 패킹을 결합해 캐시 지역성을 개선하고 메모리 대역폭 병목을 완화한다. CPU 전용 고성능 matmul 구현에서 핵심적인 최적화 수단이다.
프리-노름 디코더 스택(Pre-norm decoder stack): — 프리-노름 디코더 스택은 각 서브레이어에 LayerNorm을 적용한 후 잔차 연결을 더하는 구조로, 깊은 트랜스포머에서 학습 안정성이 향상된다. 이 구조는 디코더 블록을 쌓아 causal self-attention과 MLP를 결합하는 GPT 스타일 모델 구현에 쓰인다. 학습 중 그래디언트 흐름과 수렴 거동에 영향을 미치므로 설계 선택이 중요하다.

코드 예제

text

null

본문에 코드 블록이 포함되지 않음

언급된 리소스

GitHubTensorLib repository