Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2026-09-02 12:09:36 +02:00
Code Issues Packages Projects Releases Wiki Activity

[to #42322933] Support multi-machine data and tensor parallel finetuning

Browse Source 
Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11682479
This commit is contained in:
hemu.zp
2023-02-28 18:47:17 +08:00
committed by wenmeng.zwm
parent 51cf0b16a8
commit ceeb85f10f
18 changed files with 341 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
modelscope/models/nlp/gpt3/configuration.py
View File

@@ -21,6 +21,83 @@ logger = logging.get_logger()
class GPT3Config(PretrainedConfig):
r"""
Configuration classes for GPT-3 model.
Class attributes:
- **model_type** (`str`) -- An identifier for the model type, serialized into the JSON file, can be used to recreate
the correct object in [`~transformers.AutoConfig`].
Args:
vocab_size (`int`, *optional*, defaults to 25600):
Vocabulary size of the GPT model. Defines the number of different
tokens that can be represented by the `inputs_ids` passed when
calling [`GPT3Model`].
hidden_size (`int`, *optional*, defaults to 768):
Dimensionality of the decoder layers and the pooler layer.
ffn_hidden_size (`int`, *optional*, defaults to None):
Dimensionality of the ffn layer, None defaults to four times the hidden_size.
num_hidden_layers (`int`, *optional*, defaults to 12):
Number of hidden layers in the Transformer decoder.
num_attention_heads (`int`, *optional*, defaults to 12):
Number of attention heads for each attention layer in the
Transformer decoder.
intermediate_size (`int`, *optional*, defaults to 3072):
Dimensionality of the "intermediate" (often named feed-forward)
layer in the Transformer decoder.
hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
The non-linear activation function (function or string) in the
decoder and pooler. If string, `"gelu"`, `"relu"`, `"silu"` and
`"gelu_new"` are supported.
hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
The dropout probability for all fully connected layers in the
embeddings, decoder, and pooler.
attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
The dropout ratio for the attention probabilities.
max_position_embeddings (`int`, *optional*, defaults to 512):
The maximum sequence length that this model might ever be used with.
Typically set this to something large just in case (e.g., 512 or
1024 or 2048).
type_vocab_size (`int`, *optional*, defaults to 2):
The vocabulary size of the `token_type_ids` passed when calling
[`GPT3Model`].
layernorm_epsilon (`float`, *optional*, defaults to 1e-12):
The epsilon used by the layer normalization layers.
bias_gelu_fusion (`bool`, *optional*, defaults to True):
Whether to use gelu activation function when mixing bias.
fp32_residual_connection (`bool`, *optional*, defaults to False):
Whether to use fp32 for residual connection
between layers to improve accuracy.
sequence_parallel (`bool`, *optional*, defaults to False):
Whether to use sequence parallel during training.
bf16 (`bool`, *optional*, defaults to `False`):
Whether to use bf16 16-bit (mixed) precision training instead of 32-bit training.
Requires Ampere or higher NVIDIA architecture or using CPU (no_cuda).
This is an experimental API and it may change.
fp16 (`bool`, *optional*, defaults to `False`):
Whether to use fp16 16-bit (mixed) precision training instead of 32-bit training.
apply_query_key_layer_scaling (`bool`, *optional*, defaults to `True`):
Whether to scale query and key layer parameters during training.
init_method_std (`float`, *optional*, defaults to `0.02`):
The standard deviation of the normal distribution for initialization process.
eod_id (`int`, *optional*, defaults to `1`):
The end of text label for tokenizer, also indicates the end of the generation.
tokens_to_generate (`int`, *optional*, defaults to 100):
Number of tokens to generate.
top_k (`int`, *optional*, defaults to 0):
Number of highest probability vocabulary tokens to keep for
top-k-filtering that will be used by default in
the `generate` method of the model.
top_p (`float`, *optional*, defaults to 0.9):
Value that will be used by default in the `generate` method of the model
for `top_p`. If set to float < 1,
only the most probable tokens with probabilities that add up to `top_p`
or higher are kept for generation.
temperature (`float`, *optional*, defaults to 1.0):
The value used to module the next token probabilities that will be used
by default in the `generate` method of the model. Must be strictly positive.
"""
model_type = 'gpt3'
@@ -53,10 +130,11 @@ class GPT3Config(PretrainedConfig):
hidden_dropout=0.1,
init_method_std=0.02,
# generate
eod_id=7,
eod_id=1,
tokens_to_generate=100,
top_k=0,
top_p=0.9,
temperature=1.0,
**kwargs):
super().__init__(layer_norm_eps=layernorm_epsilon, **kwargs)
@@ -95,6 +173,7 @@ class GPT3Config(PretrainedConfig):
self.tokens_to_generate = tokens_to_generate
self.top_k = top_k
self.top_p = top_p
self.temperature = temperature
TORCH_MAJOR = int(torch.__version__.split('.')[0])
TORCH_MINOR = int(torch.__version__.split('.')[1])

