Merge pull request #1125 from aws-neuron/release_222

Release 2.22.0 (#2216)

Author: natemail-aws

_ext/neuron_tag.py

@@ -22,9 +22,9 @@
                 'general/announcements/index',
                 'frameworks/tensorflow/tensorflow-neuron/'
                 ]
-add_trn1_tag = ['frameworks/neuron-customops/','frameworks/torch/inference-torch-neuronx', 'libraries/nemo-megatron/']
-# add_trn2_tag = ['']
-add_neuronx_tag = ['frameworks/torch/torch-neuronx/','frameworks/tensorflow/tensorflow-neuronx/','frameworks/torch/inference-torch-neuronx/','libraries/transformers-neuronx/','libraries/neuronx-distributed/','neuronx-distributed/nxd-training', 'general/setup/tensorflow-neuronx']
+add_trn1_tag = ['frameworks/neuron-customops/','frameworks/torch/inference-torch-neuronx', 'libraries/nemo-megatron/', 'libraries/nxd-training/']
+add_trn2_tag = ['libraries/nxd-training/']
+add_neuronx_tag = ['frameworks/torch/torch-neuronx/','frameworks/tensorflow/tensorflow-neuronx/','frameworks/torch/inference-torch-neuronx/','libraries/transformers-neuronx/','libraries/neuronx-distributed/','libraries/nxd-training', 'general/setup/tensorflow-neuronx']
 clear_inf1_tag = ['general/arch/neuron-features/neuron-caching',
                 'general/arch/neuron-features/eager-debug-mode',
                 'general/arch/neuron-features/collective-communication-operations',
@@ -82,7 +82,8 @@
                  'general/arch/neuron-hardware/trn2-arch',
                  'general/arch/neuron-hardware/neuron-core-v3',
                  '/general/announcements/neuron2.x/announce-neuron-trn2',
-                 '/general/arch/neuron-features/logical-neuroncore-config'
+                 '/general/arch/neuron-features/logical-neuroncore-config',
+                 'libraries/nxd-training/'
                 ]
 
 clear_inf2_tag = ['frameworks/torch/torch-neuronx/training',
@@ -99,7 +100,8 @@
                   'general/arch/neuron-hardware/trn2-arch',
                   'general/arch/neuron-hardware/neuron-core-v3', 
                   '/general/announcements/neuron2.x/announce-neuron-trn2',
-                  '/general/arch/neuron-features/logical-neuroncore-config'
+                  '/general/arch/neuron-features/logical-neuroncore-config',
+                  'libraries/nxd-training/'
                ]
 
 
@@ -265,6 +267,9 @@ def run(self):
                 return_instances.append('Inf2')
                 return_instances.append('Trn1')
                 return_instances.append('Trn2')
+
+        if cur_file=='general/appnotes/neuronx-distributed/introducing-nxdt-training':
+            return_instances = ['Trn1','Trn2']
 
         # generate text from instances list if the list is not empty.
         return_instances = sorted(set(return_instances))

compiler/neuronx-cc/api-reference-guide/neuron-compiler-cli-reference-guide.rst

@@ -122,7 +122,7 @@ Common parameters for the Neuron CLI:
 
     - ``llm-training``: Enable the compiler to perform optimizations applicable to large language model (LLMS) training runs that  shard parameters, gradients, and optimizer states across data-parallel workers. This is equivalent to the previously documented option argument value of ``NEMO``, which will be deprecated in a future release.
 
-  - :option:`--logical-nc-config <shard_degree>`: Instructs the compiler to shard the input graph across physical NeuronCore accelerators. Possible numeric values are {1, 2}. (only available on trn2; Default: ``2``)
+  - :option:`--logical-nc-config <shard_degree>`: Instructs the compiler to shard the input graph across physical NeuronCore accelerators. Possible numeric values are {1, 2}. (Only available on trn2; Default: ``2``)
 
     Valid values:
 
@@ -141,7 +141,7 @@ Common parameters for the Neuron CLI:
 
   - :option:`--enable-mixed-precision-accumulation`: Perform intermediate calculations of accumulation operators (such as softmax and layernorm) in FP32 and cast the result to the model-designated datatype. This improves the operator's resulting accuracy.
 
