resolve conflicts

Signed-off-by: ericharper <complex451@gmail.com>

.github/workflows/import-test.yml

@@ -6,25 +6,24 @@ on:
     paths:
       - "**"
 
+# Check https://hub.docker.com/r/pytorch/pytorch/tags for latest tags
 jobs:
-  ci-import-check:
-    runs-on: ubuntu-latest
 
-    # Check https://hub.docker.com/r/pytorch/pytorch/tags for latest tags
+  test-asr-imports:
+    runs-on: ubuntu-latest
     container:
-      image: pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime
-
+      image: pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
     steps:
-    - uses: actions/checkout@v2
-
+    - name: Checkout repo
+      uses: actions/checkout@v2
     - name: Update base dependencies
       run: |
         apt-get update && apt-get install -y build-essential
         apt-get install -y libsndfile1 make
-
     - name: Install nemo dependencies
       id: nemo-wheel
       run:  |
+        pip install Cython
         # install test requirements
         pip install -r requirements/requirements_test.txt
         # Build nemo as a wheel
@@ -33,7 +32,6 @@ jobs:
         # Preserve wheel location
         DIST_FILE=$(find ./dist -name "*.whl" | head -n 1)
         echo "::set-output name=DIST_FILE::${DIST_FILE}"
-
     - name: Test ASR Domain Imports
       run: |
         # Install NeMo Domain
@@ -43,6 +41,29 @@ jobs:
         # Uninstall NeMo
         pip uninstall -y nemo_toolkit
 
+  test-tts-imports:
+    runs-on: ubuntu-latest
+    container:
+      image: pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+    steps:
+    - name: Checkout repo
+      uses: actions/checkout@v2
+    - name: Update base dependencies
+      run: |
+        apt-get update && apt-get install -y build-essential
+        apt-get install -y libsndfile1 make
+    - name: Install nemo dependencies
+      id: nemo-wheel
+      run:  |
+        pip install Cython
+        # install test requirements
+        pip install -r requirements/requirements_test.txt
+        # Build nemo as a wheel
+        pip install build
+        python -m build --no-isolation --wheel
+        # Preserve wheel location
+        DIST_FILE=$(find ./dist -name "*.whl" | head -n 1)
+        echo "::set-output name=DIST_FILE::${DIST_FILE}"
     - name: Test TTS Domain Imports
       run: |
         # Install NeMo Domain
@@ -52,6 +73,29 @@ jobs:
         # Uninstall NeMo
         pip uninstall -y nemo_toolkit
 
+  test-nlp-imports:
+    runs-on: ubuntu-latest
+    container:
+      image: pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
+    steps:
+    - name: Checkout repo
+      uses: actions/checkout@v2
+    - name: Update base dependencies
+      run: |
+        apt-get update && apt-get install -y build-essential
+        apt-get install -y libsndfile1 make
+    - name: Install nemo dependencies
+      id: nemo-wheel
+      run:  |
+        pip install Cython
+        # install test requirements
+        pip install -r requirements/requirements_test.txt
+        # Build nemo as a wheel
+        pip install build
+        python -m build --no-isolation --wheel
+        # Preserve wheel location
+        DIST_FILE=$(find ./dist -name "*.whl" | head -n 1)
+        echo "::set-output name=DIST_FILE::${DIST_FILE}"
     - name: Test NLP Domain Imports
       run: |
         # Install NeMo Domain

Dockerfile

@@ -72,6 +72,11 @@ WORKDIR /tmp/nemo
 COPY requirements .
 RUN for f in $(ls requirements*.txt); do pip3 install --disable-pip-version-check --no-cache-dir -r $f; done
 
+# install flash attention dependencies
+RUN pip install flash-attn
+# pinned triton version for flash-attention https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attn_triton.py#L3
+RUN pip install triton==2.0.0.dev20221202 
+
 # install k2, skip if installation fails
 COPY scripts /tmp/nemo/scripts/
 RUN INSTALL_MSG=$(/bin/bash /tmp/nemo/scripts/speech_recognition/k2/setup.sh); INSTALL_CODE=$?; \