Multi-GPU-Arch pre-compilation of operators not supported

[Flamefire]

Part of host code relies on __CUDA_ARCH__ which must not be done. When compiling for multiple GPU architectures that can cause different evaluations leading to ODR violations, or the CUDA equivalent of that.

I noticed this in a test: pytest DeepSpeed/tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py -k 'test_DS4Sci_EvoformerAttention[tensor_shape0-dtype0]'

Depending on the order of CUDA archs in TORCH_CUDA_ARCH_LIST the test fails with cudaInvalidDeviceFunc (in CUDA 12.8 or later it is cudaErrorInvalidResourceHandle instead)

Reproducible by compiling DeepSpeed with pre-compiled operators:
TORCH_CUDA_ARCH_LIST='8.0;7.0' DS_BUILD_OPS=0 DS_BUILD_EVOFORMER_ATTN=1 pip install .

Then running the above test on A100 (CUDA CC 8.0)
Swapping the archs to ``TORCH_CUDA_ARCH_LIST='7.0;8.0'` makes it succeed. However I expect it will then fail on SM 7.0 devices.

I traced this to attention_impl_template and attention_back_impl_template which both have the same issue:
This declaration: typename std::enable_if<!CheckArch<arch, scalar_t>::value>::type attention_impl_template

CheckArch::value depends on the value of __CUDA_ARCH__:

DeepSpeed/csrc/deepspeed4science/evoformer_attn/gemm_kernel_utils.h

Line 49 in dbc1b07

static constexpr bool compiler_cc = arch::kMinComputeCapability * 10 <= __CUDA_ARCH__;

So depending on which compile iteration is done it will yield once true and once false. So when linking the multiple object files together the calling function (attention_impl) will have differing implementations -> ODR violation.
The order of -gencode flags to nvcc seems to also affect the ordering in the fatbin file which could be another reason why it may or may not find the correct function.

I made a small reproducer to show the issue:

#include <cuda_runtime.h>
#include <cstdio>
#include <cstdlib>

template<typename T>
__global__ void test_kernel(T *out) {}

template<typename arch>
struct Check {
  static constexpr bool value =
#if defined(__CUDA_ARCH__)
    __CUDA_ARCH__ >= 800;
#else
    true;
#endif
};

template <typename arch>
typename std::enable_if<!Check<arch>::value>::type foo() {}
template <typename arch>
typename std::enable_if<Check<arch>::value>::type foo()
{
  cudaFuncAttributes attr;
  auto func = test_kernel<float>;
  auto e = cudaFuncGetAttributes(&attr, func);
  if (e != cudaSuccess) {
    fprintf(stderr, "%s\n", cudaGetErrorString(e));
    std::exit(1);
  }
}

int main() {
    foo<bool>();
    printf("SUCCESS\n");
    return 0;
}

• nvcc repro.cu -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_70,code=sm_70 -o repro_80_first && ./repro_80_first FAILS
• nvcc repro.cu -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_80,code=sm_80 -o repro_70_first && ./repro_70_first WORKS

My understanding is that __CUDA_ARCH__ is meant to be used in __host__ __device__ or __device__ code and must not affect host code.
I only isolated this single issue but there might be other instances where __CUDA_ARCH__ is used incorrectly, especially through secondary uses like #if __CUDA_ARCH__ --> #define FOO and later in host code #ifdef FOO

For JIT compiled operations this isn't usually an issue because there is usually only one type of GPU available on the current system so only a single -gencode will be passed:

DeepSpeed/op_builder/builder.py

Lines 635 to 643 in dbc1b07

	if self.jit_mode:
	# Compile for underlying architectures since we know those at runtime
	for i in range(torch.cuda.device_count()):
	CC_MAJOR, CC_MINOR = torch.cuda.get_device_capability(i)
	cc = f"{CC_MAJOR}.{CC_MINOR}"
	if cc not in ccs:
	ccs.append(cc)
	ccs = sorted(ccs)
	ccs[-1] += '+PTX'

Although it may cause problems.

DeepSpeed/op_builder/builder.py

Lines 565 to 572 in dbc1b07

	# Torch will try and apply whatever CCs are in the arch list at compile time,
	# we have already set the intended targets ourselves we know that will be
	# needed at runtime. This prevents CC collisions such as multiple __half
	# implementations. Stash arch list to reset after build.
	torch_arch_list = None
	if "TORCH_CUDA_ARCH_LIST" in os.environ:
	torch_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST")
	os.environ["TORCH_CUDA_ARCH_LIST"] = ""

Unsetting TORCH_CUDA_ARCH_LIST may cause PyTorch to add all archs:

/tmp/lib/python3.11/site-packages/torch/utils/cpp_extension.py:1967: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.

And indeed: The JIT nvcc commandline shows multiple -gencode=arch=compute_80,code=compute_80 flags. I counted 4, 2 after --threads=8, likely added by DeepSpeed.

[tohtana]

Hi @Flamefire,
From my understanding, the inconsistency between DS_EVOFORMER_GPU_ARCH and TORCH_CUDA_ARCH_LIST is causing this issue.
I’ve opened #7872 to address this. Could you please check if it resolves the problem on your end?

[Flamefire]

Is my understanding of the code correct: DS_EVOFORMER_GPU_ARCH=7.0 && TORCH_CUDA_ARCH_LIST=7.0;8.0 will be accepted and build for both architectures but only the 7.0 code paths? That might work.

Maybe a static_assert at places where __CUDA_ARCH__ is used could catch any remaining or new uses of wrong -DGPU_ARCH

Why does DS_EVOFORMER_GPU_ARCH exist? Couldn't it just use the lowest value of TORCH_CUDA_ARCH_LIST?

[tohtana]

Great questions, here is my understanding and the intention in #7872.

• DS_EVOFORMER_GPU_ARCH=7.0 with TORCH_CUDA_ARCH_LIST='7.0;8.0' is accepted. It emits both sm_70 and sm_80 slices, but Evoformer is compiled with GPU_ARCH=70 family semantics (not Ampere-specialized semantics). So this is a compatibility-first mode and can work for fp16. But it will lose Ampere-specific behavior.
• static_assert around __CUDA_ARCH__/GPU_ARCH is reasonabl. The fix in Fix Evoformer arch filtering consistency for mixed targets (#7863) #7872 tried to address the root trigger at build time.
• DS_EVOFORMER_GPU_ARCH exists to control Evoformer kernel-family selection explicitly (including JIT/cross-compile/no-local-GPU cases). Using lowest value in TORCH_CUDA_ARCH_LIST would implicitly downgrade kernel family in mixed lists (e.g. 7.0;8.0 => forced 70), which wouldn't use Ampere-specific implemetation even on A100.

[Flamefire]

• Using lowest value in TORCH_CUDA_ARCH_LIST would implicitly downgrade kernel family in mixed lists (e.g. 7.0;8.0 => forced 70), which wouldn't use Ampere-specific implemetation even on A100.

It would indeed use the sm70 implementation (template) but at least compiled for A100. Not sure how much this looses.

If I read your PR right than the current behavior of DS_EVOFORMER_GPU_ARCH=8.0 with TORCH_CUDA_ARCH_LIST='7.0;8.0' is that it doesn't work on 7.0 GPUs as no kernels for that will be compiled. So it is (at build time) silently ignoring 7.0 in TORCH_CUDA_ARCH_LIST and error at runtime.

Not sure what the best option here is. Maybe require setting DS_EVOFORMER_GPU_ARCH when precompiling to make this more visible and mention that in the documentation. For JIT compilation the local GPU specs can be automatically used. At least that seems what the other DeepSpeed ops do.

[tohtana]

Thank you, @Flamefire. I think your concern is valid.
But the issue is that we need a non-trivial change in the template dispatch to address the root cause.

With #7863, build logs now show:

[WARNING] Filtered compute capabilities ['7.0'] below Evoformer floor 80

when mixed architectures are specified. So it's no longer fully silent. I think it is aligned with your suggestion (though users can still miss it).
At least, the fix also addresses the original problematic case: TORCH_CUDA_ARCH_LIST='8.0;7.0' DS_BUILD_OPS=0 DS_BUILD_EVOFORMER_ATTN=1 pip install .

Some clarification on behavior:

It would indeed use the sm70 implementation (template) but at least compiled for A100. Not sure how much this looses.

I'm not sure if I got your point, but I don't think it is a good choice to lose Ampere specialization by default.

For JIT compilation the local GPU specs can be automatically used. At least that seems what the other DeepSpeed ops do.

I think the current code works this way (as long as there aren't heterogeneous GPUs). Both DS_EVOFORMER_GPU_ARCH and TORCH_CUDA_ARCH_LIST are taken from the local GPUs in the env by default.

[Flamefire]

So it's no longer fully silent. I think it is aligned with your suggestion (though users can still miss it).

That's what I fear. I ran into this when packaging DeepSpeed for an HPC system where the same installation may be used for multiple GPU gens. Everything else "simply works" with TORCH_CUDA_ARCH_LIST='8.0;7.0'
If it wasn't for an inexplicable error in one of the tests I wouldn't have taken a second look.
And in the (semi-)automated workflow a warning like that would likely go unnoticed unless something fails hard later.

But I'm not sure anything else uses completely different kernel template specialisations for different GPUs.
My first impression was that it would compile kernels for all elements of TORCH_CUDA_ARCH_LIST and choose at runtime which kernel to call. But it is only if GPU >= DS_EVOFORMER_GPU_ARCH single_kernel; else error

I'm not sure if I got your point, but I don't think it is a good choice to lose Ampere specialisation by default.

With TORCH_CUDA_ARCH_LIST='8.0;7.0 DS_EVOFORMER_GPU_ARCH=7.0` it will instantiate the arch70 kernel and compile it for sm70 and sm80. So nvcc can choose "better" instructions when this fatbinary is run on an A100.
But it will not use the different implementation done here for A100

[tohtana]

Yes, that is a good summary of what we have right now, and I agree with your concern for HPC packaging workflows.

I'm not sure if we currently have the bandwidth for a larger dispatch redesign. One workaround would be adding a "strict mode" to raise a hard error when requested targets are pruned. (e.g. add DS_EVOFORMER_STRICT_ARCH env var). Do you think it would help?

[Flamefire]

Yes that would indeed help.

FWIW: I contacted NVIDIA support about this issue. So far it seems to be invalid to use __CUDA_ARCH__ here at all, i.e. for selecting kernel templates. So it might be better to use just GPU_ARCH.
But I'm still waiting for a clarification on this specific use case.