-  - :option:`--enable-saturate-infinity`: Convert +/- infinity values to MAX/MIN_FLOAT for compiler-introduced matrix-multiply transpose computations that have a high risk of generating Not-a-Number (NaN) values. There is a potential performance impact during model execution when this conversion is enabled.
+  - :option:`--enable-saturate-infinity`: Convert +/- infinity values to MAX/MIN_FLOAT for compiler-introduced matrix-multiply transpose computations that have a high risk of generating Not-a-Number (NaN) values. There is a potential performance impact during model execution when this conversion is enabled. (Only needed on trn1; while the trn2 compiler will accept this flag for compatibility reasons, it has no effect on the compilation.)
 
   - :option:`--enable-fast-context-switch`: Optimize for faster model switching rather than execution latency.
       This option will defer loading some weight constants until the start of model execution. This results in overall faster system performance when your application switches between models frequently on the same Neuron Core (or set of cores).

conf.py

@@ -195,7 +195,7 @@
 
 #top_banner_message="<span>&#9888;</span><a class='reference internal' style='color:white;' href='https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/setup/setup-troubleshooting.html#gpg-key-update'>  Neuron repository GPG key for Ubuntu installation has expired, see instructions how to update! </a>"
 
-top_banner_message="Neuron 2.21.1 is released! check <a class='reference internal' style='color:white;' href='https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/index.html#latest-neuron-release'> What's New  </a> and <a class='reference internal' style='color:white;' href='https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/announcements/index.html'> Announcements  </a>"
+top_banner_message="Neuron 2.22.0 is released! check <a class='reference internal' style='color:white;' href='https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/index.html#latest-neuron-release'> What's New  </a> and <a class='reference internal' style='color:white;' href='https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/announcements/index.html'> Announcements  </a>"
 
 html_theme = "sphinx_book_theme"
 html_theme_options = {

dlami/index.rst

@@ -26,9 +26,31 @@ comes pre-installed with all the Neuron libraries including Neuron compiler and
   Tensorflow-neuron 2.10 (inf1) released in SDK v2.20.2 is not compatible with the latest runtime in v2.21 SDK. 
   Code that compiles will face runtime errors with the latest SDK 2.21.1 version.
 
-  Neuron team is aware of this issue and it will be fixed in the next minor release.
+  Neuron team is aware of this issue and we will ship a single-framework AMI for TF 2.10 inf1 in a future release.
+
+  You can use multi-framework DLAMIs from Neuron SDK v2.20.0 for inf1 workloards to avoid this issue. For example:
+
+  Deep Learning AMI Neuron (Ubuntu 22.04/AL2023) 20241027
+
+  | Ubuntu22: ami-017ff4652165fd617
+  | AL2023: ami-06fdb253ce8a32239
+
+  .. code-block:: shell
+    
+     aws ec2 run-instances --image-id <ami-id>
+  
+  
+  Alternatively, you can use the latest Neuron DLAMIs on Ubuntu and run this command as a work-around:
 
-  Please refer to `this page <https://github.com/aws-neuron/aws-neuron-sdk/issues/1071>`_ for more information on the issue and a temporary work-around.
+  .. code-block:: shell
+
+     sudo apt-get remove -y aws-neuronx-dkms aws-neuronx-collectives aws-neuronx-runtime-lib aws-neuronx-tools
+     sudo apt-get install aws-neuronx-dkms=2.18.* -y
+     sudo apt-get install aws-neuronx-collectives=2.22.* -y
+     sudo apt-get install aws-neuronx-runtime-lib=2.22.* -y
+     sudo apt-get install aws-neuronx-tools=2.19.* -y
+  
+  https://github.com/aws-neuron/aws-neuron-sdk/issues/1071 for more information on the issue.
 
 Multi Framework DLAMIs supported
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -44,10 +66,11 @@ Multi Framework DLAMIs supported
       - DLAMI Name
 
     * - Ubuntu 22.04
-      - Inf2, Trn1, Trn1n, Trn2
+      - Inf1, Inf2, Trn1, Trn1n, Trn2
       - Deep Learning AMI Neuron (Ubuntu 22.04)
+
     * - Amazon Linux 2023
-      - Inf2, Trn1, Trn1n, Trn2
+      - Inf1, Inf2, Trn1, Trn1n, Trn2
       - Deep Learning AMI Neuron (Amazon Linux 2023)
 
 
@@ -72,18 +95,21 @@ Virtual Environments pre-installed
 
     * - PyTorch 2.5 NxD Inference, Torch NeuronX
       - /opt/aws_neuronx_venv_pytorch_2_5_nxd_inference
+    
+    * - Transformers NeuronX (PyTorch 2.5)
+      - /opt/aws_neuronx_venv_pytorch_2_5_transformers
 
     * - JAX 0.4 NeuronX
       - /opt/aws_neuronx_venv_jax_0_4
 
     * - Tensorflow 2.10 NeuronX
       - /opt/aws_neuronx_venv_tensorflow_2_10
 
-    * - Transformers NeuronX (PyTorch 2.5)
-      - /opt/aws_neuronx_venv_pytorch_2_5_transformers
-
     * - Tensorflow 2.10 Neuron (Inf1)
       - /opt/aws_neuron_venv_tensorflow_2_10_inf1
+    
+    * - PyTorch 1.13 Neuron (Inf1)
+      - /opt/aws_neuron_venv_pytorch_1_13_inf1
 
 You can easily get started with the multi-framework DLAMI through AWS console by following this :ref:`setup guide <setup-ubuntu22-multi-framework-dlami>`. If you are looking to
 use the Neuron DLAMI in your cloud automation flows, Neuron also supports :ref:`SSM parameters <ssm-parameter-neuron-dlami>` to easily retrieve the latest DLAMI id.
@@ -109,15 +135,30 @@ Single Framework DLAMIs supported
       - Neuron Instances Supported
       - DLAMI Name
 
-    * - Tensorflow 2.10
+    * - PyTorch 2.5
       - Ubuntu 22.04
-      - Inf1, Inf2, Trn1, Trn1n, Trn2
-      - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 22.04)
+      - Inf2, Trn1, Trn1n, Trn2
+      - Deep Learning AMI Neuron PyTorch 2.5 (Ubuntu 22.04) 
+
+    * - PyTorch 2.5
+      - Amazon Linux 2023
+      - Inf2, Trn1, Trn1n, Trn2
+      - Deep Learning AMI Neuron PyTorch 2.5 (Amazon Linux 2023)
 
     * - Tensorflow 2.10
