google/paxml

어떤 프로젝트인가요?

핵심 기능

• 01dataclass와 Fiddle을 기반으로 구성 가능한 실험 및 하이퍼파라미터를 제공합니다.
• 02pjit/SPMD, pmap, TPU 멀티슬라이스 구성을 위한 병렬화를 지원합니다.
• 031B, 16B 및 GPT-3 XL 스타일 모델을 위한 C4 학습 예제 구성을 제공합니다.
• 04SeqIO, Lingvo 어댑터, 사용자 정의 입력 등 다양한 데이터 입력 방식을 지원합니다.
• 05TPU v4 Pod에서 Model FLOPs Utilization을 벤치마크합니다.

활용 사례

• Google Cloud TPU v4 Pod에서 대규모 decoder-only 언어 모델을 학습합니다.
• 모델 및 병렬화 설정을 구성해 재현 가능한 JAX 기반 실험을 실행합니다.
• 다양한 모델 및 칩 규모에서 학습 효율성(MFU)을 벤치마크합니다.

빠른 시작

1. 1

   Cloud TPU VM을 생성합니다(먼저 환경 변수를 설정하세요).

   gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR
2. 2

   TPU VM에 SSH로 접속합니다.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
3. 3

   TPU용 JAX와 함께 Pax를 PyPI에서 설치합니다.

   python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
4. 4

   예제 테스트 모델을 실행합니다.

   python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>

사용 전 알아둘 점

• !Cloud TPU 하드웨어용으로 설계되었으며, README의 설정 방법과 예제는 TPU VM을 전제로 합니다.
• !설정하려면 GCP 프로젝트와 스토리지 버킷 같은 Google Cloud 리소스가 필요합니다.

토픽

자주 묻는 질문

paxml은(는) 무엇인가요?

Pax는 주로 TPU에서 대규모 언어 모델을 대상으로 머신러닝 실험을 구성하고 실행하는 JAX 기반 프레임워크입니다. Pax(Paxml)는 JAX를 기반으로 머신러닝 실험을 구성하고 실행하는 프레임워크입니다. 완전히 구성 가능한 실험과 병렬화를 중심으로 설계되었으며, decoder-only Transformer 언어 모델 같은 대규모 모델 학습을 목표로 합니다. 모델, 태스크, 러너는 dataclass 기반 하이퍼파라미터로 정의하며, 이 하이퍼파라미터는 함께 제공되는 Praxis 레이어 라이브러리를 기반으로 하고 Fiddle 구성을 지원합니다.

paxml은(는) 무료로 사용할 수 있나요? 라이선스는 무엇인가요?

paxml은(는) Apache-2.0 라이선스의 오픈소스입니다. 상업용 제품에 사용하기 전에 라이선스 조건을 확인하세요.

paxml은(는) 아직 유지보수되고 있나요?

활발히 유지보수 중. paxml의 마지막 커밋은 2026년 10월 26일입니다.

paxml은(는) 어디에 사용할 수 있나요?

Google Cloud TPU v4 Pod에서 대규모 decoder-only 언어 모델을 학습합니다. 모델 및 병렬화 설정을 구성해 재현 가능한 JAX 기반 실험을 실행합니다. 다양한 모델 및 칩 규모에서 학습 효율성(MFU)을 벤치마크합니다.

paxml는 어떻게 설치하나요?

Cloud TPU VM을 생성합니다(먼저 환경 변수를 설정하세요).: gcloud compute tpus tpu-vm create $TPU_NAME --zone=$ZONE --version=$VERSION --project=$PROJECT --accelerator-type=$ACCELERATOR TPU VM에 SSH로 접속합니다.: gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE TPU용 JAX와 함께 Pax를 PyPI에서 설치합니다.: python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html 예제 테스트 모델을 실행합니다.: python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=gs://<your-bucket>

paxml은(는) 어디서 받을 수 있나요?

소스 코드와 설치 방법은 GitHub에 있습니다: https://github.com/google/paxml