PyTorch XLA 성능 프로파일링

개요

이 가이드에서는 PyTorch에서 Cloud TPU 성능 도구 및 측정항목 자동 분석 기능을 사용하는 방법을 설명합니다. 이러한 도구는 학습 워크로드 성능을 디버깅하고 최적화하는 데 도움이 됩니다.

개념

PyTorch / XLA를 처음 사용하는 경우 PyTorch API_GUIDE문제 해결 문서를 참조하세요. Cloud TPU의 경우 개념 문서를 참조하세요.

TPU 노드 + PyTorch/XLA 프로파일링

필요한 Cloud TPU 관련 리소스를 만들고 초기화합니다.

  1. TPU 리소스에 사용할 프로젝트 ID, Cloud Storage 버킷, 영역의 변수를 만듭니다.

    export PROJECT_ID=PROJECT_ID
    export BUCKET_NAME=BUCKET_NAME
    export ZONE=ZONE
    
    gcloud --project=$PROJECT_ID compute project-info add-metadata \
    --metadata BUCKET_NAME=$BUCKET_NAME
    
  2. Compute Engine VM 인스턴스 만들기 모든 Python 스크립트와 모델이 여기에 저장됩니다.

    gcloud compute instances create profiler-tutorial-vm \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --machine-type=n1-standard-16 \
      --image-project=ml-images \
      --image-family=torch-xla \
      --boot-disk-size=300GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    
  3. TPU 리소스를 만듭니다.

    gcloud compute tpus create profiler-tutorial-tpu \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --network=default \
     --version=pytorch-1.8 \
     --accelerator-type=v3-8
    
  4. Cloud Storage 버킷을 만듭니다.

    gsutil CLI가 아직 설치되어 있지 않으면 먼저 설치합니다. 설치 안내를 참조하세요.

  5. gsutil mb 명령어를 사용하여 모든 프로파일링 아티팩트가 저장되는 Cloud Storage 버킷을 만듭니다. 리전 및 버킷 이름 변수를 학습에 사용할 값으로 바꿉니다.

    gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}
    

    각 항목의 의미는 다음과 같습니다.

    • REGION은 Cloud TPU를 만든 리전입니다(예: europe-west4).
  6. Cloud TPU 프로젝트의 서비스 계정을 만듭니다.

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    이 명령어는 다음 형식의 Cloud TPU 서비스 계정을 반환합니다.

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

    예를 들면 다음과 같습니다. service-164006649440@cloud-tpu.iam.gserviceaccount.com

  7. 서비스 계정을 내보내고 스토리지 버킷에 대한 서비스 계정 권한을 부여합니다. account-number를 서비스 계정 생성 출력에서 반환된 PROJECT_NUMBER로 바꿉니다.

    export SERVICE_ACCOUNT=service-ACCOUNT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME}
    gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}
    

텐서보드 설정

  1. ssh를 통해 VM의 VM 전달 포트 9001을 로컬 머신의 9001 포트에 연결합니다. 이 포트는 로컬 브라우저에서 텐서보드 UI를 여는 데 사용됩니다.

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
  2. 텐서보드 설치 전용 conda 환경을 만듭니다.

    conda create -y -n tensorboard python=3.6
    conda activate tensorboard
    pip install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
  3. Compute Engine VM에서 텐서보드 서버를 실행한 다음 로컬 머신에서 http://localhost:9001/#profile을 방문하여 서버에 연결을 시도하여 설치를 테스트합니다.

     # Get bucket name
     BUCKET_NAME=$(curl "http://metadata.google.internal/computeMetadata/v1/project/attributes/BUCKET_NAME" -H "Metadata-Flavor: Google")
    
     tensorboard --logdir gs://${BUCKET_NAME} --port 9001
    

로컬 머신에서 http://localhost:9001/#profile을 방문하면 다음과 같이 표시됩니다.

이미지

http://localhost:9001을 방문한 경우 오른쪽 상단의 업로드 버튼 옆에 있는 드롭다운에서 프로필 옵션을 선택하여 위 프로필 페이지에 액세스할 수 있습니다.

이미지

모델 프로파일링

텐서보드 서버를 활성 상태로 유지하려면 로컬 머신에서 새 터미널 창을 시작하고 ssh를 통해 GCE VM에 다시 연결합니다(이번에는 포트 전달에 -L옵션을 사용하지 않음).

  1. 새 터미널 창에서 프로젝트 ID, 스토리지 버킷 환경, 영역 변수를 다시 내보냅니다. 새로운 셸이기 때문입니다.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. ssh를 통해 VM에 연결합니다.

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
    conda activate torch-xla-1.8
    PROJECT_ID=$(curl "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
    export TPU_IP_ADDRESS=$(gcloud compute tpus describe profiler-tutorial-tpu --zone=${ZONE} --project=${PROJECT_ID} \
    --format="value(ipAddress)")
    echo TPU_IP_ADDRESS=${TPU_IP_ADDRESS}
    export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    
  3. 통합 테스트가 환경에서 엔드 투 엔드 방식으로 작동하는지 확인합니다.

    python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
    >
  4. 학습을 시작하기 전에 /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py에서 다음 줄을 수정하세요.

    다음 줄을 변경합니다.

    accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    수신:
    accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    train_mnist에서 위의 두 인수는 동적 그래프 및 텐서 가져오기를 인위적으로 유발하여 나중에 자동 측정항목 분석 섹션에서 살펴봅니다. 지금은 TPU만 프로파일링하므로 다음 예시는 정격 성능으로 실행됩니다.

  5. 서버 프로파일링에 사용되는 학습 실행을 시작합니다.

    PT_XLA_DEBUG=1 XLA_HLO_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU (서버) 프로파일링

학습이 실행되면 http://localhost:9001/#profile을 방문하고 다음 안내에 따라 프로필을 캡처합니다.

이미지

다음 페이지가 자동으로 새로고침됩니다.

이미지

지원되는 도구는 왼쪽 창의 도구 드롭다운 아래에 표시됩니다.

  • 개요 페이지(TensorFlow/TPU와 PyTorch / XLA TPU는 근본적으로 다른 입력 파이프라인 분석기 표시는 포함하지 않음)
  • 메모리 뷰어
  • Op Profile
  • pod 뷰어
  • TensorFlow 통계(프레임워크 수준 통계(예: PyTorch 통계))
  • Trace 뷰어(Chrome 브라우저 필요)

개요 페이지

이 페이지에서는 캡처된 프로필의 개요를 보여줍니다. 이 예시에서는 MNIST 데이터 세트에서 작은 모델을 학습시키므로 매우 유휴 시간이 높습니다.

메모리 뷰어

텐서 및 HLO PyTorch 작업당 사용되는 기기 메모리(HBM)를 표시합니다. 메모리 뷰어는 HLO 모듈별 메모리 보기를 캡처하므로 기기 데이터 할당(입력 및 라벨 일괄 전송) 그래프와 같은 모듈이 있습니다. 특정 HLO 모듈의 메모리 사용량을 보려면 왼쪽의 호스트 드롭다운에서 모듈을 선택합니다.

이미지

선택한 HLO 모듈을 표시하면 모듈 실행 HBM 사용 공간 타임라인을 전체적으로 볼 수 있습니다. 할당 크기, 프로그램 실행 크기, 패딩 크기에 따라 정렬됩니다.

이미지

그런 다음 버퍼 할당 위로 마우스를 가져가면 각 버퍼 할당을 추가로 검사할 수 있습니다. 예를 들어 기기 HBM을 가장 많이 차지하는 할당을 확인하는 방법은 다음과 같습니다.

이미지

위의 예시에서 (1)은 사용자 코드가 추가한 torch_xla.debug.profiler.Trace 주석에 해당합니다. 이 줄에 해당하는 test_profile_mp_mnist.py 코드를 검사합니다.

   class MNIST(nn.Module):
   ...
     def forward(self, x):
       with xp.Trace('conv1'):
         x = F.relu(F.max_pool2d(self.conv1(x), 2))
         x = self.bn1(x)
   ...
   

또한 test_mnist 작업 네임스페이스에서 이 HLO 모듈은 xp.trace('test_mnist') 컨텍스트 관리자가 있으므로 eval 루프에 해당함을 알 수 있습니다.

XLA 작업 프로필

작업 프로필은 프로파일링 기간 동안 실행된 XLA 작업의 성능 통계를 표시하는 Cloud TPU 도구입니다. 작업 프로필에 표시되는 내용은 다음과 같습니다.

  • 애플리케이션에서 Cloud TPU를 얼마나 잘 사용하는가(카테고리별로 작업에 소요된 시간의 비율 및 TPU FLOPS 사용률)
  • 가장 많은 시간을 소비하는 작업. 이러한 작업은 잠재적 최적화 대상입니다.
  • 작업을 사용하는 형태, 패딩, 표현식을 포함한 각 작업의 세부정보

작업 프로필을 사용하여 적절한 최적화 대상을 찾을 수 있습니다. 예를 들어 모델의 TPU 최고 FLOPS가 5%에 불과한 경우 이 도구를 사용하여 실행에 가장 많은 시간이 소비되는 XLA 작업과 이러한 작업이 소비하는 TPU FLOPS가 얼만큼인지를 식별할 수 있습니다.

이미지

각 항목에 대한 설명:

  1. 개요 섹션. Cloud TPU 사용률을 표시하고 최적화 방안을 제안합니다.
  2. 제어판. 테이블에 표시되는 작업 수, 표시되는 작업, 정렬 방법을 설정할 수 있는 컨트롤을 포함합니다.
  3. 작업 테이블. XLA 작업과 관련된 상위 TensorFlow 작업 카테고리를 나열하는 테이블입니다. Cloud TPU 사용량 비율에 따라 작업이 정렬됩니다.
  4. 작업 세부정보 카드 테이블의 작업으로 마우스를 가져갈 때 표시되는 작업에 대한 세부정보. FLOPS 사용률, 작업이 사용된 표현식, 작업 레이아웃(핏)이 포함됩니다.

pod 뷰어

pod 뷰어 도구에 대한 자세한 설명은 TPU 도구를 참조하세요.

프레임워크 통계(TensorFlow/PyTorch 통계)

프레임워크 통계는 TPU 기기 및 TPU 호스트에서 실행되는 상세한 PyTorch 및 XRT 작업 통계 분석을 제공합니다.

Trace 뷰어

Trace 뷰어는 Cloud TPU 성능 분석 도구입니다. Trace 뷰어는 Chrome trace 이벤트 프로파일링 뷰어를 사용하므로 Chrome 브라우저를 사용해야 합니다.

Trace 뷰어에 표시되는 타임라인은 다음과 같은 정보를 제공합니다.

  • TensorFlow 모델에서 실행한 작업 기간
  • 작업을 실행한 시스템의 부분(TPU 또는 호스트 머신). 일반적으로 PyTorch/XLA의 경우 호스트 머신이 주로 컴파일 및 버퍼 할당/할당 해제에서 작동하는 반면 TPU는 실제 모델 학습을 실행합니다.
  • Trace 뷰어를 사용하면 모델의 성능 문제를 파악한 후 해결을 위한 조치를 취할 수 있습니다. 상세히 살펴보면 실행하는 데 가장 오래 걸리는 PyTorch / XLA 작업을 파악할 수 있습니다.

xp.Trace(NAME) 주석을 추가하여 모델의 특정 부분이 실행되는 데 걸리는 시간을 측정하기 위해 trace를 직접 추가할 수 있습니다. 예를 들어 다음 trace에 표시되는 내용입니다.

이미지

  1. test_profile_mp_mnist.py의 모델 코드에서 명시적인 사용자 주석으로 생성됩니다.
  2. 실행된 PyTorch 작업(사전 감소)
  3. PyTorch/XLA 자동 생성 HLO 모듈 이름
  4. 기기에서 실행 XLA 작업(통합)

자세한 내용은 Trace 뷰어에 대한 일반 TPU 문서를 참조하세요. 입력 파이프라인 및 다른 TensorFlow 관련 부분 섹션은 이 문서의 내용과 관련이 없으므로 무시합니다.

PyTorch/XLA 클라이언트 프로파일링

모델 실행이 진행되는 동안 TPU 측을 프로파일링할 때와 비슷하게 이제 학습 중에 PyTorch/XLA 클라이언트 측을 프로파일링합니다. 클라이언트 측에서 사용되는 주요 모니터링 도구는 Trace 뷰어입니다.

학습 스크립트에서 프로파일링 서버를 시작해야 합니다. 예시는 텐서보드에서 쿼리하여 trace를 캡처할 수 있는 항목을 참조하세요.

이미지

여러 프로세스의 trace를 캡처하기 위해 각 프로세스는 다른 포트에서 서버 프로파일링(예: 기본 포트 번호에 xm.get_ordinal() 추가)을 한 다음 ','로 연결된 localhost:port 목록을 제공하기 시작합니다. 텐서보드는 한 번에 여러 프로세스의 trace를 보는 것을 지원하지 않으므로 각 프로세스에 서로 다른 호스트 드롭다운이 표시됩니다.

다음 다이어그램은 샘플 trace를 보여줍니다.

이미지

TPU trace에 서로 다른 네임스페이스 trace를 추가하는 방법과 마찬가지로 동일한 API를 사용하여 클라이언트 측 trace(xp.Trace(NAME))에 추가할 수 있습니다. 이 모델은 크기가 작고 작은 MNIST 이미지와 함께 사용되므로 단계 시간은 짧으며 항상 균일하지는 않습니다. 연습으로 trace를 추가하고 test_train_mp_imagenet.py --fake_data의 예시와 유사한 프로파일러 서버를 시작하여 ResNet50의 trace를 가져올 수 있습니다.

trace에는 검사할 수 있는 추가 메타데이터가 있습니다. 예를 들어 TransferToServer 및 TransferFromServer trace는 전송 및 수신되는 정확한 텐서 수와 총 크기를 표시합니다.

이미지

XLA 그래프 컴파일의 경우 문제 진단에 도움이 될 수 있는 그래프 해시를 확인할 수 있습니다.

이미지

또한 텐서보드 UI를 통해 프로파일링하는 대신 PyTorch / XLA: xp.trace()에서 TPU와 클라이언트를 프로그래매틱 방식으로 프로파일링하기 위한 API도 제공합니다.

자동 측정항목 분석

이 섹션에서는 디버깅 모드를 사용하여 다음과 같은 성능 문제를 감지하는 방법을 보여줍니다.

  • 동적 그래프/연속 컴파일
  • 매우 느린 그래프 컴파일
  • 매우 느린 지연 실행
  • 빈번한 XLA→CPU 전송
  • RAM 교체를 호스트하는 반복된 기기 HBM
  • 반복된 HBM 조각 모음
  • 감소되지 않은 aten:: 작업

학습을 시작하기 전에 /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py에서 다음 줄을 되돌립니다.

다음 줄을 변경합니다.

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   
수신:
   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

이러한 변경사항은 컴파일 및 텐서 가져오기를 인위적으로 유발합니다. dynamic_graph=True는 각 단계의 배치 크기를 인위적으로 변경하여 XLA 감소 그래프를 각각의 단계와 재컴파일 시 다르게 변경합니다. fetch_often=True는 모든 단계에서 loss.item() 호출을 삽입하여 각 단계에서 기기의 텐서 값을 가져와 성능이 저하됩니다.

학습 스크립트 예시 실행:

   PT_XLA_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --fake_data --num_cores=1
   

디버깅 시 디버깅 프로세스를 간소화하므로 --num_cores=1을 사용하여 실행하는 것이 좋습니다. 샘플 출력은 다음과 같이 표시됩니다.

Epoch 1 train begin 01:18:05
| Training Device=xla:1/0 Step=0 Loss=0.00000 Rate=1905.00 GlobalRate=1904.72 Time=01:18:05
pt-xla-profiler: TransferFromServerTime too frequent: 3 counts during 3 steps
pt-xla-profiler: TransferFromServerTime too frequent: 4 counts during 4 steps
pt-xla-profiler: TransferFromServerTime too frequent: 5 counts during 5 steps
pt-xla-profiler: TransferFromServerTime too frequent: 6 counts during 6 steps
pt-xla-profiler: TransferFromServerTime too frequent: 7 counts during 7 steps
pt-xla-profiler: TransferFromServerTime too frequent: 8 counts during 8 steps
pt-xla-profiler: TransferFromServerTime too frequent: 9 counts during 9 steps
pt-xla-profiler: TransferFromServerTime too frequent: 10 counts during 10 steps
pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps

pt-xla-profiler 프리픽스가 있는 줄은 자동 측정항목 분석 출력에 해당합니다. 이 예시에서는 TransferFromServerTime이 단계당 한 번으로 너무 자주 표시된 것을 볼 수 있습니다. 이것은 모든 단계에서 loss.item() 값을 가져오는 학습 루프 때문입니다. 또한 그래프의 동적 모양으로 인해 모델을 반복해서 다시 컴파일할 때 CompileTime이 너무 빈번함 경고가 표시될 수 있습니다. 이러한 유형의 문제를 일으키는 코드 스니펫은 다음과 같습니다. test_profile_mp_mnist.py

    for step, (data, target) in enumerate(loader):
      if dynamic_graph:
        # The batch dimension is different every step.
        index = max(-step, -flags.batch_size + 1)  # non-empty
        data, target = data[:-index, :, :, :], target[:-index]
      ...
      if fetch_often:
        # Fetch tensor value from XLA:TPU to CPU every step.
        loss_i = loss.item()

그런 다음 학습 스크립트에서 Ctrl^C를 사용할 수 있고 마지막에 감소되지 않은 작업 요약이 표시됩니다. aten::_local_scalar_dense는 XLA 텐서를 CPU 컨텍스트에 다시 가져오는 데 해당하는 특수 작업입니다.

이 보고서에는 aten::_local_scalar_dense 작업이 호출되는 두 개의 주요 위치가 있으며, 둘 다 loss.item() 소스 코드와 일치합니다.

  • test/test_profile_mp_mnist.py:158
  • test/test_profile_mp_mnist.py:61
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=27):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:158)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: FRAME (count=2):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   _train_update (test/test_profile_mp_mnist.py:61)
pt-xla-profiler:    (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:700)
pt-xla-profiler:   _run_step_closures (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:709)
pt-xla-profiler:   mark_step (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:723)
pt-xla-profiler:   __exit__ (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/debug/profiler.py:153)
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:162)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

