huggingface / huggingface/trl

Introduce a Minimal Backend Generation Interface

Open
#5,193 0 comments 0 reactions 1 assignee Claimed by @rycerzes View on GitHub
Dominant language
Python
Stars
19.3k
Forks
3k
Avg merge
1d 20h
Merged PRs (30d)
194

Description

### As part of #5119

**Motivation:**
After sub-issue #5121 decouples rollout dispatch, `_generate_single_turn` still has
three backend-specific branches inlined in the method body. These branches share
common structure (apply chat template -> call backend -> extract ids/logprobs)
but differ in backend-specific details. Encapsulating them behind a common
interface would:

- Reduce the cognitive load of reading _generate_single_turn
- Make it possible to add new backends without modifying the trainer
- Allow openenv/utils.py to call the generation backend without knowing which
one is active
- Align TRL with the patterns used by other frameworks

**Scope:**

Define a minimal internal interface (not necessarily a full ABC) that all
generation paths implement:

```python
class GenerationBackend(Protocol):
def generate(
self,
prompts: list,
num_generations: int,
processing_class,
generation_config: GenerationConfig | None = None,
) -> tuple[list[list[int]], list[list[int]], list[list[float]] | None]:
"""
Returns (prompt_ids, completion_ids, logprobs).
Chat templating happens inside this method.
"""
...

def sync_weights(self) -> None:
"""Sync model weights from trainer to inference engine (no-op for in-process backends)."""
...
```

CC: @albertvillanova

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.