Repository: google/jax

Issue:
Doesn't compile on Nvidia Jetson Nano



Discussion:
There's no action we can take here without more information, given we don't have access to such a device. What happens when you try? Can you share the logs from the build?
> There's no action we can take here without more information, given we don't have access to such a device. What happens when you try? Can you share the logs from the build?

Actually, I got it working! Instructions in my message below. 

Main problem with the Jetsons is that they're very slow on updating their Ubuntu version so they're still on 18.04 and so Python 3.6. By the way these are great devices (RPi on steroids basically) and I think they have a great future. The linked instructions above should also work on the higher end Jetsons such as Xavier AGX, Xavier NX, and TX1/2. 
### INSTRUCTIONS FOR INSTALLING GOOGLE JAX ML ON 4GB JETSON NANO WITH SSD ###

SORTING OUT SWAP SPACE

Before you do any compiling of Jax, It’s really important to have plenty of swap space, because the Jetson Nano's 4GB is not nearly enough. Even 5GB of swap space is not enough, because gcc got killed for out of memory. I’ve gone with 10GB. For this reason you will almost certainly have to be running your Jetson off a USB3.0 SSD rather than the SD-CARD (If you want to do this, JetsonHacks has a good tutorial and helper utilities: [Jetson Nano - Boot from USB - JetsonHacks](https://www.jetsonhacks.com/2021/03/10/jetson-nano-boot-from-usb/) HOWEVER, do NOT take the final bit of advice to remove the SD-CARD. With the latest jetpack this causes SSD boot to fail. You have to keep the SD card in at the same time as the SSD is plugged into USB).

Unfortunately, the Jetson is natively configured to use 2 GB of ZRAM swap, which is almost useless (other than a mild amount of compression). You need to get rid of ZRAM:

`sudo systemctl disable nvzramconfig`

Probably best to reboot the Nano after this step, just in case (I did). Then use htop to check that you have no swap. Finally, create a 10GB swap file using these instructions (make sure you replace `1G` with `10G` in the fallocate stage).

