Parallelization on CPU only

[npiroozan]

Hello,

I hope you are doing well.

If possible, I have a question regarding the way the model works with LAMMPS on CPU only.

When I run the trained model (for Silicon) with LAMMPS, I do not see any usage of neighbor lists or OMP parallelization in the simulation. It seems like it is only MPI parallelized. I think I'm doing something wrong, but would you be so kind as to review the Training Script and the LAMMPS Script? Thank you very much.

Training Script:

{
"model": {
"type_map": ["Si"],
"data_stat_nbatch": 10,

    "descriptor" :{
        "type":             "se_e2_a",
        "sel":              [46],
        "rcut_smth":        0.50,
        "rcut":             6.0,
        "resnet_dt":        false,
        "precision":        "bfloat16"
    },

    "fitting_net" : {
        "type":             "ener",
        "neuron":           [240, 240, 240],
        "activation_function": "tanh",
        "precision":        "bfloat16",
        "resnet_dt":        true,
        "seed":             1
    }
},

"learning_rate" :{
    "scale_by_worker": "none",
    "type":         "exp",
    "decay_steps":  472,
    "start_lr":     0.024905919858219574,
    "stop_lr":      5.484673842346183e-9
},

"loss" :{
    "type":         "ener",
    "start_pref_e": 0.02,
    "limit_pref_e": 1,
    "start_pref_f": 1000,
    "limit_pref_f": 1,
    "start_pref_v": 0,
    "limit_pref_v": 0
},

"training" : {
    "training_data": {
        "systems":          ["/localdisk/npirooza/Input_Data/Silicon_Original/data_0/", "/localdisk/npirooza/Input_Data/Silicon_Original/data_1/", "/localdisk/npirooza/Input_Data/Silicon_Original/data_2/", "/localdisk/npirooza/Input_Data/Silicon_Original/data_3/", "/localdisk/npirooza/Input_Data/Silicon_Original/data_4", "/localdisk/npirooza/Input_Data/Silicon_Original/data_5", "/localdisk/npirooza/Input_Data/Silicon_Original/data_6", "/localdisk/npirooza/Input_Data/Silicon_Original/data_7", "/localdisk/npirooza/Input_Data/Silicon_Original/data_8"],
        "batch_size":       32
    },

    "validation_data":{
        "systems":          ["/localdisk/npirooza/Input_Data/Silicon_Original/data_9"],
        "batch_size":       32
    },

    "mixed_precision": {
            "output_prec": "float32",
            "compute_prec": "bfloat16"
    },

    "numb_steps":   32000,
    "seed":         1,
    "disp_file":    "lcurve.out",
    "disp_freq":    100,
    "save_freq":    1000000,
    "numb_test":    10,
    "save_ckpt":    "model.ckpt",
    "disp_training": true,
    "time_training": true,
    "tensorboard":  false,
    "tensorboard_log_dir": "log",
    "tensorboard_freq": 100,
    "enable_profiler": false,
    "profiling": false,
    "profiling_file": "timeline.json"
}

}

LAMMPS Input Script:

Initialize

units metal
dimension 3
boundary p p p
atom_style atomic
#plugin load libdeepmd_lmp.so
package intel 0 omp 1 mode double

System

lattice diamond 5.4307 orient x 1 0 0 orient y 0 1 0 orient z 0 0 1
region sim block 0 7 0 8 0 9
create_box 1 sim
create_atoms 1 region sim
mass 1 28.0855

Force Field Selection

pair_style deepmd graph-compress-high.pb
pair_coeff * *
timestep 0.002

Neighbor List

neighbor 2.5 bin
neigh_modify delay 10 check yes

Define Outputs

thermo 100
thermo_style custom step temp pe ke etotal
log log.energy

Dynamics

Minimize Energy

min_style cg
minimize 1.0e-8 1.0e-8 5000 10000
reset_timestep 0

Thermalization

velocity all create 300.0 65897 mom yes rot yes

fix 1 all langevin 300.0 300.0 1 32456
fix 2 all nve
run 25000
unfix 1
unfix 2

fix solid all nvt temp 375.0 375.0 2.5 drag 2.0
run 50000
unfix solid

Output LAMMPS:

MPI task timing breakdown:
Section | min time | avg time | max time |%varavg| %total

Pair | 29.41 | 39.333 | 51.029 | 63.7 | 68.44
Neigh | 0 | 0 | 0 | 0.0 | 0.00
Comm | 4.116 | 15.128 | 25.341 | 103.7 | 26.32
Output | 0.018891 | 0.019464 | 0.025107 | 0.3 | 0.03
Modify | 1.1884 | 2.4839 | 4.4482 | 38.7 | 4.32
Other | | 0.503 | | | 0.88

Thank you very much!

[njzjz]

I do not see any usage of neighbor lists or OMP parallelization in the simulation. It seems like it is only MPI parallelized.

Could you explain how you saw that?

[njzjz]

According to the TensorFlow blog, oneDNN needs to be enabled by TF_ENABLE_ONEDNN_OPTS=1. Is this helpful?

[npiroozan]

Thank you very much. Will reinstall Tensorflow C++ api with mkl to ensure we are no longer using Eigen.

This, along with your additional suggestion I believe will resolve the issue of OMP parallelism.

Is that correct?

[npiroozan]

OpenMP parallelism is working and fixed the lammps flag issue.

There seems to be a large increase in cache memory stalls when multiple threads/core are used.

This is what resulted in a decrease in throughput. Have you seen this issue with CPUs?

[njzjz]

LAMMPS uses a clever distributed-memory MPI parallelism. Each processor stores information for the subset of atoms within its subdomain. So, there is only a little communication between processors. See the LAMMPS paper for details.
OpenMP parallelism is within the TensorFlow OPs. In principle, OpenMP may need more communication to exchange variables.

[npiroozan]

Thank you very much. Definitely there seems to be areas of improvement for performance on CPUs. Perhaps we can discuss further offline?