[Bug] Evaluation metrics logged to training in wandb

[AJain9199]

1. Did you update? pip install --upgrade unsloth unsloth_zoo
   Yes

2. Colab or Kaggle or local / cloud
   Colab

3. Number GPUs used, use nvidia-smi
   1 A100 80GB

4. Which notebook? Please link!
   Code

5. Which Unsloth version, TRL version, transformers version, PyTorch version?

unsloth==2025.12.6
unsloth_zoo==2025.12.5
trl==0.25.0, 0.24.0, 0.23.1 (all tested)
torch==2.8.0
torchao=0.13.0
vllm==0.11.0
flashinfer-python==0.3.1.post1.

6. Which trainer? SFTTrainer, GRPOTrainer etc
   GRPO

This basic GRPO script reproduces the issue

Hi. When using GRPO with wandb, the eval rewards are logged with the train logs. This makes it extremely hard to track evaluation progress.

Running the script, the eval log is:

{'eval_loss': -0.0030808576848357916, .. 'rewards/dummy_reward/mean': 0.5854228019714356,  ...}

But it is logged with the train reward:

Rewards are not logged in the evaluation pane:

Thanks!

[numb3r33]

@danielhanchen @shimmyshimmer

I did some debugging and found that during evaluation, Trainer.evaluate() correctly calls model.eval() to set model.training=False. However, Unsloth's patched _generate_and_score_completions function calls self.model.for_training() (rl_replacements.py#L267, L273) but never restores the original mode afterward. This leaves model.training=True even during evaluation.

When TRL's log() method runs, it checks self.model.training to determine which metrics bucket to use. Since model.training=True, eval metrics get logged to the "train" bucket instead of "eval".

Proposed Fix:

Save the original training state (_was_training = self.model.training) before calling for_training() (around L260)
Restore it at the end of _generate_and_score_completions before return output.

Example:

Modify the existing replacement_lines to save the mode

replacement_lines = """
    batch_size = self.args.per_device_train_batch_size if mode == "train" else self.args.per_device_eval_batch_size
    _was_training = self.model.training
    try:
        # TRL 0.23.1 and below path
        if not has_images:
            # Left pad prompt before calculation old and ref hidden states
            prompt_completion_ids = left_pack_padding(prompt_completion_ids, self.processing_class.pad_token_id)
        self.model.for_training()
    except:
        # TRL 0.24.0 and below path
        if images is None:
            # Left pad prompt before calculation old and ref hidden states
            prompt_completion_ids = left_pack_padding(prompt_completion_ids, self.processing_class.pad_token_id)
        self.model.for_training()"""

Add this new replacement near the end of the function (before return function):

function = function.replace(
    "        return output",  # 8 spaces before 'return'
    """        if not _was_training:
            self.model.for_inference()
        return output"""
)

I've tested this fix locally and confirmed eval metrics now appear correctly in wandb's evaluation panel. Would you like me to open a PR?

[Datta0]

Hey @numb3r33
Thanks for the initial investigation. This is very valuable and I can see what is potentially happening here...

I'd strongly encourage you to raise the PR if you want to and I'll be happy to review it asap (please do tag me).
I think this is a very good thing to fix. A test script/colab notebook to showcase your changes would be awesome as well ( I know it'd be just setting report_to="wandb" in any of our existing notebooks but it'd be easier for us to approve)

[numb3r33]

Thanks @Datta0 for the feedback, happy to raise a PR and would share a script to validate

[danielhanchen]

@numb3r33 Thanks! Definitely if you can make a PR we'll merge it in asap - appreciate it!

[numb3r33]

@Datta0 I have raised a PR and sharing a Colab Notebook to validate results. Attaching a screenshot of wandb panel and link to the workspace. Please take a look and let me know if any more info or changes are required