Skip to content

Fixed use_amp bug #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yashc73080 merged 3 commits into from
Mar 4, 2026
Merged

Fixed use_amp bug #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
13f2a06
Amarel ray script
yashc73080 Feb 11, 2026
481c38a
Updated Amarel GCC version
yashc73080 Feb 11, 2026
fbd6a30
Fixed use_amp bug
yashc73080 Mar 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 1 addition & 11 deletions EIANN/rules/dendritic_loss.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,7 @@
class DendriticLossBias(BiasLearningRule):
def step(self):
# self.population.bias.data += self.learning_rate * self.population.dendritic_state
bias_update += self.learning_rate * self.population.plateau

# Convert to same dtype as bias for AMP compatibility
if self.population.network.use_amp and bias_update.dtype != self.population.bias.dtype:
bias_update = bias_update.to(self.population.bias.dtype)

self.population.bias.data += bias_update
self.population.bias.data += self.learning_rate * self.population.plateau


class DendriticLoss_6(LearningRule):
Expand All @@ -31,9 +25,5 @@ def step(self):
delta_weight = torch.outer(
torch.clamp(self.projection.post.forward_dendritic_state.detach().clone(), min=-1, max=1),
torch.clamp(self.projection.pre.forward_prev_activity, min=0, max=1))

# Ensure delta_weight has same dtype as weight for AMP compatibility
if self.projection.post.network.use_amp and delta_weight.dtype != self.projection.weight.dtype:
delta_weight = delta_weight.to(self.projection.weight.dtype)

self.projection.weight.data += self.sign * self.learning_rate * delta_weight
69 changes: 69 additions & 0 deletions EIANN/simulate/jobscripts/run_EIANN_gpu_amarel_mnist_ray.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
#!/bin/bash

# Export the current date and time for job labeling
export DATE=$(date +%Y%m%d_%H%M%S)
export LABEL="$1"
export JOB_NAME=eiann_gpu_mnist_ray_"$LABEL"_"$DATE"

# Environment variables to optimize performance
export OMP_NUM_THREADS=4
export MKL_NUM_THREADS=1
export NUMEXPR_NUM_THREADS=1
export OPENBLAS_NUM_THREADS=1
export CUDA_DEVICE_MAX_CONNECTIONS=1
export WANDB_START_METHOD=thread

# Create directories for logs and scratch data
mkdir -p /scratch/${USER}/logs/eiann
mkdir -p /scratch/${USER}/data/eiann

# Submit the job
sbatch <<EOT
#!/bin/bash
#SBATCH -J $JOB_NAME
#SBATCH -o /scratch/${USER}/logs/eiann/$JOB_NAME.%j.o
#SBATCH -e /scratch/${USER}/logs/eiann/$JOB_NAME.%j.e
#SBATCH --requeue
#SBATCH --ntasks=1
#SBATCH --nodes=1
#SBATCH --partition=gpu
#SBATCH --gres=gpu:3
#SBATCH --mem=80G
#SBATCH --cpus-per-task=15
#SBATCH --time=02:00:00
#SBATCH --mail-user=yc1376@scarletmail.rutgers.edu
#SBATCH --mail-type=ALL

module purge
module use /projects/community/modulefiles
module load gcc/10.2.0-bz186
module load cuda/11.7

set -x

cd \$HOME/EIANN/

source ~/miniconda/etc/profile.d/conda.sh
conda activate eiann5
LD_LIBRARY_PATH=\$CONDA_PREFIX/lib:\$LD_LIBRARY_PATH

Comment on lines +48 to +49
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LD_LIBRARY_PATH is assigned but not exported. Since python ... is launched as a subprocess, it will only see the updated LD_LIBRARY_PATH if it is exported. Use export LD_LIBRARY_PATH=... (or remove the line if it’s unnecessary).

Copilot uses AI. Check for mistakes.
# Start Ray (local mode)
ray start --head --num-gpus=3 --num-cpus=15

