Skip to content

Commit 64aa542

Browse files
hjh0119Jintao-Huang
authored andcommitted
[grpo] remove data collator to top-level to avoid pickle error in spawn mode (#4582)
* remove data collator to top-level * comment
1 parent f1e8821 commit 64aa542

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

swift/trainers/rlhf_trainer/grpo_trainer.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,10 @@ class DataCache:
7474
outputs: List[Dict] = field(default_factory=list)
7575

7676

77+
def identity_data_collator(features):
78+
return features
79+
80+
7781
class GRPOTrainer(RLHFTrainerMixin, SwiftMixin, HFGRPOTrainer):
7882
executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
7983

@@ -173,10 +177,7 @@ def __init__(self,
173177
self.completion_length_limit_scope = args.completion_length_limit_scope
174178
model.warnings_issued['estimate_tokens'] = True
175179

176-
def data_collator(features): # No data collation is needed in GRPO
177-
return features
178-
179-
kwargs['data_collator'] = data_collator
180+
kwargs['data_collator'] = identity_data_collator # No data collation is needed in GRPO
180181
self.shuffle_dataset = args.dataset_shuffle
181182

182183
self.use_vllm = args.use_vllm

0 commit comments

Comments
 (0)