본문으로 건너뛰기

LLM 출력의 수치 발산을 막는 JAX Homeostatic Kernel

JAX와 PyTorch 사이에 물리 기반 순전파 필터를 배치해 LLM 출력의 수치 발산을 제어하려는 PoC다.

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

TL;DR

이 저장소는 Llama·Mistral 같은 확률적 LLM의 출력 스트림을 JAX 기반 2세대 Homeostatic Kernel로 감싸 수치 발산과 NaN 전파를 줄이려는 PoC입니다. PyTorch CUDA Tensor를 CUDA Array Interface v3로 JAX에 직접 연결하고, stop_gradient로 역전파 그래프를 차단하며, Laplacian·Burgers 점성 소산·skewness·L2 정규화를 순전파 필터에 결합합니다. README는 0-byte 복사, 32-byte 정렬, VRAM O(1), 통신 지연 100% 은폐 같은 목표를 제시하지만 이를 뒷받침하는 실행 결과나 비교 벤치마크는 제공하지 않습니다. 따라서 실제 배포용 성능·안전성보다 구조적 아이디어와 검증 항목을 담은 실험 단계 저장소로 보는 편이 타당합니다.

섹션별 상세

01
이 저장소는 확률적 출력에 의존하는 1세대 LLM과 수치 안정성을 담당하는 2세대 Homeostatic Kernel을 분리하는 PoC 구조를 제안합니다. LLM은 다음 토큰 예측과 추상 개념 합성을 맡고, 별도 커널은 입력과 출력 사이에서 시간적 인과성 및 수치 상태를 조정하는 방식입니다. README는 이를 생물학적 뇌의 추상 추론과 항상성 제어를 나눈 구조에 대응시키지만, 실제 검증 결과보다 설계 개념과 모듈 구성이 중심입니다.
02
입력 센서 로그와 시스템 교란은 adapters/interface/manifold.py를 거쳐 4D manifold의 정적 가상 뷰와 32-byte 정렬 구조로 변환됩니다. 이후 CUDA Array Interface v3를 이용해 PyTorch Tensor의 물리 포인터를 JAX/XLA 배열 공간에 직접 연결하고, LLM이 생성한 확률적 스트림은 kernel/physics_filter.py의 출력 필터로 전달됩니다. 이 흐름은 CPU/RAM을 거치지 않는 0-byte 복사와 비동기 통신·계산 중첩을 목표로 하지만, README 안에는 해당 지연 시간이나 처리량을 측정한 결과가 제시되지 않습니다.
03
autograd_free.py는 jax.lax.stop_gradient를 입력 경계와 출력 경계에 배치해 역전파 추적 그래프가 시간축을 따라 누적되지 않도록 설계됩니다. 저장소의 논리는 원시 입력을 격리한 뒤 순전파 계산만 통과시키고, donate_argnums=(1,) 및 사용자 정의 norm 연산으로 메모리 재사용을 유도하는 방식입니다. 이에 따라 문서는 문맥 길이나 시간 tick 수와 무관한 VRAM 공간 복잡도 O(1)을 목표로 제시하지만, test_memory_o1.py의 실제 실행 수치나 비교 기준은 본문에 포함하지 않습니다.
04
physics_filter.py는 Neumann 경계 패딩, Laplacian 기반 곡률 계산, Burgers 점성 소산, Schrödinger 기반 potential barrier를 순서대로 연결합니다. 곡률 κ에 동적 계수 σ_dynamic을 곱해 장벽 U_barrier를 만들고, T=exp(-2√(2m·U_barrier)/ℏ_eff) 형태의 전송 계수를 사용해 변동이 클수록 신호를 더 강하게 감쇠시키는 계산 흐름입니다. README는 이 필터가 NaN이나 급격한 통계 변동을 줄인다고 설명하지만, 실제 LLM 환각률 또는 안정성 개선 폭을 보여주는 벤치마크는 제시하지 않습니다.
05
허용 오차 영역에 들어온 작은 잡음은 Casimir 압력 식과 MUX 선택 연산을 이용해 elastic baseline으로 되돌리는 경로를 거칩니다. 정규화 거리 d=|X|+ε에서 ε=10^-6을 사용하고, P_casimir가 1/δ^4·X보다 큰 경우 X_elastic baseline을 선택해 NaN 전파를 차단하도록 구성했습니다. 이어서 3차 모멘트로 계산한 skewness를 α·S만큼 보정하고 L2 norm으로 재정규화해 분산된 4D sharding 축 사이의 신호 크기를 맞추는 흐름입니다.
06
interface/silicon_mux.py는 ((size + 7) & ~7) 패딩 식으로 32-byte 하드웨어 stride 정렬을 시도하고, if-else 대신 0.0f와 1.0f 마스크를 이용한 jax.lax.add와 jax.lax.mul 조합을 사용합니다. async_scheduler.py는 데이터 의존성이 없다고 판단한 통신을 jax.lax.psum all-reduce와 계산 사이에 겹치고, asyncio.Lock 생성은 실제 트래픽 경계까지 늦추는 구조입니다. 문서는 이 설계가 NCCL 동기화 비용을 100% 숨긴다고 표현하지만, 다중 노드 환경에서의 측정 로그나 재현 조건은 제공하지 않습니다.
07
저장소 구조에는 kernel, interface, adapters, tests 네 영역이 있고, 테스트 파일은 CAD 경계 수렴, 무한 루프 메모리 프로파일링, 7축 로봇 궤적 이상 검사를 대상으로 합니다. HuggingFace의 CausalLMOutput과 상용 API 스트림을 adapters 계층에서 받아 CUDA 포인터 브리지와 물리 필터를 통과시키는 연결점도 포함됩니다. 따라서 현재 결과물은 LLM 자체를 학습시키는 구현보다 기존 모델의 출력 스트림을 하드웨어 친화적 수치 필터로 감싸려는 PoC에 가깝고, 핵심 성능·안전 주장은 별도 실험으로 확인해야 합니다.

