Repository: huggingface/trl
Pull request: 📇 GRPO: print completions to console and update docs
# What does this PR do?

- Update `GRPOConfig` to replace `log_completions` with `log_completions_steps`
- Add `print_prompt_completions_sample()` utility function for rich console logging
- Modify `GRPOTrainer` to additionally print 5 random prompt-completion pairs every log_completions_steps steps

Fixes #2948 

## Before submitting
- [x] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Problem statement:
GPRO: Expand log_completions logic or update docs
### Feature request

Currently, the GRPOConfig docstring for `log_completions` is:
> log_completions (bool, optional, defaults to False) — Whether to log the completions during training.

Which is a bit misleading, given that these are only logged when Weights and Biases is active.

Suggestion 1: Change the docs:
> log_completions (bool, optional, defaults to False) — If True, will try to log generated (prompt, completion) pairs to `wandb`, if it's available.

Suggestion 2: Log to `wandb` if possible, but also log to console. Since `rich` is a requirement, it could be a rich table:

```python
from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.text import Text

def print_output_sample(prompts: list[str], completions: list[str]):
    """Print out a sample of model completions."""
    console = Console()
    table = Table(show_header=True, header_style="bold white", expand=True, padding=(0, 1, 1, 0))
    table.add_column("Prompt", style="bright_yellow")
    for s, p in zip(source, prediction, strict=True):
        table.add_row(Text(s), Text(p))
    panel = Panel(table, expand=False, title="Output Sample", border_style="bold white")
    console.print(panel)
```

In that case, it may be useful to let `log_completions` be an int to specify a custom interval and avoid the console being spammed.

### Motivation

It's only right

### Your contribution

See suggestions

Hints from issue discussion:
Cool! WDYT of 

```python

from rich.console import Console
from rich.panel import Panel
from rich.table import Table
from rich.text import Text

def print_output_sample(prompts: list[str], completions: list[str], step: int) -> None:
    """Print out a sample of model completions."""
    console = Console()
    table = Table(show_header=True, header_style="bold white", expand=True, padding=(0, 1, 1, 0))
    table.add_column("Prompt", style="bright_yellow")
    table.add_column("Completion", style="bright_green")
    for s, p in zip(prompts, completions, strict=True):
        table.add_row(Text(s), Text(p))
    panel = Panel(table, expand=False, title=f"Step {step}", border_style="bold white")
    console.print(panel)

print_output_sample(["Hello, my name is", "The weather is", "I am feeling"], [" John", " sunny", " happy"], 1)
print_output_sample(["Hello, my name is", "The weather is", "I am feeling"], [" John", " sunny", " happy"], 2)
```

```console
╭───────────── Step 1 ─────────────╮
│ ┏━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓ │
│ ┃Prompt            ┃Completion ┃ │
│ ┃                  ┃           ┃ │
│ ┡━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩ │
│ │Hello, my name is │ John      │ │
│ │                  │           │ │
│ │The weather is    │ sunny     │ │
│ │                  │           │ │
│ │I am feeling      │ happy     │ │
│ │                  │           │ │
│ └──────────────────┴───────────┘ │
╰──────────────────────────────────╯
╭───────────── Step 2 ─────────────╮
│ ┏━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓ │
│ ┃Prompt            ┃Completion ┃ │
│ ┃                  ┃           ┃ │
│ ┡━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩ │
│ │Hello, my name is │ John      │ │
│ │                  │           │ │
│ │The weather is    │ sunny     │ │
│ │                  │           │ │
│ │I am feeling      │ happy     │ │
│ │                  │           │ │
│ └──────────────────┴───────────┘ │
╰──────────────────────────────────╯
```

Are you willing to open a PR?
@qgallouedec Looks good! I'll open a PR soon 🙂

Commit message: Add rich availability check and use fallback in print_prompt_completions_sample when rich is not available

Patch to review:
diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py
index 8386098982..fc4d2cbb92 100644
--- a/trl/trainer/grpo_config.py
+++ b/trl/trainer/grpo_config.py
@@ -112,9 +112,8 @@ class GRPOConfig(TrainingArguments):
             set `sync_ref_model=True`.
 
         > Parameters that control the logging
