update twinkle dpo

Author: Yunnglin

cookbook/client/twinkle/self_host/dpo.py

@@ -15,10 +15,9 @@
 from peft import LoraConfig
 
 from twinkle import get_logger
-from twinkle.dataset import DatasetMeta
+from twinkle.dataset import Dataset, DatasetMeta
 from twinkle_client import init_twinkle_client
 from twinkle.dataloader import DataLoader
-from twinkle.dataset import LazyDataset
 from twinkle_client.model import MultiLoraTransformersModel
 from twinkle.loss import DPOLoss
 from twinkle.metric import DPOMetric
@@ -65,7 +64,7 @@
 
 def create_dpo_dataset():
     """Create DPO dataset with positive/negative format."""
-    dataset = LazyDataset(dataset_meta=DatasetMeta(dataset_id, data_slice=range(6000)))
+    dataset = Dataset(DatasetMeta(dataset_id, data_slice=range(600)))
     dataset.set_template('Qwen3_5Template', model_id=f'ms://{base_model}', max_length=max_length)
     dataset.map(
         EmojiDPOProcessor,
@@ -75,7 +74,7 @@ def create_dpo_dataset():
     )
     # DPO preprocessor returns {'positive': [...], 'negative': [...]}
     # batch_encode handles this format automatically
-    dataset.encode(batched=True)
+    dataset.encode()
     return dataset
 
 
@@ -179,7 +178,7 @@ def train():
         # Get reference outputs using base model (without LoRA adapter)
         # disable_lora=True tells the model to skip LoRA and use base weights
         ref_outputs = model.forward_only(inputs=dpo_batch, disable_lora=True)
-        model.forward_backward(inputs=dpo_batch, ref_outputs=ref_outputs)
+        model.forward_backward(inputs=dpo_batch, ref_outputs=ref_outputs.result)
         model.clip_grad_and_step()
 
         optim_step += 1

src/twinkle/loss/dpo.py

@@ -310,6 +310,8 @@ def __call__(
             reference_rejected_logps = ref_rejected_logps.to(device=device, dtype=dtype)
         elif ref_logps is not None:
             # Per-token reference log probs provided, need to align and sum
+            if not torch.is_tensor(ref_logps):
+                ref_logps = torch.as_tensor(ref_logps)
             ref_logps_aligned = self._align_logps(ref_logps, labels.shape, device, dtype)
             ref_chosen, ref_rejected = self._split_chosen_rejected(ref_logps_aligned)
             reference_chosen_logps = self._compute_sequence_logps(ref_chosen, chosen_labels)

src/twinkle/metric/dpo.py

@@ -50,6 +50,9 @@ def _align_logps(self, logps, target_shape, device, dtype):
             Aligned tensor with shape matching target_shape
         """
         import torch
+
+        if not torch.is_tensor(logps):
+            logps = torch.as_tensor(logps)
         logps = logps.to(device=device, dtype=dtype)
         batch_size, src_len = logps.shape
         _, target_len = target_shape

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

@@ -8,6 +8,7 @@
 
 from twinkle import remote_class, remote_function
 from twinkle.data_format import InputFeature, Trajectory
+from twinkle.infra import collect_tensor_dict
 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,
@@ -119,7 +120,13 @@ def tinker_load(self, checkpoint_dir: str, **kwargs):
     # Twinkle-native methods (InputFeature/Trajectory-based I/O)
     # ------------------------------------------------------------------
 
-    @remote_function(dispatch='slice_dp', collect='mean')
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
+    def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs):
+        """Forward-only for twinkle-native clients (InputFeature/Trajectory I/O)."""
+        output = super().forward_only(inputs=inputs, **kwargs)
+        return to_cpu_safe_output(output)
+
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
     def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]],
                          **kwargs):
         """Forward+backward for twinkle-native clients (InputFeature/Trajectory I/O)."""

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

@@ -11,6 +11,7 @@
 
 from twinkle import remote_class, remote_function
 from twinkle.data_format import InputFeature, Trajectory
+from twinkle.infra import collect_tensor_dict
 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,
@@ -106,7 +107,13 @@ def tinker_load(self, checkpoint_dir: str, **kwargs):
     # Twinkle-native methods (InputFeature/Trajectory-based I/O)
     # ------------------------------------------------------------------
 
-    @remote_function(dispatch='slice_dp', collect='mean')
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
+    def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]], **kwargs):
+        """Forward-only for twinkle-native clients (InputFeature/Trajectory I/O)."""
+        output = super().forward_only(inputs=inputs, **kwargs)
+        return to_cpu_safe_output(output)
+
+    @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
     def forward_backward(self, *, inputs: Union[InputFeature, List[InputFeature], Trajectory, List[Trajectory]],
                          **kwargs):
         """Forward+backward for twinkle-native clients (InputFeature/Trajectory I/O)."""

src/twinkle_client/common/serialize.py

@@ -2,6 +2,7 @@
 import json
 from numbers import Number
 from peft import LoraConfig
+from pydantic import BaseModel
 from typing import Any, Mapping
 
 from twinkle.dataset import DatasetMeta
@@ -56,6 +57,9 @@ def serialize_object(obj) -> str:
         }
         filtered_dict['_TWINKLE_TYPE_'] = 'LoraConfig'
         return json.dumps(filtered_dict, ensure_ascii=False)
+    elif isinstance(obj, BaseModel):
+        # Pydantic models: convert to dict for JSON serialization by requests
+        return obj.model_dump(mode='json')
     elif isinstance(obj, Mapping):
         return json.dumps(obj, ensure_ascii=False)
     elif isinstance(obj, basic_types):