용어 해설

Homeostatic Kernel
LLM의 확률적 출력에서 발생하는 수치 발산을 후처리 단계에서 제어하려는 전방향 필터 구조입니다. 입력과 출력 사이에 별도 커널을 배치해 시간적 인과성, 수치 안정성, 메모리 사용량을 관리하는 개념으로 제시됩니다.
stop_gradient
자동미분 그래프에서 특정 값의 역전파 경로를 끊는 연산입니다. 이 저장소는 입력과 출력 경계에 jax.lax.stop_gradient를 배치해 과거 활성값 축적을 막고 VRAM 공간 복잡도를 O(1)로 유지한다고 설명합니다.
Burgers 방정식(Burgers' Equation)
유체의 비선형 흐름과 점성 소산을 표현하는 편미분방정식입니다. 저장소에서는 입력 스트림의 Laplacian 기반 곡률과 결합해 급격한 수치 변동을 감쇠하는 물리 필터의 일부로 사용합니다.
CUDA Array Interface v3
CUDA 메모리에 있는 배열의 포인터, shape, dtype 같은 정보를 프레임워크 사이에서 공유하는 Python 사양입니다. 이 프로젝트는 PyTorch Tensor의 VRAM 주소를 JAX 배열로 직접 넘겨 CPU 메모리 복사와 중간 객체 생성을 피하려 합니다.
Fused Multiply-Add
곱셈과 덧셈을 하나의 연산 흐름으로 결합하는 하드웨어 명령 방식입니다. 저장소는 조건 분기 대신 0.0f와 1.0f 마스크를 곱하고 더하는 형태로 변환해 GPU warp divergence와 분기 명령을 줄인다고 설명합니다.

기술

  • JAX
  • XLA
  • PyTorch
  • CuPy
  • HuggingFace
  • Llama
  • Mistral
  • CUDA Array Interface v3
  • DLPack
  • jax.lax.stop_gradient
  • jax.lax.psum
  • asyncio.Lock
  • NCCL
  • CUDA
  • FMA
  • SASS/PTX

활용 사례

  • 실시간 센서 로그의 수치 안정화
  • CAD 기하 공차와 경계 수렴 검사
  • 7축 로봇 궤적의 이상 신호 차단
  • HuggingFace CausalLMOutput 후처리
  • 상용 API 텍스트 스트림의 동시성 제어
  • 분산 GPU 추론에서 통신과 계산의 중첩
AI 분석 전체 내용 보기

AI 요약 · 북마크 · 개인 피드 설정 — 무료

출처 · 인용 안내

원문 발행 2026. 08. 25.수집 2026. 08. 25.출처 타입 RSS

인용 시 "요약 출처: AI Trends (aitrends.kr)"를 표기하고, 사실 확인은 원문 보기 기준으로 진행해 주세요. 자세한 기준은 운영 정책을 참고해 주세요.