Investigating JAX traces

개요

하나의 캡처에 대한 7단계 조사: 판결, 조정, 스윔레인, 커널 → HLO → Python 라인, 루프라인, 변경 사항, 신뢰 목록.

모든 정보는 사용자가 제공한 캡처에서 가져옵니다: <host>.xplane.pb 및/또는 <host>.trace.json.gz. 워크로드를 다시 실행하거나, XLA를 덤프하거나, 다른 프로파일러를 실행하지 마십시오. 트레이스가 답을 제공할 수 없는 경우, "신뢰할 수 없는 목록"에 추가합니다.

목표

보고서를 읽은 후, 독자는 다음을 알 수 있습니다:

(a) 가정과 함께 추정치로, 스텝을 더 빠르게 만들기 위해 변경해야 할 사항;
(b) 이 트레이스가 속성을 지원할 수 없는 위치와 그 개수.

도구

스크립트 경로는 이 스킬의 기본 디렉토리(스킬이 로드될 때 표시되는 디렉토리)를 기준으로 합니다. 스크래치 디렉토리에서 실행하십시오: 두 스크립트 모두 출력을 현재 디렉토리에 씁니다.

B=<this skill's base directory>
python3 -P $B/../analyzing-jax-launch-overhead/scripts/skip_stats.py <capture>/<host>.trace.json.gz
python3 -P $B/scripts/attribute_kernels.py <capture dir> [--program ID] [--user-re REGEX]

skip_stats.py (표준 라이브러리만 사용): 판결, 스트림별 시간 예산, 유휴 시간 동안의 호스트 이벤트, SKIP 논문 지표. 읽는 방법: analyzing-jax-launch-overhead 스킬 참조.
attribute_kernels.py (jax 및 xprof 패키지 필요): 모든 GPU 커널 → HLO 명령어 → op_name → Python 라인. .xplane.pb만 읽고, 먼저 ./attribution/xplane/에 복사합니다. XProf가 HLO 프로토를 그 옆에 쓰기 때문입니다. kernels.csv, ops.csv, lines.csv, how.csv, summary.txt 및 프로그램의 HLO 텍스트를 ./attribution/에 씁니다.
- 정확(exact): 호출당 하나의 커널이 실행되는 경우, 트레이서 자체의 hlo_op.
- 이름(name): 이름이 해당 윈도우에서 유일한 후보인 그래프 커널.
- 순서(order): 공유 커널, 라이브러리 GEMM, 복사 및 컬렉티브, 스케줄 순서로 일치.
- 없음(none): 남은 후보가 없음.
- 윈도우는 인접한 정확 커널 사이의 명령어들로, 알려진 반복 횟수로 while 루프가 확장된 프로그램 스케줄에 따릅니다.
- 검증: MiniGPT 캡처에서, 명령 버퍼를 끈 동일 프로그램 실행과 100% 일치하는 파이프라인으로 147,840개 모든 커널에 대해 동일한 명령어를 선택했습니다.

다섯 가지 기법과 각각의 도구

분야	기법	여기서의 도구
성능 엔지니어링	임계 경로; 루프라인	`skip_stats.py`의 스트림 예산; 올바른 피크만 사용하는 루프라인(트랩 참조)
분산 추적	상관 ID로 호스트 및 GPU 이벤트 결합; 지연 작업 찾기	`skip_stats.py` → `rows.csv`; GPU 간 스텝 종료 시간의 분포
컴파일러	컴파일러 패스를 통한 소스 출처	`attribute_kernels.py` → `lines.csv`
포렌식 회계	총계가 독립적인 레코드에서 일치해야 함	2단계의 검사
정보 설계	GPU 스트림당 하나의 정렬된 스윔레인	한 스텝, 모든 GPU × 스트림, 하나의 시간 축에서

단계

다음은 주어진 파일에 대해 모두 실행되는 7단계 계획입니다. 새 캡처만이 답할 수 있는 질문(호스트 호출 내부 동작, 수정의 효과, 배치 스윕 등)은 7단계의 "새 캡처 필요" 목록에 추가되며, 캡처 여부는 사용자가 결정합니다.

