未验证 提交 26e5ecab 编写于 作者: J Jeff Rasley 提交者: GitHub

Initial commit of setup

上级 b8301c2c
"""
Copyright 2020 The Microsoft DeepSpeed Team
DeepSpeed library
Create a new wheel via the following command: python setup.py bdist_wheel
The wheel will be located at: dist/*.whl
"""
import torch
from setuptools import setup, find_packages
from torch.utils.cpp_extension import CUDAExtension, BuildExtension
cmdclass = {}
ext_modules = []
cmdclass['build_ext'] = BuildExtension
ext_modules.append(
CUDAExtension(name='fused_lamb_cuda',
sources=['csrc/fused_lamb_cuda.cpp',
'csrc/fused_lamb_cuda_kernel.cu'],
extra_compile_args={
'cxx': [
'-O3',
],
'nvcc': ['-O3',
'--use_fast_math']
}))
setup(name='deepspeed',
version='0.1',
description='DeepSpeed library',
author='DeepSpeed Team',
author_email='deepspeed@microsoft.com',
url='http://aka.ms/deepspeed',
packages=find_packages(exclude=[
"docker",
"third_party",
"csrc"
]),
scripts=[
'bin/deepspeed',
'bin/ds'
],
classifiers=['Programming Language :: Python :: 3.6'],
ext_modules=ext_modules,
cmdclass=cmdclass)
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册