Support for CDNA 2 Architecture for Frontier

[zixianwang2022]

Hi.

I have been trying to use the library at scale (up to 1024 gpus) on Frontier (MI250X). Is it possible to have support for CDNA 2 architecture as well? Especially the GroupGEMM kernel, which currently utilizes fp8, which is unsupportive on MI250X.

[xiaobochen-amd]

@zixianwang2022 Hi, could you share which operators are required for your workload?
We haven't ported Primus-Turbo to CDNA2 (MI250X) yet, and knowing your needs will help us assess the effort to support it.

[zixianwang2022]

Hi @xiaobochen-amd ,

Thank you for your thoughtfulness. Currently we are looking to utilize just the bf16/fp16 ck-group-gemm for our fine-grained moe training (more than 128 experts).

We are able now to just run the groupgemm kernel alone on MI250X by adding a new instance here and rebuild after some modifications to the setup.py to include CDNA2. We could also open a PR if you would like.

The rest of the features (DeepEP, etc) are not as urgent as the groupgemm for our use-case, but it would be great to have it available in the near future.

[xiaobochen-amd]

@zixianwang2022

We’d be very happy to have a PR from you! ：）

Also, may I check if you are from AMD as well? If so, we could sync offline with the Team for more efficient communication.

[zixianwang2022]

My colleague is working on it, and he will open it once everything is ready to go.

I am currently a phd student at UIUC advised my Prof. Minjia Zhang, but I did intern at AMD twice : ). I am very happy to assist the project you are working on and any potential collaborations in the future! My email is zixianwang4@illinois.edu if you need anything from our lab.

[xiaobochen-amd]

We are very happy to collaborate with community contributors.
In the near future, Primus-Turbo will focus on three major areas: performance optimization, compute-communication overlap, and low-precision.
If you have any topics you’re interested, please feel free to reach out.

My email: Xiaobo.Chen@amd.com

[zixianwang2022]

Hi @xiaobochen-amd ,

Is there a publicly available Dockerfile to build the primus-turbo env instead of pulling from docker hub?

[xiaobochen-amd]

@zixianwang2022
Turbo does not have a Dockerfile at the moment. However, the Primus Docker image we build includes Turbo by default.
You can refer to this Dockerfile as an example:
https://github.com/ROCm/unified-training-dockers/blob/pipeline/primus_training/docker/Dockerfile#L366

[zixianwang2022]

@zixianwang2022 Turbo does not have a Dockerfile at the moment. However, the Primus Docker image we build includes Turbo by default. You can refer to this Dockerfile as an example: https://github.com/ROCm/unified-training-dockers/blob/pipeline/primus_training/docker/Dockerfile#L366

I don't have access to the github, and it appears as 404 on my end. Would it be possible to publicly access it?

[xiaobochen-amd]

@zixianwang2022 Turbo does not have a Dockerfile at the moment. However, the Primus Docker image we build includes Turbo by default. You can refer to this Dockerfile as an example: https://github.com/ROCm/unified-training-dockers/blob/pipeline/primus_training/docker/Dockerfile#L366

I don't have access to the github, and it appears as 404 on my end. Would it be possible to publicly access it?

Sorry, I just realized this is an internal repo.

You can refer to the following code:

ARG TURBO_REPO=https://github.com/AMD-AGI/Primus-Turbo.git
ARG TURBO_COMMIT=main
ENV MAX_JOBS=128
ENV GPU_ARCHS="gfx942;gfx950"

RUN git clone ${TURBO_REPO} --recursive \
    && cd Primus-Turbo \
    && git checkout ${TURBO_COMMIT} \
    && git submodule update --init --recursive \
    && pip3 install -r requirements.txt \
    && pip3 install --no-build-isolation . -v \
    && cd ..

[zixianwang2022]

Thanks a lot!