73
modelscope/models/nlp/gpt3/distributed_gpt3.py
View File

@@ -13,14 +13,13 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import io
import math
import os
from os import path as osp
from collections import OrderedDict
from typing import Callable, Dict, List, Optional, Union
import torch
from megatron_util import mpu
from megatron_util import get_args, mpu
from megatron_util.global_vars import get_global_memory_buffer
from megatron_util.model import (AttnMaskType, Float16Module, LayerNorm,
bias_gelu_impl)
@@ -29,11 +28,9 @@ from torch import nn
from torch.nn import functional as F
from transformers.modeling_utils import PreTrainedModel
from modelscope.fileio import File
from modelscope.models import TorchModel
from modelscope.models.nlp.gpt3 import GPT3Config
from modelscope.outputs import TextGenerationModelOutput, TokenGeneratorOutput
from modelscope.utils.checkpoint import weights_to_cpu
from modelscope.utils.megatron_utils import init_megatron_util
from modelscope.utils.nlp.load_checkpoint import pre_load
@@ -948,21 +945,6 @@ def split_state_dict(state_dict: Dict[str, torch.Tensor], model: GPT3Model,
return state_dict
def save_checkpoint(model: torch.nn.Module, filename: str, **kwargs) -> None:
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
model = model.module
checkpoint = {'module': weights_to_cpu(model.state_dict())}
mp_rank = mpu.get_tensor_model_parallel_rank()
filename = osp.join(
osp.dirname(filename), 'model',
'mp_rank_{:02d}'.format(mp_rank) + '_model_states.pt')
with io.BytesIO() as f:
torch.save(checkpoint, f)
File.write(f.getvalue(), filename)
class DistributedGPT3(TorchModel):
def __init__(self,
@@ -992,7 +974,8 @@ class DistributedGPT3(TorchModel):
self.dist_model = model
tensor_ws = mpu.get_tensor_model_parallel_world_size()
ckpt_ws = kwargs.pop('checkpoint_model_parallel_size', tensor_ws)
ckpt_ws = get_args().get('checkpoint_tensor_model_parallel_size',
tensor_ws)
ckpt_rank = mpu.get_tensor_model_parallel_rank() * ckpt_ws // tensor_ws
load_model = pre_load(ckpt_rank, model_dir, tag=path_load_tag)
load_model = split_state_dict(load_model, model, tensor_ws // ckpt_ws)
@@ -1011,8 +994,8 @@ class DistributedGPT3(TorchModel):
attention_mask=None,
position_ids=None,
labels=None,
prompt_length=None,
is_pair=(False, )):
prompts_len=None,
inputs_len=None):
logits, losses = self.dist_model(
tokens,
@@ -1026,10 +1009,15 @@ class DistributedGPT3(TorchModel):
self.inference_params.sequence_len_offset += tokens.size(1)
else:
loss_mask = torch.ones(
tokens.size(), dtype=torch.float, device=tokens.device)
if is_pair[0]:
for i, length in enumerate(prompt_length):
loss_mask[i, :length] = 0
labels.size(), dtype=torch.float, device=tokens.device)
if inputs_len is None:
for i, l in enumerate(prompts_len):
loss_mask[i, l:] = 0
else:
for i, l in enumerate(inputs_len):
loss_mask[i, l - 1:] = 0
for i, l in enumerate(prompts_len):
loss_mask[i, :l - 1] = 0
losses = losses.float()
loss_mask = loss_mask.view(-1).float()
@@ -1039,15 +1027,15 @@ class DistributedGPT3(TorchModel):
def sample(self,
tokens,
temperature=1.0,
prompts_len=None,
use_eod_token_for_early_termination=True,
stop_on_double_eol=False,
stop_on_eol=False,
**kwargs):
batch_size = tokens.size(0)
lengths = kwargs.pop(
'prompt_length',
torch.tensor([tokens.size(1)], device=tokens.device))
lengths = prompts_len
if lengths is None:
lengths = torch.tensor([tokens.size(1)], device=tokens.device)
pads = torch.ones(
batch_size, self.config.tokens_to_generate,
device=tokens.device).long() * self.config.eod_id
@@ -1096,9 +1084,9 @@ class DistributedGPT3(TorchModel):
last_token_logits = logits[:, -1, :]
new_sample = sample(
last_token_logits,
top_k=self.config.top_k,
top_p=self.config.top_p,
temperature=temperature,
top_k=kwargs.pop('top_k', self.config.top_k),
top_p=kwargs.pop('top_p', self.config.top_p),
temperature=kwargs.pop('temperature', self.config.temperature),
vocab_size=self.config.vocab_size)
# If a prompt length is smaller or equal th current context
@@ -1257,6 +1245,11 @@ class DistributedGPT3(TorchModel):
def state_dict(self, destination=None, prefix='', keep_vars=False):
return self.dist_model.state_dict(destination, prefix, keep_vars)
def load_state_dict(self,
state_dict: 'OrderedDict[str, torch.Tensor]',
strict: bool = True):
return self.dist_model.load_state_dict(state_dict, strict)
def save_pretrained(self,
target_folder: Union[str, os.PathLike],
save_checkpoint_names: Union[str, List[str]] = None,
@@ -1265,13 +1258,15 @@ class DistributedGPT3(TorchModel):
**kwargs):
# DistributedPipeline type is different from task name
config['pipeline']['type'] = 'gpt3-generation'
# a temp fix for master_ip, master_port and rank
# can be removed after refactoring megatron_util
for unused_key in ('master_ip', 'master_port', 'rank'):
config['model'].pop(unused_key, None)
config['model'].pop('rank', None)
config['megatron'].pop('checkpoint_tensor_model_parallel_size', None)
tp_size = get_args().tensor_model_parallel_size
pp_size = get_args().pipeline_model_parallel_size
config['megatron']['world_size'] = tp_size * pp_size
return super().save_pretrained(target_folder, save_checkpoint_names,
save_checkpoint, config, **kwargs)
save_function, config, **kwargs)
class BeamHypotheses:

12
modelscope/models/nlp/gpt3/text_generation.py
View File

@@ -1,4 +1,5 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from collections import OrderedDict
from typing import Dict
import torch
@@ -9,6 +10,7 @@ from modelscope.models.base import Tensor, TorchModel
from modelscope.models.builder import MODELS
from modelscope.models.nlp.gpt3 import GPT3Model
from modelscope.utils.constant import Tasks
from modelscope.utils.hub import read_config
__all__ = ['GPT3ForTextGeneration']
@@ -27,7 +29,7 @@ class GPT3ForTextGeneration(TorchModel):
# Temporarily compatible with DistributedGPT3 and GPT3Model,
# the base/large model based on GPT3Model will be replaced in the future,
# and GPT3Model will be deprecated
if 'world_size' in kwargs:
if 'megatron' in read_config(model_dir):
from modelscope.models.nlp import DistributedGPT3
self.model = DistributedGPT3(model_dir, **kwargs)
else:
@@ -66,3 +68,11 @@ class GPT3ForTextGeneration(TorchModel):
if not isinstance(self.model, GPT3Model):
return self.model.save_pretrained(*args, **kwargs)
return super().save_pretrained(*args, **kwargs)
def state_dict(self, destination=None, prefix='', keep_vars=False):
return self.model.state_dict(destination, prefix, keep_vars)
def load_state_dict(self,
state_dict: 'OrderedDict[str, Tensor]',
strict: bool = True):
return self.model.load_state_dict(state_dict, strict)

6
modelscope/models/nlp/gpt3/tokenizer.py
View File

@@ -32,7 +32,7 @@ class JiebaBPETokenizer:
except ImportError:
raise ImportError(
'You need to install jieba to use JiebaTokenizer. '
'See https://pypi.org/project/rjieba/ for installation.')
'See https://pypi.org/project/jieba/ for installation.')
self.jieba = jieba
self.new_line = self.vocab['\n']
self.sep_token = self.vocab['<sep>']
@@ -64,7 +64,9 @@ class JiebaBPETokenizer:
return self.tokenizer.encode(
text, is_pretokenized=False, add_special_tokens=True).ids
def detokenize(self, token_ids):
def detokenize(self, token_ids: List[int], early_stop: bool = True) -> str:
if early_stop and self.sep_token in token_ids:
token_ids = token_ids[:token_ids.index(self.sep_token)]
text = self.tokenizer.decode(token_ids, skip_special_tokens=True)
return text

12
modelscope/pipelines/base.py
View File

