Add support for inf2 dependency handling

Author: Naman Nandan

requirements/neuronx.txt

@@ -0,0 +1,10 @@
+numpy==1.21.6
+protobuf==3.20.3
+torch==1.13.1+cpu
+torchvision==0.14.1+cpu
+torchtext==0.14.1
+torchaudio==0.13.1+cpu
+torchdata==0.5.1
+neuronx-cc
+torch-neuronx
+transformers-neuronx

ts_scripts/install_dependencies.py

@@ -67,6 +67,13 @@ def install_python_packages(self, cuda_version, requirements_file_path, nightly)
             gpu_requirements_file = os.path.join("requirements", "common_gpu.txt")
             os.system(f"{sys.executable} -m pip install -U -r {gpu_requirements_file}")
 
+        # Install dependencies for Inferentia2
+        if args.neuronx:
+            neuronx_requirements_file = os.path.join("requirements", "neuronx.txt")
+            os.system(
+                f"{sys.executable} -m pip install -U -r {neuronx_requirements_file}"
+            )
+
     def install_node_packages(self):
         os.system(
             f"{self.sudo_cmd}npm install -g newman newman-reporter-htmlextra markdown-link-check"
@@ -193,6 +200,11 @@ def get_brew_version():
         choices=["cu92", "cu101", "cu102", "cu111", "cu113", "cu116", "cu117", "cu118"],
         help="CUDA version for torch",
     )
+    parser.add_argument(
+        "--neuronx",
+        action="store_true",
+        help="Install dependencies for inferentia2 support",
+    )
     parser.add_argument(
         "--environment",
         default="prod",