Trace 뷰어 도구

Trace 뷰어를 사용하여 프로파일링 세션 중에 발생한 이벤트의 타임라인을 시각화할 수 있습니다. 호스트 (CPU) 및 가속기 (GPU 또는 TPU)와 같은 시스템의 여러 부분에서 모델이 실행한 작업의 시간을 표시합니다. 이를 통해 모델이 하드웨어 리소스를 사용하는 방식을 파악하고, 성능 병목 현상을 식별하고, 더 빠른 실행을 위해 모델을 최적화할 수 있습니다. Trace 뷰어 UI는 chrome://tracing에서 사용되는 UI를 기반으로 하므로 Chrome 브라우저를 사용해야 합니다.

지원되는 플랫폼

TPU: 지원됨

GPU: 지원됨

GPU용 Trace 뷰어는 GPU 칩 및 내부 스트림별로 타임라인의 이벤트를 그룹화하지만 TPU의 경우 TPU 칩 및 코어별로 타임라인의 이벤트를 구성합니다. 두 경우 모두 Trace Viewer는 가속기가 연결된 호스트 CPU의 스레드별 타임라인에 이벤트도 표시합니다.

타임라인과 상호작용

트레이스 뷰어는 타임라인을 탐색하고 검사하는 데 사용할 수 있는 여러 도구와 방법을 제공합니다.

• 탐색: 다음 단축키를 사용할 수 있습니다.
  • W: 확대합니다.
  • S: 축소합니다.
  • A: 왼쪽으로 화면 이동
  • D: 오른쪽으로 화면 이동
