Skip to content

Commit d88bada

Browse files
hjh0119Jintao-Huang
authored andcommitted
[grpo] fix pickle data collator (#4562)
1 parent 97d646b commit d88bada

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

swift/trainers/rlhf_trainer/grpo_trainer.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -172,7 +172,11 @@ def __init__(self,
172172
self.max_completion_length = args.max_completion_length
173173
self.completion_length_limit_scope = args.completion_length_limit_scope
174174
model.warnings_issued['estimate_tokens'] = True
175-
kwargs['data_collator'] = lambda features: features
175+
176+
def data_collator(features): # No data collation is needed in GRPO
177+
return features
178+
179+
kwargs['data_collator'] = data_collator
176180
self.shuffle_dataset = args.dataset_shuffle
177181

178182
self.use_vllm = args.use_vllm

0 commit comments

Comments
 (0)