이제 아래 스크립트에서 자동 측정항목 분석을 실행합니다. 여기에는 감소되지 않은 작업 _ctc_loss가 포함됩니다(1.8 출시 버전 기준).

PT_XLA_DEBUG=1 python <<EOF
import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t = torch.randn(50, 16, 20).log_softmax(2).to(dev)
target = torch.randint(low=1, high=20, size=(16, 30), dtype=torch.long).to(dev)
input_lengths = torch.full(size=(16,), fill_value=50, dtype=torch.long).to(dev)
target_lengths = torch.randint(low=10, high=30, size=(16,), dtype=torch.long).to(dev)

for _ in range(10):
  loss = torch.nn.CTCLoss()(t, target, input_lengths, target_lengths)
  xm.mark_step()
EOF

PT_XLA_DEBUG=1로 위의 스크립트를 실행하면 출력은 다음과 비슷합니다.

…
pt-xla-profiler: TransferFromServerTime too frequent: 30 counts during 10 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=10):
pt-xla-profiler: Unlowered Op: "_ctc_loss"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   ctc_loss (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/functional.py:2305)
pt-xla-profiler:   forward (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/loss.py:1593)
pt-xla-profiler:   _call_impl (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/module.py:889)
pt-xla-profiler:    (:11)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

자동 측정항목 분석기로 인해 STDIN의 11번째 줄(예: torch.nn.CTCLoss()가 있는 줄)이 이와 같이 감소되지 않은 작업을 유발하고 있습니다. 현재 ctc_loss 작업은 감소하지 않았기 때문에 위의 보고서가 표시됩니다. 또한 실행 전에 텐서가 처음에는 XLA:TPU에서 있기 때문에 TransferFromServerTime에 관한 일부 경고가 표시될 수도 있지만, 작업이 감소하지 않으므로 먼저 XLA 텐서를 CPU로 다시 전송하고 CPU에서 aten:: 작업을 실행하고 다시 전송해야 합니다.

