ZIPSLICER 📁✂️

대규모 PyTorch 체크포인트 (checkpoints)의 점진적 로딩 (incremental loading)을 위한 라이브러리<br>
작성자가 직접 쓴 블로그 포스트 소개 읽기

개요 (Synopsis)

import torch
import zipslicer

...

이 예제와 단위 테스트 (unit-tests)를 실행하세요:

python examples/example_resnet18.py

pytest -o log_cli=true --capture=tee-sys -p no:asyncio

체크포인트의 호환성을 테스트하세요:

python tests/test_checkpoint_readonly.py your_magnificent_checkpoint.pth

모든 항목이 초록색으로 표시되면 정상적으로 작동할 것입니다.

요구 사항 (Prerequisites)

• 지원되는 Python 및 torch 버전: python-3.10 + torch-(1.11,1.12,stable) python-3.11 + torch:stable
• 일반적으로, zipslicer는 충분히 최신 버전의 PyTorch 설치 환경에서 작동해야 합니다. 포함된 안전한 테스트를 사용하여 사용자의 PyTorch 및 체크포인트와 zipslicer 간의 호환성을 확인하십시오. 이 라이브러리는 순수 Python 라이브러리이므로, 특정 CPU 아키텍처는 중요하지 않습니다.
• vanilla torch.save(...)를 통해 모델의 state_dict를 저장하여 생성된 체크포인트 - Torch는 ZIP 압축을 사용하지 않으므로 기본 설정으로도 충분할 것입니다.
• 점진적으로 로드된 체크포인트를 활용할 수 있는 애플리케이션 - 즉, 애플리케이션이 루프 내에서 즉시 모든 state_dict.items()를 로드하기만 한다면 이 라이브러리를 사용하는 의미가 크지 않습니다. 코드가 state_dict.keys()를 읽고 (필요한 경우 state_dict.get_meta(k)도 포함) 이를 지능적으로 사용하여 한 번에 state_dict[k] 텐서(tensors)의 일부 집합만 처리하도록 작성해야 합니다. 일반적인 영감을 얻으려면 이것 (HF)과 이것 (arxiv)을 읽어보시기 바랍니다.

추가적인 엔지니어링을 더한다면, 하룻밤 정도의 시간을 기다릴 의향이 있다는 전제하에 BLOOM-176B나 FLAN-T5-XXL과 같은 대규모 언어 모델 (Large Language Models (LLMs))을 가정용 미드레인지 GPU 한 장에서 실행하는 것이 가능할 것입니다. 대규모 배치 (large batch) 환경에서는, 예를 들어 일련의 문서들을 임베딩 (embeddings)으로 처리하는 것과 같이 실질적인 의미를 가질 수도 있습니다.

설치 (Install)

일반적으로 zipslicer/zipslicer 디렉토리를 프로젝트의 소스 트리 (source tree)로 복사하는 것만으로 충분합니다.

만약 출처가 불분명한 실행 모듈에 대해 공식적인 의례 중심의 설치 프로세스를 선호하신다면, 곧 이 부티크 소프트웨어 모듈을 pip를 통해 설치할 수 있는 기능이 제공될 예정입니다: pip install zipslicer

참고 사항 (Notes)

• 이 라이브러리는 체크포인트 (checkpoints)에서 PyTorch 텐서 (tensors)를 읽는 용도로만 사용됩니다. 쓰는 기능 (writing)은 향후 작업 과제로 남겨두었습니다.
• 로드된 state_dict에 쓰는 것은 권장되지 않지만, 작동은 할 것입니다. 다만, 현재로서는 키 (keys)를 반복(iterating)하는 동안 이 업데이트가 반영될 것이라고 기대하며 작업을 수행하는 것은 피해야 합니다.
• 아마도 더 중요한 점은, 범용 피클 (general-purpose pickles)은 지원되지 않는다는 것입니다. 이 라이브러리의 설계상 신경망 클래스 인스턴스 (neural network class instances) 전체를 로드할 수는 없습니다. 보통 이는 필요하지 않으며, PyTorch 공식 문서에서도 모델 직렬화 (model serialization)를 위해 state_dict를 사용할 것을 권장합니다. 저희는 state_dict를 지원합니다.
• 일부 희귀한 텐서 타입 (예: PyTorch 양자화 텐서 (quantized tensors) - 정상 작동하는 정수 텐서 (integer tensors)와 혼동하지 마십시오)은 아직 지원되지 않습니다. 이 점이 불편하시다면 이슈 (issues)를 통해 경험을 공유해 주세요.
• 저희는 HF safetensors 프로젝트에 경의를 표하지만, 그들의 방식과 비교했을 때 저희의 접근 방식은 상당한 시간과 저장 공간이 소요되는 체크포인트 변환 (checkpoint conversion) 과정을 필요로 하지 않는다는 점에 유의하십시오.

사실, 두 접근 방식은 서로 보완적일 수 있습니다. safetensors로 변환하기 위해서는 PyTorch 체크포인트(checkpoint)에서 텐서(tensor)를 어떤 방식으로든 로드해야 하며, 기본 로딩 메커니즘은 가용 RAM(Random Access Memory)에 의해 제한되기 때문입니다.

고려 중인 향후 기능들

이러한 기능 중 일부에 관심이 있으시다면, 이슈(issue)를 생성해 주세요:

• 텐서 슬라이스(tensor slices)의 효율적인 로딩 - 샤딩된 배포(sharded deployments) 환경에서 텐서 병렬성(tensor parallelism)을 구현하기 위함
• 네트워크를 통한 소스 체크포인트(source checkpoint) 접근
• 체크포인트에 인플레이스(in-place)로 쓰기
• 다른 체크포인트 형식으로의 점진적 변환 (Incremental conversion)