update

Author: Yunnglin

cookbook/client/tinker/transformer/grpo.py

@@ -182,68 +182,98 @@ def main():
             step += 1
             continue
 
-        # ========== 6. Training step ==========
-        # Select samples with positive advantages for training
-        # Weight them by their advantage value for GRPO-style optimization
-        training_data = []
-        for i, seq in enumerate(all_sequences):
-            if advantages[i] <= 0:
-                continue
-            # Build a Datum from the completion tokens
-            # Prompt tokens: weight=0 (don't compute loss on prompt)
-            # Completion tokens: weight=advantage (advantage-weighted SFT)
-            prompt_feature = prompts[i // NUM_GENERATIONS]
-            prompt_ids = prompt_feature['input_ids']
-            if hasattr(prompt_ids, 'tolist'):
-                prompt_ids = prompt_ids.tolist()
-
-            full_tokens = prompt_ids + list(seq.tokens)
-            prompt_weights = [0.0] * len(prompt_ids)
-            # Scale completion weights by normalized advantage
-            completion_weights = [float(advantages[i])] * len(seq.tokens)
-
-            # Shift by one for next-token prediction
-            input_tokens = full_tokens[:-1]
-            target_tokens = full_tokens[1:]
-            weights = (prompt_weights + completion_weights)[1:]
-
-            datum = types.Datum(
-                model_input=types.ModelInput.from_ints(input_tokens),
-                loss_fn_inputs={
-                    'target_tokens': target_tokens,
-                    'weights': weights,
-                },
-            )
-            training_data.append(datum)
+    # Train the policies with the Advantage-Regularized policy 
+    # gradient (GRPO) loss function.
+    # 
+    # The GRPO loss function requires:
+    # 1. logprobs: The log probabilities of the tokens under the current policy
+    # 2. advantages: The advantage values for each completion
+    # 
+    # The training data is constructed with:
+    # - model_input: The full prompt + completion tokens
+    # - target_tokens: The shifted tokens for next-token prediction
+    # - logprobs: The log probabilities from the sampling step
+    # - advantages: The computed advantage values
+    training_data = []
+    for i, seq in enumerate(all_sequences):
+        # Build a Datum from the completion tokens with logprobs and advantages
+        prompt_feature = prompts[i // NUM_GENERATIONS]
+        prompt_ids = prompt_feature['input_ids']
+        if hasattr(prompt_ids, 'tolist'):
+            prompt_ids = prompt_ids.tolist()
+
+        full_tokens = prompt_ids + list(seq.tokens)
+        
+        # Shift by one for next-token prediction
+        input_tokens = full_tokens[:-1]
+        target_tokens = full_tokens[1:]
+        
+        # Get logprobs from the sampling result
+        logprobs = seq.logprobs if seq.logprobs else [0.0] * len(seq.tokens)
+        # Pad logprobs to match full sequence length (prompt + completion)
+        # Prompt positions get 0.0 logprobs (no loss computed on prompt)
+        padded_logprobs = [0.0] * len(prompt_ids) + logprobs
+        
+        # Get advantage for this sequence
+        advantage = float(advantages[i])
+        
+        # Pad advantages to match full sequence length
+        # Only completion tokens get the advantage value, prompt gets 0.0
+        padded_advantages = [0.0] * len(prompt_ids) + [advantage] * len(seq.tokens)
+        
+        # Verify lengths match
+        assert len(input_tokens) == len(target_tokens) == len(padded_logprobs) == len(padded_advantages), \
+            f"Length mismatch: input={len(input_tokens)}, target={len(target_tokens)}, " \
+            f"logprobs={len(padded_logprobs)}, advantages={len(padded_advantages)}"
+
+        datum = types.Datum(
+            model_input=types.ModelInput.from_ints(input_tokens),
+            loss_fn_inputs={
+                'target_tokens': target_tokens,
+                'logprobs': types.TensorData.from_numpy(np.array(padded_logprobs, dtype=np.float32)),
+                'advantages': types.TensorData.from_numpy(np.array(padded_advantages, dtype=np.float32)),
+            },
+        )
+        training_data.append(datum)
 
         if not training_data:
             logger.info(
-                f"Step {step}: No positive-advantage samples, skipping")
+                f"Step {step}: No training data constructed, skipping")
             step += 1
             continue
 
-        # Forward-backward pass with cross-entropy on advantage-weighted data
+        # Forward-backward pass with importance_sampling (GRPO) loss
+        # The training data already contains logprobs and advantages for the GRPO loss
         fwdbwd_future = training_client.forward_backward(
-            training_data, "cross_entropy")
+            training_data, "importance_sampling")
         optim_future = training_client.optim_step(
             types.AdamParams(learning_rate=LEARNING_RATE))
-
+        
         fwdbwd_result = fwdbwd_future.result()
         optim_result = optim_future.result()
 
-        # Compute weighted average loss for monitoring
-        logprobs = np.concatenate(
-            [output['logprobs'].tolist()
-             for output in fwdbwd_result.loss_fn_outputs])
-        weights = np.concatenate(
-            [d.loss_fn_inputs['weights'].tolist() for d in training_data])
-        loss_per_token = -np.dot(logprobs, weights) / max(weights.sum(), 1e-8)
+        # Compute metrics from the forward-backward result
+        # For importance_sampling, we get logprobs and elementwise_loss
+        logprobs_list = []
+        elementwise_losses = []
+        for output in fwdbwd_result.loss_fn_outputs:
+            if output.get('logprobs') is not None:
+                logprobs_list.append(output['logprobs'].to_numpy())
+            if output.get('elementwise_loss') is not None:
+                elementwise_losses.append(output['elementwise_loss'].to_numpy())
+        
+        # Compute average loss per token (weighted by advantages)
+        if elementwise_losses:
+            all_losses = np.concatenate(elementwise_losses)
+            avg_loss = np.mean(all_losses) if len(all_losses) > 0 else 0.0
+        else:
+            avg_loss = 0.0
 
         gc.collect()
 
         # ========== 7. Log ==========
         log_dict = metrics.calculate()
-        log_dict['train/loss_per_token'] = loss_per_token
+        log_dict['train/loss_per_token'] = float(avg_loss)
         log_dict['train/frac_reward_zero_std'] = frac_zero_std
         log_dict['train/num_training_samples'] = len(training_data)
         logger.info(f"Step {step}: {log_dict}")

src/twinkle/server/tinker/common/__init__.py

@@ -1,3 +1,4 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from .datum import datum_to_input_feature, input_feature_to_datum
+from .transformers_model import _extract_rl_fields_from_inputs as extract_rl_fields_from_inputs
 from twinkle.utils import exists, requires

src/twinkle/server/tinker/common/datum.py

@@ -22,6 +22,17 @@ def datum_to_input_feature(datum: types.Datum) -> InputFeature:
         labels = datum.loss_fn_inputs['target_tokens'].to_numpy()
 
         input_feature['labels'] = np.where(weights > 0, labels, -100).tolist()
+    
+    # 3. Handle importance_sampling specific fields
+    # 'logprobs' -> 'old_logps' (for GRPO loss)
+    if 'logprobs' in datum.loss_fn_inputs:
+        old_logps = datum.loss_fn_inputs['logprobs'].to_numpy().tolist()
+        input_feature['old_logps'] = old_logps
+    
+    # 'advantages' -> 'advantages' (for GRPO loss)
+    if 'advantages' in datum.loss_fn_inputs:
+        advantages = datum.loss_fn_inputs['advantages'].to_numpy().tolist()
+        input_feature['advantages'] = advantages
 
     return input_feature
 

src/twinkle/server/tinker/common/megatron_model.py

@@ -2,12 +2,53 @@
 import numpy as np
 import torch
 from tinker import types
-from typing import List, TYPE_CHECKING
+from typing import List, TYPE_CHECKING, Tuple, Optional, Any
 from twinkle import remote_class, remote_function
 from twinkle.utils import exists, requires
 from .datum import datum_to_input_feature
 from .io_utils import create_checkpoint_manager
 
+
+def _extract_rl_fields_from_inputs(
+    input_features: List[dict], 
+    kwargs: dict
+) -> Tuple[Optional[List], Optional[List], dict]:
+    """Extract old_logps and advantages from input features and kwargs.
+    
+    This function handles the common logic for extracting reinforcement learning
+    fields (old_logps and advantages) from both input features and kwargs.
+    
+    Args:
+        input_features: List of input feature dictionaries
+        kwargs: Keyword arguments dictionary
+        
+    Returns:
+        Tuple of (old_logps, advantages, updated_kwargs)
+    """
+    # Extract from kwargs first (higher priority)
+    old_logps = kwargs.pop('old_logps', None)
+    advantages = kwargs.pop('advantages', None)
+    
+    # If not in kwargs, check input features
+    if old_logps is None:
+        old_logps_list = [inp.get('old_logps') for inp in input_features if inp.get('old_logps') is not None]
+        if old_logps_list:
+            old_logps = old_logps_list
+    
+    if advantages is None:
+        advantages_list = [inp.get('advantages') for inp in input_features if inp.get('advantages') is not None]
+        if advantages_list:
+            advantages = advantages_list
+    
+    # Prepare kwargs for loss function
+    loss_kwargs = kwargs.copy()
+    if old_logps is not None:
+        loss_kwargs['old_logps'] = old_logps
+    if advantages is not None:
+        loss_kwargs['advantages'] = advantages
+        
+    return old_logps, advantages, loss_kwargs
+
 if TYPE_CHECKING:
     from twinkle.model.megatron import MultiLoraMegatronModel as _MegatronBase
 elif exists('megatron_core'):
@@ -83,9 +124,12 @@ def forward_backward(self, *, inputs: List[types.Datum], **kwargs):
         # Convert Datum to InputFeature
         input_features = [datum_to_input_feature(datum) for datum in inputs]
 
+        # Extract old_logps and advantages using common utility
+        old_logps, advantages, loss_kwargs = _extract_rl_fields_from_inputs(input_features, kwargs)
+        
         adapter_name = kwargs.get('adapter_name')
         # Megatron forward_backward returns loss directly
-        loss = super().forward_backward(inputs=input_features, **kwargs)
+        loss = super().forward_backward(inputs=input_features, **loss_kwargs)
 
         # Get logits from outputs
         optimizer_config = self.optimizer_group.get(adapter_name)

src/twinkle/server/tinker/common/transformers_model.py

@@ -1,11 +1,52 @@
 import torch
 from tinker import types
-from typing import List
+from typing import List, Tuple, Optional, Any
 from twinkle.model import MultiLoraTransformersModel
 from twinkle import remote_class, remote_function
 from .datum import datum_to_input_feature
 from .io_utils import create_checkpoint_manager
 
+
+def _extract_rl_fields_from_inputs(
+    input_features: List[dict], 
+    kwargs: dict
+) -> Tuple[Optional[List], Optional[List], dict]:
+    """Extract old_logps and advantages from input features and kwargs.
+    
+    This function handles the common logic for extracting reinforcement learning
+    fields (old_logps and advantages) from both input features and kwargs.
+    
+    Args:
+        input_features: List of input feature dictionaries
+        kwargs: Keyword arguments dictionary
+        
+    Returns:
+        Tuple of (old_logps, advantages, updated_kwargs)
+    """
+    # Extract from kwargs first (higher priority)
+    old_logps = kwargs.pop('old_logps', None)
+    advantages = kwargs.pop('advantages', None)
+    
+    # If not in kwargs, check input features
+    if old_logps is None:
+        old_logps_list = [inp.get('old_logps') for inp in input_features if inp.get('old_logps') is not None]
+        if old_logps_list:
+            old_logps = old_logps_list
+    
+    if advantages is None:
+        advantages_list = [inp.get('advantages') for inp in input_features if inp.get('advantages') is not None]
+        if advantages_list:
+            advantages = advantages_list
+    
+    # Prepare kwargs for loss function
+    loss_kwargs = kwargs.copy()
+    if old_logps is not None:
+        loss_kwargs['old_logps'] = old_logps
+    if advantages is not None:
+        loss_kwargs['advantages'] = advantages
+        
+    return old_logps, advantages, loss_kwargs
+
 @remote_class()
 class TwinkleCompatTransformersModel(MultiLoraTransformersModel):
     """
@@ -45,6 +86,11 @@ def forward(self, *, inputs: List[types.Datum], **kwargs):
         # Convert Datum to InputFeature
         input_features = [datum_to_input_feature(datum) for datum in inputs]
 
+        # Extract old_logps and advantages using common utility
+        old_logps, advantages, loss_kwargs = _extract_rl_fields_from_inputs(input_features, kwargs)
+        # Update kwargs for forward pass (exclude loss-specific fields)
+        kwargs.update({k: v for k, v in loss_kwargs.items() if k not in ['old_logps', 'advantages']})
+       
         outputs = super().forward(inputs=input_features, **kwargs)
         logits = outputs['logits'].detach().cpu()  # shape (batch_size, seq_len, vocab_size)
         results = self._get_forward_output(inputs, logits)
@@ -54,6 +100,11 @@ def forward(self, *, inputs: List[types.Datum], **kwargs):
     def forward_only(self, *, inputs: List[types.Datum], **kwargs):
         # Convert Datum to InputFeature
         input_features = [datum_to_input_feature(datum) for datum in inputs]
+        
+        # Extract old_logps and advantages using common utility
+        old_logps, advantages, loss_kwargs = _extract_rl_fields_from_inputs(input_features, kwargs)
+        # Update kwargs for forward pass (exclude loss-specific fields)
+        kwargs.update({k: v for k, v in loss_kwargs.items() if k not in ['old_logps', 'advantages']})
 
         outputs = super().forward_only(inputs=input_features, **kwargs)
         logits = outputs['logits'].detach().cpu()  # shape (batch_size, seq_len, vocab_size)
@@ -62,9 +113,34 @@ def forward_only(self, *, inputs: List[types.Datum], **kwargs):
 
     @remote_function(collect='mean')
     def calculate_loss(self, **kwargs):
-        loss = super().calculate_loss(**kwargs)
+        # Extract old_logps and advantages using common utility (for importance_sampling loss)
+        # Note: We don't need the input_features here since this is called separately
+        old_logps, advantages, loss_kwargs = _extract_rl_fields_from_inputs([], kwargs)
+            
+        loss = super().calculate_loss(**loss_kwargs)
         return loss
 
+    @remote_function(dispatch='slice_dp', collect='flatten')
+    def forward_backward(self, *, inputs: List[types.Datum], **kwargs):
+        # Convert Datum to InputFeature
+        input_features = [datum_to_input_feature(datum) for datum in inputs]
+        
+        # Extract old_logps and advantages using common utility
+        old_logps, advantages, loss_kwargs = _extract_rl_fields_from_inputs(input_features, kwargs)
+        
+        # Forward pass
+        outputs = super().forward(inputs=input_features, **kwargs)
+        
+        # Calculate loss with extra parameters
+        loss = super().calculate_loss(**loss_kwargs)
+        
+        # Backward pass
+        super().backward(**kwargs)
+        
+        logits = outputs['logits'].detach().cpu()  # shape (batch_size, seq_len, vocab_size)
+        results = self._get_forward_output(inputs, logits)
+        return results, loss
+
     @remote_function()
     def step(self, *, adam_params: types.AdamParams, **kwargs):
         # Gradient clipping

src/twinkle/server/tinker/model.py

@@ -362,17 +362,34 @@ async def _do_forward_backward():
 
                     if self.use_megatron:
                         # Megatron uses combined forward_backward, no separate backward/calculate_loss
+                        # Set loss first based on loss_fn
+                        if loss_fn == 'cross_entropy':
+                            self.model.set_loss('CrossEntropyLoss',
+                                                adapter_name=adapter_name)
+                        elif loss_fn == 'importance_sampling':
+                            self.model.set_loss('GRPOLoss',
+                                                adapter_name=adapter_name,
+                                                epsilon=0.2,  # Default GRPO epsilon
+                                                beta=0.0)     # No KL penalty by default
+                        else:
+                            raise ValueError(
+                                f'Unsupported loss function {loss_fn}')
+                        
                         output, loss = self.model.forward_backward(
                             inputs=datum_list,
                             adapter_name=adapter_name,
                             **loss_fn_config)
                     else:
                         # Transformers uses separate forward, calculate_loss, backward
-                        # When use_megatron is True, we don't need to set the loss
-                        # Set loss first
+                        # Set loss first based on loss_fn
                         if loss_fn == 'cross_entropy':
                             self.model.set_loss('CrossEntropyLoss',
                                                 adapter_name=adapter_name)
+                        elif loss_fn == 'importance_sampling':
+                            self.model.set_loss('GRPOLoss',
+                                                adapter_name=adapter_name,
+                                                epsilon=0.2,  # Default GRPO epsilon
+                                                beta=0.0)     # No KL penalty by default
                         else:
                             raise ValueError(
                                 f'Unsupported loss function {loss_fn}')
@@ -382,8 +399,9 @@ async def _do_forward_backward():
                         loss = self.model.calculate_loss(adapter_name=adapter_name,
                                                          **loss_fn_config)
                         self.model.backward(adapter_name=adapter_name)
+                    output_type = 'ImportanceSamplingLossReturn' if loss_fn == 'importance_sampling' else 'CrossEntropyLossReturn'
                     return types.ForwardBackwardOutput(
-                        loss_fn_output_type='CrossEntropyLossReturn',
+                        loss_fn_output_type=output_type,
                         loss_fn_outputs=output,
                         metrics={'loss:avg': loss},
                     )