[Lucid] 파라미터와 버퍼 시스템

🧲 파라미터와 버퍼 시스템

Lucid에서 학습 가능한 값(파라미터)과 러닝 스탯/상수(버퍼)를 명확히 분리한 이유는 두 가지다. 첫째, autograd가 추적해야 하는 Tensor와 그렇지 않은 Tensor를 구분해 메모리/성능을 최적화하기 위해서다. 둘째, state_dict와 디바이스 이동 시 저장/로드/이동 대상이 다르므로, 클래스 차원에서 역할을 분리하는 것이 유지보수에 유리하기 때문이다. PyTorch의 nn.Parameter/버퍼 개념을 그대로 벤치마킹하여 Lucid의 Parameter/Buffer 클래스를 설계했다.

🧭 설계 개요

• 파일: lucid/nn/parameter.py
• 핵심 목표: Tensor 인터페이스는 동일하게 유지하되, 생성 시 requires_grad/keep_grad 플래그를 고정해 학습 여부를 결정한다.
• 자동 등록: nn.Module.__setattr__에서 타입을 감지해 _parameters/_buffers 레지스트리에 자동 등록. 파라미터/버퍼가 아닌 일반 Tensor는 등록하지 않는다.
• 추적/비추적 분기: Parameter는 grad를 추적·누적하고, Buffer는 추적하지 않으며 grad를 보존하지 않는다. 이 구분은 optimizer, state_dict, device 이동 모두에 반영된다.

Lucid는 이 설계를 통해 “데이터의 성격”만 정하면 나머지(등록, 저장, 이동, 추적 여부)는 프레임워크가 자동으로 처리하는 흐름을 만들고자 했다. 파라미터는 학습 경로에 완전히 묶이고, 버퍼는 학습에 필요하지만 미분하지 않는 러닝 스탯이나 상수를 안전하게 보관한다.

🧱 Parameter 클래스

class Parameter(Tensor):
    def __init__(self, data, dtype=None, device="cpu"):
        if isinstance(data, Tensor):
            data = data.data
        super().__init__(data, requires_grad=True, keep_grad=True, dtype=dtype, device=device)

• 역할: 학습 가능한 텐서. 항상 requires_grad=True, keep_grad=True로 생성해 backward 후에도 grad를 보존.
• 데이터 승격: 입력이 Tensor이면 .data로 한 단계 벗겨 실데이터만 넘긴다(그래프 공유 방지).
• 디바이스/ dtype: 모듈의 to() 호출 시 Parameter가 자동 이동되며, 생성 시에도 명시 가능.

구현 메모

• 파라미터를 일반 Tensor로 잘못 생성하면 autograd는 추적하지만 state_dict에는 누락될 수 있다. 이를 방지하기 위해 register_parameter/__setattr__에서 타입을 강제 체크한다.
• keep_grad=True로 설정해 backward 이후에도 grad를 유지, optimizer 단계에서 사용한다. 필요 시 param.grad.zero()나 optimizer가 직접 초기화한다.
• dtype/device를 모듈 생성 시점에 맞추기 위해, 모듈 내부에서는 to() 호출 전에 버퍼와 동일하게 현재 모듈 디바이스를 기본값으로 사용한다.

사용 예시

class Linear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(lucid.randn(out_features, in_features))
        self.bias = nn.Parameter(lucid.zeros(out_features))
    def forward(self, x):
        return x @ self.weight.T + self.bias

속성 대입만으로 자동 등록되므로 별도 registry 조작이 필요 없다.

🧱 Buffer 클래스

class Buffer(Tensor):
    def __init__(self, data, dtype=None, device="cpu"):
        if isinstance(data, Tensor):
            data = data.data
        super().__init__(data, requires_grad=False, keep_grad=False, dtype=dtype, device=device)

• 역할: 러닝 스탯(예: BatchNorm의 running_mean/var)이나 상수 텐서를 보관. requires_grad=False로 autograd 그래프에서 제외, grad도 보존하지 않는다.
• 데이터 승격: Parameter와 동일하게 입력이 Tensor면 .data를 사용해 그래프 공유를 끊는다.
• 디바이스/ dtype: 모듈의 to() 호출로 이동하며, 생성 시 선택 가능.

사용 규칙