대신 pt-xla-profiler 출력을 파일에 쓰려면 PT_XLA_DEBUG=1PT_XLA_DEBUG_FILE=$PATH_TO_FILE을 설정합니다.

삭제

VM을 종료한 후 다음 명령어를 실행하여 TPU, VM, Cloud Storage 버킷을 삭제합니다.

(vm)$ exit
gcloud compute instances delete profiler-tutorial-vm \
  --zone=${ZONE} \
  --project=${PROJECT_ID}
gcloud compute tpus delete profiler-tutorial-tpu \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --async
gsutil rm -fr gs://${BUCKET_NAME}

TPU VM + PyTorch/XLA 프로파일링

이 섹션에서는 TPU VM 아키텍처를 사용하여 PyTorch/XLA를 프로파일링합니다.

환경 변수 내보내기

  1. 프로젝트 ID와 TPU 리소스에 사용할 영역의 변수를 만듭니다.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    

Cloud TPU 만들기

TPU VM 사용자 가이드를 참고하고 설정 후 torch, torch_xla, torchvision, tensorboard가 사전 설치된 상태로 제공되는 v3-8 TPU VM을 만드세요.

  1. TPU 리소스를 만듭니다.

    gcloud alpha compute tpus tpu-vm create profiler-tutorial-tpu-vm \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --version=v2-alpha \
     --accelerator-type=v3-8
    