-      - Ubuntu 20.04
-      - Inf2, Trn1, Trn1n
-      - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 20.04)
+      - Ubuntu 22.04
+      - Inf2, Trn1, Trn1n, Trn2
+      - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 22.04) 
+    
+    * - Tensorflow 2.10 (Inf1)
+      - Ubuntu 22.04
+      - Inf1
+      - Deep Learning AMI Neuron TensorFlow 2.10 Inf1 (Ubuntu 22.04) 
+    
+    * - PyTorch 1.13 (Inf1)
+      - Ubuntu 22.04
+      - Inf1
+      - Deep Learning AMI Neuron PyTorch 1.13 Inf1 (Ubuntu 22.04)
 
 
 Virtual Environments pre-installed
@@ -133,19 +174,36 @@ Virtual Environments pre-installed
       - Neuron Libraries supported
       - Virtual Environment
 
-    * - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 22.04)
-      - tensorflow-neuronx
-      - /opt/aws_neuronx_venv_tensorflow_2_10
+    * - Deep Learning AMI Neuron PyTorch 2.5 (Ubuntu 22.04, Amazon Linux 2023) 
+      - PyTorch 2.5 Torch NeuronX, NxD Core
+      - /opt/aws_neuronx_venv_pytorch_2_5
 
-    * - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 20.04)
-      - tensorflow-neuronx
-      - /opt/aws_neuron_venv_tensorflow_2_10
+    * - Deep Learning AMI Neuron PyTorch 2.5 (Ubuntu 22.04, Amazon Linux 2023) 
+      - PyTorch 2.5 NxD Training, Torch NeuronX
+      - /opt/aws_neuronx_venv_pytorch_2_5_nxd_training
 
+    * - Deep Learning AMI Neuron PyTorch 2.5 (Ubuntu 22.04, Amazon Linux 2023) 
+      - PyTorch 2.5 NxD Inference, Torch NeuronX
+      - /opt/aws_neuronx_venv_pytorch_2_5_nxd_inference
+
+    * - Deep Learning AMI Neuron PyTorch 2.5 (Ubuntu 22.04, Amazon Linux 2023) 
+      - Transformers NeuronX PyTorch 2.5
+      - /opt/aws_neuronx_venv_pytorch_2_5_transformers 
+
+    * - Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 22.04) 
+      - Pytorch Neuron (Inf1)
+      - /opt/aws_neuron_venv_pytorch_1_13_inf1
+  
+    * - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 22.04) 
+      - Tensorflow Neuronx
+      - /opt/aws_neuronx_venv_tensorflow_2_10
+    
     * - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 22.04)
