v6e TPU에서 MaxDiffusion 추론
이 튜토리얼에서는 TPU v6e에서 MaxDiffusion 모델을 제공하는 방법을 보여줍니다. 이 튜토리얼에서는 Stable Diffusion XL 모델을 사용하여 이미지를 생성합니다.
시작하기 전에
칩 4개가 있는 TPU v6e를 프로비저닝할 준비를 합니다.
- Google 계정에 로그인합니다. 아직 계정이 없다면 새 계정을 만듭니다.
- Google Cloud 콘솔의 프로젝트 선택기 페이지에서 Google Cloud 프로젝트를 선택하거나 만듭니다.
- Google Cloud 프로젝트에 결제를 사용 설정합니다. Google Cloud를 사용하려면 결제가 필요합니다.
- gcloud alpha 구성요소를 설치합니다.
다음 명령어를 실행하여 최신 버전의
gcloud
구성요소를 설치합니다.gcloud components update
Cloud Shell을 사용하여 다음
gcloud
명령어를 통해 TPU API를 사용 설정합니다. Google Cloud 콘솔에서도 사용 설정할 수 있습니다.gcloud services enable tpu.googleapis.com
TPU VM의 서비스 ID를 만듭니다.
gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE
TPU 서비스 계정을 만들고 Google Cloud 서비스에 대한 액세스 권한을 부여합니다.
서비스 계정을 사용하면 Google Cloud TPU 서비스가 다른 Google Cloud 서비스에 액세스할 수 있습니다. 사용자 관리형 서비스 계정이 권장됩니다. 다음 가이드를 따라 역할을 만들고 부여합니다. 다음 역할이 필요합니다.
- TPU 관리자: TPU를 만드는 데 필요합니다.
- 스토리지 관리자: Cloud Storage에 액세스하는 데 필요합니다.
- 로그 작성자: Logging API를 사용하여 로그를 작성하는 데 필요
- 모니터링 측정항목 작성자: Cloud Monitoring에 측정항목을 쓰는 데 필요
Google Cloud로 인증하고 Google Cloud CLI의 기본 프로젝트와 영역을 구성합니다.
gcloud auth login gcloud config set project PROJECT_ID gcloud config set compute/zone ZONE
용량 확보
Cloud TPU 영업팀 또는 계정팀에 문의하여 TPU 할당량을 요청하고 용량에 관해 문의하세요.
Cloud TPU 환경 프로비저닝
v6e TPU는 GKE, GKE 및 XPK 또는 대기열에 추가된 리소스로 프로비저닝할 수 있습니다.
기본 요건
- Google Cloud 프로젝트 내에서 액세스할 수 있는 최대 칩 수를 지정하는
TPUS_PER_TPU_FAMILY
할당량이 프로젝트에 충분한지 확인합니다. - 이 튜토리얼은 다음 구성으로 테스트되었습니다.
- Python
3.10 or later
- Nightly 소프트웨어 버전:
- nightly JAX
0.4.32.dev20240912
- nightly LibTPU
0.1.dev20240912+nightly
- nightly JAX
- 안정적인 소프트웨어 버전:
- JAX +
v0.4.35
의 JAX 라이브러리
- JAX +
- Python
- 프로젝트에 다음에 충분한 TPU 할당량이 있는지 확인합니다.
- TPU VM 할당량
- IP 주소 할당량
- 하이퍼디스크 균형 할당량
- 사용자 프로젝트 권한
- XPK와 함께 GKE를 사용하는 경우 XPK를 실행하는 데 필요한 권한은 사용자 또는 서비스 계정의 Cloud 콘솔 권한을 참고하세요.
TPU v6e 프로비저닝
gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \ --node-id TPU_NAME \ --project PROJECT_ID \ --zone ZONE \ --accelerator-type v6e-4 \ --runtime-version v2-alpha-tpuv6e \ --service-account SERVICE_ACCOUNT
list
또는 describe
명령어를 사용하여 큐에 추가된 리소스의 상태를 쿼리합니다.
gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE}
대기열에 추가된 리소스 요청 상태의 전체 목록은 대기열에 추가된 리소스 문서를 참고하세요.
SSH를 사용하여 TPU에 연결
gcloud compute tpus tpu-vm ssh TPU_NAME
Conda 환경 만들기
Miniconda의 디렉터리를 만듭니다.
mkdir -p ~/miniconda3
Miniconda 설치 프로그램 스크립트를 다운로드합니다.
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
Miniconda를 설치합니다.
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
Miniconda 설치 프로그램 스크립트를 삭제합니다.
rm -rf ~/miniconda3/miniconda.sh
PATH
변수에 Miniconda를 추가합니다.export PATH="$HOME/miniconda3/bin:$PATH"
~/.bashrc
를 새로고침하여PATH
변수에 변경사항을 적용합니다.source ~/.bashrc
새 Conda 환경을 만듭니다.
conda create -n tpu python=3.10
Conda 환경을 활성화합니다.
source activate tpu
MaxDiffusion 설정
MaxDiffusion 저장소를 클론하고 MaxDiffusion 디렉터리로 이동합니다.
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
mlperf-4.1
브랜치로 전환합니다.git checkout mlperf4.1
MaxDiffusion을 설치합니다.
pip install -e .
종속 항목을 설치합니다.
pip install -r requirements.txt
JAX를 설치합니다.
pip install -U --pre jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
추가 종속 항목을 설치합니다.
pip install huggingface_hub==0.25 absl-py flax tensorboardX google-cloud-storage torch tensorflow transformers
이미지 생성
TPU 런타임을 구성할 환경 변수를 설정합니다.
LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"
src/maxdiffusion/configs/base_xl.yml
에 정의된 프롬프트와 구성을 사용하여 이미지를 생성합니다.python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"
이미지가 생성되면 TPU 리소스를 삭제해야 합니다.
삭제
TPU를 삭제합니다.
gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \ --project PROJECT_ID \ --zone ZONE \ --force \ --async