refactor(sunjx): refactor loss-filter implementation

lightrft/trainer/fast_exp_maker.py

@@ -53,6 +53,7 @@
 from lightrft.utils.remote_rm_utils import remote_rm_fn
 from lightrft.utils import Timer, get_current_device
 from .utils import RunningMoments, compute_clip_fraction, get_cpgd_advantages_returns, fire_sampling
+from .filter_weight import FilterWeightManagerBuilder
 
 # ============================================================================
 # Data Structures
@@ -951,6 +952,19 @@ def __init__(self, *args, packing_samples: bool = False, processor=None, **kwarg
             packing_samples=self.packing_samples,
         )
 
+        # Initialize filter-weight manager
+        self.filter_weight_manager = self._init_filter_weight_manager()
+
+    def _init_filter_weight_manager(self):
+        """
+        Initialize filter-weight manager from strategy args.
+
+        :return: FilterWeightManager instance
+        :rtype: FilterWeightManager
+        """
+        args = self.strategy.args
+        return FilterWeightManagerBuilder.from_args(args, packing_samples=self.packing_samples)
+
     # ========================================================================
     # Public API Methods
     # ========================================================================
@@ -1020,12 +1034,32 @@ def make_experience_list(
 
         # ========== Stage 3: Model Inference ==========
         Timer.start('  make_experience')
-        experiences = self._make_experience_list_by_model(all_samples)
+        experiences, outputs = self._make_experience_list_by_model(all_samples)
         Timer.stop('  make_experience')
 
         # ========== Stage 4: Shard-Parallel Postprocessing ==========
         experiences = self.strategy.sp_data_processor.postprocess(experiences)
 
+        # ========== Stage 4.5: Apply Filter-Weight Framework ==========
+        if self.filter_weight_manager is not None and (
+            self.filter_weight_manager.filters or self.filter_weight_manager.weights
+        ):
+            # Compute metrics from outputs
+            current_step = getattr(self.strategy, "global_step", None)
+            metrics = self.filter_weight_manager.compute_metrics(outputs, current_step=current_step)
+
+            # Apply filters and weights
+            experiences, sample_weights = self.filter_weight_manager.apply_to_experiences(
+                experiences, metrics
+            )
+
+            # Store sample weights in experience info for later use
+            sample_idx = 0
+            for exp in experiences:
+                batch_size = len(exp.sequences)
+                exp.info["sample_weights"] = sample_weights[sample_idx:sample_idx + batch_size]
+                sample_idx += batch_size
+
         # ========== Stage 5: Reward Processing ==========
         experiences, rewards = self._process_experiences(  # GRPO's -mean / std operation is performed in this method
             experiences, generate_kwargs.get("max_new_tokens", 1024)
@@ -1360,7 +1394,16 @@ def _process_experiences(
         rewards = torch.cat([exp.info["reward"] for exp in experiences])
 
         # ========== Overlong Sequence Penalty ==========
-        if config.overlong_buffer:
+        # Use new filter_weight framework if enabled, otherwise use legacy logic
+        from .filter_weight import ResponseLengthFilter
+        use_filter_weight = (
+            self.filter_weight_manager is not None
+            and self.filter_weight_manager.filters
+            and any(isinstance(f, ResponseLengthFilter) for f in self.filter_weight_manager.filters)
+        )
+
+        if config.overlong_buffer and not use_filter_weight:
+            # Legacy overlong buffer penalty (only if not using filter_weight framework)
             expected_len = max_new_tokens - config.overlong_buffer_len
             actual_lens = torch.cat([exp.action_mask.sum(dim=1) for exp in experiences])
             exceed_len = actual_lens - expected_len
@@ -1393,7 +1436,16 @@ def _process_experiences(
 
         elif config.advantage_estimator in ["group_norm", "grpo"]:
             # Group normalization with optional dynamic filtering
-            if config.dynamic_sampling:
+            # Use new filter_weight framework if enabled, otherwise use legacy logic
+            from .filter_weight import RewardValueFilter
+            use_dynamic_filter = (
+                self.filter_weight_manager is not None
+                and self.filter_weight_manager.filters
+                and any(isinstance(f, RewardValueFilter) for f in self.filter_weight_manager.filters)
+            )
+
+            if config.dynamic_sampling and not use_dynamic_filter:
+                # Legacy dynamic sampling (only if not using filter_weight framework)
                 step_size = config.n_samples_per_prompt // config.micro_train_batch_size
                 for i in range(0, len(experiences), step_size):
                     chunk = experiences[i:i + step_size]
@@ -1545,7 +1597,7 @@ def _compute_advantages_and_returns(
     def _make_experience_list_by_model(
         self,
         all_samples: List[Union[Samples, SamplesVL]],
-    ) -> List[Union[Experience, ExperienceVL]]:
+    ) -> Tuple[List[Union[Experience, ExperienceVL]], List[_SamplesOutput]]:
         """
         Batch forward pass through all models to create experiences.
 
@@ -1607,7 +1659,8 @@ def _make_experience_list_by_model(
         self.reward_engine.compute_rewards(outputs, vlm_mode, device)
 
         # ========== Stage 5: Assemble Experiences ==========
-        return [self._pack_experience(output, vlm_mode) for output in outputs]
+        experiences = [self._pack_experience(output, vlm_mode) for output in outputs]
+        return experiences, outputs
 
     def _preprocess_sample(
         self,

lightrft/trainer/filter_weight/__init__.py

@@ -0,0 +1,169 @@
+"""
+Filter and Weight Module
+
+Unified interface for sample filtering and loss weighting in RLHF.
+
+This module provides a three-layer architecture for managing sample filtering
+and loss weighting:
+
+1. **Metrics Layer**: Compute sample-level metrics (entropy, difficulty, staleness, etc.)
+2. **Filter Layer**: Filter samples based on metrics (keep/discard decisions)
+3. **Weight Layer**: Compute per-sample loss weights based on metrics
+
+The FilterWeightManager provides a high-level API to orchestrate these components.
+
+Example usage:
+    ```python
+    from lightrft.trainer.filter_weight import (
+        FilterWeightManager,
+        ResponseLengthFilter,
+        DifficultyWeighting,
+    )
+
+    # Create manager
+    manager = FilterWeightManager(
+        filters=[ResponseLengthFilter(max_length=1024)],
+        weights=[(DifficultyWeighting(mode="prioritized"), 1.0)],
+        enable_metrics={"difficulty": True}
+    )
+
+    # Compute metrics
+    metrics = manager.compute_metrics(outputs)
+
+    # Apply to experiences
+    experiences, weights = manager.apply_to_experiences(experiences, metrics)
+    ```
+
+Author: LightRLHF Team
+"""
+
+# Metrics
+from .metrics import (
+    SampleMetrics,
+    MetricsComputer,
+)
+
+# Filters
+from .filters import (
+    SampleFilter,
+    ResponseLengthFilter,
+    RewardValueFilter,
+    EntropyFilter,
+    DifficultyFilter,
+    CompositeFilter,
+    PercentileFilter,
+)
+
+# Weights
+from .weights import (
+    LossWeighting,
+    ResponseLengthWeighting,
+    EntropyWeighting,
+    DifficultyWeighting,
+    StalenessWeighting,
+    RewardMagnitudeWeighting,
+    CompositeWeighting,
+    UniformWeighting,
+)
+
+# Manager
+from .manager import (
+    FilterWeightManager,
+    FilterWeightManagerBuilder,
+)
+
+
+__all__ = [
+    # ========== Metrics ==========
+    "SampleMetrics",
+    "MetricsComputer",
+
+    # ========== Filters ==========
+    "SampleFilter",
+    "ResponseLengthFilter",
+    "RewardValueFilter",
+    "EntropyFilter",
+    "DifficultyFilter",
+    "CompositeFilter",
+    "PercentileFilter",
+
+    # ========== Weights ==========
+    "LossWeighting",
+    "ResponseLengthWeighting",
+    "EntropyWeighting",
+    "DifficultyWeighting",
+    "StalenessWeighting",
+    "RewardMagnitudeWeighting",
+    "CompositeWeighting",
+    "UniformWeighting",
+
+    # ========== Manager ==========
+    "FilterWeightManager",
+    "FilterWeightManagerBuilder",
+]
+
+
+# Version info
+__version__ = "1.0.0"
+__author__ = "LightRLHF Team"
+
+
+def get_version():
+    """Get the version of the filter_weight module."""
+    return __version__
+
+
+# Quick access functions for common use cases
+def create_length_filter(max_length: int = 1024, **kwargs):
+    """
+    Quick function to create a response length filter.
+
+    Args:
+        max_length: Maximum response length
+        **kwargs: Additional arguments for ResponseLengthFilter
+
+    Returns:
+        ResponseLengthFilter instance
+    """
+    return ResponseLengthFilter(max_length=max_length, **kwargs)
+
+
+def create_difficulty_weighting(mode: str = "prioritized", alpha: float = 0.6, **kwargs):
+    """
+    Quick function to create difficulty weighting.
+
+    Args:
+        mode: Weighting mode ("prioritized" or "curriculum")
+        alpha: Prioritization exponent
+        **kwargs: Additional arguments for DifficultyWeighting
+
+    Returns:
+        DifficultyWeighting instance
+    """
+    return DifficultyWeighting(mode=mode, alpha=alpha, **kwargs)
+
+
+def create_manager_from_args(args, packing_samples: bool = False):
+    """
+    Quick function to create FilterWeightManager from training arguments.
+
+    Args:
+        args: Training arguments
+        packing_samples: Whether samples are packed
+
+    Returns:
+        FilterWeightManager instance
+    """
+    return FilterWeightManagerBuilder.from_args(args, packing_samples)
+
+
+# Add convenience imports at module level
+__all__.extend([
+    "get_version",
+    "create_length_filter",
+    "create_difficulty_weighting",
+    "create_manager_from_args",
+])
+
+
+