mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-04-29 20:25:29 +00:00
9d02590c9a
* refactor: separate log_probs fn from Actor forward fn * refactor: separate generate fn from Actor class * feat: update unwrap_model and get_base_model * unwrap_model returns model not wrapped by Strategy * get_base_model returns HF model for Actor, Critic and RewardModel * feat: simplify Strategy.prepare * style: remove get_base_model method of Actor * perf: tokenize text in batches * refactor: move calc_action_log_probs to utils of model * test: update test with new forward fn * style: rename forward fn args * fix: do not unwrap model in save_model fn of naive strategy * test: add gemini test for train_prompts * fix: fix _set_default_generate_kwargs
40 lines
1.7 KiB
Python
40 lines
1.7 KiB
Python
import torch
|
|
from coati.models.generation import generate_with_actor
|
|
from coati.models.utils import calc_action_log_probs, compute_reward, normalize
|
|
|
|
from .base import Experience, ExperienceMaker
|
|
|
|
|
|
class NaiveExperienceMaker(ExperienceMaker):
|
|
"""
|
|
Naive experience maker.
|
|
"""
|
|
|
|
@torch.no_grad()
|
|
def make_experience(self, input_ids: torch.Tensor, **generate_kwargs) -> Experience:
|
|
self.actor.eval()
|
|
self.critic.eval()
|
|
self.initial_model.eval()
|
|
self.reward_model.eval()
|
|
|
|
sequences, attention_mask, action_mask = generate_with_actor(self.actor,
|
|
input_ids,
|
|
return_action_mask=True,
|
|
**generate_kwargs)
|
|
num_actions = action_mask.size(1)
|
|
|
|
actor_output = self.actor(sequences, attention_mask)
|
|
action_log_probs = calc_action_log_probs(actor_output, sequences, num_actions)
|
|
base_model_output = self.initial_model(sequences, attention_mask)
|
|
base_action_log_probs = calc_action_log_probs(base_model_output, sequences, num_actions)
|
|
value = self.critic(sequences, action_mask, attention_mask)
|
|
r = self.reward_model(sequences, attention_mask)
|
|
reward = compute_reward(r, self.kl_coef, action_log_probs, base_action_log_probs, action_mask=action_mask)
|
|
|
|
advantage = reward - value
|
|
# TODO(ver217): maybe normalize adv
|
|
if advantage.ndim == 1:
|
|
advantage = advantage.unsqueeze(-1)
|
|
|
|
return Experience(sequences, action_log_probs, value, reward, advantage, attention_mask, action_mask)
|