Docker 컨테이너에서 Cloud TPU 애플리케이션 실행

Docker 컨테이너를 사용하면 코드와 필요한 모든 종속 항목을 배포 가능한 패키지 하나로 결합하여 애플리케이션을 쉽게 구성할 수 있습니다. TPU VM 내에서 Docker 컨테이너를 실행하면 Cloud TPU 애플리케이션을 간단하게 구성하고 공유할 수 있습니다. 이 문서에서는 Cloud TPU에서 지원하는 각 ML 프레임워크에서 Docker 컨테이너를 설정하는 방법을 설명합니다.

Docker 컨테이너에서 TensorFlow 모델 학습

TPU 기기

1. 현재 디렉터리에 이름이 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
   

2. Cloud Storage 버킷 만들기

   gsutil mb -c standard -l europe-west4 gs://your-bucket-name
   

3. 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
   

4. Dockerfile을 TPU VM에 복사합니다.

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
   

5. TPU VM에 SSH를 통해 연결

   gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
   

6. Docker 이미지를 빌드합니다.

   sudo docker build -t your-image-name .
   

7. Docker 컨테이너를 시작합니다.

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
   

8. 환경 변수 설정하기

   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
   

9. 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 포드

1. 현재 디렉터리에 이름이 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
   

2. 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
   

3. Dockerfile을 TPU VM에 복사합니다.

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
   

4. TPU VM에 SSH를 통해 연결

   gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
   

5. Docker 이미지를 빌드합니다.

   sudo docker build -t your-image-name .
   

6. Docker 컨테이너를 시작합니다.

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
   

7. 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"
   

학습 스크립트가 완료되면 리소스를 삭제해야 합니다.

1. exit를 입력하여 Docker 컨테이너에서 종료합니다.
2. TPU VM을 종료하려면 exit를 입력합니다.
3. TPU VM을 삭제합니다.

     $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
   

Docker 컨테이너에서 PyTorch 모델 학습

참고 현재 Docker 컨테이너에서 PyTorch를 실행하는 것은 단일 기기 TPU(v2-8, v3-8)에서만 지원됩니다.

1. 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
   

2. TPU VM에 SSH를 통해 연결합니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
   

3. TPU VM에서 컨테이너를 시작합니다.

   sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm bash
   

4. libtpu를 설치합니다.

   pip install torch_xla[tpuvm]
   export XRT_TPU_CONFIG="localservice;0;localhost:51011"
   

5. PyTorch XLA 저장소를 클론합니다.

   git clone --recursive https://github.com/pytorch/xla.git
   

6. ImageNet을 학습시킵니다.

   python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
   

학습 스크립트가 완료되면 리소스를 삭제해야 합니다.

1. exit를 입력하여 Docker 컨테이너에서 종료합니다.
2. TPU VM을 종료하려면 exit를 입력합니다.
3. TPU VM을 삭제합니다.

   $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
   

Docker 컨테이너에서 JAX 모델 학습

TPU 기기

1. TPU VM 만들기

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-vm-base
   

2. TPU VM에 SSH를 통해 연결합니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a
   

3. TPU VM에서 Docker 데몬을 시작합니다.

   sudo systemctl start docker
   

4. Docker 컨테이너 시작

   sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
   

5. JAX를 설치합니다.

   pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

6. FLAX를 설치합니다.

   pip install --upgrade clu
   git clone https://github.com/google/flax.git
   pip install --user -e flax
   

7. 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
   

학습 스크립트가 완료되면 리소스를 삭제해야 합니다.

1. exit를 입력하여 Docker 컨테이너에서 종료합니다.
2. TPU VM을 종료하려면 exit를 입력합니다.

3. 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 이미지를 만드는 방법을 보여줍니다.

1. 터미널 프롬프트에서 다음 환경 변수를 설정합니다.

   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 이름으로 바꿉니다.

2. 현재 디렉터리에 이름이 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
   

3. Docker 이미지를 빌드합니다.

   docker build -t $IMAGE .
   

4. Artifact Registry로 푸시하기 전에 Docker 이미지에 태그를 추가합니다. Artifact Registry 작업에 대한 자세한 내용은 컨테이너 이미지로 작업하기를 참조하세요.

   docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG
   

5. Docker 이미지를 Artifact Registry로 푸시합니다.

   docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
   

6. TPU VM을 만듭니다.

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone $ZONE \
   --accelerator-type $ACCEL_TYPE \
   --version tpu-vm-base
   

7. 모든 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"
   

8. 모든 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"
   

9. 모든 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 코드 실행