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': {