대규모 임베딩 모델 (LEM)용 SparseCore 심층 분석

SparseCore는 불규칙하고 스파스한 메모리 액세스 및 컴퓨팅이 포함된 워크로드의 고성능 가속을 위해 설계된 특수 타일형 프로세서로, 특히 고대역폭 메모리(HBM)에 저장된 대규모 데이터 세트에서 유용합니다. 임베딩 조회와 같은 작업에 탁월한 성능을 발휘하며, 다양한 기타 동적이고 희소한 워크로드를 가속화하는 기능도 제공합니다.

1. SparseCore 소개

주요 아키텍처 기능:

타일 아키텍처: 병렬 처리를 지원하는 여러 컴퓨팅 타일 (각 타일은 자체 로컬 메모리 및 처리 장치가 있는 완전한 Dataflow 단위)로 구성됩니다.
동적 실행: 희소 데이터에 중요한 데이터 종속 제어 흐름과 메모리 액세스를 기본적으로 지원합니다.
벡터 처리: 효율적인 계산을 위해 소형 벡터 작업 (하드웨어 버전에 따라 8요소 또는 16요소)을 활용합니다.
중앙 집중식 제어: 단일 SparseCore 시퀀서가 모든 타일에서 작업을 조정하여 동기화된 작업을 보장합니다.
데이터 요약 지원: 정렬, 필터링, 접두사 합계와 같은 작업에 유용한 전문 교차 레인 작업이 포함됩니다.
메모리 계층 구조: 대규모 데이터 세트를 저장하는 데 HBM을 전략적으로 활용하고 자주 액세스하는 데이터를 스테이징하는 데 로컬 스크래치패드 메모리 (SPMEM)를 활용하여 HBM 지연 시간을 크게 줄입니다.

사양 한눈에 보기:

속성	TPU v4	TPU v5p	Trillium
SparseCores/Chip	4	4	2
타일/SparseCore	16	16	16
SIMD 너비	8	8	8 (F32) 16 (BF16)
HBM 용량	32GiB	96GiB	32GiB

2. SparseCore 호스트 전처리

효과적인 데이터 준비는 SparseCore 성능에 매우 중요하며, 호스트 사전 처리가 중요한 역할을 합니다. 여기에는 다음과 같은 몇 가지 주요 기능이 포함됩니다.

데이터 변환:
- 원시 입력 데이터에 필요한 변환을 적용합니다.
- ID 변환을 관리합니다. 이는 특히 기능 또는 테이블 스태킹을 처리할 때 중요합니다.
- 입력 데이터를 다음 섹션에 자세히 설명된 좌표 (COO) 희소 형식으로 변환합니다.
- 칩에서 사용할 수 있는 다양한 SparseCore에 효율적으로 분산되도록 데이터를 파티셔닝합니다.
한도 유효성 검사:
- 입력 데이터의 특성 (예: ID 수)이 max_ids_per_partition, max_unique_ids_per_partition과 같은 SparseCore의 사전 정의된 운영 제한을 준수하는지 확인합니다.
- 입력 데이터가 이러한 한도를 초과하면 호스트 전처리 레이어가 제약 조건에 맞는 더 작은 미니 배치로 데이터를 분할할 수 있습니다.
데이터 전송:
- 처리되고 검증된 데이터를 TPU의 고대역폭 메모리 (HBM)에 효율적으로 복사하여 SparseCore 실행을 준비합니다.

표 스태킹 이해하기:

표 스태킹은 여러 임베딩 표를 논리적으로 결합하여 임베딩 조회 효율성을 높이는 중요한 최적화 기법입니다. 이 프로세스는 일반적으로 기본 ML 프레임워크에서 자동으로 처리합니다.

기능 스태킹: 여러 개별 기능이 동일한 기본 삽입 테이블을 공유하는 경우에 발생합니다. 일반적인 예는 서로 다른 컨텍스트의 우편번호와 같은 다양한 범주형 특성에 단일 삽입 사전이 사용되는 것입니다.
표 스태킹: 이 시나리오에서는 서로 다른 여러 삽입 표가 함께 스택됩니다. 동일한 삽입 차원과 최적화 프로그램 구성을 공유하는 테이블은 종종 그룹화됩니다.

테이블 스태킹의 주요 이점은 이러한 스택 테이블에 대한 작업의 효과적인 배치 크기를 늘릴 수 있다는 것입니다. 이렇게 하면 계산 오버헤드가 줄어들고 칩 간 통신(ICI) 지연 시간을 숨기는 데 효과적일 수 있습니다. 최적의 성능을 위해 적당한 수의 스택 테이블(일반적으로 5~100개)을 사용하는 것이 좋습니다.

3. COO 텐서로 변환

데이터가 SparseCore에 의해 처리되기 전에 일반적으로 좌표 (COO) 희소 텐서 형식으로 변환됩니다. COO 형식은 일반적으로 세 개의 배열을 사용하여 희소 행렬을 효율적으로 표현하는 방법입니다.

row_ids: 0이 아닌 각 요소의 행 색인이 포함된 배열입니다. 일괄 처리의 맥락에서 이는 일괄 처리 차원에 해당하는 경우가 많습니다.
col_ids: 0이 아닌 각 요소의 열 색인이 포함된 배열입니다. 임베딩의 경우 특성 또는 ID 값인 경우가 많습니다.
values (선택사항): 해당 (row, col) 좌표에 있는 0이 아닌 요소의 실제 값을 보유하는 배열입니다. ID 수와 관련된 한도 계산 (나중에 설명)의 경우 이러한 값 (이득)은 고려되지 않는 경우가 많습니다.

실제 예시:

ID 배치를 나타내는 입력 희소 행렬을 고려해 보세요.

[
    [id_A],                 // Sample 0
    [id_A, id_B, id_C],     // Sample 1
    [id_B, id_B, id_D],     // Sample 2 (note duplicate id_B)
]

COO 형식으로 변환한 후 (그리고 동일한 샘플 내에서 ID를 중복 삭제한 후)

row_ids = [0, 1, 1, 1, 2, 2]
col_ids = [id_A, id_A, id_B, id_C, id_B, id_D]