텐서보드 서버 시작

  1. VM에 SSH로 연결하고 tensorboard-plugin-profile을 설치하고 텐서보드 서버를 시작합니다.

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
      pip3 install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
      tensorboard --logdir ./tensorboard --port 9001
    

로컬 머신의 http://localhost:9001에서 텐서보드 출력이 뜨면 다음과 같이 표시됩니다.

이미지

또한 http://localhost:9001에서 텐서보드 출력을 확인할 때 오른쪽 상단의 업로드 버튼 옆에 있는 드롭다운에서 프로필 옵션을 선택하여 위 프로필 페이지에 액세스할 수 있습니다.

이미지

모델 프로파일링

개발 환경의 터미널 창에서 위와 동일한 환경 변수 및 ssh를 TPU VM으로 내보냅니다.

  1. 새로운 셸이므로, 새 터미널 창에서 프로젝트 ID와 영역 변수를 다시 내보냅니다.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. ssh를 통해 VM에 연결합니다.

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
  3. PyTorch/XLA 저장소를 클론하고 e2e 테스트를 실행합니다.

      git clone -b r1.8 https://github.com/pytorch/xla
      export XRT_TPU_CONFIG="localservice;0;localhost:51011"
      python3 xla/test/test_profiler.py  # takes <1 min
    >
  4. 학습을 시작하기 전에 xla/test/test_profile_mp_mnist.py에서 다음 줄을 수정합니다.

    다음과 같이 변경하세요.

        accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    수신:
        accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    train_mnist에서 위의 두 인수는 동적 그래프 및 텐서 가져오기를 인위적으로 유발하여 나중에 자동 측정항목 분석 섹션에서 살펴봅니다. 지금은 TPU만 프로파일링하므로 다음 예시는 정격 성능으로 실행됩니다.

  5. 학습 실행을 시작합니다.

     XLA_HLO_DEBUG=1 python3 xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU + 클라이언트 프로파일링

