Test (example/llm): Refactor and add basic tests for the LLM entry-po…

…int (#1002)

.github/workflows/examples_llm_pytest.yml

@@ -0,0 +1,65 @@
+name: Examples LLM Pytest
+
+on:
+  push:
+    branches: [ master, dev ]
+  pull_request:
+    types:
+      - review_requested
+
+jobs:
+  build:
+    runs-on: ${{ matrix.platform }}
+    strategy:
+      fail-fast: false
+
+
+      matrix:
+        python_version: ['3.8', '3.9']
+        pytorch_version: ['2.2.2', '2.3.1', '2.4.0']
+        platform: ['windows-latest', 'ubuntu-latest', 'macos-latest']
+        jit_status: ['jit_disabled', 'jit_enabled']
+
+
+        exclude:
+          - pytorch_version: '1.9.1'
+            platform: 'macos-latest'
+
+          - pytorch_version: '1.9.1'
+            jit_status: 'jit_enabled'
+
+
+
+    if: ${{ !github.event.pull_request.draft }}
+    steps:
+
+    - name: Checkout repo
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+       python-version: ${{ matrix.python_version }}
+
+    - name: Install Nox dependencies
+      shell: bash
+      run: pip install -r requirements/requirements-nox.txt
+
+    - name: Install update
+      shell: bash
+      run: sudo apt-get update
+      if: startsWith(runner.os, 'Linux') == true
+
+    - name: Install libsndfile and libgomp1 on Ubuntu
+      shell: bash
+      run: sudo apt-get install -y libsndfile-dev libgomp1
+      if: startsWith(runner.os, 'Linux') == true
+
+    - name: Install libomp on macOS
+      shell: bash
+      run: brew install libomp
+      if: startsWith(runner.os, 'macOS') == true
+
+    - name: Run Nox session for brevitas_examples pytest
+      shell: bash
+      run: nox -v -s tests_brevitas_examples_llm-${{ matrix.python_version }}\(${{ matrix.jit_status }}\,\ pytorch_${{ matrix.pytorch_version }}\)

.github/workflows/gen_github_actions.py

@@ -8,6 +8,7 @@
 BASE_YML_REDUCED_TEMPLATE = 'base_reduced.yml.template'
 PYTEST_YML = 'pytest.yml'
 EXAMPLES_PYTEST_YML = 'examples_pytest.yml'
+EXAMPLES_LLM_PYTEST_YML = 'examples_llm_pytest.yml'
 DEVELOP_INSTALL_YML = 'develop_install.yml'
 FINN_INTEGRATION_YML = 'finn_integration.yml'
 ORT_INTEGRATION_YML = 'ort_integration.yml'
@@ -25,6 +26,10 @@
                      ('pytorch_version', list(PYTORCH_LIST_REDUCED)),
                      ('platform', PLATFORM_LIST_REDUCED)])
 
+EXAMPLES_LLM_PYTEST_MATRIX_REDUCED = od([('python_version', list(PYTHON_VERSIONS_REDUCED)),
+                                         ('pytorch_version', list(
+                                             ('2.4.0',))), ('platform', PLATFORM_LIST_REDUCED)])
+
 FINN_MATRIX_REDUCED = od([('python_version', list(PYTHON_VERSIONS_REDUCED)),
                           ('pytorch_version', list(PYTORCH_LIST_REDUCED)),
                           ('platform', PLATFORM_LIST_REDUCED)])
@@ -61,6 +66,11 @@
 MATRIX = od([('python_version', list(PYTHON_VERSIONS)), ('pytorch_version', list(PYTORCH_VERSIONS)),
              ('platform', PLATFORM_LIST)])
 