@@ -411,7 +411,7 @@ class DistributedPipeline(Pipeline):
else:
self.model_dir = snapshot_download(model)
self.cfg = read_config(self.model_dir)
self.world_size = self.cfg.model.world_size
self.world_size = self._get_world_size(self.cfg)
self.model_pool = None
self.device_name = 'cpu'
self.device = create_device(self.device_name)
@@ -423,17 +423,17 @@ class DistributedPipeline(Pipeline):
self.model_pool = Pool(self.world_size)
master_ip = '127.0.0.1' if 'master_ip' not in kwargs else kwargs[
'master_ip']
os.environ['MASTER_ADDR'] = master_ip
master_port = '29500' if 'master_port' not in kwargs else kwargs[
'master_port']
from modelscope.utils.torch_utils import _find_free_port, _is_free_port
if not _is_free_port(int(master_port)):
master_port = str(_find_free_port())
os.environ['MASTER_PORT'] = master_port
self.model_pool.map(
partial(
self.__class__._instantiate_one,
model_dir=self.model_dir,
master_ip=master_ip,
master_port=master_port,
**self.cfg.model,
**kwargs), ranks)
self.models = []
@@ -487,6 +487,12 @@ class DistributedPipeline(Pipeline):
"""
pass
def _get_world_size(self, cfg: Config) -> int:
m_world_size = cfg.safe_get('megatron.world_size')
if m_world_size is None:
return cfg.safe_get('model.world_size')
return m_world_size
def collate_fn(data, device):
"""Prepare the input just before the forward function.

3
modelscope/pipelines/nlp/mglm_text_summarization_pipeline.py
View File

@@ -1,5 +1,6 @@
# Copyright (c) 2022 Zhipu.AI
import os
from typing import Any, Dict, Optional, Union
from modelscope.metainfo import Pipelines
@@ -30,6 +31,8 @@ class MGLMTextSummarizationPipeline(Pipeline):
self.model.eval()
if preprocessor is None:
preprocessor = MGLMSummarizationPreprocessor()
from modelscope.utils.torch_utils import _find_free_port
os.environ['MASTER_PORT'] = str(_find_free_port())
super().__init__(model=model, preprocessor=preprocessor, **kwargs)
# define the forward pass

60
modelscope/preprocessors/nlp/text_generation_preprocessor.py
View File

@@ -2,7 +2,7 @@
import os
import os.path as osp
from typing import Any, Dict, Optional, Union
from typing import Any, Dict, List, Optional, Union
import numpy as np
import torch
@@ -223,14 +223,6 @@ class TextGenerationJiebaPreprocessor(TextGenerationPreprocessorBase):
"""
return self.tokenizer.detokenize(tokens)
def _truncate(self, array: np.ndarray) -> np.ndarray:
if len(array) < self.max_length:
return np.pad(
array, (0, self.max_length - len(array)),
constant_values=self.tokenizer.eod)
else:
return array[:self.max_length]
def _tokenize_text(self, sequence1, sequence2=None, **kwargs):
"""Tokenize the text.
@@ -246,18 +238,46 @@ class TextGenerationJiebaPreprocessor(TextGenerationPreprocessorBase):
'input_ids':
torch.tensor(self.tokenizer.tokenize(sequence1)).unsqueeze_(0)
}
# continue write train: | inputs | <sep> |
# input & output train: | inputs | outputs | <sep> |
else:
tokens = self.tokenizer.tokenize(sequence1)
prompt_length = min(len(tokens), self.max_length - 1)
if sequence2 is not None:
tokens += self.tokenizer.tokenize(sequence2)
tokens = self._truncate(np.array(tokens))
return {
'tokens': tokens[:-1],
'labels': tokens[1:],
'prompt_length': prompt_length,
'is_pair': int(sequence2 is not None),
}
input_tokens = self.tokenizer.tokenize(sequence1)
if sequence2 is None:
return self._only_input(input_tokens)
else:
return self._input_and_output(
input_tokens, self.tokenizer.tokenize(sequence2))
def _only_input(self, input_tokens: List[int]) -> Dict[str, Any]:
prompts_len = len(input_tokens)
input_tokens.append(self.tokenizer.sep_token)
tokens = self._truncate(np.asarray(input_tokens))
return {
'tokens': tokens[:-1],
'labels': tokens[1:],
'prompts_len': min(prompts_len, self.max_length),
}
def _input_and_output(self, input_tokens: List[int],
output_tokens: List[int]) -> Dict[str, Any]:
tokens = input_tokens[:]
tokens.extend(output_tokens)
tokens.append(self.tokenizer.sep_token)
inputs_len = len(tokens)
tokens = self._truncate(np.asarray(tokens))
return {
'tokens': tokens[:-1],
'labels': tokens[1:],
'prompts_len': min(len(input_tokens), self.max_length),
'inputs_len': min(inputs_len, self.max_length),
}
def _truncate(self, array: np.ndarray) -> np.ndarray:
if len(array) < self.max_length:
return np.pad(
array, (0, self.max_length - len(array)), constant_values=0)
else:
return array[:self.max_length]
@PREPROCESSORS.register_module(

1
modelscope/trainers/hooks/checkpoint_hook.py
View File

@@ -14,6 +14,7 @@ from modelscope.utils.checkpoint import (load_checkpoint, save_checkpoint,
save_configuration)
from modelscope.utils.constant import LogKeys, ModelFile
from modelscope.utils.logger import get_logger
from modelscope.utils.megatron_utils import is_megatron_initialized
from modelscope.utils.torch_utils import is_master
from .builder import HOOKS
from .hook import Hook

6
modelscope/trainers/hooks/logger/tensorboard_hook.py
View File

@@ -41,7 +41,7 @@ class TensorboardHook(LoggerHook):
self.out_dir = out_dir
self.skip_keys = skip_keys
@master_only
@master_only()
def before_run(self, trainer):
super(TensorboardHook, self).before_run(trainer)
try:
@@ -59,7 +59,7 @@ class TensorboardHook(LoggerHook):
f'tensorboard files will be saved to {self.out_dir}')
self.writer = SummaryWriter(self.out_dir)
@master_only
@master_only()
def log(self, trainer):
if len(trainer.visualization_buffer.output) > 0:
self.visualization_log(trainer)
@@ -112,6 +112,6 @@ class TensorboardHook(LoggerHook):
# avoiding repeated writing of the same image buffer every self.interval
trainer.visualization_buffer.clear_output()
@master_only
@master_only()
def after_run(self, trainer):
self.writer.close()

