mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-04-07 20:30:54 +00:00
add readme for tutorials
This commit is contained in:
parent
a030ec47c6
commit
f0c43f8d76
1 changed files with 38 additions and 0 deletions
38
notebooks/README.md
Normal file
38
notebooks/README.md
Normal file
|
@ -0,0 +1,38 @@
|
|||
# Particle Mesh Simulation with JAXPM on Multi-GPU and Multi-Host Systems
|
||||
|
||||
This collection of notebooks demonstrates how to perform Particle Mesh (PM) simulations using **JAXPM**, leveraging JAX for efficient computation on multi-GPU and multi-host systems. Each notebook progressively covers different setups, from single-GPU simulations to advanced, distributed, multi-host simulations across multiple nodes.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. **[Single-GPU Particle Mesh Simulation](01-Introduction.ipynb)**
|
||||
- Introduction to basic PM simulations on a single GPU.
|
||||
- Uses JAXPM to run simulations with absolute particle positions and Cloud-in-Cell (CIC) painting.
|
||||
|
||||
2. **[Advanced Particle Mesh Simulation on a Single GPU](02-Advanced_usage.ipynb)**
|
||||
- Explores Lagrangian Perturbation Theory (LPT) simulations.
|
||||
- Introduces weighted field projections, boundary handling, and adjustments to halo sizes.
|
||||
|
||||
3. **[Multi-GPU Particle Mesh Simulation with Halo Exchange](03-MultiGPU_PM_Halo.ipynb)**
|
||||
- Extends PM simulation to multi-GPU setups with halo exchange.
|
||||
- Uses sharding and device mesh configurations to manage distributed data across GPUs.
|
||||
|
||||
4. **[Multi-GPU Particle Mesh Simulation with Advanced Solvers](04-MultiGPU_PM_Solvers.ipynb)**
|
||||
- Compares different ODE solvers (Leapfrog and Dopri5) in multi-GPU simulations.
|
||||
- Highlights performance, memory considerations, and solver impact on simulation quality.
|
||||
|
||||
5. **[Multi-Host Particle Mesh Simulation](05-MultiHost_PM.ipynb)**
|
||||
- Extends PM simulations to multi-host, multi-GPU setups for large-scale simulations.
|
||||
- Guides through job submission, device initialization, and retrieving results across nodes.
|
||||
|
||||
## Getting Started
|
||||
|
||||
Each notebook includes installation instructions and guidelines for configuring JAXPM and required dependencies. Follow the setup instructions in each notebook to ensure an optimal environment.
|
||||
|
||||
## Requirements
|
||||
|
||||
- **JAXPM** (included in the installation commands within notebooks)
|
||||
- **Diffrax** for ODE solvers
|
||||
- **JAX** with CUDA support for multi-GPU or TPU setups
|
||||
- **SLURM** for job scheduling on clusters (if running multi-host setups)
|
||||
|
||||
> **Note**: These notebooks are tested on the **Jean Zay** supercomputer and may require configuration changes for different HPC clusters.
|
Loading…
Add table
Reference in a new issue