![|32x32](https://aws1.discourse-cdn.com/nvidia/original/3X/e/1/e178dca23fe6a7631d2e99960270c84ff800013e.png) [linuxize.com – 28 Nov 18](https://linuxize.com/post/how-to-add-swap-space-on-ubuntu-18-04/)

![|690x365](https://aws1.discourse-cdn.com/nvidia/optimized/3X/1/f/1f1105be3413a3bbe857513c995217141b9dc6ba_2_690x365.jpeg)

### [How to Add Swap Space on Ubuntu 18.04](https://linuxize.com/post/how-to-add-swap-space-on-ubuntu-18-04/)


Now that you have a 10GB swap file (check it with htop), there’s one more thing I recommend: unless you have a 4-amp or greater power supply (most are 3A), move the Jetson Nano over to its low power 5 watt mode. That’s because the MAXN 10watt mode will cause compilation on all 4 cores at 100%, for many hours, and especially if you have a fan and and SSD, the power draw might be bigger than your power supply can handle. Certainly was the case for me where it crashed twice mid-compile. So use the Nvidia menu in the taskbar to turn it over to 5 watt mode. No more problems. 

CUDA

Now ensure your CUDA PATH and LD_LIBRARY_PATH are set for CUDA support. Add these to your `.bashrc` :

```
export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
```

And re-source it with `source ~/.bashrc` .

FIXING PYTHON

Python 3.69 and the related stack on the Jetson will result in a build, but it won’t work. You have to run a later version of Python, and I succeeded using Python 3.9. However note that you don’t want to replace the Jetson’s system Python 3.69 as the default python3 as this will break a lot of stuff. Instead, we’ll use a **virtualenv** . But before we do that, let’s install Python 3.9 but without making it the default Python:

```
sudo add-apt-repository ppa:deadsnakes/ppa 
sudo apt update
sudo apt install python3.9 python3.9-dev
```

This will install `python3.9` in your /usr/bin directory. However you’ll notice that when you type `python3` you still get the system python 3.69 which is what we want. To use python 3.9 you have to explicitly run `python3.9` . However we’ll need a bunch of other things for python 3.9 and for that we’ll use a virtual env. To install virtual env we must first install pip3:

```
sudo apt install python3-pip
```

Now:

```
sudo pip3 install virtualenv
```

Remember to sudo here because we need it for all users.

Next we create a python3.9 virtualenv, here I’ll create a “py39”:

```
virtualenv -p /usr/bin/python3.9 py39
```

Activate it:

```
source ./py39/bin/activate
```

This is now virgin territory with no numpy scipy or anything (even though they’re installed on the system python). So:

```
python3 -m pip install numpy scipy six wheel
```

BUILD

Now you’re ready to install `jaxlib` from source. Apt install necessary prerequisites:

```
sudo apt install g++
```

Clone the jax repository and cd into it:

```
git clone https://github.com/google/jax
cd jax
```

You’re finally ready to compile:

```
python3 build/build.py --enable_cuda
```

This step will take 12 hours. It’s a humongous codebase. Your Jetson’s swap usage will go above 5GB at times.

Once compiled:

`pip3 install dist/*.whl `

… but all this was just to get jaxlib up and running, not jax itself. There’s one last step:

```
pip3 install -e .
```

Done.

ADDENDUM:
Couple of things: 
a) The first time you do anything with Jax on the Nano, (such as using `PRNGKey` function) it'll take a _really_ long time for it to respond. You might think it's crashed. It hasn't. It's doing some weird warmup in the background but once that's done, it'll work fine. 

b) I highly recommend running LXDE and not Unity as this will free up a full gigabyte of RAM. Remember the Jetsons use system ram for the GPU as well (https://www.jetsonhacks.com/2020/11/07/save-1gb-of-memory-use-lxde-on-your-jetson/)
Awesome!

Let us know if there's something we could do to improve our documentation.

I think the things we could improve here might be: (a) to support cross-compilation so you can build jaxlib on some other beefier machine but target it for the Jetson device, and (b) look into the long delay on first computation. (a) might work right now — we've never tried. For (b), if you had a Python profile of running a script that runs a single op that might tell us something. It's not obvious to me what would take a lot of time but only on the first execution.
> It's not obvious to me what would take a lot of time but only on the first execution.

In other frameworks this is PTX jit-compilation for kernels that weren't compiled for the right GPU architecture, but JAX uses its own JIT and doesn't ship many precompiled kernels at all. Maybe NCCL (which contains the vast majority of CUDA code that JAX compiles from source) was built for the wrong set of compute capabilities?
> Awesome!
> 
> Let us know if there's something we could do to improve our documentation.
> 
> I think the things we could improve here might be: (a) to support cross-compilation so you can build jaxlib on some other beefier machine but target it for the Jetson device, and (b) look into the long delay on first computation. (a) might work right now — we've never tried. For (b), if you had a Python profile of running a script that runs a single op that might tell us something. It's not obvious to me what would take a lot of time but only on the first execution.

Yeah I thought the documentation was pretty good actually and the bazel-based build worked really smoothly. Certainly not always the case when building other libraries on Nvidia's Jetson family because they are "quirky" in some respects. Yes not sure why the warmup takes so long and I haven't played around enough with it in the last few days but as soon as I get more insight I'll post it here, with a profile dump. I can confirm though that Jax is working  well as I'm getting a 5x speedup on big dot products versus the native, NEON-optimized numpy. 5x is great if not huge, but it certainly proves it's working and means people can play with CUDA-enabled JAX for 99 bucks (though of course, colab is even cheaper!). 

Resolving patch:
diff --git a/build/build.py b/build/build.py
--- a/build/build.py
+++ b/build/build.py
@@ -417,6 +417,12 @@ def main():
       parser,
       "enable_rocm",
       help_str="Should we build with ROCm enabled?")
+  add_boolean_argument(
+      parser,
+      "enable_nccl",
+      default=True,
+      help_str="Should we build with NCCL enabled? Has non effect for non-CUDA "
+               "builds.")
   parser.add_argument(
       "--cuda_path",
       default=None,
@@ -505,6 +511,7 @@ def main():
       print("CUDA version: {}".format(args.cuda_version))
     if args.cudnn_version:
       print("CUDNN version: {}".format(args.cudnn_version))
+    print("NCCL enabled: {}".format("yes" if args.enable_nccl else "no"))
 
   print("TPU enabled: {}".format("yes" if args.enable_tpu else "no"))
 
@@ -545,6 +552,8 @@ def main():
   if args.enable_cuda:
     config_args += ["--config=cuda"]
     config_args += ["--define=xla_python_enable_gpu=true"]
+    if not args.enable_nccl:
+      config_args += ["--config=nonccl"]
   if args.enable_tpu:
     config_args += ["--define=with_tpu_support=true"]
   if args.enable_rocm:


Test patch: