| import os.path |
|
|
| import torch |
| from header import * |
|
|
|
|
| class DeepSpeedAgent: |
|
|
| def __init__(self, model, args): |
| super(DeepSpeedAgent, self).__init__() |
| self.args = args |
| self.model = model |
|
|
| self.print_model_parameters() |
| self.writer = SummaryWriter(args['log_path']) |
|
|
| self.load_parameters(self.args['save_path'], self.args['stage']) |
| |
| |
| ds_params = json.load(open(self.args['ds_config_path'])) |
| ds_params['scheduler']['params']['total_num_steps'] = self.args['total_steps'] |
| ds_params['scheduler']['params']['warmup_num_steps'] = max(10, int( |
| self.args['total_steps'] * self.args['warmup_rate'])) |
| self.ds_engine, self.optimizer, _, _ = deepspeed.initialize( |
| model=self.model, |
| model_parameters=self.model.parameters(), |
| config_params=ds_params, |
| dist_init_required=True, |
| args=types.SimpleNamespace(**args) |
| ) |
|
|
| @torch.no_grad() |
| def predict(self): |
| self.ds_engine.module.eval() |
| output = self.ds_engine.generate(self.args) |
| return output |
|
|
| def train_model(self, batch, current_step=0, pbar=None): |
| self.ds_engine.module.train() |
|
|
| loss, mle_acc, mse_loss = self.ds_engine(batch) |
|
|
| self.writer.add_scalar('loss', loss, current_step) |
| self.writer.add_scalar('mle_acc', mle_acc, current_step) |
| |
| |
| |
| |
| if isinstance(mse_loss, torch.Tensor): |
| self.writer.add_scalar('mse_loss', mse_loss, current_step) |
| else: |
| pass |
| |
|
|
| self.ds_engine.backward(loss) |
| self.ds_engine.step() |
| |
| pbar.set_description(f'[!] loss: {round(loss.item(), 4)}; token_acc: {round(mle_acc * 100, 2)}') |
| pbar.update(1) |
| if self.args['local_rank'] == 0 and self.args['log_path'] and current_step % self.args['logging_step'] == 0: |
| elapsed = pbar.format_dict['elapsed'] |
| rate = pbar.format_dict['rate'] |
| remaining = (pbar.total - pbar.n) / rate if rate and pbar.total else 0 |
| remaining = str(datetime.timedelta(seconds=remaining)) |
| logging.info( |
| f'[!] progress: {round(pbar.n / pbar.total, 5)}; remaining time: {remaining}; loss: {round(loss.item(), 4)}; token_acc: {round(mle_acc * 100, 2)}') |
| |
| mle_acc *= 100 |
| return mle_acc |
|
|
| def save_model(self, path, current_step): |
| """ |
| this function also save the trainable parameters and specific name parameters |
| """ |
| param_grad_dic = { |
| k: v.requires_grad for (k, v) in self.ds_engine.module.named_parameters() |
| } |
| state_dict = self.ds_engine.module.state_dict() |
| checkpoint = OrderedDict() |
| for k, v in self.ds_engine.module.named_parameters(): |
| if v.requires_grad: |
| checkpoint[k] = v |
| if 'gen_text_hidden_fcs' in k: |
| checkpoint[k] = v |
| if 'gen_text_hidden_fcs_video' in k: |
| checkpoint[k] = v |
| if 'gen_text_hidden_fcs_audio' in k: |
| checkpoint[k] = v |
| if 'llama_proj' in k: |
| checkpoint[k] = v |
| torch.save(checkpoint, f'{path}/pytorch_model.pt') |
| |
| self.model.llama_tokenizer.save_pretrained(path) |
| |
| self.model.llama_model.config.save_pretrained(path) |
| print(f'[!] save model into {path}') |
|
|
| def print_model_parameters(self, use_4bit=False): |
| """ |
| Prints the number of trainable parameters in the model. |
| """ |
| trainable_params = 0 |
| all_param = 0 |
| lora = 0 |
| image = 0 |
| video = 0 |
| audio = 0 |
| linear = 0 |
| llama = 0 |
| imagebind = 0 |
| for name, param in self.model.named_parameters(): |
| num_params = param.numel() |
| |
| if num_params == 0 and hasattr(param, "ds_numel"): |
| num_params = param.ds_numel |
|
|
| if 'lora' in name: |
| lora += num_params |
| elif 'gen_text_hidden_fcs_video' in name: |
| video += num_params |
| elif 'gen_text_hidden_fcs_audio' in name: |
| audio += num_params |
| elif 'gen_text_hidden_fcs' in name: |
| image += num_params |
| elif 'llama_proj' in name: |
| linear += num_params |
| elif 'llama_model' in name: |
| llama += num_params |
| elif 'visual_encoder' in name: |
| imagebind += num_params |
| else: |
| pass |
|
|
| all_param += num_params |
| if param.requires_grad: |
| trainable_params += num_params |
| if use_4bit: |
| trainable_params /= 2 |
| print( |
| f"all params: {all_param:,d} || trainable params: {trainable_params:,d} || trainable%: {100 * trainable_params / all_param}" |
| ) |
| print(f'lora params: {lora:,d} || video params: {video:,d} || audio params: {audio:,d} || image params: {image:,d}') |
| print(f'linear params: {linear:,d} || imagebind params: {imagebind:,d} || llama params: {llama:,d}') |
|
|
| def load_parameters(self, path, stage=3): |
| if os.path.exists(os.path.join(path, 'pytorch_model.pt')): |
| print('loading parameters from {}'.format(self.args['save_path'])) |
| delta_ckpt = torch.load(f'{path}/pytorch_model.pt', map_location=torch.device('cuda')) |
| checkpoint = OrderedDict() |
| if stage == 3: |
| for k, v in delta_ckpt.items(): |
| if 'llama_model.model.embed_tokens.weight' in k: |
| checkpoint['llama_model.base_model.model.model.embed_tokens.weight'] = v |
| elif 'llama_model.lm_head.weight' in k: |
| checkpoint['llama_model.base_model.model.lm_head.weight'] = v |
| else: |
| checkpoint[k] = v |
| else: |
| checkpoint = delta_ckpt |
| self.model.load_state_dict(checkpoint, strict=False) |
|
|
|
|
