Merge pull request #41 from modelscope/fix_npu_grpo

Fix npu grpo

Author: addsubmuldiv

cookbook/grpo/lora_npu.py

@@ -9,14 +9,15 @@
 from twinkle.model import TransformersModel
 from twinkle.reward import MathReward
 from twinkle.sampler import VLLMSampler, TorchSampler
-from twinkle.sampler.types import SamplingParams
+from twinkle.sampler.types import SamplingParams, SampleResponse
 from twinkle.weight_loader import NativeLoader
 from twinkle.rl import compute_advantages
 
 # Environment variable setup
 os.environ.setdefault('TRUST_REMOTE_CODE', '1')
 os.environ.setdefault('TWINKLE_SEED', '42')
 os.environ.setdefault('TWINKLE_FULL_DETERMINISM', '1')
+os.environ.setdefault('RAY_TMPDIR', os.path.expanduser('~/tmp/ray'))
 
 # Training configuration
 use_ref_model = os.environ.get('TWINKLE_USE_REF_MODEL', '1') != '0'
@@ -129,6 +130,29 @@ def get_sampling_params(eos_token_ids) -> SamplingParams:
     )
 
 
+def build_trajectories_from_sample_response(sample_response: SampleResponse, batch_list, tokenizer):
+    """Convert sampler output into GRPO trajectories."""
+    if not sample_response or not getattr(sample_response, 'sequences', None):
+        return []
+    if not batch_list:
+        return []
+
+    trajectories = []
+    for i, seq in enumerate(sample_response.sequences):
+        src_batch = batch_list[i % len(batch_list)]
+        src_messages = [dict(msg) for msg in src_batch.get('messages', [])]
+        if src_messages and src_messages[-1].get('role') == 'assistant':
+            # Remove reference answer and append sampled assistant reply.
+            src_messages = src_messages[:-1]
+
+        response_text = tokenizer.decode(seq.tokens, skip_special_tokens=True) if tokenizer is not None else ''
+        trajectories.append({
+            'messages': src_messages + [{'role': 'assistant', 'content': response_text}],
+            'user_data': list(src_batch.get('user_data', [])),
+        })
+    return trajectories
+
+
 def debug_print_rollout(step, trajectories, ground_truths, rewards=None):
     """Debug helper that prints rollout intermediates (sampling, rewards, etc.).
 
@@ -182,6 +206,19 @@ def debug_print_rollout(step, trajectories, ground_truths, rewards=None):
         )
 
 
+def _collect_sample_responses(results):
+    """Custom collect function to merge multiple SampleResponse objects."""
+    if not results:
+        return SampleResponse(sequences=[])
+    if len(results) == 1:
+        return results[0]
+    all_sequences = []
+    for resp in results:
+        if resp is not None and hasattr(resp, 'sequences'):
+            all_sequences.extend(resp.sequences)
+    return SampleResponse(sequences=all_sequences)
+
+
 @remote_class()
 class ActorGroup:
 
@@ -226,7 +263,7 @@ def __init__(self, engine_args=None, lora_config=None, adapter_name=None, **kwar
         self.adapter_name = adapter_name
         self.lora_config = lora_config
 
-    @remote_function(collect='flatten')
+    @remote_function(collect=_collect_sample_responses)
     def sample(self, batch, sampling_params: SamplingParams = None):
         return self.sampler.sample(batch, sampling_params=sampling_params, adapter_name=self.adapter_name)
 
@@ -293,6 +330,11 @@ def train():
     )
 
     eos_token_ids = get_eos_token_ids()
+    try:
+        from transformers import AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    except Exception:
+        tokenizer = None
 
     engine_args = {
         'model': model_path,
@@ -339,13 +381,18 @@ def train():
             batch_list = [batch]
         else:
             batch_list = list(batch)
-        ground_truths = batch_list.copy()
-        
         sampling_params = get_sampling_params(eos_token_ids)
 
-        trajectories = actor_group.sample(batch_list, sampling_params)
-        if callable(trajectories):
-            trajectories = trajectories()
+        sample_response = actor_group.sample(batch_list, sampling_params)
+        if callable(sample_response):
+            sample_response = sample_response()
+        trajectories = build_trajectories_from_sample_response(sample_response, batch_list, tokenizer)
+        if not trajectories:
+            print(f'[step {step}] empty sampled trajectories, skip.', flush=True)
+            continue
+
+        # Expand ground truths to align with sampled trajectory count.
+        ground_truths = [batch_list[i % len(batch_list)] for i in range(len(trajectories))]
 
         ref_logits = None
         if use_ref_model:
@@ -357,14 +404,19 @@ def train():
             else:
                 ref_logits = ref_outputs['logits'] if isinstance(ref_outputs, dict) else ref_outputs.logits
 
-        rewards = reward.calculate(trajectories, ground_truths)
+        rewards = reward(trajectories, ground_truths)
         if callable(rewards):
             rewards = rewards()
 
-        # Updated: compute advantages from rewards and store in trajectory
-        advantages = compute_advantages(rewards, num_generations=num_generations)
+        effective_num_generations = num_generations if len(rewards) % num_generations == 0 else 1
+        scale = 'group' if effective_num_generations > 1 else 'batch'
+        advantages = compute_advantages(
+            rewards,
+            num_generations=effective_num_generations,
+            scale=scale,
+        )
         for trajectory, advantage in zip(trajectories, advantages.tolist()):
-            trajectory['advantages'] = advantage
+            trajectory['advantages'] = float(advantage)
 
         # Debug: print reward statistics (enable via TWINKLE_DEBUG=1)
         debug_print_rollout(step, trajectories, ground_truths, rewards=rewards)
@@ -383,5 +435,6 @@ def train():
             break
 
 
+
 if __name__ == '__main__':
     train()

src/twinkle/rl/__init__.py

@@ -2,3 +2,35 @@
 from .base import Advantage
 from .grpo import GRPOAdvantage
 from .rloo import RLOOAdvantage
+
+
+# TODO: Temporary helpers added to unblock cookbook/grpo examples.
+# Each call creates a new Advantage instance, not suitable for production.
+# Remove once the framework provides a proper advantage computation API.
+def compute_advantages(rewards, num_generations=1, scale='group', **kwargs):
+    """Backward-compatible helper for GRPO advantage computation."""
+    return GRPOAdvantage()(
+        rewards=rewards,
+        num_generations=num_generations,
+        scale=scale,
+        **kwargs,
+    )
+
+
+def compute_advantages_rloo(rewards, num_generations=1, scale='group', **kwargs):
+    """Backward-compatible helper for RLOO advantage computation."""
+    return RLOOAdvantage()(
+        rewards=rewards,
+        num_generations=num_generations,
+        scale=scale,
+        **kwargs,
+    )
+
+
+__all__ = [
+    'Advantage',
+    'GRPOAdvantage',
+    'RLOOAdvantage',
+    'compute_advantages',
+    'compute_advantages_rloo',
+]

src/twinkle/sampler/vllm_engine.py

@@ -160,7 +160,16 @@ def _create_engine(self):
 
         logger.info(f"VLLMEngine initialized: model={self.model_id}")
         return engine
-    
+
+    def shutdown(self):
+        """Shutdown the underlying vLLM AsyncLLM engine."""
+        if hasattr(self, 'engine') and self.engine is not None:
+            try:
+                self.engine.shutdown()
+                logger.info("VLLMEngine shutdown completed.")
+            except Exception as e:
+                logger.warning(f"VLLMEngine shutdown error: {e}")
+
     async def get_tokenizer(self):
         """Get the tokenizer asynchronously."""
         if self._tokenizer is None:

src/twinkle/sampler/vllm_sampler.py

@@ -20,6 +20,7 @@
     - Results are collected via collect='flatten' (merged into single list)
 """
 import asyncio
