Skip to content

Commit 2360d73

Browse files
committed
Merge branch 'pip_package' into 'main'
Add pip package for megatron.core See merge request ADLR/megatron-lm!598
2 parents a07695f + f2b2fde commit 2360d73

File tree

5 files changed

+136
-8
lines changed

5 files changed

+136
-8
lines changed

megatron/core/README.md

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Megatron Core is a library for efficient and scalable training of transformer based models.

megatron/core/package_info.py

+23
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
2+
3+
4+
MAJOR = 0
5+
MINOR = 1
6+
PATCH = 0
7+
PRE_RELEASE = ''
8+
9+
# Use the following formatting: (major, minor, patch, pre-release)
10+
VERSION = (MAJOR, MINOR, PATCH, PRE_RELEASE)
11+
12+
__shortversion__ = '.'.join(map(str, VERSION[:3]))
13+
__version__ = '.'.join(map(str, VERSION[:3])) + ''.join(VERSION[3:])
14+
15+
__package_name__ = 'megatron_core'
16+
__contact_names__ = 'NVIDIA'
17+
__contact_emails__ = '[email protected]' # use NeMo Email
18+
__homepage__ = 'https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/' # use NeMo homepage
19+
__repository_url__ = 'https://github.com/NVIDIA/Megatron-LM/megatron/core'
20+
__download_url__ = 'https://github.com/NVIDIA/Megatron-LM/releases'
21+
__description__ = 'Megatron Core - a library for efficient and scalable training of transformer based models'
22+
__license__ = 'BSD-3'
23+
__keywords__ = 'deep learning, machine learning, gpu, NLP, NLU, language, transformer, nvidia, pytorch, torch'

megatron/core/requirements.txt

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
torch

megatron/core/utils.py

+2
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
2+
13
"""Utility functions used throughout Megatron core"""
24
from functools import reduce
35
import operator

setup.py

+109-8
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,111 @@
11
from setuptools import setup, find_packages
22

3-
setup(
4-
name="megatron.core",
5-
version="0.1",
6-
description="Core components of Megatron.",
7-
packages=find_packages(
8-
include=("megatron.core")
9-
)
10-
)
3+
"""Setup for pip package."""
4+
5+
import importlib.util
6+
import os
7+
import setuptools
8+
9+
spec = importlib.util.spec_from_file_location('package_info', 'megatron/core/package_info.py')
10+
package_info = importlib.util.module_from_spec(spec)
11+
spec.loader.exec_module(package_info)
12+
13+
14+
__contact_emails__ = package_info.__contact_emails__
15+
__contact_names__ = package_info.__contact_names__
16+
__description__ = package_info.__description__
17+
__download_url__ = package_info.__download_url__
18+
__homepage__ = package_info.__homepage__
19+
__keywords__ = package_info.__keywords__
20+
__license__ = package_info.__license__
21+
__package_name__ = package_info.__package_name__
22+
__repository_url__ = package_info.__repository_url__
23+
__version__ = package_info.__version__
24+
25+
26+
if os.path.exists('megatron/core/README.md'):
27+
with open("megatron/core/README.md", "r", encoding='utf-8') as fh:
28+
long_description = fh.read()
29+
long_description_content_type = "text/markdown"
30+
31+
else:
32+
long_description = 'See ' + __homepage__
33+
long_description_content_type = "text/plain"
34+
35+
36+
###############################################################################
37+
# Dependency Loading #
38+
# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #
39+
40+
def req_file(filename, folder="megatron/core"):
41+
with open(os.path.join(folder, filename), encoding='utf-8') as f:
42+
content = f.readlines()
43+
# you may also want to remove whitespace characters
44+
# Example: `\n` at the end of each line
45+
return [x.strip() for x in content]
46+
47+
install_requires = req_file("requirements.txt")
48+
49+
###############################################################################
50+
51+
setuptools.setup(
52+
name=__package_name__,
53+
# Versions should comply with PEP440. For a discussion on single-sourcing
54+
# the version across setup.py and the project code, see
55+
# https://packaging.python.org/en/latest/single_source_version.html
56+
version=__version__,
57+
description=__description__,
58+
long_description=long_description,
59+
long_description_content_type=long_description_content_type,
60+
# The project's main homepage.
61+
url=__repository_url__,
62+
download_url=__download_url__,
63+
# Author details
64+
author=__contact_names__,
65+
author_email=__contact_emails__,
66+
# maintainer Details
67+
maintainer=__contact_names__,
68+
maintainer_email=__contact_emails__,
69+
# The licence under which the project is released
70+
license=__license__,
71+
classifiers=[
72+
# How mature is this project? Common values are
73+
# 1 - Planning
74+
# 2 - Pre-Alpha
75+
# 3 - Alpha
76+
# 4 - Beta
77+
# 5 - Production/Stable
78+
# 6 - Mature
79+
# 7 - Inactive
80+
'Development Status :: 5 - Production/Stable',
81+
# Indicate who your project is intended for
82+
'Intended Audience :: Developers',
83+
'Intended Audience :: Science/Research',
84+
'Intended Audience :: Information Technology',
85+
# Indicate what your project relates to
86+
'Topic :: Scientific/Engineering',
87+
'Topic :: Scientific/Engineering :: Mathematics',
88+
'Topic :: Scientific/Engineering :: Image Recognition',
89+
'Topic :: Scientific/Engineering :: Artificial Intelligence',
90+
'Topic :: Software Development :: Libraries',
91+
'Topic :: Software Development :: Libraries :: Python Modules',
92+
'Topic :: Utilities',
93+
# Pick your license as you wish (should match "license" above)
94+
'License :: OSI Approved :: BSD License',
95+
# Supported python versions
96+
'Programming Language :: Python :: 3',
97+
'Programming Language :: Python :: 3.8',
98+
'Programming Language :: Python :: 3.9',
99+
# Additional Setting
100+
'Environment :: Console',
101+
'Natural Language :: English',
102+
'Operating System :: OS Independent',
103+
],
104+
packages=['megatron.core', 'megatron.core.pipeline_parallel', 'megatron.core.tensor_parallel'],
105+
install_requires=install_requires,
106+
107+
# Add in any packaged data.
108+
include_package_data=True,
109+
# PyPI package information.
110+
keywords=__keywords__,
111+
)

0 commit comments

Comments
 (0)