이 변환은 SparseCore가 작업을 처리하고 분배하는 방식에 있어 기본적입니다. 특히 col_ids는 ID가 속한 특정 SparseCore 파티션을 결정하는 데 중요하며 효율적인 샤딩과 조회를 지원합니다.

4. SparsecoreConfig: 고급 API

프레임워크별 삽입 API:

JAX: https://github.com/jax-ml/jax-tpu-embedding
TensorFlow: https://www.tensorflow.org/recommenders/api_docs/python/tfrs/layers/embedding/TPUEmbedding
Keras: https://keras.io/keras_rs/api/embedding_layers/distributed_embedding

SparsecoreConfig 또는 XLA 플래그와 같은 동등한 메커니즘은 다양한 SparseCore 동작을 제어하는 상위 수준 인터페이스 역할을 합니다. 이러한 매개변수를 철저히 이해하는 것은 효과적인 성능 조정과 모델의 올바른 작동을 보장하는 데 매우 중요합니다.

disable_table_stacking: bool = False
- 설명: 이 플래그는 자동 표 스태킹이 프레임워크에서 표를 스태킹하는 것을 방지하는지 제어합니다. 이렇게 하면 오버헤드가 증가하고 칩 간 상호 연결 (ICI) 지연 시간을 숨기는 기능이 저하되어 성능이 저하될 수 있습니다.
- 기본값: False (프레임워크에서 지원하는 경우 테이블 스태킹이 기본적으로 사용 설정됨을 의미)
max_ids_per_chip_per_sample: int = 64
- 설명: 이 파라미터는 모든 테이블에서 집계된 입력 배치에 있는 하나의 샘플에서 단일 칩이 처리할 수 있는 총 삽입 ID의 전역 상한을 설정합니다. 더 세분화된 테이블별 또는 파티션별 한도가 고려되기 전에 칩 수준에서 리소스를 관리하는 메커니즘입니다. 이 값의 미세 조정은 일반적으로 특정 모델 특성과 전체 시스템 용량에 따라 달라집니다.
- 기본값: 64.
max_ids_per_table: Optional[Dict[str, int]] = None
- 설명: 이 매개변수는 모든 SparseCore에서 모든 파티션을 고려하여 각 논리 테이블에 대해 처리할 수 있는 최대 삽입 ID 수 (중복 포함 가능)를 지정합니다. 이는 max_ids_per_partition보다 넓은 한도입니다. 표 T이 P 파티션으로 나뉘는 경우 이 한도는 모든 P 파티션으로 전달되는 ID의 합계에 적용됩니다. 이는 max_ids_per_partition_per_sample 및 전체 배치 크기와 관련이 있는 경우가 많습니다.
- 설정: 일반적으로 max_ids_per_partition가 정의된 제한 파일 (예: xla_sparse_core_max_ids_file 플래그 사용)을 사용하여 구성됩니다. 이 테이블 수준 개념은 이러한 파티션 수준 제한 (max_ids 및 max_uniques)을 설정하는 방법입니다.
- 기본값: None (명시적으로 제공되지 않은 경우 파티션별 한도 또는 기타 구성에서 값을 추론할 수 있음)
max_unique_ids_per_table: Optional[Dict[str, int]] = None
- 설명: max_ids_per_table와 유사하지만 이 매개변수는 각 논리 테이블의 고유 ID 최대 수를 지정합니다. 이는 고유 ID 처리 및 후속 벡터 작업에 사용되는 온디바이스 버퍼의 크기를 적절하게 조정하는 데 중요한 설정입니다.
- 설정: 제한 파일에 정의되거나 max_unique_ids_per_partition_per_sample에서 파생되기도 합니다.
- 기본값: None.
allow_id_dropping: bool = False
- 설명: 이 불리언 플래그는 입력 데이터에서 발견된 ID 수 (관찰된 제한)가 컴파일 중에 설정된 제한 (예: max_ids_per_partition)을 초과할 때 ID 삭제를 제어합니다.
  - True인 경우: 한도를 초과하게 되는 ID는 자동으로 삭제됩니다. 일반적으로 파티션 내 ID는 정렬된 순서로 처리되며 지정된 미니 배치 한도를 초과하는 ID는 삭제됩니다. 이렇게 하면 프로그램이 계속 실행될 수 있지만 모델 정확도에 부정적인 영향을 미칠 수 있습니다.
  - False인 경우: 오류가 트리거되고 관찰된 제한이 컴파일된 제한을 초과하면 프로세스가 종료될 수 있습니다. 이 접근 방식을 사용하면 모든 데이터가 처리되지만 한도를 더 보수적으로 구성해야 합니다.
- 기본값: False (데이터가 자동으로 삭제되는 대신 오버플로 시 오류가 발생함).
initialize_tables_on_host: bool = True
- 설명: 이 플래그는 삽입 테이블이 호스트 CPU에서 초기화된 후 TPU의 고대역폭 메모리 (HBM)로 전송되는지 여부를 결정합니다. 표는 호스트에서 초기화하는 것이 표준 관행입니다. 이 값을 True로 설정하면 이 규칙을 따릅니다. False로 설정되면 기기 내 초기화 메커니즘이 암시되며, 이는 성능에 미치는 영향이나 특정 초기화 사전 요구사항이 다를 수 있습니다.
enable_fast_table_initialization: bool = False
- 설명: TPU에서 직접 테이블을 초기화합니다. 이렇게 하면 모델 시작 시간을 줄일 수 있습니다.

5. 성능을 위한 파이프라인

파이프라이닝은 TensorCore (TC)와 SparseCore(SC)에서 작업을 동시에 실행할 수 있는 성능 최적화 기법입니다. 이러한 계산을 중첩하면 전체 처리량을 크게 개선할 수 있습니다.

메커니즘: 희소 삽입 조회(SC에서 처리)와 밀도 높은 레이어 계산(TC에서 처리)이 포함된 표준 학습 단계에서 파이프라인을 사용하면 TC가 동일한 단계 i의 다른 부분을 동시에 처리하거나 인접한 단계(예: i-1 또는 i+1)의 일부를 처리하는 동안 SC가 단계 i의 일부(예: 순방향 또는 역방향 패스)를 처리할 수 있습니다.
그라디언트에 미치는 영향: SparseCore가 '오래된' 그라디언트에서 작동할 수 있습니다. 예를 들어 i 단계의 역전파 단계에서 계산된 그라데이션은 i+2 단계까지 완전히 업데이트되지 않고 SC에 표시되지 않을 수 있습니다.
성능과 숫자 간의 균형: 이렇게 중복 실행하면 속도가 크게 향상되어 기기 단계 시간이 최대 2배까지 개선될 수 있습니다. 하지만 오래된 그라데이션을 사용하여 발생하는 숫자 (embedding_weights)의 미묘한 변화는 모델 수렴 동작이나 최종 달성된 정확도에 영향을 미칠 수 있습니다. 이 트레이드오프의 허용 여부는 모델에 따라 크게 달라지며 경험적 검증이 필요한 경우가 많습니다.
제어 플래그: 파이프라인은 tf_xla_disable_full_embedding_pipelining로 제어할 수 있습니다. 이 플래그를 true로 설정하면 전체 파이프라인 (TensorCore와 SparseCore 계산 중복)이 사용 중지되고 false로 설정하면 (또는 플래그의 의미가 false일 때 사용 설정됨을 의미하는 경우) 활성화됩니다.

개념적 파이프라인 흐름:

파이프라인이 없는 경우 (간소화된 순차적 흐름):

Loop: SC/F_i -> TC/F_i -> TC/B_i -> SC/B_i
파이프라인 (간소화된 중복 흐름):
```
Time ->
Step i:   SC/F_i | TC/F_i | TC/B_i | SC/B_i
Step i+1:          SC/F_i+1| TC/F_i+1| TC/B_i+1| SC/B_i+1
```
참고: 하드웨어와 컴파일러에 구현된 실제 파이프라인 단계는 데이터 종속성을 관리하고 정확성을 보장하기 위해 사전 루프, 기본 실행 루프, 사후 루프를 포함하는 경우가 많아 더 복잡할 수 있습니다.

6. XLA의 역할

XLA (가속 선형 대수학)는 일반적으로 TensorFlow와 같은 프레임워크에서 가져온 고급 계산 그래프를 TPU에 맞게 고도로 최적화된 기계어 코드로 변환하는 도메인별 컴파일러입니다. 여기에는 SparseCore로 전송되는 작업의 명령어 생성도 포함됩니다.

SparseCore 컨텍스트의 주요 기능은 다음과 같습니다.

희소 작업 컴파일: XLA는 삽입 조회 작업 (예: SparseDenseMatmulOp) 및 기타 희소 계산을 실행 가능한 하위 수준 SparseCore 프로그램으로 컴파일합니다.
한도 통합: 구성된 작동 한도(예: max_ids_per_partition, max_unique_ids_per_partition, xla_sparse_core_max_ids_file와 같은 플래그로 지정된 한도 파일을 통해 제공되는 경우가 많음)를 활용하여 기기 내 메모리 버퍼의 크기를 정적으로 결정하고 할당합니다(특히 SPMEM 내).
타겟팅된 최적화: XLA는 SparseCore 아키텍처를 위해 특별히 설계된 최적화 모음을 실행합니다. 여기에는 효율성을 극대화하기 위한 명령 일정, 메모리 레이아웃 변환, 작업 병합이 포함될 수 있습니다.
플래그를 사용한 제어: SparseCore 동작, 조정 매개변수, 최적화 전략의 여러 측면이 XLA 플래그 (예: 제한 추정의 경우 xla_sparse_core_estimate_max_ids, 디버깅의 경우 xla_sc_detect_nan)를 통해 노출되고 제어됩니다.

오픈소스 상태:

현재 Sparsecore 구현은 내부용이며 libtpu.so를 사용하여 제공됩니다.

오류 보고 및 진단:

SparseCore 구성 또는 리소스 제약 조건과 관련된 컴파일 실패는 XLA:TPU 컴파일 시간 오류로 나타나는 경우가 많습니다. 이러한 오류 메시지는 사용 가능한 SPMEM에 비해 너무 높은 한도가 설정되었거나 지원되지 않는 구성이 사용되는 등의 문제에 관한 유용한 정보를 제공할 수 있습니다.

7. 제한사항이 SparseCore의 테이블에 적용되는 방식

SparseCore에서 '한도'는 주로 사용 가능한 SparseCore에 샤딩 (분산)된 각 테이블의 파티션별 설정 두 가지를 참조하는 기본 구성 매개변수입니다.

max_ids_per_partition: 단일 컴퓨팅 단계 내에서 단일 SparseCore가 특정 테이블의 특정 파티션에 전송하거나 처리할 것으로 예상되는 총 ID 수(중복 포함)의 최대값을 정의합니다.
max_unique_ids_per_partition: 이는 단일 SparseCore가 전송하거나 처리할 것으로 예상되는 고유 ID의 최대 개수를 정의합니다.

물리적 테이블 레이아웃 및 처리로의 변환:

