diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py
index b2fb30dfc0..4658b16e7c 100644
--- a/deepmd/entrypoints/test.py
+++ b/deepmd/entrypoints/test.py
@@ -116,7 +116,13 @@ def test(
 
         # create data class
         tmap = dp.get_type_map() if dp.model_type == "ener" else None
-        data = DeepmdData(system, set_prefix, shuffle_test=shuffle_test, type_map=tmap)
+        data = DeepmdData(
+            system,
+            set_prefix,
+            shuffle_test=shuffle_test,
+            type_map=tmap,
+            sort_atoms=False,
+        )
 
         if dp.model_type == "ener":
             err = test_ener(
diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
index 8442f84156..485079b08d 100644
--- a/deepmd/utils/data.py
+++ b/deepmd/utils/data.py
@@ -42,6 +42,9 @@ class DeepmdData:
             Data modifier that has the method `modify_data`
     trn_all_set
             Use all sets as training dataset. Otherwise, if the number of sets is more than 1, the last set is left for test.
+    sort_atoms : bool
+            Sort atoms by atom types. Required to enable when the data is directly feeded to
+            descriptors except mixed types.
     """
 
     def __init__(
@@ -53,6 +56,7 @@ def __init__(
         optional_type_map: bool = True,
         modifier=None,
         trn_all_set: bool = False,
+        sort_atoms: bool = True,
     ):
         """Constructor."""
         root = DPPath(sys_path)
@@ -102,6 +106,7 @@ def __init__(
         if type_map is None and self.type_map is None and self.mixed_type:
             raise RuntimeError("mixed_type format must have type_map!")
         # make idx map
+        self.sort_atoms = sort_atoms
         self.idx_map = self._make_idx_map(self.atom_type)
         # train dirs
         self.test_dir = self.dirs[-1]
@@ -586,7 +591,10 @@ def _load_type_mix(self, set_name: DPPath):
     def _make_idx_map(self, atom_type):
         natoms = atom_type.shape[0]
         idx = np.arange(natoms)
-        idx_map = np.lexsort((idx, atom_type))
+        if self.sort_atoms:
+            idx_map = np.lexsort((idx, atom_type))
+        else:
+            idx_map = idx
         return idx_map
 
     def _load_type_map(self, sys_path: DPPath):
diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py
index 6d65165456..0bfe6b7c70 100644
--- a/deepmd/utils/data_system.py
+++ b/deepmd/utils/data_system.py
@@ -46,6 +46,7 @@ def __init__(
         trn_all_set=False,
         sys_probs=None,
         auto_prob_style="prob_sys_size",
+        sort_atoms: bool = True,
     ):
         """Constructor.
 
@@ -84,7 +85,10 @@ def __init__(
                                 the list of systems is devided into blocks. A block is specified by `stt_idx:end_idx:weight`,
                                 where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system,
                                 the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional
-        to the number of batches in the system.
+                to the number of batches in the system.
+        sort_atoms : bool
+            Sort atoms by atom types. Required to enable when the data is directly feeded to
+            descriptors except mixed types.
         """
         # init data
         self.rcut = rcut
@@ -101,6 +105,7 @@ def __init__(
                     optional_type_map=optional_type_map,
                     modifier=modifier,
                     trn_all_set=trn_all_set,
+                    sort_atoms=sort_atoms,
                 )
             )
         # check mix_type format
diff --git a/source/tests/test_dp_test.py b/source/tests/test_dp_test.py
index df1b51db0d..a07706acfe 100644
--- a/source/tests/test_dp_test.py
+++ b/source/tests/test_dp_test.py
@@ -180,8 +180,6 @@ def test_1frame(self):
             detail_file=detail_file,
             atomic=False,
         )
-        # TODO: see #2721
-        idx_map = np.lexsort((np.arange(len(self.atype)), self.atype))
         pred_e = np.loadtxt(detail_file + ".e.out", ndmin=2)[0, 1]
         pred_f = np.loadtxt(detail_file + ".f.out", ndmin=2)[:, 3:6]
         pred_v = np.loadtxt(detail_file + ".v.out", ndmin=2)[:, 9:18]
@@ -189,7 +187,7 @@ def test_1frame(self):
         pred_v_peratom = np.loadtxt(detail_file + ".v_peratom.out", ndmin=2)[:, 9:18]
         self.assertAlmostEqual(pred_e, np.sum(self.expected_e), places=default_places)
         np.testing.assert_almost_equal(
-            pred_f, self.expected_f.reshape(-1, 3)[idx_map], decimal=default_places
+            pred_f, self.expected_f.reshape(-1, 3), decimal=default_places
         )
         np.testing.assert_almost_equal(
             pred_v,