UCC（Unified Communication X）を使ってみる

UCCとは、集合通信操作APIであり、例えばOpenMPIから呼び出すことができます。

構成
OpenMPIから呼び出す
どの通信レイヤ（Transport Layer）が使用されるか？
- 通信レイヤのスコア変更
参考

構成

./src
|-- coll_patterns  -----  集団通信パターンの実装（例: リング、ツリーベースの通信など）
|-- coll_score  --------  集団通信のスコアリング（性能評価）を行うロジック
|-- components
|   |-- base
|   |-- cl  ------------  集団通信（Collectives Layer）の実装（例: basic, hier, doca_urom）
|   |   |-- basic
|   |   |-- doca_urom  -  NVIDIA DOCA（Data Center on a Chip Architecture）向けのプラグイン
|   |   `-- hier  ------  階層的アルゴリズム
|   |-- ec  ------------  エラーコレクション（Error Collection）関連（例: base, cpu, cuda, rocm）
|   |-- mc  ------------  メモリコピー（Memory Copy）関連（例: base, cpu, cuda, rocm）
|   |-- tl  ------------  通信レイヤ（Transport Layer）の実装
|   |   |-- cuda
|   |   |   |-- allgather
|   |   |   |-- ...
|   |   |   `-- reduce_scatterv
|   |   |-- mlx5
|   |   |   |-- alltoall
|   |   |   `-- mcast
|   |   |       `-- p2p
|   |   |-- nccl
|   |   |   `-- allgatherv
|   |   |-- rccl
|   |   |   `-- allgatherv
|   |   |-- self
|   |   |-- sharp
|   |   `-- ucp  -------  OpenUCXコンポーネントで、低レイテンシかつ高スループットの通信を実現するための抽象化された通信レイヤ
|   |       |-- coll_plugins
|   |       |   `-- example
|   |       |-- allgather
|   |       |-- ...
|   |       `-- scatterv
|   `-- topo
|-- core  --------------  UCCのコアロジック（初期化、終了処理、基本的な操作）
|-- schedule  ----------  集団通信操作のスケジューリングロジック
|-- ucc  ---------------  api: UCCの外部API（アプリケーションが利用するインターフェース）
|   `-- api
`-- utils
    |-- arch  ----------  arch: CPUアーキテクチャ（例: x86_64, aarch64）に特化した実装
    `-- profile  -------  プロファイリング関連コード

ROCm（Radeon Open Compute）は、AMDが開発したオープンソースの高性能コンピューティング（HPC）およびGPUコンピューティング向けのソフトウェアプラットフォーム

OpenMPIから呼び出す

仮想環境でテスト。

Dockerfile

FROM ubuntu:24.04
RUN apt update
RUN DEBIAN_FRONTEND=noninteractive TZ=Asia/Tokyo apt install -y tzdata build-essential openssh-server vim git flex zlib1g zlib1g-dev autoconf automake libtool libnuma1
RUN rm -rf /var/lib/apt/lists/*

RUN mkdir /var/run/sshd
RUN echo 'root:tmp' | chpasswd
RUN echo 'PermitRootLogin yes' >> /etc/ssh/sshd_config
RUN echo 'PasswordAuthentication yes' >> /etc/ssh/sshd_config

WORKDIR /app

RUN git clone https://github.com/openucx/ucx.git
RUN cd ucx && ./autogen.sh && ./configure --prefix=/usr && make -j && make install

RUN git clone https://github.com/openucx/ucc.git
RUN cd ucc && ./autogen.sh && ./configure --prefix=/usr && make -j && make install

RUN git clone https://github.com/open-mpi/ompi.git
RUN cd ompi && git submodule update --init --recursive && ./autogen.pl && ./configure --with-ucx --with-ucc --prefix=/usr && make -j && make install

CMD ["/usr/sbin/sshd", "-D"]

docker build -f Dockerfile . -t ompi_dev:latest
docker run -d -shm-size 1g --name ompi_dev ompi_dev:latest

コンテナ作成時に、共有メモリ(shared memory)を増やプログラムの実行でエラーが。

Cのテストコード

#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

int main(int argc, char** argv) {
    int rank, size;
    int send_data, *recv_data;
    int error_code;

    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    send_data = rank + 1;
    recv_data = (int*)malloc(size * sizeof(int));

    MPI_Allgather(&send_data, 1, MPI_INT, recv_data, 1, MPI_INT, MPI_COMM_WORLD);

    printf("Rank %d received data: ", rank);
    for (int i = 0; i < size; i++) printf("%d ", recv_data[i]);
    printf("\n");

    free(recv_data);
    MPI_Finalize();

    return 0;
}

実行コマンド

mpicc test_allgather.c
OMPI_MCA_coll_ucc_enable=1 \
OMPI_MCA_coll_ucc_priority=100 \
UCC_LOG_LEVEL=debug \
mpirun --allow-run-as-root -n 8 ./a.out

もしくは

mpicc test_allgather.c
UCC_LOG_LEVEL=debug \
mpirun --allow-run-as-root -n 8 \
--mca coll_ucc_enable 1 \
--mca coll_ucc_priority 100 ./a.out

どの通信レイヤ（Transport Layer）が使用されるか？

UCC_LOG_LEVELをINFOで実行すると次のようなログが出力されます。これはそれぞれの集団通信について{サイズ}:通信レイヤ:スコアのスコアマップを表しており、集団通信とデータサイズに応じてどの通信レイヤが使用されるかをこれで確認することができます。

       ucc_team.c:471  UCC  INFO  ===== COLL_SCORE_MAP (team_id 32768, size 2) =====
ucc_coll_score_map.c:206  UCC  INFO  Allgather:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..4095}:TL_UCP:10 {4K..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Allgatherv:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Allreduce:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..4095}:TL_UCP:10 {4K..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Alltoall:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..257}:TL_UCP:10 {258..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Alltoallv:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Barrier:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Bcast:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Fanin:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Fanout:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Gather:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Gatherv:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Reduce:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Reduce_scatter:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Reduce_scatterv:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
ucc_coll_score_map.c:206  UCC  INFO  Scatterv:
ucc_coll_score_map.c:206  UCC  INFO        Host: {0..inf}:TL_UCP:10 
       ucc_team.c:474  UCC  INFO  ================================================

通信レイヤのスコア変更

にあるように UCC_TL_<NAME>_TUNE=token1#token2#...#tokenN, 環境変数を使用します。

例えば、UCC_TL_UCP_TUNE=allgather:0 とすればallgatherでUCP通信レイヤが使用されなくなります。また、UCC_TL_UCP_TUNE=allgather:inf とすればallgatherでUCP通信レイヤが必ず使用されることになります。

参考

直接的な参考ではないですが、下記の連載が勉強になる。

freak-da.hatenablog.com

Sabrou-mal サブロウ丸

主にプログラミングと数学

UCC（Unified Communication X）を使ってみる

構成

OpenMPIから呼び出す

どの通信レイヤ（Transport Layer）が使用されるか？

通信レイヤのスコア変更

参考