+import atexit
 import logging
 import os
 import threading
@@ -137,7 +138,10 @@ def __init__(
         )
 
         VLLMLoraWeights().patch(self)
-    
+
+        self._shutdown_called = False
+        atexit.register(self.shutdown)
+
     def _run_event_loop(self):
         """Run the event loop in background thread."""
         asyncio.set_event_loop(self._async_loop)
@@ -409,3 +413,30 @@ def wake_up(self, tags: List[str] = None, reload_weights: bool = False) -> None:
                   Required after level 2 sleep which discards weights.
         """
         self._run_in_loop(self.engine.wake_up(tags=tags, reload_weights=reload_weights))
+
+    def shutdown(self):
+        """Gracefully shutdown the vLLM engine and background event loop.
+
+        Registered via atexit so it runs automatically on process exit,
+        before GC destroys objects in unpredictable order. Safe to call
+        multiple times (idempotent).
+        """
+        if self._shutdown_called:
+            return
+        self._shutdown_called = True
+
+        # 1. Shutdown vLLM engine (stops EngineCore process and output_handler)
+        try:
+            if hasattr(self, 'engine') and self.engine is not None:
+                self.engine.shutdown()
+        except Exception as e:
+            logger.warning(f"VLLMSampler engine shutdown error: {e}")
+
+        # 2. Stop the background event loop and join thread
+        try:
+            if hasattr(self, '_async_loop') and self._async_loop.is_running():
+                self._async_loop.call_soon_threadsafe(self._async_loop.stop)
+            if hasattr(self, '_async_thread') and self._async_thread.is_alive():
+                self._async_thread.join(timeout=5)
+        except Exception as e:
+            logger.warning(f"VLLMSampler event loop shutdown error: {e}")