• 도구 선택기: 도구 선택기에서 도구를 클릭하거나 상응하는 단축키를 사용할 수 있습니다.
  • 선택 도구 (1 또는 !): 이벤트를 클릭하여 선택하고 세부정보 창에서 세부정보를 확인합니다. Ctrl+클릭을 통해 여러 이벤트를 선택하여 요약을 확인합니다.
  • 화면 이동 도구 (2 또는 @): 드래그하여 타임라인을 가로 또는 세로로 이동합니다.
  • 확대/축소 도구 (3 또는 #): 드래그하여 타임라인의 특정 영역을 확대합니다.
  • 타이밍 도구 (4 또는 $): 드래그하여 시간 간격을 표시합니다. 표시된 간격의 길이가 표시됩니다. 'm' 키를 사용하여 선택한 부분을 표시하고 총 시간을 확인할 수도 있습니다.
• 선택한 이벤트로 확대 (f): 이벤트를 하나 이상 선택하고 'f' 키를 눌러 타임라인의 해당 부분을 빠르게 확대합니다. 이는 특정 학습 단계에 집중하는 데 유용합니다.

대량의 트레이스 이벤트를 처리할 때는 Trace Viewer가 스트리밍 모드로 작동합니다. 즉, 지도 애플리케이션이 작동하는 것처럼 타임라인을 화면 이동하고 확대/축소할 때 데이터가 필요에 따라 로드됩니다. 데이터를 로드하는 속도보다 빠르게 확대하면 로드가 완료될 때까지 데이터가 저해상도로 표시될 수 있습니다.

Trace 뷰어 인터페이스 구성요소

Trace 뷰어의 주요 UI 구성요소는 다음과 같습니다.

• 시간 축은 상단에서 가로로 표시되며 trace의 시작을 기준으로 한 시간을 보여줍니다.
• 타임라인은 섹션 및 트랙별로 정렬되며 라벨은 왼쪽 세로축에 표시됩니다. 각 섹션은 처리 요소 (예: 기기 노드 또는 호스트 스레드)를 나타내며 확장하거나 접을 수 있습니다. 각 섹션에는 특정 활동의 타임라인인 트랙이 있습니다.
• 이벤트는 타임라인 트랙에 있는 색상이 지정된 직사각형 블록으로, 작업 기간 또는 학습 단계와 같은 메타 이벤트를 나타냅니다. 이벤트의 색상에는 특별한 의미가 없습니다.
• 세부정보 창에는 타임라인 창에서 선택한 이벤트에 관한 추가 정보(예: 이름, 시작 시간, 지속 시간)가 표시됩니다.

일반적인 섹션 및 트랙

Trace Viewer는 다음과 같은 섹션과 트랙을 제공합니다.

• 다음과 같은 트랙이 있는 TPU 노드당 하나의 섹션입니다.
  • XLA 모듈: 실행 중인 XLA 프로그램입니다.
  • XLA 작업: TPU 코어에서 실행된 XLA HLO 작업을 표시합니다. 각 상위 수준 프레임워크 작업 (예: JAX, TensorFlow, PyTorch)은 하나 이상의 XLA 작업으로 변환된 후 TPU에서 실행되도록 컴파일됩니다. XLA 작업을 클릭하면 그래프 뷰어에 작업 링크와 함께 시작/중지 시간, 시간, 소스 스택 트레이스 (프레임워크 및 컴파일러에서 제공하는 경우)와 같은 작업에 관한 추가 정보가 표시됩니다.
  • XLA TraceMe: 추적하려는 논리적 작업 단위를 설명하는 코드의 사용자 지정 주석입니다. 주석을 추가하지 않았더라도 여기에 데이터가 표시될 수 있습니다. 이러한 데이터는 일반적으로 XLA (예: 배리어 코어) 또는 XProf 자체 (예: 삭제된 트레이스 항목)에서 추가합니다.
  • 단계: 사용자 프로그램 또는 프레임워크에 적절하게 주석이 달린 경우 해당 TPU 코어에서 실행되는 학습 단계의 지속 시간을 나타냅니다.
  • 프레임워크 작업: 사용자 프로그램 또는 프레임워크에 적절하게 주석이 달린 경우 TPU 코어에서 실행된 프레임워크 작업 (예: JAX, TensorFlow, PyTorch)을 표시합니다.
  • 프레임워크 이름 범위: 각 프레임워크 작업에 대해 스택 트레이스의 시각화입니다. 간단히 하기 위해 이 트랙은 단일 기기에만 표시됩니다.
  • 소스 코드: 실행 중인 소스 코드의 경로입니다(프로필에 있는 경우).
  • 스칼라 단위: TPU의 경우 스칼라 단위에서 실행되는 이벤트로, 프로필에 있는 경우 표시됩니다.
  • TensorCore 동기화 플래그: TPU의 동기화 메커니즘으로, 프로필에 있는 경우 표시됩니다.
  • 호스트 오프로드: 호스트 메모리와 가속기 메모리 간에 데이터를 비동기식으로 이동하는 작업입니다. 일반적으로 XLA 작업 행에 데이터 전송을 위해 준비되는 가속기를 나타내는 상응하는 시작 작업과 중지 작업이 표시됩니다 (예: 전송 기간 동안 소스/대상 메모리 영역을 '사용 중'으로 표시). 동시에 실행되는 오프로드 작업이 여러 개인 경우 호스트 오프로드 행이 여러 개 있을 수 있으므로 trace 뷰어에서 여러 이벤트를 동시에 표시해야 합니다.
• Sparsecore 노드당 하나의 섹션: 일부 TPU 세대 (예: TPU v5p 및 TPU v6e)에는 집적 컴퓨팅 MXU 단위 외에 하나 이상의 SparseCore 단위가 장착되어 있습니다. 이러한 코어와 연결된 모듈, 작업, TraceMes가 이 섹션에 표시됩니다.
• GPU 노드당 하나의 섹션으로 다음과 같은 트랙이 포함됩니다.
  • XLA 모듈, 프레임워크 작업, 프레임워크 이름 범위, 단계, 소스 코드 이는 모두 TPU 섹션과 유사합니다.
  • 스트림당 트랙 1개. 스트림 이름에는 스트림에서 실행되는 작업 유형 (Memcpy, Compute 등)에 관한 정보도 포함됩니다.
  • GPU에서는 XLA TraceMe가 지원되지 않습니다.
  • XLA 작업은 GPU 섹션에 표시되지만 스트림 데이터에서 파생되므로 현재 항상 정확하지는 않습니다. 따라서 XLA 작업을 여러 스트림에서 실행되는 실제 커널에 N:M 매핑하고 하드웨어의 여러 SM에 여러 스트림을 동적으로 예약하는 GPU의 실행 모델을 완전히 설명할 수 없습니다.
• 호스트 머신의 CPU에서 실행되는 각 구성요소(예: 스레드 풀 1개)에 대한 섹션 1개(스레드 풀의 경우 스레드당 트랙 1개) 프로필 수집 중에 Python 트레이스가 사용 설정된 경우 여기에서 트레이스를 확인할 수도 있습니다.

TPU의 XLA 작업과 GPU의 스트림 데이터만 수집된 프로필에 직접 연결됩니다. 다른 모든 행은 컴파일러에서 제공하는 선택적 '부채널' 정보, 선택적 사용자 주석 또는 XProf에서 적용한 휴리스틱을 포함하는 '파생된 행'입니다. 따라서 이러한 파생된 선은 특정 프로필에 표시되기도 하고 표시되지 않을 수도 있습니다.

기타 기능

• '활동 찾기' 검색창을 사용하여 특정 활동 이름을 검색할 수 있습니다. 현재 이 기능은 전체 트레이스가 아닌 화면에 표시되는 시간 범위 내에서만 검색합니다.
• 흐름 이벤트: 상단 표시줄에서 '흐름 이벤트' 버튼을 클릭하여 이 옵션을 사용 설정하면 시각화가 추가되어 한 대화목록 또는 대화목록 줄의 이벤트가 다른 대화목록 줄의 이벤트에 연결됩니다. 예를 들어 가속기의 작업을 큐에 추가하거나 실행하는 호스트의 작업에서 해당 작업을 실행하는 가속기의 작업으로 화살표가 그려질 수 있습니다. XProf는 사용자 주석, 내장 휴리스틱, 다양한 구성요소 (예: CUPTI 드라이버, 커널 실행 ID, TPU 런타임 정보 등)을 가져옵니다.
• XLA 연산을 클릭하면 세부정보 창에 추가 정보가 표시됩니다. 예를 들어 그래프 뷰어 도구의 op에 연결됩니다. 소스 코드 또는 Python 스택 트레이스, 이 XLA 연산이 생성된 원인이 된 프레임워크 연산 등에 대한 포인터도 제공할 수 있습니다(프로필에 있는 경우). FLOPS (op에서 실행된 부동 소수점 연산 수) 및 op에서 액세스한 바이트도 표시될 수 있습니다. 이 정보는 프로필의 런타임 정보가 아닌 컴파일 중에 XLA에서 정적으로 획득됩니다.