Skip to content

Fixed build issues with Pytorch extensions #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wenchenvincent merged 1 commit into from
Nov 9, 2022
Merged

Fixed build issues with Pytorch extensions #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
313 changes: 227 additions & 86 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,33 +14,48 @@
from setuptools.command.build_ext import build_ext
from distutils.version import LooseVersion
from distutils.file_util import copy_file
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME, ROCM_HOME, IS_HIP_EXTENSION

path = os.path.dirname(os.path.realpath(__file__))

with open(path + "/VERSION", "r") as f:
te_version = f.readline()

def get_cuda_bare_metal_version(cuda_dir):
raw_output = subprocess.check_output(
[cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
)
if IS_HIP_EXTENSION:
raw_output = subprocess.check_output(
[cuda_dir + "/bin/hipcc", "--version"], universal_newlines=True
)
else:
raw_output = subprocess.check_output(
[cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
)
output = raw_output.split()
release_idx = output.index("release") + 1
if IS_HIP_EXTENSION:
release_idx = output.index("version:") + 1
else:
release_idx = output.index("release") + 1
release = output[release_idx].split(".")
bare_metal_major = release[0]
bare_metal_minor = release[1][0]
return raw_output, bare_metal_major, bare_metal_minor


def append_nvcc_threads(nvcc_extra_args):
_, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
return nvcc_extra_args + ["--threads", "4"]
if IS_HIP_EXTENSION:
_, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(ROCM_HOME)
##TODO: Figure out which hipcc version starts to support this parallel compilation
return nvcc_extra_args + ["-parallel-jobs=4"]
else:
_, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
if int(bare_metal_major) >= 11 and int(bare_metal_minor) >= 2:
return nvcc_extra_args + ["--threads", "4"]
return nvcc_extra_args


def extra_gencodes(cc_flag):
if IS_HIP_EXTENSION:
return
_, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(CUDA_HOME)
if int(bare_metal_major) >= 11:
cc_flag.append("-gencode")
Expand All @@ -51,21 +66,38 @@ def extra_gencodes(cc_flag):


def extra_compiler_flags():
return [
"-O3",
"-gencode",
"arch=compute_70,code=sm_70",
"-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT16_OPERATORS__",
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT162_OPERATORS__",
"-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
"-I./transformer_engine/common/layer_norm/",
"--expt-relaxed-constexpr",
"--expt-extended-lambda",
"--use_fast_math",
]
if IS_HIP_EXTENSION:
return [
"-O3",
#"-gencode",
#"arch=compute_70,code=sm_70",
"-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT16_OPERATORS__",
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT162_OPERATORS__",
"-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
"-I./transformer_engine/common/layer_norm/",
#"--expt-relaxed-constexpr",
#"--expt-extended-lambda",
#"--use_fast_math",
]
else:
return [
"-O3",
"-gencode",
"arch=compute_70,code=sm_70",
"-U__CUDA_NO_HALF_OPERATORS__",
"-U__CUDA_NO_HALF_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT16_OPERATORS__",
"-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
"-U__CUDA_NO_BFLOAT162_OPERATORS__",
"-U__CUDA_NO_BFLOAT162_CONVERSIONS__",
"-I./transformer_engine/common/layer_norm/",
"--expt-relaxed-constexpr",
"--expt-extended-lambda",
"--use_fast_math",
]


cc_flag = []
Expand Down Expand Up @@ -137,75 +169,184 @@ def __init__(self, name, cmake_path, sources, **kwargs):
)
)

ext_modules.append(
CUDAExtension(
name="scaled_upper_triang_masked_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
if IS_HIP_EXTENSION:
## This is a workaround for the issue of hipify causing duplicated object file name
## scaled_upper_triang_masked_softmax_cuda.cu --> scaled_upper_triang_masked_softmax_hip.hip
## scaled_upper_triang_masked_softmax.cpp --> scaled_upper_triang_masked_softmax_hip.cpp
## CMake or Ninja will automatically generate build rules using scale_upper_triang_masked_softmax_hip.o
## as the object file name for both source files, resulting in a conflict.
## Currently, I don't know a way to let CMake or Ninja use more unique names. Thus changed one source file
## name to work around this issue.
os.system('cp %s %s' %(
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax_cuda.cu",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax_cuda_cuda.cu",
)
)
)
ext_modules.append(
CUDAExtension(
name="scaled_upper_triang_masked_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax_cuda_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
)
)
)

ext_modules.append(
CUDAExtension(
name="scaled_masked_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
os.system('cp %s %s' %(
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax_cuda.cu",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax_cuda_cuda.cu",
)
)
)
ext_modules.append(
CUDAExtension(
name="scaled_masked_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax_cuda_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
)
)
)

ext_modules.append(
CUDAExtension(
name="scaled_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
os.system('cp %s %s' %(
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax_cuda.cu",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax_cuda_cuda.cu",
)
)
)
ext_modules.append(
CUDAExtension(
name="scaled_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax_cuda_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
)
)
else:
ext_modules.append(
CUDAExtension(
name="scaled_upper_triang_masked_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_upper_triang_masked_softmax_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
)
)
)

ext_modules.append(
CUDAExtension(
name="scaled_masked_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_masked_softmax_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
)
)

ext_modules.append(
CUDAExtension(
name="scaled_softmax_cuda",
sources=[
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax.cpp",
),
os.path.join(
path,
"transformer_engine/pytorch/csrc/fused_softmax/scaled_softmax_cuda.cu",
),
],
extra_compile_args={
"cxx": ["-O3"],
"nvcc": append_nvcc_threads(extra_compiler_flags() + cc_flag),
},
include_dirs=[
os.path.join(path, "transformer_engine/pytorch/csrc/fused_softmax")
],
)
)

def get_cmake_bin():
cmake_bin = "cmake"
Expand Down
4 changes: 3 additions & 1 deletion transformer_engine/common/recipe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@
"""This module provides predefined FP8 recipes."""
from __future__ import annotations
from enum import Enum
from typing import Literal, Optional, Union, Callable, NamedTuple
#from typing import Literal, Optional, Union, Callable, NamedTuple
from typing import Optional, Union, Callable, NamedTuple
from typing_extensions import Literal
from pydantic.dataclasses import dataclass


Expand Down
6 changes: 6 additions & 0 deletions transformer_engine/pytorch/csrc/common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,13 @@
#include <torch/torch.h>
#include <cuda.h>
#include <cuda_runtime.h>

#ifndef USE_ROCM
#include <cuda_bf16.h>
#else
#include <hip/hip_bfloat16.h>
#endif

@HaiShaw HaiShaw Nov 9, 2022

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

could add hip_bfloat16.h in else case.

@wenchenvincent wenchenvincent Nov 9, 2022

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Hai. This was before merging your enablement of bf16 and I forgot to modify it. Will modify and merge.

#include <stdexcept>
#include <memory>
#include <iomanip>
Expand Down