• 버퍼는 학습하지 않지만 state_dict에는 저장된다. load_state_dict 시 타입을 유지한 채 복원된다.
• autograd를 타지 않으므로, 연산에 참여해도 그래프를 확장하지 않는다.
• 러닝 스탯(예: BatchNorm의 running_mean/var), 드롭아웃의 마스크 재사용, 정규화 계수 등 학습 중 값이 바뀔 수 있지만 grad가 필요 없는 값에 적합하다.
• register_buffer로 None을 등록해 조건부로 존재하는 버퍼 패턴도 지원한다(PyTorch 동일).

버퍼는 학습 그래프 밖에 있지만 모델 상태의 일부다. 예를 들어 BatchNorm의 running_mean/var는 학습 동안 업데이트되고, eval에서는 고정된다. 이 값이 제대로 저장/로드/이동되지 않으면 재현성이 깨지거나 추론 품질이 떨어진다. 따라서 버퍼를 파라미터와 동일한 state_dict 경로로 관리하면서도 autograd와 optimizer에서는 완전히 분리해둔 것이 핵심이다.

사용 예시

class BN(nn.Module):
    def __init__(self, num_features, momentum=0.1, eps=1e-5):
        super().__init__()
        self.running_mean = nn.Buffer(lucid.zeros(num_features))
        self.running_var = nn.Buffer(lucid.ones(num_features))
        self.momentum = momentum; self.eps = eps
    def forward(self, x):
        # training 시 running 통계 업데이트, eval 시 고정
        ...

🔗 nn.Module과의 연결

• __setattr__ 자동 등록: Parameter는 _parameters, Buffer는 _buffers에 저장된다. 동일 이름 재할당 시 이전 레지스트리에서 제거해 충돌을 방지.
• state_dict/load_state_dict: 파라미터와 버퍼를 prefix 기반 키로 평탄화하여 저장하고, 로드 시 Tensor로 감싼 뒤 .data 교체로 복원.
• to(device): 모듈의 to가 자신의 파라미터/버퍼를 우선 이동(recurse=False)하고, 이후 자식 모듈을 재귀 이동한다.
• train/eval: 버퍼는 대개 모드에 따라 갱신 여부가 달라진다(BN running stats 등). 모듈의 모드 전파로 버퍼 사용 패턴이 결정된다.

파라미터와 버퍼는 등록 후 state_dict에서 키로 표현되는 순간부터 모델 트리의 일부가 된다. 키 구조는 모듈 계층을 반영하므로, 동일한 모델을 재현하려면 키-값 쌍이 완전히 일치해야 한다. 이 흐름을 끊지 않기 위해 등록·저장·로드·이동이 모두 동일한 규약을 따른다.

🧮 grad 처리와 메모리

• Parameter는 항상 grad를 추적/보존하므로 학습 후 optimizer 단계에서 grad가 필요하다.
• Buffer는 grad를 추적하지 않으므로 그래프 메모리를 늘리지 않는다. 중간 계산에 포함되더라도 requires_grad=False로 처리되어 backward 경로에 참여하지 않는다.
• 두 타입 모두 .data를 직접 노출하므로, 실수로 in-place 수정 시 autograd가 감지하지 못할 수 있다. 학습 중 값 변경은 연산을 통해 수행하거나, 필요 시 .detach() 후 새 Tensor를 만들어 교체한다.
• optimizer는 module.parameters() 순회 결과만 업데이트 대상으로 삼는다. 버퍼는 optimizer가 건드리지 않는다.
• grad 누적 정책: Lucid는 PyTorch와 동일하게 backward를 여러 번 호출하면 grad가 누적된다. 파라미터 grad를 초기화하려면 명시적으로 zero()나 optimizer의 zero_grad가 필요하다.

실무적으로는 파라미터의 grad를 언제 초기화하고, 버퍼는 어느 시점에 업데이트할지 명확히 해야 한다. Lucid는 PyTorch와 동일한 규칙을 따르므로, 기존 파이프라인(optimizer.zero_grad → forward → backward → step)을 그대로 적용할 수 있다. 버퍼는 gradient 경로에 포함되지 않으니 오직 모듈 로직에서만 값을 갱신한다.

🧳 state_dict 직렬화와 필터링