학습이 실행되면 http://localhost:9001에서 텐서보드 출력을 보고 다음 안내에 따라 프로필을 캡처합니다.

이미지

그러면 다음 페이지가 다시 로드됩니다.

이미지

현재까지 TPU VM 설정에서는 trace 뷰어만 선택되므로 도구 드롭다운에서 trace_viewer를 선택하고 trace를 조사합니다. TPU VM 설정에서 '클라이언트' 측 및 TPU 기기 측 trace를 하나의 전체보기로 확인할 수 있습니다.

이미지

삭제

  1. VM을 종료한 후 다음 명령어를 실행하여 TPU, VM, Cloud Storage 버킷을 삭제합니다.

    (vm)$ exit
    

만든 TPU VM을 삭제합니다.

  1. Cloud TPU 및 Compute Engine 리소스를 삭제합니다.

    $ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} --zone=${ZONE}
    
  2. 다음 명령어를 실행하여 리소스가 삭제되었는지 확인합니다. 삭제하는 데 몇 분 정도 걸릴 수 있습니다. 다음과 같은 응답이 나타나면 인스턴스가 성공적으로 삭제되었다는 의미입니다.

    $ gcloud alpha compute tpus tpu-vm list --project ${PROJECT_ID} --zone=${ZONE}
    
    Listed 0 items.