python EIANN/simulate/run_EIANN_mnist_ray2.py \
--network-config-file-name=20231129_EIANN_2_hidden_mnist_van_bp_relu_SGD_config_G_complete_optimized.yaml \
--data-dir=/scratch/\${USER}/data/eiann \
--num-seeds=5

ray stop
EOT

# Submit job:
# cd $HOME/EIANN/EIANN/simulate/jobscripts
# sbatch run_EIANN_gpu_amarel_mnist_ray.sh van_bp

# See logs:
# cd /scratch/$USER/logs/eiann

# See the progress
# watch -n 1 squeue -u $USER
4 changes: 3 additions & 1 deletion EIANN/simulate/jobscripts/run_EIANN_gpu_bridges_mnist_ray.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,4 +39,6 @@ ray stop
# See logs:
# cd /ocean/projects/bio250022p/$USER/logs/EIANN

# interact -p GPU-shared --gres=gpu:v100-32:3 --mem=80G --ntasks-per-node=15 -t 01:00:00 -A bio250022p
# interact -p GPU-shared --gres=gpu:v100-32:3 --mem=80G --ntasks-per-node=15 -t 01:00:00 -A bio250022p

# Whatever we have here works: eiann_gpu_mnist_ray.37358310.o
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This comment hard-codes a specific job output file (eiann_gpu_mnist_ray.37358310.o), which will quickly become outdated and may confuse future users. Suggest removing it or replacing it with a generic pointer to the %j-based log naming pattern used above.

Suggested change
# Whatever we have here works: eiann_gpu_mnist_ray.37358310.o
# Log files will be named like: eiann_gpu_mnist_ray.<jobid>.o / eiann_gpu_mnist_ray.<jobid>.e (see %j above)

Copilot uses AI. Check for mistakes.
6 changes: 3 additions & 3 deletions EIANN/simulate/jobscripts/run_EIANN_gpu_multi_node_bridges_mnist_ray.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
#!/bin/bash -l
#SBATCH -J eiann_gpu_mnist_ray
#SBATCH -o /ocean/projects/bio250022p/chennawa/logs/EIANN/eiann_gpu_mnist_ray.%j.o
#SBATCH -e /ocean/projects/bio250022p/chennawa/logs/EIANN/eiann_gpu_mnist_ray.%j.e
#SBATCH -J eiann_multi_gpu_mnist_ray
#SBATCH -o /ocean/projects/bio250022p/chennawa/logs/EIANN/eiann_multi_gpu_mnist_ray.%j.o
#SBATCH -e /ocean/projects/bio250022p/chennawa/logs/EIANN/eiann_multi_gpu_mnist_ray.%j.e
Comment on lines +3 to +4
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the SLURM output/error paths, the directory includes a hard-coded username (/chennawa/) but later in the script you reference logs under $USER. This mismatch makes the script non-portable and may fail due to permissions or missing directories. Consider switching these SBATCH paths to use $USER (and/or ensure the directories exist).

Suggested change
#SBATCH -o /ocean/projects/bio250022p/chennawa/logs/EIANN/eiann_multi_gpu_mnist_ray.%j.o
#SBATCH -e /ocean/projects/bio250022p/chennawa/logs/EIANN/eiann_multi_gpu_mnist_ray.%j.e
#SBATCH -o /ocean/projects/bio250022p/$USER/logs/EIANN/eiann_multi_gpu_mnist_ray.%j.o
#SBATCH -e /ocean/projects/bio250022p/$USER/logs/EIANN/eiann_multi_gpu_mnist_ray.%j.e

Copilot uses AI. Check for mistakes.
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --partition=GPU
Expand Down
2 changes: 1 addition & 1 deletion EIANN/simulate/jobscripts/simulate_EIANN_gpu_amarel_mnist.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ sbatch <<EOT

module purge
module use /projects/community/modulefiles
module load gcc/5.4
module load gcc/10.2.0-bz186
module load cuda/11.7

set -x
Expand Down
Loading