-      - tensorflow-neuron (Inf1)
+      - Tensorflow Neuron (Inf1)
       - /opt/aws_neuron_venv_tensorflow_2_10_inf1
-
-You can easily get started with the single framework DLAMI through AWS console by following one of the corresponding setup guides . If you are looking to
+    
+    
+You can easily get started with the single framework DLAMI through AWS console by following one of the corresponding setup guides . If you are looking to 
 use the Neuron DLAMI in your cloud automation flows , Neuron also supports :ref:`SSM parameters <ssm-parameter-neuron-dlami>` to easily retrieve the latest DLAMI id.
 
 Neuron Base DLAMI
@@ -166,14 +224,14 @@ Base DLAMIs supported
       - Neuron Instances Supported
       - DLAMI Name
 
+    * - Amazon Linux 2023
+      - Inf1, Inf2, Trn1n, Trn1, Trn2 
+      - Deep Learning Base Neuron AMI (Amazon Linux 2023)
+
     * - Ubuntu 22.04
-      - Inf1, Inf2, Trn1, Trn1n
+      - Inf1, Inf2, Trn1n, Trn1, Trn2 
       - Deep Learning Base Neuron AMI (Ubuntu 22.04)
 
-    * - Ubuntu 20.04
-      - Inf1, Inf2, Trn1, Trn1n
-      - Deep Learning Base Neuron AMI (Ubuntu 20.04)
-
 
 .. _ssm-parameter-neuron-dlami:
 
@@ -221,12 +279,24 @@ SSM Parameter Prefix
     * - Deep Learning AMI Neuron (Amazon Linux 2023)
       - /aws/service/neuron/dlami/multi-framework/amazon-linux-2023
 
+    * - Deep Learning AMI Neuron PyTorch 2.5 (Ubuntu 22.04)
+      - /aws/service/neuron/dlami/pytorch-2.5/ubuntu-22.04
+    
+    * - Deep Learning AMI Neuron PyTorch 2.5 (Amazon Linux 2023)
+      - /aws/service/neuron/dlami/pytorch-2.5/amazon-linux-2023
+
+    * - Deep Learning AMI Neuron PyTorch 1.13 Inf1 (Ubuntu 22.04)
+      - /aws/service/neuron/dlami/pytorch-1.13-inf1/ubuntu-22.04
+
     * - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 22.04)
       - /aws/service/neuron/dlami/tensorflow-2.10/ubuntu-22.04
 
     * - Deep Learning AMI Neuron TensorFlow 2.10 (Ubuntu 20.04)
       - /aws/service/neuron/dlami/tensorflow-2.10/ubuntu-20.04
 
+    * - Deep Learning Base Neuron AMI (Amazon Linux 2023)
+      - /aws/service/neuron/dlami/base/amazon-linux-2023
+   
     * - Deep Learning Base Neuron AMI (Ubuntu 22.04)
       - /aws/service/neuron/dlami/base/ubuntu-22.04
 

frameworks/jax/api-reference-guide/index.rst

@@ -0,0 +1,12 @@
+.. _jax-neuronx-api-reference-guide:
+
+API Reference Guide for JAX Neuronx
+====================================================
+
+.. toctree::
+    :maxdepth: 1
+    :hidden:
+
+    /frameworks/jax/api-reference-guide/neuron-envvars
+
+* :ref:`jax-neuronx-envvars`

frameworks/jax/api-reference-guide/neuron-envvars.rst