26
modelscope/trainers/nlp/gpt3_trainer.py
View File

@@ -1,14 +1,9 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from collections.abc import Mapping
from typing import Any, Dict, List
import torch
from megatron_util import mpu
from modelscope.metainfo import Trainers
from modelscope.models import TorchModel
from modelscope.models.nlp import GPT3ForTextGeneration
from modelscope.trainers.builder import TRAINERS
from modelscope.trainers.nlp_trainer import NlpEpochBasedTrainer
@@ -23,12 +18,6 @@ class GPT3Trainer(NlpEpochBasedTrainer):
cfg.model.rank = int(os.environ.get('RANK', 0))
return cfg
def train_step(self, model: TorchModel, inputs: Mapping):
keys = list(inputs.keys())
datatype = torch.int64
inputs = mpu.broadcast_data(keys, inputs, datatype)
return super().train_step(model, inputs)
def _decode(self, tokens):
tokenizer = self.eval_preprocessor.tokenizer
return tokenizer.detokenize(tokens.tolist())
@@ -37,28 +26,25 @@ class GPT3Trainer(NlpEpochBasedTrainer):
model = self.model.module if self._dist else self.model
model.eval()
if self._is_pair(data):
if 'inputs_len' in data:
return self._generate_eval(model, data)
else:
return self._forward_eval(model, data)
@staticmethod
def _is_pair(data: Dict[str, Any]) -> bool:
return 'is_pair' in data and bool(data['is_pair'][0])
def _generate_eval(self, model: GPT3ForTextGeneration,
data: Dict[str, Any]) -> Dict[str, Any]:
data['do_sample'] = False
# Force greedy decoding in non-open tasks
data.update(top_k=1, top_p=0.)
result = model.generate(data)
prompt_length: List[int] = data['prompt_length']
prompts_len: List[int] = data['prompts_len']
result['preds'] = [
self._decode(seq[skip_len:])
for seq, skip_len in zip(result['sequences'], prompt_length)
for seq, skip_len in zip(result['sequences'], prompts_len)
]
data['tgts'] = [
self._decode(seq[skip_len - 1:])
for seq, skip_len in zip(data['labels'], prompt_length)
for seq, skip_len in zip(data['labels'], prompts_len)
]
return result

8
modelscope/trainers/trainer.py
View File

