Hyppää sisältöön

Significant changes to Puhti & Mahti authentication coming in April! Read about the SSH key and multi-factor authentication requirements.

tags: - Free

JAX

JAX on Autogradin ja XLAn yhdistelmä, joka on suunniteltu korkean suorituskyvyn koneoppimistutkimukseen.

Available

Tällä hetkellä tuetut JAX-versiot:

Versio Moduuli Puhti Mahti LUMI Huomautukset
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/ROCm:n kautta sekä suuren määrän muita Python-paketteja, joita käytetään yleisesti yhdessä JAXin kanssa. Python-versio vaihtelee moduulijulkaisujen välillä. Moduulit, joiden JAX-versio on enintään 0.4.23, käyttävät Python 3.9:ää, ja myöhemmät moduulit käyttävät Python 3.12:ta.

Versiot LUMIssa, merkitty "*":lla, ovat edelleen kokeellisia ja niille on rajoitettu tuki. Jotkin alla kuvatuista ominaisuuksista eivät ehkä toimi niiden kanssa. Huomaa, että JAX on saatavilla myös LUMIn ohjelmistokirjastossa LUMIn käyttäjätukitiimin ylläpitämänä (ei CSC:n tutkimustuki).

JAX-moduuli tulee kahdessa muodossa alkaen versiosta 0.4.30:

  • "Pieni" maku, joka sisältää JAXin CUDA 12.2 GPU-tuen kanssa sekä cpulle ja Tensorflowlle, jotta niiden tiedonlataustoimintoja voi käyttää. Näissä käytetään nimeämismuotoa jax/<version>-small, tai voit käyttää jax/small ladataksesi oletusversion (viimeisimmän).
  • "Täysi" maku, joka sisältää monia yleisesti käytettyjä koneoppimispaketteja, jotka perustuvat JAXiin -- voit käyttää pip list selvittääksesi, mitä paketteja sisältyy. Näissä käytetään nimeämismuotoa jax/<version>, tai voit yksinkertaisesti käyttää jax ladataksesi oletusversion (viimeisimmän).

Note

Koska JAX julkaisee uusia versioita jokseenkin epäsäännöllisellä aikataululla, emme tee kaikkia uusia versioita heti saataville. Sen sijaan pyrimme päivittämään järjestelmiemme saatavilla olevan JAX-version noin kuuden kuukauden välein, tavoitteena helmikuu ja elokuu, parhaan kykymme mukaan.

Kaikki moduulit perustuvat Apptainerin (aiemmin tunnettu nimellä Singularity) käyttämiin kontteihin. Kuutta tykky on tarjottu niin, että yleiset komennot kuten python, python3, pip ja pip3 ja asennettujen pakettien tarjoamat komennot toimivat normaalisti. Muiden komentojen osalta saatat joutua lisäämään etuliitteeksi apptainer_wrapper exec. Lisätietoja löydät CSC:n yleisistä ohjeista Apptainer konttien käyttöön.

Additional packages

Jos huomaat, että jokin paketti puuttuu, voit usein asentaa sen itse käyttäen pip install. On suositeltavaa käyttää Python-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ä palvelupisteeseemme.

License

JAX on lisensoitu Apache License 2.0.

Usage

Käyttääksesi oletusversiota (uusin) Puhtilla tai Mahtilla, alusta se komennolla:

module load jax

tai 

module load jax/small

CSC:n asentaman JAXin käyttöön LUMIssa:

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

Note

Suosittelemme, että skripteissäsi kiinnität version, jotta tulevien päivitysten muutokset eivät riko skriptejä, esim.: 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 ei ole tarpeen ladata erikseen cuda, cudnn tai rocm -moduuleita!

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

module avail jax

Note

Huomaa, että sisäänkirjautumissolmut eivät ole raskaaseen laskentaan, käytä sen sijaan slurm-eräajoja. Katso [ohjeemme eräajot] (../computing/running/getting-started.md).

Note

Älä ** lue suurta määrää tiedostoja jakaen tiedostojärjestelmästä**, käytä nopeaa paikallista levyä tai paketoi tietosi suurempiin tiedostoihin sen sijaan! Katso koneoppimisen oppaamme tietojen tallennusosa saadaksesi lisätietoja.

More information