[ERNIE-M] Fix layerwise decay ratio to model_zoo/ernie-m (PaddlePaddle#4892)

Author: Yam0214

model_zoo/ernie-m/run_classifier.py

@@ -24,11 +24,9 @@
 from datasets import load_dataset
 from paddle.io import Dataset
 from paddle.metric import Accuracy
-from paddle.optimizer import AdamW
 
 import paddlenlp
 from paddlenlp.data import DataCollatorWithPadding
-from paddlenlp.ops.optimizer import layerwise_lr_decay
 from paddlenlp.trainer import (
     PdArgumentParser,
     Trainer,
@@ -39,7 +37,6 @@
     AutoModelForSequenceClassification,
     AutoTokenizer,
     ErnieMForSequenceClassification,
-    LinearDecayWithWarmup,
 )
 from paddlenlp.utils.log import logger
 
@@ -220,40 +217,6 @@ def compute_metrics(p):
         accu = metric.accumulate()
         return {"accuracy": accu}
 
-    n_layers = model.ernie_m.config["num_hidden_layers"]
-    warmup = training_args.warmup_steps if training_args.warmup_steps > 0 else training_args.warmup_ratio
-    if training_args.do_train:
-        num_training_steps = (
-            training_args.max_steps
-            if training_args.max_steps > 0
-            else len(train_ds) // training_args.train_batch_size * training_args.num_train_epochs
-        )
-    else:
-        num_training_steps = 10
-
-    lr_scheduler = LinearDecayWithWarmup(training_args.learning_rate, num_training_steps, warmup)
-
-    # Generate parameter names needed to perform weight decay.
-    # All bias and LayerNorm parameters are excluded.
-    decay_params = [p.name for n, p in model.named_parameters() if not any(nd in n for nd in ["bias", "norm"])]
-    # Construct dict
-    name_dict = dict()
-    for n, p in model.named_parameters():
-        name_dict[p.name] = n
-
-    simple_lr_setting = partial(layerwise_lr_decay, model_args.layerwise_decay, name_dict, n_layers)
-
-    optimizer = AdamW(
-        learning_rate=lr_scheduler,
-        beta1=0.9,
-        beta2=0.999,
-        epsilon=training_args.adam_epsilon,
-        parameters=model.parameters(),
-        weight_decay=training_args.weight_decay,
-        apply_decay_param_fun=lambda x: x in decay_params,
-        lr_ratio=simple_lr_setting,
-    )
-
     trainer = Trainer(
         model=model,
         args=training_args,
@@ -262,9 +225,41 @@ def compute_metrics(p):
         eval_dataset=eval_ds if training_args.do_eval else None,
         tokenizer=tokenizer,
         compute_metrics=compute_metrics,
-        optimizers=[optimizer, lr_scheduler],
+        # optimizers=[optimizer, lr_scheduler],
     )
 
+    def using_layerwise_lr_decay(layerwise_decay, model, training_args):
+        """
+        Generate parameter names needed to perform weight decay.
+        All bias and LayerNorm parameters are excluded.
+        """
+        # params_list = [{"params": param, "learning_rate": lr * decay_ratio}, ... ]
+        params_list = []
+        n_layers = model.config.num_hidden_layers
+        for name, param in model.named_parameters():
+            ratio = 1.0
+            param_to_train = {"params": param, "dygraph_key_name": name}
+            if any(nd in name for nd in ["bias", "norm"]):
+                param_to_train["weight_decay"] = 0.0
+            else:
+                param_to_train["weight_decay"] = training_args.weight_decay
+
+            if "encoder.layers" in name:
+                idx = name.find("encoder.layers.")
+                layer = int(name[idx:].split(".")[2])
+                ratio = layerwise_decay ** (n_layers - layer)
+            elif "embedding" in name:
+                ratio = layerwise_decay ** (n_layers + 1)
+
+            param_to_train["learning_rate"] = ratio
+
+            params_list.append(param_to_train)
+        return params_list
+
+    params_to_train = using_layerwise_lr_decay(model_args.layerwise_decay, model, training_args)
+
+    trainer.set_optimizer_grouped_parameters(params_to_train)
+
     checkpoint = None
     if training_args.resume_from_checkpoint is not None:
         checkpoint = training_args.resume_from_checkpoint
@@ -320,6 +315,7 @@ def compute_metrics(p):
         paddlenlp.transformers.export_model(
             model=model_to_save, input_spec=input_spec, path=model_args.export_model_dir
         )
+        trainer.tokenizer.save_pretrained(model_args.export_model_dir)
 
 
 if __name__ == "__main__":