Fix and improve env build options (#105)

Summary:

Test Plan:

README.md

@@ -28,8 +28,25 @@ Please file an issue and/or be patient while we make more of our code public!
 
 ## Quick start
 
-The following commands launch a SLURM job that creates an environment for Meta Lingua.
+There are three ways you can create your environment.
-The env creation should take around 5 minutes without counting downloads.
+
+### Option 1: conda + pip
+
+Run these commands in your terminal or a script:
+
+```bash
+git clone https://github.com/facebookresearch/blt
+cd blt
+conda create -n blt python=3.12
+conda activate blt
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
+pip install ninja
+pip install -v -U git+https://github.com/facebookresearch/xformers.git@de742ec3d64bd83b1184cc043e541f15d270c148
+pip install -r requirements.txt
+conda activate blt
+```
+
+### Option 2: Slurm Job to Build Env
 
 ```bash
 git clone https://github.com/facebookresearch/blt
@@ -46,6 +63,18 @@ Once that is done you can activate the environment
 conda activate blt_<date>
 ```
 
+### Options 3 (experimental, reproducible): uv
+
+Run the following to install the env using [uv](https://docs.astral.sh/uv/).
+The main benefit of this method is that the build is reproducible since there is a lock file.
+
+```bash
+uv pip install --group pre_build --no-build-isolation
+uv pip install --group compile_xformers --no-build-isolation
+uv sync
+uv run python demo.py "A BLT has"
+```
+
 ## Downloading HF Model Weights and Generating Text
 
 We have released weights on HF for the [BLT 1B Model](https://huggingface.co/facebook/blt-1b) and [BLT 7B Model](https://huggingface.co/facebook/blt-7b).

pyproject.toml

@@ -1,3 +1,56 @@
+[project]
+name = "blt"
+version = "0.1.0"
+description = "BLT"
+readme = "README.md"
+requires-python = "==3.12.*"
+dependencies = [
+    "altair>=5.5.0",
+    "datatrove>=0.5.0",
+    "fsspec>=2024.6.1",
+    "huggingface-hub==0.30.*",
+    "lm-eval>=0.4.8",
+    "luigi>=3.6.0",
+    "numpy>=2.1.2",
+    "omegaconf>=2.3.0",
+    "orjson>=3.10.18",
+    "pydantic>=2.11.4",
+    "pynvml>=12.0.0",
+    "rich>=14.0.0",
+    "s3fs>=2024.6.1",
+    "scipy>=1.15.2",
+    "sentencepiece>=0.2.0",
+    "submitit>=1.5.2",
+    "tiktoken>=0.8.0",
+    "typer>=0.15.3",
+    "viztracer>=1.0.3",
+    "wandb>=0.19.10",
+    "xformers",
+]
+
+[[tool.uv.index]]
+name = "torch-nightly-cu121"
+url = "https://download.pytorch.org/whl/nightly/cu121"
+
+
+[tool.uv.sources]
+torch = {index = "torch-nightly-cu121"}
+xformers = { git = "https://github.com/facebookresearch/xformers.git", rev = "de742ec3d64bd83b1184cc043e541f15d270c148" }
+
+[dependency-groups]
+pre_build = [
+    "setuptools",
+    "ninja",
+    "torch==2.6.0.dev20241112",
+]
+compile_xformers = ['xformers']
+
+
+[tool.uv]
+no-build-isolation-package = ["xformers"]
+index-strategy = "unsafe-best-match"
+override-dependencies = ["torch==2.6.0.dev20241112"]
+
 [tool.isort]
 profile = "black"
 known_bytelatent = "bytelatent"

requirements.txt

@@ -5,7 +5,6 @@ rouge-score
 sacrebleu
 sentencepiece
 tiktoken
-fsspec
 blobfile
 wandb
 viztracer
@@ -21,4 +20,4 @@ submitit
 typer
 rich
 fsspec[full]
-orjson
+huggingface-hub==0.30.*

setup/create_env.sh

@@ -26,15 +26,16 @@ env_prefix=blt_$current_date
 # Create the conda environment
 
 source $CONDA_ROOT/etc/profile.d/conda.sh
-conda create -n $env_prefix python=3.11 -y -c anaconda
+conda create -n $env_prefix python=3.12 -y
 conda activate $env_prefix
 
 echo "Currently in env $(which python)"
 
 # Install packages
-pip install torch==2.5.0 xformers --index-url https://download.pytorch.org/whl/cu121
+pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
 pip install ninja
-pip install --requirement requirements.txt
+pip install -v -U git+https://github.com/facebookresearch/xformers.git@de742ec3d64bd83b1184cc043e541f15d270c148
+pip install -r requirements.txt
 
 # End timer
 end_time=$(date +%s)