Neuron SDK Release 2.24.0 (#1181)

Author: ivashkst

dlami/index.rst

@@ -89,20 +89,20 @@ Virtual Environments pre-installed
     * - Neuron Framework/Libraries supported
       - Virtual Environment
 
-    * - PyTorch 2.6 Torch NeuronX, NxD Core
-      - /opt/aws_neuronx_venv_pytorch_2_6
+    * - PyTorch 2.7 Torch NeuronX, NxD Core
+      - /opt/aws_neuronx_venv_pytorch_2_7
 
-    * - PyTorch 2.6 NxD Training, Torch NeuronX
-      - /opt/aws_neuronx_venv_pytorch_2_6_nxd_training
+    * - PyTorch 2.7 NxD Training, Torch NeuronX
+      - /opt/aws_neuronx_venv_pytorch_2_7_nxd_training
 
-    * - PyTorch 2.6 NxD Inference, Torch NeuronX
-      - /opt/aws_neuronx_venv_pytorch_2_6_nxd_inference
+    * - PyTorch 2.7 NxD Inference, Torch NeuronX
+      - /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference
 
-    * - Transformers NeuronX (PyTorch 2.6)
-      - /opt/aws_neuronx_venv_pytorch_2_6_transformers
+    * - Transformers NeuronX (PyTorch 2.7)
+      - /opt/aws_neuronx_venv_pytorch_2_7_transformers
 
-    * - JAX 0.5 NeuronX
-      - /opt/aws_neuronx_venv_jax_0_5
+    * - JAX 0.6 NeuronX
+      - /opt/aws_neuronx_venv_jax_0_6
 
     * - Tensorflow 2.10 NeuronX
       - /opt/aws_neuronx_venv_tensorflow_2_10
@@ -114,7 +114,7 @@ Virtual Environments pre-installed
       - /opt/aws_neuron_venv_pytorch_1_13_inf1
 
 
-Within the PyTorch 2.6 NxD Training virtual environment, we have included a setup script that installs required dependencies for the package. To run this script,
+Within the PyTorch 2.7 NxD Training virtual environment, we have included a setup script that installs required dependencies for the package. To run this script,
 activate the virtual environment and run ``setup_nxdt.sh`` and this will run :ref:`the setup steps here <nxdt_installation_guide>`.
 
 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 
@@ -140,25 +140,25 @@ Single Framework DLAMIs supported
       - Neuron Instances Supported
       - DLAMI Name
 
-    * - PyTorch 2.6
+    * - PyTorch 2.7
       - Ubuntu 22.04
       - Inf2, Trn1, Trn1n, Trn2
-      - Deep Learning AMI Neuron PyTorch 2.6 (Ubuntu 22.04) 
+      - Deep Learning AMI Neuron PyTorch 2.7 (Ubuntu 22.04) 
 
-    * - PyTorch 2.6
+    * - PyTorch 2.7
       - Amazon Linux 2023
       - Inf2, Trn1, Trn1n, Trn2
-      - Deep Learning AMI Neuron PyTorch 2.6 (Amazon Linux 2023)
+      - Deep Learning AMI Neuron PyTorch 2.7 (Amazon Linux 2023)
 
-    * - JAX 0.5
+    * - JAX 0.6
       - Ubuntu 22.04
       - Inf2, Trn1, Trn1n, Trn2
-      - Deep Learning AMI Neuron JAX 0.5 (Ubuntu 22.04) 
+      - Deep Learning AMI Neuron JAX 0.6 (Ubuntu 22.04) 
 
-    * - JAX 0.5
+    * - JAX 0.6
       - Amazon Linux 2023
       - Inf2, Trn1, Trn1n, Trn2
-      - Deep Learning AMI Neuron JAX 0.5 (Amazon Linux 2023) 
+      - Deep Learning AMI Neuron JAX 0.6 (Amazon Linux 2023) 
 
     * - Tensorflow 2.10
       - Ubuntu 22.04
@@ -189,25 +189,25 @@ Virtual Environments pre-installed
       - Neuron Libraries supported
       - Virtual Environment
 
-    * - Deep Learning AMI Neuron PyTorch 2.6 (Ubuntu 22.04, Amazon Linux 2023) 
-      - PyTorch 2.6 Torch NeuronX, NxD Core
-      - /opt/aws_neuronx_venv_pytorch_2_6
+    * - Deep Learning AMI Neuron PyTorch 2.7 (Ubuntu 22.04, Amazon Linux 2023) 
+      - PyTorch 2.7 Torch NeuronX, NxD Core
+      - /opt/aws_neuronx_venv_pytorch_2_7
 
-    * - Deep Learning AMI Neuron PyTorch 2.6 (Ubuntu 22.04, Amazon Linux 2023) 
-      - PyTorch 2.6 NxD Training, Torch NeuronX
-      - /opt/aws_neuronx_venv_pytorch_2_6_nxd_training
+    * - Deep Learning AMI Neuron PyTorch 2.7 (Ubuntu 22.04, Amazon Linux 2023) 
+      - PyTorch 2.7 NxD Training, Torch NeuronX
+      - /opt/aws_neuronx_venv_pytorch_2_7_nxd_training
 
-    * - Deep Learning AMI Neuron PyTorch 2.6 (Ubuntu 22.04, Amazon Linux 2023) 
-      - PyTorch 2.6 NxD Inference, Torch NeuronX
-      - /opt/aws_neuronx_venv_pytorch_2_6_nxd_inference
+    * - Deep Learning AMI Neuron PyTorch 2.7 (Ubuntu 22.04, Amazon Linux 2023) 
+      - PyTorch 2.7 NxD Inference, Torch NeuronX
+      - /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference
 
-    * - Deep Learning AMI Neuron PyTorch 2.6 (Ubuntu 22.04, Amazon Linux 2023) 
-      - Transformers NeuronX PyTorch 2.6
-      - /opt/aws_neuronx_venv_pytorch_2_6_transformers 
+    * - Deep Learning AMI Neuron PyTorch 2.7 (Ubuntu 22.04, Amazon Linux 2023) 
+      - Transformers NeuronX PyTorch 2.7
+      - /opt/aws_neuronx_venv_pytorch_2_7_transformers 
 
-    * - Deep Learning AMI Neuron JAX 0.5 (Ubuntu 22.04, Amazon Linux 2023) 
-      - JAX NeuronX 0.5
-      - /opt/aws_neuronx_venv_jax_0_5
+    * - Deep Learning AMI Neuron JAX 0.6 (Ubuntu 22.04, Amazon Linux 2023) 
+      - JAX NeuronX 0.6
+      - /opt/aws_neuronx_venv_jax_0_6
 
     * - Deep Learning AMI Neuron PyTorch 1.13 (Ubuntu 22.04) 
       - Pytorch Neuron (Inf1)
@@ -298,36 +298,30 @@ 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.6 (Ubuntu 22.04)
-      - /aws/service/neuron/dlami/pytorch-2.6/ubuntu-22.04
+    * - Deep Learning AMI Neuron PyTorch 2.7 (Ubuntu 22.04)
+      - /aws/service/neuron/dlami/pytorch-2.7/ubuntu-22.04
 
-    * - Deep Learning AMI Neuron PyTorch 2.6 (Amazon Linux 2023)
-      - /aws/service/neuron/dlami/pytorch-2.6/amazon-linux-2023
+    * - Deep Learning AMI Neuron PyTorch 2.7 (Amazon Linux 2023)
+      - /aws/service/neuron/dlami/pytorch-2.7/amazon-linux-2023
 
-    * - Deep Learning AMI Neuron JAX 0.5 (Ubuntu 22.04)
-      - /aws/service/neuron/dlami/jax-0.5/ubuntu-22.04
+    * - Deep Learning AMI Neuron JAX 0.6 (Ubuntu 22.04)
+      - /aws/service/neuron/dlami/jax-0.6/ubuntu-22.04
 
-    * - Deep Learning AMI Neuron JAX 0.5 (Amazon Linux 2023)
-      - /aws/service/neuron/dlami/jax-0.5/amazon-linux-2023
+    * - Deep Learning AMI Neuron JAX 0.6 (Amazon Linux 2023)
+      - /aws/service/neuron/dlami/jax-0.6/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
 
-    * - Deep Learning Base Neuron AMI (Ubuntu 20.04)
-      - /aws/service/neuron/dlami/base/ubuntu-20.04
-
 
 For example to find the latest DLAMI id for Multi-Framework DLAMI (Ubuntu 22) you can use the following
 

frameworks/torch/torch-neuronx/tutorials/training/finetune_hftrainer.rst

@@ -62,15 +62,15 @@ First, paste the following script into your terminal to create a “run.sh” fi
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_single_worker_training.sh
    :language: shell
-   :lines: 7-27
+   :lines: 7-28
 
 We optionally precompile the model and training script using neuron_parallel_compile to warm up the persistent
 graph cache (Neuron Cache) such that the actual run has fewer compilations (faster run
 time):
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_single_worker_training.sh
    :language: shell
-   :lines: 30
+   :lines: 31
 
 Please ignore the results from this precompile run as it is only for
 extracting and compiling the XLA graphs.
@@ -86,7 +86,7 @@ additional compilations.
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_single_worker_training.sh
    :language: shell
-   :lines: 32
+   :lines: 33
 
 If precompilation was not done, the first execution of ./run.sh will be slower due to serial compilations. Rerunning the same script a second time would show quicker execution as the compiled graphs will be already cached in persistent cache.
 
@@ -108,23 +108,23 @@ Paste the following script into your terminal to create a “run_2w.sh” file a
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_multi_worker_training_code.sh
    :language: shell
-   :lines: 7-27
+   :lines: 7-28
 
 Again, we optionally precompile the model and training script using neuron_parallel_compile to warm up the persistent
 graph cache (Neuron Cache), ignoring the results from this precompile run as it is only for
 extracting and compiling the XLA graphs:
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_multi_worker_training_code.sh
    :language: shell
-   :lines: 30
+   :lines: 31
 
 Precompilation is optional and only needed to be done once unless hyperparameters such as batch size are modified.
 After the optional precompilation, the actual run will be faster with minimal
 additional compilations.
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_multi_worker_training_code.sh
    :language: shell
-   :lines: 32
+   :lines: 33
 
 During run, you will now notice that the "Total train batch size" is now 16 and the "Total optimization steps" is now half the number for one worker training.
 
@@ -149,19 +149,19 @@ Paste the following script into your terminal to create a “run_converted.sh”
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_converted_checkpoint_training.sh
    :language: shell
-   :lines: 38-59
+   :lines: 38-60
 
 If it is the first time running with ``bert-large-uncased`` model or if hyperparameters have changed, then the optional one-time precompilation step can save compilation time:
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_converted_checkpoint_training.sh
    :language: shell
-   :lines: 62
+   :lines: 63
 
 If you have run the single worker training in a previous section, then you can skip the precompilation step and just do:
 
 .. literalinclude:: tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_converted_checkpoint_training.sh
    :language: shell
-   :lines: 65
+   :lines: 66
 
 .. _workarounds_for_older_versions:
 
@@ -234,6 +234,7 @@ The following are currently known issues:
 - Variable input sizes: When fine-tune models such as dslim/bert-base-NER using the `token-classification example <https://github.com/huggingface/transformers/tree/main/examples/pytorch/token-classification>`__, you may encounter timeouts (lots of "socket.h:524 CCOM WARN Timeout waiting for RX" messages) and execution hang. This occurs because NER dataset has different sample sizes, which causes many recompilations and compiled graph (NEFF) reloads. Furthermore, different data parallel workers can execute different compiled graph. This multiple-program multiple-data behavior is currently unsupported. To workaround this issue, please pad to maximum length using the Trainer API option ``--pad_to_max_length``.
 - When running HuggingFace GPT fine-tuning with transformers version >= 4.21.0 and using XLA_USE_BF16=1 or XLA_DOWNCAST_BF16=1, you might see NaNs in the loss immediately at the first step. This issue occurs due to large negative constants used to implement attention masking (https://github.com/huggingface/transformers/pull/17306). To workaround this issue, please use transformers version <= 4.20.0.
 - When using Trainer API option --bf16, you will see "RuntimeError: No CUDA GPUs are available". To workaround this error, please add "import torch; torch.cuda.is_bf16_supported = lambda: True" to the Python script (i.e. run_glue.py). (Trainer API option --fp16 is not yet supported).
+- When using latest HuggingFace transformers version, you may see "ValueError: Your setup doesn't support bf16/gpu." To fix this, please use ``--use_cpu True`` in your scripts.
 
 The following are resolved issues:
 

frameworks/torch/torch-neuronx/tutorials/training/tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_converted_checkpoint_training.sh

@@ -47,6 +47,7 @@ NEURON_RT_STOCHASTIC_ROUNDING_EN=1 torchrun --nproc_per_node=2 ./run_glue.py \\
 --do_train \\
 --do_eval \\
 --bf16 \\
+--use_cpu True \\
 --max_seq_length 128 \\
 --per_device_train_batch_size 8 \\
 --learning_rate 2e-5 \\

frameworks/torch/torch-neuronx/tutorials/training/tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_multi_worker_training_code.sh

@@ -15,6 +15,7 @@ NEURON_RT_STOCHASTIC_ROUNDING_EN=1 torchrun --nproc_per_node=2 ./run_glue.py \\
 --do_train \\
 --do_eval \\
 --bf16 \\
+--use_cpu True \\
 --max_seq_length 128 \\
 --per_device_train_batch_size 8 \\
 --learning_rate 2e-5 \\

frameworks/torch/torch-neuronx/tutorials/training/tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_setup_code.sh

@@ -2,8 +2,9 @@
 set -eExuo
 
 # Install packages and clone transformers
-export HF_VER=4.44.0
-pip install -U transformers==$HF_VER datasets evaluate scikit-learn
+export HF_VER=4.52.0
+export ACC_VER=1.7.0
+pip install -U transformers==$HF_VER accelerate==$ACC_VER datasets evaluate scikit-learn
 cd ~/
 git clone https://github.com/huggingface/transformers --branch v$HF_VER
-cd ~/transformers/examples/pytorch/text-classification
+cd ~/transformers/examples/pytorch/text-classification

frameworks/torch/torch-neuronx/tutorials/training/tutorial_source_code/bert_mrpc_finetuning/bert_mrpc_finetuning_single_worker_training.sh

@@ -15,6 +15,7 @@ NEURON_RT_STOCHASTIC_ROUNDING_EN=1 torchrun --nproc_per_node=1 ./run_glue.py \\
 --do_train \\
 --do_eval \\
 --bf16 \\
+--use_cpu True \\
 --max_seq_length 128 \\
 --per_device_train_batch_size 8 \\
 --learning_rate 2e-5 \\

general/appnotes/index.rst

@@ -81,5 +81,6 @@ Neuron Application Notes
          .. toctree::
             :maxdepth: 1
 
+            /general/appnotes/torch-neuronx/introducing-pytorch-2-7
             /general/appnotes/torch-neuronx/introducing-pytorch-2-6
             /general/appnotes/torch-neuronx/introducing-pytorch-2-x

general/appnotes/torch-neuronx/introducing-pytorch-2-6.rst

@@ -41,7 +41,7 @@ See :ref:`migrate_to_pytorch_2.6` for changes needed to use PyTorch NeuronX 2.6.
 How can I install PyTorch NeuronX 2.6?
 --------------------------------------------
 
-To install PyTorch NeuronX 2.6 please follow the :ref:`setup-torch-neuronx` guides for Amazon Linux 2023 and Ubuntu 22 AMI. Please also refer to the Neuron multi-framework DLAMI :ref:`setup guide <setup-ubuntu22-multi-framework-dlami>` for Ubuntu 22 with a pre-installed virtual environment for PyTorch NeuronX 2.6 that you can use to easily get started. PyTorch NeuronX 2.6 can be installed using the following:
+To install PyTorch NeuronX 2.6 please follow the :ref:`setup-torch-neuronx` guides for Amazon Linux 2023 and Ubuntu 22 AMI. Please also refer to the Neuron multi-framework DLAMI :ref:`setup guide <setup-ubuntu22-multi-framework-dlami>` for Ubuntu 22 with a pre-installed virtual environment for PyTorch NeuronX 2.6 that you can use to get started. PyTorch NeuronX 2.6 can be installed using the following:
 
 .. code::
 
@@ -66,7 +66,7 @@ To migrate the training scripts from PyTorch NeuronX 2.5 to PyTorch NeuronX 2.6,
 
 .. note::
 
-    ``xm`` below refers to ``torch_xla.core.xla_model`` and ``xr`` refers to ``torch_xla.runtime``
+    ``xm`` below refers to ``torch_xla.core.xla_model``, ``xr`` refers to ``torch_xla.runtime``, and ``xmp`` refers to ``torch_xla.distributed.xla_multiprocessing``
 
 * The environment variables ``XLA_DOWNCAST_BF16`` and ``XLA_USE_BF16`` are deprecated (warning when used) and will be removed in an upcoming release. Please switch to automatic mixed-precision or use ``model.to(torch.bfloat16)`` command to convert model to BF16 format. (see :ref:`migration_from_xla_downcast_bf16`)
 * The functions ``xm.xrt_world_size()``, ``xm.xla_model.get_ordinal()``, and ``xm.xla_model.get_local_ordinal()`` are deprecated (warning when used). Please switch to ``xr.world_size()``, ``xr.global_ordinal()``, and ``xr.local_ordinal()`` respectively as replacements.
@@ -123,20 +123,21 @@ Warning "XLA_DOWNCAST_BF16 will be deprecated after the 2.6 release, please down
 Environment variables ``XLA_DOWNCAST_BF16`` and ``XLA_USE_BF16`` are deprecated (warning when used). Please switch to automatic mixed-precision or use ``model.to(torch.bfloat16)`` command to cast model to BF16. (see :ref:`migration_from_xla_downcast_bf16`)
 
 
-AttributeError: <module 'torch_xla.core.xla_model' ... does not have the attribute 'xrt_world_size'
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+WARNING:root:torch_xla.core.xla_model.xrt_world_size() will be removed in release 2.7. is deprecated. Use torch_xla.runtime.world_size instead.
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+This is a warning that ``torch_xla.core.xla_model.xrt_world_size()`` will be removed in a future release. Please switch to using ``torch_xla.runtime.world_size`` instead.
 
-This is an error that ``torch_xla.core.xla_model.xrt_world_size()`` is removed in torch-xla version 2.7. Please switch to using ``torch_xla.runtime.world_size()`` instead.
 
-AttributeError: <module 'torch_xla.core.xla_model' ... does not have the attribute 'get_ordinal'
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+WARNING:torch_xla.core.xla_model.xla_model.get_ordinal() will be removed in release 2.7. is deprecated. Use torch_xla.runtime.global_ordinal instead.
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-This is an error that ``torch_xla.core.xla_model.xla_model.get_ordinal()`` is removed in torch-xla version 2.7. Please switch to using ``torch_xla.runtime.global_ordinal()`` instead.
+This is a warning that ``torch_xla.core.xla_model.xla_model.get_ordinal()`` will be removed in a future release. Please switch to using ``torch_xla.runtime.global_ordinal`` instead.
 
-AttributeError: <module 'torch_xla.core.xla_model' ... does not have the attribute 'get_local_ordinal'
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+WARNING:torch_xla.core.xla_model.xla_model.get_local_ordinal() will be removed in release 2.7. is deprecated. Use torch_xla.runtime.local_ordinal instead.
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-This is an error that ``torch_xla.core.xla_model.xla_model.get_local_ordinal()`` is removed in torch-xla version 2.7. Please switch to using ``torch_xla.runtime.local_ordinal()`` instead.
+This is a warning that ``torch_xla.core.xla_model.xla_model.get_local_ordinal()`` will be removed in a future release. Please switch to using ``torch_xla.runtime.local_ordinal`` instead.
 
 
 Socket Error: Socket failed to bind