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

Merge pull request #759 from TKiteRunner/fix/grpo-oom-reorder-rewards 

fix: resolve CUDA OOM in train_grpo.py on GPUs with <=8GB VRAM
pull/762/head
jingyaogong 2026-05-06 15:00:17 +08:00 committed by GitHub
parent 5020dc9dd4 10776417aa
commit e73a407f7a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
trainer/train_grpo.py
View File

@ -83,14 +83,19 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod
max_new_tokens=args.max_gen_len, max_new_tokens=args.max_gen_len,
temperature=0.8, temperature=0.8,
) )
outputs = rollout_result.output_ids outputs = rollout_result.output_ids.clone()
completion_ids = rollout_result.completion_ids completion_ids = rollout_result.completion_ids.clone()
completions = rollout_result.completions completions = rollout_result.completions
old_per_token_logps = rollout_result.per_token_logps.to(args.device) old_per_token_logps = rollout_result.per_token_logps.to(args.device).detach()
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)
# 在 policy/ref 前向占用显存之前先算 rewards，此时显存最充裕
torch.cuda.empty_cache()
rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen]
torch.cuda.empty_cache()
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)
@ -99,7 +104,6 @@ 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)):