Skip to content
体验新版
项目
组织
正在加载...
登录
切换导航
打开侧边栏
Greenplum
DeepSpeed
提交
240ea97b
D
DeepSpeed
项目概览
Greenplum
/
DeepSpeed
上一次同步 大约 1 年
通知
10
Star
0
Fork
0
代码
文件
提交
分支
Tags
贡献者
分支图
Diff
Issue
0
列表
看板
标记
里程碑
合并请求
0
DevOps
流水线
流水线任务
计划
Wiki
0
Wiki
分析
仓库
DevOps
项目成员
Pages
D
DeepSpeed
项目概览
项目概览
详情
发布
仓库
仓库
文件
提交
分支
标签
贡献者
分支图
比较
Issue
0
Issue
0
列表
看板
标记
里程碑
合并请求
0
合并请求
0
Pages
DevOps
DevOps
流水线
流水线任务
计划
分析
分析
仓库分析
DevOps
Wiki
0
Wiki
成员
成员
收起侧边栏
关闭侧边栏
动态
分支图
创建新Issue
流水线任务
提交
Issue看板
体验新版 GitCode,发现更多精彩内容 >>
未验证
提交
240ea97b
编写于
9月 10, 2020
作者:
J
Jeff Rasley
提交者:
GitHub
9月 10, 2020
浏览文件
操作
浏览文件
下载
电子邮件补丁
差异文件
only add 1bit adam reqs if mpi is installed, update cond build for cpu-adam (#400)
上级
b29229bf
变更
1
隐藏空白更改
内联
并排
Showing
1 changed file
with
22 addition
and
16 deletion
+22
-16
setup.py
setup.py
+22
-16
未找到文件。
setup.py
浏览文件 @
240ea97b
...
...
@@ -10,8 +10,10 @@ The wheel will be located at: dist/*.whl
import
os
import
torch
import
shutil
import
subprocess
import
warnings
import
cpufeature
from
setuptools
import
setup
,
find_packages
from
torch.utils.cpp_extension
import
CUDAExtension
,
BuildExtension
,
CppExtension
...
...
@@ -27,47 +29,53 @@ install_requires = fetch_requirements('requirements/requirements.txt')
dev_requires
=
fetch_requirements
(
'requirements/requirements-dev.txt'
)
sparse_attn_requires
=
fetch_requirements
(
'requirements/requirements-sparse-attn.txt'
)
onebit_adam_requires
=
fetch_requirements
(
'requirements/requirements-1bit-adam.txt'
)
# If MPI is available add 1bit-adam requirements
if
torch
.
cuda
.
is_available
():
onebit_adam_requires
.
append
(
f
"cupy-cuda
{
torch
.
version
.
cuda
.
replace
(
'.'
,
''
)[:
3
]
}
"
)
install_requires
+=
onebit_adam_requires
if
shutil
.
which
(
'ompi_info'
)
or
shutil
.
which
(
'mpiname'
):
onebit_adam_requires
=
fetch_requirements
(
'requirements/requirements-1bit-adam.txt'
)
onebit_adam_requires
.
append
(
f
"cupy-cuda
{
torch
.
version
.
cuda
.
replace
(
'.'
,
''
)[:
3
]
}
"
)
install_requires
+=
onebit_adam_requires
# Constants for each op
LAMB
=
"lamb"
TRANSFORMER
=
"transformer"
SPARSE_ATTN
=
"sparse-attn"
ADAM
=
"cpu-adam"
CPU_
ADAM
=
"cpu-adam"
# Build environment variables for custom builds
DS_BUILD_LAMB_MASK
=
1
DS_BUILD_TRANSFORMER_MASK
=
10
DS_BUILD_SPARSE_ATTN_MASK
=
100
DS_BUILD_ADAM_MASK
=
1000
DS_BUILD_
CPU_
ADAM_MASK
=
1000
DS_BUILD_AVX512_MASK
=
10000
# Allow for build_cuda to turn on or off all ops
DS_BUILD_ALL_OPS
=
DS_BUILD_LAMB_MASK
|
DS_BUILD_TRANSFORMER_MASK
|
DS_BUILD_SPARSE_ATTN_MASK
|
DS_BUILD_ADAM_MASK
|
DS_BUILD_AVX512_MASK
DS_BUILD_ALL_OPS
=
DS_BUILD_LAMB_MASK
|
DS_BUILD_TRANSFORMER_MASK
|
DS_BUILD_SPARSE_ATTN_MASK
|
DS_BUILD_
CPU_
ADAM_MASK
|
DS_BUILD_AVX512_MASK
DS_BUILD_CUDA
=
int
(
os
.
environ
.
get
(
'DS_BUILD_CUDA'
,
1
))
*
DS_BUILD_ALL_OPS
# Set default of each op based on if build_cuda is set
OP_DEFAULT
=
DS_BUILD_CUDA
==
DS_BUILD_ALL_OPS
DS_BUILD_ADAM
=
int
(
os
.
environ
.
get
(
'DS_BUILD_ADAM'
,
OP_DEFAULT
))
*
DS_BUILD_ADAM_MASK
DS_BUILD_CPU_ADAM
=
int
(
os
.
environ
.
get
(
'DS_BUILD_CPU_ADAM'
,
OP_DEFAULT
))
*
DS_BUILD_CPU_ADAM_MASK
DS_BUILD_LAMB
=
int
(
os
.
environ
.
get
(
'DS_BUILD_LAMB'
,
OP_DEFAULT
))
*
DS_BUILD_LAMB_MASK
DS_BUILD_TRANSFORMER
=
int
(
os
.
environ
.
get
(
'DS_BUILD_TRANSFORMER'
,
OP_DEFAULT
))
*
DS_BUILD_TRANSFORMER_MASK
DS_BUILD_SPARSE_ATTN
=
int
(
os
.
environ
.
get
(
'DS_BUILD_SPARSE_ATTN'
,
0
))
*
DS_BUILD_SPARSE_ATTN_MASK
DS_BUILD_AVX512
=
int
(
os
.
environ
.
get
(
'DS_BUILD_AVX512'
,
0
))
*
DS_BUILD_AVX512_MASK
OP_DEFAULT
))
*
DS_BUILD_SPARSE_ATTN_MASK
DS_BUILD_AVX512
=
int
(
os
.
environ
.
get
(
'DS_BUILD_AVX512'
,
cpufeature
.
CPUFeature
[
'AVX512f'
]))
*
DS_BUILD_AVX512_MASK
# Final effective mask is the bitwise OR of each op
BUILD_MASK
=
(
DS_BUILD_LAMB
|
DS_BUILD_TRANSFORMER
|
DS_BUILD_SPARSE_ATTN
|
DS_BUILD_ADAM
)
|
DS_BUILD_
CPU_
ADAM
)
install_ops
=
dict
.
fromkeys
([
LAMB
,
TRANSFORMER
,
SPARSE_ATTN
,
ADAM
],
False
)
install_ops
=
dict
.
fromkeys
([
LAMB
,
TRANSFORMER
,
SPARSE_ATTN
,
CPU_
ADAM
],
False
)
if
BUILD_MASK
&
DS_BUILD_LAMB
:
install_ops
[
LAMB
]
=
True
if
BUILD_MASK
&
DS_BUILD_ADAM
:
install_ops
[
ADAM
]
=
True
if
BUILD_MASK
&
DS_BUILD_
CPU_
ADAM
:
install_ops
[
CPU_
ADAM
]
=
True
if
BUILD_MASK
&
DS_BUILD_TRANSFORMER
:
install_ops
[
TRANSFORMER
]
=
True
if
BUILD_MASK
&
DS_BUILD_SPARSE_ATTN
:
...
...
@@ -103,9 +111,7 @@ if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4):
version_ge_1_5
=
[
'-DVERSION_GE_1_5'
]
version_dependent_macros
=
version_ge_1_1
+
version_ge_1_3
+
version_ge_1_5
import
cpufeature
cpu_info
=
cpufeature
.
CPUFeature
SIMD_WIDTH
=
''
if
cpu_info
[
'AVX512f'
]
and
DS_BUILD_AVX512
:
SIMD_WIDTH
=
'-D__AVX512__'
...
...
@@ -133,7 +139,7 @@ if BUILD_MASK & DS_BUILD_LAMB:
}))
## Adam ##
if
BUILD_MASK
&
DS_BUILD_ADAM
:
if
BUILD_MASK
&
DS_BUILD_
CPU_
ADAM
:
ext_modules
.
append
(
CUDAExtension
(
name
=
'deepspeed.ops.adam.cpu_adam_op'
,
sources
=
[
...
...
编辑
预览
Markdown
is supported
0%
请重试
或
添加新附件
.
添加附件
取消
You are about to add
0
people
to the discussion. Proceed with caution.
先完成此消息的编辑!
取消
想要评论请
注册
或
登录