update

Author: Yunnglin

cookbook/client/tinker/self_host/short_math_grpo.py

@@ -217,7 +217,7 @@ def main():
     from tinker import ServiceClient
     service_client = ServiceClient(
         base_url='http://localhost:8000',
-        api_key=os.environ.get('MODELSCOPE_TOKEN')
+        api_key='EMPTY_TOKEN'
     )
 
     logger.info('Creating LoRA training client...')

cookbook/client/twinkle/self_host/grpo.py

@@ -103,12 +103,13 @@ def train():
     model.set_loss('GRPOLoss', epsilon=0.2, beta=0.0)
 
     # Set optimizer and LR scheduler
-    model.set_optimizer('AdamW', lr=LEARNING_RATE)
-    model.set_lr_scheduler(
-        'CosineWarmupScheduler',
-        num_warmup_steps=500,
-        num_training_steps=MAX_STEPS,
-    )
+    model.set_optimizer('Adam', lr=LEARNING_RATE)
+    # Set LR scheduler (if server use megatron, don't support set self.lr_scheduler)
+    # model.set_lr_scheduler(
+    #     'CosineWarmupScheduler',
+    #     num_warmup_steps=500,
+    #     num_training_steps=MAX_STEPS,
+    # )
 
     # Set processor and template for encoding inputs
     model.set_processor('InputProcessor')

src/twinkle/server/model/backends/common.py

@@ -5,9 +5,10 @@
 import numpy as np
 import re
 import torch
+from collections.abc import Mapping
 from numbers import Number
 from tinker import types
-from typing import List
+from typing import Any, List
 
 from twinkle import DeviceMesh
 from twinkle.template import Template
@@ -58,6 +59,33 @@ def collect_forward_backward_results(results, device_mesh: DeviceMesh):
     return [all_outputs, avg_loss]
 
 
+def to_cpu_safe_output(obj: Any) -> Any:
+    """Convert nested model outputs into CPU-safe Python objects for HTTP transport.
+
+    Recursively walks tensors, numpy arrays, mappings and sequences,
+    converting each tensor/array to a plain Python scalar or list so
+    Ray can serialise the result without requiring CUDA on the driver.
+    """
+    from twinkle.utils import torch_util
+
+    if isinstance(obj, torch.Tensor):
+        tensor = torch_util.to_local_tensor(obj).detach().cpu()
+        if tensor.numel() == 1:
+            return tensor.item()
+        return tensor.tolist()
+    if isinstance(obj, np.ndarray):
+        if obj.size == 1:
+            return obj.item()
+        return obj.tolist()
+    if isinstance(obj, np.generic):
+        return obj.item()
+    if isinstance(obj, Mapping):
+        return {key: to_cpu_safe_output(value) for key, value in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [to_cpu_safe_output(value) for value in obj]
+    return obj
+
+
 def clean_metrics(metrics: dict) -> dict:
 
     def _to_float(v):

src/twinkle/server/model/backends/megatron_model.py

@@ -4,12 +4,14 @@
 """
 import torch
 from tinker import types
-from typing import TYPE_CHECKING, Any, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
 
 from twinkle import remote_class, remote_function
+from twinkle.data_format import InputFeature, Trajectory
 from twinkle.model.megatron import MultiLoraMegatronModel
 from twinkle.server.common.datum import datum_to_input_feature, extract_rl_feature
-from twinkle.server.model.backends.common import TwinkleCompatModelBase, clean_metrics, collect_forward_backward_results
+from twinkle.server.model.backends.common import (TwinkleCompatModelBase, clean_metrics,
+                                                  collect_forward_backward_results, to_cpu_safe_output)
 
 
 @remote_class(execute='all')
@@ -112,3 +114,14 @@ def tinker_load(self, checkpoint_dir: str, **kwargs):
             return super().load(name=resolved.checkpoint_name, output_dir=resolved.checkpoint_dir, **kwargs)
         else:
             return super().load(name=resolved.checkpoint_name, **kwargs)
+
+    # ------------------------------------------------------------------
+    # Twinkle-native methods (InputFeature/Trajectory-based I/O)
+    # ------------------------------------------------------------------
+
+    @remote_function(dispatch='slice_dp', collect='mean')
+    def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]],
+                         **kwargs):
+        """Forward+backward for twinkle-native clients (InputFeature/Trajectory I/O)."""
+        output = super().forward_backward(inputs=inputs, **kwargs)
+        return to_cpu_safe_output(output)

src/twinkle/server/model/backends/transformers_model.py

@@ -6,17 +6,15 @@
 - TwinkleCompatTransformersModel: handles both tinker (Datum-based I/O) via /tinker/*
   endpoints and twinkle-native (InputFeature/Trajectory-based I/O) via /twinkle/* endpoints.
 """
-import numpy as np
-import torch
-from collections.abc import Mapping
 from tinker import types
-from typing import Any, List, Union
+from typing import List, Union
 
 from twinkle import remote_class, remote_function
 from twinkle.data_format import InputFeature, Trajectory
 from twinkle.model import MultiLoraTransformersModel
 from twinkle.server.common.datum import datum_to_input_feature, extract_rl_feature
-from twinkle.server.model.backends.common import TwinkleCompatModelBase, clean_metrics, collect_forward_backward_results
+from twinkle.server.model.backends.common import (TwinkleCompatModelBase, clean_metrics,
+                                                  collect_forward_backward_results, to_cpu_safe_output)
 
 
 @remote_class()
@@ -28,32 +26,6 @@ class TwinkleCompatTransformersModel(MultiLoraTransformersModel, TwinkleCompatMo
     - Twinkle-native I/O (InputFeature / Trajectory) via /twinkle/* endpoints.
     """
 
-    # ------------------------------------------------------------------
-    # Shared helper: CPU-safe serialisation for HTTP transport
-    # ------------------------------------------------------------------
-
-    @staticmethod
-    def _to_cpu_safe_output(obj: Any) -> Any:
-        """Convert nested outputs into CPU-safe Python objects for HTTP transport."""
-        from twinkle.utils import torch_util
-
-        if isinstance(obj, torch.Tensor):
-            tensor = torch_util.to_local_tensor(obj).detach().cpu()
-            if tensor.numel() == 1:
-                return tensor.item()
-            return tensor.tolist()
-        if isinstance(obj, np.ndarray):
-            if obj.size == 1:
-                return obj.item()
-            return obj.tolist()
-        if isinstance(obj, np.generic):
-            return obj.item()
-        if isinstance(obj, Mapping):
-            return {key: TwinkleCompatTransformersModel._to_cpu_safe_output(value) for key, value in obj.items()}
-        if isinstance(obj, (list, tuple)):
-            return [TwinkleCompatTransformersModel._to_cpu_safe_output(value) for value in obj]
-        return obj
-
     # ------------------------------------------------------------------
     # Tinker-compat methods (Datum-based I/O)
     # ------------------------------------------------------------------
@@ -135,4 +107,4 @@ def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Tr
                          **kwargs):
         """Forward+backward for twinkle-native clients (InputFeature/Trajectory I/O)."""
         output = super().forward_backward(inputs=inputs, **kwargs)
-        return self._to_cpu_safe_output(output)
+        return to_cpu_safe_output(output)