Spaces:
Runtime error
Runtime error
File size: 2,957 Bytes
8366b03 |
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 |
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"]}')
# load config parameters of deepspeed
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):
# only save trainable model 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
torch.save(checkpoint, f'{path}/pytorch_model.pt')
# save tokenizer
self.model.llama_tokenizer.save_pretrained(path)
# save configuration
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)
|