zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-06-27 19:31:20 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

smoll fix

Browse Source
This commit is contained in:
Goekdeniz-Guelmez 2025-02-26 15:21:57 +01:00
parent ef6ff92add
commit fab2dc2688
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
llms/mlx_lm/tuner/grpo_trainer.py
View File

@ -279,7 +279,8 @@ def grpo_loss(
reward_weights = mx.array(reward_weights, dtype=mx.float32) reward_weights = mx.array(reward_weights, dtype=mx.float32)
else: else:
reward_weights = mx.ones(len(reward_funcs), dtype=mx.float32) reward_weights = mx.ones(len(reward_funcs), dtype=mx.float32)
rewards = (rewards * mx.expand_dims(reward_weights, 0)).sum(axis=1)
rewards = (rewards * mx.expand_dims(reward_weights, 0)).sum(axis=1)
# Reshape rewards and compute advantages # Reshape rewards and compute advantages
rewards_reshaped = rewards.reshape(batch_size, group_size) rewards_reshaped = rewards.reshape(batch_size, group_size)