-
-        log_completions (`bool`, *optional*, defaults to `False`):
-            Whether to log the completions during training.
+        log_completions_steps (`int`, *optional*, defaults to `None`):
+            If set, will print a sample of completions every `log_completions_steps` steps. Also logs to `wandb`.
     """
 
     # Parameters that control the model and reference model
@@ -267,7 +266,9 @@ class GRPOConfig(TrainingArguments):
     )
 
     # Parameters that control the logging
-    log_completions: bool = field(
-        default=False,
-        metadata={"help": "Whether to log the completions during training."},
+    log_completions_steps: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "If set, will print a sample of completions every `log_completions_steps` steps. Also logs to `wandb`."
+        },
     )
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 6dd9d0ecee..0bc51a6ded 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import random
 import textwrap
 import warnings
 from collections import defaultdict
@@ -48,7 +49,13 @@
 from ..models import create_reference_model, prepare_deepspeed, unwrap_model_for_generation
 from .callbacks import SyncRefModelCallback
 from .grpo_config import GRPOConfig
-from .utils import generate_model_card, get_comet_experiment_url, pad, selective_log_softmax
+from .utils import (
+    generate_model_card,
+    get_comet_experiment_url,
+    pad,
+    print_prompt_completions_sample,
+    selective_log_softmax,
+)
 
 
 if is_peft_available():
@@ -248,7 +255,7 @@ def __init__(
         self,
         model: Union[str, PreTrainedModel],
         reward_funcs: Union[RewardFunc, list[RewardFunc]],
-        args: GRPOConfig = None,
+        args: Optional[GRPOConfig] = None,
         train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
         eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None,
         processing_class: Optional[PreTrainedTokenizerBase] = None,
@@ -391,7 +398,7 @@ def data_collator(features):  # No data collation is needed in GRPO
 
         # Initialize the metrics
         self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
-        self.log_completions = args.log_completions
+        self.log_completions_steps = args.log_completions_steps
 
         super().__init__(
             model=model,
@@ -839,22 +846,34 @@ def _generate_and_score_completions(
         self._metrics[mode]["reward_std"].append(std_grouped_rewards.mean().item())
 
         if (
-            self.log_completions
-            and self.state.global_step % self.args.logging_steps == 0
-            and "wandb" in self.args.report_to
+            self.log_completions_steps is not None
+            and self.log_completions_steps > 0
+            and self.state.global_step % self.log_completions_steps == 0
         ):
-            import pandas as pd
-
-            # For logging
-            table = {
-                "step": [str(self.state.global_step)] * len(rewards),
-                "prompt": gather_object(prompts_text),
-                "completion": gather_object(completions_text),
-                "reward": rewards.tolist(),
-            }
-            df = pd.DataFrame(table)
-
-            if wandb.run is not None and self.accelerator.is_main_process:
+            prompts_to_log: list[str] = gather_object(prompts_text)
+            completions_to_log: list[str] = gather_object(completions_text)
+            pairs_to_log = random.sample(list(zip(prompts_to_log, completions_to_log)), 5)
+            print_prompt_completions_sample(
+                [prompt for prompt, _ in pairs_to_log],
+                [completion for _, completion in pairs_to_log],
+                self.state.global_step,
+            )
+            if (
+                self.args.report_to
+                and "wandb" in self.args.report_to
+                and wandb.run is not None
+                and self.accelerator.is_main_process
+            ):
+                import pandas as pd
+
+                # For logging
+                table = {
+                    "step": [str(self.state.global_step)] * len(rewards),
+                    "prompt": prompts_to_log,
+                    "completion": completions_to_log,
+                    "reward": rewards.tolist(),
+                }
+                df = pd.DataFrame(table)
                 wandb.log({"completions": wandb.Table(dataframe=df)})
 
         return {
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
index 853ba1f3ca..d9cebcec22 100644
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -31,8 +31,6 @@
 from accelerate import Accelerator, PartialState
 from accelerate.state import AcceleratorState
 from huggingface_hub import ModelCard, ModelCardData
-from rich.console import Console
-from rich.table import Table
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import IterableDataset
 from transformers import (
@@ -52,9 +50,16 @@
     is_torch_xpu_available,
 )
 
+from ..import_utils import is_rich_available
 from ..trainer.model_config import ModelConfig
 
 
+if is_rich_available():
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.table import Table
+    from rich.text import Text
+
 if is_comet_available():
     import comet_ml
 
@@ -1683,3 +1688,38 @@ def selective_log_softmax(logits, index):
             per_token_logps.append(row_per_token_logps)
         per_token_logps = torch.stack(per_token_logps)
     return per_token_logps
+
+
+def print_prompt_completions_sample(prompts: list[str], completions: list[str], step: int) -> None:
+    """Print out a sample of model completions to the console.
+
+    This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring
+    model outputs during training.
+
+    If `rich` is not available, uses the regular print function.
+
+    Args:
+        prompts (list[str]): List of input prompts.
+        completions (list[str]): List of model-generated completions corresponding to the prompts.
+        step (int): Current training step number, used in the output title.
+
+    Returns:
+        None: Prints the formatted table to the console.
+    """
+    if is_rich_available():
+        console = Console()
+        table = Table(show_header=True, header_style="bold white", expand=True, padding=(0, 1, 1, 0))
+        table.add_column("Prompt", style="bright_yellow")
+        table.add_column("Completion", style="bright_green")
+        for s, p in zip(prompts, completions, strict=True):
+            table.add_row(Text(s), Text(p))
+        panel = Panel(table, expand=False, title=f"Step {step}", border_style="bold white")
+        console.print(panel)
+    else:
+        # Fallback to regular print when rich is not available
+        print(f"\n===== Step {step} =====")
+        for i, (s, p) in enumerate(zip(prompts, completions, strict=True)):
+            print(f"\n--- Sample {i + 1} ---")
+            print(f"Prompt:     {s}")
+            print(f"Completion: {p}")
+        print("\n" + "=" * 30 + "\n")


Code review comments:
diff_hunk: @@ -839,22 +846,34 @@ def _generate_and_score_completions(
         self._metrics[mode]["reward_std"].append(std_grouped_rewards.mean().item())
 
         if (
-            self.log_completions
-            and self.state.global_step % self.args.logging_steps == 0
-            and "wandb" in self.args.report_to
+            self.log_completions_steps is not None
+            and self.log_completions_steps > 0
+            and self.state.global_step % self.log_completions_steps == 0
         ):
-            import pandas as pd
-
-            # For logging
-            table = {
-                "step": [str(self.state.global_step)] * len(rewards),
-                "prompt": gather_object(prompts_text),
-                "completion": gather_object(completions_text),
-                "reward": rewards.tolist(),
-            }
-            df = pd.DataFrame(table)
-
-            if wandb.run is not None and self.accelerator.is_main_process:
+            prompts_to_log: list[str] = gather_object(prompts_text)
+            completions_to_log: list[str] = gather_object(completions_text)
+            pairs_to_log = random.sample(list(zip(prompts_to_log, completions_to_log)), 5)
path: trl/trainer/grpo_trainer.py
text: @user1:
similarly, I'd go for full logging. We can still add a param to limit the number of lines to be logged later.

diff_hunk: @@ -1683,3 +1688,38 @@ def selective_log_softmax(logits, index):
             per_token_logps.append(row_per_token_logps)
         per_token_logps = torch.stack(per_token_logps)
     return per_token_logps
+
+
+def print_prompt_completions_sample(prompts: list[str], completions: list[str], step: int) -> None:
+    """Print out a sample of model completions to the console.
+
+    This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring
+    model outputs during training.
+
+    If `rich` is not available, uses the regular print function.
+
+    Args:
+        prompts (list[str]): List of input prompts.
+        completions (list[str]): List of model-generated completions corresponding to the prompts.
+        step (int): Current training step number, used in the output title.
+
+    Returns:
+        None: Prints the formatted table to the console.
+    """
+    if is_rich_available():
+        console = Console()
+        table = Table(show_header=True, header_style="bold white", expand=True, padding=(0, 1, 1, 0))
+        table.add_column("Prompt", style="bright_yellow")
+        table.add_column("Completion", style="bright_green")
+        for s, p in zip(prompts, completions, strict=True):
+            table.add_row(Text(s), Text(p))
+        panel = Panel(table, expand=False, title=f"Step {step}", border_style="bold white")
+        console.print(panel)
+    else:
+        # Fallback to regular print when rich is not available
+        print(f"\n===== Step {step} =====")
+        for i, (s, p) in enumerate(zip(prompts, completions, strict=True)):
+            print(f"\n--- Sample {i + 1} ---")
+            print(f"Prompt:     {s}")
+            print(f"Completion: {p}")
+        print("\n" + "=" * 30 + "\n")
path: trl/trainer/utils.py
text: @user1:
I think you can drop the fallback

diff_hunk: @@ -1683,3 +1688,38 @@ def selective_log_softmax(logits, index):
             per_token_logps.append(row_per_token_logps)
         per_token_logps = torch.stack(per_token_logps)
     return per_token_logps
+
+
+def print_prompt_completions_sample(prompts: list[str], completions: list[str], step: int) -> None:
+    """Print out a sample of model completions to the console.
+
+    This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring
+    model outputs during training.
+
+    If `rich` is not available, uses the regular print function.
+
+    Args:
+        prompts (list[str]): List of input prompts.
+        completions (list[str]): List of model-generated completions corresponding to the prompts.
+        step (int): Current training step number, used in the output title.
+
+    Returns:
+        None: Prints the formatted table to the console.
+    """
path: trl/trainer/utils.py
text: @user1:
```suggestion
    """
    Print out a sample of model completions to the console.

    This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring model outputs
    during training. It requires the `rich` library to be installed.

    Args:
        prompts (`list[str]`):
            List of prompts.
        completions (`list[str]`):
            List of completions corresponding to the prompts.
        step (`int`):
            Current training step number, used in the output title.

    Example:
    ```python
    >>> from trl.trainer.utils import print_prompt_completions_sample
    >>> prompts = ["The sky is", "The sun is"]
    >>> completions = [" blue.", " in the sky."]
    >>> print_prompt_completions_sample(prompts, completions, step=42)
    ╭────────── Step 42 ──────────╮
    │ ┏━━━━━━━━━━━┳━━━━━━━━━━━━━┓ │
    │ ┃Prompt     ┃Completion   ┃ │
    │ ┃           ┃             ┃ │
    │ ┡━━━━━━━━━━━╇━━━━━━━━━━━━━┩ │
    │ │The sky is │ blue.       │ │
    │ │           │             │ │
    │ │The sun is │ in the sky. │ │
    │ │           │             │ │
    │ └───────────┴─────────────┘ │
    ╰─────────────────────────────╯
    ```
    """
```

diff_hunk: @@ -1683,3 +1688,38 @@ def selective_log_softmax(logits, index):
             per_token_logps.append(row_per_token_logps)
         per_token_logps = torch.stack(per_token_logps)
     return per_token_logps
+
+
+def print_prompt_completions_sample(prompts: list[str], completions: list[str], step: int) -> None:
+    """Print out a sample of model completions to the console.
+
+    This function creates a nicely formatted table showing prompt-completion pairs, useful for monitoring
+    model outputs during training.
+
+    If `rich` is not available, uses the regular print function.
+
+    Args:
+        prompts (list[str]): List of input prompts.
+        completions (list[str]): List of model-generated completions corresponding to the prompts.
+        step (int): Current training step number, used in the output title.
+
+    Returns:
+        None: Prints the formatted table to the console.
+    """
+    if is_rich_available():
+        console = Console()
+        table = Table(show_header=True, header_style="bold white", expand=True, padding=(0, 1, 1, 0))
+        table.add_column("Prompt", style="bright_yellow")
+        table.add_column("Completion", style="bright_green")
+        for s, p in zip(prompts, completions, strict=True):
+            table.add_row(Text(s), Text(p))
path: trl/trainer/utils.py
text: @user1:
```suggestion
        for prompt, completion in zip(prompts, completions, strict=True):
            table.add_row(Text(prompt), Text(completion))
```