File size: 9,335 Bytes
b4942cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 |
import json
import os
import pathlib
import deepspeed
import torch
import transformers
from deepspeed import get_accelerator
from torch.utils.data import ConcatDataset
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModel, AutoConfig
from transformers import Trainer
from transformers.integrations.deepspeed import unset_hf_deepspeed_config, set_hf_deepspeed_config
from callback import TuneTauCallback, MonitorCallback
from ovis.model.configuration_ovis import OvisConfig
from ovis.model.modeling_ovis import Ovis
from ovis.train.arguments import ModelArguments, TrainingArguments
from ovis.train.dataset.caption_dataset import CaptionDataset
from ovis.train.dataset.conversation_dataset import ConversationDataset
from ovis.train.dataset.multimodal_dataset import DataCollatorForMultimodalDataset
from ovis.util.constants import BEGIN_LINE, END_LINE
from ovis.util.utils import smart_unit, rank0_print
def train():
# parse args
parser = transformers.HfArgumentParser(
(ModelArguments, TrainingArguments))
model_args, training_args = parser.parse_args_into_dataclasses()
# save args to checkpoint dir
with training_args.main_process_first(local=False):
if training_args.process_index == 0:
def args2dict(args):
return {k: str(v) for k, v in args.__dict__.items()}
args_log = json.dumps(dict(
model_args=args2dict(model_args),
training_args=args2dict(training_args)
), ensure_ascii=False, indent=2)
print(args_log)
os.makedirs(training_args.output_dir, exist_ok=True)
with open(os.path.join(training_args.output_dir, 'model_training_args.json'), 'w',
encoding='utf-8') as f:
f.write(args_log + '\n')
# construct or load ovis model
if not training_args.ovis_pretrained_path: # construct model (S1)
# 1. construct ovis config
ovis_config = OvisConfig(
multimodal_max_length=model_args.multimodal_max_length,
conversation_formatter_class=model_args.conversation_formatter_class,
llm_attn_implementation=model_args.llm_attn_implementation
)
# 2. load pretrained llm and text tokenizer
attn_kwargs = dict()
if model_args.llm_attn_implementation:
attn_kwargs['attn_implementation'] = model_args.llm_attn_implementation
llm = AutoModelForCausalLM.from_pretrained(model_args.llm_name_or_path, **attn_kwargs)
text_tokenizer = AutoTokenizer.from_pretrained(model_args.llm_name_or_path)
if text_tokenizer.pad_token_id is None and model_args.pad_token_id is not None:
text_tokenizer.pad_token_id = model_args.pad_token_id
# 3. construct visual tokenizer
# deepspeed zero.Init with bfloat16 fail for visual_tokenizer, so temporarily disable zero.Init here
unset_hf_deepspeed_config()
if training_args.visual_tokenizer_pretrained_path is not None:
visual_tokenizer = AutoModel.from_pretrained(
training_args.visual_tokenizer_pretrained_path,
image_processor_name_or_path=training_args.visual_tokenizer_pretrained_path
)
else:
visual_tokenizer_config = AutoConfig.for_model(
model_type=model_args.visual_tokenizer_type + "_visual_tokenizer",
vocab_size=model_args.visual_vocab_size,
tokenize_function=model_args.visual_tokenize_function,
tau=model_args.visual_tau,
depths=model_args.visual_depths,
drop_cls_token=model_args.visual_drop_cls_token,
hidden_stride=model_args.visual_hidden_stride,
)
visual_tokenizer = AutoModel.from_config(visual_tokenizer_config, train_from_scratch=True)
visual_tokenizer = visual_tokenizer.to(
device=torch.device(get_accelerator().device_name(os.getenv("LOCAL_RANK"))))
if getattr(training_args, 'hf_deepspeed_config', None) is not None:
set_hf_deepspeed_config(training_args.hf_deepspeed_config)
# 4. construct ovis model
model = Ovis(ovis_config, llm=llm, text_tokenizer=text_tokenizer, visual_tokenizer=visual_tokenizer,
train_from_scratch=True)
else: # load pretrained ovis model
model, loading_info = Ovis.from_pretrained(training_args.ovis_pretrained_path,
multimodal_max_length=model_args.multimodal_max_length,
output_loading_info=True)
rank0_print(BEGIN_LINE)
rank0_print(f'Loading info of Ovis:\n{loading_info}')
rank0_print(END_LINE)
training_args.vte_re_init = False
model.get_llm().config.use_cache = False
model.config.use_cache = False
text_tokenizer = model.get_text_tokenizer()
rank0_print(BEGIN_LINE)
rank0_print(f'model.config:\n{model.config}')
rank0_print(END_LINE)
# maybe re-init vte
if training_args.vte_re_init:
with deepspeed.zero.GatheredParameters([model.get_wte().weight]):
mean = model.get_wte().weight.mean().item()
std = model.get_wte().weight.std().item()
rank0_print(f'Statistics of embedding table of LLM: {mean=}, {std=}')
model.re_init_vte(mean, std)
# select train modules
model.requires_grad_(False)
for module in training_args.train_modules.split('|'):
if module == 'all':
model.requires_grad_(True)
elif module == 'llm':
model.get_llm().requires_grad_(True)
elif module == 'visual_tokenizer':
model.get_visual_tokenizer().requires_grad_(True)
elif module == 'visual_tokenizer.backbone':
model.get_visual_tokenizer().get_backbone().requires_grad_(True)
elif module.startswith('visual_tokenizer.backbone.layer.'):
layer_index = int(module[len('visual_tokenizer.backbone.layer.'):])
layer = model.get_visual_tokenizer().get_backbone_layer(layer_index)
layer.requires_grad_(True)
elif module == 'visual_tokenizer.head':
model.get_visual_tokenizer().get_head().requires_grad_(True)
elif module == 'vte':
model.get_vte().requires_grad_(True)
else:
raise ValueError(f'Invalid train module name: {module}')
rank0_print(BEGIN_LINE)
rank0_print('Parameters to train:')
for name, param in model.named_parameters():
if param.requires_grad:
rank0_print(name)
rank0_print(f'LLM\'s attn implementation: {model.get_llm().config._attn_implementation}')
rank0_print(END_LINE)
# construct data module
datasets = []
dataset_info_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
f'dataset/{training_args.dataset_info}.json')
with open(dataset_info_path, 'r', encoding='utf-8') as f:
dataset_info = json.load(f)
for name in training_args.dataset_names.split('|'):
info = dataset_info[name]
data_format = info['data_format']
if data_format == 'caption':
dataset = CaptionDataset(name, info, model, training_args)
elif data_format == 'conversation':
dataset = ConversationDataset(name, info, model, training_args)
else:
raise ValueError(f'Invalid data format `{data_format}` for dataset `{name}`')
datasets.append(dataset)
data_module = dict(
train_dataset=ConcatDataset(datasets),
data_collator=DataCollatorForMultimodalDataset(text_tokenizer)
)
# train
train_callbacks = [MonitorCallback]
if model_args.visual_tokenize_function == 'gumbel_argmax':
train_callbacks.append(TuneTauCallback)
trainer = Trainer(
model=model,
args=training_args,
callbacks=train_callbacks,
**data_module
)
rank0_print(BEGIN_LINE)
rank0_print('Dataset sample tensor:')
rank0_print(data_module['train_dataset'][0])
rank0_print(END_LINE)
rank0_print(BEGIN_LINE)
rank0_print('Dataset sample input_ids decoding:')
rank0_print(text_tokenizer.decode([x for x in data_module['train_dataset'][0]['input_ids'] if x >= 0]))
rank0_print(END_LINE)
rank0_print(BEGIN_LINE)
rank0_print('Dataset sample labels decoding:')
rank0_print(text_tokenizer.decode([x for x in data_module['train_dataset'][0]['labels'] if x >= 0]))
rank0_print(END_LINE)
rank0_print(BEGIN_LINE)
rank0_print(f'#param of model: {smart_unit(model.num_parameters())}')
rank0_print(f'#param of llm: {smart_unit(model.get_llm().num_parameters())}')
rank0_print(f'#param of visual_tokenizer: {smart_unit(model.get_visual_tokenizer().num_parameters())}')
rank0_print(f'#param of vte: {smart_unit(model.get_vte().weight.numel())}')
rank0_print(END_LINE)
if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
trainer.train(resume_from_checkpoint=True)
else:
trainer.train()
trainer.save_state()
# save model
model.get_llm().config.use_cache = True
model.config.use_cache = True
trainer.save_model()
if __name__ == '__main__':
train()
|