Operator Learning

---

This repository contains implementations, experimental scripts, and documentation related to operator learning—a machine learning paradigm focused on learning mappings between infinite-dimensional function spaces, such as solving parametric partial differential equations (PDEs) using neural operators like the Fourier Neural Operator (FNO).

---

Installation

✅ Base Environment

• Python
• OpenMPI
• CUDA
• PyTorch
• Python dependencies listed in requirements.txt

🔧 Setup Instructions

git clone https://github.com/chelseajohn/operator_learning.git
cd operator_learning
python3 -m venv venv
source venv/bin/activate
python -m pip install -e .

---

🚀 Distributed Data Parallel (DDP) Training

Model training can be accelerated using Distributed Data Parallel (DDP), implemented with PyTorch DDP. The model is replicated across GPUs, trained in parallel on different data samples, and synchronized at each step.

🔧 Enable DDP in Configuration

Edit your config.yaml to include:

parallel_strategy:
  ddp: True              # Enable Distributed Data Parallel
  gpus_per_node: 4       # Number of GPUs per node

Particle Parallel Training For PIC (TP)

For PIC problem, the data can be parallelised along particle dimension (tensor parallel/ input sharding) in addition to time step dimension which is parallelised with DDP as (timestep/dp_ranks, channels, nparticle/tp_ranks)

🔧 Enable DDP + TP in Configuration

Edit your config.yaml to include:

parallel_strategy:
  ddp: true     # Distributed Data Parallel
  tp : true     # Tensor Parallel(Only for PIC-DSE)
  tp_size: 2    # must be divisible by nParticle          
  gpus_per_node: 4

🖥️ Launching Distributed Jobs via SLURM

Distributed jobs can be launched using torchrun within a SLURM environment:

##### Network parameters #####
MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
MASTER_PORT=6000
GPUS_PER_NODE=4

###### JSC System specifics #####
export SRUN_CPUS_PER_TASK=${SLURM_CPUS_PER_TASK}
export CUDA_VISIBLE_DEVICES=0,1,2,3
export OMP_NUM_THREADS=1
export NCCL_SOCKET_IFNAME=ib0
MASTER_ADDR="${MASTER_ADDR}i"

srun python -u -m torch.distributed.run   \
    --nproc_per_node $GPUS_PER_NODE   \
    --nnodes $SLURM_JOB_NUM_NODES  \
    --node_rank $SLURM_PROCID   \
    --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT \
    --rdzv_backend c10d    \
    --rdzv_conf=is_host=$(if ((SLURM_NODEID)); then echo False; else echo True; fi)  \
    --max_restarts 0     \
    --tee 3 \
    $python_file  

🔍 Profiling

To enable profiling using Nsight systems or Torch profiler add a profile section into the config:

profile:
  enableProfiler: True
  profiler: nsys               # or torch
  profileDir: profileLog
  evalOnly: false             # profile only inference using torch profiler 
  input_shape: [1,3,10000]    # relevant only if eval_only=true
  use_amp: 1                  # enable mixed precision
  benchmark: true             # run profiling in benchmark mode
  compile_eval: 1             # use torch.compile on model inference
  compile_train: 0            # use torch.compile on training (set to 1, if compile_eval=0)
  comipile_mode: default      # compiler mode 

Problems Solved

• Rayleigh-Bénard convection in 2D/3D using FNO, with datasets generated via the pseudo-spectral solver Dedalus for 2D and pySDC for 3D.
• Particle-In-Cell (1D/2D/3D) algorithm using FNO for electric field prediction