LHS 비용 모델

요약

이 페이지에서는 지연 시간 숨기기 스케줄러에서 사용하는 비용 모델의 내부를 설명합니다. 모델을 조정하는 데 관심이 있다면 조정 섹션으로 바로 이동하세요.

지연 시간 숨기기 스케줄러 (LHS)는 실제 경과 시간을 최소화하는 방식으로 HLO DAG를 예약하는 컴파일러 패스입니다.

이러한 결정은 성능 표와 분석 모델을 혼합하여 사용하는 통합 비용 모델을 기반으로 합니다. 특히 XLA는 GEMM 및 빠른 상호 연결 집합의 성능 표를 삽입하고 다른 경우에는 분석 네트워킹 및 융합 비용 모델을 사용합니다. 문서의 나머지 부분에서는 이러한 기능의 작동 방식을 개략적으로 설명합니다.

실적 표 - ICI 집단

성능 표는 수집기와 인터폴레이터라는 두 가지 주요 구성요소로 구성됩니다.

수집기

수집기는 집단 작업의 성능 표를 생성하는 C++ 도구입니다. 개별 HLO 작업 (예: all-gather, all-reduce)를 정적으로 정의된 매개변수 공간에 걸쳐 실행합니다.

작동 방식

이 도구는 지정된 클러스터에 대해 다양한 집단 작업, 전송 크기, 전송 스키마를 스위핑합니다. 기존 멀티 호스트 HLO 러너 인프라와 ExecutionProfile 데이터를 사용하여 생성된 HLO를 실행하고 성능 측정항목을 수집합니다.

데이터 수집 매개변수

지연 시간 표는 다음 매개변수의 교차 곱에 대해 수집됩니다.

  • Collective Type:
    • all-reduce
    • all-gather
    • reduce-scatter
  • 전송 크기:
    • 1024B에서 2GiB까지의 로그 스케일 (예: 1024B, 2048B, 4096B, ...)
  • 전송 방식:
    • rail-aligned
    • non-rail-aligned

이 스위프는 기기가 2, 4, 8개인 노드 내 클러스터에 대해 실행됩니다.

출력

수집 실행의 결과는 .pbtxt 형식의 지연 시간 테이블입니다(플랫폼당 약 116KB).

보간기

인터폴레이터는 생성된 성능 표를 사용하여 컴파일 중에 런타임 추정치를 제공하는 컴파일러 구성요소입니다.

내부 데이터 구조

초기화 시 인터폴레이터는 성능 표를 지도로 처리합니다. 이 맵은 (collective_type, transfer_scheme) 튜플을 키로 사용합니다.

각 키와 연결된 값은 2D 유클리드 평면입니다. 이 평면은 두 축을 기반으로 수집기가 측정한 네트워크 처리량을 색인화합니다.

  1. 전송 크기입니다.
  2. 관련된 기기 수입니다.

조회 및 보간

컴파일러가 집단 연산을 발견하면 인터폴레이터는 다음 단계를 실행합니다.

  1. 작업의 (collective_type, transfer_scheme)를 맵 키로 사용하여 올바른 2D 처리량 평면을 식별합니다.
  2. 그런 다음 작업의 (transfer_size, num_devices)를 쿼리 지점으로 사용하여 해당 2D 평면 내에서 가중 평균 검색 (유클리드 거리 기반)을 사용합니다.
  3. 이 조회 결과는 단일 고유 네트워크 처리량 값입니다.

근거: 처리량 및 외삽

시스템은 원시 지연 시간 대신 네트워크 처리량을 저장하도록 설계되었습니다. 이 디자인 선택으로 테이블에 명시적으로 표시되지 않은 전송 크기의 성능을 외삽하는 작업이 크게 간소화됩니다.

지연 시간 표에서 집합 크기 S의 네트워크 대역폭 포화를 포착하는 경우 해당 시점의 처리량 T이 최대로 간주됩니다. 크기가 S' > S인 새 집단의 경우 런타임은 다음과 같이 추정할 수 있습니다.

\[\text{EstimatedTime}(S') = \frac{S'}{T_{\text{saturated} } }\]

이를 통해 모델은 수집기가 측정한 최대 크기인 2GiB보다 큰 집단에서도 모든 크기의 집단의 성능을 추정할 수 있습니다.

  • 최대 처리량을 과소평가합니다.
  • 따라서 대규모 전송의 런타임을 과대평가합니다.