판결 및 호스트 측 원인: skip_stats.py 실행.
조정. 각 검사는 통과하거나 신뢰 목록에 추가됩니다:
- 장치별 GPU 이벤트: json = xplane;
- 모든 GPU 이벤트는 하나의 호스트 실행 호출과 결합;
- 스텝당 커널 수는 일정;
- 모든 GPU가 동일한 작업 수행(스텝당 동일한 커널 수 또는 차이 설명 가능);
- 각 스트림에 대해 compute + collective + memcpy + idle 점유율이 합쳐져 span이 됩니다.
스윔레인. 중간 스텝에 대해, 모든 GPU × 스트림을 하나의 시간 축에 놓고, 실행 스레드의 호스트 이벤트도 추가합니다. 이는 임계 경로와 지연 작업(가장 늦게 완료되는 GPU와 해당 GPU가 대기하는 것)을 보여줍니다.
커널 → HLO → Python 라인: attribute_kernels.py 실행.
- 각 일치 클래스의 GPU 시간 점유율과 후보가 없는 커널 보고.
- lines.csv와 ops.csv를 사용하여 라인별 및 연산별 테이블 생성.
- --user-re는 사용자 코드로 간주할 프레임을 선택합니다. 기본값은 site-packages와 jax 외부의 첫 번째 프레임입니다.
루프라인, HLO 형태에서 FLOPs와 바이트를 얻을 수 있는 연산에 대해서만: attribute_kernels.py가 ./attribution/에 쓰는 <program>(<id>).hlo.txt를 사용합니다. 연산의 데이터 타입에 대한 피크를 벤더 사양에서 인용하여 *읽기(read)*로 표시합니다. 트레이스는 피크를 측정할 수 없습니다.
변경 사항. 원인을 임계 경로 점유율 기준으로 순위를 매깁니다. 각각에 대해, 해당 원인이 사라졌을 때의 스텝 시간을 추정합니다(예: 임계 스트림에서 host-idle 제거). *추정값(inferred, not run)*으로 표시하고, 사용자가 테스트할 가설로서 조정 항목을 명명합니다(donate_argnums, 명령 버퍼 끄기, 더 적거나 더 큰 컬렉티브 등).
이 트레이스가 신뢰할 수 없는 위치: 각 항목을 개수 또는 점유율과 함께 나열합니다. 그런 다음 "새 캡처 필요" 목록을 추가합니다: 각 미해결 질문과 이를 해결할 캡처 유형.

트랩 (ROCm MI355X 캡처에서 측정)

XProf의 op 뷰는 그래프 커널을 누락합니다. hlo_stats는 4-GPU 캡처에서 GPU 커널 시간의 37.6%만을 포함했습니다. 나머지 62.4%(명령 버퍼 커널)는 kernel_stats에서 op "XlaModule"로 나타납니다. hlo_stats, op_profile 또는 roofline_model을 인용하기 전에 적용 범위를 확인하십시오.
gfx950에 대한 XProf의 루프라인 피크는 데이터 타입과 관계없이 2,517 TFLOP/s 및 3,624 GiB/s입니다. FP32 GEMM은 FP32 피크가 필요합니다.
RCCL/NCCL 커널 시간에는 피어 대기 시간이 포함됩니다. 이는 전송 시간이 아닙니다.
비동기 컬렉티브: 해당 hlo_op(예: all-gather.34)는 캡처에 저장된 HLO에 없습니다. attribute_kernels.py는 레이블을 유지하고 순서에 따라 일치하는 collective-start 명령어에서 Python 라인을 가져옵니다.
ROCm memcpy 행은 kind:Unknown size:0을 표시할 수 있습니다.
Python 트레이서는 호스트 시간을 부풀립니다.
하나의 프로그램이 두 번 나타날 수 있습니다(서로 다른 ID를 가진 두 개의 .hlo_proto.pb 파일). attribute_kernels.py는 xplane에서 가장 많은 GPU 시간을 사용하는 프로그램 ID를 사용합니다. json에는 program_id가 없습니다.

보고서

마크다운: 결과를 먼저 제시.
숫자: 각 숫자는 스텝 및 GPU 간 범위와 함께 측정값(measured), 읽기 값(read) 또는 *추정값(inferred)*으로 표시.
섹션: 판결 및 원인, 조정, GPU/스트림별 테이블, 상위 연산 및 Python 라인, 변경 사항(추정치), 신뢰 목록.
사실 확인: 전송 전에 새로운 하위 에이전트가 테이블을 기준으로 각 주장을 확인.
HTML 페이지: 사용자가 원하는 경우, writing-trace-report-pages 스킬이 이 보고서와 테이블로부터 HTML 페이지를 구축.

JAX 트레이스 조사