diff --git a/README.md b/README.md index c92902c..01ddf3e 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,16 @@ For the 40b model: ``` python -m evo2.test.test_evo2_generation --model_name evo2_40b ``` +### Example conda environment + +For users who prefer a reproducible setup, we provide an example conda environment file (`environment.yml`) that pins Python, CUDA, PyTorch, Transformer Engine, FlashAttention and Vortex. + +You can create and activate the environment with: + +```bash +conda env create -f environment.yml +conda activate evo2 +python -m evo2.test.test_evo2_generation --model_name evo2_7b ### Docker @@ -100,6 +110,14 @@ Once inside the container: ```bash python -m evo2.test.test_evo2_generation --model_name evo2_7b ``` +### Troubleshooting (installation) + +Some common installation errors and suggested fixes: + +- `ModuleNotFoundError: No module named 'transformer_engine'` + Make sure the PyTorch extension for Transformer Engine is installed, for example: + ```bash + pip install "transformer_engine[pytorch]" ## Usage diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..5c1370d --- /dev/null +++ b/environment.yml @@ -0,0 +1,16 @@ +name: evo2 +channels: + - nvidia + - conda-forge + - pytorch + - defaults +dependencies: + - python=3.12 + - pip + - pytorch=2.5.1 + - pytorch-cuda=12.1 + - transformer-engine-torch=2.3.0 + - pip: + - flash-attn==2.8.0.post2 --no-build-isolation + - vtx>=0.0.8 + - evo2