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

이 가이드에서는 Cloud TPU에서 Facebook Research DLRM을 학습시키는 방법을 설명합니다.

목표

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

비용

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

  • Compute Engine
  • Cloud TPU

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

시작하기 전에

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

  1. Google 계정으로 로그인합니다.

    아직 계정이 없으면 새 계정을 등록하세요.

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

    프로젝트 선택기 페이지로 이동

  3. Google Cloud 프로젝트에 결제가 사용 설정되어 있는지 확인합니다. 프로젝트에 결제가 사용 설정되어 있는지 확인하는 방법을 알아보세요.

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

Compute Engine 인스턴스 설정

  1. Cloud Shell 창을 엽니다.

    Cloud Shell 열기

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

    export PROJECT_ID=project-id
    
  3. Cloud TPU를 만들려는 프로젝트를 사용하도록 gcloud 명령줄 도구를 구성합니다.

    gcloud config set project ${PROJECT_ID}
    
  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-1.6  \
    --accelerator-type=v3-8
    
  2. Cloud TPU 리소스의 IP 주소를 식별합니다.

    (vm) $ gcloud compute tpus list --zone=us-central1-a
    

PyTorch 환경 만들기 및 구성

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

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

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

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

  1. 다음을 실행하여 TPU 호환 DLRM을 가져옵니다.

    (vm) $ git clone --recursive https://github.com/pytorch-tpu/examples.git
    
  2. 종속 항목을 설치합니다.

    (vm) $ pip install onnx
    
  3. 임의 데이터에서 모델을 실행합니다.

    (vm) $ python 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 데이터 세트에서 학습

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

    여기의 안내에 따라 Criteo Kaggle 데이터 세트에서 데이터 세트를 다운로드합니다. 다운로드가 완료되면 dac.tar.gz 파일의 압축을 풀고 ./criteo-kaggle/ 디렉터리에 저장합니다.

     (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 파일이 생성되면 학습 실행을 종료하려 하므로 하나의 코어에서만 실행하고 있습니다. 사전 처리하는 데 시간이 오래 걸릴 수 있습니다(1시간 이상).

    (vm) $ python 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 \
        --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 \
        --test-freq=101376 \
        --use-tpu \
        --num-indices-per-lookup=1 \
        --num-indices-per-lookup-fixed \
        --tpu-model-parallel-group-len 1 \
        --tpu-metrics-debug \
        --tpu-cores=1
    

    다음과 같은 로그가 표시되면 데이터 세트가 사전 처리된 것이며 실행(Ctrl^C)을 종료하고 pkill python을 실행하여 유출된 프로세스가 없는지 확인할 수 있습니다.

    Load 6548659/6548660  Split: 0  Label True: 0  Stored: 0
    Saved /home/jysohn/test/train_day_0.npz!
    Load 6548659/6548660  Split: 1  Label True: 1  Stored: 1
    Saved /home/jysohn/test/train_day_1.npz!
    Load 3742283/6548660  Split: 2  Label True: 0  Stored: 0
    ...
    time/loss/accuracy (if enabled): 2020-08-26 13:17:18.139566
    Finished training it 1024/306968 of epoch 0, -1.00 ms/it, loss 0.518172, accuracy 75.720 %, 131072 samples, @ 2020-08-26 13:18:04.515570
    
  3. 사전 처리가 성공했는지 확인합니다.

    criteo-kaggle 디렉터리 내에 생성된 kaggleAdDisplayChallenge_processed.npz가 표시됩니다.

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

    (vm) $ python 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 \
        --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% 이상의 정확도로 3시간 이상 내에 완료되어야 합니다.

삭제

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

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

    (vm) $ exit
    

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

  2. Cloud Shell에서 gcloud 명령줄 도구를 사용하여 Compute Engine 인스턴스를 삭제합니다.

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

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

다음 단계

PyTorch Colab 사용: