support cp ,fix qwen3.5 gdn sp

cookbook/transformers/sp_fsdp_dense.py

@@ -19,9 +19,10 @@
     device_type=Platform.get_platform().device_prefix(),
 )]
 
-# FSDP + SP validation over 4 GPUs: dp=2, fsdp=2 (SP only affects input slicing)
+# FSDP + sequence-parallel validation over 4 GPUs: dp=2, fsdp=2.
+# In Transformers route, ulysses_size is the total sequence-parallel degree.
 device_mesh = DeviceMesh(
-    device_type='cuda',
+    device_type=Platform.get_platform().device_prefix(),
     mesh=np.arange(4).reshape(2, 2),
     mesh_dim_names=('dp', 'fsdp'),
     ulysses_size=2,

cookbook/transformers/sp_fsdp_dense.sh

@@ -1,5 +1,6 @@
 #!/bin/bash
-# To enabele sequence parallelism, please set ulysses_size > 1
+# To enable Transformers sequence parallelism, please set ulysses_size > 1.
+# ulysses_size is interpreted as the total sequence-parallel degree.
 # device_mesh = DeviceMesh(
 #     device_type="cuda",
 #     mesh=np.arange(4).reshape(2, 2),

src/twinkle/data_format/output.py

@@ -19,14 +19,16 @@ class ModelOutput(TypedDict, total=False):
         logits: The logits output by the model.
         loss: The loss calculated by the model.
         logps: The log-probabilities of correct tokens by the model.
+        num_tokens: The token denominator associated with ``loss``.
     """
     logits: Optional[OutputType]
     loss: Optional[OutputType]
     logps: Optional[OutputType]
+    num_tokens: Optional[OutputType]
 
 
 class LossOutput(TypedDict, total=False):
-    """The output structure for the Losses"""
+    """The output structure for the Losses."""
 
     loss: Optional[OutputType]
-    num_tokens: Optional[int]
+    num_tokens: Optional[OutputType]

src/twinkle/metric/loss.py

@@ -26,12 +26,16 @@ def accumulate(self, inputs: Union[InputFeature, List[InputFeature]], outputs: M
         loss = outputs['loss']
         loss_reduction = kwargs.get('loss_reduction', 'mean')
         if loss_reduction == 'sum':
-            if not isinstance(inputs, list):
-                inputs = [inputs]
-            for input in inputs:
-                # `Transformers` models may use reduction=sum, to average grads before step
-                labels = input['labels']
-                self.num_tokens += (labels >= 0).sum().item()
+            output_num_tokens = outputs.get('num_tokens')
+            if output_num_tokens is not None:
+                self.num_tokens += output_num_tokens.item() if hasattr(output_num_tokens, 'item') else output_num_tokens
+            else:
+                if not isinstance(inputs, list):
+                    inputs = [inputs]
+                for input in inputs:
+                    # Fallback for losses that do not expose an explicit token denominator in outputs.
+                    labels = input['labels']
+                    self.num_tokens += (labels >= 0).sum().item()
         grad_norm = kwargs.get('grad_norm')
         if grad_norm is not None:
             self.grad_norm = grad_norm