Check layer types for Optimizer construction (#10598)

* Check layer types for Optimizer construction

* Duplicate class

tests/test_trainer_utils.py

@@ -30,8 +30,23 @@ if is_torch_available():
         DistributedTensorGatherer,
         LabelSmoother,
         LengthGroupedSampler,
+        get_parameter_names,
     )
 
+    class TstLayer(torch.nn.Module):
+        def __init__(self, hidden_size):
+            super().__init__()
+            self.linear1 = torch.nn.Linear(hidden_size, hidden_size)
+            self.ln1 = torch.nn.LayerNorm(hidden_size)
+            self.linear2 = torch.nn.Linear(hidden_size, hidden_size)
+            self.ln2 = torch.nn.LayerNorm(hidden_size)
+            self.bias = torch.nn.Parameter(torch.zeros(hidden_size))
+
+        def forward(self, x):
+            h = self.ln1(torch.nn.functional.relu(self.linear1(x)))
+            h = torch.nn.functional.relu(self.linear2(x))
+            return self.ln2(x + h + self.bias)
+
 
 @require_torch
 class TrainerUtilsTest(unittest.TestCase):
@@ -117,3 +132,12 @@ class TrainerUtilsTest(unittest.TestCase):
         self.assertEqual(lengths[indices_process_0[0]], 50)
         # The indices should be a permutation of range(100)
         self.assertEqual(list(sorted(indices_process_0 + indices_process_1)), list(range(100)))
+
+    def test_get_parameter_names(self):
+        model = torch.nn.Sequential(TstLayer(128), torch.nn.ModuleList([TstLayer(128), TstLayer(128)]))
+        # fmt: off
+        self.assertEqual(
+            get_parameter_names(model, [torch.nn.LayerNorm]),
+            ['0.linear1.weight', '0.linear1.bias', '0.linear2.weight', '0.linear2.bias', '0.bias', '1.0.linear1.weight', '1.0.linear1.bias', '1.0.linear2.weight', '1.0.linear2.bias', '1.0.bias', '1.1.linear1.weight', '1.1.linear1.bias', '1.1.linear2.weight', '1.1.linear2.bias', '1.1.bias']
+        )
+        # fmt: on