Cloud TPU에서 PyTorch를 사용하여 DLRM 학습


이 튜토리얼에서는 Cloud TPU에서 Facebook Research DLRM을 학습시키는 방법을 설명합니다.

목표

  • PyTorch 환경 만들기 및 구성
  • 허위 데이터로 학습 작업 실행
  • (선택사항) Criteo Kaggle 데이터 세트에서 학습

비용

이 문서에서는 비용이 청구될 수 있는 다음과 같은 Google Cloud 구성요소를 사용합니다.

  • Compute Engine
  • Cloud TPU

프로젝트 사용량을 기준으로 예상 비용을 산출하려면 가격 계산기를 사용하세요. Google Cloud를 처음 사용하는 사용자는 무료 체험판을 사용할 수 있습니다.

시작하기 전에

이 튜토리얼을 시작하기 전에 Google Cloud 프로젝트가 올바르게 설정되었는지 확인하세요.

  1. Google Cloud 계정에 로그인합니다. Google Cloud를 처음 사용하는 경우 계정을 만들고 Google 제품의 실제 성능을 평가해 보세요. 신규 고객에게는 워크로드를 실행, 테스트, 배포하는 데 사용할 수 있는 $300의 무료 크레딧이 제공됩니다.
  2. Google Cloud Console의 프로젝트 선택기 페이지에서 Google Cloud 프로젝트를 선택하거나 만듭니다.

    프로젝트 선택기로 이동

  3. Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

  4. Google Cloud Console의 프로젝트 선택기 페이지에서 Google Cloud 프로젝트를 선택하거나 만듭니다.

    프로젝트 선택기로 이동

  5. Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다.

  6. 이 둘러보기에서는 비용이 청구될 수 있는 Google Cloud 구성요소를 사용합니다. 예상 비용은 Cloud TPU 가격 책정 페이지에서 확인하세요. 리소스 사용을 마쳤으면 불필요한 비용이 청구되지 않도록 생성한 리소스를 삭제하세요.

Compute Engine 인스턴스 설정

  1. Cloud Shell 창을 엽니다.

    Cloud Shell 열기

  2. 프로젝트 ID의 변수를 만듭니다.

    export PROJECT_ID=project-id
    
  3. Cloud TPU를 만들려는 위치의 프로젝트가 사용되도록 Google Cloud CLI를 구성합니다.

    gcloud config set project ${PROJECT_ID}
    

    새 Cloud Shell VM에서 이 명령어를 처음 실행하면 Authorize Cloud Shell 페이지가 표시됩니다. 페이지 하단에 있는 Authorize를 클릭하여 gcloud에서 사용자 인증 정보로 Google Cloud API를 호출하도록 허용합니다.

  4. Cloud Shell에서 이 튜토리얼에 필요한 Compute Engine 리소스를 실행합니다. 참고: Criteo Kaggle 데이터 세트에서 학습하는 경우 n1-highmem-96 machine-type을 사용할 수 있습니다.

    gcloud compute instances create dlrm-tutorial \
    --zone=us-central1-a \
    --machine-type=n1-standard-64 \
    --image-family=torch-xla \
    --image-project=ml-images  \
    --boot-disk-size=200GB \
    --scopes=https://www.googleapis.com/auth/cloud-platform
    
  5. 새 Compute Engine 인스턴스에 연결합니다.

    gcloud compute ssh dlrm-tutorial --zone=us-central1-a
    

Cloud TPU 리소스 시작

  1. Compute Engine 가상 머신에서 다음 명령어를 사용하여 Cloud TPU 리소스를 시작합니다.

    (vm) $ gcloud compute tpus create dlrm-tutorial \
    --zone=us-central1-a \
    --network=default \
    --version=pytorch-2.0  \
    --accelerator-type=v3-8
    
  2. Cloud TPU 리소스의 IP 주소를 식별합니다.

    (vm) $ gcloud compute tpus describe dlrm-tutorial --zone=us-central1-a
    

