Cloud TPU VM에서 모델 프로파일링

프로파일링은 Cloud TPU에서 모델의 성능을 최적화하는 주요 도구 중 하나입니다. 핵심 프로파일링 도구는 XProf라고 하며 OpenXLA/XProf GitHub 저장소에서 사용할 수 있습니다. XProf는 JAX, Pytorch XLA, Tensorflow/Keras를 비롯한 모든 XLA 기반 프레임워크의 프로파일링을 지원합니다.

프로필 캡처

XProf를 사용하려면 먼저 모델의 프로필을 캡처해야 합니다. 프로필을 캡처하는 방법에는 두 가지가 있습니다.

1. 프로그래매틱 캡처
2. 온디맨드 캡처(수동 캡처)

프로그래매틱 캡처를 사용하려면 코드에서 프로필을 캡처할 위치를 지정하기 위해 모델 코드에 주석을 달아야 합니다. 일반적으로 몇 가지 학습 단계의 프로필을 캡처하거나 모델 내에서 특정 코드 블록을 프로파일링합니다. JAX, Pytorch XLA, TensorFlow 등 다양한 프레임워크에서 트레이스를 캡처하는 방법에는 API 기반 트레이스 시작 또는 중지, 컨텍스트 관리자 기반 등 여러 가지가 있습니다. MaxText와 같은 상위 수준 프레임워크에서는 MaxText를 실행할 때 profiler=xplane 플래그를 사용 설정하기만 하면 프로필 캡처를 사용 설정할 수 있습니다.

애드혹 방식으로 프로필을 캡처하려는 경우이거나 프로그래매틱 프로필 캡처를 사용 설정하지 않은 경우에는 온디맨드 프로필 캡처를 사용할 수 있습니다. 이는 실행 중에 모델 측정항목에 문제가 표시되었을 때 해당 문제를 진단하기 위해 특정 기간 동안 해당 시점의 프로필을 캡처하려는 경우에 유용합니다.

온디맨드 프로필 캡처를 사용 설정하려면 코드 내에서 XProf 서버를 시작해야 합니다. profiler.start_server를 사용 설정하면 프로필 캡처를 시작하기 위해 온디맨드 캡처 트리거를 수신 대기하는 XProf 서버가 ML 워크로드에서 시작됩니다. Tensorboard UI를 통해 또는 XProfiler 도구를 사용하여 CLI를 통해 온디맨드 프로필 캡처를 트리거할 수 있습니다.

다양한 프레임워크에 대해 프로그래매틱 및 온디맨드 캡처를 사용 설정하는 방법은 다음을 참고하세요.

• JAX
• Tensorflow

프로그래매틱 프로필 캡처와 온디맨드 프로필 캡처 모두 캡처된 프로필을 저장할 위치를 지정해야 합니다. /profiles/run-name과 유사한 경로를 사용하여 TPU VM의 디렉터리에 프로필을 저장하거나 gs://bucket-name/run-name/과 유사한 경로를 사용하여 Cloud Storage를 사용할 수 있습니다.

위에서 언급한 경로를 사용하면 프로필이 다음 경로에 캡처됩니다.

/profiles/run-name/plugins/profile/session1/

또는

gs://bucket-name/run-name/plugins/profile/session1/

세션은 모델 실행 중에 프로필을 한 번 캡처한 것을 나타냅니다. 예를 들어 1~3단계에서 학습 실행의 프로필을 캡처한 후 8~10단계에서 프로필을 캡처하면 동일한 실행의 프로필이지만 1~3단계의 첫 번째 캡처는 session1이고 8~10단계의 두 번째 캡처는 session2입니다.

각 실행 아래에는 다른 날짜 스탬프가 표시되어 세션이 구분됩니다. 프로그래매틱 방식, 온디맨드 방식 또는 두 가지를 혼합하여 여러 세션에서 프로필을 캡처할 수 있습니다. 모델 학습의 여러 부분에서 프로필을 비교하려는 경우에 유용합니다. 예를 들어 학습 시작 시의 프로필과 학습 종료 시의 프로필을 비교합니다.

Tensorboard로 프로필 시각화

Tensorboard에서 XProf 프로필을 보려면 tensorboard-plugin-profile 플러그인을 설치해야 합니다.

캡처한 프로필을 Google Cloud에서 시각화하려면 TPU VM에서 Tensorboard를 사용하거나 XProfiler 도구를 사용하여 VM 또는 GKE 포드에서 Tensorboard를 호스팅하는 것이 좋습니다.

TPU VM에서 프로필을 캡처한 경우 해당 TPU VM에서 Tensorboard를 사용하여 캡처된 프로필을 볼 수 있습니다.

TPU VM에 XProf와 TensorBoard를 설치하려면 다음 명령어를 실행하면 됩니다.

pip install tensorboard_plugin_profile tensorboard
pip install xprof

Tensorboard로 캡처한 프로필을 보려면 다음을 실행하세요.

$ tensorboard --logdir=profiles/run-name

또는

$ tensorboard --logdir=profiles

여러 실행에서 여러 프로필을 로드하려면 Tensorboard를 특정 실행의 프로필이 있는 하위 디렉터리가 아닌 모든 실행과 프로필이 포함된 루트 디렉터리로 지정하세요.

Google Cloud에서 Tensorboard 프로필을 시각화하는 Cloud-Diagnostics-XProf 라이브러리

cloud-diagnostics-xprof(XProfiler 라이브러리라고도 함)를 사용하면 Tensorboard를 호스팅하고 Google Cloud에서 프로필을 시각화하기가 더 쉬워집니다. 자세한 내용은 cloud-diagnostics-xprof GitHub 저장소를 참고하세요.

