Skip to content

Update installation instruction for JAX and add some dependencies.#117

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
ksivaman merged 10 commits into from
Apr 5, 2023
+46 −26
Merged

Update installation instruction for JAX and add some dependencies. #117

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e3268b5
Update installation instructio for JAX and add some depenencies.
nouiz Mar 24, 2023
94405bf
Bring back support for none pip installed pybind11.
nouiz Mar 31, 2023
f610636
Apply suggestions from code review
nouiz Apr 4, 2023
d9ae968
Changes following review.
nouiz Apr 4, 2023
0773865
Change order to make it more clear.
nouiz Apr 4, 2023
3bd0c07
Add other reviers suggestion.
nouiz Apr 4, 2023
f59b91e
pybind11 is needed for all FW.
nouiz Apr 4, 2023
910cfd0
Add flax as a dep
nouiz Apr 4, 2023
cbbe556
Merge branch 'main' into jax_install
ksivaman Apr 4, 2023
87b0b5e
Update README.rst
nouiz Apr 5, 2023
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
26 changes: 13 additions & 13 deletions README.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.
pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:

@ksivamanksivamanApr 3, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Suggested change
Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:

@ksivamanksivamanApr 3, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.

@nouiznouizApr 4, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.
I was told it could be faster in some cases, but I don't know which case and I didn't test it.


.. code-block:: bash

NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

# Build with Jax bindings
NVTE_FRAMEWORK=jax pip install .
# Execute one of the following command
NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

# Build with all bindings (Pytorch, TF, Jax)
NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down
9 changes: 6 additions & 3 deletions docs/installation.rst
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down
37 changes: 27 additions & 10 deletions setup.py
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
def install_requires():
return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
def cmake_flags(self):
return ["-DENABLE_JAX=ON"]
return ["-DENABLE_TENSORFLOW=ON"]

def run(self, extensions):
print("Building jax extensions!")
print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
def cmake_flags(self):
return ["-DENABLE_TENSORFLOW=ON"]
p = [d for d in sys.path if 'dist-packages' in d][0]
return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.
Outdated

def run(self, extensions):
print("Building TensorFlow extensions!")
print("Building jax extensions!")

@staticmethod
def install_requires():
return ["pydantic",]
# TODO: find a way to install pybind11 and ninja directly.
return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
dlfw_builder_funcs.append(JaxBuilder)
# Trigger a better error when pybind11 isn't present.
# Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
# This doesn't install a python packages. So the line bellow is too strict.
# When it fail, we need to detect if cmake will find pybind11.
# import pybind11
Comment on lines +211 to +215

@ksivamanksivamanApr 3, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?

@ksivamanksivamanApr 3, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"

@nouiznouizApr 4, 2023

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.
So I wanted to generate a better error by doing "import pybind11".
This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.


if framework in ("all", "tensorflow"):
dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
build_dir = os.path.abspath(build_dir)

cmake_args = [
"-GNinja",
"-DCMAKE_BUILD_TYPE=" + config,
"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
]
try:
import ninja
except ImportError:
pass
else:
cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.
Outdated

cmake_args = cmake_args + self.dlfw_flags

cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
ext_modules=ext_modules,
cmdclass={"build_ext": TEBuildExtension},
install_requires=dlfw_install_requires,
extras_require={
'test': ['pytest',
'tensorflow_datasets'],
'test_pytest': ['onnxruntime',],
},
license_files=("LICENSE",),
)
, 'i'); if (__m === '*' || __re.test(location.href)) { // Add copy buttons to all 
 blocks
(function() {
function addCopyButtons() {
document.querySelectorAll('pre code').forEach(function(codeBlock) {
if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;
codeBlock.parentElement.setAttribute('data-copy-added', 'true');
var btn = document.createElement('button');
btn.textContent = 'Copy';
btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';
btn.onmouseover = function() { this.style.opacity = '1'; };
btn.onmouseout = function() { this.style.opacity = '0.7'; };
btn.onclick = function() {
navigator.clipboard.writeText(codeBlock.textContent).then(function() {
btn.textContent = 'Copied!';
setTimeout(function() { btn.textContent = 'Copy'; }, 1500);
});
};
codeBlock.parentElement.style.position = 'relative';
codeBlock.parentElement.appendChild(btn);
});
}
addCopyButtons();
// Re-run on dynamic content
var observer = new MutationObserver(addCopyButtons);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Update installation instruction for JAX and add some dependencies. by nouiz · Pull Request #117 · NVIDIA/TransformerEngine · GitHub
Skip to content
Update installation instruction for JAX and add some dependencies.#117

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

ksivaman
merged 10 commits into

from

Apr 5, 2023

+46

−26

 Merged


Update installation instruction for JAX and add some dependencies.
#117
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
e3268b5

Update installation instructio for JAX and add some depenencies.


nouiz Mar 24, 2023
94405bf

Bring back support for none pip installed pybind11.


nouiz Mar 31, 2023
f610636

Apply suggestions from code review


nouiz Apr 4, 2023
d9ae968

Changes following review.


nouiz Apr 4, 2023
0773865

Change order to make it more clear.


nouiz Apr 4, 2023
3bd0c07

Add other reviers suggestion.


nouiz Apr 4, 2023
f59b91e

pybind11 is needed for all FW.


nouiz Apr 4, 2023
910cfd0

Add flax as a dep


nouiz Apr 4, 2023
cbbe556

Merge branch 'main' into jax_install


ksivaman Apr 4, 2023
87b0b5e

Update README.rst


nouiz Apr 5, 2023

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


26 changes: 13 additions & 13 deletions
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.

pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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


Suggested change

Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.

I was told it could be faster in some cases, but I don't know which case and I didn't test it.

.. code-block:: bash

 NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
 pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

 # Build with Jax bindings
 NVTE_FRAMEWORK=jax pip install .
 # Execute one of the following command
 NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
 NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

 # Build with all bindings (Pytorch, TF, Jax)
 NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down

9 changes: 6 additions & 3 deletions
docs/installation.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down

37 changes: 27 additions & 10 deletions
setup.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
 def install_requires():
 return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_JAX=ON"]
 return ["-DENABLE_TENSORFLOW=ON"]

 def run(self, extensions):
 print("Building jax extensions!")
 print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_TENSORFLOW=ON"]
 p = [d for d in sys.path if 'dist-packages' in d][0]
 return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 def run(self, extensions):
 print("Building TensorFlow extensions!")
 print("Building jax extensions!")

 @staticmethod
 def install_requires():
 return ["pydantic",]
 # TODO: find a way to install pybind11 and ninja directly.
 return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
 dlfw_builder_funcs.append(JaxBuilder)
 # Trigger a better error when pybind11 isn't present.
 # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
 # This doesn't install a python packages. So the line bellow is too strict.
 # When it fail, we need to detect if cmake will find pybind11.
 # import pybind11

Comment on lines

+211
 to 
+215

@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.

So I wanted to generate a better error by doing "import pybind11".

This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.

if framework in ("all", "tensorflow"):
 dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
 dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
 build_dir = os.path.abspath(build_dir)

 cmake_args = [
 "-GNinja",
 "-DCMAKE_BUILD_TYPE=" + config,
 "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
 ]
 try:
 import ninja
 except ImportError:
 pass
 else:
 cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 cmake_args = cmake_args + self.dlfw_flags

 cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
 ext_modules=ext_modules,
 cmdclass={"build_ext": TEBuildExtension},
 install_requires=dlfw_install_requires,
 extras_require={
 'test': ['pytest',
 'tensorflow_datasets'],
 'test_pytest': ['onnxruntime',],
 },
 license_files=("LICENSE",),
)

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Force GitHub README to respect dark mode
(function() {
var style = document.createElement('style');
style.textContent = '
.markdown-body {
color-scheme: dark light;
}
.markdown-body pre { background: #161b22 !important; }
.markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }
.markdown-body table th, .markdown-body table td { border-color: #30363d !important; }
.markdown-body img { background: #0d1117; }
.markdown-body blockquote { border-left-color: #8b949e; }
.markdown-body hr { border-color: #30363d; }
';
document.head.appendChild(style);
})();
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Update installation instruction for JAX and add some dependencies. by nouiz · Pull Request #117 · NVIDIA/TransformerEngine · GitHub
Skip to content
Update installation instruction for JAX and add some dependencies.#117

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

ksivaman
merged 10 commits into

from

Apr 5, 2023

+46

−26

 Merged


Update installation instruction for JAX and add some dependencies.
#117
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
e3268b5

Update installation instructio for JAX and add some depenencies.


nouiz Mar 24, 2023
94405bf

Bring back support for none pip installed pybind11.


nouiz Mar 31, 2023
f610636

Apply suggestions from code review


nouiz Apr 4, 2023
d9ae968

Changes following review.


nouiz Apr 4, 2023
0773865

Change order to make it more clear.


nouiz Apr 4, 2023
3bd0c07

Add other reviers suggestion.


nouiz Apr 4, 2023
f59b91e

pybind11 is needed for all FW.


nouiz Apr 4, 2023
910cfd0

Add flax as a dep


nouiz Apr 4, 2023
cbbe556

Merge branch 'main' into jax_install


ksivaman Apr 4, 2023
87b0b5e

Update README.rst


nouiz Apr 5, 2023

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


26 changes: 13 additions & 13 deletions
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.

pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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


Suggested change

Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.

I was told it could be faster in some cases, but I don't know which case and I didn't test it.

.. code-block:: bash

 NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
 pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

 # Build with Jax bindings
 NVTE_FRAMEWORK=jax pip install .
 # Execute one of the following command
 NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
 NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

 # Build with all bindings (Pytorch, TF, Jax)
 NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down

9 changes: 6 additions & 3 deletions
docs/installation.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down

37 changes: 27 additions & 10 deletions
setup.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
 def install_requires():
 return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_JAX=ON"]
 return ["-DENABLE_TENSORFLOW=ON"]

 def run(self, extensions):
 print("Building jax extensions!")
 print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_TENSORFLOW=ON"]
 p = [d for d in sys.path if 'dist-packages' in d][0]
 return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 def run(self, extensions):
 print("Building TensorFlow extensions!")
 print("Building jax extensions!")

 @staticmethod
 def install_requires():
 return ["pydantic",]
 # TODO: find a way to install pybind11 and ninja directly.
 return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
 dlfw_builder_funcs.append(JaxBuilder)
 # Trigger a better error when pybind11 isn't present.
 # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
 # This doesn't install a python packages. So the line bellow is too strict.
 # When it fail, we need to detect if cmake will find pybind11.
 # import pybind11

Comment on lines

+211
 to 
+215

@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.

So I wanted to generate a better error by doing "import pybind11".

This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.

if framework in ("all", "tensorflow"):
 dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
 dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
 build_dir = os.path.abspath(build_dir)

 cmake_args = [
 "-GNinja",
 "-DCMAKE_BUILD_TYPE=" + config,
 "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
 ]
 try:
 import ninja
 except ImportError:
 pass
 else:
 cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 cmake_args = cmake_args + self.dlfw_flags

 cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
 ext_modules=ext_modules,
 cmdclass={"build_ext": TEBuildExtension},
 install_requires=dlfw_install_requires,
 extras_require={
 'test': ['pytest',
 'tensorflow_datasets'],
 'test_pytest': ['onnxruntime',],
 },
 license_files=("LICENSE",),
)

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Highlight search terms from Google/DuckDuckGo/Bing referrer
(function() {
var ref = document.referrer;
var terms = [];
if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {
var url = new URL(ref);
var q = url.searchParams.get('q') || url.searchParams.get('p');
if (q) {
terms = q.split(/\s+/).filter(function(t) { return t.length > 2; });
}
}
if (terms.length === 0) return;
var style = document.createElement('style');
style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';
document.head.appendChild(style);
function highlight(node) {
if (node.nodeType === 3) { // text node
var text = node.textContent;
var found = false;
terms.forEach(function(term) {
var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\]\\]/g, '\\') + ')', 'gi');
if (regex.test(text)) {
found = true;
var frag = document.createDocumentFragment();
var parts = text.split(regex);
parts.forEach(function(part, i) {
if (i % 2 === 0) {
frag.appendChild(document.createTextNode(part));
} else {
var span = document.createElement('span');
span.className = 'userscript-highlight';
span.textContent = part;
frag.appendChild(span);
}
});
node.parentNode.replaceChild(frag, node);
}
});
} else if (node.nodeType === 1 && node.childNodes) { // element
var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];
if (!skipTags.includes(node.tagName)) {
Array.from(node.childNodes).forEach(highlight);
}
}
}
highlight(document.body);
// Re-highlight on dynamic content
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1 || node.nodeType === 3) highlight(node);
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Update installation instruction for JAX and add some dependencies. by nouiz · Pull Request #117 · NVIDIA/TransformerEngine · GitHub
Skip to content
Update installation instruction for JAX and add some dependencies.#117

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

ksivaman
merged 10 commits into

from

Apr 5, 2023

+46

−26

 Merged


Update installation instruction for JAX and add some dependencies.
#117
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
e3268b5

Update installation instructio for JAX and add some depenencies.


nouiz Mar 24, 2023
94405bf

Bring back support for none pip installed pybind11.


nouiz Mar 31, 2023
f610636

Apply suggestions from code review


nouiz Apr 4, 2023
d9ae968

Changes following review.


nouiz Apr 4, 2023
0773865

Change order to make it more clear.


nouiz Apr 4, 2023
3bd0c07

Add other reviers suggestion.


nouiz Apr 4, 2023
f59b91e

pybind11 is needed for all FW.


nouiz Apr 4, 2023
910cfd0

Add flax as a dep


nouiz Apr 4, 2023
cbbe556

Merge branch 'main' into jax_install


ksivaman Apr 4, 2023
87b0b5e

Update README.rst


nouiz Apr 5, 2023

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


26 changes: 13 additions & 13 deletions
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.

pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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


Suggested change

Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.

I was told it could be faster in some cases, but I don't know which case and I didn't test it.

.. code-block:: bash

 NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
 pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

 # Build with Jax bindings
 NVTE_FRAMEWORK=jax pip install .
 # Execute one of the following command
 NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
 NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

 # Build with all bindings (Pytorch, TF, Jax)
 NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down

9 changes: 6 additions & 3 deletions
docs/installation.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down

37 changes: 27 additions & 10 deletions
setup.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
 def install_requires():
 return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_JAX=ON"]
 return ["-DENABLE_TENSORFLOW=ON"]

 def run(self, extensions):
 print("Building jax extensions!")
 print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_TENSORFLOW=ON"]
 p = [d for d in sys.path if 'dist-packages' in d][0]
 return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 def run(self, extensions):
 print("Building TensorFlow extensions!")
 print("Building jax extensions!")

 @staticmethod
 def install_requires():
 return ["pydantic",]
 # TODO: find a way to install pybind11 and ninja directly.
 return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
 dlfw_builder_funcs.append(JaxBuilder)
 # Trigger a better error when pybind11 isn't present.
 # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
 # This doesn't install a python packages. So the line bellow is too strict.
 # When it fail, we need to detect if cmake will find pybind11.
 # import pybind11

Comment on lines

+211
 to 
+215

@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.

So I wanted to generate a better error by doing "import pybind11".

This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.

if framework in ("all", "tensorflow"):
 dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
 dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
 build_dir = os.path.abspath(build_dir)

 cmake_args = [
 "-GNinja",
 "-DCMAKE_BUILD_TYPE=" + config,
 "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
 ]
 try:
 import ninja
 except ImportError:
 pass
 else:
 cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 cmake_args = cmake_args + self.dlfw_flags

 cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
 ext_modules=ext_modules,
 cmdclass={"build_ext": TEBuildExtension},
 install_requires=dlfw_install_requires,
 extras_require={
 'test': ['pytest',
 'tensorflow_datasets'],
 'test_pytest': ['onnxruntime',],
 },
 license_files=("LICENSE",),
)

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Strip utm_, fbclid, gclid, etc. from all links on page
(function() {
var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',
'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',
'ref', 'ref_src', 'source', 'medium', 'campaign'];
function cleanUrl(url) {
try {
var u = new URL(url, window.location.origin);
var changed = false;
trackingParams.forEach(function(p) {
if (u.searchParams.has(p)) {
u.searchParams.delete(p);
changed = true;
}
});
return changed ? u.toString() : url;
} catch (e) {
return url;
}
}
function cleanLinks() {
document.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
cleanLinks();
var observer = new MutationObserver(function(mutations) {
mutations.forEach(function(m) {
m.addedNodes.forEach(function(node) {
if (node.nodeType === 1) {
if (node.tagName === 'A') cleanLinks();
node.querySelectorAll('a[href]').forEach(function(a) {
var clean = cleanUrl(a.href);
if (clean !== a.href) a.href = clean;
});
}
});
});
});
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '
Update installation instruction for JAX and add some dependencies. by nouiz · Pull Request #117 · NVIDIA/TransformerEngine · GitHub
Skip to content
Update installation instruction for JAX and add some dependencies.#117

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

ksivaman
merged 10 commits into

from

Apr 5, 2023

+46

−26

 Merged


Update installation instruction for JAX and add some dependencies.
#117
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
e3268b5

Update installation instructio for JAX and add some depenencies.


nouiz Mar 24, 2023
94405bf

Bring back support for none pip installed pybind11.


nouiz Mar 31, 2023
f610636

Apply suggestions from code review


nouiz Apr 4, 2023
d9ae968

Changes following review.


nouiz Apr 4, 2023
0773865

Change order to make it more clear.


nouiz Apr 4, 2023
3bd0c07

Add other reviers suggestion.


nouiz Apr 4, 2023
f59b91e

pybind11 is needed for all FW.


nouiz Apr 4, 2023
910cfd0

Add flax as a dep


nouiz Apr 4, 2023
cbbe556

Merge branch 'main' into jax_install


ksivaman Apr 4, 2023
87b0b5e

Update README.rst


nouiz Apr 5, 2023

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


26 changes: 13 additions & 13 deletions
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.

pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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


Suggested change

Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.

I was told it could be faster in some cases, but I don't know which case and I didn't test it.

.. code-block:: bash

 NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
 pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

 # Build with Jax bindings
 NVTE_FRAMEWORK=jax pip install .
 # Execute one of the following command
 NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
 NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

 # Build with all bindings (Pytorch, TF, Jax)
 NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down

9 changes: 6 additions & 3 deletions
docs/installation.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down

37 changes: 27 additions & 10 deletions
setup.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
 def install_requires():
 return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_JAX=ON"]
 return ["-DENABLE_TENSORFLOW=ON"]

 def run(self, extensions):
 print("Building jax extensions!")
 print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_TENSORFLOW=ON"]
 p = [d for d in sys.path if 'dist-packages' in d][0]
 return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 def run(self, extensions):
 print("Building TensorFlow extensions!")
 print("Building jax extensions!")

 @staticmethod
 def install_requires():
 return ["pydantic",]
 # TODO: find a way to install pybind11 and ninja directly.
 return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
 dlfw_builder_funcs.append(JaxBuilder)
 # Trigger a better error when pybind11 isn't present.
 # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
 # This doesn't install a python packages. So the line bellow is too strict.
 # When it fail, we need to detect if cmake will find pybind11.
 # import pybind11

Comment on lines

+211
 to 
+215

@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.

So I wanted to generate a better error by doing "import pybind11".

This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.

if framework in ("all", "tensorflow"):
 dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
 dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
 build_dir = os.path.abspath(build_dir)

 cmake_args = [
 "-GNinja",
 "-DCMAKE_BUILD_TYPE=" + config,
 "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
 ]
 try:
 import ninja
 except ImportError:
 pass
 else:
 cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 cmake_args = cmake_args + self.dlfw_flags

 cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
 ext_modules=ext_modules,
 cmdclass={"build_ext": TEBuildExtension},
 install_requires=dlfw_install_requires,
 extras_require={
 'test': ['pytest',
 'tensorflow_datasets'],
 'test_pytest': ['onnxruntime',],
 },
 license_files=("LICENSE",),
)

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Auto-enable theater mode on YouTube
(function() {
function tryTheater() {
var btn = document.querySelector('button[aria-label="Theater mode"], ytd-player #player button[title="Theater mode"]');
if (btn && !btn.classList.contains('activated')) {
btn.click();
}
}
// Try immediately
tryTheater();
// Try after navigation (SPA)
var lastUrl = location.href;
setInterval(function() {
if (location.href !== lastUrl) {
lastUrl = location.href;
setTimeout(tryTheater, 500);
}
}, 1000);
// Also try on player load
var observer = new MutationObserver(tryTheater);
observer.observe(document.body, { childList: true, subtree: true });
})();
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Update installation instruction for JAX and add some dependencies. by nouiz · Pull Request #117 · NVIDIA/TransformerEngine · GitHub
Skip to content
Update installation instruction for JAX and add some dependencies.#117

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

ksivaman
merged 10 commits into

from

Apr 5, 2023

+46

−26

 Merged


Update installation instruction for JAX and add some dependencies.
#117
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
e3268b5

Update installation instructio for JAX and add some depenencies.


nouiz Mar 24, 2023
94405bf

Bring back support for none pip installed pybind11.


nouiz Mar 31, 2023
f610636

Apply suggestions from code review


nouiz Apr 4, 2023
d9ae968

Changes following review.


nouiz Apr 4, 2023
0773865

Change order to make it more clear.


nouiz Apr 4, 2023
3bd0c07

Add other reviers suggestion.


nouiz Apr 4, 2023
f59b91e

pybind11 is needed for all FW.


nouiz Apr 4, 2023
910cfd0

Add flax as a dep


nouiz Apr 4, 2023
cbbe556

Merge branch 'main' into jax_install


ksivaman Apr 4, 2023
87b0b5e

Update README.rst


nouiz Apr 5, 2023

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


26 changes: 13 additions & 13 deletions
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.

pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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


Suggested change

Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.

I was told it could be faster in some cases, but I don't know which case and I didn't test it.

.. code-block:: bash

 NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
 pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

 # Build with Jax bindings
 NVTE_FRAMEWORK=jax pip install .
 # Execute one of the following command
 NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
 NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

 # Build with all bindings (Pytorch, TF, Jax)
 NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down

9 changes: 6 additions & 3 deletions
docs/installation.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down

37 changes: 27 additions & 10 deletions
setup.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
 def install_requires():
 return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_JAX=ON"]
 return ["-DENABLE_TENSORFLOW=ON"]

 def run(self, extensions):
 print("Building jax extensions!")
 print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_TENSORFLOW=ON"]
 p = [d for d in sys.path if 'dist-packages' in d][0]
 return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 def run(self, extensions):
 print("Building TensorFlow extensions!")
 print("Building jax extensions!")

 @staticmethod
 def install_requires():
 return ["pydantic",]
 # TODO: find a way to install pybind11 and ninja directly.
 return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
 dlfw_builder_funcs.append(JaxBuilder)
 # Trigger a better error when pybind11 isn't present.
 # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
 # This doesn't install a python packages. So the line bellow is too strict.
 # When it fail, we need to detect if cmake will find pybind11.
 # import pybind11

Comment on lines

+211
 to 
+215

@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.

So I wanted to generate a better error by doing "import pybind11".

This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.

if framework in ("all", "tensorflow"):
 dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
 dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
 build_dir = os.path.abspath(build_dir)

 cmake_args = [
 "-GNinja",
 "-DCMAKE_BUILD_TYPE=" + config,
 "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
 ]
 try:
 import ninja
 except ImportError:
 pass
 else:
 cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 cmake_args = cmake_args + self.dlfw_flags

 cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
 ext_modules=ext_modules,
 cmdclass={"build_ext": TEBuildExtension},
 install_requires=dlfw_install_requires,
 extras_require={
 'test': ['pytest',
 'tensorflow_datasets'],
 'test_pytest': ['onnxruntime',],
 },
 license_files=("LICENSE",),
)

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Remove or un-stick sticky/fixed headers that block content
(function() {
function unstick() {
document.querySelectorAll('header, nav, [role="banner"], .header, .navbar, .sticky, .fixed-top, [style*="position: fixed"], [style*="position:sticky"]').forEach(function(el) {
if (el.style.position === 'fixed' || el.style.position === 'sticky' || getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {
el.style.position = 'static';
el.style.top = 'auto';
el.style.zIndex = 'auto';
}
});
}
unstick();
var observer = new MutationObserver(unstick);
observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });
})();
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '
Update installation instruction for JAX and add some dependencies. by nouiz · Pull Request #117 · NVIDIA/TransformerEngine · GitHub
Skip to content
Update installation instruction for JAX and add some dependencies.#117

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

ksivaman
merged 10 commits into

from

Apr 5, 2023

+46

−26

 Merged


Update installation instruction for JAX and add some dependencies.
#117
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
e3268b5

Update installation instructio for JAX and add some depenencies.


nouiz Mar 24, 2023
94405bf

Bring back support for none pip installed pybind11.


nouiz Mar 31, 2023
f610636

Apply suggestions from code review


nouiz Apr 4, 2023
d9ae968

Changes following review.


nouiz Apr 4, 2023
0773865

Change order to make it more clear.


nouiz Apr 4, 2023
3bd0c07

Add other reviers suggestion.


nouiz Apr 4, 2023
f59b91e

pybind11 is needed for all FW.


nouiz Apr 4, 2023
910cfd0

Add flax as a dep


nouiz Apr 4, 2023
cbbe556

Merge branch 'main' into jax_install


ksivaman Apr 4, 2023
87b0b5e

Update README.rst


nouiz Apr 5, 2023

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


26 changes: 13 additions & 13 deletions
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.

pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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


Suggested change

Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.

I was told it could be faster in some cases, but I don't know which case and I didn't test it.

.. code-block:: bash

 NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
 pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

 # Build with Jax bindings
 NVTE_FRAMEWORK=jax pip install .
 # Execute one of the following command
 NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
 NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

 # Build with all bindings (Pytorch, TF, Jax)
 NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down

9 changes: 6 additions & 3 deletions
docs/installation.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down

37 changes: 27 additions & 10 deletions
setup.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
 def install_requires():
 return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_JAX=ON"]
 return ["-DENABLE_TENSORFLOW=ON"]

 def run(self, extensions):
 print("Building jax extensions!")
 print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_TENSORFLOW=ON"]
 p = [d for d in sys.path if 'dist-packages' in d][0]
 return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 def run(self, extensions):
 print("Building TensorFlow extensions!")
 print("Building jax extensions!")

 @staticmethod
 def install_requires():
 return ["pydantic",]
 # TODO: find a way to install pybind11 and ninja directly.
 return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
 dlfw_builder_funcs.append(JaxBuilder)
 # Trigger a better error when pybind11 isn't present.
 # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
 # This doesn't install a python packages. So the line bellow is too strict.
 # When it fail, we need to detect if cmake will find pybind11.
 # import pybind11

Comment on lines

+211
 to 
+215

@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.

So I wanted to generate a better error by doing "import pybind11".

This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.

if framework in ("all", "tensorflow"):
 dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
 dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
 build_dir = os.path.abspath(build_dir)

 cmake_args = [
 "-GNinja",
 "-DCMAKE_BUILD_TYPE=" + config,
 "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
 ]
 try:
 import ninja
 except ImportError:
 pass
 else:
 cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 cmake_args = cmake_args + self.dlfw_flags

 cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
 ext_modules=ext_modules,
 cmdclass={"build_ext": TEBuildExtension},
 install_requires=dlfw_install_requires,
 extras_require={
 'test': ['pytest',
 'tensorflow_datasets'],
 'test_pytest': ['onnxruntime',],
 },
 license_files=("LICENSE",),
)

, 'i');
if (__m === '*' || __re.test(location.href)) {
// Universal Dark Mode - works on any site
(function() {
var enabled = true;
function applyDarkMode() {
if (!enabled) return;
// Create style element if it doesn't exist
var style = document.getElementById('universal-dark-mode-style');
if (!style) {
style = document.createElement('style');
style.id = 'universal-dark-mode-style';
document.head.appendChild(style);
}
// Dark mode CSS - inverts colors but preserves images/video
style.textContent = '
/* Invert everything except media */
html {
filter: invert(1) hue-rotate(180deg) !important;
background: #1a1a2e !important;
}
/* Restore images, videos, iframes, canvas */
img, video, iframe, canvas, svg, picture, [style*="background-image"] {
filter: invert(1) hue-rotate(180deg) !important;
}
/* Preserve specific elements that should not be inverted */
.no-dark-mode, .no-dark-mode *,
[data-theme="light"], [data-theme="light"],
.ace_editor, .ace_editor *,
.CodeMirror, .CodeMirror *,
.monaco-editor, .monaco-editor *,
.markdown-body pre, .markdown-body pre *,
.highlight, .highlight *,
pre code, pre code * {
filter: none !important;
}
/* Fix common UI elements */
.modal, .popup, .dropdown-menu, .tooltip, .popover {
filter: invert(1) hue-rotate(180deg) !important;
background: #2d2d44 !important;
border-color: #444 !important;
}
/* Scrollbars */
::-webkit-scrollbar { background: #1a1a2e !important; }
::-webkit-scrollbar-thumb { background: #444 !important; }
::-webkit-scrollbar-thumb:hover { background: #555 !important; }
/* Selection */
::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }
';
}
function removeDarkMode() {
var style = document.getElementById('universal-dark-mode-style');
if (style) style.remove();
}
// Toggle with Alt+Shift+D
document.addEventListener('keydown', function(e) {
if (e.altKey && e.shiftKey && e.key === 'D') {
e.preventDefault();
enabled = !enabled;
if (enabled) {
applyDarkMode();
console.log('[Universal Dark Mode] Enabled');
} else {
removeDarkMode();
console.log('[Universal Dark Mode] Disabled');
}
}
});
// Apply on load
applyDarkMode();
// Re-apply on dynamic content
var observer = new MutationObserver(function(mutations) {
if (enabled && !document.getElementById('universal-dark-mode-style')) {
applyDarkMode();
}
});
observer.observe(document.head, { childList: true });
console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');
})();
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();
Update installation instruction for JAX and add some dependencies. by nouiz · Pull Request #117 · NVIDIA/TransformerEngine · GitHub
Skip to content
Update installation instruction for JAX and add some dependencies.#117

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

ksivaman
merged 10 commits into

from

Apr 5, 2023

+46

−26

 Merged


Update installation instruction for JAX and add some dependencies.
#117
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


10 commits


Select commit
Hold shift + click to select a range
e3268b5

Update installation instructio for JAX and add some depenencies.


nouiz Mar 24, 2023
94405bf

Bring back support for none pip installed pybind11.


nouiz Mar 31, 2023
f610636

Apply suggestions from code review


nouiz Apr 4, 2023
d9ae968

Changes following review.


nouiz Apr 4, 2023
0773865

Change order to make it more clear.


nouiz Apr 4, 2023
3bd0c07

Add other reviers suggestion.


nouiz Apr 4, 2023
f59b91e

pybind11 is needed for all FW.


nouiz Apr 4, 2023
910cfd0

Add flax as a dep


nouiz Apr 4, 2023
cbbe556

Merge branch 'main' into jax_install


ksivaman Apr 4, 2023
87b0b5e

Update README.rst


nouiz Apr 5, 2023

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


26 changes: 13 additions & 13 deletions
README.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -162,32 +162,32 @@ Transformer Engine comes preinstalled in the pyTorch container on
From source
^^^^^^^^^^^

First, install the prequisites.
For JAX and Tensorflow, pybind11 must be installed:

.. code-block:: bash

apt-get install ninja-build pybind11-dev
Comment thread
ptrendx marked this conversation as resolved.

pip install pybind11

Clone the repository and inside it type:
Then, you can install this optional dependency:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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


Suggested change

Then, you can install this optional dependency:
Then, the following optional dependencies can be installed:
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Stating the purpose of these dependencies would be useful, since they are 'optional'.
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

Agreed. I wanted to do that. But it isn't clear to me what ninja bring.

I was told it could be faster in some cases, but I don't know which case and I didn't test it.

.. code-block:: bash

 NVTE_FRAMEWORK=all pip install . # Building with all frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Building with pyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Building with JAX only.
 pip install ninja

You can also specify which framework bindings to build. The default is pytorch only.
Install TE (optionally specifying the framework):

.. code-block:: bash

# Build with TensorFlow bindings
NVTE_FRAMEWORK=tensorflow pip install .
git clone https://github.com/NVIDIA/TransformerEngine.git
cd TransformerEngine

 # Build with Jax bindings
 NVTE_FRAMEWORK=jax pip install .
 # Execute one of the following command
 NVTE_FRAMEWORK=all pip install . # Build TE for all supported frameworks.
 NVTE_FRAMEWORK=pytorch pip install . # Build TE for PyTorch only.
 NVTE_FRAMEWORK=jax pip install . # Build TE for JAX only.
 NVTE_FRAMEWORK=tensorflow pip install . # Build TE for TensorFlow only.

 # Build with all bindings (Pytorch, TF, Jax)
 NVTE_FRAMEWORK=all pip install .
If the framework is not explicitly specified, TE will be built for PyTorch only.

User Guide
----------
Expand Down

9 changes: 6 additions & 3 deletions
docs/installation.rst

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -29,9 +29,12 @@ pip - from GitHub
Additional Prerequisites
^^^^^^^^^^^^^^^^^^^^^^^^

1. `CMake <https://cmake.org/>`__ version 3.18 or later
2. `pyTorch <https://pytorch.org/>`__ with GPU support
3. `Ninja <https://ninja-build.org/>`__
1. `CMake <https://cmake.org/>`__ version 3.18 or later.
2. [For pyTorch support] `pyTorch <https://pytorch.org/>`__ with GPU support.
3. [For JAX support] `JAX <https://github.com/google/jax/>`__ with GPU support, version >= 0.4.7.
4. [For TensorFlow support] `TensorFlow <https://www.tensorflow.org/>`__ with GPU support.
5. `pybind11`: `pip install pybind11`.
6. [Optional] `Ninja <https://ninja-build.org/>`__: `pip install ninja`.

Installation (stable release)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Expand Down

37 changes: 27 additions & 10 deletions
setup.py

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -160,23 +160,24 @@ def run(self, extensions):
 def install_requires():
 return ["flash-attn @ git+https://github.com/ksivaman/flash-attention.git@hopper",]

class JaxBuilder(FrameworkBuilderBase):
class TensorFlowBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_JAX=ON"]
 return ["-DENABLE_TENSORFLOW=ON"]

 def run(self, extensions):
 print("Building jax extensions!")
 print("Building TensorFlow extensions!")

class TensorFlowBuilder(FrameworkBuilderBase):
class JaxBuilder(FrameworkBuilderBase):
 def cmake_flags(self):
 return ["-DENABLE_TENSORFLOW=ON"]
 p = [d for d in sys.path if 'dist-packages' in d][0]
 return ["-DENABLE_JAX=ON", "-DCMAKE_PREFIX_PATH="+p]
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 def run(self, extensions):
 print("Building TensorFlow extensions!")
 print("Building jax extensions!")

 @staticmethod
 def install_requires():
 return ["pydantic",]
 # TODO: find a way to install pybind11 and ninja directly.
 return ['cmake', 'flax']

ext_modules = []
dlfw_builder_funcs = []
Expand DownExpand Up@@ -207,11 +208,16 @@ def install_requires():

if framework in ("all", "jax"):
 dlfw_builder_funcs.append(JaxBuilder)
 # Trigger a better error when pybind11 isn't present.
 # Sadly, if pybind11 was installed with `apt -y install pybind11-dev`
 # This doesn't install a python packages. So the line bellow is too strict.
 # When it fail, we need to detect if cmake will find pybind11.
 # import pybind11

Comment on lines

+211
 to 
+215

@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

This is not clear, what is the goal of this comment?
@ksivamanksivamanApr 3, 2023

Copy link


Copy Markdown

Member
Choose a reason for hiding this comment

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

Is it to answer the question, "Why we're getting pybind11 via cmake?"
@nouiznouizApr 4, 2023

Copy link


Copy Markdown

CollaboratorAuthor
Choose a reason for hiding this comment

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

If pybind11 isn't installed, the error message is ugly and hard to read.

So I wanted to generate a better error by doing "import pybind11".

This work when we install pybind11 via pip, but it doesn't work when installed via apt-get.

if framework in ("all", "tensorflow"):
 dlfw_builder_funcs.append(TensorFlowBuilder)

dlfw_install_requires = []
dlfw_install_requires = ['pydantic']
for builder in dlfw_builder_funcs:
 dlfw_install_requires = dlfw_install_requires + builder.install_requires()

Expand DownExpand Up@@ -272,10 +278,16 @@ def build_extensions(self) -> None:
 build_dir = os.path.abspath(build_dir)

 cmake_args = [
 "-GNinja",
 "-DCMAKE_BUILD_TYPE=" + config,
 "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(config.upper(), build_dir),
 ]
 try:
 import ninja
 except ImportError:
 pass
 else:
 cmake_args.append("-GNinja")
Comment thread
timmoon10 marked this conversation as resolved.


Outdated


 cmake_args = cmake_args + self.dlfw_flags

 cmake_build_args = ["--config", config]
Expand DownExpand Up@@ -399,5 +411,10 @@ def get_outputs(self):
 ext_modules=ext_modules,
 cmdclass={"build_ext": TEBuildExtension},
 install_requires=dlfw_install_requires,
 extras_require={
 'test': ['pytest',
 'tensorflow_datasets'],
 'test_pytest': ['onnxruntime',],
 },
 license_files=("LICENSE",),
)