PyTorch 환경 만들기 및 구성

  1. conda 환경을 시작합니다.

    (vm) $ conda activate torch-xla-2.0
    
  2. Cloud TPU 리소스의 환경 변수를 구성합니다.

    (vm) $ export TPU_IP_ADDRESS=ip-address
    
    (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    

허위 데이터로 학습 작업 실행

  1. 종속 항목을 설치합니다.

    (vm) $ pip install onnx
    
  2. 임의 데이터에서 모델을 실행합니다. 5~10분 정도 걸립니다.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
        --arch-embedding-size=1000000-1000000-1000000-1000000-1000000-1000000-1000000-1000000 \
        --arch-sparse-feature-size=64 \
        --arch-mlp-bot=512-512-64 \
        --arch-mlp-top=1024-1024-1024-1 \
        --arch-interaction-op=dot \
        --lr-num-warmup-steps=10 \
        --lr-decay-start-step=10 \
        --mini-batch-size=2048 \
        --num-batches=1000 \
        --data-generation='random' \
        --numpy-rand-seed=727 \
        --print-time \
        --print-freq=100 \
        --num-indices-per-lookup=100 \
        --use-tpu \
        --num-indices-per-lookup-fixed \
        --tpu-model-parallel-group-len=8 \
        --tpu-metrics-debug \
        --tpu-cores=8
    

(선택사항) Criteo Kaggle 데이터 세트에서 학습

이 단계는 선택사항입니다. Criteo Kaggle 데이터 세트에서 학습시키려는 경우에만 이 단계를 수행해야 합니다.

  1. 데이터 세트를 다운로드합니다.

    여기의 안내에 따라 Criteo Kaggle 데이터 세트에서 데이터 세트를 다운로드합니다. 다운로드가 완료되면 ./criteo-kaggle/ 디렉터리에 dac.tar.gz 파일을 복사합니다. tar -xzvf 명령어를 사용하여 ./critero-kaggle 디렉터리에서 tar.gz 파일의 콘텐츠를 추출합니다.

     (vm) $ mkdir criteo-kaggle
     (vm) $ cd criteo-kaggle
     (vm) $ # Download dataset from above link here.
     (vm) $ tar -xzvf dac.tar.gz
     (vm) $ cd ..
    
  2. 데이터 세트를 사전 처리합니다.

    Criteo 데이터 세트를 사전 처리하려면 이 스크립트를 시작합니다. 이 스크립트는 kaggleAdDisplayChallenge_processed.npz라는 파일을 생성하고 데이터 세트를 사전 처리하는 데 3시간 이상 걸립니다.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_data_pytorch.py \
        --data-generation=dataset \
        --data-set=kaggle \
        --raw-data-file=criteo-kaggle/train.txt \
        --mini-batch-size=128 \
        --memory-map \
        --test-mini-batch-size=16384 \
        --test-num-workers=4
    
  3. 사전 처리가 성공했는지 확인합니다.

    criteo-kaggle 디렉터리에 kaggleAdDisplayChallenge_processed.npz 파일이 표시됩니다.

  4. 사전 처리된 Criteo Kaggle 데이터 세트에서 학습 스크립트를 실행합니다.

    (vm) $ python /usr/share/torch-xla-2.0/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
        --arch-sparse-feature-size=16 \
        --arch-mlp-bot="13-512-256-64-16" \
        --arch-mlp-top="512-256-1" \
        --data-generation=dataset \
        --data-set=kaggle \
        --raw-data-file=criteo-kaggle/train.txt \
        --processed-data-file=criteo-kaggle/kaggleAdDisplayChallenge_processed.npz \
        --loss-function=bce \
        --round-targets=True \
        --learning-rate=0.1 \
        --mini-batch-size=128 \
        --print-freq=1024 \
        --print-time \
        --test-mini-batch-size=16384 \
        --test-num-workers=4 \
        --memory-map \
        --test-freq=101376 \
        --use-tpu \
        --num-indices-per-lookup=1 \
        --num-indices-per-lookup-fixed \
        --tpu-model-parallel-group-len 8 \
        --tpu-metrics-debug \
        --tpu-cores=8
    

    학습은 78.75% 이상의 정확도로 2시간 이후에 완료되어야 합니다.

삭제

만든 리소스를 사용한 후에는 계정에 불필요한 비용이 청구되지 않도록 삭제를 수행하세요.

  1. Compute Engine 인스턴스에서 연결을 해제합니다.

    (vm) $ exit
    

    프롬프트가 user@projectname으로 바뀌면 Cloud Shell에 있는 것입니다.

  2. Cloud Shell에서 Google Cloud CLI를 사용하여 Compute Engine 인스턴스를 삭제합니다.

    $ gcloud compute instances delete dlrm-tutorial --zone=us-central1-a
    
  3. Google Cloud CLI를 사용하여 Cloud TPU 리소스를 삭제합니다.

    $ gcloud compute tpus delete dlrm-tutorial --zone=us-central1-a
    

다음 단계

PyTorch Colab 사용: