ML-IAP interface Allegro+lammps bug about "fix langevin" for controlling temperature locally

[liuyenfu]

After trying the Allegro_0.7.1+lammps with ML-IAP interface, using the "fix langevin" to set different temperatures for different regions of a system for calculate the thermal conductivity. The local temperatures cannot be established and with the wrong f_fixID which is used to compute the accumulated input heat energy with time. The lammps can run but with physical non-sense reseut. Saying that there is strange bug becuase i have compared with previous non-ML-IAP versions (the AOTI version) Allegro+lammps with the same lammps script but just different pair_style and the previouse version functions correctly. The problem may be about the ML model or the lammps ML_IAP interface. Also found other ML-IAP interfaced ML+lammps (MACE_0.3.13+ML-IAP interfaced lammps) have the same problem. The problem may more likely be the lammps ML-IAP interface but still cannot conclusively comfirm which part of the process the problem comes form.

[cw-tan]

Hi @liuyenfu ,

Thanks for the report. Could you please provide more exact details.

1. For the LAMMPS ML-IAP run, please send your LAMMPS input file and possible job submission script (just want to see how you call the lmp executable).
2. For the pair_nequip_allegro run, please send your LAMMPS input file and possible job submission script (just want to see how you call the lmp executable). Also, just confirming that you're using pair_allegro?
3. How many MPI ranks / GPUs / nodes are you using? (may or may not be an important detail, but useful to know)
4. Could you please provide more details on the problem "The local temperatures cannot be established and with the wrong f_fixID which is used to compute the accumulated input heat energy with time."

Thanks.

[liuyenfu]

hi~ i give a describing for the situation first. this is a lammps script to compute the thermal conductivity with langevin method.

with all the same dataset, hyperparameters etc. models were trained saperately by Allegro_0.7.1 and Allegro_0.7.0,
then just using different
pair_style
pair_coeff
in the lammps scrpt for the ML-IAP and AOTI version.
Then the lammps run ressults are totally different while the AOTI version is physically correct, the ML-IAP version is physically non-sense by

1. the c_Thot c_cold cannot reach the goal temp 325K and 275K (while for the AOTI version, it's very easy to reach the goal temp)
2. the accumulated heat computed, f_heat f_cool, is non-sense for the ML-IAP version

when using the ML_IAP:

pair_style   mliap unified /home/liuyenfu/Allegro-0.7.1/model.nequip.lmp.pt 0
pair_coeff    * * C

with the running command:

mpirun -np 2 "/home/liuyenfu/lammps-Allegro_0.7.1_mliap/lammps/build-mliap/lmp"  -in in.langevin-Allegro-0.7.1-mliap  -k on g 2 -sf kk -pk kokkos newton on neigh half

on 2 H200

when using the AOTI:

pair_style      allegro
pair_coeff      * *   /home/liuyenfu/allegro-0.7.0/model.nequip.pt2 C

with the running command:

"/home/liuyenfu/lammps-Allegro_0.7.0/lammps/build-AOTI-H200/lmp" -in in.langevin-Allegro-0.7.0

on 1 H200

======= example of the test script ========

boundary f f f

units metal
atom_style atomic
atom_modify map yes
newton on

read_data   ( some material structure file )

pair_style   mliap unified /home/liuyenfu/Allegro-0.7.1/model.nequip.lmp.pt 0
pair_coeff    * * C


or ( pair_style      allegro
       pair_coeff      * *   /home/liuyenfu/allegro-0.7.0/model.nequip.pt2 C  )



region hot block INF INF INF INF ${hot_region_from} ${hot_region_end}
region cold block INF INF INF INF ${cold_region_from} ${cold_region_end}
compute Thot all temp/region hot
compute Tcold all temp/region cold

velocity all create 300 87287

print "------------------stage1------------------"

fix 1 all nvt temp 300 300 0.1
thermo 100
run 500
velocity all scale 300
unfix 1

langevin equilibration run
print "------------------------------stage2--------------------------------"

fix 1 all nve
fix heat all langevin 325 325 0.05 59804 tally yes
fix cool all langevin 275 275 0.05 287859 tally yes
fix_modify heat temp Thot
fix_modify cool temp Tcold

thermo_style custom step temp c_Thot c_cold f_heat f_cool

thermo 100
run 15000
=================================END===================================

[cw-tan]

Could you please try controlling for number of GPUs used just to isolate that as a potential route to triggering the error?

[liuyenfu]

hi~ i also try the ML-IAP version with only one GPU by command:

mpirun -np 1 "/home/liuyenfu/lammps-Allegro_0.7.1_mliap/lammps/build-mliap/lmp" -in in.langevin-Allegro-0.7.1-mliap -k on g 1 -sf kk -pk kokkos newton on neigh half

but the rusults are the same as the two GPUs case just the speed is slower.

[cw-tan]

One other thing you could try is to use pair_nequip on a single GPU (see https://github.com/mir-group/pair_nequip_allegro/blob/d40561852ac7236826e7cb6d654896fb132e349c/tests/test_python_repro_nequip.py#L48). The implementation is slightly different from pair_allegro, but I suspect the outcome will be similar to pair_allegro for this problem. But that'll be another data point to say that there's something wrong with our ML-IAP interface.

[liuyenfu]

hi, I make a more conclusive test for the ML-IAP interface that could tell there is really strange bugs for the ML-IAP interface.

I train a Allegro 0.7.1 model and then compile the last.ckpt with two ways:

one by
1)

nequip-compile \
    path/to/last.ckpt \
    path/to/compiled_model.nequip.pth \
    --device [cpu|cuda] \
    --mode torchscript

for running plain version lammps compiled with no AOTI, no triton, no KOKKOS, etc. ( just with very basic cmake ../cmake)

and second i compile with exactly the same last.ckpt file for the ML-IAP version by
2)

