사례 연구: Chip과 Cluster 리덕션

Chip과 Cluster 리덕션은 서로 다른 Chip 또는 Cluster 차원을 따라 저장된 부분 값을 결합합니다. 예제는 두 차원 모두에 대해 ReduceScatter, AllGather, 이 둘의 조합, 직접 AllReduce를 구현합니다. 별도의 Chip 예제는 Butterfly AllReduce를 구현합니다.

형태 변환

[dimension], [local] 표기법은 차원 형태와 각 차원 인덱스에 로컬로 저장된 형태를 구분합니다. R을 리덕션할 차원, S를 분배하거나 수집할 로컬 축, P...와 Q...를 각각 S 앞과 뒤의 로컬 형태라고 합시다. ReduceScatter와 AllGather에서는 R과 S의 크기가 N으로 같아야 합니다. Broadcast(N)은 해당 차원을 따라 있는 N개의 인덱스가 모두 같은 값을 가진다는 뜻입니다. 표는 논리적 결과를 나타내며, 제거되지 않는 로컬 축의 순서를 유지합니다.

ReduceScatter는 R을 리덕션하고 그 결과인 S 축을 분배합니다. 차원 인덱스 s는 S의 좌표 s를 보유합니다. AllGather는 모든 차원 인덱스에서 모든 S 좌표를 수집하여 그 분배를 되돌립니다. AllReduce는 R을 리덕션하고 결과를 복제하면서 로컬 형태를 유지합니다. 직접 순환 리덕션, ReduceScatter 후 AllGather를 수행하는 방식, 버터플라이 교환은 모두 같은 AllReduce 형태 변환을 가집니다. 이 변환은 Chip 차원과 Cluster 차원 모두에 적용됩니다.

부동소수점 리덕션 순서

Broadcast(N)은 논리적 배치에 대한 보장입니다. 순환 리덕션은 0 + 1 + 2 + 3과 1 + 2 + 3 + 0 같은 순서를 사용하므로, 독립적으로 계산된 부동소수점 복제본은 서로 다를 수 있습니다. ReduceScatter 후의 AllGather는 각 샤드를 한 번만 계산한 뒤 수집하므로 해당 복제본은 동일합니다. 예제는 반올림 순서가 적용되지 않는 i32를 사용합니다.

재분배 프리미티브

예제는 지연된 SRAM 재분배 계획을 사용합니다. to_dm 이전의 shuffle 단계와 slice 단계는 하나의 DMA로 융합됩니다. 반면 asymmetric_chip_slice와 asymmetric_cluster_slice를 직접 호출하면 어느 차원으로도 데이터를 이동하지 않고 서브 컨텍스트를 사용합니다.

동기화

컴파일러는 to_dm이 최종 재분배 계획을 구체화할 때 그 계획에서 동기화를 도출합니다. 소스 Chip이 바뀌면 ChipSync가 삽입되고, 소스 Cluster가 바뀌면 ClusterSync가 삽입됩니다. 두 차원을 모두 바꾸는 융합 계획에는 두 동기화 프로토콜이 모두 적용됩니다. 로컬 DmTensor 슬라이싱은 둘 중 어느 것도 추가하지 않습니다.

컴파일러는 원격 DMA 접근 전에 목적지 준비 동기화를, 접근 후에 완료 동기화를 생성합니다. 자기 자신으로 향하는 Chip 경로는 Chip 간 송신이나 수신을 생성하지 않습니다.

이 개수는 예제에 명시된 재분배 호출을 나타냅니다. HBM 전송과 로컬 Vector Engine 작업에도 각자의 의존성이 있지만, Chip 재분배 동기화나 Cluster 재분배 동기화를 추가하지는 않습니다.

Chip 예제

예제에서는 Chip 4개로 구성된 A 축과 B 축, 그리고 256바이트 [C, D] 페이로드를 사용합니다. Chip 인덱스와 샤드 인덱스의 범위는 0부터 3까지입니다.

use furiosa_opt_std::prelude::*;

axes![A = 4, B = 4, C = 8, D = 8, I = 2, X = 2];