@@ -38,6 +38,7 @@ from modelscope.utils.data_utils import to_device
from modelscope.utils.device import create_device
from modelscope.utils.file_utils import func_receive_dict_inputs
from modelscope.utils.logger import get_logger
from modelscope.utils.megatron_utils import is_megatron_initialized
from modelscope.utils.registry import build_from_cfg
from modelscope.utils.torch_utils import (get_dist_info, get_local_rank,
init_dist, is_dist, is_master,
@@ -683,6 +684,13 @@ class EpochBasedTrainer(BaseTrainer):
find_unused_parameters=True,
device_ids=[torch.cuda.current_device()])
if is_megatron_initialized():
from megatron_util import mpu
dp_cfg.update({
'output_device': torch.cuda.current_device(),
'process_group': mpu.get_data_parallel_group()
})
return build_parallel(dp_cfg)
def unwrap_module(self, model) -> Union[nn.Module, TorchModel]:

29
modelscope/trainers/utils/inference.py
View File

@@ -10,7 +10,9 @@ import torch
from torch import distributed as dist
from tqdm import tqdm
from modelscope.utils.constant import DistributedParallelType
from modelscope.utils.data_utils import to_device
from modelscope.utils.megatron_utils import is_megatron_initialized
from modelscope.utils.torch_utils import (broadcast, get_dist_info, is_master,
make_tmp_dir)
@@ -187,9 +189,10 @@ def evaluate_batch(trainer, data, metric_classes, vis_closure):
def get_metric_values(metric_classes):
rank, world_size = get_dist_info()
_, world_size = get_dist_info()
metric_values = {}
if rank == 0:
if is_master(
DistributedParallelType.DP if is_megatron_initialized() else None):
for metric_cls in metric_classes:
metric_values.update(metric_cls.evaluate())
if world_size > 1:
@@ -218,16 +221,17 @@ def collect_results_cpu(result_part, tmpdir=None):
rank, world_size = get_dist_info()
if tmpdir is None:
tmpdir = make_tmp_dir()
if not os.path.exists(tmpdir) and is_master():
if not os.path.exists(tmpdir) and is_master(DistributedParallelType.TP):
os.makedirs(tmpdir)
dist.barrier()
# dump the part result to the dir
with open(os.path.join(tmpdir, f'part_{rank}.pkl'), 'wb') as f:
pickle.dump(result_part, f)
if is_master(DistributedParallelType.TP):
with open(os.path.join(tmpdir, f'part_{rank}.pkl'), 'wb') as f:
pickle.dump(result_part, f)
dist.barrier()
# collect all parts
if rank != 0:
if not is_master():
return None
else:
# load results of all parts from tmp dir
@@ -261,14 +265,19 @@ def collect_results_gpu(result_part):
Returns:
list: The collected results.
"""
rank, world_size = get_dist_info()
_, world_size = get_dist_info()
group = None
if is_megatron_initialized():
from megatron_util import mpu
group = mpu.get_data_parallel_group()
# dump result part to tensor with pickle
part_tensor = torch.tensor(
bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
# gather all result part tensor shape
shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
shape_list = [shape_tensor.clone() for _ in range(world_size)]
dist.all_gather(shape_list, shape_tensor)
dist.all_gather(shape_list, shape_tensor, group)
# padding result part tensor to max length
shape_max = torch.tensor(shape_list).max()
part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
@@ -277,9 +286,9 @@ def collect_results_gpu(result_part):
part_tensor.new_zeros(shape_max) for _ in range(world_size)
]
# gather all result part
dist.all_gather(part_recv_list, part_send)
dist.all_gather(part_recv_list, part_send, group)
if rank == 0:
if is_master():
part_list = []
for recv, shape in zip(part_recv_list, shape_list):
part_result = pickle.loads(recv[:shape[0]].cpu().numpy().tobytes())

4
modelscope/utils/checkpoint.py
View File

@@ -20,6 +20,7 @@ from modelscope.fileio import File, LocalStorage
from modelscope.utils.config import Config, JSONIteratorEncoder
from modelscope.utils.constant import ConfigFields, ModelFile
from modelscope.utils.logger import get_logger
from modelscope.utils.torch_utils import is_master
logger = get_logger()
@@ -585,7 +586,8 @@ def save_pretrained(model,
ignore_file_set = set(origin_file_to_be_ignored)
ignore_file_set.add(ModelFile.CONFIGURATION)
ignore_file_set.add('.*')
if hasattr(model, 'model_dir') and model.model_dir is not None:
if hasattr(model,
'model_dir') and model.model_dir is not None and is_master():
copytree(
model.model_dir,
target_folder,

7
modelscope/utils/constant.py
View File

@@ -509,3 +509,10 @@ class MetaDataFields:
DatasetVisibilityMap = {1: 'private', 3: 'internal', 5: 'public'}
class DistributedParallelType(object):
"""Parallel Strategies for Distributed Models"""
DP = 'data_parallel'
TP = 'tensor_model_parallel'
PP = 'pipeline_model_parallel'

22
modelscope/utils/megatron_utils.py
View File

@@ -1,12 +1,5 @@
# Copyright (c) Alibaba, Inc. and its affiliates.
from typing import Optional
from megatron_util import initialize_megatron
from modelscope.utils.config import Config
from modelscope.utils.hub import read_config
_DEFAULT_CFG_WITH_MODEL_TYPE = {
'gpt-moe': {
'version': 'moe',
@@ -24,10 +17,13 @@ _DEFAULT_CFG_WITH_MODEL_TYPE = {
},
}
_IS_MEGATRON_INITIALIZED = False
def init_megatron_util(cfg=None, model_dir=None, **kwargs):
from modelscope.utils.hub import read_config
from megatron_util import initialize_megatron
def init_megatron_util(cfg: Optional[Config] = None,
model_dir: Optional[str] = None,
**kwargs):
assert not (cfg is None and model_dir is None), \
'cfg and model_dir cannot both be None when initializing megatron_util'
if cfg is None:
@@ -44,3 +40,9 @@ def init_megatron_util(cfg: Optional[Config] = None,
if model_type in _DEFAULT_CFG_WITH_MODEL_TYPE else {}
megatron_cfg.update(kwargs)
initialize_megatron(megatron_cfg)
global _IS_MEGATRON_INITIALIZED
_IS_MEGATRON_INITIALIZED = True
def is_megatron_initialized() -> bool:
return _IS_MEGATRON_INITIALIZED

2
modelscope/utils/nlp/load_checkpoint.py
View File

@@ -66,7 +66,7 @@ def pre_load(mp_rank, load_dir, tag=''):
load_path = _get_ckpt_name(mp_rank, load_dir, tag)
checkpoint = torch.load(
load_path, map_location=lambda storage, loc: storage)
return checkpoint['module']
return checkpoint['module'] if 'module' in checkpoint else checkpoint
def _load_checkpoint(model,

50
modelscope/utils/torch_utils.py
View File

@@ -14,6 +14,9 @@ import torch
import torch.multiprocessing as mp
from torch import distributed as dist
from modelscope.utils.constant import DistributedParallelType
from modelscope.utils.megatron_utils import is_megatron_initialized
def _find_free_port() -> str:
# Copied from https://github.com/facebookresearch/detectron2/blob/main/detectron2/engine/launch.py # noqa: E501
@@ -107,14 +110,12 @@ def _init_dist_slurm(backend: str, port: Optional[int] = None) -> None:
def get_dist_info() -> Tuple[int, int]:
if is_dist():
try:
group = None
if is_megatron_initialized():
from megatron_util import mpu
assert mpu.model_parallel_is_initialized()
rank = mpu.get_data_parallel_rank()
world_size = mpu.get_data_parallel_world_size()
except (ImportError, AssertionError):
rank = dist.get_rank()
world_size = dist.get_world_size()
group = mpu.get_data_parallel_group()
rank = dist.get_rank(group)
world_size = dist.get_world_size(group)
else:
rank = 0
world_size = 1
@@ -160,18 +161,37 @@ def is_dist():
return dist.is_available() and dist.is_initialized()
def is_master():
return dist.get_rank() == 0 if is_dist() else True
def is_master(group=None):
if isinstance(group, str):
group = _parse_parallel_group(group)
return dist.get_rank(group) == 0 if is_dist() else True
def master_only(func: Callable) -> Callable:
def _parse_parallel_group(group: str):
from megatron_util import mpu
if group == DistributedParallelType.DP:
return mpu.get_data_parallel_group()
if group == DistributedParallelType.TP:
return mpu.get_tensor_model_parallel_group()
if group == DistributedParallelType.PP:
return mpu.get_pipeline_model_parallel_group()
raise ValueError(
f"Wrong group '{group}'. Supported groups are '{DistributedParallelType.DP}', "
f"'{DistributedParallelType.TP}' or '{DistributedParallelType.PP}'")
@functools.wraps(func)
def wrapper(*args, **kwargs):
if is_master():
return func(*args, **kwargs)
return wrapper
def master_only(group=None):
def decorate(func: Callable) -> Callable:
@functools.wraps(func)
def wrapper(*args, **kwargs):
if is_master(group):
return func(*args, **kwargs)
return wrapper
return decorate
def make_tmp_dir():

122
tests/trainers/test_finetune_gpt3.py
View File

@@ -7,9 +7,10 @@ import unittest
from modelscope.metainfo import Trainers
from modelscope.msdatasets import MsDataset
from modelscope.trainers import build_trainer
from modelscope.utils.test_utils import DistributedTestCase, test_level
class TestFinetuneTextGeneration(unittest.TestCase):
class TestFinetuneTextGeneration(DistributedTestCase):
def setUp(self):
print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
@@ -21,62 +22,6 @@ class TestFinetuneTextGeneration(unittest.TestCase):
shutil.rmtree(self.tmp_dir)
super().tearDown()
@unittest.skip(
'skip since the test requires multiple GPU and takes a long time to run'
)
def test_finetune_poetry(self):
dataset_dict = MsDataset.load('chinese-poetry-collection')
train_dataset = dataset_dict['train'].remap_columns(
{'text1': 'src_txt'})
eval_dataset = dataset_dict['test'].remap_columns({'text1': 'src_txt'})
max_epochs = 10
tmp_dir = './gpt3_poetry'
num_warmup_steps = 100
def noam_lambda(current_step: int):
current_step += 1
return min(current_step**(-0.5),
current_step * num_warmup_steps**(-1.5))
def cfg_modify_fn(cfg):
cfg.train.lr_scheduler = {
'type': 'LambdaLR',
'lr_lambda': noam_lambda,
'options': {
'by_epoch': False
}
}
cfg.train.optimizer = {'type': 'AdamW', 'lr': 3e-4}
cfg.train.dataloader = {
'batch_size_per_gpu': 16,
'workers_per_gpu': 1
}
cfg.train.hooks.append({
'type': 'EvaluationHook',
'by_epoch': True,
'interval': 1
})
cfg.evaluation.dataloader = {
'batch_size_per_gpu': 8,
'workers_per_gpu': 1
}
cfg.evaluation.metrics = 'ppl'
return cfg
kwargs = dict(
model='damo/nlp_gpt3_text-generation_1.3B',
train_dataset=train_dataset,
eval_dataset=eval_dataset,
max_epochs=max_epochs,
work_dir=tmp_dir,
cfg_modify_fn=cfg_modify_fn)
# Construct trainer and train
trainer = build_trainer(
name=Trainers.gpt3_trainer, default_args=kwargs)
trainer.train()
@unittest.skip(
'skip since the test requires multiple GPU and takes a long time to run'
)
@@ -137,6 +82,69 @@ class TestFinetuneTextGeneration(unittest.TestCase):
name=Trainers.gpt3_trainer, default_args=kwargs)
trainer.train()
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
def test_single_finetune_portry(self):
finetune_poetry()
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
def test_multi_finetune_portry(self):
self.start(
finetune_poetry, num_gpus=4, work_dir=self.tmp_dir, dp_tp=True)
# TODO: add gpt3 trainer predict unittest
def finetune_poetry(dp_tp=False):
dataset_dict = MsDataset.load('chinese-poetry-collection')
train_dataset = dataset_dict['train'].remap_columns({'text1': 'src_txt'})
eval_dataset = dataset_dict['test'].remap_columns({'text1': 'src_txt'})
max_epochs = 2
tmp_dir = './gpt3_poetry'
num_warmup_steps = 100
def noam_lambda(current_step: int):
current_step += 1
return min(current_step**(-0.5),
current_step * num_warmup_steps**(-1.5))
def cfg_modify_fn(cfg):
cfg.train.lr_scheduler = {
'type': 'LambdaLR',
'lr_lambda': noam_lambda,
'options': {
'by_epoch': False
}
}
cfg.train.optimizer = {'type': 'AdamW', 'lr': 3e-4}
cfg.train.dataloader = {'batch_size_per_gpu': 2, 'workers_per_gpu': 1}
cfg.train.hooks.append({
'type': 'EvaluationHook',
'by_epoch': True,
'interval': 1
})
cfg.evaluation.dataloader = {
'batch_size_per_gpu': 8,
'workers_per_gpu': 1
}
cfg.evaluation.metrics = 'ppl'
cfg.train.train_iters_per_epoch = 10
if dp_tp:
cfg.megatron = {'world_size': 4, 'tensor_model_parallel_size': 2}
return cfg
kwargs = dict(
model='damo/nlp_gpt3_text-generation_1.3B',
train_dataset=train_dataset,
eval_dataset=eval_dataset,
max_epochs=max_epochs,
work_dir=tmp_dir,
cfg_modify_fn=cfg_modify_fn)
# Construct trainer and train
trainer = build_trainer(name=Trainers.gpt3_trainer, default_args=kwargs)
trainer.train()
if __name__ == '__main__':
unittest.main()