nequip-prepare-lmp-mliap \
  /absolute/last.ckpt  \
  /absolute/path/to/output.nequip.lmp.pt \
  --modifiers enable_CuEquivarianceContracter  enable_TritonContracter 

for running the ML-IAP interfaced lammps

Those two comes form exactly the same model (last.ckpt) but just different model compiling ways.
then i run a simple same nvt: (fix 1 all nvt temp 300 300 0.1) lammps simulation.

=================================================================================

The plain version behaves normally (thought slower and memory cost) like:

NequIP/Allegro is using input precision d and output precision d
NequIP/Allegro: Loading model from /home/liuyenfu/Allegro-0.7.1-Carbon/Allegro071-C_7A_lmax3_ly2_f64-32_CA9_D3-NEP_D3-E0-cu/outputs/2025-08-14/17-10-11/Allegro071-C_7A_lmax3_ly2_f64-32_CA9_D3-NEP_D3-E0-cu-epoch500.nequip.pth
NequIP/Allegro: Freezing TorchScript model...
Type mapping:
NequIP/Allegro type | NequIP/Allegro name | LAMMPS type | LAMMPS name
0 | C | 1 | C
Neighbor list info ...
update: every = 1 steps, delay = 0 steps, check = yes
max neighbors/atom: 2000, page size: 100000
master list distance cutoff = 9
ghost atom cutoff = 9
binsize = 4.5, bins = 23 23 100
1 neighbor lists, perpetual/occasional/extra = 1 0 0
(1) pair allegro, perpetual
attributes: full, newton on, ghost
pair build: full/bin/ghost
stencil: full/ghost/bin/3d
bin: standard
Setting up Verlet run ...
Unit style : metal
Current step : 0
Time step : 0.001
Per MPI rank memory allocation (min/avg/max) = 3.376 | 3.376 | 3.376 Mbytes
Step Temp E_pair E_mol TotEng Press
0 300 -9148.1611 0 -9106.94 295.13801
100 287.43131 -9147.042 0 -9107.5479 -175.41913
200 415.52377 -9169.358 0 -9112.2635 26.561809
300 270.0317 -9154.1299 0 -9117.0265 214.47927
400 319.21109 -9164.6545 0 -9120.7938 -87.872547
500 299.63101 -9164.7478 0 -9123.5774 -50.032376
600 286.33118 -9164.6303 0 -9125.2874 120.94988
700 298.94828 -9166.8669 0 -9125.7903 -53.146963
800 289.78729 -9164.0228 0 -9124.205 -81.058993
900 296.97551 -9163.3668 0 -9122.5613 73.77792
1000 316.46047 -9165.6286 0 -9122.1458 -63.682438

=================================================================================

while the ML-IAP version goes crazy like :

Neighbor list info ...
update: every = 1 steps, delay = 0 steps, check = yes
max neighbors/atom: 2000, page size: 100000
master list distance cutoff = 9
ghost atom cutoff = 9
binsize = 9, bins = 12 12 50
1 neighbor lists, perpetual/occasional/extra = 1 0 0
(1) pair mliap/kk, perpetual
attributes: full, newton on, kokkos_device
pair build: full/bin/kk/device
stencil: full/bin/3d
bin: kk/device
Setting up Verlet run ...
Unit style : metal
Current step : 0
Time step : 0.001
/home/liuyenfu/.conda/envs/allegro_071_mliap_tr271_cu126_cxx11/lib/python3.10/site-packages/nequip/integrations/lammps_mliap/lmp_mliap_wrapper.py:178: UserWarning: Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.
Consider using tensor.detach() first. (Triggered internally at /pytorch/torch/csrc/autograd/generated/python_variable_methods.cpp:835.)
lmp_data.energy = nequip_total_energy
Per MPI rank memory allocation (min/avg/max) = 231.5 | 231.5 | 231.5 Mbytes
Step Temp E_pair E_mol TotEng Press
0 300 -9148.1587 0 -9106.9377 -79.297354
100 103641.82 -5255.9926 0 8984.7604 2832.0119
200 10558.1 -5196.7912 0 -3746.071 -133.64315
300 13709.474 -4854.1952 0 -2970.465 -394.03189
400 16036.475 -4557.3333 0 -2353.8649 379.30914
500 6831.5228 -4880.376 0 -3941.7006 -67.489759
600 16419.679 -4346.3085 0 -2090.1866 81.498152
700 158691.93 -4627.9813 0 17176.851 4725.181
800 36212.557 -4608.1607 0 367.57255 968.15914
900 14540.336 -4675.6922 0 -2677.7986 339.44379
1000 4228.7012 -4678.2609 0 -4097.2224 -135.18055
Loop time of 34.0561 on 1 procs for 1000 steps with 1064 atoms

The Temp go crazy and the simulation collaps soon while in the OVITO movie the system looks like suddenly explosion and suddenly stop (very strange not make physical sense with Newton's law) .

=================================================================================

Those two results using exactly the same trained model just with different model compiling ways for different interface to run the same NVT simulation but with totally fifferent simulation results !!

[cw-tan]

One comment I can make is that you shouldn't be composing modifiers since they replace the same tensor product operations. Use one of the other. I doubt fixing this will fix the problem, but just for the record.

--modifiers enable_CuEquivarianceContracter  enable_TritonContracter 

[wcwitt]

Hi @liuyenfu, thanks for sticking with it. You mentioned the differences are for region-dependent thermostatting. Does that mean you were able to confirm normal behavior with a standard thermostat?

[cw-tan]

If you could send us the exact LAMMPS input file and job submission script (for the call signature), that would be helpful. Also, might be helpful if you have a training config file for us to try and reproduce this behavior.

[liuyenfu]

One comment I can make is that you shouldn't be composing modifiers since they replace the same tensor product operations. Use one of the other. I doubt fixing this will fix the problem, but just for the record.

--modifiers enable_CuEquivarianceContracter enable_TritonContracter

Hi, thanks for the indication. After using only one modifier (or no modifier), the NVT simulation become normal and there is no problem anymore about the NVT. However, the error of region-dependent thermostattin with "fix langevin" still exist no matter using which modifier.

[cw-tan]

Hi @liuyenfu just checking if you've learned anything new about this failure mode? Thanks!