표 샤딩 전략: 삽입된 표는 일반적으로 시스템의 모든 SparseCore에서 '모듈로 샤딩'됩니다. 즉, 각 SparseCore는 각 테이블의 어휘 (행)의 고유한 하위 집합을 담당합니다. ID j는 일반적으로 k = j % num_total_sparse_cores와 같은 공식을 기반으로 SparseCore_k에 할당됩니다.
'파티션' 정의: 이 맥락에서 '파티션'은 단일 SparseCore가 조회를 처리하는 임베딩 테이블의 특정 세그먼트를 의미합니다.
SPMEM 버퍼 할당: 이러한 제한은 XLA 컴파일러가 기기 내 스크래치패드 메모리(SPMEM) 내에서 버퍼의 크기를 정적으로 조정하고 할당하는 데 사용됩니다. 버퍼는 지정된 max_ids 및 max_unique_ids 제한까지 지정된 파티션의 ID와 관련된 모든 필수 데이터를 처리할 수 있도록 SPMEM에 로드할 수 있도록 크기가 조정됩니다. 이는 특히 파티션 내 중복 ID를 줄이는 것과 같은 요소별이 아닌 계산 (예: 압축된 희소 행 (CSR) 표현을 만드는 경우)에 중요합니다. 이러한 계산에서는 해당 파티션의 ID에 대한 전체 관련 데이터 세트를 빠른 메모리에서 바로 사용할 수 있어야 합니다.
컴파일된 한도와 관찰된 한도:
- 관찰된 한도: 처리되는 입력 데이터를 기반으로 런타임 중에 각 파티션에 대해 발견된 실제 ID 수입니다.
- 관찰된 한도가 컴파일된 한도를 초과하면 ID 삭제 (allow_id_dropping가 사용 설정된 경우) 또는 오류가 발생할 수 있습니다.
한도 계산: 적절한 한도를 결정하는 프로세스에는 입력 데이터 분포에 대한 신중한 분석이 필요합니다. 특정 테이블 (T1라고 함. 더 큰 스택 테이블 T의 일부일 수 있음)의 경우 다음을 충족해야 합니다.
1. 입력 배치 (예: 형태가 [BatchSize, MaxSequenceLength]인 2D SparseTensor)는 처음에 사용 가능한 SparseCore에 걸쳐 분할됩니다. 예를 들어 TensorCore가 SparseCore 2개와 페어링된 경우 각 SparseCore는 [BatchSize/2, MaxSequenceLength] 모양의 하위 배치를 수신할 수 있습니다.
2. 이 하위 배치(sub-batch)는 COO 형식으로 변환되어 row_ids 및 col_ids를 생성합니다.
3. 동일한 샘플 내의 중복 ID (즉, 동일한 row_id 및 col_id가 있는 항목)가 삭제됩니다.
4. 샘플 내의 각 고유 col_id에 대해 이 ID를 담당하는 타겟 SparseCore는 mod-sharding 규칙 target_sc_id = col_id % num_total_sparse_cores을 사용하여 결정됩니다.
5. 각 target_sc_id로 전송되는 총 ID 수 (ids_per_sparse_core[target_sc_id]++)와 고유 ID 수 (target_sc_id의 고유성을 확인한 후의 unique_ids_per_sparse_core[target_sc_id]++)가 유지됩니다.
6. 그런 다음 표 T1의 max_ids_per_partition이 max(ids_per_sparse_core_array)로 설정됩니다.
7. 마찬가지로 테이블 T1의 max_unique_ids_per_partition이 max(unique_ids_per_sparse_core_array)로 설정됩니다.
8. 표 T1가 누적 표의 구성요소인 경우 모든 구성 표의 통계를 합산하기 전에 회전 또는 이동과 같은 추가 변환이 ID 분포에 적용될 수 있습니다. 이렇게 하면 칩 간에 부하를 분산하는 데 도움이 됩니다.

이러한 한도를 올바르게 설정하는 것은 균형을 맞추는 작업입니다. 한도를 낮추면 단계별로 처리해야 하는 데이터가 줄고 SPMEM 압력이 감소하므로 성능이 향상될 수 있지만, 너무 낮게 설정하면 미니 배치 수가 지나치게 많아지거나 원치 않는 ID 삭제가 발생할 수 있습니다.

8. 각 SparseCore의 통신 방식

특히 삽입 조회를 위한 ID 목록을 처리하는 맥락에서 SparseCore 통신은 다음과 같은 여러 조정된 메커니즘에 의존합니다.

모듈 샤딩 및 암시적 라우팅:
- 임베딩 테이블은 시스템의 모든 SparseCore에 걸쳐 모드 샤딩됩니다.
- 호스트가 입력 데이터 배치를 제공하면(이후 col_ids를 포함한 COO 형식으로 사전 처리됨) col_id 값이 특정 ID(target_sc_id = col_id % num_total_sparse_cores)를 담당하는 SparseCore를 결정하는 데 사용됩니다.
- 각 SparseCore는 할당된 어휘 파티션에 매핑되는 ID의 하위 집합만 효과적으로 수신하고 처리합니다. 호스트 전처리 단계는 각 SparseCore가 관련 ID를 쉽게 식별하고 작동할 수 있도록 데이터를 준비하는 데 중요합니다.
호스트별 데이터 분포:
- 호스트 전처리 로직은 전체 입력 배치에 파티션을 적용하고 row_ids 및 col_ids의 관련 부분을 각 SparseCore에서 직접 액세스할 수 있는 메모리 (HBM) 또는 SparseCore가 필요한 데이터를 가져올 공유 HBM에 배포합니다 (관련 기능이나 가중치가 있는 경우 함께).
Intra-SparseCore 처리:
- SparseCore는 지정된 테이블 파티션의 ID 집합을 수신하면 이러한 ID의 중복 삭제 및 해당 임베딩 벡터 수집과 같은 작업을 실행합니다. 이는 주로 SparseCore의 자체 타일 내에서 실행되고 로컬 SPMEM을 활용하는 로컬 컴퓨팅입니다.
SparseCore 간 통신 (All-to-All):
- 초기 처리 단계 (예: 임베딩 조회) 후에는 'all-to-all' 통신 패턴을 사용하여 SparseCore 간에 결과를 결합하거나 재분배할 수 있습니다 (예: 모든 원래 샘플 위치에 해당하는 입력을 예상하는 TensorCore 레이어에 활성화를 공급하기 전). 이는 원래 입력 배치에 병렬 처리를 위해 분산된 경우 활성화의 전체 세트를 재구성하는 데 중요합니다.
TensorCore와의 통신:
- SparseCore는 TensorCore와 통신하여 삽입 활성화 (정방향 전달 중)를 전송하고 그라데이션 (역방향 전달 중)을 수신합니다. 이 상호작용은 XLA 컴파일 프로그램에 의해 조정되며 중간 버퍼로 HBM이 자주 사용됩니다. 파이프라인 전략 (앞서 설명함)은 이 SC-TC 통신의 타이밍과 동기화에 큰 영향을 미칩니다.

기본적으로 적절한 SparseCore에 ID를 초기 '분배'하는 작업은 샤딩 스키마와 호스트 전처리 단계에서 대부분 처리합니다. 후속 통신에는 로컬 데이터에서 작동하는 SparseCore가 포함되며, TensorCore에서 추가 처리하기 전에 SparseCore 간에 데이터를 전역적으로 교환하거나 재정렬해야 하는 경우 all-to-all과 같은 집단 통신 작업이 이어질 수 있습니다.

9. SparseCore 메모리 관리

각 SparseCore는 여러 가지 고유한 유형의 메모리를 효율적으로 관리하여 계산을 실행합니다.