+EXAMPLES_LLM_PYTEST_PYTORCH_VERSIONS = ('2.2.2', '2.3.1', '2.4.0')
+EXAMPLES_LLM_PYTEST_MATRIX = od([('python_version', list(PYTHON_VERSIONS)),
+                                 ('pytorch_version', list(EXAMPLES_LLM_PYTEST_PYTORCH_VERSIONS)),
+                                 ('platform', PLATFORM_LIST)])
+
 FINN_MATRIX = od([('python_version', list(PYTHON_VERSIONS)),
                   ('pytorch_version', list(PYTORCH_VERSIONS)), ('platform', FINN_PLATFORM_LIST)])
 
@@ -80,6 +90,13 @@
             'nox -v -s tests_brevitas_examples_cpu-${{ matrix.python_version }}\(${{ matrix.jit_status }}\,\ pytorch_${{ matrix.pytorch_version }}\)'
         )]),]
 
+EXAMPLES_LLM_PYTEST_STEP_LIST = [
+    od([('name', 'Run Nox session for brevitas_examples pytest'), ('shell', 'bash'),
+        (
+            'run',
+            'nox -v -s tests_brevitas_examples_llm-${{ matrix.python_version }}\(${{ matrix.jit_status }}\,\ pytorch_${{ matrix.pytorch_version }}\)'
+        )]),]
+
 FINN_INTEGRATION_STEP_LIST = [
     od([('name', 'Install protobuf on Ubuntu'), ('shell', 'bash'),
         ('run', 'sudo apt-get install protobuf-compiler libprotoc-dev'),
@@ -167,6 +184,23 @@ def gen_examples_pytest_yml():
     pytest.gen_yaml(BASE_YML_REDUCED_TEMPLATE, 'reduced_' + EXAMPLES_PYTEST_YML)
 
 
+def gen_examples_llm_pytest_yml():
+    pytest = Action(
+        'Examples LLM Pytest',
+        EXCLUDE_LIST + JIT_EXCLUDE_LIST,
+        combine_od_list([EXAMPLES_LLM_PYTEST_MATRIX, PYTEST_MATRIX_EXTRA]),
+        EXAMPLES_LLM_PYTEST_STEP_LIST,
+        STRATEGY)
+    pytest.gen_yaml(BASE_YML_TEMPLATE, EXAMPLES_LLM_PYTEST_YML)
+    pytest = Action(
+        'Examples LLM Pytest',
+        EXCLUDE_LIST,
+        combine_od_list([EXAMPLES_LLM_PYTEST_MATRIX_REDUCED, PYTEST_MATRIX_EXTRA_REDUCED]),
+        EXAMPLES_LLM_PYTEST_STEP_LIST,
+        STRATEGY)
+    pytest.gen_yaml(BASE_YML_REDUCED_TEMPLATE, 'reduced_' + EXAMPLES_LLM_PYTEST_YML)
+
+
 def gen_test_develop_install_yml():
     test_develop_install = Action(
         'Test develop install', EXCLUDE_LIST, MATRIX, TEST_INSTALL_DEV_STEP_LIST, STRATEGY)
@@ -243,6 +277,7 @@ def gen_test_brevitas_end_to_end():
 if __name__ == '__main__':
     gen_pytest_yml()
     gen_examples_pytest_yml()
+    gen_examples_llm_pytest_yml()
     gen_test_develop_install_yml()
     gen_test_brevitas_finn_integration()
     gen_test_brevitas_ort_integration()

.github/workflows/reduced_examples_llm_pytest.yml

@@ -0,0 +1,64 @@
+name: Examples LLM Pytest
+
+on:
+  pull_request:
+    types:
+      - opened
+      - reopened
+      - synchronize
+      - ready_for_review
+
+
+jobs:
+  build:
+    runs-on: ${{ matrix.platform }}
+    strategy:
+      fail-fast: false
+
+
+      matrix:
+        python_version: ['3.8']
+        pytorch_version: ['2.4.0']
+        platform: ['ubuntu-latest']
+        jit_status: ['jit_disabled']
+
+
+        exclude:
+          - pytorch_version: '1.9.1'
+            platform: 'macos-latest'
+
+
+
+    if: ${{ !github.event.pull_request.draft }}
+    steps:
+
+    - name: Checkout repo
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+       python-version: ${{ matrix.python_version }}
+
+    - name: Install Nox dependencies
+      shell: bash
+      run: pip install -r requirements/requirements-nox.txt
+
+    - name: Install update
+      shell: bash
+      run: sudo apt-get update
+      if: startsWith(runner.os, 'Linux') == true
+
+    - name: Install libsndfile and libgomp1 on Ubuntu
+      shell: bash
+      run: sudo apt-get install -y libsndfile-dev libgomp1
+      if: startsWith(runner.os, 'Linux') == true
+
+    - name: Install libomp on macOS
+      shell: bash
+      run: brew install libomp
+      if: startsWith(runner.os, 'macOS') == true
+
+    - name: Run Nox session for brevitas_examples pytest
+      shell: bash
+      run: nox -v -s tests_brevitas_examples_llm-${{ matrix.python_version }}\(${{ matrix.jit_status }}\,\ pytorch_${{ matrix.pytorch_version }}\)

noxfile.py

@@ -7,15 +7,20 @@
 
 import nox
 from packaging import version
+from packaging.version import parse
 
 sys.path.append(os.path.join(os.path.dirname(__file__), os.path.join('.', '.github', 'workflows')))
+from gen_github_actions import EXAMPLES_LLM_PYTEST_PYTORCH_VERSIONS
 from gen_github_actions import JIT_STATUSES
 from gen_github_actions import PYTHON_VERSIONS
 from gen_github_actions import PYTORCH_VERSIONS
 
 IS_OSX = system() == 'Darwin'
-PYTORCH_STABLE_WHEEL_SRC = 'https://download.pytorch.org/whl/torch_stable.html'
+PYTORCH_STABLE_WHEEL_SRC = 'https://download.pytorch.org/whl/cpu'
+PYTORCH_STABLE_WHEEL_SRC_LEGACY = 'https://download.pytorch.org/whl/torch_stable.html'
 PYTORCH_IDS = tuple([f'pytorch_{i}' for i in PYTORCH_VERSIONS])
+EXAMPLES_LLM_PYTEST_PYTORCH_IDS = tuple([
+    f'pytorch_{i}' for i in EXAMPLES_LLM_PYTEST_PYTORCH_VERSIONS])
 JIT_IDS = tuple([f'{i}'.lower() for i in JIT_STATUSES])
 LSTM_EXPORT_MIN_PYTORCH = '1.10.1'
 
@@ -26,14 +31,21 @@
     '1.12.1': '0.13.1',
     '1.13.0': '0.14.0',
     '2.0.1': '0.15.2',
-    '2.1.0': '0.16.0'}
+    '2.1.0': '0.16.0',
+    '2.2.2': '0.17.2',
+    '2.3.1': '0.18.1',
+    '2.4.0': '0.19.0'}
 
 PARSED_TORCHVISION_VERSION_DICT = {version.parse(k): v for k, v in TORCHVISION_VERSION_DICT.items()}
 
 
 def install_pytorch(pytorch, session):
     if not IS_OSX:
-        cmd = [f'torch=={pytorch}+cpu', '-f', PYTORCH_STABLE_WHEEL_SRC]
+        if parse(pytorch) < parse('2.4.0'):
+            cmd = [f'torch=={pytorch}+cpu', '-f', PYTORCH_STABLE_WHEEL_SRC_LEGACY]
+        else:
+            cmd = [f'torch=={pytorch}', '--index-url', PYTORCH_STABLE_WHEEL_SRC]
+
     else:
         cmd = [f'torch=={pytorch}']
     session.install(*cmd)
@@ -42,11 +54,18 @@ def install_pytorch(pytorch, session):
 def install_torchvision(pytorch, session):
     torchvision = PARSED_TORCHVISION_VERSION_DICT[version.parse(pytorch)]
     if not IS_OSX:
-        cmd = [
-            f'torch=={pytorch}+cpu',  # make sure correct pytorch version is kept
-            f'torchvision=={torchvision}+cpu',
-            '-f',
-            PYTORCH_STABLE_WHEEL_SRC]
+        if parse(pytorch) < parse('2.4.0'):
+            cmd = [
+                f'torch=={pytorch}+cpu',  # make sure correct pytorch version is kept
+                f'torchvision=={torchvision}+cpu',
+                '-f',
+                PYTORCH_STABLE_WHEEL_SRC_LEGACY]
+        else:
+            cmd = [
+                f'torch=={pytorch}',
+                f'torchvision=={torchvision}',
+                '--index-url',
+                PYTORCH_STABLE_WHEEL_SRC]
     else:
         cmd = [f'torch=={pytorch}', f'torchvision=={torchvision}']
     session.install(*cmd)
@@ -105,7 +124,25 @@ def tests_brevitas_examples_cpu(session, pytorch, jit_status):
     install_pytorch(pytorch, session)
     install_torchvision(pytorch, session)  # For CV eval scripts
     session.install('--upgrade', '.[test, tts, stt, vision]')
-    session.run('pytest', '-n', 'logical', 'tests/brevitas_examples')
+    session.run(
+        'pytest',
+        '-n',
+        'logical',
+        '--ignore-glob',
+        'tests/brevitas_examples/*llm*',
+        'tests/brevitas_examples')
+
+
+@nox.session(python=PYTHON_VERSIONS)
+@nox.parametrize(
+    "pytorch", EXAMPLES_LLM_PYTEST_PYTORCH_VERSIONS, ids=EXAMPLES_LLM_PYTEST_PYTORCH_IDS)
+@nox.parametrize("jit_status", JIT_STATUSES, ids=JIT_IDS)
+def tests_brevitas_examples_llm(session, pytorch, jit_status):
+    session.env['BREVITAS_JIT'] = '{}'.format(int(jit_status == 'jit_enabled'))
+    install_pytorch(pytorch, session)
+    install_torchvision(pytorch, session)  # Optimum seems to require torchvision
+    session.install('-e', '.[test, llm, export]')
+    session.run('pytest', '-n', 'logical', '-k', 'llm', 'tests/brevitas_examples/test_llm.py')
 
 
 @nox.session(python=PYTHON_VERSIONS)

pytest.ini

@@ -7,3 +7,6 @@ log_cli_format = %(message)s
 
 # pytest-mock should use Pypi's mock rather than Python's built-in
 mock_use_standalone_module = true
+
+markers =
+    llm: mark a test which tests brevitas_examples/llm

requirements/requirements-llm.txt

@@ -0,0 +1,3 @@
+optimum-amd[brevitas] @ git+https://github.com/huggingface/optimum-amd.git@main
+tqdm
+transformers

setup.py

@@ -39,6 +39,7 @@ def read_requirements(filename):
         "test": read_requirements('requirements-test.txt'),
         "tts": read_requirements('requirements-tts.txt'),
         "stt": read_requirements('requirements-stt.txt'),
+        "llm": read_requirements('requirements-llm.txt'),
         "vision": read_requirements('requirements-vision.txt'),
         "finn_integration": read_requirements('requirements-finn-integration.txt'),
         "ort_integration": read_requirements('requirements-ort-integration.txt')},

src/brevitas_examples/common/generative/quantize.py

@@ -153,7 +153,8 @@
                         'sym': Int8DynamicActPerGroupFloat}}},
             'po2_scale': {
                 'stats': {
-                    'per_group': MXInt8Act}}}},
+                    'per_group': {
+                        'sym': MXInt8Act}}}}},
     'float': {
         'static': {
             'float_scale': {
@@ -175,7 +176,8 @@
         'dynamic': {
             'po2_scale': {
                 'stats': {
-                    'per_group': MXFloat8e4m3Act}}}},
+                    'per_group': {
+                        'sym': MXFloat8e4m3Act}}}}},
     'float_fnuz': {
         'static': {
             'float_scale': {