@@ -0,0 +1,69 @@
+.. _jax-neuronx-envvars:
+
+JAX NeuronX Environment Variables
+======================================
+
+Environment variables allow modifications to JAX NeuronX behavior
+without requiring code change to user script. It is recommended to set
+them in code or just before invoking the python process, such as
+``NEURON_RT_VISIBLE_CORES=8 python3 <script>`` to avoid inadvertently
+changing behavior for other scripts. Environment variables specific to
+JAX Neuronx are:
+
+``NEURON_CC_FLAGS``
+
+-  Compiler options. Full compiler options are described in the :ref:`mixed-precision-casting-options`.
+
+``XLA_FLAGS``
+
+- When set to ``"--xla_dump_hlo_snapshots --xla_dump_to=<dir>"``, this environmental variable enables dumping snapshots in ``<dir>`` directory. See :ref:`torch-neuronx-snapshotting` section for more information. The snapshotting interface for JAX and Pytorch are identical.
+- When set to ``"--xla_dump_hlo_as_text --xla_dump_hlo_as_proto --xla_dump_to=<dir> --xla_dump_hlo_pass_re='.*'"``, this environmental variable enables dumping HLOs in proto and text formats after each XLA pass. The dumped ``*.hlo.pb`` files are in HloProto format.
+
+``NEURON_FORCE_PJRT_PLUGIN_REGISTRATION``
+
+- When ``NEURON_FORCE_PJRT_PLUGIN_REGISTRATION=1``, the Neuron PJRT plugin will be registered in JAX regardless of the instance type.
+
+``NEURON_RUN_TRIVIAL_COMPUTATION_ON_CPU``
+
+-  When ``NEURON_RUN_TRIVIAL_COMPUTATION_ON_CPU=1``, the Neuron PJRT plugin will compile and execute "trivial" computations on CPU instead of Neuron cores. A "trivial" computation is defined as an HLO program that does not contain any collective-compute instructions. The HLO program will be compiled by the XLA CPU compiler and outputs of the computation will be allocated on Neuron cores. The following HLO instructions are considered as collective-compute instructions.
+
+    - ``all-gather``
+    - ``all-gather-done``
+    - ``all-gather-start``
+    - ``all-reduce-done``
+    - ``all-reduce-start``
+    - ``all-to-all``
+    - ``collective-permute``
+    - ``partition-id``
+    - ``replica-id``
+    - ``recv``
+    - ``recv-done``
+    - ``reduce-scatter``
+    - ``send``
+    - ``send-done``
+
+``NEURON_PJRT_PROCESSES_NUM_DEVICES``
+
+- Should be set to a comma-separated list stating the number of NeuronCores used by each worker process. It is used to construct a global device array with its size equal to the sum of the list. This gets reported to the XLA PJRT runtime when requested. Must be set for multi-process executions. It can be used in conjunction with ``NEURON_RT_VISIBLE_CORES`` to expose a limited number of NeuronCores to each worker process. If ``NEURON_RT_VISIBLE_CORES`` is not set, it should be set to available number of NeuronCores on the host. ``NEURON_PJRT_PROCESSES_NUM_DEVICES`` must be less than or equal to ``NEURON_RT_VISIBLE_CORES``.
+
+``NEURON_PJRT_PROCESS_INDEX``
+
+- An integer stating the index (or rank) of the current worker process. This is required for multi-process environments where all workers need to know information on all participating processes. Must be set for multi-process executions. The value should be between ``0`` and ``NEURON_PJRT_PROCESS_INDEX - 1``.
+
+``NEURON_RT_STOCHASTIC_ROUNDING_EN`` **[Neuron Runtime]**
+
+- When ``NEURON_RT_STOCHASTIC_ROUNDING_EN=1``, JAX Neuron will use stochastic rounding instead of
+  round-nearest-even for all internal rounding operations when casting from FP32 to a reduced precision data type (FP16, BF16, FP8, TF32).
+  This feature has been shown to improve
+  training convergence for reduced precision training jobs. 
+  To switch to round-nearest-even mode, set ``NEURON_RT_STOCHASTIC_ROUNDING_EN=0``.
+
+``NEURON_RT_STOCHASTIC_ROUNDING_SEED`` **[Neuron Runtime]**
+
+- Sets the seed for the random number generator used in stochastic rounding (see previous section). If this environment variable is not set, the seed is set to 0 by default. Please set ``NEURON_RT_STOCHASTIC_ROUNDING_SEED`` to a fixed value to ensure reproducibility between runs.
+
+``NEURON_RT_VISIBLE_CORES`` **[Neuron Runtime]**
+
+- Integer range of specific NeuronCores needed by the process (for example, 0-3 specifies NeuronCores 0, 1, 2, and 3). Use this environment variable when launching processes to limit the launched process to specific consecutive NeuronCores.
+
+Additional Neuron runtime environment variables are described in :ref:`nrt-configuration`.

frameworks/jax/index.rst

@@ -15,6 +15,8 @@ a tested combination of the ``jax-neuronx``, ``jax``, ``jaxlib``, ``libneuronxla
 
         /frameworks/jax/setup/jax-setup
         /frameworks/jax/setup/jax-neuronx-known-issues
+        /frameworks/jax/api-reference-guide/index
 
 * :ref:`jax-neuron-setup`
-* :ref:`jax_neuron-known-issues`
+* :ref:`jax-neuron-known-issues`
+* :ref:`jax-neuronx-api-reference-guide`