初始化项目，由ModelHub XC社区提供模型
Model: Tencent-Hunyuan/Hy-MT2-7B-GGUF Source: Original Platform
This commit is contained in:
ModelHub XC
360
train/deepspeed_support/train_dense.py
Normal file
360
train/deepspeed_support/train_dense.py
Normal file
@@ -0,0 +1,360 @@
|
||||
# Copyright 2024 Tencent Inc. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
|
||||
# and OPT implementations in this library. It has been modified from its
|
||||
# original forms to accommodate minor architectural differences compared
|
||||
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""
|
||||
Training script for HunYuan Dense models (1.8B, 7B).
|
||||
|
||||
This script is adapted from the original finetune.py for dense models,
|
||||
with improvements from the new training framework (train.py for MoE models).
|
||||
|
||||
Key differences from train.py (MoE version):
|
||||
- No MoE-related patches (router dtype fix, expert key rename, etc.)
|
||||
- Supports model_size parameter to handle different tokenizer formats
|
||||
- 7B model uses different special tokens than 1.8B model
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
import json
|
||||
import torch
|
||||
import shutil
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Dict, Literal
|
||||
|
||||
import transformers
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import Trainer, TrainerCallback
|
||||
from peft import LoraConfig, get_peft_model, PeftModel
|
||||
from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR
|
||||
from transformers.modeling_utils import unwrap_model
|
||||
|
||||
|
||||
def print_args(args, name='arguments'):
|
||||
"""Print arguments."""
|
||||
if torch.distributed.get_rank() == 0:
|
||||
print(f'------------------------ {name} ------------------------', flush=True)
|
||||
str_list = []
|
||||
for arg in vars(args):
|
||||
dots = '.' * (48 - len(arg))
|
||||
str_list.append(' {} {} {}'.format(arg, dots, getattr(args, arg)))
|
||||
for arg in sorted(str_list, key=lambda x: x.lower()):
|
||||
print(arg, flush=True)
|
||||
print(f'-------------------- end of {name} ---------------------', flush=True)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
use_flash_attn: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Enable FlashAttention-2 for faster training."}
|
||||
)
|
||||
use_lora: bool = field(default=False, metadata={"help": "Enable Lora for faster training."})
|
||||
hidden_size: int = field(default=2048, metadata={"help": "The hidden size of the model."})
|
||||
num_layers: int = field(default=32, metadata={"help": "The number of layers of the model."})
|
||||
num_attention_heads: int = field(default=16, metadata={"help": "The number of attention heads of the model."})
|
||||
intermediate_size: int = field(default=6144, metadata={"help": "The intermediate size of the model."})
|
||||
num_key_value_heads: int = field(default=4, metadata={"help": "The number of key-value heads in GQA."})
|
||||
use_qk_norm: bool = field(default=False, metadata={"help": "Whether to use qk norm."})
|
||||
tie_word_embeddings: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "Whether to tie the word embeddings of the encoder and the decoder."}
|
||||
)
|
||||
lora_rank: int = field(default=64, metadata={"help": "The rank of lora."})
|
||||
lora_alpha: int = field(default=128, metadata={"help": "Lora alpha"})
|
||||
lora_dropout: float = field(default=0.0, metadata={"help": "Lora dropout"})
|
||||
train_attention_params_only: bool = field(default=False, metadata={
|
||||
"help": "Whether to train attention parameters only."}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataArguments:
|
||||
model_size: Literal["0.5B", "1.8B", "4B", "7B"] = field(
|
||||
default="1.8B",
|
||||
metadata={"help": "Select the model size from ['0.5B', '1.8B', '4B', '7B']. "
|
||||
"This affects the tokenizer special tokens used for loss masking."}
|
||||
)
|
||||
train_data_file: str = field(default=None, metadata={"help": "Path to the training data."})
|
||||
max_seq_length: int = field(
|
||||
default=4096,
|
||||
metadata={"help": "The max sequence length of the model inputs after tokenization."}
|
||||
)
|
||||
use_dummy_data: bool = field(default=False, metadata={"help": "Use dummy data."})
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingArguments(transformers.TrainingArguments):
|
||||
cache_dir: Optional[str] = field(default=None)
|
||||
optim: str = field(default="adamw_torch")
|
||||
model_max_length: int = field(
|
||||
default=4096,
|
||||
metadata={"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."},
|
||||
)
|
||||
tokenizer_name_or_path: Optional[str] = field(default=None)
|
||||
model_name_or_path: Optional[str] = field(default=None)
|
||||
min_lr: float = field(
|
||||
default=1e-6,
|
||||
metadata={"help": "The minimum learning rate at the end of the cosine decay."}
|
||||
)
|
||||
|
||||
|
||||
IGNORE_INDEX = -100
|
||||
|
||||
|
||||
class DummyDataset(Dataset):
|
||||
def __init__(self, tokenizer, max_seq_length=512, length=1000):
|
||||
self.tokenizer = tokenizer
|
||||
self.max_seq_length = max_seq_length
|
||||
self.length = length
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
|
||||
def __getitem__(self, index):
|
||||
tokens = torch.randint(0, self.tokenizer.vocab_size, (self.max_seq_length,))
|
||||
return {'input_ids': tokens, 'labels': tokens}
|
||||
|
||||
|
||||
class SFTDataset(Dataset):
|
||||
def __init__(self, data_file, tokenizer, max_seq_length=4096, model_size="1.8B"):
|
||||
self.tokenizer = tokenizer
|
||||
self.max_seq_length = max_seq_length
|
||||
self.model_size = model_size
|
||||
self.data_list = self.load_data(data_file)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data_list)
|
||||
|
||||
def load_data(self, data_file):
|
||||
logging.info('Loading data: {}'.format(data_file))
|
||||
with open(data_file, 'r', encoding='utf8') as f:
|
||||
data_list = f.readlines()
|
||||
logging.info("there are {} data in dataset".format(len(data_list)))
|
||||
return data_list
|
||||
|
||||
def encode_data(self, data_dict):
|
||||
model_inputs = {}
|
||||
template_output = self.tokenizer.apply_chat_template(
|
||||
data_dict['messages'], tokenize=True, return_dict=False
|
||||
)
|
||||
if isinstance(template_output, list) and len(template_output) > 0 and isinstance(template_output[0], list):
|
||||
template_output = template_output[0]
|
||||
message_tokens = torch.tensor(template_output, dtype=torch.long)
|
||||
|
||||
# Note: The 7B model uses a different vocabulary/special tokens than other models.
|
||||
if self.model_size == "7B":
|
||||
sep_token_id = self.tokenizer.convert_tokens_to_ids('<|extra_0|>')
|
||||
eos_token_id = self.tokenizer.convert_tokens_to_ids('<|eos|>')
|
||||
else:
|
||||
sep_token_id = self.tokenizer.convert_tokens_to_ids('<｜hy_Assistant｜>')
|
||||
eos_token_id = self.tokenizer.convert_tokens_to_ids('<｜hy_place▁holder▁no▁2｜>')
|
||||
|
||||
# Find assistant reply boundaries
|
||||
loss_token_begins = (message_tokens == sep_token_id).nonzero(as_tuple=True)[0].tolist()
|
||||
loss_token_ends = (message_tokens == eos_token_id).nonzero(as_tuple=True)[0].tolist()
|
||||
message_labels = torch.tensor([IGNORE_INDEX] * message_tokens.shape[0])
|
||||
for begin_idx, end_idx in zip(loss_token_begins, loss_token_ends):
|
||||
# Compute loss from sep_token to eos_token (inclusive)
|
||||
message_labels[begin_idx:end_idx + 1] = message_tokens[begin_idx:end_idx + 1]
|
||||
|
||||
input_ids = message_tokens.to(torch.long)
|
||||
labels = message_labels.to(torch.long)
|
||||
|
||||
input_ids = input_ids[:self.max_seq_length]
|
||||
labels = labels[:self.max_seq_length]
|
||||
|
||||
pad_token_id = self.tokenizer.pad_token_id
|
||||
attention_mask = [1 if val != pad_token_id else 0 for val in input_ids]
|
||||
model_inputs["input_ids"] = input_ids
|
||||
model_inputs["attention_mask"] = torch.tensor(attention_mask, dtype=torch.bool)
|
||||
model_inputs["labels"] = labels
|
||||
|
||||
return model_inputs
|
||||
|
||||
def __getitem__(self, index):
|
||||
data = self.data_list[index]
|
||||
data = json.loads(data)
|
||||
model_inputs = self.encode_data(data)
|
||||
return model_inputs
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataCollatorForSupervisedDataset(object):
|
||||
"""Collate examples for supervised fine-tuning."""
|
||||
|
||||
tokenizer: transformers.PreTrainedTokenizer
|
||||
|
||||
def __call__(self, instances):
|
||||
input_ids = [instance['input_ids'] for instance in instances]
|
||||
labels = [instance['labels'] for instance in instances]
|
||||
pad_token_id = self.tokenizer.pad_token_id
|
||||
input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=pad_token_id)
|
||||
labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)
|
||||
return dict(
|
||||
input_ids=input_ids,
|
||||
labels=labels,
|
||||
attention_mask=input_ids.ne(pad_token_id),
|
||||
)
|
||||
|
||||
|
||||
def make_supervised_data_module(tokenizer, data_args) -> Dict:
|
||||
"""Make dataset and collator for supervised fine-tuning."""
|
||||
if data_args.use_dummy_data:
|
||||
train_dataset = DummyDataset(tokenizer, data_args.max_seq_length)
|
||||
else:
|
||||
train_dataset = SFTDataset(
|
||||
tokenizer=tokenizer,
|
||||
data_file=data_args.train_data_file,
|
||||
max_seq_length=data_args.max_seq_length,
|
||||
model_size=data_args.model_size,
|
||||
)
|
||||
data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
|
||||
return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
|
||||
|
||||
|
||||
# Copy tokenizer and config files to each checkpoint directory for self-contained inference
|
||||
class CustomSaveCallback(TrainerCallback):
|
||||
def on_save(self, args, state, control, **kwargs):
|
||||
if torch.distributed.get_rank() == 0:
|
||||
output_dir = os.path.join(args.output_dir, f"{PREFIX_CHECKPOINT_DIR}-{state.global_step}")
|
||||
|
||||
# Copy tokenizer files to checkpoint directory
|
||||
tokenizer_files = [
|
||||
'generation_config.json',
|
||||
'hy.tiktoken',
|
||||
'tokenizer_config.json',
|
||||
'tokenization_hy.py',
|
||||
'tokenizer.json',
|
||||
'special_tokens_map.json',
|
||||
'chat_template.jinja',
|
||||
'config.json',
|
||||
]
|
||||
src_dir = args.tokenizer_name_or_path or args.model_name_or_path
|
||||
for fname in tokenizer_files:
|
||||
src = os.path.join(src_dir, fname)
|
||||
if os.path.isfile(src):
|
||||
shutil.copy(src, os.path.join(output_dir, fname))
|
||||
|
||||
return control
|
||||
|
||||
|
||||
def train():
|
||||
parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
print_args(model_args, 'model arguments')
|
||||
print_args(data_args, 'data arguments')
|
||||
print_args(training_args, 'training arguments')
|
||||
|
||||
tokenizer = transformers.AutoTokenizer.from_pretrained(
|
||||
training_args.tokenizer_name_or_path,
|
||||
trust_remote_code=True
|
||||
)
|
||||
|
||||
init_kwargs = {}
|
||||
if model_args.use_flash_attn:
|
||||
init_kwargs["attn_implementation"] = "flash_attention_2"
|
||||
if training_args.bf16:
|
||||
init_kwargs["torch_dtype"] = torch.bfloat16
|
||||
elif training_args.fp16:
|
||||
init_kwargs["torch_dtype"] = torch.float16
|
||||
|
||||
# Load model from pretrained weights
|
||||
if training_args.model_name_or_path is not None and os.path.exists(training_args.model_name_or_path):
|
||||
print(f"Initializing model from local file: {training_args.model_name_or_path}")
|
||||
model = transformers.AutoModelForCausalLM.from_pretrained(
|
||||
training_args.model_name_or_path,
|
||||
trust_remote_code=True,
|
||||
**init_kwargs
|
||||
)
|
||||
else:
|
||||
raise FileNotFoundError(
|
||||
f"Model path {training_args.model_name_or_path} is invalid or does not exist. "
|
||||
f"Dense model training requires pre-trained weights."
|
||||
)
|
||||
|
||||
if model_args.train_attention_params_only:
|
||||
for name, param in model.named_parameters():
|
||||
if 'self_attn' not in name:
|
||||
param.requires_grad = False
|
||||
|
||||
if model_args.use_lora:
|
||||
# Define LoRA configuration
|
||||
lora_config = LoraConfig(
|
||||
r=model_args.lora_rank,
|
||||
lora_alpha=model_args.lora_alpha,
|
||||
lora_dropout=model_args.lora_dropout,
|
||||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
)
|
||||
model = get_peft_model(model, lora_config)
|
||||
|
||||
data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
|
||||
# Tell Trainer not to attempt DataParallel
|
||||
model.is_parallelizable = True
|
||||
model.model_parallel = True
|
||||
|
||||
training_args.lr_scheduler_kwargs = {
|
||||
'min_lr_rate': training_args.min_lr / training_args.learning_rate,
|
||||
}
|
||||
|
||||
# -----------------------------------------------------------------------
|
||||
# Fix: DeepSpeed ZeRO-3 + gradient checkpointing compatibility.
|
||||
#
|
||||
# PyTorch's torch.utils.checkpoint with use_reentrant=False (the default
|
||||
# in transformers) performs strict metadata checks on recomputed tensors
|
||||
# during backward. Under ZeRO-3, parameters are all-gathered during the
|
||||
# first forward pass (shape=[full_size]) but may be partitioned back
|
||||
# (shape=[0]) when the checkpoint recomputes, causing a CheckpointError.
|
||||
#
|
||||
# Setting use_reentrant=True avoids this strict metadata check.
|
||||
# -----------------------------------------------------------------------
|
||||
if training_args.gradient_checkpointing and training_args.deepspeed:
|
||||
training_args.gradient_checkpointing_kwargs = {"use_reentrant": True}
|
||||
|
||||
trainer = Trainer(
|
||||
model=model,
|
||||
processing_class=tokenizer,
|
||||
args=training_args,
|
||||
callbacks=[CustomSaveCallback],
|
||||
**data_module
|
||||
)
|
||||
model.config.use_cache = False
|
||||
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
train()
|
||||
Reference in New Issue
Block a user