스크래치패드 메모리 (SPMEM):
- 네이처: 각 SparseCore에서만 사용할 수 있는 비교적 작지만 매우 빠른 로컬 SRAM입니다. SPMEM은 캐시가 아닙니다. 사용은 XLA 컴파일러에 의해 명시적으로 관리되고 오케스트레이션됩니다.
- 목적: SPMEM은 '기회에 따라 데이터를 스테이징'하는 데 사용됩니다. 여기에는 지속적인 SC 계산에 필요한 입력, 출력, 중간 결과가 포함됩니다. SPMEM에 데이터를 스테이징하면 HBM 액세스와 일반적으로 관련된 높은 지연 시간이 크게 줄어듭니다.
- 크기 조정: '한도' 섹션에서 설명한 대로 SPMEM 버퍼는 컴파일 시간에 정적으로 크기가 조정됩니다. 이 크기는 max_ids_per_partition 및 max_unique_ids_per_partition과 같은 매개변수를 기반으로 합니다. 이 정적 할당을 통해 테이블 파티션의 특정 작업 (예: CSR 감소)에 대해 파티션 ID에 필요한 모든 데이터 (정의된 한도까지)가 SPMEM에 맞게 됩니다.
- 컴파일러 최적화: XLA 컴파일러는 정교한 최적화를 통합하여 HBM 지연 시간을 효과적으로 숨기고 성능을 극대화하기 위해 SPMEM에 스테이징해야 하는 데이터의 양과 특정 데이터 요소를 정확하게 결정합니다.
- 동적 할당 제약 조건: SparseCore 컴파일러는 현재 동적 스크래치패드 할당을 지원하지 않습니다. 이는 한도를 신중하게 구성하여 정적 크기 조정이 얼마나 중요한지 보여줍니다.
고대역폭 메모리 (HBM):
- Nature: 모든 SparseCore, TensorCore, 호스트 시스템에서 액세스할 수 있는 대규모 공유 메모리 리소스입니다. 기본 삽입 테이블은 HBM에 저장됩니다.
- 스택 사용량: SparseCore 작업에는 제한된 SPMEM에 맞지 않거나 처리 파이프라인의 더 큰 단계 간에 전달해야 하는 중간 결과를 위해 HBM의 임시 스토리지가 필요한 경우가 많습니다. 정방향 및 역방향 패스 중 HBM 스택 사용량은 다음과 같이 추정할 수 있습니다.
  - 전달 패스 HBM 스택 (단일 테이블) ≈ (2 * feature_width + 1) * max_unique_nz_per_row * logical_replica_count * 4바이트
  - 역방향 전달 HBM 스택 (단일 테이블) ≈ 3 * feature_width * max_unique_nz_per_row * logical_replica_count * 4바이트
- 힙 사용량: HBM은 호스트에서 관리하는 힙도 수용합니다. 힙은 밀도 높은 레이어 가중치, 모델에서 사용되는 상수, 미리 가져온 입력 데이터와 같은 데이터를 저장합니다. 힙 사용량은 호스트가 데이터를 프리패치하는 단계 수에 따라 증가하는 경향이 있습니다 (maximum_parallel_iterations 플래그로 제어됨). 더 많은 프리패치를 사용하면 호스트-기기 전송과 기기 계산이 중복되어 성능이 향상될 수 있지만 더 많은 HBM이 소비됩니다.
- HBM 최적화를 위한 직렬화: xla_sc_num_serialized_tables_to_optimize_hbm 플래그는 특정 시점에 HBM 스택 메모리에 '라이브'로 유지되는 테이블의 데이터 수를 제어하는 메커니즘을 제공합니다. 이 숫자를 늘리면 더 많은 테이블의 처리가 효과적으로 직렬화되어 최대 HBM 스택 사용량이 줄어들 수 있지만 병렬 처리가 줄어들어 성능이 저하될 수 있습니다.
벡터 메모리 (VMEM):
- VMEM은 TC (TensorCore)에서만 사용하는 로컬 스크래치패드 메모리입니다. VMEM은 SparseCore에서 직접 관리하지는 않지만 SC가 주로 TensorCore를 통해 상호작용하는 메모리 생태계의 필수적인 부분입니다.

전반적인 메모리 관리 전략:

SparseCore의 핵심 메모리 관리 전략은 SparseCore 타일에서 활발하게 처리되는 '핫' 데이터에 작고 빠른 SPMEM을 사용하여 느린 HBM에 대한 액세스를 최소화하는 것입니다. 구성된 한도는 SPMEM이 오버플로되지 않도록 하는 기본 메커니즘입니다. HBM은 SPMEM 용량을 초과하거나 다양한 처리 단위 또는 파이프라인 단계에서 공유해야 하는 대규모 임베딩 테이블과 임시 데이터를 저장하는 데 사용됩니다. XLA 컴파일러는 이러한 아키텍처 원칙과 사용자가 구성한 제한에 따라 모든 데이터 이동과 버퍼 할당을 오케스트레이션합니다.

10. 성능 및 메모리 병목 현상

SparseCore로 최적의 성능을 달성하려면 잠재적인 병목 현상과 이를 해결하는 방법을 명확하게 이해해야 합니다. 이러한 문제는 호스트에서, SparseCore 자체 내에서 또는 TensorCore와의 상호작용에서 발생할 수 있습니다.

일반적인 성능 병목 현상:

