Clean up the installation instruction. (#159)

* Clean up the installation instruction. We where telling to install the dev version in the README.
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>

* Typos
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Frederic Bastien <fbastien@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

README.rst

@@ -162,29 +162,8 @@ Transformer Engine comes preinstalled in the pyTorch container on
 From source
 ^^^^^^^^^^^
 
-For JAX and Tensorflow, pybind11 must be installed:
+`See the installation guide <https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/installation.html>`_.
 
-.. code-block:: bash
-
-  pip install pybind11
-
-Then, you can install this optional dependency:
-
-.. code-block:: bash
-
-  pip install ninja
-
-Install TE (optionally specifying the framework):
-
-.. code-block:: bash
-
-  # Execute one of the following command
-  NVTE_FRAMEWORK=all pip install git+https://github.com/NVIDIA/TransformerEngine.git@main        # Build TE for all supported frameworks.
-  NVTE_FRAMEWORK=pytorch pip install git+https://github.com/NVIDIA/TransformerEngine.git@main    # Build TE for PyTorch only.
-  NVTE_FRAMEWORK=jax pip install git+https://github.com/NVIDIA/TransformerEngine.git@main        # Build TE for JAX only.
-  NVTE_FRAMEWORK=tensorflow pip install git+https://github.com/NVIDIA/TransformerEngine.git@main # Build TE for TensorFlow only.
-
-If the framework is not explicitly specified, TE will be built for PyTorch only.
 
 User Guide and Examples
 -----------------------

docs/installation.rst

@@ -29,7 +29,7 @@ pip - from GitHub
 Additional Prerequisites
 ^^^^^^^^^^^^^^^^^^^^^^^^
 
-1. `CMake <https://cmake.org/>`__ version 3.18 or later.
+1. `CMake <https://cmake.org/>`__ version 3.18 or later: `pip install cmake`.
 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.
@@ -43,7 +43,11 @@ Execute the following command to install the latest stable version of Transforme
 
 .. code-block:: bash
 
-   pip install --upgrade git+https://github.com/NVIDIA/TransformerEngine.git@stable
+  # Execute one of the following commands
+  NVTE_FRAMEWORK=pytorch pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable    # Build TE for PyTorch only. The default.
+  NVTE_FRAMEWORK=jax pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable        # Build TE for JAX only.
+  NVTE_FRAMEWORK=tensorflow pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable # Build TE for TensorFlow only.
+  NVTE_FRAMEWORK=all pip install git+https://github.com/NVIDIA/TransformerEngine.git@stable        # Build TE for all supported frameworks.
 
 Installation (development build)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -58,5 +62,8 @@ Execute the following command to install the latest development build of Transfo
 
 .. code-block:: bash
 
-   pip install --upgrade git+https://github.com/NVIDIA/TransformerEngine.git@main
-
+  # Execute one of the following commands
+  NVTE_FRAMEWORK=pytorch pip install git+https://github.com/NVIDIA/TransformerEngine.git@main    # Build TE for PyTorch only. The default.
+  NVTE_FRAMEWORK=jax pip install git+https://github.com/NVIDIA/TransformerEngine.git@main        # Build TE for JAX only.
+  NVTE_FRAMEWORK=tensorflow pip install git+https://github.com/NVIDIA/TransformerEngine.git@main # Build TE for TensorFlow only.
+  NVTE_FRAMEWORK=all pip install git+https://github.com/NVIDIA/TransformerEngine.git@main        # Build TE for all supported frameworks.