• Parameter/Buffer 모두 기본적으로 저장 대상이다. keep_vars=False면 numpy 배열로 직렬화해 그래프와 분리, keep_vars=True면 Tensor 그대로 저장해 디버깅용으로 사용할 수 있다.
• _state_dict_pass_attr에 이름을 넣으면 저장 시 제외된다. 학습과 무관한 캐시나 임시 버퍼를 건너뛸 때 사용.
• 로드 시 strict 옵션으로 누락/불일치 키를 검증, PyTorch와 동일한 에러 메시지 스타일로 보고.

직렬화 경로

1) state_dict 호출 → 파라미터/버퍼를 prefix 키로 채움.
2) _state_dict_pass_attr에 해당하는 키 제거.
3) 디스크/네트워크에 저장.
4) load_state_dict → numpy/리스트를 Tensor로 감싸 device 맞춰 .data 교체.

직렬화 경로가 명확해야 모델 재현성이 확보된다. Lucid는 state_dict를 단순한 OrderedDict[str, array] 형태로 만들고, 필요 시 Tensor 그대로 보관하는 옵션을 둬서 디버깅에 유연성을 확보했다. strict 검증을 통해 훈련/추론 코드 불일치도 조기에 감지할 수 있다.

🛰 디바이스/타입 전파와 다중 백엔드 고려

• Parameter/Buffer 생성 시 기본 device는 "cpu"로 설정하지만, 모듈 단위로 to("gpu") 호출 시 전체 트리를 순회하며 이동한다.
• MLX(Metal) 백엔드에서도 동일한 인터페이스를 유지하기 위해 device를 문자열로 통일했다. dtype 역시 생성자와 to()에서 명시적으로 설정 가능.
• 백엔드별 차이를 감추기 위해 파라미터/버퍼는 가능한 한 동일한 Tensor API를 사용하며, NumPy/MLX 간 데이터 변환은 Tensor 내부에 캡슐화된다.

디바이스 전파는 종종 버그가 발생하는 영역이다. Lucid에서는 모듈이 자신의 파라미터/버퍼를 먼저 이동시키고, 이후 자식 모듈을 이동시킨다. 이렇게 하면 중복 이동을 피하면서도 트리 전체가 일관된 디바이스를 갖는다. dtype 변경도 to()에 동일하게 반영할 수 있도록 API를 단순화했다.

🛠 구현 시 어려움과 해결

1. 그래프 공유 문제: 기존 Tensor를 Parameter/Buffer로 감쌀 때 그래프가 공유되어 중간 grad가 엉킬 수 있었다. → 생성자에서 Tensor 입력이면 .data로 unwrap.
2. 등록 누락: 일반 Tensor를 속성에 대입하면 학습 대상인데도 state_dict에 빠지는 문제. → __setattr__에서 타입 체크를 엄격히 하고, 필요 시 register_parameter/register_buffer를 사용하도록 문서화.
3. 디바이스 불일치: 모듈 이동 시 자식 모듈이 중복 이동하거나 놓치는 경우. → 모듈 to()에서 자기 파라미터/버퍼만 이동하고, 이후 자식 모듈을 순회해 한 번씩만 이동하도록 정리.
4. state_dict 뷰: 저장된 numpy 뷰가 원본과 연결되는 문제. → .numpy() 복사본을 반환해 안전하게 직렬화.
5. 버퍼 업데이트 규약: 학습 모드에서만 갱신해야 하는 버퍼(BN running stats 등)를 어떻게 식별할지 → 버퍼 자체는 플래그를 가지지 않지만 모듈 모드(training)를 활용해 업데이트 경로를 모듈에서 제어.
6. None 처리: 조건부로 존재하는 파라미터/버퍼를 허용하려면 None을 허용해야 한다. → 등록 API에서 None을 허용하고, state_dict에서 자동 스킵되도록 했다.

✅ 정리

Lucid의 파라미터/버퍼 시스템은 PyTorch의 개념을 그대로 가져오되, NumPy/MLX 백엔드에 맞게 단순화했다. Parameter는 항상 학습 대상이며 grad를 보존하고, Buffer는 학습 대상이 아니지만 state_dict에 포함된다. nn.Module의 자동 등록/저장/이동 로직과 결합되어, 모델 정의 시 “데이터의 성격만” 지정하면 나머지는 프레임워크가 처리한다. 이 설계 덕분에 모델 구현자는 파라미터/버퍼 구분만 명확히 하면 되고, 저장·로드·디바이스 이동·학습 모드 전환이 일관되게 동작한다.