Repo for "The Grand Illusion: The Myth of Software Portability and Implications for ML Progress."
For most environments we use poetry
The base install command is
poetry install
However we have different commands depending on your environment
to install use:
poetry run install --with m1
to install use:
poetry run install --with intel
to install use:
poetry run install --with vm
For jax gpus we use vms running cudnn 8.2 and use the following wheel poetry run pip install https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.7+cuda11.cudnn82-cp38-cp38-manylinux2014_x86_64.whl
poetry run pip install "jax[tpu]>=0.3.13" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
poetry run python main.py
The tests require an environment variable specifying the correct device. In most environments you can run the tests with the following command
DEVICE=$device poetry run pytest
where $device would be the different devices cpu, gpu, tpu
So on a local machine the command would be:
DEVICE=cpu poetry run pytest
And on a gpu VM it would be:
DEVICE=gpu poetry run pytest
However due to the quirks of the TPU vms we are not using poetry. So there the command would be:
DEVICE=tpu pytest
If in any environment you want to see print commands you add you can append -s on your command like so:
DEVICE=cpu poetry run pytest -s