| from header import * |
|
|
| class DeepSpeedAgent: |
| |
| def __init__(self, model, args): |
| super(DeepSpeedAgent, self).__init__() |
| self.args = args |
| self.model = model |
| if args['stage'] == 2: |
| self.load_stage_1_parameters(args["delta_ckpt_path"]) |
| print(f'[!] load stage 1 checkpoint from {args["delta_ckpt_path"]}') |
|
|
| |
| 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, batch): |
| self.model.eval() |
| string = self.model.generate_one_sample(batch) |
| return string |
|
|
| def train_model(self, batch, current_step=0, pbar=None): |
| self.ds_engine.module.train() |
| loss, mle_acc = self.ds_engine(batch) |
|
|
| 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): |
| |
| 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 |
| 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 load_stage_1_parameters(self, path): |
| delta_ckpt = torch.load(path, map_location=torch.device('cpu')) |
| self.model.load_state_dict(delta_ckpt, strict=False) |
|
|
