JAX

JAX on Autograd ja XLA yhdistettynä suorituskykyiseen koneoppimistutkimukseen.

Saatavuus

Tällä hetkellä tuetut JAX-versiot:

Version	Module	Puhti	Mahti	LUMI	Notes
0.5.0	`jax/0.5.0`	default	default	-	all packages
	`jax/0.5.0-small`	X	X	-	framework only
0.4.38	`jax/0.4.38`	X	X	default*	all packages
	`jax/0.4.38-small`	X	X	X*	framework only
0.4.30	`jax/0.4.30`	X	X	X*	all packages
	`jax/0.4.30-small`	X	X	X*	framework only
0.4.23	`jax/0.4.23-py3.9`	X	X	X*
0.4.20	`jax/0.4.20`	X	X	X*
0.4.18	`jax/0.4.18`	-	-	X*
0.4.14	`jax/0.4.14`	X	X	-
0.4.13	`jax/0.4.13`	X	X	-
0.4.1	`jax/0.4.1`	X	X	-
0.3.13	`jax/0.3.13`	X	X	-

Moduulit sisältävät JAXin Pythonille GPU-tuella CUDA:n/ROCm:n kautta sekä laajan joukon muita Python-paketteja, joita käytetään yleisesti yhdessä JAXin kanssa. Python-versio vaihtelee moduulijulkaisujen välillä. Moduulit, joissa JAX-versio on enintään 0.4.23, käyttävät Python 3.9:ää, myöhemmät moduulit käyttävät Python 3.12:ta.

LUMIssa "*" merkityt versiot ovat edelleen kokeellisia ja niiden tuki on rajallinen. Osa alla kuvatuista ominaisuuksista ei välttämättä toimi niissä. Huomaa, että JAX on saatavilla myös LUMI Software Libraryssa, jota ylläpitää LUMI User Support Team (ei CSC:n tutkimustuki).

Versiosta 0.4.30 alkaen JAX-moduulista on kaksi versiota:

"small"-versio sisältää JAXin CUDA 12.2 GPU-tuella sekä PyTorchin ja Tensorflown vain CPU:ta käyttävät versiot, jotta niiden datanlataustyökaluja voidaan käyttää. Näiden nimeämiskäytäntö on jax/<version>-small, tai voit käyttää jax/small ladataksesi oletusversion (uusimman).
"full"-versio sisältää monia yleisesti käytettyjä JAXiin perustuvia koneoppimispaketteja -- voit tarkistaa pip list -komennolla täydellisen luettelon kaikista mukana olevista paketeista. Näiden nimeämiskäytäntö on jax/<version>, tai voit yksinkertaisesti käyttää jax ladataksesi oletusversion (uusimman).

Note

Koska JAX julkaisee uusia versioita hieman epäsäännöllisellä aikataululla, emme tuo kaikkia uusia versioita heti saataville. Sen sijaan pyrimme päivittämään järjestelmissämme saatavilla olevan JAX-version noin kuuden kuukauden välein, tavoitteenamme helmikuu ja elokuu, parhaamme mukaan.

Kaikki moduulit perustuvat Apptaineria käyttäviin kontteihin (aiemmin tunnettu nimellä Singularity). tykky:n kautta on tarjottu wrapper-skriptit, jotta yleiset komennot kuten python, python3, pip ja pip3 sekä asennettujen pakettien tarjoamat komennot toimisivat normaalisti. Muiden komentojen kohdalla ne voi olla tarpeen suorittaa etuliitteellä apptainer_wrapper exec. Lisätietoja on CSC:n yleisissä ohjeissa Apptainer-konttien suorittamisesta.

Lisäpaketit

Jos huomaat, että jokin paketti puuttuu, voit usein asentaa sen itse komennolla pip install. On suositeltavaa käyttää Pythonin virtuaaliympäristöjä. Katso Python-dokumentaatiostamme lisätietoja pakettien asentamisesta itse. Jos mielestäsi jokin tärkeä paketti pitäisi sisällyttää CSC:n tarjoamaan moduuliin, ota yhteyttä asiakastukeemme.

Lisenssi

JAX on lisensoitu Apache License 2.0:lla.

Käyttö

Jos haluat käyttää oletusversiota (uusinta) Puhdissa tai Mahdissa, ota se käyttöön komennolla:

module load jax

tai

module load jax/small

Jos haluat käyttää CSC:n asentamaa JAXia LUMIssa:

module use /appl/local/csc/modulefiles/
module load jax  # jax/small

Note

Suosittelemme kiinnittämään version skripteissäsi, jotta tulevien päivitysten muutokset eivät riko skriptejä, esimerkiksi: module load jax/0.5.0 tai module load jax/0.5.0-small

Huomaa, että JAX-moduulit sisältävät jo vastaavat CUDA- ja cuDNN- tai ROCm-kirjastot, joten mitään cuda-, cudnn- tai rocm-moduuleja ei tarvitse ladata erikseen!

Tämä näyttää kaikki saatavilla olevat JAX-versiot:

module avail jax

Note

Huomaa, että kirjautumissolmut eivät ole tarkoitettu raskaaseen laskentaan, käytä sen sijaan slurm-eräajoja. Katso ohjeemme eräajojärjestelmän käyttöön.

Note

Älä lue valtavaa määrää tiedostoja jaetusta tiedostojärjestelmästä, vaan käytä nopeaa paikallista levyä tai paketoi datasi suurempiin tiedostoihin! Katso datan tallennusta käsittelevä osio koneoppimisoppaassamme saadaksesi lisätietoja.