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

[fix] prompt length calculate

pull/825/head
jingyaogong 2025-11-15 18:25:37 +08:00
parent d334a6f9f4
commit c8727733f3
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
trainer/train_ppo.py
View File

@ -124,7 +124,7 @@ def ppo_train_epoch(epoch, loader, iters, old_actor_model, ref_model, actor_sche
prompts = batch["prompt"] # list[str], length B prompts = batch["prompt"] # list[str], length B
enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True,
max_length=args.max_seq_len).to(args.device) # input_ids: [B, P], attention_mask: [B, P] max_length=args.max_seq_len).to(args.device) # input_ids: [B, P], attention_mask: [B, P]
prompt_lengths = enc.attention_mask.sum(dim=1) # [B] prompt_lengths = torch.full((enc.input_ids.size(0),), enc.input_ids.shape[1], dtype=torch.long, device=enc.input_ids.device) # [B]
with torch.no_grad(): with torch.no_grad():
# DDP 模型需要使用 .module 访问 generate 方法 # DDP 模型需要使用 .module 访问 generate 方法