Running JAX

Author

Marie-Hélène Burle

This sections covers important considerations on how to use resources efficiently when training neural networks and instructions on how to run code during this course.

Resource efficient workflow

Code prototyping

Python being an interpreted language, it makes sense to prototype code in an interactive fashion. There are many options of this, including:

  • launching the Python REPL (which finally saw a little refresh with Python 3.13 but still remains very austere)
  • using the much more powerful IPython shell,
  • using the even more powerful ptpython (prompt toolkit) shell,
  • using the previous two combined (ptpython integrates with IPython thanks to its ptipython executable),
  • using Emacs as a Python IDE,
  • using JupyterLab.

Training a model requires a lot of resources. You might need multiple GPUs or entire nodes. It would be silly to have so much resource sit idle for hours while you are typing in a Jupyter notebook or thinking about your code.

The answer is to prototype code in an interactive environment at a very small scale (e.g. on a tiny subsample of data) until you have a program (a script) that works.

Scaling things up

Once you are confident that your code is good, you can scale it up on more hardware (not all at once, do multiple tests at increasingly larger scale so that you don’t wait for three weeks for results that don’t work).

For this part, it is best to SSH into a cluster and launch a batch Slurm job.

One of the great things about JAX is that the same code runs on any device so you can test the code on your machine on CPUs and then run it as is on a clusters on GPU.

Our workflow for this course

We will mostly use a JupyterHub during this course to play with snippets of code. When we get to really trying to train a model, we will log in to our training cluster via SSH and submit jobs to the Slurm scheduler.

Accessing our JupyterHub

Note that this JupyterHub is temporary and will be destroyed at the end of the course.

  • Go to the etherpad shared during the course to claim a username,
  • go to the URL of the JupyterHub for this course,
  • sign in with the username you claimed and the password we gave you,
  • the server options are good as they are unless you want to bump the time a little (e.g. to 1.5h),
  • press Start,
  • start a Python notebook: click on the button Python 3 in the Notebook section (top row of buttons).

The packages for this course are already installed in the JupyterHub.

If you don’t need all the time you asked for after all, you should log out (the resources you are using on a cluster are shared amongst many people and when resources are allocated to you, they aren’t available to other people. So it is a good thing not to ask for unnecessary resources and have them sit idle when others could be using them).

To log out, click on File in the top menu and select Log out at the very bottom.

If you would like to make a change to the information you entered on the server option page after you have pressed start, log out in the same way, log back in, edit the server options, and press start again.

Logging in through SSH

Step 1: get the info

During the course, we will give you 3 pieces of information:

  • a link to a list of usernames,
  • the hostname for our temporary training cluster,
  • the password to access that cluster.

Step 2: claim a username

Add your first name or a pseudo next to a free username on the list to claim it.

Your username is the name that was already on the list, NOT what you wrote next to it (which doesn’t matter at all and only serves at signalling that this username is now taken).

Your username will look like userxx—xx being 2 digits—with no space and no capital letter.

Step 3: run the ssh command

Linux users:   open the terminal emulator of your choice.
macOS users:   open “Terminal”.

Then type:

ssh userxx@hostname

and press Enter.

  • Replace userxx by your username (e.g. user09).
  • Replace hostname by the hostname we will give you the day of the workshop.

When asked:

Are you sure you want to continue connecting (yes/no/[fingerprint])?

Answer: “yes”.

We suggest using the free version of MobaXterm, a software that comes with a terminal emulator and a GUI interface for SSH sessions.

Here is how to install MobaXterm:

  • download the “Installer edition” to your computer (green button to the right),
  • unzip the file,
  • double-click on the .msi file to launch the installation.

Here is how to log in with MobaXterm:

  • open MobaXterm,
  • click on Session (top left corner),
  • click on SSH (top left corner),
  • fill in the Remote host * box with the cluster hostname we gave you,
  • tick the box Specify username,
  • fill in the box with the username you selected (e.g. user09),
  • press OK,
  • when asked Are you sure you want to continue connecting (yes/no/[fingerprint])?, answer: “yes”.

Here is a live demo.

Step 4: enter the password

When prompted, enter the password we gave you.

You will not see anything happen as you type the password. This is normal and it is working, so keep on typing the password.

This is called blind typing and is a Linux safety feature. It can be unsettling at first not to get any feed-back while typing as it really looks like it is not working. Type slowly and make sure not to make typos.

Then press Enter.

Am I logged in?

To know whether or not you are logged in, look at your prompt: it should look like the following (with your actual username):

[userxx@login1 ~]$

Troubleshooting

Problems logging in are almost always due to typos. If you cannot log in, retry slowly, entering your password carefully.

How do I log out?

You can log out by pressing Ctl+d.