리덕션 예제는 다음 Vector Engine 덧셈을 공통으로 사용합니다:

/// Adds two `[C, D]` DM tensors with the Vector Engine.
pub fn add_pair(
    device: &mut Device,
    lhs: &DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]>,
    rhs: &DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]>,
    output: &mut DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]>,
) {
    device
        .main
        .begin_interleaved::<I, _, _, _, _, _>(lhs.view(), rhs.view())
        .fetch::<m![C, I], m![D]>()
        .collect::<m![C, I], m![D]>()
        .vector_init()
        .vector_intra_slice_unzip::<I, m![C]>()
        .vector_clip_zip(ClipBinaryOpI32::AddFxp)
        .vector_final()
        .commit_trim::<m![D]>()
        .commit_view(output.view_mut());
}

ReduceScatter

입력은 A를 여러 Chip에 걸쳐 배치하고 [B, C, D]는 로컬에 유지합니다. 각 chip_shuffle().chip_slice().to_dm() 체인은 B 좌표 하나를 선택하여 대상 Chip으로 라우팅합니다. 그런 다음 세 개의 Vector Engine 체인이 라우팅된 [C, D] 값 네 개를 합산합니다. 여기서 in[p, s]는 소스 Chip p가 샤드 s에 대해 가진 값이며, out[s] = sum_p in[p, s]입니다.

/// Chip `i` keeps shard `i` of the sum over all chips.
#[device(chip = 4)]
pub fn reduce_scatter(
    device: &mut Device,
    hbm: &HbmTensor<i32, m![A], m![B, C, D]>,
) -> HbmTensor<i32, m![A], m![C, D]> {
    let full: DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![B, C, D]> = hbm.to_dm(&mut device.tdma);
    reduce_scatter_dm(device, &full).to_hbm(&mut device.tdma)
}

결과는 샤딩된 상태로 유지됩니다. 대상 Chip s가 out[s]를 소유합니다.

AllGather

입력은 소스 Chip s에 [C, D] 샤드 하나를 제공하며, 이를 in[s]로 표기합니다. to_dm() 한 번으로 Chip 차원을 로컬 요소 형태로 옮기고, 그 형태를 모든 Chip에 브로드캐스트합니다.

/// Every chip receives every shard along `A`.
#[device(chip = 4)]
pub fn all_gather(device: &mut Device, hbm: &HbmTensor<i32, m![A], m![C, D]>) -> HbmTensor<i32, m![4], m![A, C, D]> {
    let shard: DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]> = hbm.to_dm(&mut device.tdma);
    all_gather_dm::<m![A], m![A, C, D]>(&mut device.tdma, &shard).to_hbm(&mut device.tdma)
}

모든 대상이 모든 샤드를 받습니다.

ReduceScatter 이후의 AllGather

이 예제는 ReduceScatter 결과를 AllGather에 바로 입력합니다.

/// Composes ReduceScatter and AllGather so every chip receives the reduced shards.
#[device(chip = 4)]
pub fn all_gather_after_reduce_scatter(
    device: &mut Device,
    hbm: &HbmTensor<i32, m![A], m![B, C, D]>,
) -> HbmTensor<i32, m![4], m![B, C, D]> {
    let full: DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![B, C, D]> = hbm.to_dm(&mut device.tdma);
    let shard = reduce_scatter_dm(device, &full);
    // Target chip b owns coordinate b, so changing the axis name preserves wire order.
    let shard = unsafe { shard.reshape::<m![B], m![1 # 2], m![1 # 256], m![C, D]>() };
    all_gather_dm::<m![B], m![B, C, D]>(&mut device.tdma, &shard).to_hbm(&mut device.tdma)
}

최종 배치는 위에 나온 AllGather 상태와 같습니다.

AllReduce

입력은 Chip p에 형태가 [C, D]인 부분합 in[p] 하나를 제공합니다. 원본 텐서가 항등 회전을 제공하고, 세 번의 Chip 셔플이 나머지 순환 회전을 제공합니다. 네 회전을 모두 더하면 모든 대상 Chip에 out[target] = sum_i in[i]가 남습니다.

/// Sums four cyclic chip rotations so every chip receives the full reduction.
#[device(chip = 4)]
pub fn all_reduce(device: &mut Device, hbm: &HbmTensor<i32, m![A], m![C, D]>) -> HbmTensor<i32, m![A], m![C, D]> {
    let round0: DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]> = hbm.to_dm(&mut device.tdma);
    let round1 = round0.chip_shuffle([1, 2, 3, 0]).to_dm::<m![C, D]>(&mut device.tdma);
    let round2 = round0.chip_shuffle([2, 3, 0, 1]).to_dm::<m![C, D]>(&mut device.tdma);
    let round3 = round0.chip_shuffle([3, 0, 1, 2]).to_dm::<m![C, D]>(&mut device.tdma);

    let mut sum01 = DmTensor::new();
    let mut sum012 = DmTensor::new();
    let mut sum = DmTensor::new();
    add_pair(device, &round0, &round1, &mut sum01);
    add_pair(device, &round2, &sum01, &mut sum012);
    add_pair(device, &round3, &sum012, &mut sum);
    sum.to_hbm(&mut device.tdma)
}

Butterfly AllReduce

입력은 Chip p에 형태가 [C, D]인 부분 값 in[p] 하나를 제공합니다. 셔플 후 더하기를 두 라운드 거치면 모든 Chip에 sum_i in[i]가 남습니다. 여기서 i는 모든 소스 chip을 가리킵니다.

/// Reduces across chips with XOR-paired butterfly rounds.
#[device(chip = 4)]
pub fn butterfly_all_reduce(
    device: &mut Device,
    hbm: &HbmTensor<i32, m![A], m![C, D]>,
) -> HbmTensor<i32, m![A], m![C, D]> {
    let partials: DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]> = hbm.to_dm(&mut device.tdma);
    let swapped = partials.chip_shuffle([1, 0, 3, 2]).to_dm::<m![C, D]>(&mut device.tdma);
    let mut first: DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]> = DmTensor::new();
    add_pair(device, &swapped, &partials, &mut first);

    let swapped = first.chip_shuffle([2, 3, 0, 1]).to_dm::<m![C, D]>(&mut device.tdma);
    let mut sum: DmTensor<i32, m![A], m![1 # 2], m![1 # 256], m![C, D]> = DmTensor::new();
    add_pair(device, &swapped, &first, &mut sum);
    sum.to_hbm(&mut device.tdma)
}

Cluster 예제

Cluster 예제에서는 동일한 256바이트 [C, D] 페이로드를 갖는, 요소가 두 개인 A 축과 B 축을 사용합니다. Cluster 인덱스와 샤드 인덱스의 범위는 0부터 1까지입니다.

use furiosa_opt_std::prelude::*;

axes![A = 2, B = 2, C = 8, D = 8, I = 2];

Cluster 예제에서는 이에 대응하는 두 Cluster 간 덧셈을 사용합니다:

/// Adds two `[C, D]` DM tensors with the Vector Engine.
pub fn add_pair(
    device: &mut Device,
    lhs: &DmTensor<i32, m![1], m![A], m![1 # 256], m![C, D]>,
    rhs: &DmTensor<i32, m![1], m![A], m![1 # 256], m![C, D]>,
    output: &mut DmTensor<i32, m![1], m![A], m![1 # 256], m![C, D]>,
) {
    device
        .main
        .begin_interleaved::<I, _, _, _, _, _>(lhs.view(), rhs.view())
        .fetch::<m![C, I], m![D]>()
        .collect::<m![C, I], m![D]>()
        .vector_init()
        .vector_intra_slice_unzip::<I, m![C]>()
        .vector_clip_zip(ClipBinaryOpI32::AddFxp)
        .vector_final()
        .commit_trim::<m![D]>()
        .commit_view(output.view_mut());
}

ReduceScatter

첫 번째 로컬 cluster_slice는 B 좌표 [0, 1]을 선택하며, Cluster마다 값을 하나씩 선택합니다. 두 번째는 [1, 0]을 선택하고, cluster_swap().to_dm()은 해당 값들을 반대편 Cluster로 이동합니다. 그런 다음 Vector Engine 체인이 로컬 [C, D] 값과 이동된 [C, D] 값을 더합니다. 여기서 in[c, s]는 샤드 s에 대한 소스 Cluster c의 값이며, out[s] = sum_c in[c, s]입니다.

/// Cluster `i` keeps shard `i` of the sum over both clusters.
#[device(chip = 1)]
pub fn reduce_scatter(
    device: &mut Device,
    hbm: &HbmTensor<i32, m![1], m![A, B, C, D]>,
) -> HbmTensor<i32, m![1], m![A, C, D]> {
    let full: DmTensor<i32, m![1], m![A], m![1 # 256], m![B, C, D]> = hbm.to_dm(&mut device.tdma);
    reduce_scatter_dm(device, &full).to_hbm(&mut device.tdma)
}

AllGather

입력은 소스 Cluster s에 in[s]로 표기되는 [C, D] 샤드 하나를 제공합니다. to_dm() 한 번으로 Cluster 차원을 로컬 요소 형태로 이동하고, 그 형태를 두 Cluster 모두에 브로드캐스트합니다.

/// Both clusters receive every shard along `A`.
#[device(chip = 1)]
pub fn all_gather(device: &mut Device, hbm: &HbmTensor<i32, m![1], m![A, C, D]>) -> HbmTensor<i32, m![1], m![A, C, D]> {
    let shard: DmTensor<i32, m![1], m![A], m![1 # 256], m![C, D]> = hbm.to_dm(&mut device.tdma);
    all_gather_dm::<m![A], m![A, C, D]>(&mut device.tdma, &shard).to_hbm(&mut device.tdma)
}

ReduceScatter 후 AllGather

이 예제는 Cluster ReduceScatter 결과를 AllGather에 바로 입력합니다.

/// Composes ReduceScatter and AllGather so both clusters receive the reduced shards.
#[device(chip = 1)]
pub fn all_gather_after_reduce_scatter(
    device: &mut Device,
    hbm: &HbmTensor<i32, m![1], m![A, B, C, D]>,
) -> HbmTensor<i32, m![1], m![B, C, D]> {
    let full: DmTensor<i32, m![1], m![A], m![1 # 256], m![B, C, D]> = hbm.to_dm(&mut device.tdma);
    let shard = reduce_scatter_dm(device, &full);
    // Target cluster b owns coordinate b, so changing the axis name preserves wire order.
    let shard = unsafe { shard.reshape::<m![1], m![B], m![1 # 256], m![C, D]>() };
    all_gather_dm::<m![B], m![B, C, D]>(&mut device.tdma, &shard).to_hbm(&mut device.tdma)
}

AllReduce

입력은 Cluster c에 형태가 [C, D]인 부분합 in[c] 하나를 제공합니다. 한 번의 Cluster 스왑으로 두 부분합을 교환하고, 한 번의 덧셈으로 두 Cluster 모두에 완전한 리덕션 결과가 남습니다.

/// Adds the local and swapped cluster partials so both clusters receive the full reduction.
#[device(chip = 1)]
pub fn all_reduce(device: &mut Device, hbm: &HbmTensor<i32, m![1], m![A, C, D]>) -> HbmTensor<i32, m![1], m![A, C, D]> {
    let partials: DmTensor<i32, m![1], m![A], m![1 # 256], m![C, D]> = hbm.to_dm(&mut device.tdma);
    let remote = partials.cluster_swap().to_dm::<m![C, D]>(&mut device.tdma);

    let mut sum: DmTensor<i32, m![1], m![A], m![1 # 256], m![C, D]> = DmTensor::new();
    add_pair(device, &partials, &remote, &mut sum);
    sum.to_hbm(&mut device.tdma)
}

DMA 레이아웃 요구 사항

정렬, 활성 인덱스, 256바이트 연속 구간 요구 사항은 재분배 연산을 참조하십시오. 이 예제들은 m![B, C, D]를 사용하며, 그 내부의 m![C, D] 구간은 정확히 256바이트의 i32 데이터입니다.