Upload diffusion_llm/training/dapo.py with huggingface_hub

diffusion_llm/training/dapo.py

@@ -0,0 +1,187 @@
+"""AAM Diffusion LLM — DAPO Training
+
+Decoupled Clip & Dynamic Sampling Policy Optimization (Yu et al., 2025).
+Four improvements over GRPO:
+1. Decoupled Clip (asymmetric epsilon)
+2. Dynamic Sampling (filter zero-variance groups)
+3. Token-Level Policy Gradient Loss
+4. Overlong Filtering
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DAPOConfig:
+    clip_ratio_low: float = 0.2
+    clip_ratio_high: float = 0.28
+    dynamic_sampling: bool = True
+    token_level_loss: bool = True
+    overlong_filter: bool = True
+    max_response_length: int = 2048
+    num_responses_per_prompt: int = 8
+    kl_coefficient: float = 0.1
+    discount_factor: float = 1.0
+    use_reward_normalization: bool = True
+    use_advantage_normalization: bool = True
+    learning_rate: float = 1e-6
+    reference_model_freeze: bool = True
+    entropy_coefficient: float = 0.01
+    max_grad_norm: float = 1.0
+    temperature: float = 0.7
+    reward_shaping: str = "centered"
+
+    def __post_init__(self) -> None:
+        if self.clip_ratio_low <= 0:
+            raise ValueError(f"clip_ratio_low must be positive, got {self.clip_ratio_low}")
+        if self.clip_ratio_high <= 0:
+            raise ValueError(f"clip_ratio_high must be positive, got {self.clip_ratio_high}")
+        if self.num_responses_per_prompt < 2:
+            raise ValueError(f"num_responses_per_prompt must be >= 2, got {self.num_responses_per_prompt}")
+
+
+class DAPOTrainer:
+    """DAPO Trainer for AAM Diffusion LLM."""
+
+    def __init__(
+        self,
+        config: DAPOConfig,
+        policy_model: nn.Module,
+        reference_model: Optional[nn.Module] = None,
+        reward_fn: Optional[Callable] = None,
+        optimizer: Optional[torch.optim.Optimizer] = None,
+    ) -> None:
+        self.config = config
+        self.policy_model = policy_model
+        self.reward_fn = reward_fn
+
+        if reference_model is not None:
+            self.reference_model = reference_model
+        elif config.kl_coefficient > 0:
+            self.reference_model = copy.deepcopy(policy_model)
+        else:
+            self.reference_model = None
+
+        if self.reference_model is not None and config.reference_model_freeze:
+            for param in self.reference_model.parameters():
+                param.requires_grad = False
+
+        trainable_params = [p for p in policy_model.parameters() if p.requires_grad]
+        self.optimizer = optimizer or torch.optim.AdamW(
+            trainable_params, lr=config.learning_rate, betas=(0.9, 0.95), weight_decay=0.01,
+        )
+
+        self.device = next(policy_model.parameters()).device
+
+    def compute_dapo_loss(
+        self,
+        log_probs: torch.Tensor,
+        old_log_probs: torch.Tensor,
+        ref_log_probs: torch.Tensor,
+        rewards: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> Tuple[torch.Tensor, Dict[str, float]]:
+        cfg = self.config
+
+        log_ratio = log_probs - old_log_probs
+        ratio = torch.exp(log_ratio)
+
+        advantages = self._compute_advantages(rewards)
+        advantages_expanded = advantages.unsqueeze(-1).expand_as(log_probs) if advantages.dim() == 1 else advantages
+
+        clipped_ratio = torch.clamp(ratio, 1.0 - cfg.clip_ratio_low, 1.0 + cfg.clip_ratio_high)
+
+        surr1 = ratio * advantages_expanded
+        surr2 = clipped_ratio * advantages_expanded
+
+        if cfg.token_level_loss:
+            per_token_loss = -torch.min(surr1, surr2) * attention_mask
+            num_valid_tokens = attention_mask.sum(dim=-1, keepdim=True).clamp(min=1)
+            policy_loss = (per_token_loss.sum(dim=-1) / num_valid_tokens.squeeze(-1)).mean()
+        else:
+            per_token_loss = -torch.min(surr1, surr2) * attention_mask
+            seq_loss = per_token_loss.sum(dim=-1) / attention_mask.sum(dim=-1).clamp(min=1)
+            policy_loss = seq_loss.mean()
+
+        kl_penalty = torch.tensor(0.0, device=log_probs.device)
+        if ref_log_probs is not None and cfg.kl_coefficient > 0:
+            kl_per_token = torch.exp(log_probs) * (log_probs - ref_log_probs) * attention_mask
+            kl_penalty = cfg.kl_coefficient * (kl_per_token.sum(dim=-1) / attention_mask.sum(dim=-1).clamp(min=1)).mean()
+
+        entropy = torch.tensor(0.0, device=log_probs.device)
+        if cfg.entropy_coefficient > 0:
+            per_token_entropy = -torch.exp(log_probs) * log_probs * attention_mask
+            entropy = (per_token_entropy.sum(dim=-1) / attention_mask.sum(dim=-1).clamp(min=1)).mean()
+
+        loss = policy_loss + kl_penalty - cfg.entropy_coefficient * entropy
+
+        with torch.no_grad():
+            metrics = {
+                "dapo/policy_loss": policy_loss.item(),
+                "dapo/kl_penalty": kl_penalty.item() if isinstance(kl_penalty, torch.Tensor) else kl_penalty,
+                "dapo/entropy": entropy.item() if isinstance(entropy, torch.Tensor) else entropy,
+                "dapo/loss": loss.item(),
+                "dapo/mean_reward": rewards.mean().item(),
+            }
+
+        return loss, metrics
+
+    def _compute_advantages(self, rewards: torch.Tensor) -> torch.Tensor:
+        cfg = self.config
+        if cfg.use_reward_normalization and rewards.numel() > 1:
+            rewards = self._shape_rewards(rewards, cfg.reward_shaping)
+        advantages = rewards.clone()
+        if cfg.use_advantage_normalization and advantages.numel() > 1:
+            adv_std = advantages.std()
+            if adv_std > 1e-8:
+                advantages = (advantages - advantages.mean()) / (adv_std + 1e-8)
+        return advantages
+
+    def _shape_rewards(self, rewards: torch.Tensor, strategy: str) -> torch.Tensor:
+        if strategy == "raw":
+            return rewards
+        if strategy == "centered":
+            return rewards - rewards.mean()
+        if strategy == "rank_based":
+            sorted_indices = rewards.argsort()
+            ranks = torch.zeros_like(rewards, dtype=torch.float32)
+            ranks[sorted_indices] = torch.arange(len(rewards), dtype=torch.float32, device=rewards.device) / max(len(rewards) - 1, 1)
+            return 2.0 * ranks - 1.0
+        return rewards
+
+    def filter_prompts(
+        self,
+        prompts: List[str],
+        responses: List[List[str]],
+        rewards: torch.Tensor,
+    ) -> Tuple[List[str], List[List[str]], torch.Tensor, Dict[str, int]]:
+        if not self.config.dynamic_sampling:
+            return prompts, responses, rewards, {"filtered": 0, "total": len(prompts)}
+
+        if rewards.dim() == 1:
+            has_variance = rewards > 1e-6
+        else:
+            reward_std_per_prompt = rewards.std(dim=-1)
+            has_variance = reward_std_per_prompt > 1e-6
+
+        valid_indices = has_variance.nonzero(as_tuple=True)[0]
+        if len(valid_indices) == 0:
+            return prompts, responses, rewards, {"filtered": len(prompts), "total": len(prompts)}
+
+        filtered_prompts = [prompts[i] for i in valid_indices]
+        filtered_responses = [responses[i] for i in valid_indices]
+        filtered_rewards = rewards[valid_indices]
+        num_filtered = len(prompts) - len(valid_indices)
+
+        return filtered_prompts, filtered_responses, filtered_rewards, {"filtered": num_filtered, "total": len(prompts)}