Docker 컨테이너에서 Cloud TPU 애플리케이션 실행
Docker 컨테이너를 사용하면 코드와 필요한 모든 종속 항목을 배포 가능한 패키지 하나로 결합하여 애플리케이션을 쉽게 구성할 수 있습니다. TPU VM 내에서 Docker 컨테이너를 실행하면 Cloud TPU 애플리케이션을 간단하게 구성하고 공유할 수 있습니다. 이 문서에서는 Cloud TPU에서 지원하는 각 ML 프레임워크에서 Docker 컨테이너를 설정하는 방법을 설명합니다.
Docker 컨테이너에서 TensorFlow 모델 학습
TPU 기기
현재 디렉터리에 이름이
Dockerfile
인 파일을 만들고 다음 텍스트를 붙여넣습니다.FROM python:3.8 RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so RUN git clone https://github.com/tensorflow/models.git WORKDIR ./models RUN pip install -r official/requirements.txt ENV PYTHONPATH=/models
Cloud Storage 버킷 만들기
gsutil mb -c standard -l europe-west4 gs://your-bucket-name
TPU VM을 만듭니다.
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-tf-2.11.0
Dockerfile을 TPU VM에 복사합니다.
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
TPU VM에 SSH를 통해 연결
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker 이미지를 빌드합니다.
sudo docker build -t your-image-name .
Docker 컨테이너를 시작합니다.
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
환경 변수 설정하기
export STORAGE_BUCKET=gs://your-bucket-name export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x export ZONE=europe-west4-a export TPU_NAME=local
ResNet을 학습시킵니다.
python3 official/vision/train.py \ --tpu=local \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"
TPU 포드
현재 디렉터리에 이름이
Dockerfile
인 파일을 만들고 다음 텍스트를 붙여넣습니다.FROM python:3.8 RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.0/libtpu.so -o /lib/libtpu.so RUN git clone https://github.com/tensorflow/models.git WORKDIR ./models RUN pip install -r official/requirements.txt ENV PYTHONPATH=/models
TPU VM을 만듭니다.
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v3-32 \ --version=tpu-vm-tf-2.11.0-pod
Dockerfile을 TPU VM에 복사합니다.
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
TPU VM에 SSH를 통해 연결
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker 이미지를 빌드합니다.
sudo docker build -t your-image-name .
Docker 컨테이너를 시작합니다.
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
ResNet을 학습시킵니다.
python3 official/vision/train.py \ --tpu=your-tpu-name \ --experiment=resnet_imagenet \ --mode=train_and_eval \ --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \ --model_dir=${MODEL_DIR} \ --params_override="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
학습 스크립트가 완료되면 리소스를 삭제해야 합니다.
exit
를 입력하여 Docker 컨테이너에서 종료합니다.- TPU VM을 종료하려면
exit
를 입력합니다. - TPU VM을 삭제합니다.
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Docker 컨테이너에서 PyTorch 모델 학습
Cloud TPU VM을 만듭니다.
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-pt-1.13
TPU VM에 SSH를 통해 연결합니다.
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
TPU VM에서 컨테이너를 시작합니다.
sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm bash
libtpu를 설치합니다.
pip install torch_xla[tpuvm] export XRT_TPU_CONFIG="localservice;0;localhost:51011"
PyTorch XLA 저장소를 클론합니다.
git clone --recursive https://github.com/pytorch/xla.git
ImageNet을 학습시킵니다.
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
학습 스크립트가 완료되면 리소스를 삭제해야 합니다.
exit
를 입력하여 Docker 컨테이너에서 종료합니다.- TPU VM을 종료하려면
exit
를 입력합니다. - TPU VM을 삭제합니다.
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
Docker 컨테이너에서 JAX 모델 학습
TPU 기기
TPU VM 만들기
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-vm-base
TPU VM에 SSH를 통해 연결합니다.
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
TPU VM에서 Docker 데몬을 시작합니다.
sudo systemctl start docker
Docker 컨테이너 시작
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
JAX를 설치합니다.
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
FLAX를 설치합니다.
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
FLAX MNIST 학습 스크립트 실행
cd flax/examples/mnist python3 main.py --workdir=/tmp/mnist \ --config=configs/default.py \ --config.learning_rate=0.05 \ --config.num_epochs=5
학습 스크립트가 완료되면 리소스를 삭제해야 합니다.
exit
를 입력하여 Docker 컨테이너에서 종료합니다.- TPU VM을 종료하려면
exit
를 입력합니다. TPU VM을 삭제합니다.
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
TPU 포드
TPU Pod에서 JAX 코드를 실행하는 경우 모든 TPU 워커에서 JAX 코드를 동시에 실행해야 합니다. 이를 위한 한 가지 방법은 gcloud compute tpus tpu-vm ssh
명령어를 --worker=all
및 --command
플래그와 함께 사용하는 것입니다. 다음 절차는 각 TPU 워커를 더 쉽게 설정할 수 있도록 Docker 이미지를 만드는 방법을 보여줍니다.
터미널 프롬프트에서 다음 환경 변수를 설정합니다.
export USER=your-user-id export PROJECT=your-project-name export REPO_NAME=your-repo-name export IMAGE_NAME=your-image-name export TAG=your-tag export TPU_NAME=your-tpu-name export ZONE=europe-west4-a export ACCEL_TYPE=v2-8
다음과 같이 바꿉니다.
your-user-id
를 사용자 ID로 바꿉니다.your-project-name
을 Google Cloud 프로젝트 이름으로 바꿉니다.your-repo-name
을 Artifact Registry 저장소 이름으로 바꿉니다.your-image-name
을 Docker 이미지 이름으로 바꿉니다.your-tag
를 Docker 이미지의 태그로 바꿉니다.your-tpu-name
을 TPU 이름으로 바꿉니다.
현재 디렉터리에 이름이
Dockerfile
인 파일을 만들고 다음 텍스트를 붙여넣습니다.FROM google-cloud-cli:latest RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html RUN pip install --upgrade clu RUN git clone https://github.com/google/flax.git RUN pip install --user -e flax WORKDIR ./flax/examples/mnist
Docker 이미지를 빌드합니다.
docker build -t $IMAGE .
Artifact Registry로 푸시하기 전에 Docker 이미지에 태그를 추가합니다. Artifact Registry 작업에 대한 자세한 내용은 컨테이너 이미지로 작업하기를 참조하세요.
docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG
Docker 이미지를 Artifact Registry로 푸시합니다.
docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
TPU VM을 만듭니다.
gcloud compute tpus tpu-vm create $TPU_NAME \ --zone $ZONE \ --accelerator-type $ACCEL_TYPE \ --version tpu-vm-base
모든 TPU 워커에서 Artifact Registry의 Docker 이미지를 가져옵니다.
gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="sudo usermod -a -G docker $USER" gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet" gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="docker pull europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag"
모든 TPU 워커에서 컨테이너를 실행합니다.
gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag bash"
모든 TPU 워커에서 학습 스크립트를 실행합니다.
gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \ --zone $ZONE \ --command="docker exec --privileged $CONTAINER_NAME python3 main.py --workdir=/tmp/mnist \ --config=configs/default.py \ --config.learning_rate=0.05 \ --config.num_epochs=5"
학습 스크립트가 완료되면 리소스를 삭제해야 합니다. 다음 명령어를 사용하여 TPU VM을 삭제합니다.
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
다음 단계
- Cloud TPU 튜토리얼
- Cloud TPU 사용자 가이드
- Cloud TPU 시스템 아키텍처
- TPU Pod 슬라이스에서 TensorFlow 코드 실행
- TPU Pod 슬라이스에서 JAX 코드 실행