(
tokenizer: "PreTrainedTokenizerBase",
trajectory_groups: list[TrajectoryGroup],
allow_training_without_logprobs: bool,
scale_rewards: bool,
shuffle_group_trajectories: bool = True,
image_processor: BaseImageProcessor | None = None,
)
| 69 | |
| 70 | |
| 71 | def tokenize_trajectory_groups( |
| 72 | tokenizer: "PreTrainedTokenizerBase", |
| 73 | trajectory_groups: list[TrajectoryGroup], |
| 74 | allow_training_without_logprobs: bool, |
| 75 | scale_rewards: bool, |
| 76 | shuffle_group_trajectories: bool = True, |
| 77 | image_processor: BaseImageProcessor | None = None, |
| 78 | ) -> Generator["TokenizedResult", None, None]: |
| 79 | for group in trajectory_groups: |
| 80 | if not group: |
| 81 | continue |
| 82 | results: list[TokenizedResult] = [] |
| 83 | # Calculate GRPO group mean and standard deviation |
| 84 | reward_mean = sum(trajectory.reward for trajectory in group) / len(group) |
| 85 | reward_std = math.sqrt( |
| 86 | sum((trajectory.reward - reward_mean) ** 2 for trajectory in group) |
| 87 | / len(group) |
| 88 | ) |
| 89 | for trajectory in group: |
| 90 | # Calculate GRPO advantage for this trajectory |
| 91 | advantage = trajectory.reward - reward_mean |
| 92 | if scale_rewards: |
| 93 | advantage /= reward_std + 1e-6 |
| 94 | # Skip trajectories with no advantage |
| 95 | if advantage == 0: |
| 96 | continue |
| 97 | trajectory_results: list[TokenizedResult] = [] |
| 98 | for history in [ |
| 99 | History( |
| 100 | messages_and_choices=trajectory.messages_and_choices, |
| 101 | tools=trajectory.tools, |
| 102 | ), |
| 103 | *trajectory.additional_histories, |
| 104 | ]: |
| 105 | if result := tokenize_trajectory( |
| 106 | tokenizer, |
| 107 | image_processor, |
| 108 | history, |
| 109 | advantage, |
| 110 | allow_training_without_logprobs, |
| 111 | trajectory, |
| 112 | ): |
| 113 | trajectory_results.append(result) |
| 114 | weight = 1 / ( |
| 115 | sum(sum(result.assistant_mask) for result in trajectory_results) + 1e-6 |
| 116 | ) |
| 117 | for result in trajectory_results: |
| 118 | result.weight = weight |
| 119 | results.extend(trajectory_results) |
| 120 | # Choose a random prompt id |
| 121 | prompt_id = random.randint(-(2**63), 2**63 - 1) |
| 122 | # Find the longest shared prefix |
| 123 | # TODO: Potentially support multiple prompts per group |
| 124 | # Initial thought is to sort the results by token_ids and then |
| 125 | # successively group prompts with the same prefix. |
| 126 | prompt_length = len( |
| 127 | list( |
| 128 | takewhile( |
no test coverage detected