chore: import upstream snapshot with attribution

tests/train/test_vllm_importance_sampling_basic.py

@@ -0,0 +1,485 @@
+"""
+Basic tests for vLLM Importance Sampling implementation
+
+This test file verifies the core functionality of the vLLM IS correction,
+including the IS weight computation and metrics calculation.
+
+Reference: verl/verl/trainer/ppo/rollout_corr_helper.py
+"""
+
+import torch
+
+
+class MockAccelerator:
+    """Mock accelerator for testing metrics gathering"""
+
+    def __init__(self, device='cpu'):
+        self.device = device
+
+    def gather_for_metrics(self, tensor):
+        # In testing, just return the tensor as-is
+        return tensor
+
+
+class MockGRPOTrainer:
+    """Mock GRPO trainer for testing IS methods"""
+
+    def __init__(self, mode='token_truncate', threshold=2.0):
+        self.rollout_importance_sampling_mode = mode
+        self.rollout_importance_sampling_threshold = threshold
+        self.accelerator = MockAccelerator()
+
+    def _compute_sequence_level_ratios(self, is_ratio: torch.Tensor, completion_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Helper function to compute sequence-level importance sampling ratios.
+
+        Args:
+            is_ratio: Token-level IS ratios, shape [B, T]
+            completion_mask: Boolean mask for completion tokens, shape [B, T]
+
+        Returns:
+            Sequence-level ratios as geometric mean of token-level ratios
+        """
+        log_ratio = torch.log(is_ratio.clamp(min=1e-10))
+        seq_log_ratios = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)
+        seq_ratios = torch.exp(seq_log_ratios)
+
+        return seq_ratios
+
+    def _apply_rollout_importance_sampling(self, rollout_log_ratio: torch.Tensor,
+                                           completion_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Apply vLLM importance sampling correction using one of four modes.
+
+        Args:
+            rollout_log_ratio: log(π_θ / π_rollout) per token, shape [B, T]
+            completion_mask: Boolean mask for completion tokens, shape [B, T]
+
+        Returns:
+            IS weights to multiply with loss, same shape as rollout_log_ratio
+        """
+        mode = self.rollout_importance_sampling_mode
+        threshold = self.rollout_importance_sampling_threshold
+
+        # Clamp log_ratio to prevent numerical overflow from padding values (-1e10)
+        # A log_ratio of 20 corresponds to exp(20) ≈ 485 million, which is already extreme
+        SAFETY_BOUND = 20.0
+        rollout_log_ratio_safe = torch.clamp(rollout_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
+
+        # Compute importance sampling ratios: exp(log_ratio)
+        is_ratio = torch.exp(rollout_log_ratio_safe)
+
+        if mode == 'token_truncate':
+            # Token-level truncated IS: clip ratios from above at threshold
+            is_weights = torch.clamp(is_ratio, max=threshold)
+
+        elif mode == 'token_mask':
+            # Token-level masked IS: mask out tokens with ratio > threshold
+            is_weights = torch.where(is_ratio <= threshold, is_ratio, torch.zeros_like(is_ratio))
+
+        elif mode == 'sequence_truncate':
+            # Sequence-level truncated IS: compute sequence-level ratio and clip
+            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
+            clipped_seq_ratios = torch.clamp(seq_ratios, max=threshold)
+
+            is_weights = clipped_seq_ratios.unsqueeze(-1).expand_as(is_ratio)
+
+        elif mode == 'sequence_mask':
+            # Sequence-level masked IS: mask entire sequences with ratio > threshold
+            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
+            seq_mask = (seq_ratios <= threshold).float()
+
+            # Apply mask to original token-level ratios
+            is_weights = is_ratio * seq_mask.unsqueeze(-1)
+        else:
+            return is_ratio
+
+        return is_weights
+
+    def _compute_is_correction_metrics(
+        self,
+        vllm_log_ratio: torch.Tensor,
+        is_weights: torch.Tensor,
+        completion_mask: torch.Tensor,
+    ) -> dict:
+        """
+        Compute importance sampling correction metrics (ess, clipped_frac, is_weight_mean).
+        Only called when rollout_importance_sampling_mode is enabled.
+
+        Args:
+            vllm_log_ratio: Log ratio log(π_policy / π_rollout), shape [B, T]
+            is_weights: Importance sampling weights after correction, shape [B, T]
+            completion_mask: Boolean mask for completion tokens, shape [B, T]
+
+        Returns:
+            Dictionary with IS-specific metrics:
+                - is_weight_mean: Mean of IS weights
+                - ess: Effective Sample Size = 1 / E[(w_i / E[w_i])²]
+                - clipped_frac: Fraction of clipped/masked samples
+        """
+        metrics = {}
+        SAFETY_BOUND = 20.0
+        threshold = self.rollout_importance_sampling_threshold
+        threshold_lower = 1.0 / threshold  # Default lower threshold (reciprocal of upper)
+
+        # Helper function for masked mean
+        def masked_mean(x, mask):
+            return (x * mask).sum() / mask.sum().clamp(min=1.0)
+
+        # Compute IS ratio with safety bounds
+        log_ratio_safe = torch.clamp(vllm_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
+        is_ratio = torch.exp(log_ratio_safe)
+
+        # 1. IS weight statistics
+        mean_is_weight = masked_mean(is_weights, completion_mask)
+        metrics['is_weight_mean'] = self.accelerator.gather_for_metrics(mean_is_weight).nanmean().item()
+
+        # 2. Compute Effective Sample Size (ESS) for IS weights
+        # ESS = 1 / E[(w_i / E[w_i])²] (using clamped weights for stability)
+        # This measures how many "effective" independent samples we have after IS weighting
+        weights_for_ess = is_weights.clamp(min=threshold_lower, max=threshold)
+        mean_for_ess = masked_mean(weights_for_ess, completion_mask)
+        is_weights_normalized = weights_for_ess / (mean_for_ess + 1e-8)  # Avoid division by zero
+        ess = 1.0 / masked_mean(is_weights_normalized.square(), completion_mask).clamp(min=1e-10)
+        metrics['ess'] = self.accelerator.gather_for_metrics(ess).nanmean().item()
+
+        # 3. Fraction of clipped/masked samples
+        if self.rollout_importance_sampling_mode in ['token_truncate', 'token_mask']:
+            # Token-level
+            if self.rollout_importance_sampling_mode == 'token_truncate':
+                clipped_frac = masked_mean((is_ratio > threshold).float(), completion_mask)
+            else:  # token_mask
+                clipped_frac = masked_mean((is_weights == 0).float(), completion_mask)
+            metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item()
+        else:
+            # Sequence-level (both truncate and mask)
+            seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask)
+            clipped_frac = (seq_ratios > threshold).float().mean()
+            metrics['clipped_frac'] = self.accelerator.gather_for_metrics(clipped_frac).nanmean().item()
+
+        return metrics
+
+
+class TestVLLMImportanceSampling:
+    """Test suite for vLLM Importance Sampling"""
+
+    def test_token_truncate_basic(self):
+        """Test token-level truncated IS"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Create mock data: [batch=2, seq_len=4]
+        # Log ratios that will produce ratios [0.5, 1.5, 3.0, 5.0]
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0], [0.8, 1.2, 2.5, 4.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Check truncation at threshold=2.0
+        assert is_weights.shape == vllm_log_ratio.shape
+        assert torch.allclose(is_weights[0, 0], torch.tensor(0.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 1], torch.tensor(1.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 2], torch.tensor(2.0), atol=1e-5)  # Truncated
+        assert torch.allclose(is_weights[0, 3], torch.tensor(2.0), atol=1e-5)  # Truncated
+
+    def test_token_mask_basic(self):
+        """Test token-level masked IS"""
+        trainer = MockGRPOTrainer(mode='token_mask', threshold=2.0)
+
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Check masking: ratio > threshold should be 0
+        assert torch.allclose(is_weights[0, 0], torch.tensor(0.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 1], torch.tensor(1.5), atol=1e-5)
+        assert torch.allclose(is_weights[0, 2], torch.tensor(0.0), atol=1e-5)  # Masked
+        assert torch.allclose(is_weights[0, 3], torch.tensor(0.0), atol=1e-5)  # Masked
+
+    def test_sequence_truncate_basic(self):
+        """Test sequence-level truncated IS"""
+        trainer = MockGRPOTrainer(mode='sequence_truncate', threshold=2.0)
+
+        # First sequence has high ratios, second has low ratios
+        vllm_log_ratio = torch.log(
+            torch.tensor([
+                [3.0, 3.0, 3.0, 3.0],  # geometric mean=3.0 > 2.0
+                [1.0, 1.0, 1.0, 1.0]
+            ]))  # geometric mean=1.0 < 2.0
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # First sequence should be truncated to 2.0 for all tokens
+        assert torch.allclose(is_weights[0, :], torch.tensor(2.0), atol=1e-5)
+        # Second sequence should remain 1.0
+        assert torch.allclose(is_weights[1, :], torch.tensor(1.0), atol=1e-5)
+
+    def test_sequence_mask_basic(self):
+        """Test sequence-level masked IS"""
+        trainer = MockGRPOTrainer(mode='sequence_mask', threshold=2.0)
+
+        vllm_log_ratio = torch.log(
+            torch.tensor([
+                [3.0, 3.0, 3.0, 3.0],  # geometric mean=3.0 > 2.0
+                [1.0, 1.0, 1.0, 1.0]
+            ]))  # geometric mean=1.0 < 2.0
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # First sequence should be completely masked (0)
+        # Note: sequence_mask multiplies is_ratio by 0, so all tokens become 0
+        assert torch.allclose(is_weights[0, :], torch.tensor(0.0), atol=1e-5)
+        # Second sequence should keep original ratios (1.0 * 1.0 = 1.0)
+        assert torch.allclose(is_weights[1, :], torch.tensor(1.0), atol=1e-5)
+
+    def test_threshold_sensitivity(self):
+        """Test different threshold values"""
+        vllm_log_ratio = torch.log(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        # Test threshold=1.5
+        trainer_low = MockGRPOTrainer(mode='token_truncate', threshold=1.5)
+        is_weights_low = trainer_low._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Test threshold=3.5
+        trainer_high = MockGRPOTrainer(mode='token_truncate', threshold=3.5)
+        is_weights_high = trainer_high._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Lower threshold should truncate more
+        truncated_low = (is_weights_low < torch.exp(vllm_log_ratio)).sum()
+        truncated_high = (is_weights_high < torch.exp(vllm_log_ratio)).sum()
+        assert truncated_low > truncated_high
+
+    def test_completion_mask(self):
+        """Test that completion mask is respected"""
+        trainer = MockGRPOTrainer(mode='sequence_truncate', threshold=2.0)
+
+        vllm_log_ratio = torch.log(torch.tensor([[3.0, 3.0, 3.0, 3.0]]))
+        # Mask out last two tokens
+        completion_mask = torch.tensor([[1.0, 1.0, 0.0, 0.0]])
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Should only consider masked tokens for sequence ratio calculation
+        # With only first two tokens (both 3.0), geometric mean=3.0, truncated to 2.0
+        assert torch.allclose(is_weights[0, :2], torch.tensor(2.0), atol=1e-5)
+
+    def test_edge_cases(self):
+        """Test edge cases"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Case 1: All ratios below threshold
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.0, 1.5]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+        assert torch.allclose(is_weights, torch.exp(vllm_log_ratio), atol=1e-5)
+
+        # Case 2: All ratios above threshold
+        vllm_log_ratio = torch.log(torch.tensor([[3.0, 4.0, 5.0]]))
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask[:, :3])
+        assert torch.allclose(is_weights, torch.tensor(2.0), atol=1e-5)
+
+        # Case 3: Empty mask
+        vllm_log_ratio = torch.log(torch.tensor([[1.0, 2.0, 3.0]]))
+        completion_mask = torch.zeros_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+        # Should still compute but result may not be meaningful
+        assert is_weights.shape == vllm_log_ratio.shape
+
+    def test_safety_bound(self):
+        """Test that extreme log ratios are clamped"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Create extreme log ratios that would overflow without clamping
+        vllm_log_ratio = torch.tensor([[100.0, -100.0, 0.0]])  # exp(100) would overflow
+        completion_mask = torch.ones_like(vllm_log_ratio)
+
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        # Should not have inf or nan
+        assert torch.isfinite(is_weights).all()
+        # Large positive log_ratio should be clamped to threshold
+        assert is_weights[0, 0] <= 2.0
+        # Large negative log_ratio should result in small positive value
+        assert is_weights[0, 1] > 0
+
+
+class TestISCorrectionMetrics:
+    """Test suite for IS correction metrics"""
+
+    def test_ess_uniform_weights(self):
+        """Test ESS with uniform weights (should be close to 1.0)"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Uniform weights of 1.0
+        vllm_log_ratio = torch.zeros((2, 4))  # exp(0) = 1.0
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = torch.ones_like(vllm_log_ratio)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # ESS should be 1.0 for uniform weights
+        assert abs(metrics['ess'] - 1.0) < 0.01
+        # Mean weight should be 1.0
+        assert abs(metrics['is_weight_mean'] - 1.0) < 0.01
+        # No clipping for uniform weights
+        assert metrics['clipped_frac'] == 0.0
+
+    def test_ess_varied_weights(self):
+        """Test ESS with varied weights (should be < 1.0)"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # Varied weights
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.0, 1.5, 2.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = torch.tensor([[0.5, 1.0, 1.5, 2.0]])
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # ESS should be less than 1.0 for non-uniform weights
+        assert metrics['ess'] < 1.0
+        assert metrics['ess'] > 0.0
+
+    def test_clipped_frac_token_truncate(self):
+        """Test clipped_frac for token_truncate mode"""
+        trainer = MockGRPOTrainer(mode='token_truncate', threshold=2.0)
+
+        # 2 out of 4 tokens exceed threshold
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # 2/4 = 0.5 tokens clipped
+        assert abs(metrics['clipped_frac'] - 0.5) < 0.01
+
+    def test_clipped_frac_token_mask(self):
+        """Test clipped_frac for token_mask mode"""
+        trainer = MockGRPOTrainer(mode='token_mask', threshold=2.0)
+
+        # 2 out of 4 tokens exceed threshold
+        vllm_log_ratio = torch.log(torch.tensor([[0.5, 1.5, 3.0, 5.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # 2/4 = 0.5 tokens masked (is_weights == 0)
+        assert abs(metrics['clipped_frac'] - 0.5) < 0.01
+
+    def test_clipped_frac_sequence_level(self):
+        """Test clipped_frac for sequence-level modes"""
+        trainer = MockGRPOTrainer(mode='sequence_truncate', threshold=2.0)
+
+        # First sequence exceeds threshold, second doesn't
+        vllm_log_ratio = torch.log(torch.tensor([[3.0, 3.0, 3.0, 3.0], [1.0, 1.0, 1.0, 1.0]]))
+        completion_mask = torch.ones_like(vllm_log_ratio)
+        is_weights = trainer._apply_rollout_importance_sampling(vllm_log_ratio, completion_mask)
+
+        metrics = trainer._compute_is_correction_metrics(vllm_log_ratio, is_weights, completion_mask)
+
+        # 1/2 = 0.5 sequences clipped
+        assert abs(metrics['clipped_frac'] - 0.5) < 0.01
+
+
+class TestOffpolicyMetrics:
+    """Test suite for off-policy diagnostic metrics"""
+
+    def test_kl_divergence_same_policy(self):
+        """Test KL divergence when policies are identical"""
+        # When per_token_logps == rollout_per_token_logps, KL should be 0
+        per_token_logps = torch.tensor([[-1.0, -2.0, -1.5, -0.5]])
+        rollout_per_token_logps = per_token_logps.clone()
+        completion_mask = torch.ones_like(per_token_logps)
+
+        # Helper function for masked mean
+        def masked_mean(x, mask, axis=None):
+            if axis is None:
+                return (x * mask).sum() / mask.sum().clamp(min=1.0)
+            else:
+                return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0)
+
+        # KL = E[log(π_rollout) - log(π_training)]
+        kl = masked_mean(rollout_per_token_logps - per_token_logps, completion_mask)
+
+        assert abs(kl.item()) < 1e-6
+
+    def test_k3_kl_estimator(self):
+        """Test K3 KL estimator"""
+        per_token_logps = torch.tensor([[-1.0, -2.0, -1.5, -0.5]])
+        rollout_per_token_logps = torch.tensor([[-1.1, -1.9, -1.6, -0.4]])
+        completion_mask = torch.ones_like(per_token_logps)
+
+        def masked_mean(x, mask, axis=None):
+            if axis is None:
+                return (x * mask).sum() / mask.sum().clamp(min=1.0)
+            else:
+                return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0)
+
+        # K3 estimator: E[exp(log_ratio) - log_ratio - 1]
+        log_ratio = per_token_logps - rollout_per_token_logps
+        log_ratio *= completion_mask
+        k3_kl_matrix = torch.exp(log_ratio) - log_ratio - 1
+        k3_kl = masked_mean(k3_kl_matrix, completion_mask)
+
+        # K3 KL should be non-negative
+        assert k3_kl.item() >= 0
+
+    def test_chi2_divergence(self):
+        """Test χ² divergence calculation"""
+        per_token_logps = torch.tensor([[-1.0, -2.0]])
+        rollout_per_token_logps = torch.tensor([[-1.5, -1.5]])
+        completion_mask = torch.ones_like(per_token_logps)
+
+        def masked_mean(x, mask, axis=None):
+            if axis is None:
+                return (x * mask).sum() / mask.sum().clamp(min=1.0)
+            else:
+                return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0)
+
+        SAFETY_BOUND = 20.0
+        log_ratio = per_token_logps - rollout_per_token_logps
+        log_ratio_safe = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)
+        rho_token = torch.exp(log_ratio_safe)
+        rho_squared_token = rho_token.square()
+        chi2_token = masked_mean(rho_squared_token, completion_mask) - 1.0
+
+        # χ² should be >= -1 (can be negative if E[ρ²] < 1)
+        assert chi2_token.item() >= -1.0
+
+
+if __name__ == '__main__':
+    # Run tests manually
+    import sys
+
+    test_classes = [
+        ('TestVLLMImportanceSampling', TestVLLMImportanceSampling),
+        ('TestISCorrectionMetrics', TestISCorrectionMetrics),
+        ('TestOffpolicyMetrics', TestOffpolicyMetrics),
+    ]
+
+    failed_tests = []
+
+    for class_name, test_class in test_classes:
+        print(f'\n=== {class_name} ===')
+        test_instance = test_class()
+
+        test_methods = [m for m in dir(test_instance) if m.startswith('test_')]
+
+        for method_name in test_methods:
+            try:
+                print(f'Running {method_name}...')
+                getattr(test_instance, method_name)()
+                print(f'✓ {method_name} passed')
+            except Exception as e:
+                print(f'✗ {method_name} failed: {e}')
+                failed_tests.append(f'{class_name}.{method_name}')
+
+    if failed_tests:
+        print(f'\nFailed tests: {failed_tests}')
+        sys.exit(1)
+    else:
+        print('\nAll tests passed!')