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

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

중요: TPU VM에서 Docker 포드를 최대 256개까지 실행할 수 있습니다.

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

TPU 기기

1. 현재 디렉터리에 이름이 Dockerfile인 파일을 만들고 다음 텍스트를 붙여넣습니다.

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.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.16.1-pjrt
   

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
   

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"
   

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

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

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

TPU Pod

1. 현재 디렉터리에 이름이 Dockerfile인 파일을 만들고 다음 텍스트를 붙여넣습니다.

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.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.16.1-pod-pjrt
   

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=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*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
   

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

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

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

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

TPU 기기

1. Cloud TPU VM을 만듭니다.

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

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

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

3. 야간 PyTorch/XLA 이미지를 사용하여 TPU VM에서 컨테이너를 시작합니다.

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

   참고: 이 명령어를 실행하면 터미널이 실행 중인 컨테이너에 연결되었음을 나타내는 명령 프롬프트가 변경됩니다.

4. TPU 런타임 구성

   PyTorch/XLA 런타임은 PJRT와 XRT라는 두 가지 옵션이 있습니다. XRT를 사용할 이유가 없으면 PJRT를 사용하는 것이 좋습니다. 다양한 런타임 구성에 대해 자세히 알아보려면 XRT를 사용해야 하는 이유를 참조하세요. 다양한 런타임 구성에 대한 자세한 내용은 PJRT 런타임 문서를 참조하세요.

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

   export XRT_TPU_CONFIG="localservice;0;localhost:51011"
   

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

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

6. ResNet50을 학습시킵니다.

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

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

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

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

TPU Pod

TPU Pod에서 PyTorch 코드를 실행하는 경우 모든 TPU 워커에서 코드를 동시에 실행해야 합니다. 이를 위한 한 가지 방법은 gcloud compute tpus tpu-vm ssh 명령어를 --worker=all 및 --command 플래그와 함께 사용하는 것입니다. 다음 절차에서는 각 TPU 워커를 더 쉽게 설정할 수 있도록 Docker 이미지를 만드는 방법을 보여줍니다.

1. TPU VM을 만듭니다.

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=us-central2-b \
   --accelerator-type=v4-32 \
   --version=tpu-ubuntu2204-base
   

2. Docker 그룹에 현재 사용자를 추가합니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=us-central2-b \
   --worker=all \
   --command="sudo usermod -a -G docker $USER"
   

3. 모든 TPU 워커의 컨테이너에서 학습 스크립트를 실행합니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"
   

   Docker 명령어 플래그:

   • --rm은 프로세스가 종료된 후 컨테이너를 삭제합니다.
   • --privileged는 TPU 기기를 컨테이너에 노출합니다.
   • --net=host는 모든 컨테이너 포트를 TPU VM에 바인딩하여 포드의 호스트 간 통신을 허용합니다.
   • -e는 환경 변수를 설정합니다.

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

다음 명령어를 사용하여 TPU VM을 삭제합니다.

$ gcloud compute tpus tpu-vm delete your-tpu-name \
  --zone=us-central2-b

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-ubuntu2204-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] -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. exit를 입력하여 TPU VM을 종료합니다.

3. TPU VM을 삭제합니다.

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

TPU Pod

TPU Pod에서 JAX 코드를 실행하는 경우 모든 TPU 워커에서 JAX 코드를 동시에 실행해야 합니다. 이를 위한 한 가지 방법은 gcloud compute tpus tpu-vm ssh 명령어를 --worker=all 및 --command 플래그와 함께 사용하는 것입니다. 다음 절차에서는 각 TPU 워커를 더 쉽게 설정할 수 있도록 Docker 이미지를 만드는 방법을 보여줍니다.

1. 현재 디렉터리에 이름이 Dockerfile인 파일을 만들고 다음 텍스트를 붙여넣습니다.

   FROM python:3.8
   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
   

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

   docker build -t your-image-name .
   

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

   docker tag your-image-name europe-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
   

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

   docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
   

5. TPU VM을 만듭니다.

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

6. 모든 TPU 워커에서 Artifact Registry의 Docker 이미지를 가져옵니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="sudo usermod -a -G docker ${USER}"
   

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"
   

7. 모든 TPU 워커에서 컨테이너를 실행합니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   zone=europe-west4-a \
   --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image:your-tag bash"
   

8. 모든 TPU 워커에서 학습 스크립트를 실행합니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5"
   

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

1. 모든 작업자의 컨테이너를 종료합니다.

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker kill your-container-name"
   

2. 다음 명령어를 사용하여 TPU VM을 삭제합니다.

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

다음 단계

• Cloud TPU 튜토리얼
• TPU 관리
• Cloud TPU 시스템 아키텍처
• TPU Pod 슬라이스에서 TensorFlow 코드 실행
• TPU Pod 슬라이스에서 JAX 코드 실행