호스트 병목 현상:
- 문제: 호스트 CPU가 데이터를 전처리하고 TPU에 충분히 빠르게 피드하지 못하여 SparseCore 및 TensorCore가 제대로 활용되지 않을 수 있습니다. 이는 자주 발생하는 성능 제한 요소입니다.
- 완화: 호스트 CPU 사용률과 입력 파이프라인 측정항목을 모니터링합니다. 호스트 측 데이터 로드 및 사전 처리 루틴을 최적화합니다 (COO 변환 팁 참고). maximum_parallel_iterations 플래그를 조정하여 데이터 미리 가져오기를 미세 조정합니다.
최적화되지 않은 TC/SC 동기화 (파이프라인 부족):
- 문제: TensorCore와 SparseCore 간의 파이프라인이 사용 중지되었거나 효율적으로 작동하지 않으면 한 단위가 다른 단위를 기다리는 데 상당한 시간을 소비하여 전체 시스템 처리량이 감소할 수 있습니다.
- 완화: 파이프라인이 사용 설정되어 있는지 확인합니다 (예: tf_xla_disable_full_embedding_pipelining = false 또는 이에 상응하는 항목).
한도로 인한 병목 현상:
- 문제:
  - 한도가 너무 낮음: 과도한 미니 배치 (엄격한 한도를 충족하기 위해 입력 배치를 여러 개의 작은 하위 배치로 분할)가 트리거될 수 있습니다. 이렇게 하면 정확성은 유지되지만 각 미니 배치에 처리 오버헤드가 발생하여 전체 실행 속도가 느려질 수 있습니다. allow_id_dropping가 true인 경우 너무 낮은 한도로 인해 ID가 삭제되어 모델 정확도에 영향을 미칠 수도 있습니다.
  - 한도가 너무 높음 (하지만 여전히 적합): 한도가 매우 높으면 미니 배치가 방지될 수 있지만 실제 데이터 특성이 이러한 최대값에 거의 도달하지 않는 경우 불필요하게 SPMEM 압력이 증가할 수 있습니다. 또한 엄격하게 필요한 것보다 더 큰 HBM 스택 사용으로 이어질 수도 있습니다.
  - 컴파일 실패: 구성된 제한에 사용 가능한 실제 메모리보다 더 많은 SPMEM 또는 HBM 스택이 필요한 경우 컴파일이 실패합니다.
- 완화: 한도가 올바르게 설정되어 있는지 확인합니다.
데이터 분포 왜곡:
- 문제: 특정 SparseCore 파티션이 다른 파티션에 비해 불균형적으로 많은 ID를 지속적으로 수신하는 경우 (ID 분포가 좋지 않음을 나타냄) 과부하된 SparseCore가 성능 병목 현상이 됩니다.
- 완화: 미니 배치 프로세스 중에 ID를 셔플하면 특히 '핫' 사용자 테이블이 있는 스택 테이블의 경우 이 문제를 완화할 수 있습니다. ID 분포를 신중하게 분석하여 적절하고 균형 잡힌 테이블별 한도를 설정하세요.
표 스태킹 문제:
- 문제:
  - 스택된 테이블이 너무 적음: ICI 지연 시간을 효과적으로 숨기거나 처리 오버헤드를 적절하게 줄이는 데 충분하지 않을 수 있습니다.
  - 테이블이 너무 많이 쌓임: 관리하기 어려워지거나 사용 가능한 리소스 한도를 초과하는 매우 큰 논리 테이블이 생성될 수 있습니다.
- 완화:
  - 스태킹에 적합한 테이블 수를 확인합니다. 일반적인 가이드라인에서는 스태킹을 위해 5~100개의 테이블을 사용하는 것이 적절하다고 제안합니다.
비효율적인 숫자/양자화:
- 문제: BF16 또는 양자화된 정수와 같은 낮은 정밀도 형식이 충분하고 계산 속도가 더 빠른 경우 전체 FP32 정밀도를 사용하면 성능 병목 현상이 발생할 수 있습니다.
- 완화: 정밀도가 낮은 옵션을 살펴봅니다. 하지만 양자화 자체에 오버헤드가 있으며 모델 정확도를 유지하려면 양자화 매개변수를 신중하게 조정해야 할 수 있습니다.
HBM 대역폭 포화:
- 문제: 매우 작은 특징 너비 (높은 패딩 오버헤드 발생), 비효율적인 메모리 액세스 패턴 또는 매우 많은 조회수로 인해 HBM으로 또는 HBM에서 데이터가 과도하게 이동하면 사용 가능한 HBM 대역폭이 포화될 수 있습니다.
- 완화: TPU 수를 조정하면 HBM 대역폭 포화를 방지할 수 있습니다.

일반적인 메모리 병목 현상:

SPMEM 오버플로 (컴파일 실패):
- 문제: max_ids_per_partition 및 max_unique_ids_per_partition가 너무 높게 설정된 경우 XLA 컴파일러가 충분한 SPMEM을 할당하지 못하여 "Fixed size allocations (...) do not fit in TileSpmem (...)"과 같은 컴파일 오류가 발생할 수 있습니다. 또한 타일 SPMEM 내에서 스테이징 수집 피연산자에 (sample_count * feature_width) / kNumTiles (여기서 kNumTiles은 SC당 타일 수)이 너무 크면 "Gather operand too large..."와 같은 오류가 발생할 수 있습니다.
- 완화: 배치 크기를 줄이거나 처리에 사용되는 칩 수를 늘립니다.
HBM 스택 오버플로 (런타임 또는 컴파일):
- 문제: feature_width, max_unique_nz_per_row, logical_replica_count의 조합으로 인해 사용 가능한 HBM을 초과하는 HBM 스택 메모리 요구사항이 발생하면 런타임 또는 컴파일 중에 메모리 부족 (OOM) 오류가 발생할 수 있습니다.
- 완화: 테이블 처리를 직렬화하여 HBM 스택 사용량을 줄이도록 xla_sc_num_serialized_tables_to_optimize_hbm 플래그를 조정합니다 (일반적으로 성능 비용이 발생함).
HBM 힙 소진:
- 문제: 주로 매우 큰 밀도층 가중치, 메모리에 저장된 수많은 상수 또는 지나치게 적극적인 입력 프리패치 (높은 maximum_parallel_iterations)로 인해 발생합니다.
- 완화: XProf Memory Viewer와 같은 도구를 사용하여 힙 사용량을 모니터링합니다.
패딩 오버헤드:
- 문제: 삽입 테이블이 특성 차원에서 32B 정렬 (8개의 부동 소수점과 동일)되도록 패딩됩니다. 따라서 작은 특징 너비 (예: 1개의 부동 소수점)는 상당한 패딩 오버헤드 (예: 할당된 버퍼 공간의 7/8이 패딩)를 발생시켜 HBM이 낭비됩니다. 테이블의 어휘 차원도 시스템의 SparseCore 수의 배수가 되도록 패딩됩니다. 하지만 어휘 크기가 충분히 큰 테이블의 경우 이 영향은 일반적으로 미미합니다.

성능과 메모리에 영향을 미치는 일반적인 요인:

토폴로지: 사용 가능한 칩의 수와 상호 연결 아키텍처입니다.
배치 크기: SparseCore당 sample_count에 직접적인 영향을 미치며, 이는 메모리 소비 및 컴퓨팅 부하에 영향을 미칩니다.
데이터 형식 지정: 최적의 성능을 위해서는 효율적인 온디바이스 데이터 레이아웃이 중요합니다.

11. SparseCore 프로필 분석

성능 프로필을 분석하는 것은 SparseCore 워크로드 내에서 병목 현상을 식별하고 최적화 기회를 찾는 데 중요한 단계입니다.

트레이스 획득:
- XProf와 같은 프로파일링 도구를 활용하여 모델이 학습되거나 추론을 실행하는 동안 자세한 실행 트레이스를 캡처합니다. 이 트레이스는 호스트, TensorCore, SparseCore에서 발생하는 작업의 타임라인을 제공합니다.
트레이스 뷰어 검사 (예: XProf 또는 TensorBoard):
- 호스트 활동: 호스트의 활동을 자세히 살펴봅니다. TPU 활동에 상당한 격차가 있나요? 이러한 격차는 호스트가 병목 현상으로 인해 데이터를 충분히 빠르게 제공하지 못함을 나타낼 수 있습니다. 입력 파이프라인의 성능을 분석합니다.
- TensorCore (TC) 및 SparseCore (SC) 활동:
  - TC와 SC의 실행 타임라인을 확인합니다. 병렬로 작동하여 효과적인 파이프라인을 나타내나요? 아니면 한 단위가 다른 단위를 기다리며 유휴 상태로 있는 기간이 긴가요?
  - SC와 TC 모두에서 가장 많은 시간을 소비하는 작업 (가장 오래 실행되는 작업)을 식별합니다.
  - 시각적 트레이스 출력(시간이 지남에 따라 다양한 작업을 나타내는 색상 블록을 표시하는 경우가 많음(예: TPU:0 SparseCore 1 (pid 1005)))은 지배적인 작업과 유휴 기간을 시각적으로 식별하는 데 매우 유용합니다.
- 단계 시간 분석: 전체 단계 시간을 관찰하고 호스트 처리, SC 계산, TC 계산 간에 어떻게 분산되거나 분류되는지 파악합니다.
메모리 분석 (XProf 메모리 뷰어):
- 힙 사용량: XProf의 '메모리 뷰어' 탭과 같은 도구를 사용하여 HBM 힙 사용량을 검사합니다. 이를 통해 대형 모델 가중치, 상수 또는 지나치게 적극적인 입력 프리패치가 과도한 양의 HBM을 소비하는지 확인할 수 있습니다. --vmodule=best_fit_allocator=1와 같은 플래그를 사용 설정하면 최대 힙 사용량의 로그가 제공될 수 있습니다.
- 스택 사용량 (간접): 직접 HBM 스택 프로파일링은 복잡할 수 있지만 메모리 부족 오류가 발생하고 힙 사용량이 적절해 보이면 HBM 스택 소진 (과도한 제한이나 기능 너비로 인한 경우가 많음)이 강력한 용의자입니다. HBM 스택 사용량에 제공된 수식을 사용하면 이를 추정할 수 있습니다.
특정 패턴 찾기:
- 미니 배치: 한도를 자주 초과하는 경우 트레이스에서 미니 배치 증거를 확인할 수 있습니다 (예: 전역 배치 크기에 예상되는 것보다 작은 SC 작업 수가 더 많음). 이는 로그에서 추론하거나 특정 작업의 호출 수를 관찰하여 확인할 수 있습니다.
- ID 삭제: ID 삭제가 사용 설정되어 있고 삭제가 발생하는 경우 시스템 로그에 이 사실이 표시될 수 있습니다. 이는 구성된 한도가 입력 데이터에 너무 제한적이라는 명확한 신호이기도 합니다.
- 컴파일 시간: 특히 피드백 기반 최적화 (FDO)가 사용 설정되어 있고 제한을 자주 조정하는 경우 컴파일 시간이 길어지면 전체 학습 시간에 상당한 오버헤드가 추가될 수 있습니다.
플래그 및 구성과 연관:
- 프로필에서 관찰된 동작을 SparseCore 구성 (한도 파일의 설정, XLA 플래그)과 연결합니다. 예를 들어 xla_sc_num_serialized_tables_to_optimize_hbm이 높은 값으로 설정된 경우 SC 성능은 느려지지만 HBM 스택 소비는 낮아질 수 있습니다.
반복 프로세스:
- 프로파일링은 반복적인 개선 프로세스인 경우가 많습니다. 특정 변경사항 (한도 조정, 기능 사용 설정 또는 사용 중지)을 적용하고 새 프로필을 캡처한 다음 이전 프로필과 비교하여 수정사항의 영향을 확인합니다.

12. 일반 디버깅 플래그

SparseCore 실행과 관련된 문제를 디버깅하는 데 도움이 되도록 여러 플래그를 사용 설정할 수 있습니다. 이러한 검사를 사용 설정하면 성능 페널티가 발생하는 경우가 많으므로 일반적으로 프로덕션 실행에서는 사용 중지해야 합니다.

