Shortcuts

lightrft.trainer.fast_exp_maker

FastExperienceMaker Module

This module provides an optimized experience maker for RLHF (Reinforcement Learning from Human Feedback) that supports high-performance inference backends like VLLM and SGLang. It extends the base NaiveExperienceMaker with enhanced features for multimodal data processing, reward computation, and advantage estimation.

Key Features:
  • VLLM/SGLang backend support for efficient text generation

  • Multimodal (vision-language) data processing

  • Multiple advantage estimation methods (GAE, RLOO, REINFORCE, Group Norm)

  • Flexible reward model composition with custom reward functions

  • Sample packing support for improved training efficiency

  • Running reward normalization and advantage whitening

Classes:

MultimodalDataProcessor: Handles preprocessing of mixed text/image data RewardComputationEngine: Manages reward model inference and aggregation FastExperienceMaker: Main experience generation class

FastExperienceMaker

class lightrft.trainer.fast_exp_maker.FastExperienceMaker(*args, packing_samples: bool = False, processor=None, **kwargs)[source]

Experience maker backed by vLLM or SGLang rollout.

The class coordinates rollout, shard-local preprocessing, Actor/Reference/Critic inference, reward aggregation, and advantage computation. It supports text and visual-language batches, optional sample packing, and runtime reward statistics.

Parameters:
  • args – Positional arguments forwarded to NaiveExperienceMaker.

  • packing_samples – Whether to pack multiple sequences into one batch.

  • processor – Optional multimodal processor.

  • kwargs – Keyword arguments forwarded to NaiveExperienceMaker.

__init__(*args, packing_samples: bool = False, processor=None, **kwargs)[source]

Initialize FastExperienceMaker.

Parameters:
  • args (tuple) – Positional arguments for NaiveExperienceMaker

  • packing_samples (bool) – Enable sample packing for efficiency

  • processor (Optional[Any]) – Multimodal processor (required for VLM models)

  • kwargs (dict) – Keyword arguments for NaiveExperienceMaker

generate_samples(all_prompts: List[str], all_images: List | None = None, all_videos: List | None = None, images_num: List[int] | None = None, videos_num: List[int] | None = None, all_references: List[str] | None = None, all_labels: List | None = None, **generate_kwargs) List[Samples]

Generate samples using the inference engine (VLLM or SGLang).

This method handles:
  • Sampling parameter configuration

  • Multimodal data processing

  • Inference engine invocation

  • Output processing into Samples format

Parameters:
  • all_prompts (List[str]) – List of text prompts

  • all_images (Optional[List]) – Optional images for VLM

  • images_num (Optional[List[int]]) – Number of images per prompt

  • all_references (Optional[List[str]]) – Reference texts

  • all_labels (Optional[List]) – Sample labels

  • all_videos (Optional[List]) – Optional videos for VLM

  • videos_num (Optional[List[int]]) – Number of videos per prompt

  • generate_kwargs (dict) – Generation parameters (temperature, max_new_tokens, etc.)

Returns:

List of Samples or SamplesVL objects

Return type:

List[Union[Samples, SamplesVL]]

get_advantages_and_returns(values: torch.Tensor, rewards: torch.Tensor, action_mask: torch.Tensor, gamma: float, lambd: float) Tuple[torch.Tensor, torch.Tensor, float][source]

Compute advantages and returns using Generalized Advantage Estimation (GAE).

Extends parent method with advantage whitening and clipping.

Parameters:
  • values (torch.Tensor) – Value estimates from critic

  • rewards (torch.Tensor) – Reward signals

  • action_mask (torch.Tensor) – Mask for valid action positions

  • gamma (float) – Discount factor

  • lambd (float) – GAE lambda parameter

Returns:

Tuple of (advantages, returns, advantage_clip_fraction)

Return type:

Tuple[torch.Tensor, torch.Tensor, float]

make_experience_list(all_prompts: List[str], all_images: List | None = None, all_videos: List | None = None, all_references: List[str] | None = None, all_labels: List | None = None, **generate_kwargs) List[ExperienceVL]

Generate a list of experiences from prompts and optional multimodal inputs.

This is the main entry point for experience generation. It orchestrates the entire pipeline from sampling to advantage computation.

Parameters:
  • all_prompts (List[str]) – List of text prompts

  • all_images (Optional[List]) – Optional images for multimodal generation

  • all_references (Optional[List[str]]) – Optional reference texts for evaluation

  • all_labels (Optional[List]) – Optional labels for samples

  • all_videos (Optional[List]) – Optional videos for multimodal generation

  • generate_kwargs (dict) – Generation parameters (temperature, max_new_tokens, etc.)

Returns:

List of Experience or ExperienceVL objects with computed advantages and returns

Return type:

List[Union[Experience, ExperienceVL]]