wip

Signed-off-by: Raphael Glon <[email protected]>

Author: oOraph

dockerfiles/pytorch/neuronx.sh

@@ -8,7 +8,10 @@ apt-get update -y \
     gnupg2 \
     wget
 
-echo "deb https://apt.repos.neuron.amazonaws.com jammy main" > /etc/apt/sources.list.d/neuron.list
+. /etc/os-release
+tee /etc/apt/sources.list.d/neuron.list > /dev/null <<EOF
+deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main
+EOF
 wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -
 
 apt-get update -y \
@@ -20,11 +23,11 @@ apt-get update -y \
 
 pip install -U pip
 
-pip3 install neuronx-cc==2.12.68.0 \
-    torch-neuronx==1.13.1.1.13.1 \
-    transformers-neuronx==0.9.474 \
+pip3 install neuronx-cc==2.* \
+    torch-neuronx==2.* \
+    transformers-neuronx\
     --extra-index-url=https://pip.repos.neuron.amazonaws.com
 
-pip3 install optimum[neuronx,diffusers]
+pip3 install --extra-index-url=https://pip.repos.neuron.amazonaws.com optimum[neuronx,diffusers]
 
-pip install ".[st,torch1]"
+pip install ".[st,torch-neuronx]"

setup.py

@@ -17,8 +17,8 @@
     "wheel==0.42.0",
     "setuptools==69.1.0",
     "cmake==3.28.3",
-    "transformers[sklearn,sentencepiece, audio, vision]==4.38.2",
-    "huggingface_hub==0.20.3",
+    "transformers[sklearn,sentencepiece, audio, vision]>=4.36.0",
+    "huggingface_hub==0.23.0",
     "orjson",
     # vision
     "Pillow",
@@ -40,7 +40,7 @@
 extras["diffusers"] = ["diffusers==0.26.3", "accelerate==0.27.2"]
 extras["torch"] = ["torch==2.2.0", "torchvision", "torchaudio"]
 # For neuronx
-extras["torch1"] = ["torch==1.13.1", "torchvision", "torchaudio"]
+extras["torch-neuronx"] = ["torch-neuronx", "torchvision", "torchaudio"]
 extras["tensorflow"] = ["tensorflow"]
 extras["test"] = [
     "pytest==7.2.1",