ID 확인 (범위 외):
- 플래그: xla_sparse_core_enable_id_bound_check = true
- 목적: 입력 데이터의 삽입 ID가 지정된 삽입 테이블에 정의된 유효한 어휘 범위에 속하지 않는지 감지하기 위해 호스트 시스템에서 검사를 사용 설정합니다. 이렇게 하면 잘못되거나 손상된 입력 데이터와 관련된 문제를 포착하는 데 도움이 됩니다.
NaN 검사기:
- 플래그: xla_sc_detect_nan = true
- 목적: SparseCore에서 처리된 부동 소수점 데이터 내에서 NaN (숫자가 아님) 값의 감지를 사용 설정합니다. 다양한 컴파일러 패스의 입력 또는 출력에서 NaN이 감지되면 이 플래그로 인해 오류가 발생합니다. 이러한 오류는 일반적으로 NaN이 발생한 위치에 관한 정보를 제공합니다.
범위 검사기 (메모리 액세스):
- 플래그: xla_sc_assert_level=bounds
- 목적: 이 플래그는 동적 검사를 포함하도록 메모리 액세스 명령어 (예: VMEM 로드/저장 및 DMA 작업)를 다시 작성하는 ASAN (AddressSanitizer) 스타일 도구를 사용 설정합니다. 이러한 검사는 메모리 액세스가 타겟 메모리 영역의 할당된 경계 내에 있는지 확인합니다.
- 동작: 범위를 벗어난 메모리 액세스가 감지되면 실행이 실패합니다.
- 주의: 이 검사기는 검사기가 완전히 이해하지 못하는 복잡한 스트라이드 액세스 패턴과 같은 이유로 거짓양성을 생성할 수 있습니다. 이 변환은 백엔드 컴파일 프로세스의 후반 단계에서 적용됩니다.
버퍼 검사기 (메모리 손상):
- 플래그:
  - xla_tpu_buffer_contents_sanitizer_config='cores_to_sanitize: [TC, SC_SCS, SC_TILE], sanitizer_mode: LOCAL_ONLY'
  - xla_tpu_verify_launch_id_across_cores=true
- 목적: 이러한 플래그는 메모리 버퍼가 관련 없는 작업으로 인해 실수로 손상되거나 덮어쓰이지 않도록 하는 데 도움이 됩니다. 버퍼 소독기는 버퍼의 콘텐츠를 검사하여 예상치 못한 변경사항이 없는지 확인합니다.

13. 양자화 지원

SparseCore의 SparseDenseMatmulOp는 32비트 부동 소수점 (FP32) 및 정수 데이터 유형을 모두 사용하여 삽입 테이블에 대한 작업을 지원하도록 설계되었습니다. 모델 학습은 일반적으로 임베딩 테이블에 FP32 정밀도를 사용하여 실행되지만 사후 학습 양자화 (PTQ)를 적용할 수 있습니다. PTQ를 사용하면 추론에 더 낮은 정밀도 데이터 유형 (예: 8비트 정수)을 사용할 수 있으므로 성능이 향상되고 메모리 사용 공간이 줄어들 수 있습니다.

시뮬레이션된 양자화:

SparseDenseMatmulOp는 '시뮬레이션된 양자화'를 실행하도록 구성할 수 있습니다. 이 작동 모드에서는 임베딩 벡터가 먼저 낮은 정밀도로 양자화된 후 후속 계산에 사용되기 전에 다시 높은 정밀도 (예: FP32)로 역양자화됩니다. 이 기법을 사용하면 양자화 노이즈의 영향을 고려하면서 모델을 학습할 수 있습니다. 시뮬레이션된 양자화로 학습하면 추론을 위해 완전히 양자화되었을 때 최종 모델의 정확도를 높일 수 있습니다.

`SparseDenseMatmulOp`의 구성 속성 (양자화용):

quantization_config_num_buckets = 256
- 이 속성은 32비트 부동 소수점 숫자가 양자화될 개별 버킷 또는 수준의 수를 지정합니다. 예를 들어 8비트 정수로 양자화할 때는 일반적으로 2^8 =256개의 버킷을 지정합니다.
quantization_config_low = -X.X
- 이 속성은 양자화 범위의 최소 부동 소수점 값을 정의합니다. 이 지정된 최솟값 미만의 입력 값은 양자화 중에 이 최솟값으로 잘립니다.
quantization_config_high = Y.Y
- 이 속성은 양자화 범위의 최대 부동 소수점 값을 정의합니다. 이 지정된 최댓값을 초과하는 입력값은 양자화 중에 이 최댓값으로 잘립니다.

숫자 및 파이프라인 상호작용:

TensorCore와 SparseCore 간의 파이프라인이 사용 설정되어 있는지에 따라 모델의 수치 동작이 달라질 수 있습니다. 파이프라인이 활성 상태인 경우 SparseCore에서 처리한 그라데이션이 '오래된' (이전 반복에서) 것일 수 있습니다. 이는 양자화 프로세스와 상호작용하여 모델 학습 역학 또는 최종 정확도에 영향을 미칠 수 있습니다.

14. 예정된 기능 및 최근 개선사항

SparseCore 생태계는 지속적으로 개발되고 개선됩니다.

로드맵:

샘플 차원 미니 배치:
- 이는 기존 어휘 차원 미니 배치 기능을 보완하는 기능으로 계획되었습니다.
- 이렇게 하면 샘플 차원을 따라 삽입 입력을 추가로 파티셔닝할 수 있습니다. 이는 한 번에 샘플의 하위 집합에서 조회를 필터링하고 처리할 수 있는 온디바이스 루프를 도입하여 달성할 수 있습니다. 이러한 기능은 매우 큰 샘플별 ID 수를 관리하거나 처리 단위 간 부하 분산을 개선하는 데 유용할 수 있습니다.
행당 정수가 8개 미만인 삽입 지원 개선 (작은 기능 너비):
- 현재 설계에서는 8개의 부동 소수점 (32바이트에 해당) 미만인 임베딩 기능 너비에 상당한 패딩을 사용하는 경우가 많습니다. 이 패딩으로 인해 HBM이 낭비되고 컴퓨팅 리소스가 제대로 활용되지 않을 수 있습니다. 향후 개선사항은 작은 특징 측정기준이 있는 테이블의 이러한 비효율성을 완화하는 것을 목표로 합니다.

최근 개선사항:

HBM에서 피연산자 수집 스테이징:
- 이 최적화는 일부 수집 작업 입력 또는 출력을 더 큰 HBM에 스테이징할 수 있도록 하여 공유 스크래치패드 메모리 (SPMEM)의 압력을 줄이는 데 도움이 됩니다.
스택 메모리 사용량 감소:
- 전반적인 성능이나 처리량에 부정적인 영향을 미치지 않으면서 SparseCore 작업 중에 HBM 스택 메모리 소비를 줄이도록 개선사항이 구현되었습니다.

이러한 개선사항은 더 광범위한 스파스 워크로드에 대해 SparseCore의 성능, 메모리 효율성, 운영 유연성을 개선하는 데 중점을 둡니다.

대규모 임베딩 모델 (LEM)용 SparseCore 심층 분석 컬렉션을 사용해 정리하기 내 환경설정을 기준으로 콘텐츠를 저장하고 분류하세요.