일반적으로 XLA:GPU팀에서 성능 표를 유지관리하지만 사용자가 자체 표를 제공하기로 결정한 경우 표를 생성하는 사용자는 표가 대표성을 가지며 타겟 하드웨어의 대역폭 포화 영역에서 측정을 포함해야 합니다.

실적 표 - GEMM

GEMM 지연 시간 표는 집단 시스템과 마찬가지로 수집기와 인터폴레이터라는 두 가지 구성요소로 지원됩니다.

수집기

collector는 일반 행렬 곱셈 (GEMM)의 성능 표를 계산하는 C++ 도구입니다. HLO dot 작업 수준에서 행렬 곱셈의 성능을 측정합니다.

작동 방식

이 도구는 GEMM 차원 (배치, 두 개의 비수축 차원, 하나의 수축 차원)과 데이터 유형의 정적 공간을 스위프합니다.

  • 기본 데이터 유형: LHS = bf16,f32, RHS = bf16,f32, OUT = bf16,f32
  • 인프라: HLO 작업 프로파일러를 재사용합니다.

수집 매개변수

지연 시간 표는 다음 측정기준의 교차 제품에 대해 수집됩니다.

  • batch: {1, 2, 4}
  • m (비수축): {256, 512, ..., 4096}
  • n (비수축): {256, 512, ..., 4096}
  • k (축소): {256, 512, ..., 4096}

출력 및 저장소

전체 스윕은 인터폴레이터에서 사용할 수 있는 .pbtxt 지연 시간 테이블을 생성합니다.

보간기

인터폴레이터는 생성된 표를 사용하여 GEMM 성능을 추정하는 컴파일러 구성요소입니다.

근거: FLOPS 포화

수집된 지연 시간 테이블을 통해 인터폴레이터는 각 항목의 FLOPS를 재구성할 수 있습니다.

\[\text{FLOPS} = \frac{2 \times b \times m \times n \times k}{\text{runtime} }\]

핵심 통계는 특정 시점에서 FLOPS가 포화된다는 것입니다. 즉, 특정 행렬 모양을 초과하면 하드웨어가 최고 FLOPS에 도달합니다. 이 포화 상태를 통해 집단에 사용되는 것과 동일한 외삽법을 사용할 수 있습니다.

조회 및 보간

인터폴레이터는 표 데이터에서 4D 유클리드 공간을 빌드합니다. 성능 추정치를 제공하기 위해 이 4D 공간 내에서 가중 평균 보간을 실행합니다. 특정 데이터 유형의 테이블이 없는 경우 휴리스틱으로 각 측정기준이 바이트 수로 정규화됩니다.

분석 비용 모델 - DCN

S형 곡선 집단 비용 모델

S곡선 모델은 완전 분석적 네트워킹 루프라인 모델입니다.

개요

이 모델은 고정된 네트워크 속성 집합을 기반으로 집단 작업의 성능을 추정하도록 설계되었습니다.

모델 입력

모델에는 다음 두 가지 카테고리의 입력이 필요합니다.

  1. 고정 네트워크 속성 (사용자 정의):

    • 집단 실행 오버헤드
    • NIC 속도
    • RTT (왕복 시간)

    기본적으로 XLA는 플랫폼을 자동 감지하고 가장 일반적인 아키텍처의 값을 사용합니다. 이러한 속성은 사용자가 구성할 수 있습니다. 자세한 내용은 튜닝 섹션을 참고하세요.

  2. 집단별 입력:

    • 집단 유형 (예: AllGather, ReduceScatter)
    • 전송 크기
    • 통신에 참여한 노드 수

통합

S-곡선 모델은 XLA:GPU에 통합되어 Hopper 및 Blackwell에서 사용되고 있습니다.

분석 비용 모델 - 융합

다른 커널의 경우 GPU 성능 비용 모델을 사용하여 적절한 런타임을 추정합니다. 자세한 내용은 여기에서 확인하실 수 있습니다.

조정

올바른 XLA 플래그를 실행하여 S-곡선 모델을 조정할 수 있습니다. 기본 구성은 대부분의 경우에 충분하지만 다른 경우에는 모델 제어가 노출됩니다.

export NIC_SPEED_GBPS=... # NIC speed per GPU in Gigabytes
export GPUS_PER_NODE=... # Num of GPUs per cluster interconnected with fast network (e.g. NVLINK)
export XLA_FLAGS=--xla_gpu_analytical_latency_estimator_options="nic_speed_gbps=$NIC_SPEED_GBPS,gpus_per_node=$GPUS_PER_NODE"