diff --git a/noxfile.py b/noxfile.py index 0d6680460..f71510979 100644 --- a/noxfile.py +++ b/noxfile.py @@ -142,6 +142,8 @@ def tests_brevitas_examples_llm(session, pytorch, jit_status): install_pytorch(pytorch, session) install_torchvision(pytorch, session) # Optimum seems to require torchvision session.install('-e', '.[test, llm, export]') + session.install( + 'optimum-amd[brevitas] @ git+https://github.com/huggingface/optimum-amd.git@main') session.run('pytest', '-n', 'logical', '-k', 'llm', 'tests/brevitas_examples/test_llm.py') diff --git a/requirements/requirements-llm.txt b/requirements/requirements-llm.txt index 7070cc9c6..8935cb159 100644 --- a/requirements/requirements-llm.txt +++ b/requirements/requirements-llm.txt @@ -1,3 +1,3 @@ -optimum-amd[brevitas] @ git+https://github.com/huggingface/optimum-amd.git@main +# optimum-amd[brevitas] @ git+https://github.com/huggingface/optimum-amd.git@main tqdm transformers diff --git a/src/brevitas_examples/llm/README.md b/src/brevitas_examples/llm/README.md index 838eaed03..457c74804 100644 --- a/src/brevitas_examples/llm/README.md +++ b/src/brevitas_examples/llm/README.md @@ -2,12 +2,12 @@ ## Requirements -- transformers (from source) +- transformers - datasets - torch_mlir (optional for torch-mlir based export) -- optimum (from source) -- optimum-amd (WIP, install brevitas-compatibility branch) -- accelerate (from source) +- optimum +- optimum-amd (install from main) +- accelerate ## Run