fix: fix wrong call time for gradient AllReduce

Author: Tridu33

examples/parallel/bert_imdb_finetune/bert_imdb_finetune_cpu_mindnlp_trainer_npus_same.py

@@ -1,20 +1,28 @@
 #!/usr/bin/env python
 # coding: utf-8
 """
-python bert_imdb_finetune_cpu_mindnlp_trainer.py
+unset MULTI_NPU && python bert_imdb_finetune_cpu_mindnlp_trainer_npus_same.py
 bash bert_imdb_finetune_npu_mindnlp_trainer.sh
 """
 
+import mindspore
+from mindspore.dataset import transforms
+from mindnlp.engine import Trainer
+from mindnlp.dataset import load_dataset
+
+from mindnlp.accelerate.utils.constants import accelerate_distributed_type
+from mindnlp.accelerate.utils.dataclasses import DistributedType
+
 def main():
-    import mindspore
-    from mindspore.dataset import transforms
-    from mindnlp.engine import Trainer
-    from mindnlp.dataset import load_dataset
+    """demo
 
+    Returns:
+        desc: _description_
+    """
     imdb_ds = load_dataset('imdb', split=['train', 'test'])
     imdb_train = imdb_ds['train']
-    imdb_test = imdb_ds['test']
     imdb_train.get_dataset_size()
+
     from mindnlp.transformers import AutoTokenizer
     # tokenizer
     tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
@@ -41,15 +49,10 @@ def tokenize(text):
             dataset = dataset.padded_batch(batch_size, pad_info={'input_ids': (None, tokenizer.pad_token_id),
                                                                 'token_type_ids': (None, 0),
                                                                 'attention_mask': (None, 0)})
-
         return dataset
 
-    # split train dataset into train and valid datasets
-    imdb_train, imdb_val = imdb_train.split([0.7, 0.3])
 
     dataset_train = process_dataset(imdb_train, tokenizer, shuffle=True)
-    dataset_val = process_dataset(imdb_val, tokenizer)
-    dataset_test = process_dataset(imdb_test, tokenizer)
 
     next(dataset_train.create_tuple_iterator())
 
@@ -62,36 +65,21 @@ def tokenize(text):
 
     training_args = TrainingArguments(
         output_dir="bert_imdb_finetune_cpu",
-        evaluation_strategy="epoch",
         save_strategy="epoch",
         logging_strategy="epoch",
-        load_best_model_at_end=True,
         num_train_epochs=2.0,
         learning_rate=2e-5
     )
-    training_args = training_args.set_optimizer(name="adamw", beta1=0.8) # OptimizerNames.SGD
-
-    from mindnlp import evaluate
-    import numpy as np
-    metric = evaluate.load("accuracy")
-    def compute_metrics(eval_pred):
-        logits, labels = eval_pred
-        predictions = np.argmax(logits, axis=-1)
-        return metric.compute(predictions=predictions, references=labels)
-    
+    training_args = training_args.set_optimizer(name="adamw", beta1=0.8) # 手动指定优化器，OptimizerNames.SGD
+
     trainer = Trainer(
         model=model,
         args=training_args,
         train_dataset=dataset_train,
-        eval_dataset=dataset_val,
-        compute_metrics=compute_metrics
     )
     print("Start training")
     trainer.train()
 
-    print("Start checking the test set")
-    trainer.evaluate(dataset_test)
-
 if __name__ == '__main__':
     main()
 

examples/parallel/bert_imdb_finetune/bert_imdb_finetune_npu_mindnlp_trainer.sh

@@ -19,7 +19,8 @@ mkdir bert_imdb_finetune_cpu_mindnlp_trainer_npus_same
 echo "start training"
 
 export MULTI_NPU="true" 
-# unset MULTI_NPU
-msrun --worker_num=2 --local_worker_num=2 --master_port=8119 \
+export ASCEND_SLOG_PRINT_TO_STDOUT=1
+
+msrun --worker_num=2 --local_worker_num=2 --master_port=8121 \
 --log_dir=bert_imdb_finetune_cpu_mindnlp_trainer_npus_same --join=True \
---cluster_time_out=30 bert_imdb_finetune_cpu_mindnlp_trainer_npus_same.py 
+--cluster_time_out=10 bert_imdb_finetune_cpu_mindnlp_trainer_npus_same.py 

mindnlp/engine/trainer/base.py

@@ -1377,6 +1377,20 @@ def _prepare_inputs(self, inputs: Dict[str, Union[mindspore.Tensor, Any]]) -> Di
 
         return inputs
 
+
+    def update_gradient_by_distributed_type(self, model: nn.Module) -> None:
+        """update gradient by distributed_type"""
+        if accelerate_distributed_type == DistributedType.NO:
+            return
+        if accelerate_distributed_type == DistributedType.MULTI_NPU:
+            from mindspore.communication import get_group_size
+            from mindspore.communication.comm_func import all_reduce
+            rank_size = get_group_size()
+            for parameter in model.parameters():
+                new_grads_mean = all_reduce(parameter.grad) / rank_size
+                parameter.grad = new_grads_mean
+
+
     def training_step(self, model: nn.Module, inputs: Dict[str, Union[mindspore.Tensor, Any]]) -> Tuple[List[mindspore.Tensor], mindspore.Tensor]:
         """
         Perform a training step on a batch of inputs.
@@ -1399,14 +1413,6 @@ def training_step(self, model: nn.Module, inputs: Dict[str, Union[mindspore.Tens
         inputs = self._prepare_inputs(inputs)
 
         def forward(inputs):
-            if accelerate_distributed_type == DistributedType.MULTI_NPU:
-                from mindspore.communication import get_group_size
-                import mindspore.ops as msops
-                rank_size = get_group_size()
-                for parameter in model.parameters():
-                    all_reduce_sum =  msops.AllReduce(msops.ReduceOp.SUM)
-                    new_grads_mean = all_reduce_sum(parameter.grad) / rank_size
-                    parameter.grad = new_grads_mean
             return self.compute_loss(model, inputs)
 
         if getattr(self, 'grad_fn', None) is None or self.model_reload:
@@ -1416,7 +1422,7 @@ def forward(inputs):
             self.grad_fn = value_and_grad(forward, weights, attach_grads=True)
 
         loss = self.grad_fn(inputs)
-
+        self.update_gradient_by_distributed_type(model)
         return loss / self.args.gradient_accumulation_steps
 
     def compute_loss(self, model, inputs, return_outputs=False):