github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Compare commits

base: github:802c15b2b4fa1658c1ca5a2fe8414a5b052028f0
Branches Tags
github:master
github:website
github:v2
github:minimind-v1
..
compare: github:5020dc9dd4ea5ff997fb4f1e773157056c5253a0
Branches Tags
github:master
github:website
github:v2
github:minimind-v1

No commits in common. "802c15b2b4fa1658c1ca5a2fe8414a5b052028f0" and "5020dc9dd4ea5ff997fb4f1e773157056c5253a0" have entirely different histories.
802c15b2b4 ... 5020dc9dd4

2 changed files with 4 additions and 6 deletions
Show Stats Expand all files Collapse all files

3
trainer/train_agent.py
View File

@ -269,8 +269,6 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
old_per_token_logps = torch.tensor([old_logps + [0.0] * ((max_len - 1) - len(old_logps)) for _, _, _, old_logps in packed_samples], device=args.device, dtype=torch.float32) old_per_token_logps = torch.tensor([old_logps + [0.0] * ((max_len - 1) - len(old_logps)) for _, _, _, old_logps in packed_samples], device=args.device, dtype=torch.float32)
full_mask = (input_ids != tokenizer.pad_token_id).long() full_mask = (input_ids != tokenizer.pad_token_id).long()
rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch)
model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
with autocast_ctx: with autocast_ctx:
res = model_unwrapped(input_ids, attention_mask=full_mask) res = model_unwrapped(input_ids, attention_mask=full_mask)
@ -290,6 +288,7 @@ def rl_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_model
completion_mask = completion_mask * (pos <= eos_idx.unsqueeze(1)).float() completion_mask = completion_mask * (pos <= eos_idx.unsqueeze(1)).float()
token_counts = completion_mask.sum(dim=1) token_counts = completion_mask.sum(dim=1)
valid_rows = token_counts > 0 valid_rows = token_counts > 0
rewards = calculate_rewards(prompts, completions, gt_batch, tools_batch, args.num_generations, reward_model, device=args.device, turn_outputs_batch=turn_outputs_batch, unfinished_batch=unfinished_batch)
if args.debug_mode and is_main_process() and step % args.debug_interval == 0: if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
for i in range(len(messages_batch)): for i in range(len(messages_batch)):

5
trainer/train_grpo.py
View File

@ -86,13 +86,11 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
outputs = rollout_result.output_ids outputs = rollout_result.output_ids
completion_ids = rollout_result.completion_ids completion_ids = rollout_result.completion_ids
completions = rollout_result.completions completions = rollout_result.completions
old_per_token_logps = rollout_result.per_token_logps.to(args.device).detach() old_per_token_logps = rollout_result.per_token_logps.to(args.device)
prompt_lens = rollout_result.prompt_lens.to(args.device) prompt_lens = rollout_result.prompt_lens.to(args.device)
full_mask = (outputs != tokenizer.pad_token_id).long() full_mask = (outputs != tokenizer.pad_token_id).long()
logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0) logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0)
rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen]
model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model
with autocast_ctx: with autocast_ctx:
res = model_unwrapped(outputs, attention_mask=full_mask) res = model_unwrapped(outputs, attention_mask=full_mask)
@ -101,6 +99,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
with torch.no_grad(): with torch.no_grad():
ref_per_token_logps = F.log_softmax(ref_model(outputs, attention_mask=full_mask).logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos) ref_per_token_logps = F.log_softmax(ref_model(outputs, attention_mask=full_mask).logits[:, :-1, :], dim=-1).gather(2, outputs[:, 1:].unsqueeze(-1)).squeeze(-1).gather(1, logp_pos)
rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen]
if args.debug_mode and is_main_process() and step % args.debug_interval == 0: if args.debug_mode and is_main_process() and step % args.debug_interval == 0:
for i in range(len(prompts)): for i in range(len(prompts)):