Install Neuron driver and associated packages

examples/large_models/inferentia2/llama2/Readme.md

@@ -33,11 +33,6 @@ Follow the steps below to complete package installations
 sudo apt-get update
 sudo apt-get upgrade
 
-# Install Neuron libraries, SDK 2.12.2: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/prev/content.html#id8
-sudo apt-get install aws-neuronx-dkms=2.11.9.0
-sudo apt-get install aws-neuronx-collectives=2.15.16.0*
-sudo apt-get install aws-neuronx-runtime-lib=2.15.14.0*
-
 # Activate Python venv
 source /opt/aws_neuron_venv_pytorch/bin/activate
 
@@ -51,9 +46,6 @@ python ts_scripts/install_dependencies.py --neuronx --environment=dev
 # Install torchserve and torch-model-archiver
 python ts_scripts/install_from_src.py
 
-# Install additional neuron packages, SDK 2.12.2: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/prev/content.html#id8
-python -m pip install neuronx-cc==2.8.0.25 torch-neuronx==1.13.1.1.9.1 transformers-neuronx==0.5.58
-
 # Navigate to `examples/large_models/inferentia2/llama2` directory
 cd examples/large_models/inferentia2/llama2/
 

examples/large_models/inferentia2/llama2/requirements.txt

@@ -1,3 +1,3 @@
-transformers==4.31.0
-tokenizers==0.13.3
-sentencepiece==0.1.99
+transformers
+tokenizers
+sentencepiece

ts_scripts/install_dependencies.py

@@ -95,6 +95,9 @@ def install_wget(self):
     def install_numactl(self):
         pass
 
+    def install_neuronx_driver(self):
+        pass
+
 
 class Linux(Common):
     def __init__(self):
@@ -124,6 +127,29 @@ def install_numactl(self):
         if os.system("numactl --show") != 0 or args.force:
             os.system(f"{self.sudo_cmd}apt-get install -y numactl")
 
+    def install_neuronx_driver(self):
+        # Configure Linux for Neuron repository updates
+        os.system(
+            ". /etc/os-release\n"
+            + f"{self.sudo_cmd}tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF\n"
+            + "deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main\n"
+            + "EOF\n"
+            + "wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | sudo apt-key add -"
+        )
+
+        # Update OS packages
+        os.system(f"{self.sudo_cmd}apt-get update -y")
+
+        # Install OS headers
+        os.system(f"{self.sudo_cmd}apt-get install -y linux-headers-$(uname -r)")
+
+        # install Neuron Driver
+        os.system(f"{self.sudo_cmd}apt-get install -y aws-neuronx-dkms")
+
+        # Install Neuron Runtime
+        os.system(f"{self.sudo_cmd}apt-get install -y aws-neuronx-collectives")
+        os.system(f"{self.sudo_cmd}apt-get install -y aws-neuronx-runtime-lib")
+
 
 class Windows(Common):
     def __init__(self):
@@ -142,6 +168,9 @@ def install_wget(self):
     def install_numactl(self):
         pass
 
+    def install_neuronx_driver(self):
+        pass
+
 
 class Darwin(Common):
     def __init__(self):
@@ -170,6 +199,9 @@ def install_numactl(self):
         if os.system("numactl --show") != 0 or args.force:
             os.system("brew install numactl")
 
+    def install_neuronx_driver(self):
+        pass
+
 
 def install_dependencies(cuda_version=None, nightly=False):
     os_map = {"Linux": Linux, "Windows": Windows, "Darwin": Darwin}
@@ -184,6 +216,9 @@ def install_dependencies(cuda_version=None, nightly=False):
     # Sequence of installation to be maintained
     system.install_java()
 
+    if args.neuronx:
+        system.install_neuronx_driver()
+
     requirements_file = "common.txt" if args.environment == "prod" else "developer.txt"
     requirements_file_path = os.path.join("requirements", requirements_file)