feat: add data parallel of native mindspore to mindnlp.Trainer.base (#1852)

Tridu33 · web-flow · commit c3c4e2dadfac · 2024-12-17T21:14:39.000+08:00
diff --git a/examples/parallel/bert_imdb_finetune/bert_imdb_finetune_cpu_mindnlp_trainer_npus_same.py b/examples/parallel/bert_imdb_finetune/bert_imdb_finetune_cpu_mindnlp_trainer_npus_same.py
@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+# coding: utf-8
+"""
+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():
+    """demo
+
+    Returns:
+        desc: _description_
+    """
+    imdb_ds = load_dataset('imdb', split=['train', 'test'])
+    imdb_train = imdb_ds['train']
+    imdb_train.get_dataset_size()
+
+    from mindnlp.transformers import AutoTokenizer
+    # tokenizer
+    tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
+
+    def process_dataset(dataset, tokenizer, max_seq_len=256, batch_size=32, shuffle=False):
+        is_ascend = mindspore.get_context('device_target') == 'Ascend'
+        def tokenize(text):
+            if is_ascend:
+                tokenized = tokenizer(text, padding='max_length', truncation=True, max_length=max_seq_len)
+            else:
+                tokenized = tokenizer(text, truncation=True, max_length=max_seq_len)
+            return tokenized['input_ids'], tokenized['token_type_ids'], tokenized['attention_mask']
+
+        if shuffle:
+            dataset = dataset.shuffle(batch_size)
+
+        # map dataset
+        dataset = dataset.map(operations=[tokenize], input_columns="text", output_columns=['input_ids', 'token_type_ids', 'attention_mask'])
+        dataset = dataset.map(operations=transforms.TypeCast(mindspore.int32), input_columns="label", output_columns="labels")
+        # batch dataset
+        if is_ascend:
+            dataset = dataset.batch(batch_size)
+        else:
+            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
+
+
+    dataset_train = process_dataset(imdb_train, tokenizer, shuffle=True)
+
+    next(dataset_train.create_tuple_iterator())
+
+    from mindnlp.transformers import AutoModelForSequenceClassification
+
+    # set bert config and define parameters for training
+    model = AutoModelForSequenceClassification.from_pretrained('bert-base-cased', num_labels=2)
+
+    from mindnlp.engine import TrainingArguments
+    
+    training_args = TrainingArguments(
+        output_dir="bert_imdb_finetune_cpu",
+        save_strategy="epoch",
+        logging_strategy="epoch",
+        num_train_epochs=2.0,
+        learning_rate=2e-5
+    )
+    training_args = training_args.set_optimizer(name="adamw", beta1=0.8) # 手动指定优化器，OptimizerNames.SGD
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset_train,
+    )
+    print("Start training")
+    trainer.train()
+
+if __name__ == '__main__':
+    main()
+
diff --git a/examples/parallel/bert_imdb_finetune/bert_imdb_finetune_npu_mindnlp_trainer.sh b/examples/parallel/bert_imdb_finetune/bert_imdb_finetune_npu_mindnlp_trainer.sh
@@ -0,0 +1,26 @@
+#!/bin/bash
+
+echo "=========================================="
+echo "Please run the script as: "
+echo "bash bert_imdb_finetune_npu_mindnlp_trainer.sh"
+echo "==========================================="
+
+EXEC_PATH=$(pwd)
+if [ ! -d "${EXEC_PATH}/data" ]; then
+    if [ ! -f "${EXEC_PATH}/emotion_detection.tar.gz" ]; then
+        wget wget https://baidu-nlp.bj.bcebos.com/emotion_detection-dataset-1.0.0.tar.gz -O emotion_detection.tar.gz
+    fi
+    tar xvf emotion_detection.tar.gz
+fi
+export DATA_PATH=${EXEC_PATH}/data/
+
+rm -rf bert_imdb_finetune_cpu_mindnlp_trainer_npus_same
+mkdir bert_imdb_finetune_cpu_mindnlp_trainer_npus_same
+echo "start training"
+
+export MULTI_NPU="true" 
+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=10 bert_imdb_finetune_cpu_mindnlp_trainer_npus_same.py 
diff --git a/mindnlp/accelerate/__init__.py b/mindnlp/accelerate/__init__.py
@@ -5,8 +5,9 @@
     # DDPCommunicationHookType,
     # DeepSpeedPlugin,
     # DistributedDataParallelKwargs,
-    # DistributedType,
     # FullyShardedDataParallelPlugin,
+    accelerate_distributed_type,
+    DistributedType,
     # GradScalerKwargs,
     # InitProcessGroupKwargs,
     # ProfileKwargs,
diff --git a/mindnlp/accelerate/accelerator.py b/mindnlp/accelerate/accelerator.py
@@ -9,11 +9,11 @@
 
 from .state import AcceleratorState
 from .utils import (
-    DistributedType,
     MindFormersPlugin,
     is_mindformers_available,
     wait_for_everyone
 )
+from .utils import DistributedType,accelerate_distributed_type
 from ..utils import logging
 
 if is_mindformers_available():
@@ -45,7 +45,7 @@ def __init__(
         # init mindformers_plugin from env variables
         if mindformers_plugin is None:
             mindformers_plugin = (
-                MindFormersPlugin() if os.environ.get("ACCELERATE_USE_MINDFORMERS", "false") == "true" else None
+                MindFormersPlugin() if accelerate_distributed_type == DistributedType.MINDFORMERS else None
             )
         else:
             os.environ["ACCELERATE_USE_MINDFORMERS"] = "true"
@@ -104,10 +104,11 @@ def prepare(self, *args):
         """
         result = []
 
-        # Only support mindsormers now
+        # Only support mindsormers and MULTI_NPU now
         if self.distributed_type == DistributedType.MINDFORMERS:
             result = self._prepare_mindformers(*args)
-
+        elif self.distributed_type == DistributedType.MULTI_NPU:
+            pass # nothing prepare for data parallel
         return result
 
     def _prepare_mindformers(self, *args):
diff --git a/mindnlp/accelerate/state.py b/mindnlp/accelerate/state.py
@@ -1,17 +1,18 @@
 """accelerate state"""
-import os
 from functools import partial
 from contextlib import contextmanager
 from typing import Callable, Any
 from mindspore import communication
+
 try:
     from mindspore.communication.comm_func import barrier
 except:
     barrier = None
 
 from .utils import (
-    DistributedType, is_mindformers_available
+    is_mindformers_available
 )
+from ..accelerate.utils import accelerate_distributed_type, DistributedType
 
 SharedDict = dict
 
@@ -341,10 +342,13 @@ def print(self, *args, **kwargs):
             print(*args, **kwargs)
 
     def _prepare_backend(self):
-        # now mindformers only
-        if is_mindformers_available():
+        # now mindformers and mindspore data parallel only
+        if accelerate_distributed_type == DistributedType.MINDFORMERS and is_mindformers_available():
             self.backend = "hccl"
             self.distributed_type = DistributedType.MINDFORMERS
+        elif accelerate_distributed_type == DistributedType.MULTI_NPU:
+            self.backend = "hccl"
+            self.distributed_type = DistributedType.MULTI_NPU
 
     @num_processes.setter
     def num_processes(self, value):
@@ -366,10 +370,14 @@ def __init__(self, mindformers_plugin=None, **kwargs):
         if PartialState._shared_state:
             PartialState(**kwargs)
         self.__dict__.update(PartialState._shared_state)
-
-        if os.environ.get("ACCELERATE_USE_MINDFORMERS", "false") == "true":
+        # set distributed_type
+        if accelerate_distributed_type == DistributedType.MULTI_NPU:
+            self.distributed_type = DistributedType.MULTI_NPU
+        elif accelerate_distributed_type == DistributedType.MINDFORMERS:
             self.distributed_type = DistributedType.MINDFORMERS
             self.mindformers_plugin = mindformers_plugin
+        else:
+            self.distributed_type = DistributedType.NO
 
         PartialState._shared_state["distributed_type"] = self.distributed_type
 
diff --git a/mindnlp/accelerate/utils/__init__.py b/mindnlp/accelerate/utils/__init__.py
@@ -1,4 +1,5 @@
 """accelerate utils"""
+from .constants import accelerate_distributed_type
 from .dataclasses import (
     DistributedType,
     MindFormersPlugin
diff --git a/mindnlp/accelerate/utils/constants.py b/mindnlp/accelerate/utils/constants.py
@@ -0,0 +1,34 @@
+"""constants"""
+import os
+import mindspore
+import numpy
+from .dataclasses import DistributedType
+
+
+_random_seed = numpy.random.randint(1000)
+
+
+def _prepare_data_parallel_native_minspore():
+    # initialize data parallel hcc backend for data_loader and Trainer API
+    mindspore.set_auto_parallel_context(parallel_mode=mindspore.ParallelMode.DATA_PARALLEL, gradients_mean=True)
+    mindspore.communication.init()
+    mindspore.set_seed(_random_seed)
+
+
+def detect_accelerate_distributed_type():
+    """
+    detect distributed_type
+
+    Returns:
+        _type_: According to the factors such as the available parallel software and hardware environment of the current system and the user-specified parallel scheme,
+          the optimal parallel strategy is comprehensively decided in different situations.
+    """
+    if os.environ.get("MULTI_NPU", None) == "true":
+        _prepare_data_parallel_native_minspore()
+        return DistributedType.MULTI_NPU
+    if os.environ.get("ACCELERATE_USE_MINDFORMERS", "false") == "true":
+        return DistributedType.MINDFORMERS
+    else:
+        return DistributedType.NO
+
+accelerate_distributed_type = detect_accelerate_distributed_type()
diff --git a/mindnlp/accelerate/utils/dataclasses.py b/mindnlp/accelerate/utils/dataclasses.py
@@ -17,8 +17,11 @@ class DistributedType(str, enum.Enum):
 
     Values:
         - **MINDFORMERS** -- Using mindformers
+        - **NO** -- Not a distributed environment, just a single process.
+        - **MULTI_NPU** -- Distributed data parallel on multiple NPUs.
     """
 
+    MULTI_NPU = "MULTI_NPU"
     MINDFORMERS = "MINDFORMERS"
     NO = "NO"
 
diff --git a/mindnlp/dataset/load.py b/mindnlp/dataset/load.py
@@ -18,11 +18,15 @@
 """
 import os
 from typing import Union, Optional, Dict, Sequence, Mapping
-from mindspore.dataset import GeneratorDataset
 from datasets import load_dataset as hf_load
 from datasets import Dataset, IterableDataset, Split, Features, \
     DownloadConfig, DownloadMode, VerificationMode, Version
+from mindspore.dataset import GeneratorDataset
+from mindspore.communication import get_rank, get_group_size
 from mindnlp.configs import DEFAULT_ROOT
+from ..accelerate import DistributedType
+from ..accelerate.utils import accelerate_distributed_type
+
 
 class TransferIterableDataset():
     """TransferDataset for Huggingface Dataset."""
@@ -331,12 +335,19 @@ def load_dataset(
         column_names = list(raw_ds.features.keys())
         source = TransferDataset(raw_ds, column_names) if isinstance(raw_ds, Dataset) \
             else TransferIterableDataset(raw_ds, column_names)
-        ms_ds = GeneratorDataset(
-            source=source,
+        if accelerate_distributed_type == DistributedType.MULTI_NPU:
+            ms_ds = GeneratorDataset(source=source,
+            column_names=column_names,
+            shuffle=shuffle,
+            num_parallel_workers=num_proc if num_proc else 1,
+            num_shards=get_group_size(), shard_id=get_rank())
+            datasets_dict[key] = ms_ds
+        else:
+            ms_ds = GeneratorDataset(source=source,
             column_names=column_names,
             shuffle=shuffle,
             num_parallel_workers=num_proc if num_proc else 1)
-        datasets_dict[key] = ms_ds
+            datasets_dict[key] = ms_ds
 
     if len(datasets_dict) == 1:
         return datasets_dict.popitem()[1]
diff --git a/mindnlp/engine/trainer/base.py b/mindnlp/engine/trainer/base.py
@@ -45,6 +45,8 @@
     WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, SAFE_WEIGHTS_INDEX_NAME
 from ...dataset import BaseMapFunction
 from ...utils import logging, find_labels, can_return_loss
+from ...accelerate.utils import DistributedType
+from ...accelerate.utils import accelerate_distributed_type
 from ...utils.import_utils import is_safetensors_available
 from ...transformers.modeling_utils import PreTrainedModel
 from ...transformers.configuration_utils import PretrainedConfig
@@ -124,7 +126,6 @@ class Trainer:
     """
     Trainer is a simple but feature-complete training and eval loop for MindSpore, optimized for 🤗 Transformers.
     """
-    from ..utils import _get_learning_rate
     def __init__(
         self,
         model: Union[PreTrainedModel, nn.Module] = None,
@@ -284,6 +285,30 @@ def __init__(
         # Internal variables to help with automatic batch size reduction
         self._train_batch_size = args.train_batch_size
         self._created_lr_scheduler = False
+        self.actual_distributed_type = accelerate_distributed_type
+
+
+    def _get_learning_rate(self):
+        r"""
+        This function retrieves the learning rate used by the optimizer.
+        
+        Args:
+            self: An instance of the class containing the optimizer and learning rate scheduler.
+        
+        Returns:
+            The learning rate value (float) used by the optimizer.
+        
+        Raises:
+            None.
+        """
+        if isinstance(self.lr_scheduler, optim.lr_scheduler.ReduceLROnPlateau):
+            last_lr = self.optimizer.param_groups[0]["lr"]
+        else:
+            last_lr = self.lr_scheduler.get_last_lr()[0]
+        if ops.is_tensor(last_lr):
+            last_lr = last_lr.item()
+        return last_lr
+
 
     def _activate_neftune(self, model):
         r"""
@@ -1133,6 +1158,7 @@ def _inner_training_loop(
                             model.parameters(),
                             args.max_grad_norm,
                         )
+
                     # Optimizer step
                     self.optimizer.step()
 
@@ -1351,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.
@@ -1382,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):
diff --git a/mindnlp/engine/utils.py b/mindnlp/engine/utils.py
diff --git a/tests/accelerate/grad_Reduce_ut/test_grad_Reduce.py b/tests/accelerate/grad_Reduce_ut/test_grad_Reduce.py
diff --git a/tests/accelerate/grad_Reduce_ut/test_grad_Reduce.sh b/tests/accelerate/grad_Reduce_ut/test_grad_Reduce.sh

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`"""accelerate utils"""`
	`2`	`+from .constants import accelerate_distributed_type`
`2`	`3`	`from .dataclasses import (`
`3`	`4`	`DistributedType,`
`4`	`5`	`MindFormersPlugin`