From 26e5ecab9b6524844a0df454de802da870730df3 Mon Sep 17 00:00:00 2001
From: Jeff Rasley <jerasley@microsoft.com>
Date: Fri, 31 Jan 2020 15:57:11 -0800
Subject: [PATCH] Initial commit of setup

---
 setup.py | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)
 create mode 100644 setup.py

diff --git a/setup.py b/setup.py
new file mode 100644
index 000000000..4061c060e
--- /dev/null
+++ b/setup.py
@@ -0,0 +1,48 @@
+"""
+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)