cloud-diagnostics-xprof 저장소는 로컬로 XProf 및 TensorBoard를 실행할 때에 비해 다음과 같은 개선점을 제공합니다.

• XProf 및 TensorBoard 종속 항목 설정 및 패키징
• 장기 보관 및 실행 후 분석에 유용할 수 있는 Cloud Storage에 프로필을 저장합니다(연구자가 실행을 완료하면 캡처된 로컬 프로필이 삭제됨).
• Compute Engine VM 또는 GKE 포드에 TensorBoard를 프로비저닝하여 대규모 프로필과 여러 프로필을 빠르게 로드하고, 로드 속도와 비용에 대한 사용자 요구사항에 따라 머신 유형을 변경하는 옵션 제공
• 프로필 공유 및 팀원과 Google 엔지니어와의 공동작업을 위한 링크 만들기
• GKE 및 Compute Engine에서 워크로드를 더 쉽게 온디맨드 프로파일링하여 워크로드를 실행하는 호스트를 선택하여 프로필을 캡처할 수 있습니다.

XProfiler 라이브러리를 사용하기 전에 프로그래매틱 방식으로 또는 프로필 서버를 시작하여 워크로드 코드의 프로필을 캡처해야 나중에 온디맨드 프로필 캡처를 실행할 수 있습니다.

XProfiler 라이브러리를 설정하려면 Google Cloud CLI와 Python 가상 환경을 설정해야 합니다. 그런 다음 다음을 실행하면 됩니다.

pip install cloud-diagnostics-xprof

이렇게 하면 필요한 모든 XProf 및 TensorBoard 종속 항목이 설치됩니다.

다음으로 Tensorboard를 호스팅할 VM 또는 GKE 포드를 만들려면 다음 명령어를 실행합니다.

xprofiler create -z $ZONE -l $GCS_PATH

또는

xprofiler create --GKE -z $ZONE -l $GCS_PATH

$ZONE을 영역으로, $GCS_PATH를 프로필 트레이스 경로로 바꿉니다. 여러 실행의 모든 프로필 트레이스가 포함된 루트 디렉터리나 한 실행의 특정 프로필 트레이스 집합을 지정할 수 있습니다.

예를 들어 다음 디렉터리에서 프로필을 캡처하는 경우

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

루트 디렉터리의 GCS 경로를 설정할 수 있습니다(gs://bucket-name).

xprofiler create -z $ZONE -l gs://bucket-name/

Tensorboard UI에는 모든 실행 및 세션의 프로필이 run1/session1, run1/session2, run2/session1로 표시됩니다.

기본적으로 xprofiler create는 Compute Engine VM, 특히 c4-highmem-8 VM을 만듭니다. -m 플래그를 사용하여 머신 유형을 변경할 수 있습니다. Compute Engine VM 대신 GKE 포드에 TensorBoard 인스턴스를 만들려면 –GKE 플래그를 xprofiler create에 전달하면 됩니다. GKE 포드에서 TensorBoard 인스턴스를 호스팅하면 GKE에 배포된 나머지 워크로드와 함께 TensorBoard를 더 쉽게 관리할 수 있습니다.

Tensorboard를 호스팅하는 VM 또는 GKE 포드는 로컬로 호스팅되는 Tensorboard보다 Google Cloud 에서 대규모 프로필과 여러 프로필을 훨씬 빠르게 로드합니다. 벤치마킹에 따르면 기본 c4-highmem-8 VM을 사용하여 처음 로드할 때 몇 분 이내에 1GB 로드 프로필이 표시됩니다. 또한 성능 및 비용 요구사항에 따라 적절한 머신 유형을 선택할 수 있습니다.

xprofiler create를 실행하면 다음과 비슷한 출력이 표시됩니다.

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

첫 번째 옵션은 Tensorboard에서 XProf 프로필을 볼 수 있는 링크입니다. 이 링크를 사용하면 팀과 공유할 수 있으며, Google Cloud의 성능 최적화를 지원하는 Google 엔지니어와도 공유할 수 있습니다.

프로필 데이터를 저장하는 Cloud Storage 버킷에 설정된 권한에 따라 링크에 액세스할 수 있는 사용자를 관리할 수 있습니다. 어떤 이유로 링크가 작동하지 않는 경우 xprofiler connect 명령어를 사용하여 프로필을 볼 수 있도록 SSH를 사용하여 TensorBoard 인스턴스에 연결하는 방법도 제공됩니다.

워크로드 코드에서 Cloud Profiler 서버를 사용 설정한 경우 (프로그래매틱 캡처와 온디맨드 캡처 모두에 필요) 온디맨드 프로파일링을 수행하려면 다음 두 가지 방법으로 수행할 수 있습니다.

a. TensorBoard에서 프로필 캡처 버튼을 클릭합니다. 이제 워크로드가 실행되는 기기 호스트를 선택하여 프로파일링할 수 있습니다. Google은 GKE 또는 Compute Engine에서 실행되는 워크로드의 온디맨드 캡처를 지원합니다.

b. 영역, Cloud Storage 버킷, 프레임워크, 호스트 VM/포드 이름, 캡처 기간(밀리초)과 같은 정보를 제공하여 XProfiler 캡처 명령어를 사용합니다. 이는 Tensorboard UI에 입력해야 하는 정보와 동일합니다.

cloud-diagnostics-xprof 라이브러리 사용 방법에 관한 자세한 내용은 이 GitHub 페이지를 참고하세요.