|
| 1 | +--- |
| 2 | +authors: |
| 3 | + - wolfv |
| 4 | +tags: [conda-forge] |
| 5 | +image: https://i.imgur.com/pRdJaYw.png |
| 6 | +--- |
| 7 | + |
| 8 | +# GPU enabled TensorFlow builds on conda-forge |
| 9 | + |
| 10 | + |
| 11 | + |
| 12 | +Recently we've been able to add GPU-enabled TensorFlow builds to |
| 13 | +conda-forge! This was quite a journey, with multiple contributors trying |
| 14 | +different ways to convince the Bazel-based build system of TensorFlow to |
| 15 | +build CUDA-enabled packages. But we managed, [and the pull request got |
| 16 | +merged](https://github.com/conda-forge/tensorflow-feedstock/pull/157). |
| 17 | + |
| 18 | +<!--truncate--> |
| 19 | + |
| 20 | +We now have a configuration in place that creates CUDA-enabled |
| 21 | +TensorFlow builds for all conda-forge supported configurations (CUDA |
| 22 | +10.2, 11.0, 11.1, and 11.2+). Building out the CUDA packages requires |
| 23 | +beefy machines -- on a 32 core machine it still takes around 3 hours to |
| 24 | +build a single package. Our build matrix now includes 12 CUDA-enabled |
| 25 | +packages & 3 CPU packages (because we need separate packages per Python |
| 26 | +version). As one can imagine, this isn't easily possible on an average |
| 27 | +"home computer". |
| 28 | + |
| 29 | +For this purpose, we have written an Ansible playbook that lets us boot |
| 30 | +up cloud machines which then build the feedstock (using the |
| 31 | +`build-locally.py` script). Thanks to the generous support |
| 32 | +of OVH we were able to boot multiple 32-core virtual machines |
| 33 | +simultaneously to build the different TensorFlow variants. |
| 34 | + |
| 35 | +We have open-sourced the [Ansible playbook in |
| 36 | +GitHub](https://github.com/mamba-org/build-locally-ansible) and we're |
| 37 | +working towards making it (more) generally useful for other long-running |
| 38 | +builds! |
| 39 | + |
| 40 | +{width="600px"} |
| 41 | + |
| 42 | +With the TensorFlow builds in place, conda-forge now has CUDA-enabled |
| 43 | +builds for PyTorch and Tensorflow, the two most popular deep learning |
| 44 | +libraries. |
| 45 | + |
| 46 | +We are still missing Windows builds for TensorFlow (CPU & CUDA, |
| 47 | +unfortunately) and would love the community to help us out with that. |
| 48 | +There is an open PR, but it probably needs some poking in Bazel to get |
| 49 | +it to pass: |
| 50 | +<https://github.com/conda-forge/tensorflow-feedstock/pull/111>. |
| 51 | + |
| 52 | +We hope that these new GPU builds will enable many more packages to be |
| 53 | +added to the conda-forge channel! We are already looking forward to the |
| 54 | +2.6.2 and 2.7 releases of TensorFlow and to adding Windows support in |
| 55 | +the future. We hope you enjoy this work. |
| 56 | + |
| 57 | +## Installation |
| 58 | + |
| 59 | +You can now select between GPU enabled (default) and CPU packages using |
| 60 | +the `tensorflow-gpu` and `tensorflow-cpu` packages. Just run |
| 61 | + |
| 62 | +```bash |
| 63 | +mamba install tensorflow-gpu -c conda-forge |
| 64 | +# OR |
| 65 | +conda install tensorflow-gpu -c conda-forge |
| 66 | +``` |
| 67 | + |
| 68 | +When installing the `tensorflow` package, the package resolution will |
| 69 | +now default to the GPU-enabled builds of tensorflow if the local machine |
| 70 | +has a GPU (these builds can be identified by "cuda" at the beginning |
| 71 | +of the version number). Note that GPU-enabled packages can also work on |
| 72 | +CPU-only machines, but one would need to override the enviornment |
| 73 | +variable `CONDA_OVERRIDE_CUDA` like below. This could be handy if you |
| 74 | +are in a situation where your current node (e.g. login node) on an HPC |
| 75 | +does not have GPUs, but the compute nodes with GPUs do not have internet |
| 76 | +access. |
| 77 | + |
| 78 | +```bash |
| 79 | +CONDA_OVERRIDE_CUDA="11.2" conda install tensorflow cudatoolkit>=11.2 -c conda-forge |
| 80 | +# OR |
| 81 | +CONDA_OVERRIDE_CUDA="11.2" mamba install tensorflow cudatoolkit>=11.2 -c conda-forge |
| 82 | +``` |
| 83 | + |
| 84 | +Note that you should select the cudatoolkit version most appropraite for |
| 85 | +your GPU; currently, we have "10.2", "11.0", "11.1", and "11.2" |
| 86 | +builds available where the the "11.2" builds are compatible with all |
| 87 | +cudatoolkits>=11.2. You could also force a specific version of |
| 88 | +`cudatoolkit` by specifying it like above. Moreover, you could ensure |
| 89 | +you get a sepcific build of tensorflow by appending the package name |
| 90 | +like `tensorflow==2.7.0=cuda*` or `tensorflow==2.7.0=cuda112*`. If you |
| 91 | +want the slimmer "cpu-only" package, then you can install |
| 92 | +`tensorflow-cpu` directly or equivalently `tensorflow==2.7.0=cpu*`. At |
| 93 | +the time of writing (February 2022), on a machine without a GPU, one |
| 94 | +would always get the `-cpu` variant unless overriden like |
| 95 | +above. This decision has been made to allow greater accessibility for |
| 96 | +users with limited bandwidth and resources. |
| 97 | + |
| 98 | +## Thanks to |
| 99 | + |
| 100 | +- Mark Harfouche (@hmaarrfk) & Ista Zahn (@izahn) for their initial |
| 101 | + work on the TensorFlow GPU builds, and all other TensorFlow |
| 102 | + maintainers. Uwe Korn (@xhochy) for his work on the Bazel scripts & |
| 103 | + TensorFlow -- and all the other maintainers of the [TensorFlow |
| 104 | + feedstock](https://github.com/conda-forge/tensorflow-feedstock)! |
| 105 | +- NVIDIA for pushing cudatoolkit and cudnn on conda-forge that makes |
| 106 | + this possible |
| 107 | +- OVH for their generous sponsoring of large build machines that we |
| 108 | + could use to build the recipes |
| 109 | +- Bloomberg for their sponsorship of QuantStack's involvement with |
| 110 | + conda-forge |
| 111 | +- Andreas Trawoger (@atrawog) for the Ansible scripts that this is |
| 112 | + based on |
| 113 | +- Thorsten Beier (@derthorsten) and Adrien Delsalle (@adriendelsalle) |
| 114 | + for their contributions to the recipe |
0 commit comments